fix(qwen3): wait for async overlap prefill while idle#706
Conversation
| .async_prefill | ||
| .as_ref() | ||
| .ok_or_else(|| anyhow::anyhow!("no async prefill is in flight"))?; | ||
| let status = unsafe { cudarc::driver::sys::cuEventSynchronize(state.event) }; |
There was a problem hiding this comment.
This still does not satisfy #699's no-busy-spin acceptance criterion. The event being synchronized here is created at L494-498 with only CU_EVENT_DISABLE_TIMING; CUDA documents that cuEventSynchronize busy-waits unless the event has CU_EVENT_BLOCKING_SYNC. The fake executor's Condvar wait masks that production behavior.
There was a problem hiding this comment.
Addressed in 655f373: the event is now created with CU_EVENT_BLOCKING_SYNC | CU_EVENT_DISABLE_TIMING, so the idle-only cuEventSynchronize blocks the host thread instead of busy-waiting. The normal top-of-loop cuEventQuery path remains non-blocking.
| warn!("decode-overlap: async prefill wait failed: {error:#}"); | ||
| let pending = inflight_prefill_pending.take().unwrap(); | ||
| let targets = pending.iter().map(pending_failure_target).collect(); | ||
| fail_touched_requests( |
There was a problem hiding this comment.
wait_async_prefill can reach this branch when cuEventSynchronize fails. fail_touched_requests then calls drop_request, returning the pending requests' KV pages even though the failed synchronization did not establish that the prefill stream has stopped using them.
There was a problem hiding this comment.
Addressed in 655f373: a wait error now sends request errors and fail-stops the executor without calling drop_request, so unquiesced prefill KV pages cannot re-enter the live allocator. Added decode_overlap_wait_failure_does_not_release_inflight_kv. Validation: all 12 GitHub checks pass; on RTX 5090 / CUDA 12.8 the two decode-overlap tests pass and the full Qwen3 lib suite passes 75/75.
9acd6ec to
655f373
Compare
| | cudarc::driver::sys::CUevent_flags_enum::CU_EVENT_DISABLE_TIMING | ||
| as u32, | ||
| ); | ||
| cudarc::driver::sys::cuEventRecord(event, prefill_stream.0); |
There was a problem hiding this comment.
The blocking flag is fixed, but this event still cannot reliably represent GreenCtx prefill work. The worker is bound to the primary context, so cuEventCreate creates a primary-context event, while prefill_stream comes from cuGreenCtxStreamCreate. CUDA requires cuEventRecord’s event and stream to share a context; the call therefore fails in GreenCtx mode, but both create/record results are ignored. An unrecorded event represents empty work, so poll_async_prefill can treat it as immediately complete and fall through to the worker’s cuStreamSynchronize, bypassing the new blocking-event path. Could this use cuStreamGetGreenCtx + cuGreenCtxRecordEvent for GreenCtx, plain cuEventRecord for SharedSm, and check/clean up both failure paths? #695 already contains most of this dispatch; when combining them, the event also needs to retain CU_EVENT_BLOCKING_SYNC.
| completion_tokens: 0, | ||
| }); | ||
| } | ||
| return; |
There was a problem hiding this comment.
This avoids returning the KV pages to the live pool, but return still drops Qwen3Executor. Its shutdown makes the worker exit without resolve_inflight_prefill, which drops InflightPrefillState and its GPU buffers even though the failed wait explicitly did not establish prefill-stream quiescence. Those CudaSlice frees are enqueued on their owning primary stream, with no ordering behind the overlap stream. Could this abort on an untrustworthy synchronization failure, or otherwise guarantee that the in-flight allocations are never freed? The fake test only proves that drop_request was not called; it cannot cover the CUDA teardown race.
|
Addressed the latest review feedback in commit
Validation on the RTX 5090 / SM120 AutoDL host:
Please re-review the latest commit. |
8aba65d to
9ad98b7
Compare
|
Follow-up: the implementation commit was amended with the required DCO sign-off and force-with-lease pushed as |
9ad98b7 to
43fb4e5
Compare
|
DCO follow-up: the PR history has been rewritten so all four commits now carry valid |
| )?; | ||
|
|
||
| let mut green_ctx: sys::CUgreenCtx = ptr::null_mut(); | ||
| let get_ctx = unsafe { sys::cuStreamGetGreenCtx(stream, &raw mut green_ctx) }; |
There was a problem hiding this comment.
cuStreamGetGreenCtx still cannot reliably identify the GreenCtx streams created here. create_green obtains ctx_* with cuCtxFromGreenCtx but never makes either context current, while CUDA documents that cuGreenCtxStreamCreate ignores the caller's current context. In this sequence cuStreamGetGreenCtx therefore returns NULL, so this helper falls through to cuEventRecord for a GreenCtx stream. CUDA requires the event and stream to share a context, so the original CUDA_ERROR_INVALID_HANDLE can still occur and the cuGreenCtxRecordEvent branch is never reached. Could this use cuStreamGetCtx_v2(..., pGreenCtx) or retain/pass gctx_prefill explicitly? Please also run a real --decode-overlap green-ctx smoke test; the current fake scheduler tests do not exercise this driver path.
There was a problem hiding this comment.
Addressed in 62e6ba8. SplitConcurrent now carries the retained gctx_prefill handle from OverlapStreams into the worker. GreenCtx records completion with cuGreenCtxRecordEvent(gctx_prefill, event), while SharedSm continues to use cuEventRecord(event, prefill_stream), so no runtime ownership inference is required. Validation: release build, Qwen3 Clippy with -D warnings, and both decode-overlap regression tests pass. On an RTX 5090, a real green-ctx concurrent smoke completed 7/7 requests with repeated async-prefill unified steps and no invalid-handle, event record/query/synchronize, FATAL, or Xid errors. A separate stream-mode smoke completed 4/4 requests and covered the plain cuEventRecord branch.
| prefill_result, | ||
| ), | ||
| Err(error) => { | ||
| warn!("decode-overlap: async prefill wait failed: {error:#}"); |
There was a problem hiding this comment.
This comment no longer matches HEAD: a failed cuEventSynchronize now aborts inside wait_async_prefill, so the only Errs reachable here are the two state-mismatch cases ("no async prefill is in flight" / "completed without a result"), and the latter has already established quiescence. Worth rewording so the next reader doesn't reason about safety from an unreachable scenario — same for the regression test's "wait failure" framing, which the real executor can no longer produce.
There was a problem hiding this comment.
Addressed in 62e6ba8. The trait and scheduler comments now reflect the reachable semantics: CUDA query/synchronize failures fail-stop inside the real executor, and a returned Err is limited to post-wait state/result resolution after event completion established quiescence. The old fake wait-failure test was replaced with decode_overlap_missing_result_reports_chain_and_stops_scheduler, which simulates a completed event with a lost worker result.
| // may still be using these requests' KV pages. Fail the | ||
| // requests but retain executor ownership until the | ||
| // fail-stop teardown; returning them to the live pool | ||
| // here could race the in-flight kernels. |
There was a problem hiding this comment.
error.to_string() drops the anyhow chain, the warn above already logs {:#}, so consider format!("{error:#}") here too if the client message should carry the cause.
There was a problem hiding this comment.
Addressed in 62e6ba8. The client-facing error now uses format!("{error:#}"), and the regression test asserts that both the outer async-prefill context and the underlying fake worker-response cause are preserved. All 15 checks in CI run 29799900510 pass.
|
All commit attributions previously flagged on this pull request are resolved. |
62e6ba8 to
dd3cca5
Compare
Signed-off-by: Nyvo <75425811+Nyvo-io@users.noreply.github.com>
dd3cca5 to
6baaf3b
Compare
|
All maintainer feedback is addressed on final head
Validation is complete:
Please re-review when convenient. |
| + inflight_prefill_pending.as_ref().map_or(0, Vec::len)) as u64, | ||
| (deferred.len() + loading.len()) as u64, | ||
| ); | ||
| // Flush the prior step's cache changes to a router (no-op unless the |
There was a problem hiding this comment.
Thanks for making event-setup failure quiesce the prefill stream. There is still a wider ownership gap after the prefill launch, though: execute_decode, decode-stream synchronization, and build_batch_decode_request_results can return before the event and lane.inflight_prefill are installed. The scheduler then runs fail_touched_requests -> drop_request, which can return this Unified step's KV pages while the prefill stream still uses them. Could we guard the whole launched-but-not-transferred interval, synchronizing prefill on every early return and aborting if that synchronization fails? The later handoff in Qwen3Executor::execute_unified needs the same protection because rkv.apply_decode(...)? currently runs before self.async_prefill owns the event.
There was a problem hiding this comment.
Addressed in 7698fdd. The worker now arms a prefill-stream guard before entering the fallible forward and keeps it armed through decode launch/synchronization, decode-result construction, and completion-event creation. It disarms only after lane.inflight_prefill owns the launched work. The completion event is then an RAII owner across worker aggregation and rkv.apply_decode, and transfers only when self.async_prefill is installed. Any earlier return synchronizes before KV cleanup; an unprovable synchronization failure aborts. GPU temporaries created inside the forward are deferred until that synchronization, which also fixes the Rust unwind/drop-order gap found during the full PR review.
| }; | ||
| if record != sys::CUresult::CUDA_SUCCESS { | ||
| unsafe { | ||
| sys::cuEventDestroy_v2(event); |
There was a problem hiding this comment.
Small mismatch with the PR contract: both this destroy and the normal completion destroy in poll_async_prefill discard cuEventDestroy_v2's result, while the PR says event cleanup failures are checked. Since this path already has the recording error, could we preserve the destroy failure as context as well, and handle/log the completed-event cleanup result in poll_async_prefill? If cleanup errors are intentionally non-fatal after completion is proven, the PR text should say that explicitly.
There was a problem hiding this comment.
Addressed in 7698fdd. A record failure now returns both the record status and any simultaneous cuEventDestroy_v2 cleanup failure. Normal query/synchronize completion routes through checked destruction as well. Once completion is already proven, destroy failure is logged and intentionally non-fatal because it can leak an event resource but cannot race KV or temporary-buffer reuse; the PR description now states that policy explicitly. The new CPU regression covers the combined record/destroy error and event ownership cleanup.
Signed-off-by: Nyvo <75425811+Nyvo-io@users.noreply.github.com>
|
All latest maintainer feedback is addressed on head
Please re-review when convenient. |
|
Please rebase onto main before review. If you encounter conflicts, ensure tests are re-run as necessary. |
Summary
Why
A unified decode-overlap step can retire the last active request while its prefill continues on the overlap stream. With no active, deferred, or prefilling work left, the scheduler previously blocked only on
submit_rx.blocking_recv(), so the prefill result was not consumed until an unrelated request arrived.The first fix still left a wider ownership gap after prefill launch: decode launch/synchronization, result construction, worker aggregation, and executor-side decode application could fail before durable async-prefill state owned the GPU work. Those returns could let scheduler cleanup recycle KV pages or drop temporary allocations while the prefill stream still referenced them.
Implementation
The scheduler centralizes async-prefill result application and, when the overlap prefill is the only remaining work, calls
wait_async_prefillbefore considering the engine idle.SplitConcurrentcarries the Green Context owner retained byOverlapStreams. GreenCtx mode records completion withcuGreenCtxRecordEvent(gctx_prefill, event); shared-SM mode usescuEventRecord(event, prefill_stream). Events useCU_EVENT_BLOCKING_SYNC | CU_EVENT_DISABLE_TIMING.Stream guards are armed before fallible prefill and decode forwards. The prefill guard stays armed until
LocalQwen3Lane::inflight_prefillowns the launched work; an owned completion event then protects worker aggregation and executor-side decode application untilQwen3Executor::async_prefilltakes ownership. Every pre-handoff error synchronizes before scheduler cleanup can recycle KV, and an unprovable synchronization failure aborts.GPU-backed prefill temporaries, including errors raised inside the forward, move to a deferred-drop queue while a stream override is active. They are destroyed only after successful prefill-stream synchronization, preventing
cuMemFreeAsyncon the primary stream from racing kernels on the overlap stream.Event create/record/cleanup errors are checked. A record failure preserves a simultaneous
cuEventDestroy_v2cleanup failure in the returned context. After query or synchronize has already proven completion, a destroy failure is logged and treated as non-fatal: it may leak a driver resource, but cannot race KV or temporary-buffer reuse. Post-wait state/result mismatches retain the completeanyhowcause chain for the client.Validation
Final follow-up head
7698fddwas validated on RTX 5090 / CUDA 12.8:cargo fmt --all -- --check-D warningsopeninfer-serverbuildgreen-ctxconcurrent smoke: 7/7 requests completed, with 13 async-prefill unified steps and 13 completionsstreamsmoke: 4/4 requests completed, with 6 async-prefill unified steps and 6 completionsThe branch remains based on
origin/mainafd7a64. Final GitHub CI run29889818043passes all 15 checks on7698fdd, including formatting, DCO, attribution, locked Cargo metadata, CPU Clippy/tests, simulated frontend E2E, Qwen3 CUDA compile, and Qwen3 CUDA Clippy.Closes #699