diff --git a/docs/models/qwen3/green-ctx-sm-partition.md b/docs/models/qwen3/green-ctx-sm-partition.md index 763b59c5..083451f9 100644 --- a/docs/models/qwen3/green-ctx-sm-partition.md +++ b/docs/models/qwen3/green-ctx-sm-partition.md @@ -63,11 +63,13 @@ Power climbs with QPS: at QPS 8 the board still has headroom (draw oscillates ~5 ## Pitfalls -- **The Xid 31/43 hit during bring-up was a cross-stream buffer use-after-free; the supported success and checked-`Result` paths are fixed — not an open driver risk.** Prefill temporaries are allocated on `ctx.stream` but consumed by override-stream kernels; the fix is a producer fence (`green_ctx::fence_producers_before_override`, ordering each `ctx.stream` alloc/H2D ahead of its override consumer) plus RAII stream drains that hold each stream until its kernels finish before any buffer frees or KV page is released — `PrefillTempBin` (prefill) and `DecodeStreamGuard` (decode), each fail-closed (a sync failure aborts). An arbitrary panic that unwinds past a not-yet-parked prefill buffer is out of scope. +- **The Xid 31/43 hit during bring-up was a cross-stream buffer use-after-free; the supported success and checked-`Result` paths are fixed — not an open driver risk.** Prefill temporaries are allocated on `ctx.stream` but consumed by override-stream kernels; the fix is a producer fence (`green_ctx::fence_producers_before_override`, ordering each `ctx.stream` alloc/H2D ahead of its override consumer) plus fail-closed RAII ownership. `LaunchedStreamGuard` drains launched prefill/decode streams before any pre-handoff return can recycle KV, `DeferredDrop` parks GPU-backed prefill temporaries until prefill-stream quiescence is proven, and `AsyncPrefillEvent` owns the launched prefill across the worker-to-executor handoff. An arbitrary panic that unwinds past an unarmed launch guard is out of scope. - **The "`cuGreenCtxStreamCreate` itself faults on driver 590 + >16 GB resident" theory did not reproduce.** This 5090 (driver 590.48.01, ~27 GB resident — 19.3 GB of it KV) ran green-ctx clean across both 2026-06-19 sweeps and the 2026-06-20 three-mode smoke (split path fired, 0 Xid). - **`gemm_lt` is still disabled under the stream override** (`5af4fd5`, to avoid the cuBLASLt workspace Xid-31 path). So split-path decode keeps its CUDA graph but loses the per-shape Lt tuning — a remaining decode-side lever, not yet re-measured under the partition. - **A single request never exercises the split path** — smoke-test with concurrent load or you are only testing the full-SM graph. +- **Record async completion with the retained Green Context handle.** A Green Context stream does not reliably identify its owner through `cuStreamGetGreenCtx` when the derived context was never made current. The split command therefore carries `gctx_prefill` from `OverlapStreams` and records its blocking, timing-disabled event with `cuGreenCtxRecordEvent`; the shared-primary-context path still uses `cuEventRecord`. This keeps idle scheduler waits blocking and prevents an invalid or unrecorded event from bypassing prefill completion. +- **GPU launch ownership starts before a fallible forward returns.** `SplitConcurrent` arms stream-quiescence guards before both prefill and decode launches because either call may submit kernels before reporting an error. GPU-backed prefill temporaries move to a thread-local deferred-drop queue on both success and error, including errors raised inside the forward; only a successful prefill-stream synchronization drains that queue. The prefill guard remains armed through decode synchronization, decode-result construction, and completion-event setup; it disarms only after `LocalQwen3Lane::inflight_prefill` owns the launched work. The completion event then remains an armed RAII owner across worker aggregation and executor-side decode application, so any pre-handoff return synchronizes before the scheduler can recycle KV pages. A synchronization failure is fail-stop. Once query/synchronize has already proven completion, event-destroy failure is logged but intentionally non-fatal: it may leak a driver resource, but cannot race KV or temporary-buffer reuse. - **`pkill` from an ssh one-liner matches its own command line** — use `pkill -f "[t]arget/release/openinfer"`, and kill/launch in separate ssh invocations. - Build on the 5090 with `CUDA_HOME=/usr/local/cuda-13.1` (stale `/usr/local/cuda` → cuBLAS 12.9 N=1025 cliff; see `serving-perf-5090.md`). Verify `ldd target/release/openinfer | grep cublas` shows `.so.13`. diff --git a/openinfer-qwen3/src/executor.rs b/openinfer-qwen3/src/executor.rs index 71dc97c8..022f0d0a 100644 --- a/openinfer-qwen3/src/executor.rs +++ b/openinfer-qwen3/src/executor.rs @@ -505,6 +505,7 @@ fn execute_step_on_lane( decode_requests, decode_kv_views, prefill_stream, + prefill_green_ctx, decode_stream, sample_seed, } => { @@ -529,75 +530,77 @@ fn execute_step_on_lane( .map(|req| req.lora_adapter.as_deref()) .collect(); - // Declared before the bin so that on any early return (`?` or panic) - // the bin drops first: its Drop synchronizes the prefill stream before - // `prefill_logits` and the parked temporaries the bin owns are freed. - let prefill_logits; - let mut prefill_temp_bin = crate::prefill::PrefillTempBin::armed(prefill_stream.0); - { + // Sync ctx.stream: ensures all prior stream-ordered allocs and + // H2D copies are complete before green streams touch them. + lane.model.device_ctx().sync()?; + + // Arm before entering execute_prefill: a failed forward may still + // have submitted kernels. Until lane.inflight_prefill owns the + // launched work, every return path must first quiesce this stream. + let mut launched_prefill = LaunchedStreamGuard::new_prefill(prefill_stream.0); + + // Launch prefill on prefill partition stream. + let prefill_logits = { let _prefill_override = unsafe { StreamOverrideGuard::activate(prefill_stream.0) }; - let (logits, _, _) = lane.execute_prefill( + let (prefill_logits, _, _) = lane.execute_prefill( &prefill_prompts, prefill_kv_views, &prefill_lora_adapters, false, None, )?; - prefill_logits = logits; - } - // Close the prefill parking window before decode, so decode's own - // temporaries don't land in it. - prefill_temp_bin.close(); + crate::prefill::DeferredDrop::new(prefill_logits) + }; + // Launch decode on the decode partition stream. CUDA graph stays + // enabled: batch_decode captures into the split graph cache keyed on + // the active stream override, so the replayed kernel nodes stay + // pinned to the decode SM partition (CUDA PG §4.6.5 — capture stream + // determines a node's execution context). + let mut launched_decode = LaunchedStreamGuard::new(decode_stream.0, "decode"); { - let _decode_guard = DecodeStreamGuard { - stream: decode_stream.0, - }; let _decode_override = unsafe { StreamOverrideGuard::activate(decode_stream.0) }; lane.execute_decode(&decode_tokens, decode_kv_views, &decode_lora_adapters)?; } - let decode_result = - build_batch_decode_request_results(lane, decode_requests, *sample_seed)?; - - let event = lane - .model - .device_ctx() - .ctx - .new_event(None) - .map_err(|e| anyhow::anyhow!("cuEventCreate(prefill poll) failed: {e}"))?; - unsafe { - // The prefill stream may be a Green Context stream; cuEventRecord - // needs the event and stream in one context, so record via the - // green context when the stream has one (stream mode has none). - let mut gctx: cudarc::driver::sys::CUgreenCtx = std::ptr::null_mut(); - let get = cudarc::driver::sys::cuStreamGetGreenCtx(prefill_stream.0, &raw mut gctx); - let record = if get != cudarc::driver::sys::CUresult::CUDA_SUCCESS { - get - } else if gctx.is_null() { - cudarc::driver::sys::cuEventRecord(event.cu_event(), prefill_stream.0) - } else { - cudarc::driver::sys::cuGreenCtxRecordEvent(gctx, event.cu_event()) - }; - anyhow::ensure!( - record == cudarc::driver::sys::CUresult::CUDA_SUCCESS, - "recording prefill poll event failed: {record:?}" - ); - } - - lane.inflight_prefill = Some(InflightPrefillState { - temp_bin: prefill_temp_bin, - prefill_logits, - prefill_requests: prefill_requests.clone(), - sample_seed: *sample_seed, - }); + // Only sync decode stream — decode result is ready for sampling. + // Prefill continues async on GPU; polled later via event. + launched_decode.quiesce(); - Ok(WorkerStepOutcome::SplitDecodeReady { - decode: DecodeResult { - requests: decode_result, - }, - prefill_event: event, - }) + if collect_result { + // Sample decode tokens immediately. + let decode_result = + build_batch_decode_request_results(lane, decode_requests, *sample_seed)?; + + // Record after the prefill launches. Green Context streams need + // cuGreenCtxRecordEvent; plain overlap streams use cuEventRecord. + // Both creation and recording are checked before the event is + // handed to the scheduler. + let event = + crate::green_ctx::record_stream_event(prefill_stream.0, *prefill_green_ctx)?; + let prefill_event = AsyncPrefillEvent::new(event); + let prefill_requests = prefill_requests.clone(); + + // Store prefill state for deferred sync+sample. + lane.inflight_prefill = Some(InflightPrefillState { + prefill_stream: prefill_stream.0, + prefill_logits: prefill_logits.into_inner(), + prefill_requests, + sample_seed: *sample_seed, + }); + launched_prefill.disarm(); + + Ok(WorkerStepOutcome::SplitDecodeReady { + decode: DecodeResult { + requests: decode_result, + }, + prefill_event, + }) + } else { + // Non-primary worker: still need to sync prefill before returning. + launched_prefill.quiesce(); + Ok(WorkerStepOutcome::Ack) + } } StepCommand::SpeculativeVerify { requests, @@ -955,6 +958,16 @@ pub(crate) trait ModelExecutor: Send { None } + /// Block until the async prefill completes and return its result. The + /// scheduler calls this only when no other work can make progress, so an + /// event wait is preferable to polling the CUDA event in a tight loop. + /// CUDA synchronization failures fail-stop inside the real executor. A + /// returned error therefore means post-wait state/result resolution failed + /// after the completion event established stream quiescence. + fn wait_async_prefill(&mut self) -> Result { + anyhow::bail!("async prefill wait is not implemented for this executor") + } + // ── KV block-event feed (no-op unless built with the event feed on) ── /// Take the raw block-event receiver, once. `None` unless the engine was @@ -1025,7 +1038,7 @@ pub struct Qwen3Executor { overlap: Option, /// In-flight async prefill state. Populated by the SplitConcurrent step, /// consumed by `poll_async_prefill`. - async_prefill: Option, + async_prefill: Option, /// DFlash draft metadata; `Some` once a draft model is loaded into the /// primary lane. Speculative decoding is enabled iff this is set. speculative: Option, @@ -1056,6 +1069,219 @@ struct ExecutorKvEvents { pending_dropped: Vec>, } +/// State for an in-flight async prefill on the prefill overlap stream. +struct AsyncPrefillState { + event: AsyncPrefillEvent, +} + +// SAFETY: AsyncPrefillState is only accessed from the single executor/scheduler +// thread that owns the GPU context. The raw CUevent pointer is not shared. +unsafe impl Send for AsyncPrefillState {} + +/// Owns an event from worker creation through executor-side completion. +/// +/// Dropping an event before the executor installs it still establishes GPU +/// quiescence. This covers worker aggregation failures and every early return +/// in the worker-to-executor handoff without relying on callers to remember a +/// cleanup branch. +struct AsyncPrefillEvent { + raw: cudarc::driver::sys::CUevent, + #[cfg(test)] + synchronize_override: Option cudarc::driver::sys::CUresult>, + #[cfg(test)] + destroy_override: Option cudarc::driver::sys::CUresult>, +} + +impl AsyncPrefillEvent { + fn new(raw: cudarc::driver::sys::CUevent) -> Self { + Self { + raw, + #[cfg(test)] + synchronize_override: None, + #[cfg(test)] + destroy_override: None, + } + } + + #[cfg(test)] + fn with_operations( + raw: cudarc::driver::sys::CUevent, + synchronize: fn(cudarc::driver::sys::CUevent) -> cudarc::driver::sys::CUresult, + destroy: fn(cudarc::driver::sys::CUevent) -> cudarc::driver::sys::CUresult, + ) -> Self { + Self { + raw, + synchronize_override: Some(synchronize), + destroy_override: Some(destroy), + } + } + + fn raw(&self) -> cudarc::driver::sys::CUevent { + self.raw + } + + fn synchronize(&self) -> cudarc::driver::sys::CUresult { + #[cfg(test)] + if let Some(synchronize) = self.synchronize_override { + return synchronize(self.raw); + } + unsafe { cudarc::driver::sys::cuEventSynchronize(self.raw) } + } + + #[cfg(test)] + fn destroy(&self, event: cudarc::driver::sys::CUevent) -> cudarc::driver::sys::CUresult { + if let Some(destroy) = self.destroy_override { + return destroy(event); + } + unsafe { cudarc::driver::sys::cuEventDestroy_v2(event) } + } + + /// Destroy an event after query/synchronize already proved completion. + fn destroy_completed(mut self, operation: &'static str) { + let event = std::mem::replace(&mut self.raw, std::ptr::null_mut()); + #[cfg(test)] + let status = self.destroy(event); + #[cfg(not(test))] + let status = unsafe { cudarc::driver::sys::cuEventDestroy_v2(event) }; + handle_completed_prefill_event_destroy(status, operation); + } +} + +// SAFETY: The event is created on a worker thread and then moved through a +// channel to the executor thread on the same device. Access is sequential. +unsafe impl Send for AsyncPrefillEvent {} + +impl Drop for AsyncPrefillEvent { + fn drop(&mut self) { + if self.raw.is_null() { + return; + } + let status = self.synchronize(); + if status != cudarc::driver::sys::CUresult::CUDA_SUCCESS { + log::error!( + "FATAL: an untransferred async prefill event could not establish stream \ + quiescence ({status:?}); aborting" + ); + std::process::abort(); + } + #[cfg(test)] + let destroy = self.destroy(self.raw); + #[cfg(not(test))] + let destroy = unsafe { cudarc::driver::sys::cuEventDestroy_v2(self.raw) }; + handle_completed_prefill_event_destroy(destroy, "untransferred async prefill"); + self.raw = std::ptr::null_mut(); + } +} + +/// CUDA calls may submit work before returning an error. Keep this guard armed +/// across each fallible launch interval. Prefill disarms only after +/// `LocalQwen3Lane::inflight_prefill` owns it; decode explicitly quiesces before +/// its result can be applied. A prefill synchronization also drains temporaries +/// whose destruction was deferred by the launched forward. +struct LaunchedStreamGuard { + stream: cudarc::driver::sys::CUstream, + operation: &'static str, + armed: bool, + drain_prefill_deferred_drops: bool, + #[cfg(test)] + synchronize_override: + Option cudarc::driver::sys::CUresult>, +} + +impl LaunchedStreamGuard { + fn new(stream: cudarc::driver::sys::CUstream, operation: &'static str) -> Self { + Self { + stream, + operation, + armed: true, + drain_prefill_deferred_drops: false, + #[cfg(test)] + synchronize_override: None, + } + } + + fn new_prefill(stream: cudarc::driver::sys::CUstream) -> Self { + Self { + stream, + operation: "prefill", + armed: true, + drain_prefill_deferred_drops: true, + #[cfg(test)] + synchronize_override: None, + } + } + + #[cfg(test)] + fn with_synchronizer( + stream: cudarc::driver::sys::CUstream, + operation: &'static str, + drain_prefill_deferred_drops: bool, + synchronize: fn(cudarc::driver::sys::CUstream) -> cudarc::driver::sys::CUresult, + ) -> Self { + Self { + stream, + operation, + armed: true, + drain_prefill_deferred_drops, + synchronize_override: Some(synchronize), + } + } + + fn disarm(&mut self) { + self.armed = false; + } + + fn quiesce(&mut self) { + self.synchronize_or_abort(); + self.disarm(); + } + + fn synchronize_or_abort(&self) { + #[cfg(test)] + let status = if let Some(synchronize) = self.synchronize_override { + synchronize(self.stream) + } else { + unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream) } + }; + #[cfg(not(test))] + let status = unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream) }; + + if status != cudarc::driver::sys::CUresult::CUDA_SUCCESS { + let operation = self.operation; + log::error!( + "FATAL: launched {operation} work was not safely handed off and \ + cuStreamSynchronize({operation}) failed ({status:?}); aborting" + ); + std::process::abort(); + } + if self.drain_prefill_deferred_drops { + crate::prefill::drain_deferred_drops(); + } + } +} + +impl Drop for LaunchedStreamGuard { + fn drop(&mut self) { + if self.armed { + self.synchronize_or_abort(); + } + } +} + +fn handle_completed_prefill_event_destroy( + status: cudarc::driver::sys::CUresult, + operation: &'static str, +) { + if status != cudarc::driver::sys::CUresult::CUDA_SUCCESS { + // Completion is already proven, so this cannot race KV/buffer reuse. + // Keep serving while making the driver-resource leak observable. + log::error!( + "cuEventDestroy_v2({operation}) failed after async prefill completion \ + was proven ({status:?}); continuing" + ); + } +} + /// One request's in-flight CPU-tier KV prefetch. /// /// `probe` holds the GPU-hit prefix resident for the request's whole parked @@ -1662,6 +1888,21 @@ impl Qwen3Executor { Ok(()) } + /// Resolve a prefill after a fallible decode-side handoff step failed. + /// + /// The scheduler will drop the requests touched by the failed step, but the + /// worker still owns the prefill buffers until its stream is synchronized + /// and its in-flight state is consumed. Preserve the original execution + /// error while making cleanup failure visible as additional context. + fn resolve_async_prefill_after_error(&mut self, error: anyhow::Error) -> anyhow::Error { + match ::wait_async_prefill(self) { + Ok(_) => error, + Err(cleanup_error) => error.context(format!( + "failed to clean up async prefill after execution error: {cleanup_error:#}" + )), + } + } + /// vLLM-style `--no-prefix-cache`. Behaviour depends on whether offload is /// active: /// * **No offload** — classic: disable prefix matching outright, so every @@ -2696,12 +2937,13 @@ impl ModelExecutor for Qwen3Executor { } fn execute_unified(&mut self, plan: UnifiedPlan<'_>) -> Result { - // The scheduler resolves any prior async prefill before this step; a - // pending one here is a broken contract. Checked before any KV commit. - anyhow::ensure!( - self.async_prefill.is_none(), - "async prefill invariant violated: a previous prefill is still pending at a new unified step" - ); + // Scheduler admission normally prevents overlap here. Resolve first if + // a direct caller reaches this method with an older prefill in flight, + // before a new worker step can replace its worker-side state. + if self.async_prefill.is_some() { + self.wait_async_prefill() + .context("resolve previous async prefill before unified step")?; + } // Low-level callers can bypass the startup guard; LoRA prefill scratch is // unordered across ctx.stream and the overlap stream. anyhow::ensure!( @@ -2750,6 +2992,7 @@ impl ModelExecutor for Qwen3Executor { decode_requests: plan.decode_requests.to_vec(), decode_kv_views, prefill_stream: overlap.prefill_stream, + prefill_green_ctx: overlap.prefill_green_context(), decode_stream: overlap.decode_stream, sample_seed: plan.sample_seed, } @@ -2800,43 +3043,32 @@ impl ModelExecutor for Qwen3Executor { } WorkerStepOutcome::SplitDecodeReady { decode: decode_result, - prefill_event: event, + prefill_event, } => { - // Prefill may still be in flight (the `event` signals completion), - // writing this step's KV pages. An error return releases those pages, - // so sync the event first; success leaves the prefill in flight. - let applied = - decode_result - .requests - .iter() - .try_for_each(|req_result| -> Result<()> { - let rkv = self - .request_kvs - .get_mut(&req_result.request_id) - .expect("request must exist after split decode"); - rkv.apply_decode(req_result.token, self.kv_mgr.pool())?; - Ok(()) - }); - if let Err(e) = applied { - let sync = unsafe { cudarc::driver::sys::cuEventSynchronize(event.cu_event()) }; - if sync != cudarc::driver::sys::CUresult::CUDA_SUCCESS { - log::error!( - "FATAL: cuEventSynchronize(prefill) failed on the decode-error \ - path ({sync:?}); aborting to avoid releasing KV pages the \ - prefill stream may still be writing" - ); - std::process::abort(); + // SM-partition path: decode done, prefill still in-flight. + // Establish executor ownership before any fallible decode + // update. If an update fails, resolve_async_prefill_after_error + // synchronizes the event and clears the worker-side state before + // the scheduler recycles the touched requests. + self.async_prefill = Some(AsyncPrefillState { + event: prefill_event, + }); + let apply_result = (|| -> Result<()> { + for req_result in &decode_result.requests { + let rkv = self + .request_kvs + .get_mut(&req_result.request_id) + .expect("request must exist after split decode"); + rkv.apply_decode(req_result.token, self.kv_mgr.pool())?; } - let rx = self.primary.resolve_prefill()?; - rx.recv().map_err(|_| { - anyhow::anyhow!("worker dropped resolve_prefill on decode-error path") - })??; - return Err(e); + Ok(()) + })(); + if let Err(error) = apply_result { + return Err(self.resolve_async_prefill_after_error(error)); } for req_result in &decode_result.requests { self.save_sealed_blocks(req_result.request_id); } - self.async_prefill = Some(event); // Return a UnifiedResult with empty prefill — scheduler will // get prefill results via poll_async_prefill. Ok(UnifiedResult { @@ -3041,8 +3273,9 @@ impl ModelExecutor for Qwen3Executor { } fn poll_async_prefill(&mut self) -> Option { - let status = - unsafe { cudarc::driver::sys::cuEventQuery(self.async_prefill.as_ref()?.cu_event()) }; + let state = self.async_prefill.as_ref()?; + // Non-blocking check: is the prefill stream done? + let status = unsafe { cudarc::driver::sys::cuEventQuery(state.event.raw()) }; match status { cudarc::driver::sys::CUresult::CUDA_ERROR_NOT_READY => return None, cudarc::driver::sys::CUresult::CUDA_SUCCESS => {} @@ -3054,13 +3287,14 @@ impl ModelExecutor for Qwen3Executor { std::process::abort(); } } - let _ = self.async_prefill.take(); - let rx = self.primary.resolve_prefill().unwrap_or_else(|e| { - log::error!( - "FATAL: resolve_prefill failed after the async prefill completed ({e}); aborting" - ); - std::process::abort(); - }); + // Prefill is done — resolve it via the worker. + let event = self.async_prefill.take().unwrap().event; + event.destroy_completed("completed async prefill"); + + // Ask worker to sync + sample the prefill result. + let Ok(rx) = self.primary.resolve_prefill() else { + return None; + }; let result = match rx.recv() { Ok(Ok(result)) => result, Ok(Err(e)) => { @@ -3083,13 +3317,159 @@ impl ModelExecutor for Qwen3Executor { } Some(result) } + + fn wait_async_prefill(&mut self) -> Result { + let state = self + .async_prefill + .as_ref() + .ok_or_else(|| anyhow::anyhow!("no async prefill is in flight"))?; + let status = unsafe { cudarc::driver::sys::cuEventSynchronize(state.event.raw()) }; + if status != cudarc::driver::sys::CUresult::CUDA_SUCCESS { + // The event is the only proof that the prefill stream stopped using + // its temporary buffers. Returning an error would immediately drop + // the executor and race those frees, so fail-stop the process. + log::error!("FATAL: cuEventSynchronize(async prefill) failed ({status:?}); aborting"); + std::process::abort(); + } + + // `poll_async_prefill` performs the event cleanup and worker-side + // result resolution after the event is known to be complete. + self.poll_async_prefill() + .ok_or_else(|| anyhow::anyhow!("async prefill completed without a result")) + } } #[cfg(test)] mod tests { use std::collections::HashSet; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use super::AsyncPrefillEvent; + use super::LaunchedStreamGuard; use super::ensure_lora_capacity; + use crate::prefill::DeferredDrop; + + static STREAM_SYNCHRONIZATIONS: AtomicUsize = AtomicUsize::new(0); + static EVENT_SYNCHRONIZATIONS: AtomicUsize = AtomicUsize::new(0); + static EVENT_DESTROYS: AtomicUsize = AtomicUsize::new(0); + static DEFERRED_DROP_ORDER: AtomicUsize = AtomicUsize::new(0); + + struct DropAfterPrefillSync; + + impl Drop for DropAfterPrefillSync { + fn drop(&mut self) { + assert_eq!(DEFERRED_DROP_ORDER.load(Ordering::SeqCst), 1); + DEFERRED_DROP_ORDER.store(2, Ordering::SeqCst); + } + } + + fn successful_stream_synchronize( + _stream: cudarc::driver::sys::CUstream, + ) -> cudarc::driver::sys::CUresult { + STREAM_SYNCHRONIZATIONS.fetch_add(1, Ordering::SeqCst); + cudarc::driver::sys::CUresult::CUDA_SUCCESS + } + + fn successful_event_synchronize( + _event: cudarc::driver::sys::CUevent, + ) -> cudarc::driver::sys::CUresult { + EVENT_SYNCHRONIZATIONS.fetch_add(1, Ordering::SeqCst); + cudarc::driver::sys::CUresult::CUDA_SUCCESS + } + + fn successful_event_destroy( + _event: cudarc::driver::sys::CUevent, + ) -> cudarc::driver::sys::CUresult { + EVENT_DESTROYS.fetch_add(1, Ordering::SeqCst); + cudarc::driver::sys::CUresult::CUDA_SUCCESS + } + + fn ordered_stream_synchronize( + _stream: cudarc::driver::sys::CUstream, + ) -> cudarc::driver::sys::CUresult { + DEFERRED_DROP_ORDER + .compare_exchange(0, 1, Ordering::SeqCst, Ordering::SeqCst) + .expect("prefill stream must synchronize before deferred resources drop"); + cudarc::driver::sys::CUresult::CUDA_SUCCESS + } + + fn return_after_launch() -> anyhow::Result<()> { + let _guard = LaunchedStreamGuard::with_synchronizer( + std::ptr::null_mut(), + "prefill", + true, + successful_stream_synchronize, + ); + anyhow::bail!("injected post-launch failure") + } + + #[test] + fn async_prefill_ownership_guards_cover_early_return_and_handoff() { + STREAM_SYNCHRONIZATIONS.store(0, Ordering::SeqCst); + EVENT_SYNCHRONIZATIONS.store(0, Ordering::SeqCst); + EVENT_DESTROYS.store(0, Ordering::SeqCst); + + assert!(return_after_launch().is_err()); + assert_eq!(STREAM_SYNCHRONIZATIONS.load(Ordering::SeqCst), 1); + + { + let mut guard = LaunchedStreamGuard::with_synchronizer( + std::ptr::null_mut(), + "prefill", + true, + successful_stream_synchronize, + ); + guard.quiesce(); + } + assert_eq!(STREAM_SYNCHRONIZATIONS.load(Ordering::SeqCst), 2); + + { + let mut guard = LaunchedStreamGuard::with_synchronizer( + std::ptr::null_mut(), + "prefill", + true, + successful_stream_synchronize, + ); + guard.disarm(); + } + assert_eq!(STREAM_SYNCHRONIZATIONS.load(Ordering::SeqCst), 2); + + let fake_event: cudarc::driver::sys::CUevent = std::ptr::dangling_mut(); + drop(AsyncPrefillEvent::with_operations( + fake_event, + successful_event_synchronize, + successful_event_destroy, + )); + assert_eq!(EVENT_SYNCHRONIZATIONS.load(Ordering::SeqCst), 1); + assert_eq!(EVENT_DESTROYS.load(Ordering::SeqCst), 1); + + AsyncPrefillEvent::with_operations( + fake_event, + successful_event_synchronize, + successful_event_destroy, + ) + .destroy_completed("test completion"); + assert_eq!(EVENT_SYNCHRONIZATIONS.load(Ordering::SeqCst), 1); + assert_eq!(EVENT_DESTROYS.load(Ordering::SeqCst), 2); + + DEFERRED_DROP_ORDER.store(0, Ordering::SeqCst); + let guard = LaunchedStreamGuard::with_synchronizer( + std::ptr::null_mut(), + "prefill", + true, + ordered_stream_synchronize, + ); + { + let _override = unsafe { + openinfer_kernels::tensor::StreamOverrideGuard::activate(std::ptr::null_mut()) + }; + drop(DeferredDrop::new(DropAfterPrefillSync)); + } + assert_eq!(DEFERRED_DROP_ORDER.load(Ordering::SeqCst), 0); + drop(guard); + assert_eq!(DEFERRED_DROP_ORDER.load(Ordering::SeqCst), 2); + } #[test] fn lora_capacity_rejects_new_adapter_at_limit() { @@ -3125,6 +3505,9 @@ mod tests { impl Drop for Qwen3Executor { fn drop(&mut self) { + // Establish prefill quiescence before worker shutdown drops the lane's + // inflight state and its GPU buffers. + drop(self.async_prefill.take()); self.primary.shutdown(); for worker in &mut self.workers { worker.shutdown(); @@ -3181,8 +3564,7 @@ struct LocalQwen3Lane { /// Stored state for an async prefill that was launched but not yet synced. struct InflightPrefillState { - /// Field order is load-bearing: `temp_bin` must drop before `prefill_logits`. - temp_bin: crate::prefill::PrefillTempBin, + prefill_stream: cudarc::driver::sys::CUstream, prefill_logits: HiddenStates, prefill_requests: Vec, /// Per-step sampling seed captured when the prefill was launched, replayed @@ -3194,25 +3576,6 @@ struct InflightPrefillState { // owns the GPU context. It is never shared across threads. unsafe impl Send for InflightPrefillState {} -/// Drains the decode stream on scope exit (success, `?`, or panic) before this -/// step's KV pages can be released; a failed drain aborts (fail-closed). -struct DecodeStreamGuard { - stream: cudarc::driver::sys::CUstream, -} - -impl Drop for DecodeStreamGuard { - fn drop(&mut self) { - let r = unsafe { cudarc::driver::sys::cuStreamSynchronize(self.stream) }; - if r != cudarc::driver::sys::CUresult::CUDA_SUCCESS { - log::error!( - "FATAL: cuStreamSynchronize(decode) failed ({r:?}); aborting rather than \ - release KV pages the decode stream may still be reading" - ); - std::process::abort(); - } - } -} - impl LocalQwen3Lane { fn new( model: Qwen3Model, @@ -3369,12 +3732,22 @@ impl LocalQwen3Lane { /// Sync the in-flight prefill stream and sample prefill tokens. fn resolve_inflight_prefill(&mut self) -> Result { - let mut state = self + let state = self .inflight_prefill .take() .ok_or_else(|| anyhow::anyhow!("no inflight prefill to resolve"))?; - state.temp_bin.synchronize(); + // Sync prefill stream + let r = unsafe { cudarc::driver::sys::cuStreamSynchronize(state.prefill_stream) }; + if r != cudarc::driver::sys::CUresult::CUDA_SUCCESS { + // The state owns GPU buffers that may still be referenced by the + // stream. Returning would drop them without proof of quiescence. + log::error!("FATAL: cuStreamSynchronize(inflight prefill) failed ({r:?}); aborting"); + std::process::abort(); + } + + // Now safe to drop deferred GPU buffers (prefill kernels are done). + crate::prefill::drain_deferred_drops(); // Sample prefill tokens let params: Vec<&SamplingParams> = @@ -3608,6 +3981,7 @@ enum StepCommand { decode_requests: Vec, decode_kv_views: Vec, prefill_stream: crate::green_ctx::SendStream, + prefill_green_ctx: Option, decode_stream: crate::green_ctx::SendStream, sample_seed: u64, }, @@ -3721,7 +4095,7 @@ enum WorkerStepOutcome { decode: DecodeResult, /// Event recorded on prefill stream after all prefill kernels; /// query this to check if prefill is done without blocking. - prefill_event: cudarc::driver::CudaEvent, + prefill_event: AsyncPrefillEvent, }, SpeculativeVerify(VerifyResult), SpeculativeDraft(DraftResult), diff --git a/openinfer-qwen3/src/green_ctx.rs b/openinfer-qwen3/src/green_ctx.rs index 8073eeea..86025b6c 100644 --- a/openinfer-qwen3/src/green_ctx.rs +++ b/openinfer-qwen3/src/green_ctx.rs @@ -79,6 +79,15 @@ pub(crate) fn fence_producers_before_override( Ok(()) } +/// A `CUgreenCtx` wrapper that can cross the worker command channel. +/// The owning [`OverlapStreams`] keeps the context alive while a copied handle +/// is used to record the split prefill completion event. +#[derive(Clone, Copy, Debug)] +#[repr(transparent)] +pub(crate) struct SendGreenContext(pub sys::CUgreenCtx); + +unsafe impl Send for SendGreenContext {} + /// Two CUDA streams used to overlap prefill and decode within one scheduler /// step. In [`DecodeOverlap::GreenCtx`] mode they are pinned to disjoint SM /// partitions via Green Contexts; in [`DecodeOverlap::SharedSm`] mode they are @@ -138,7 +147,57 @@ fn create_primary_stream() -> Result { Ok(stream) } +/// Record an event after work submitted to `stream`. +/// +/// A Green Context stream must pass its owning context explicitly. Deriving it +/// from the stream is unreliable because `cuGreenCtxStreamCreate` ignores the +/// caller's current context, so stream context queries can report no Green +/// Context even for an SM-pinned stream. +pub(crate) fn record_stream_event( + stream: CUstream, + green_ctx: Option, +) -> Result { + let mut event: sys::CUevent = ptr::null_mut(); + check_cu( + unsafe { + sys::cuEventCreate( + &raw mut event, + sys::CUevent_flags_enum::CU_EVENT_BLOCKING_SYNC as u32 + | sys::CUevent_flags_enum::CU_EVENT_DISABLE_TIMING as u32, + ) + }, + "cuEventCreate (async prefill)", + )?; + + let record = match green_ctx { + Some(green_ctx) => unsafe { sys::cuGreenCtxRecordEvent(green_ctx.0, event) }, + None => unsafe { sys::cuEventRecord(event, stream) }, + }; + if record != sys::CUresult::CUDA_SUCCESS { + let destroy = unsafe { sys::cuEventDestroy_v2(event) }; + return Err(async_prefill_record_error(record, destroy)); + } + Ok(event) +} + +fn async_prefill_record_error(record: sys::CUresult, destroy: sys::CUresult) -> anyhow::Error { + if destroy == sys::CUresult::CUDA_SUCCESS { + anyhow::anyhow!("recording async prefill event failed: {record:?}") + } else { + anyhow::anyhow!( + "recording async prefill event failed: {record:?}; \ + cuEventDestroy_v2 cleanup also failed: {destroy:?}" + ) + } +} + impl OverlapStreams { + pub(crate) fn prefill_green_context(&self) -> Option { + self.green + .as_ref() + .map(|green| SendGreenContext(green.gctx_prefill)) + } + /// Set up the overlap streams for the given device, or `None` when overlap /// is [`DecodeOverlap::Off`] (the executor keeps its single stream). pub(crate) fn create(device_ordinal: usize, overlap: DecodeOverlap) -> Result> { @@ -349,3 +408,21 @@ impl Drop for OverlapStreams { // SAFETY: OverlapStreams is only used from the executor's single GPU worker thread. unsafe impl Send for OverlapStreams {} + +#[cfg(test)] +mod tests { + use super::async_prefill_record_error; + use super::sys; + + #[test] + fn async_prefill_record_error_preserves_cleanup_failure() { + let error = async_prefill_record_error( + sys::CUresult::CUDA_ERROR_INVALID_HANDLE, + sys::CUresult::CUDA_ERROR_DEINITIALIZED, + ) + .to_string(); + + assert!(error.contains("recording async prefill event failed: CUDA_ERROR_INVALID_HANDLE")); + assert!(error.contains("cuEventDestroy_v2 cleanup also failed: CUDA_ERROR_DEINITIALIZED")); + } +} diff --git a/openinfer-qwen3/src/lora.rs b/openinfer-qwen3/src/lora.rs index e4cdcf9f..1ad32dd3 100644 --- a/openinfer-qwen3/src/lora.rs +++ b/openinfer-qwen3/src/lora.rs @@ -165,6 +165,16 @@ pub(crate) struct DeviceLoraTokenGroup<'a> { pub(crate) token_indices_d: Option>, } +impl Drop for DeviceLoraTokenGroup<'_> { + fn drop(&mut self) { + if openinfer_kernels::tensor::has_stream_override() + && let Some(token_indices_d) = self.token_indices_d.take() + { + crate::prefill::defer_drop(token_indices_d); + } + } +} + pub(crate) fn build_lora_token_ranges<'a>( seq_lens: impl IntoIterator, adapters: impl IntoIterator>, @@ -397,9 +407,11 @@ pub(crate) fn apply_lora_projection_delta_range( if token_len == 0 { return Ok(()); } - let mut rank_out = HiddenStates::zeros(ctx, projection.a.rows, token_len)?; + let mut rank_out = + crate::prefill::DeferredDrop::new(HiddenStates::zeros(ctx, projection.a.rows, token_len)?); ops::gemm_token_range_into_checked(ctx, &projection.a, input, token_offset, &mut rank_out)?; - let mut delta = HiddenStates::zeros(ctx, projection.b.rows, token_len)?; + let mut delta = + crate::prefill::DeferredDrop::new(HiddenStates::zeros(ctx, projection.b.rows, token_len)?); ops::gemm_into_checked(ctx, &projection.b, &rank_out, &mut delta)?; ops::scaled_add_rows_token_range_into(ctx, &delta, scale, out, row_offset, token_offset) } @@ -417,11 +429,21 @@ pub(crate) fn apply_lora_projection_delta_indexed( if token_count == 0 { return Ok(()); } - let mut compact_input = HiddenStates::zeros(ctx, input.hidden_dim, token_count)?; + let mut compact_input = + crate::prefill::DeferredDrop::new(HiddenStates::zeros(ctx, input.hidden_dim, token_count)?); ops::gather_hidden_tokens_into(ctx, input, token_indices_d, token_count, &mut compact_input)?; - let mut rank_out = HiddenStates::zeros(ctx, projection.a.rows, token_count)?; + + let mut rank_out = crate::prefill::DeferredDrop::new(HiddenStates::zeros( + ctx, + projection.a.rows, + token_count, + )?); ops::gemm_into_checked(ctx, &projection.a, &compact_input, &mut rank_out)?; - let mut delta = HiddenStates::zeros(ctx, projection.b.rows, token_count)?; + let mut delta = crate::prefill::DeferredDrop::new(HiddenStates::zeros( + ctx, + projection.b.rows, + token_count, + )?); ops::gemm_into_checked(ctx, &projection.b, &rank_out, &mut delta)?; ops::scaled_add_rows_indexed_into( ctx, diff --git a/openinfer-qwen3/src/prefill.rs b/openinfer-qwen3/src/prefill.rs index 70aa7520..f04b7ff2 100644 --- a/openinfer-qwen3/src/prefill.rs +++ b/openinfer-qwen3/src/prefill.rs @@ -1,4 +1,6 @@ use std::any::Any; +use std::ops::Deref; +use std::ops::DerefMut; use anyhow::Result; use cudarc::driver::CudaSlice; @@ -14,66 +16,67 @@ use crate::lora::DeviceLoraTokenGroup; use crate::lora::build_lora_token_ranges; use crate::lora::prepare_lora_token_groups; -// Prefill temporaries free on ctx.stream but are consumed by override-stream -// kernels; `park()` defers them into the open parking window until that -// stream syncs. No window open (non-overlap) → drops in place. +// Thread-local deferred-drop queue for decode-overlap mode. Buffers pushed here +// during prefill (under stream override) are dropped later when +// `drain_deferred_drops()` runs after the prefill stream is synchronized. thread_local! { - static PREFILL_TEMP_WINDOW: std::cell::RefCell>>> = - const { std::cell::RefCell::new(None) }; + static DEFERRED_DROPS: std::cell::RefCell>> = + std::cell::RefCell::new(Vec::new()); } -fn park(val: T) { - PREFILL_TEMP_WINDOW.with(|w| { - if let Some(items) = w.borrow_mut().as_mut() { - items.push(Box::new(val)); - } - }); +/// Defer an object's drop until `drain_deferred_drops()` is called. +pub(crate) fn defer_drop(val: T) { + DEFERRED_DROPS.with(|q| q.borrow_mut().push(Box::new(val))); } -/// Owns parked prefill temporaries until the override stream is synchronized. -pub(crate) struct PrefillTempBin { - stream: Option, - items: Vec>, +/// Owns a GPU temporary that must outlive work submitted through a stream +/// override. On the ordinary model stream it behaves like a plain value. +pub(crate) struct DeferredDrop { + value: Option, + defer: bool, } -impl PrefillTempBin { - pub(crate) fn armed(stream: cudarc::driver::sys::CUstream) -> Self { - PREFILL_TEMP_WINDOW.with(|w| *w.borrow_mut() = Some(Vec::new())); +impl DeferredDrop { + pub(crate) fn new(value: T) -> Self { Self { - stream: Some(stream), - items: Vec::new(), + value: Some(value), + defer: openinfer_kernels::tensor::has_stream_override(), } } - pub(crate) fn close(&mut self) { - if let Some(mut captured) = PREFILL_TEMP_WINDOW.with(|w| w.borrow_mut().take()) { - self.items.append(&mut captured); - } + pub(crate) fn into_inner(mut self) -> T { + self.value.take().expect("deferred value already taken") } +} - /// Drains the prefill stream; aborts if synchronization fails. - pub(crate) fn synchronize(&mut self) { - self.close(); - if let Some(stream) = self.stream { - let r = unsafe { cudarc::driver::sys::cuStreamSynchronize(stream) }; - if r != cudarc::driver::sys::CUresult::CUDA_SUCCESS { - log::error!( - "FATAL: cuStreamSynchronize(prefill) failed ({r:?}); aborting rather than \ - free buffers the prefill stream may still be reading" - ); - std::process::abort(); - } - self.stream = None; - self.items.clear(); - } +impl Deref for DeferredDrop { + type Target = T; + + fn deref(&self) -> &Self::Target { + self.value.as_ref().expect("deferred value already taken") } } -impl Drop for PrefillTempBin { +impl DerefMut for DeferredDrop { + fn deref_mut(&mut self) -> &mut Self::Target { + self.value.as_mut().expect("deferred value already taken") + } +} + +impl Drop for DeferredDrop { fn drop(&mut self) { - self.synchronize(); + if self.defer + && let Some(value) = self.value.take() + { + defer_drop(value); + } } } + +/// Drop all deferred objects. Call after prefill stream sync. +pub(crate) fn drain_deferred_drops() { + DEFERRED_DROPS.with(|q| q.borrow_mut().clear()); +} use openinfer_core::tensor::DeviceContext; use openinfer_core::tensor::HiddenStates; use openinfer_kv_cache::KvView; @@ -145,23 +148,18 @@ impl Qwen3Model { let seq_len = token_ids.len(); let hidden_dim = self.config.hidden_size; - let token_ids_gpu = self - .ctx - .stream - .clone_htod(token_ids) - .map_err(|e| anyhow::anyhow!("H2D copy failed: {}", e))?; + // Copy token IDs to GPU + let token_ids_gpu = DeferredDrop::new( + self.ctx + .stream + .clone_htod(token_ids) + .map_err(|e| anyhow::anyhow!("H2D copy failed: {}", e))?, + ); - let mut out = HiddenStates::zeros(&self.ctx, hidden_dim, seq_len)?; + let mut out = DeferredDrop::new(HiddenStates::zeros(&self.ctx, hidden_dim, seq_len)?); crate::green_ctx::fence_producers_before_override(&self.ctx)?; - let launched = - ops::embedding_batch(&self.ctx, &self.embed_tokens, &token_ids_gpu, &mut out); - park(token_ids_gpu); - if let Err(e) = launched { - park(out); - return Err(e); - } - - Ok(out) + ops::embedding_batch(&self.ctx, &self.embed_tokens, &token_ids_gpu, &mut out)?; + Ok(out.into_inner()) } /// Embed a device-resident token buffer into a pre-allocated output, with @@ -422,8 +420,16 @@ impl Qwen3Model { /// Used when `echo=true` to return prompt token log-probabilities. /// Applies final RMS norm + lm_head projection in a single batched GEMM. /// Returns `HiddenStates` with shape `[vocab_size, total_tokens]`. - fn compute_all_position_logits(&self, hidden: &HiddenStates) -> Result { - let mut normed = HiddenStates::zeros(&self.ctx, hidden.hidden_dim, hidden.seq_len)?; + pub(crate) fn compute_all_position_logits( + &self, + hidden: &HiddenStates, + ) -> Result { + let mut normed = DeferredDrop::new(HiddenStates::zeros( + &self.ctx, + hidden.hidden_dim, + hidden.seq_len, + )?); + crate::green_ctx::fence_producers_before_override(&self.ctx)?; ops::rms_norm_batch_into( &self.ctx, hidden, @@ -431,7 +437,13 @@ impl Qwen3Model { self.config.rms_norm_eps, &mut normed, ); - ops::gemm(&self.ctx, self.output_projection(), &normed) + let mut logits = DeferredDrop::new(HiddenStates::zeros( + &self.ctx, + self.output_projection().rows, + hidden.seq_len, + )?); + ops::gemm_into_checked(&self.ctx, self.output_projection(), &normed, &mut logits)?; + Ok(logits.into_inner()) } /// Batched last-token logits: gather the given token columns out of @@ -443,21 +455,11 @@ impl Qwen3Model { token_indices: &[i32], ) -> Result { let n = token_indices.len(); - // Allocate all buffers up front so one producer fence orders them ahead - // of the override-stream gather/norm/GEMM. - let indices_d = self.ctx.stream.clone_htod(token_indices)?; - let mut gathered = HiddenStates::zeros(&self.ctx, hidden.hidden_dim, n)?; - let mut normed = HiddenStates::zeros(&self.ctx, hidden.hidden_dim, n)?; - let mut logits = HiddenStates::zeros(&self.ctx, self.output_projection().rows, n)?; + let indices_d = DeferredDrop::new(self.ctx.stream.clone_htod(token_indices)?); + let mut gathered = DeferredDrop::new(HiddenStates::zeros(&self.ctx, hidden.hidden_dim, n)?); crate::green_ctx::fence_producers_before_override(&self.ctx)?; - - let gather = - ops::gather_hidden_tokens_into(&self.ctx, hidden, &indices_d, n, &mut gathered); - park(indices_d); - if let Err(e) = gather { - park(gathered); - return Err(e); - } + ops::gather_hidden_tokens_into(&self.ctx, hidden, &indices_d, n, &mut gathered)?; + let mut normed = DeferredDrop::new(HiddenStates::zeros(&self.ctx, hidden.hidden_dim, n)?); ops::rms_norm_batch_into( &self.ctx, &gathered, @@ -465,15 +467,13 @@ impl Qwen3Model { self.config.rms_norm_eps, &mut normed, ); - let gemm = - ops::gemm_into_checked(&self.ctx, self.output_projection(), &normed, &mut logits); - park(gathered); - park(normed); - if let Err(e) = gemm { - park(logits); - return Err(e); - } - Ok(logits) + let mut logits = DeferredDrop::new(HiddenStates::zeros( + &self.ctx, + self.output_projection().rows, + n, + )?); + ops::gemm_into_checked(&self.ctx, self.output_projection(), &normed, &mut logits)?; + Ok(logits.into_inner()) } /// Concatenates all prompts' tokens, runs one GEMM per layer for the @@ -507,7 +507,7 @@ impl Qwen3Model { let seq_lens: Vec = prompts.iter().map(|p| p.len()).collect(); let lora_ranges = build_lora_token_ranges(seq_lens.iter().copied(), lora_adapters.iter().copied()); - let mut lora_groups = prepare_lora_token_groups(&self.ctx, &lora_ranges)?; + let lora_groups = prepare_lora_token_groups(&self.ctx, &lora_ranges)?; let start_positions: Vec = kv_views .iter() .zip(prompts.iter()) @@ -522,7 +522,7 @@ impl Qwen3Model { .iter() .map(openinfer_kv_cache::KvView::last_page_len) .collect(); - let plan = PrefillPagedPlan::from_raw_batch_with_cta_tile_q( + let plan = DeferredDrop::new(PrefillPagedPlan::from_raw_batch_with_cta_tile_q( &self.ctx, &page_indices, &last_page_lens, @@ -532,58 +532,60 @@ impl Qwen3Model { self.local_num_key_value_heads(), self.config.head_dim, PREFILL_ATTENTION_CTA_TILE_Q, - )?; + )?); let all_tokens: Vec = prompts.iter().flat_map(|p| p.iter().copied()).collect(); - let mut hidden = self.get_embeddings_batch(&all_tokens)?; - - // Failed launches may still leave kernels reading these inputs. - let result = self - .process_all_layers_batch_multi( - &mut hidden, - layout, - kv_buffer, - &plan, - &lora_groups, - capture_layer_ids, - ) - .and_then(|captured_hidden| { - let all_logits = if echo { - Some(self.compute_all_position_logits(&hidden)?) - } else { - None - }; - - let mut last_indices = Vec::with_capacity(batch_size); - let mut offset = 0usize; - for &seq_len in &seq_lens { - last_indices.push((offset + seq_len - 1) as i32); - offset += seq_len; - } - let logits = self.batch_token_logits(&hidden, &last_indices)?; - Ok((logits, all_logits, captured_hidden)) - }); + let hidden = DeferredDrop::new(self.get_embeddings_batch(&all_tokens)?); - for group in &mut lora_groups { - if let Some(indices) = group.token_indices_d.take() { - park(indices); - } + // Forward through all layers + let (hidden, captured_hidden) = self.process_all_layers_batch_multi( + hidden.into_inner(), + layout, + kv_buffer, + &plan, + &lora_groups, + capture_layer_ids, + )?; + let hidden = DeferredDrop::new(hidden); + let captured_hidden = captured_hidden.map(DeferredDrop::new); + + // All-position logits for echo (before we extract last-token logits) + let all_logits = if echo { + Some(DeferredDrop::new( + self.compute_all_position_logits(&hidden)?, + )) + } else { + None + }; + + let mut last_indices = Vec::with_capacity(batch_size); + let mut offset = 0usize; + for &seq_len in &seq_lens { + last_indices.push((offset + seq_len - 1) as i32); + offset += seq_len; } - park(hidden); - park(plan); + let logits = DeferredDrop::new(self.batch_token_logits(&hidden, &last_indices)?); - result + Ok(( + logits.into_inner(), + all_logits.map(DeferredDrop::into_inner), + captured_hidden.map(DeferredDrop::into_inner), + )) } fn process_all_layers_batch_multi( &self, - hidden: &mut HiddenStates, + hidden: HiddenStates, layout: &KvLayout, kv_buffer: &cudarc::driver::CudaSlice, plan: &PrefillPagedPlan, lora_groups: &[DeviceLoraTokenGroup<'_>], capture_layer_ids: Option<&[usize]>, - ) -> Result> { + ) -> Result<(HiddenStates, Option)> { + // Install ownership before any validation/allocation can return. If a + // later launch fails, these values move to the deferred queue while the + // caller's armed prefill guard establishes stream quiescence. + let mut hidden = DeferredDrop::new(hidden); let total_tokens = hidden.seq_len; let inter_dim = self.local_intermediate_size(); let q_dim = self.local_q_dim(); @@ -603,22 +605,22 @@ impl Qwen3Model { let mut captured_hidden = if capture_layer_ids.is_empty() { None } else { - Some(HiddenStates::zeros( + Some(DeferredDrop::new(HiddenStates::zeros( &self.ctx, self.config.hidden_size * capture_layer_ids.len(), total_tokens, - )?) + )?)) }; let mut next_capture = 0usize; - let mut bufs = PrefillBuffers::new( + let mut bufs = DeferredDrop::new(PrefillBuffers::new( &self.ctx, self.config.hidden_size, q_dim, kv_dim, inter_dim, total_tokens, - )?; + )?); crate::green_ctx::fence_producers_before_override(&self.ctx)?; @@ -627,7 +629,7 @@ impl Qwen3Model { self.forward_layer_batch_paged( layer_idx, layer, - hidden, + &mut hidden, kv_buffer, layout, plan, @@ -640,7 +642,7 @@ impl Qwen3Model { .expect("capture buffer exists when ids are non-empty"); ops::copy_hidden_rows_into( &self.ctx, - hidden, + &hidden, out, next_capture * self.config.hidden_size, )?; @@ -649,9 +651,11 @@ impl Qwen3Model { } Ok(()) })(); - park(bufs); run?; - Ok(captured_hidden) + Ok(( + hidden.into_inner(), + captured_hidden.map(DeferredDrop::into_inner), + )) } } diff --git a/openinfer-qwen3/src/scheduler.rs b/openinfer-qwen3/src/scheduler.rs index 14c15484..2e4d3d21 100644 --- a/openinfer-qwen3/src/scheduler.rs +++ b/openinfer-qwen3/src/scheduler.rs @@ -47,6 +47,7 @@ use self::resolve::resolve_step; use crate::Qwen3LoraOptions; use crate::Qwen3OffloadOptions; use crate::executor::ModelExecutor; +use crate::executor::PrefillResult; use crate::executor::Qwen3Executor; use crate::executor::RequestId; use crate::weights::Qwen3MemoryOptions; @@ -524,6 +525,34 @@ fn publish_load( }); } +/// Apply a completed decode-overlap prefill. The pending requests are kept in +/// the scheduler until this point because their first token and KV state are +/// produced by the async prefill stream. +fn apply_async_prefill_result( + executor: &mut E, + active: &mut Vec, + prefilling: &mut Vec, + tracker: &mut phase_trace::PhaseTracker, + inflight: &mut Option>, + result: PrefillResult, +) { + let pending = inflight + .take() + .expect("async prefill result without pending requests"); + info!( + "decode-overlap: async prefill completed ({} reqs)", + pending.len() + ); + let scheduled_at_unix_s = openinfer_core::engine::unix_now_s(); + let artifacts = ExecutionArtifacts::Prefill { + pending, + result, + scheduled_at_unix_s, + }; + let effects = resolve_step(&*executor, active, artifacts); + apply_effects(executor, active, prefilling, tracker, effects); +} + fn scheduler_loop( mut executor: E, mut submit_rx: mpsc::UnboundedReceiver, @@ -578,24 +607,13 @@ fn scheduler_loop( // 0. Poll in-flight async prefill (decode-overlap mode). if inflight_prefill_pending.is_some() { if let Some(prefill_result) = executor.poll_async_prefill() { - let pending = inflight_prefill_pending.take().unwrap(); - info!( - "decode-overlap: async prefill completed ({} reqs)", - pending.len() - ); - let scheduled_at_unix_s = openinfer_core::engine::unix_now_s(); - let artifacts = ExecutionArtifacts::Prefill { - pending, - result: prefill_result, - scheduled_at_unix_s, - }; - let effects = resolve_step(&executor, &active, artifacts); - apply_effects( + apply_async_prefill_result( &mut executor, &mut active, &mut prefilling, &mut tracker, - effects, + &mut inflight_prefill_pending, + prefill_result, ); } } @@ -606,14 +624,50 @@ fn scheduler_loop( } // 2. Reclaim settled prefetches, then offer fresh requests to prefetch. - let reserve_floor = admitted_future_blocks(&executor, &active, &prefilling); + let reserve_floor = admitted_future_blocks_with_inflight( + &executor, + &active, + &prefilling, + inflight_prefill_pending.as_deref().unwrap_or(&[]), + ); reclaim_ready_prefetch(&mut executor, &mut deferred, &mut loading, reserve_floor); offer_prefetch(&mut executor, &mut deferred, &mut loading, reserve_floor); // 3. Nothing active and nothing admittable → block. Prefer blocking on // an in-flight load (so its request prefills next) over a new submit; - // only truly idle (no loads either) do we block on the channel. + // an in-flight decode-overlap prefill over a new submit; only truly + // idle (no loads or GPU work) do we block on the channel. if active.is_empty() && deferred.is_empty() && prefilling.is_empty() { + if inflight_prefill_pending.is_some() { + match executor.wait_async_prefill() { + Ok(prefill_result) => apply_async_prefill_result( + &mut executor, + &mut active, + &mut prefilling, + &mut tracker, + &mut inflight_prefill_pending, + prefill_result, + ), + Err(error) => { + warn!("decode-overlap: async prefill wait failed: {error:#}"); + // CUDA event synchronization failures abort inside the + // real executor. A returned error is a post-wait state + // or result mismatch, so report it and stop this broken + // scheduler instance; executor teardown owns the + // remaining request resources. + let message = format!("{error:#}"); + for req in inflight_prefill_pending.as_ref().unwrap() { + let _ = req.token_tx.send(TokenEvent::Error { + message: message.clone(), + prompt_tokens: req.prompt_tokens.len(), + completion_tokens: 0, + }); + } + return; + } + } + continue; + } if !loading.is_empty() { let reserve_floor = admitted_future_blocks(&executor, &active, &prefilling); block_on_loading(&mut executor, &mut deferred, &mut loading, reserve_floor); @@ -637,10 +691,11 @@ fn scheduler_loop( release_rejected(&mut executor, &mut tracker, rejected); } - let admission = admit_deferred_requests( + let admission = admit_deferred_requests_with_inflight( lora_validation.accepted, &active, &prefilling, + inflight_prefill_pending.as_deref().unwrap_or(&[]), executor.block_size(), executor.available_blocks(), executor.max_request_blocks(), @@ -1128,10 +1183,22 @@ fn admitted_future_blocks( executor: &E, active: &[ActiveRequestState], prefilling: &[PendingRequest], +) -> usize { + admitted_future_blocks_with_inflight(executor, active, prefilling, &[]) +} + +fn admitted_future_blocks_with_inflight( + executor: &E, + active: &[ActiveRequestState], + prefilling: &[PendingRequest], + inflight_prefilling: &[PendingRequest], ) -> usize { let block_size = executor.block_size(); active_future_blocks(active, block_size) + prefilling_future_blocks(prefilling, block_size, |id| executor.prefetched_blocks(id)) + + inflight_prefilling_future_blocks(inflight_prefilling, block_size, |id| { + executor.prefetched_blocks(id) + }) } fn prefilling_future_blocks( @@ -1152,6 +1219,25 @@ fn prefilling_future_blocks( .sum() } +fn inflight_prefilling_future_blocks( + prefilling: &[PendingRequest], + block_size: usize, + prefetch_credit: impl Fn(RequestId) -> usize, +) -> usize { + prefilling + .iter() + .map(|req| { + let scheduled_prompt_tokens = req + .prefill_pos + .saturating_add(req.step_chunk) + .min(req.prompt_tokens.len()); + pending_lifetime_blocks(req, block_size) + .saturating_sub(blocks_needed(scheduled_prompt_tokens, block_size)) + .saturating_sub(prefetch_credit(req.request_id)) + }) + .sum() +} + /// Default for `max_prefill_tokens`: prompt tokens forwarded in a single step /// (chunked prefill). Prefill activation scratch scales with the step's total /// prompt tokens (~22 KB/token measured on Qwen3-4B), so an unbounded prefill @@ -1169,12 +1255,44 @@ fn prefilling_future_blocks( pub const DEFAULT_MAX_PREFILL_TOKENS: usize = 1024; fn admit_deferred_requests( + deferred: Vec, + active: &[ActiveRequestState], + prefilling: &[PendingRequest], + block_size: usize, + available_blocks: usize, + max_request_blocks: usize, + max_context_tokens: usize, + max_decode_batch_size: usize, + max_prefill_tokens: usize, + prefetch_credit: impl Fn(RequestId) -> usize, +) -> AdmissionOutcome { + admit_deferred_requests_with_inflight( + deferred, + active, + prefilling, + &[], + block_size, + available_blocks, + max_request_blocks, + max_context_tokens, + max_decode_batch_size, + max_prefill_tokens, + prefetch_credit, + ) +} + +fn admit_deferred_requests_with_inflight( deferred: Vec, active: &[ActiveRequestState], // Admitted requests still mid-prefill: they hold KV for their applied // chunks and will take a decode slot when they promote, so admission // must reserve both or completing chunks can overshoot capacity. prefilling: &[PendingRequest], + // A decode-overlap prefill is temporarily removed from `prefilling` while + // its GPU work is in flight, but it still owns the same KV blocks and + // eventual decode slot. Keep it in the accounting view until its result + // is applied. + inflight_prefilling: &[PendingRequest], block_size: usize, available_blocks: usize, max_request_blocks: usize, @@ -1193,10 +1311,16 @@ fn admit_deferred_requests( prefilling, block_size, &prefetch_credit, + )) + .saturating_sub(inflight_prefilling_future_blocks( + inflight_prefilling, + block_size, + &prefetch_credit, )); let mut decode_slots = max_decode_batch_size .saturating_sub(active.len()) - .saturating_sub(prefilling.len()); + .saturating_sub(prefilling.len()) + .saturating_sub(inflight_prefilling.len()); let mut pending = Vec::new(); let mut still_deferred = Vec::new(); let mut rejected = Vec::new(); diff --git a/openinfer-qwen3/src/scheduler/tests.rs b/openinfer-qwen3/src/scheduler/tests.rs index fb15f80c..a811b7fc 100644 --- a/openinfer-qwen3/src/scheduler/tests.rs +++ b/openinfer-qwen3/src/scheduler/tests.rs @@ -1,9 +1,15 @@ use std::collections::HashMap; use std::sync::Arc; +use std::sync::Barrier; +use std::sync::Condvar; use std::sync::Mutex; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::time::Duration; use std::time::Instant; +use anyhow::Context; use anyhow::Result; use openinfer_core::engine::EngineControlError; use openinfer_core::engine::LoadLoraAdapterRequest; @@ -20,10 +26,54 @@ use crate::executor::PrefillStepItem; use crate::executor::UnifiedPlan; use crate::executor::UnifiedResult; +struct DecodePause { + started: Barrier, + release: Barrier, +} + +impl DecodePause { + fn new() -> Self { + Self { + started: Barrier::new(2), + release: Barrier::new(2), + } + } +} + +struct AsyncPrefillGate { + ready: Mutex, + wake: Condvar, + wait_calls: AtomicUsize, +} + +impl AsyncPrefillGate { + fn new() -> Self { + Self { + ready: Mutex::new(false), + wake: Condvar::new(), + wait_calls: AtomicUsize::new(0), + } + } + + fn wait(&self) { + self.wait_calls.fetch_add(1, Ordering::SeqCst); + let mut ready = self.ready.lock().unwrap(); + while !*ready { + ready = self.wake.wait(ready).unwrap(); + } + } + + fn release(&self) { + *self.ready.lock().unwrap() = true; + self.wake.notify_all(); + } +} + struct FakeExecutor { block_size: usize, max_request_blocks: usize, max_context_tokens: usize, + max_decode_batch_size: usize, available_blocks: usize, held_tokens: HashMap, // Prompt progress of requests mid-chunked-prefill (mirrors the real @@ -35,6 +85,11 @@ struct FakeExecutor { dropped: Arc>>, prefetch_offers: Arc>>, stop_token: Option, + decode_overlap: bool, + decode_pause: Option>, + async_prefill_gate: Option>, + async_prefill: Option, + lose_async_prefill_result: bool, } impl FakeExecutor { @@ -43,6 +98,7 @@ impl FakeExecutor { block_size: 16, max_request_blocks, max_context_tokens: usize::MAX, + max_decode_batch_size: 64, available_blocks: max_request_blocks, held_tokens: HashMap::new(), prefill_positions: HashMap::new(), @@ -52,14 +108,40 @@ impl FakeExecutor { dropped, prefetch_offers: Arc::new(Mutex::new(Vec::new())), stop_token: None, + decode_overlap: false, + decode_pause: None, + async_prefill_gate: None, + async_prefill: None, + lose_async_prefill_result: false, } } + fn with_decode_overlap( + mut self, + decode_pause: Arc, + async_prefill_gate: Arc, + ) -> Self { + self.decode_overlap = true; + self.decode_pause = Some(decode_pause); + self.async_prefill_gate = Some(async_prefill_gate); + self + } + fn with_stop_token(mut self, token: u32) -> Self { self.stop_token = Some(token); self } + fn with_missing_async_prefill_result(mut self) -> Self { + self.lose_async_prefill_result = true; + self + } + + fn with_max_decode_batch_size(mut self, max_decode_batch_size: usize) -> Self { + self.max_decode_batch_size = max_decode_batch_size; + self + } + fn with_decode_failure(mut self) -> Self { self.fail_decode_once = true; self @@ -131,7 +213,7 @@ impl ModelExecutor for FakeExecutor { } fn max_decode_batch_size(&self) -> usize { - 64 + self.max_decode_batch_size } fn available_blocks(&self) -> usize { @@ -192,6 +274,10 @@ impl ModelExecutor for FakeExecutor { } fn execute_decode(&mut self, plan: DecodePlan<'_>) -> Result { + if let Some(pause) = &self.decode_pause { + pause.started.wait(); + pause.release.wait(); + } if !self.decode_delay.is_zero() { std::thread::sleep(self.decode_delay); } @@ -235,23 +321,64 @@ impl ModelExecutor for FakeExecutor { self.ensure_request_tokens(req.request_id, current_tokens + 1)?; } + let prefill_requests: Vec<_> = plan + .prefill_requests + .iter() + .map(|req| self.fake_prefill_result(req)) + .collect(); + let decode_requests = plan + .decode_requests + .iter() + .map(|req| DecodeRequestResult { + request_id: req.request_id, + token: 200 + req.request_id.get() as u32, + logprob: None, + }) + .collect(); + if self.decode_overlap { + self.async_prefill = Some(PrefillResult { + requests: prefill_requests.clone(), + dflash_context_captured_requests: Vec::new(), + }); + } Ok(UnifiedResult { - prefill_requests: plan - .prefill_requests - .iter() - .map(|req| self.fake_prefill_result(req)) - .collect(), - decode_requests: plan - .decode_requests - .iter() - .map(|req| DecodeRequestResult { - request_id: req.request_id, - token: 200 + req.request_id.get() as u32, - logprob: None, - }) - .collect(), + prefill_requests: if self.decode_overlap { + Vec::new() + } else { + prefill_requests + }, + decode_requests, }) } + + fn has_decode_overlap(&self) -> bool { + self.decode_overlap + } + + fn poll_async_prefill(&mut self) -> Option { + let gate = self.async_prefill_gate.as_ref()?; + if !*gate.ready.lock().unwrap() { + return None; + } + self.async_prefill.take() + } + + fn wait_async_prefill(&mut self) -> Result { + let gate = self + .async_prefill_gate + .as_ref() + .ok_or_else(|| anyhow::anyhow!("fake async prefill is not enabled"))?; + gate.wait(); + if self.lose_async_prefill_result { + self.async_prefill.take(); + let missing: Result = + Err(anyhow::anyhow!("fake worker response missing")); + return missing.context("async prefill completed without a result"); + } + self.async_prefill + .take() + .ok_or_else(|| anyhow::anyhow!("fake async prefill result was lost")) + } } #[test] @@ -441,6 +568,82 @@ fn admission_respects_decode_batch_capacity() { assert!(outcome.rejected.is_empty()); } +#[test] +fn admission_counts_inflight_prefill_decode_slot() { + let (token_tx, _rx) = TokenSink::standalone(); + let active = [ActiveRequestState { + request_id: RequestId(0), + lora_adapter: None, + token_tx, + last_token: 1, + generated_count: 1, + max_tokens: 8, + prompt_len: 16, + params: SamplingParams::default(), + logprobs: 0, + }]; + let mk = |id: u64, prompt_len, max_tokens| { + PendingRequest::from_scheduler_request(RequestId(id), request(prompt_len, max_tokens).0) + }; + let mut inflight = mk(1, 16, 1); + inflight.step_chunk = 16; + let deferred = mk(2, 16, 1); + + let outcome = admit_deferred_requests_with_inflight( + vec![deferred], + &active, + &[], + &[inflight], + 16, + 1024, + 1024, + usize::MAX, + 2, + 32, + |_| 0, + ); + + assert!( + outcome.pending.is_empty(), + "active decode plus in-flight prefill already fill decode capacity" + ); + assert_eq!(outcome.deferred[0].request_id, RequestId(2)); + assert!(outcome.rejected.is_empty()); +} + +#[test] +fn admission_charges_inflight_prefill_only_for_unscheduled_kv_tail() { + let active: [ActiveRequestState; 0] = []; + let mk = |id: u64, prompt_len, max_tokens| { + PendingRequest::from_scheduler_request(RequestId(id), request(prompt_len, max_tokens).0) + }; + let mut inflight = mk(1, 16, 17); + inflight.step_chunk = 16; + let deferred = mk(2, 16, 1); + + let outcome = admit_deferred_requests_with_inflight( + vec![deferred], + &active, + &[], + &[inflight], + 16, + 3, + 1024, + usize::MAX, + 64, + 32, + |_| 0, + ); + + assert_eq!( + outcome.pending[0].request_id, + RequestId(2), + "the in-flight request's current chunk is already out of available_blocks" + ); + assert!(outcome.deferred.is_empty()); + assert!(outcome.rejected.is_empty()); +} + #[test] fn prefill_chunking_caps_step_tokens_and_keeps_fifo_progress() { let mk = |id: u64, prompt_len, max_tokens| { @@ -839,6 +1042,223 @@ fn wait_until(timeout: Duration, mut predicate: impl FnMut() -> bool) -> bool { false } +fn try_recv_event_with_timeout( + rx: &mut openinfer_core::engine::TokenStreamReceiver, + timeout: Duration, +) -> Option { + let start = Instant::now(); + while start.elapsed() < timeout { + match rx.try_recv() { + Ok((_, TokenEvent::Scheduled { .. })) => {} + Ok((_, event)) => return Some(event), + Err(tokio::sync::mpsc::error::TryRecvError::Empty) => { + std::thread::sleep(Duration::from_millis(5)); + } + Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => return None, + } + } + None +} + +#[test] +fn decode_overlap_idle_wait_resolves_prefill_before_final_handle_drop() { + let dropped = Arc::new(Mutex::new(Vec::new())); + let decode_pause = Arc::new(DecodePause::new()); + let async_prefill_gate = Arc::new(AsyncPrefillGate::new()); + let executor = FakeExecutor::new(8, Arc::clone(&dropped)) + .with_decode_overlap(Arc::clone(&decode_pause), Arc::clone(&async_prefill_gate)); + let handle = start_with_executor(executor, 42, DEFAULT_MAX_PREFILL_TOKENS); + + // Keep the first request active after its first decode so the next request + // enters a unified decode + async-prefill step. + let (active_request, mut active_rx) = request(16, 3); + handle + .submit(active_request) + .expect("submit active request"); + assert!(matches!( + recv_skipping_scheduled(&mut active_rx), + Some(TokenEvent::Token { id: 100, .. }) + )); + + // Submit the pending request while the first decode is paused. This makes + // the next scheduler step deterministic: the active request finishes while + // the pending request's prefill is launched asynchronously. + decode_pause.started.wait(); + let (overlap_request, mut overlap_rx) = request(16, 1); + handle + .submit(overlap_request) + .expect("submit overlap request"); + decode_pause.release.wait(); + + // The fixed scheduler must wait on the async prefill event once all other + // queues are empty. The old scheduler parks on submit_rx instead, so this + // condition stays false and the test fails without the fix. + let entered_prefill_wait = wait_until(Duration::from_secs(1), || { + async_prefill_gate.wait_calls.load(Ordering::SeqCst) > 0 + }); + + let drop_finished = Arc::new(AtomicBool::new(false)); + let drop_finished_for_thread = Arc::clone(&drop_finished); + let drop_thread = std::thread::spawn(move || { + drop(handle); + drop_finished_for_thread.store(true, Ordering::SeqCst); + }); + let finished_before_prefill_release = wait_until(Duration::from_millis(100), || { + drop_finished.load(Ordering::SeqCst) + }); + + // Releasing the event lets the scheduler resolve the result before the + // final EngineHandle drop joins its thread. + async_prefill_gate.release(); + drop_thread.join().expect("scheduler thread should join"); + + assert!( + entered_prefill_wait, + "idle scheduler must wait on in-flight prefill instead of blocking only on submissions" + ); + assert!( + !finished_before_prefill_release, + "final EngineHandle drop must not abandon an in-flight prefill" + ); + assert!(matches!( + try_recv_event_with_timeout(&mut overlap_rx, Duration::from_secs(1)), + Some(TokenEvent::Token { id: 101, .. }) + )); + assert!(matches!( + try_recv_event_with_timeout(&mut overlap_rx, Duration::from_secs(1)), + Some(TokenEvent::Finished { .. }) + )); + assert!(dropped.lock().unwrap().contains(&0)); + assert!(dropped.lock().unwrap().contains(&1)); +} + +#[test] +fn decode_overlap_admission_counts_inflight_prefill_capacity() { + let dropped = Arc::new(Mutex::new(Vec::new())); + let decode_pause = Arc::new(DecodePause::new()); + let async_prefill_gate = Arc::new(AsyncPrefillGate::new()); + let executor = FakeExecutor::new(8, Arc::clone(&dropped)) + .with_decode_overlap(Arc::clone(&decode_pause), Arc::clone(&async_prefill_gate)) + .with_max_decode_batch_size(2); + let handle = start_with_executor(executor, 42, DEFAULT_MAX_PREFILL_TOKENS); + let load_watch = handle.load_watch().expect("scheduler exposes load watch"); + + let (active_request, mut active_rx) = request(16, 6); + handle + .submit(active_request) + .expect("submit active request"); + assert!(matches!( + recv_skipping_scheduled(&mut active_rx), + Some(TokenEvent::Token { id: 100, .. }) + )); + + // First decode-only step: hold the active request runnable while queuing a + // second request. The following scheduler iteration will run unified + // decode+prefill and leave that prefill in flight. + decode_pause.started.wait(); + let (overlap_request, mut overlap_rx) = request(16, 2); + handle + .submit(overlap_request) + .expect("submit overlap request"); + decode_pause.release.wait(); + + // Second decode-only step: the overlap prefill is in flight but the active + // request can still make progress, so the scheduler is not blocked in + // wait_async_prefill(). Submit the third request for admission during these + // in-flight iterations. + decode_pause.started.wait(); + let (deferred_request, mut deferred_rx) = request(16, 1); + handle + .submit(deferred_request) + .expect("submit request during async prefill"); + decode_pause.release.wait(); + + // Third decode-only step: admission has seen the third request while + // active+in-flight already fill the two decode slots. Let one more decode + // finish so the next load snapshot exposes whether the request was wrongly + // counted as running. + decode_pause.started.wait(); + decode_pause.release.wait(); + let over_admitted = wait_until(Duration::from_secs(1), || { + load_watch.borrow().num_running_reqs > 2 + }); + + // Release the fourth active decode and the async prefill so the scheduler + // can drain all requests before the handle is dropped. + decode_pause.started.wait(); + async_prefill_gate.release(); + decode_pause.release.wait(); + + assert!(matches!( + try_recv_event_with_timeout(&mut overlap_rx, Duration::from_secs(1)), + Some(TokenEvent::Token { id: 101, .. }) + )); + assert!(matches!( + try_recv_event_with_timeout(&mut deferred_rx, Duration::from_secs(1)), + Some(TokenEvent::Token { id: 102, .. }) + )); + assert!(matches!( + try_recv_event_with_timeout(&mut deferred_rx, Duration::from_secs(1)), + Some(TokenEvent::Finished { .. }) + )); + drop(handle); + + assert!( + !over_admitted, + "in-flight async prefill must reserve its scheduler slot before admitting new requests" + ); +} + +#[test] +fn decode_overlap_missing_result_reports_chain_and_stops_scheduler() { + let dropped = Arc::new(Mutex::new(Vec::new())); + let decode_pause = Arc::new(DecodePause::new()); + let async_prefill_gate = Arc::new(AsyncPrefillGate::new()); + let executor = FakeExecutor::new(8, Arc::clone(&dropped)) + .with_decode_overlap(Arc::clone(&decode_pause), Arc::clone(&async_prefill_gate)) + .with_missing_async_prefill_result(); + let handle = start_with_executor(executor, 42, DEFAULT_MAX_PREFILL_TOKENS); + + let (active_request, mut active_rx) = request(16, 3); + handle + .submit(active_request) + .expect("submit active request"); + assert!(matches!( + recv_skipping_scheduled(&mut active_rx), + Some(TokenEvent::Token { id: 100, .. }) + )); + + decode_pause.started.wait(); + let (overlap_request, mut overlap_rx) = request(16, 1); + handle + .submit(overlap_request) + .expect("submit overlap request"); + decode_pause.release.wait(); + + assert!(wait_until(Duration::from_secs(1), || { + async_prefill_gate.wait_calls.load(Ordering::SeqCst) > 0 + })); + async_prefill_gate.release(); + + let message = match try_recv_event_with_timeout(&mut overlap_rx, Duration::from_secs(1)) { + Some(TokenEvent::Error { message, .. }) => message, + other => panic!("expected async prefill result error, got {other:?}"), + }; + assert!(message.contains("async prefill completed without a result")); + assert!(message.contains("fake worker response missing")); + drop(handle); + + let dropped = dropped.lock().unwrap(); + assert!( + dropped.contains(&0), + "the completed decode request is retired" + ); + assert!( + !dropped.contains(&1), + "a request without a resolved prefill result must not be retired normally" + ); +} + #[test] fn unknown_lora_request_is_rejected_without_blocking_base_request() { let dropped = Arc::new(Mutex::new(Vec::new()));