Skip to content

fix(qwen3): wait for async overlap prefill while idle#706

Open
Nyvo-io wants to merge 2 commits into
openinfer-project:mainfrom
Nyvo-io:fix/qwen3-async-prefill-idle-699
Open

fix(qwen3): wait for async overlap prefill while idle#706
Nyvo-io wants to merge 2 commits into
openinfer-project:mainfrom
Nyvo-io:fix/qwen3-async-prefill-idle-699

Conversation

@Nyvo-io

@Nyvo-io Nyvo-io commented Jul 18, 2026

Copy link
Copy Markdown
Contributor

Summary

  • keep the scheduler from parking on new submissions while decode-overlap prefill work is still in flight
  • wait on a blocking, timing-disabled CUDA event only when no CPU or decode work can make progress
  • guard the full launched-but-not-transferred interval on both worker and executor handoffs
  • defer GPU temporary destruction until prefill-stream quiescence is proven
  • preserve fail-stop ownership when CUDA synchronization cannot prove stream quiescence

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_prefill before considering the engine idle.

SplitConcurrent carries the Green Context owner retained by OverlapStreams. GreenCtx mode records completion with cuGreenCtxRecordEvent(gctx_prefill, event); shared-SM mode uses cuEventRecord(event, prefill_stream). Events use CU_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_prefill owns the launched work; an owned completion event then protects worker aggregation and executor-side decode application until Qwen3Executor::async_prefill takes 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 cuMemFreeAsync on the primary stream from racing kernels on the overlap stream.

Event create/record/cleanup errors are checked. A record failure preserves a simultaneous cuEventDestroy_v2 cleanup 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 complete anyhow cause chain for the client.

Validation

Final follow-up head 7698fdd was validated on RTX 5090 / CUDA 12.8:

  • cargo fmt --all -- --check
  • Qwen3 Clippy with -D warnings
  • complete Qwen3 library suite: 79 passed, 0 failed
  • release openinfer-server build
  • real green-ctx concurrent smoke: 7/7 requests completed, with 13 async-prefill unified steps and 13 completions
  • real shared-SM stream smoke: 4/4 requests completed, with 6 async-prefill unified steps and 6 completions
  • no invalid-handle, event record/query/synchronize, FATAL, worker-response, or Xid errors
  • full Codex review of the complete PR plus follow-up changes: no actionable findings

The branch remains based on origin/main afd7a64. Final GitHub CI run 29889818043 passes all 15 checks on 7698fdd, including formatting, DCO, attribution, locked Cargo metadata, CPU Clippy/tests, simulated frontend E2E, Qwen3 CUDA compile, and Qwen3 CUDA Clippy.

Closes #699

Comment thread openinfer-qwen3/src/executor.rs Outdated
.async_prefill
.as_ref()
.ok_or_else(|| anyhow::anyhow!("no async prefill is in flight"))?;
let status = unsafe { cudarc::driver::sys::cuEventSynchronize(state.event) };

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread openinfer-qwen3/src/scheduler.rs Outdated
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(

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@FeathBow FeathBow self-assigned this Jul 18, 2026
@Nyvo-io
Nyvo-io force-pushed the fix/qwen3-async-prefill-idle-699 branch from 9acd6ec to 655f373 Compare July 19, 2026 06:00
@FeathBow
FeathBow self-requested a review July 19, 2026 17:19
Comment thread openinfer-qwen3/src/executor.rs Outdated
| cudarc::driver::sys::CUevent_flags_enum::CU_EVENT_DISABLE_TIMING
as u32,
);
cudarc::driver::sys::cuEventRecord(event, prefill_stream.0);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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;

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@Nyvo-io

Nyvo-io commented Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

Addressed the latest review feedback in commit 8aba65d and pushed it to the PR branch.

  • Async prefill events are now created with CU_EVENT_BLOCKING_SYNC | CU_EVENT_DISABLE_TIMING; creation, Green Context lookup, recording, and cleanup errors are checked.
  • The event helper calls cuGreenCtxRecordEvent for Green Context streams and cuEventRecord for ordinary shared-primary-context streams, so the event and stream use the same CUDA context.
  • cuEventQuery treats only CUDA_ERROR_NOT_READY as pending. Any other query/synchronize failure fail-stops the real Qwen3 executor instead of allowing Drop to free buffers whose prefill stream may still be using them.
  • If event setup fails after prefill launch, the prefill stream is synchronized before returning the error; a failed fallback synchronization aborts.

Validation on the RTX 5090 / SM120 AutoDL host:

  • cargo fmt --all -- --check
  • cargo test --release -p openinfer-qwen3 --lib decode_overlap_ (2 passed, 0 failed; the test binary compiled the new Green Context APIs)

Please re-review the latest commit.

@Nyvo-io
Nyvo-io force-pushed the fix/qwen3-async-prefill-idle-699 branch from 8aba65d to 9ad98b7 Compare July 20, 2026 02:44
@Nyvo-io

Nyvo-io commented Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

Follow-up: the implementation commit was amended with the required DCO sign-off and force-with-lease pushed as 9ad98b7 (superseding 8aba65d). The code and AutoDL validation are unchanged; the new GitHub CI run is in progress.

@Nyvo-io
Nyvo-io force-pushed the fix/qwen3-async-prefill-idle-699 branch from 9ad98b7 to 43fb4e5 Compare July 20, 2026 02:52
@Nyvo-io

Nyvo-io commented Jul 20, 2026

Copy link
Copy Markdown
Contributor Author

DCO follow-up: the PR history has been rewritten so all four commits now carry valid Signed-off-by trailers matching their authors where applicable. The new head is 43fb4e5 (superseding 9ad98b7); code content is unchanged. GitHub CI will rerun from this head.

@FeathBow
FeathBow self-requested a review July 20, 2026 18:21
Comment thread openinfer-qwen3/src/green_ctx.rs Outdated
)?;

let mut green_ctx: sys::CUgreenCtx = ptr::null_mut();
let get_ctx = unsafe { sys::cuStreamGetGreenCtx(stream, &raw mut green_ctx) };

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:#}");

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread openinfer-qwen3/src/scheduler.rs Outdated
// 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.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@github-actions

github-actions Bot commented Jul 21, 2026

Copy link
Copy Markdown

All commit attributions previously flagged on this pull request are resolved.

@Nyvo-io
Nyvo-io force-pushed the fix/qwen3-async-prefill-idle-699 branch from 62e6ba8 to dd3cca5 Compare July 21, 2026 18:00
Signed-off-by: Nyvo <75425811+Nyvo-io@users.noreply.github.com>
@Nyvo-io
Nyvo-io force-pushed the fix/qwen3-async-prefill-idle-699 branch from dd3cca5 to 6baaf3b Compare July 21, 2026 18:02
@Nyvo-io

Nyvo-io commented Jul 21, 2026

Copy link
Copy Markdown
Contributor Author

All maintainer feedback is addressed on final head 6baaf3b.

  • The idle-only wait uses a blocking, timing-disabled CUDA event. GreenCtx ownership is retained and passed explicitly to cuGreenCtxRecordEvent; SharedSm uses cuEventRecord. Event setup/cleanup is checked, and untrustworthy CUDA query/synchronize failures remain fail-stop so in-flight KV and temporary buffers cannot be released unsafely.
  • Scheduler comments and the regression test now match the reachable post-wait state/result-mismatch semantics, and the client error preserves the complete anyhow cause chain.
  • The branch is rebased onto origin/main@afd7a64 and condensed to one Nyvo-authored, Nyvo-committed, DCO-signed commit while retaining the upstream device-ordinal, batched-logprob, TP checkpoint-prefetch, and rustfmt changes.

Validation is complete:

  • final GitHub CI run 29855508474: 15/15 checks passed on 6baaf3b
  • RTX 5090 / CUDA 12.8: release build, Qwen3 Clippy with -D warnings, both deterministic scheduler regressions, and the full Qwen3 library suite passed
  • real decode-overlap smoke: green-ctx 7/7 requests and stream 4/4 requests completed with repeated async-prefill unified steps and no invalid-handle, CUDA event, FATAL, Xid, or worker-response errors

Please re-review when convenient.

@FeathBow
FeathBow self-requested a review July 21, 2026 21:33
+ 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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread openinfer-qwen3/src/green_ctx.rs Outdated
};
if record != sys::CUresult::CUDA_SUCCESS {
unsafe {
sys::cuEventDestroy_v2(event);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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>
@Nyvo-io

Nyvo-io commented Jul 22, 2026

Copy link
Copy Markdown
Contributor Author

All latest maintainer feedback is addressed on head 7698fdd after a full PR review.

  • The entire launched-but-not-transferred interval is guarded on the worker and executor handoffs, including fallible forward internals and Rust drop ordering.
  • Event record/destroy double failures preserve both statuses; completed-event destroy failures are checked, logged, and explicitly documented as non-fatal only after quiescence is proven.
  • Final RTX 5090 validation passed: format, Qwen3 Clippy -D warnings, 79/79 Qwen3 lib tests, release server build, GreenCtx 7/7 with 13 paired async-prefill steps, and shared-SM stream 4/4 with 6 paired steps; no event/FATAL/worker/Xid errors.
  • Full Codex review of the complete PR plus follow-up changes reported no actionable findings.
  • GitHub CI run 29889818043 passes all 15 checks.

Please re-review when convenient.

@FeathBow

Copy link
Copy Markdown
Collaborator

Please rebase onto main before review. If you encounter conflicts, ensure tests are re-run as necessary.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

qwen3: scheduler can park on blocking_recv while an async decode-overlap prefill is still in flight

2 participants