feat(mori-io): add sync transfer wait API - #341
Merged
Conversation
YukioZzz
added a commit
to SemiAnalysisAI/InferenceX
that referenced
this pull request
Aug 12, 2026
Three fixes to the READ path, grouped because they all narrow the same window: the decoder acting on KV that has not arrived. Applied from apply_k3_moriio_patches.sh, after the trim patch they build on. 1. start_load_kv waits for the reads it posted. READ mode returns load_kv_async=False, which tells the scheduler the KV will be in place by the time the forward runs and that the blocks may be published into the prefix cache in the same step (delay_cache_blocks tracks that flag). Nothing honoured it: the reads are posted non-blocking and wait_for_layer_load is a no-op, so the forward could read blocks the NIC was still filling, and those blocks were advertised as cached at the same time. The flag itself is left alone. READ=False / WRITE=True is MoRIIO's paradigm in every copy of this connector, so flipping it would change behaviour for every model; making the existing promise true does not. 2. A request completes only when every layer's read has completed. _pop_done_transfers inspected status_list[-1], but there is one status per layer (two for a KDA layer), and _post_read_with_backoff deliberately returns a failed status when the send queue never drains. A failure on any layer but the last was therefore reported as success and the decoder ran on partial KV. 3. Failed reads are reported to the scheduler. get_block_ids_with_load_errors now surfaces the destination blocks of a failed or timed-out read, as NIXL does, so the scheduler recomputes the affected prefix instead of leaving the request to expire on a timeout. As in NIXL, the block half is skipped for hybrid models, where the block-id channel carries [attn, mamba] halves that do not map onto token-prefix recovery. The blocking wait uses mori's batched WaitAll where the build has it (ROCm/mori#341: GIL released, condition variable instead of a 1 ms spin, one shared deadline, failure winning over in-flight) and falls back to polling otherwise, probed once and cached -- an older mori raises AttributeError on the first transfer, which is how this API broke sglang's MI35x CI. The non-blocking verdict in 2 stays a Python scan even where WaitAll exists: mori's zero-timeout wait runs PollProgress on the calling thread, which would drive the backend's progress callback from the engine thread alongside its own CQ poller. Same semantics, no new concurrency, which is worth keeping while a race is still under investigation. scripts/test_moriio_wait.py covers the batch verdicts and the blocking wait against a stubbed mori, with and without WaitAll, and asserts the wait really blocks until a status turns terminal.
YukioZzz
added a commit
to YukioZzz/vllm
that referenced
this pull request
Aug 14, 2026
…le batched wait) READ mode declares load_kv_async=False, which promises the KV is in place by the time the forward runs, but nothing kept that promise for the whole batch: reads were posted non-blocking and a request was judged by a single status, so a deep model (Kimi-K3 is ~69 KDA layers plus attention, and a KDA layer posts two reads) could be called done while an earlier layer was still in flight or had already failed. Port the completion semantics sglang uses on this same backend (sglang#26922 on top of ROCm/mori#341): - TransferBatchState + MoRIIOWrapper.poll_transfer_batch give a non-blocking verdict over every status of a request. Deliberately a Python scan even where wait_all exists: mori's zero-timeout wait runs PollProgress on the calling thread, which would drive the backend's progress callback from the engine thread alongside its own CQ poller. - waiting_for_transfer_complete blocks inside mori via IOEngine.wait_all with the GIL released, sharing one deadline across the batch and giving failure precedence over still-in-flight. The 1 ms Python spin remains as _poll_transfers_until_done for builds predating mori#341, probed once and cached -- that exact API gap is what broke sglang's MI35x CI. - wait_for_layer_load uses the batched wait instead of its own spin loop, so the forward's early layers still overlap the later layers' transfers. CUDAGraphMode.FULL cannot host a blocking wait, so start_load_kv drains the step's statuses through _await_reads_issued_this_step for that case only. - A failed or timed-out read now reports its destination blocks through get_block_ids_with_load_errors, mirroring NIXL's _handle_failed_transfer, so the scheduler recomputes the affected prefix instead of running on whatever the NIC left behind. Skipped for hybrid models, as NIXL skips it under HMA. Tests: transfer-completion unit coverage over both mori generations -- batch verdicts, the availability probe, blocking until terminal, per-status detail recovered from a batch return code, and that the non-blocking poll never calls into mori. Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com>
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds a synchronous wait API for transfers, giving callers timeout control and a batch-wait helper on top of the existing
TransferStatus::Wait().TransferStatus::WaitFor(timeoutMs)— single-status wait with timeout semantics:< 0(default): wait indefinitely== 0: poll once, non-blocking> 0: wait up to the deadline (may returnIN_PROGRESSon timeout)IOEngine::WaitAll(statuses, timeoutMs)— batch wait with failure-wins precedence; empty input returnsSUCCESS. The timeout is shared across the whole batch.Wait()now delegates toWaitFor(-1)and blocks on acondition_variableinstead of busy-spinning (woken byUpdate/SetCode).TransferStatusis now non-copyable/non-movable (it owns a mutex + condvar; all usages are by-pointer or stack locals).WaitFor/wait_allexposed;Wait/WaitFor/WaitAllrelease the GIL while waiting.Behavior notes
Wait()on XGMI uses the blocking finalize callback (unchanged, efficient).WaitFor(timeout>0)on XGMI (poll-driven, no background notifier) polls once up front then re-polls at the deadline — correct result, but completion may be observed late within the timeout window. This is intentional and covered by tests.Testing
BUILD_UMBP=1 BUILD_TESTS=1 pip install .builds clean.test_transfer_wait: 21/21 pass.tests/python/io/test_transfer_wait.py: 3/3 pass.