hipfile: feat(hipfile/python): add async stream I/O bindings

[jiejingzhangamd]

Expose the hipFile async API to the Python bindings. The C API (hipFileReadAsync / hipFileWriteAsync / hipFileStreamRegister / hipFileStreamDeregister) already exists; only the Python layer was missing -- the bindings previously had just the hipFileAsyncNotSupported enum.

• _chipfile.pxd: declare hipStream_t + the four async functions.
• _hipfile.pyx: AsyncIOHandle cdef class that owns the in/out C slots (size / file_offset / buffer_offset / bytes_done) so their addresses stay valid until the stream completes. The driver dereferences these pointers AFTER the async call returns (the transfer finishes on stream sync), so a plain Cython stack-local passed by address would be written back to a dead address -- bytes_done reads 0 and later submits can hit hipFileInvalidValue. Plus read/write async wrappers, stream register/deregister, and a supports_async() probe.
• file.py: FileHandle.read_async / write_async returning an AsyncIOHandle (keep alive past stream sync, then read bytes_done), a Stream context manager, and supports_async().
• init.py: export Stream and supports_async.

Mirrors the cuFile async pattern. Validated on gfx942 / ROCm 7.2: build + import, supports_async() == True, and a GPU async write->read round-trip (bytes_done == 4096, data byte-identical).

Motivation

Technical Details

JIRA ID

AIHIPFILE-171

Test Plan

Test Result

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[jiejingzhangamd]

Why we need it
We run a KV-cache L3 offload connector for vLLM/SGLang on MI300X: paged attention KV blocks are spilled to / loaded from an L3 tier (local NVMe or NFS) through hipFile GPU-direct. On a cache hit, a decode/prefill step has to pull hundreds of MB — multiple GB of KV from disk into the GPU KV-cache blocks before the attention kernel reads them.

With the synchronous hipFileRead, each chunk's disk→GPU load blocks the calling thread and the GPU sits idle for the duration of the I/O — the load is serialized against compute. At serving scale that I/O wait dominates the step (we saw multi-second per-step H2D stalls and a pacing-bound throughput plateau).

The fix is the standard GPUDirect-Storage overlap pattern: submit the read on a side stream and let it run concurrently with the model forward pass, gating each layer's attention on just its chunk via a HIP event — exactly what cuFileReadAsync enables on NVIDIA. The hipFile C async API already does this; only the Python layer was missing. This PR adds it so our Python connector can use the async path instead of a ctypes shim.

Usage — layerwise load overlapped with compute (our actual pattern)

import os, torch, hipfile

assert hipfile.supports_async()          # fall back to sync read() if False

# One-time setup: a dedicated copy stream + a registered GPU staging buffer.
copy_stream = torch.cuda.Stream()
gbuf = torch.empty(CHUNK_BYTES, dtype=torch.uint8, device="cuda")
buf = hipfile.Buffer(gbuf.data_ptr(), gbuf.numel(), 0); buf.register()

with hipfile.FileHandle(path, os.O_RDONLY | os.O_DIRECT) as fh, \
     hipfile.Stream(copy_stream.cuda_stream) as st:        # register stream once

    # 1) Submit every layer's chunk read on the copy stream. Each call returns
    #    immediately; the transfer runs concurrently with the compute below.
    inflight = []
    with torch.cuda.stream(copy_stream):
        for layer, (file_off, size, dst_off) in enumerate(load_plan):
            io = fh.read_async(buf, size, file_off, dst_off, st.handle)
            ev = torch.cuda.Event(); ev.record(copy_stream)
            inflight.append((layer, io, ev))   # MUST keep `io` alive (see note)

    # 2) Compute overlaps the I/O: each layer waits only for its own chunk.
    for layer, io, ev in inflight:
        torch.cuda.current_stream().wait_event(ev)   # gate attn on this read
        run_attention(layer, gbuf)                   # reads freshly-loaded KV
        assert io.bytes_done == expected_size[layer] # valid after the stream sync
write_async is symmetric (used on the save path to spill KV to L3 without blocking the forward pass).

The one correctness gotcha worth a binding (why AsyncIOHandle exists)
hipFileReadAsync takes pointer args (size_p, file_offset_p, buffer_offset_p, bytes_read_p) that the driver dereferences after the call returns — the transfer completes when the stream is synced.

Following is how it fixed:

The slots are C members of a cdef class, so their addresses are stable for the whole lifetime of the Python object — not stack-locals that die when the wrapper returns:

cdef class AsyncIOHandle:
    cdef size_t       _size
    cdef _c.hoff_t    _file_off
    cdef _c.hoff_t    _buf_off
    cdef ssize_t      _bytes_done          # driver writes here, later, on the stream
    def __cinit__(self, size_t size, _c.hoff_t file_offset, _c.hoff_t buffer_offset):
        self._size = size; self._file_off = file_offset
        self._buf_off = buffer_offset; self._bytes_done = 0
    @property
    def bytes_done(self):
        return self._bytes_done
def hipFileReadAsync(uintptr_t handle, uintptr_t buffer_base,
                     AsyncIOHandle io not None, uintptr_t stream_handle):
    cdef size_t    *size_p = &io._size      # &member of a live Python object
    cdef _c.hoff_t *foff_p = &io._file_off
    cdef _c.hoff_t *boff_p = &io._buf_off
    cdef ssize_t   *done_p = &io._bytes_done
    with nogil:
        e = _c.hipFileReadAsync(<_c.hipFileHandle_t>handle, <void*>buffer_base,
                                size_p, foff_p, boff_p, done_p,
                                <_c.hipStream_t>stream_handle)
    ...

So the driver's deferred write to bytes_read_p lands in a valid address as long as the AsyncIOHandle is alive.

Validated (gfx942 / ROCm 7.2)
Built from source + installed; import hipfile, supports_async() == True, and a GPU async write→read round-trip: bytes_done == 4096, data byte-identical.

[Copilot]

Pull request overview

This PR adds Python bindings for hipFile’s asynchronous stream-attached I/O API so Python users can submit hipFileReadAsync/hipFileWriteAsync operations and manage stream registration from the Python layer (mirroring cuFile’s async pattern).

Changes:

• Add Cython declarations and wrappers for hipFileReadAsync, hipFileWriteAsync, hipFileStreamRegister, and hipFileStreamDeregister, plus an async support probe.
• Introduce AsyncIOHandle to keep async in/out slot storage alive across stream execution.
• Add high-level Python FileHandle.read_async/write_async, Stream registration context manager, and export Stream/supports_async from the package.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 9 comments.

File	Description
projects/hipfile/python/hipfile/file.py	Adds high-level async read/write APIs and a stream registration helper class.
projects/hipfile/python/hipfile/_hipfile.pyx	Adds Cython async wrappers, AsyncIOHandle, stream register/deregister wrappers, and an async capability probe.
projects/hipfile/python/hipfile/_chipfile.pxd	Declares hipStream_t and async/stream APIs in the Cython pxd layer.
projects/hipfile/python/hipfile/init.py	Exports Stream and supports_async at the package top level.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[riley-dixon]

Hi @jiejingzhangamd ,
Thank you for the contribution to the hipFile Python bindings!
There are a couple of things I just want to mention to make sure you are aware of:

1. The Async API only supports the POSIX fallback path at this moment in time. Async support for the GPU optimized path is planned for in the future.

2. The long-term plan is to migrate these Python bindings into HIP-Python. When that happens, the internals of AsyncIOHandle will be reworked to use HIP-Python Pointer objects for storage rather than the current Cython implementation. I believe though the public API contract for read_async and write_async could be maintained.

3. If this sounds reasonable to you, I can update the hipfile PyPI package after these changes have been merged.

Please let me know your thoughts!

[jiejingzhangamd]

Thanks for the review! Pushed fixes:

• errno only captured for POSIX/C errors (err == -1) in the async read/write/register/deregister wrappers; other codes leave extra=0.
• read_async/write_async raise OSError on err == -1 (consistent with read()/write()).
• supports_async() deregisters the default stream if the probe registered it (no leaked registration).
• Stream.deregister() clears registered only after a successful deregister.
• Added size/file_offset/buffer_offset setters on AsyncIOHandle; fixed the docstring (no Stream.synchronize).

On the roadmap — the HIP-Python migration and you updating the PyPI package after merge both sound good to me.

[therock-pr-bot]

❌ PR Check — Action Required

Check	Status	Details
🌿 Branch Name	❌ Fail	Branch name does not match allowed patterns. Branch: feat/hipfile-python-async-bindings Allowed patterns: - ^users\/[A-Za-z0-9][A-Za-z0-9\-]*\/.+ - ^shared\/.+ - ^[A-Za-z0-9][A-Za-z0-9\-_]*$ - ^dependabot\/.+ - ^revert-[0-9]+-.+
📝 PR Title/Description	❌ Fail	Error: Title does not follow Conventional Commits style. Expected: start with a valid type (feat, fix, docs, …). Desired format: type(optional-scope): short description ─── Error: PR description must reference a JIRA ID, ISSUE ID, or a GitHub closing keyword. Expected: include a JIRA ID / ISSUE ID line (separator : or -, or omitted; value may be a JIRA key, a number with/without #, or a link), OR a closing keyword + issue reference. Accepted examples: • JIRA ID : TESTAUTO-6039 • JIRA ID - #330 • JIRA ID #330 • ISSUE ID : TESTUTO-3334 • ISSUE ID #3334 • ISSUE ID - TESTAUTO-3433 • ISSUE ID : https://github.com/<org_name>/<repo_name>/issues/1234 • Closes #10 • Fixes octo-org/octo-repo#100 • Resolves: #123 • #123 • https://github.com/<org_name>/<repo_name>/issues/123 Current: no valid JIRA/ISSUE/closing-keyword reference found
⛔ Forbidden Files	✅ Pass	—
🧪 Unit Test	❌ Fail	Error: Source/code files changed without an accompanying unit test. Expected: add at least one test file named like test_<name>.py / test_<name>.cpp (or <name>_test.*). Current: code file(s) changed: projects/hipfile/python/hipfile/__init__.py, projects/hipfile/python/hipfile/file.py; no test file found
🔎 pre-commit	⏳ Pending	⏳ Still running…
🚫 Draft PR	🔜 To Be Enabled	—
🚩 Feature Flag	🔜 To Be Enabled	—
📊 Code Coverage	🔜 To Be Enabled	—

⚠️ 3 policy check(s) failed. Please address the issues above before this PR can be Reviewed.

🚫 Please fix the failed policies

• ❌ Branch Name
• ❌ PR Title/Description
• ❌ Unit Test

The Not ready to Review label was added to this PR. Once all policies pass, the label is removed automatically.

📖 Need help? See the Policy FAQ for details on every check and how to fix failures.

[therock-pr-bot]

🚫 Please fix the failed policies before requesting reviews.

The following policy checks failed:

• ❌ Branch Name
• ❌ PR Title/Description
• ❌ Unit Test

The Not ready to Review label has been added to this PR.
Once all policies pass, the label will be removed automatically.

[riley-dixon]

Hi @jiejingzhangamd - just want to let you know that the hipfile package on PyPI has been updated to 0.4.0.dev0 which includes your async changes!