From fbe384b65e8fde12dff28fcdef2f4c238ddd1c4f Mon Sep 17 00:00:00 2001 From: CAICAIIs <3360776475@qq.com> Date: Thu, 10 Sep 2026 11:29:35 +0800 Subject: [PATCH 1/5] refactor(qwen35): split TP worker runtime into tp_executor/worker.rs Move the per-rank worker runtime out of the tp_executor.rs God module: TpWorker + spawn/drop lifecycle, TpStartupGate, the NCCL startup watchdog, TpWorkerState/TpWorkerPrepared and its command loop (run/respond/ execute_*/precapture/drop), the slot bookkeeping helpers, the shared decode-row view/sampling helpers, recurrent-capacity math, and the worker CublasThreadGuard binding. The entry keeps command/reply types, the executor orchestration (including the precapture sweep), plan validators, and response validation, reaching worker items via the same `use super::*` style as scheduler/{backend,tp}. Rebased over #720/#722. Behavior unchanged; the watchdog disarm test moves into worker.rs with its implementation. Signed-off-by: CAICAIIs <3360776475@qq.com> --- pegainfer-qwen35/src/tp_executor.rs | 1178 +------------------ pegainfer-qwen35/src/tp_executor/worker.rs | 1197 ++++++++++++++++++++ 2 files changed, 1201 insertions(+), 1174 deletions(-) create mode 100644 pegainfer-qwen35/src/tp_executor/worker.rs diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index e02ace2c3..dcc7a0227 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -46,13 +46,14 @@ use crate::recurrent_state::RecurrentState; use crate::weights::ModelRuntimeConfig; use crate::weights::Qwen35Model; -const TP_NCCL_STARTUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); +mod worker; + +use worker::*; + const TP_RUNTIME_STEP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300); -const TP_WORKER_SHUTDOWN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); /// The pre-capture sweep records every decode bucket per rank; the 60 s NCCL /// startup budget is far too small for that. const TP_PRECAPTURE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(600); -const TP_RUNTIME_MEMORY_RESERVE_BYTES: usize = 512 * 1024 * 1024; const TRITON_AOT_DEVICE_TABLE_LEN: usize = 16; /// One controller-barriered phase of the TP decode-graph pre-capture sweep. @@ -1070,1092 +1071,6 @@ impl Drop for Qwen35TpExecutor { } } -fn spawn_nccl_startup_watchdog() -> Result<(mpsc::SyncSender<()>, JoinHandle<()>)> { - let (done_tx, done_rx) = mpsc::sync_channel(1); - let watchdog = thread::Builder::new() - .name("qwen35-tp-nccl-startup-watchdog".into()) - .spawn(move || { - if done_rx.recv_timeout(TP_NCCL_STARTUP_TIMEOUT).is_ok() { - return; - } - eprintln!( - "Qwen3.5 TP NCCL startup did not complete within {}s; aborting", - TP_NCCL_STARTUP_TIMEOUT.as_secs() - ); - log::error!( - "Qwen3.5 TP NCCL startup did not complete within {}s; aborting", - TP_NCCL_STARTUP_TIMEOUT.as_secs() - ); - std::process::abort(); - }) - .map_err(|err| anyhow::anyhow!("failed to spawn Qwen3.5 TP NCCL watchdog: {err}"))?; - Ok((done_tx, watchdog)) -} - -#[allow(clippy::needless_pass_by_value)] -fn disarm_nccl_startup_watchdog( - done_tx: mpsc::SyncSender<()>, - watchdog: JoinHandle<()>, -) -> Result<()> { - done_tx - .send(()) - .map_err(|_| anyhow::anyhow!("Qwen3.5 TP NCCL watchdog exited unexpectedly"))?; - watchdog - .join() - .map_err(|_| anyhow::anyhow!("Qwen3.5 TP NCCL watchdog panicked")) -} - -struct TpWorker { - tx: mpsc::Sender, - handle: Option>, - done: mpsc::Receiver<()>, -} - -#[derive(Clone, Copy, Default, PartialEq, Eq)] -enum TpStartupDecision { - #[default] - Pending, - Connect, - Cancel, -} - -#[derive(Default)] -struct TpStartupGate { - decision: Mutex, - changed: Condvar, -} - -impl TpStartupGate { - fn connect(&self) { - self.set(TpStartupDecision::Connect); - } - - fn cancel(&self) { - self.set(TpStartupDecision::Cancel); - } - - fn wait(&self) -> bool { - let mut decision = self.decision.lock().unwrap_or_else(PoisonError::into_inner); - while *decision == TpStartupDecision::Pending { - decision = self - .changed - .wait(decision) - .unwrap_or_else(PoisonError::into_inner); - } - *decision == TpStartupDecision::Connect - } - - fn set(&self, next: TpStartupDecision) { - let mut decision = self.decision.lock().unwrap_or_else(PoisonError::into_inner); - if *decision == TpStartupDecision::Pending { - *decision = next; - self.changed.notify_all(); - } - } -} - -impl TpWorker { - #[allow(clippy::too_many_arguments)] - #[allow(clippy::type_complexity)] - fn spawn( - rank: usize, - world_size: usize, - model: Qwen35Model, - max_batch: usize, - max_prefill_tokens: usize, - graph_enabled: bool, - nccl_id: cudarc::nccl::safe::Id, - startup_gate: Arc, - effective_max_batch: Arc, - poison: Arc, - ) -> Result<( - Self, - mpsc::Receiver>, - mpsc::Receiver>, - )> { - let (tx, rx) = mpsc::channel(); - let (preflight_tx, preflight_rx) = mpsc::channel(); - let (startup_tx, startup_rx) = mpsc::channel(); - let (done_tx, done_rx) = mpsc::channel(); - let panic_poison = Arc::clone(&poison); - let handle = thread::Builder::new() - .name(format!("qwen35-tp-rank-{rank}")) - .spawn(move || { - let outcome = catch_unwind(AssertUnwindSafe(|| { - let prepared = TpWorkerPrepared::new( - rank, - world_size, - model, - max_batch, - max_prefill_tokens, - graph_enabled, - ); - let prepared = match prepared { - Ok((prepared, rank_max_batch)) => { - let _ = preflight_tx.send(Ok(rank_max_batch)); - prepared - } - Err(err) => { - let _ = preflight_tx.send(Err(err)); - return; - } - }; - if !startup_gate.wait() { - return; - } - let max_batch = effective_max_batch.load(Ordering::Acquire); - match prepared.connect(nccl_id, max_batch, graph_enabled, poison) { - Ok(mut state) => { - let _ = startup_tx.send(Ok(())); - state.run(rx); - } - Err(err) => { - let _ = startup_tx.send(Err(err)); - } - } - })); - if outcome.is_err() { - panic_poison.poison(format!("worker rank {rank} panicked")); - } - let _ = done_tx.send(()); - }) - .map_err(|e| anyhow::anyhow!("failed to spawn Qwen3.5 TP worker {rank}: {e}"))?; - - Ok(( - Self { - tx, - handle: Some(handle), - done: done_rx, - }, - preflight_rx, - startup_rx, - )) - } - - fn send(&self, command: TpWorkerCommand) -> Result<()> { - self.tx - .send(command) - .map_err(|_| anyhow::anyhow!("Qwen3.5 TP worker channel closed")) - } - - fn join_bounded(&mut self) { - if self.handle.is_none() { - return; - } - if self.done.recv_timeout(TP_WORKER_SHUTDOWN_TIMEOUT).is_err() { - fatal_tp_abort("Qwen3.5 TP worker did not exit during bounded shutdown"); - } - if let Some(handle) = self.handle.take() { - let _ = handle.join(); - } - } -} - -impl Drop for TpWorker { - fn drop(&mut self) { - let _ = self.tx.send(TpWorkerCommand::Shutdown); - self.join_bounded(); - } -} - -struct TpWorkerState { - rank: usize, - _world_size: usize, - max_batch: usize, - /// Before `model` on purpose: NCCL comm teardown polls until every graph - /// that recorded its collectives is destroyed, so the decode graphs must - /// drop before `model.tp_comm` (qwen3 teardown-hang precedent). - graph_state: Option, - model: Qwen35Model, - requests: Vec, - /// Graph-mode slot ownership: `slot_map[i]` is the request whose recurrent - /// state lives in `graph_state.slot_states[i]`. The scheduler owns slot - /// assignment and compaction; the worker only applies and checks them. - /// Empty in eager mode. - slot_map: Vec>, - decode_buffers: BatchDecodeBuffers35, - /// Eager decode GDR pointer tables: allocated once at capacity, refilled - /// with the live rows every step. - decode_pointer_tables: LinearStatePointerTables, - sample_scratch: pegainfer_sample::SampleScratch, - _cublas_guard: CublasThreadGuard, - poison: Arc, -} - -struct TpWorkerPrepared { - rank: usize, - world_size: usize, - max_batch: usize, - model: Qwen35Model, - decode_buffers: BatchDecodeBuffers35, - sample_scratch: pegainfer_sample::SampleScratch, - cublas_guard: CublasThreadGuard, -} - -struct TpRequestState { - request_id: RequestId, - phase: TpRequestPhase, - kv: KvState, - /// Prefill-owned recurrent state. Graph mode moves it into the decode slot - /// on the request's first decode row (`None` afterwards); the eager path - /// keeps it for the request's whole lifetime. - recurrent: Option, -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq)] -enum TpRequestPhase { - Prefilling, - Decoding, -} - -#[cfg(test)] -#[derive(Clone, Debug, PartialEq, Eq)] -struct WorkerStateSnapshot { - rank: usize, - request_count: usize, - requests: Vec<(RequestId, TpRequestPhase)>, -} - -impl TpWorkerPrepared { - fn new( - rank: usize, - world_size: usize, - model: Qwen35Model, - requested_max_batch: usize, - max_prefill_tokens: usize, - graph_enabled: bool, - ) -> Result<(Self, usize)> { - let cublas_guard = bind_worker_thread(&model)?; - let (free_bytes, total_bytes) = model - .device_ctx() - .ctx - .mem_get_info() - .map_err(|err| anyhow::anyhow!("failed to query TP rank {rank} memory: {err}"))?; - // Recurrent state is rank-local, so worker capacity math uses the - // local value-head/qkv sizes. - let recurrent_bytes = RecurrentState::allocation_bytes(model.config(), model.geometry); - let prefill_scratch_tokens = prefill_scratch_tokens(max_prefill_tokens); - let prefill_scratch_bytes = GdrChunkwiseScratch35::estimate_bytes( - model.config(), - model.geometry, - prefill_scratch_tokens, - ); - // Graph mode pre-allocates one fixed-address slot state per decode - // bucket position up front; reserve that before sizing per-request - // (prefill-transient) state capacity. The reserve must track the - // bucket of the *effective* batch, not the requested one: reserving - // for `bucket_for(requested)` can starve a tight-memory rank down to - // zero capacity. Iterate the bucket downward until it stabilises — - // the bucket only shrinks, so this converges — and clamp the fitted - // batch to the reserved bucket so the later `bucket_for(effective)` - // graph allocation never exceeds the reserve. - let max_batch = if graph_enabled { - let mut slot_bucket = bucket_for(requested_max_batch); - loop { - let reserve = slot_bucket * recurrent_bytes; - let candidate = effective_recurrent_capacity( - requested_max_batch, - free_bytes.saturating_sub(reserve), - recurrent_bytes, - TP_RUNTIME_MEMORY_RESERVE_BYTES, - prefill_scratch_bytes, - ); - let fitted = candidate.min(slot_bucket); - let next = bucket_for(fitted); - if next >= slot_bucket { - break fitted; - } - slot_bucket = next; - } - } else { - effective_recurrent_capacity( - requested_max_batch, - free_bytes, - recurrent_bytes, - TP_RUNTIME_MEMORY_RESERVE_BYTES, - prefill_scratch_bytes, - ) - }; - anyhow::ensure!( - max_batch > 0, - "Qwen3.5 TP rank {rank} has {} MiB free after fixed buffers, but one recurrent request needs {} MiB plus {} MiB runtime reserve and {} MiB prefill scratch for {} tokens", - free_bytes / (1024 * 1024), - recurrent_bytes / (1024 * 1024), - TP_RUNTIME_MEMORY_RESERVE_BYTES / (1024 * 1024), - prefill_scratch_bytes / (1024 * 1024), - prefill_scratch_tokens, - ); - log::info!( - "Qwen3.5 TP rank {rank} recurrent capacity: requested={requested_max_batch}, effective={max_batch}, per_request={:.3} MiB, free={:.0} MiB/{:.0} MiB, runtime_reserve={} MiB, prefill_tokens={}, prefill_scratch={:.0} MiB", - recurrent_bytes as f64 / 1024.0 / 1024.0, - free_bytes as f64 / 1024.0 / 1024.0, - total_bytes as f64 / 1024.0 / 1024.0, - TP_RUNTIME_MEMORY_RESERVE_BYTES / (1024 * 1024), - prefill_scratch_tokens, - prefill_scratch_bytes as f64 / 1024.0 / 1024.0, - ); - let decode_buffers = model.create_batch_decode_buffers_with_capacity(max_batch)?; - let sample_scratch = pegainfer_sample::SampleScratch::new( - model.device_ctx(), - model.config().selection_vocab, - max_batch, - )?; - Ok(( - Self { - rank, - world_size, - max_batch, - model, - decode_buffers, - sample_scratch, - cublas_guard, - }, - max_batch, - )) - } - - fn connect( - self, - nccl_id: cudarc::nccl::safe::Id, - effective_max_batch: usize, - graph_enabled: bool, - poison: Arc, - ) -> Result { - let Self { - rank, - world_size, - max_batch, - mut model, - decode_buffers, - sample_scratch, - cublas_guard, - } = self; - anyhow::ensure!( - effective_max_batch > 0 && effective_max_batch <= max_batch, - "Qwen3.5 TP rank {rank} effective max_batch {effective_max_batch} exceeds local capacity {max_batch}" - ); - let comm = cudarc::nccl::safe::Comm::from_rank( - model.device_ctx().stream.clone(), - rank, - world_size, - nccl_id, - ) - .map_err(|e| anyhow::anyhow!("failed to initialize Qwen3.5 TP NCCL rank {rank}: {e:?}"))?; - model.attach_tp_comm(comm); - let decode_pointer_tables = LinearStatePointerTables::with_capacity( - model.device_ctx(), - model.config(), - effective_max_batch, - "Qwen3.5 TP eager decode", - )?; - let (graph_state, slot_map) = if graph_enabled { - // cuBLASLt plans are thread-local: tune the decode bucket GEMMs on - // this worker thread now so plan selection never runs inside - // cuStreamBeginCapture during the pre-capture sweep. - model.tune_decode_gemm_algos()?; - let slots = bucket_for(effective_max_batch); - let graph_state = model.create_batch_decode_graph_state_with_capacity(slots)?; - (Some(graph_state), vec![None; slots]) - } else { - (None, Vec::new()) - }; - Ok(TpWorkerState { - rank, - _world_size: world_size, - max_batch: effective_max_batch, - graph_state, - model, - requests: Vec::new(), - slot_map, - decode_buffers, - decode_pointer_tables, - sample_scratch, - _cublas_guard: cublas_guard, - poison, - }) - } -} - -fn prefill_scratch_tokens(max_prefill_tokens: usize) -> usize { - max_prefill_tokens.min(PREFILL_CHUNK_LEN) -} - -fn effective_recurrent_capacity( - requested_max_batch: usize, - free_bytes: usize, - recurrent_bytes_per_request: usize, - runtime_reserve_bytes: usize, - prefill_scratch_bytes: usize, -) -> usize { - if recurrent_bytes_per_request == 0 { - return requested_max_batch; - } - requested_max_batch.min( - free_bytes - .saturating_sub(runtime_reserve_bytes) - .saturating_sub(prefill_scratch_bytes) - / recurrent_bytes_per_request, - ) -} - -impl TpWorkerState { - #[allow(clippy::needless_pass_by_value)] - fn run(&mut self, rx: mpsc::Receiver) { - while let Ok(command) = rx.recv() { - let fatal = match command { - TpWorkerCommand::Ping { resp } => { - self.respond(resp, "ping", Ok(TpWorkerReply::Ack)) - } - TpWorkerCommand::RunPrefillChunks { - chunks, - sample_seed, - start, - resp, - } => { - if start.wait() == TpCommandDecision::Cancel { - false - } else { - let result = self.execute_prefill_chunks(&chunks, sample_seed); - self.respond(resp, "prefill", result) - } - } - TpWorkerCommand::RunDecodeStep { - requests, - sample_seed, - start, - resp, - } => { - if start.wait() == TpCommandDecision::Cancel { - false - } else { - let result = self.execute_decode(&requests, sample_seed); - self.respond(resp, "decode", result) - } - } - TpWorkerCommand::RunUnifiedStep { plan, start, resp } => { - if start.wait() == TpCommandDecision::Cancel { - false - } else { - let result = self.execute_unified(&plan); - self.respond(resp, "unified step", result) - } - } - TpWorkerCommand::DropRequest { - request_id, - compaction, - start, - resp, - } => { - if start.wait() == TpCommandDecision::Cancel { - false - } else { - let result = self - .drop_request(request_id, compaction) - .map(|existed| TpWorkerReply::DropAck { existed }); - self.respond(resp, "drop request", result) - } - } - TpWorkerCommand::Precapture { phase, start, resp } => { - if start.wait() == TpCommandDecision::Cancel { - false - } else { - let result = self.precapture_phase(phase).map(|()| TpWorkerReply::Ack); - self.respond(resp, "decode graph precapture", result) - } - } - #[cfg(test)] - TpWorkerCommand::SnapshotState { resp } => { - let snapshot = WorkerStateSnapshot { - rank: self.rank, - request_count: self.requests.len(), - requests: self - .requests - .iter() - .map(|state| (state.request_id, state.phase)) - .collect(), - }; - self.respond( - resp, - "snapshot state", - Ok(TpWorkerReply::Snapshot(snapshot)), - ) - } - #[cfg(test)] - TpWorkerCommand::RemoveRequestStateForTest { request_id, resp } => { - let _ = resp.send(self.drop_request(request_id, None).unwrap_or(false)); - false - } - #[cfg(test)] - TpWorkerCommand::DisconnectForTest { ready } => { - let _ = ready.send(()); - break; - } - TpWorkerCommand::Shutdown => break, - }; - if fatal { - break; - } - } - } - - #[allow(clippy::needless_pass_by_value)] - fn respond( - &self, - resp: mpsc::Sender, - operation: &'static str, - result: Result, - ) -> bool { - match result { - Ok(reply) => { - let _ = resp.send(TpWorkerResponse { - rank: self.rank, - result: Ok(reply), - }); - false - } - Err(err) => { - let reason = self.poison.poison(format!( - "rank {} failed during {operation}: {err:#}", - self.rank - )); - let _ = resp.send(TpWorkerResponse { - rank: self.rank, - result: Err(anyhow::anyhow!(reason)), - }); - true - } - } - } - - fn execute_prefill_chunks( - &mut self, - chunks: &[TpPrefillChunkItem], - sample_seed: u64, - ) -> Result { - let requests = self.execute_prefill_rows(chunks, sample_seed)?; - if self.rank == 0 { - Ok(TpWorkerReply::Prefill(PrefillResult { requests })) - } else { - Ok(TpWorkerReply::Ack) - } - } - - fn execute_prefill_rows( - &mut self, - chunks: &[TpPrefillChunkItem], - sample_seed: u64, - ) -> Result> { - anyhow::ensure!( - !chunks.is_empty(), - "Qwen3.5 TP prefill chunk command requires at least one chunk" - ); - validate_prefill_chunks(chunks)?; - let new_requests = chunks - .iter() - .filter(|chunk| self.request_index(chunk.request_id).is_none()) - .count(); - anyhow::ensure!( - self.requests.len() + new_requests <= self.max_batch, - "Qwen3.5 TP prefill chunks would exceed worker capacity {}", - self.max_batch - ); - - let mut primary_results = Vec::new(); - let mut final_row_idx = 0usize; - for chunk in chunks { - let state_idx = self.ensure_prefill_state(chunk.request_id)?; - let state = &mut self.requests[state_idx]; - anyhow::ensure!( - state.phase == TpRequestPhase::Prefilling, - "Qwen3.5 TP request {} is already in decode state", - chunk.request_id.get() - ); - - let prompt = [chunk.prompt_tokens.as_slice()]; - let mut recurrent_refs = vec![ - state - .recurrent - .as_mut() - .expect("prefill-phase TP request owns its recurrent state"), - ]; - let logits = self.model.batch_prefill_logits( - &prompt, - std::slice::from_mut(&mut state.kv), - &mut recurrent_refs, - )?; - - if chunk.finish_prefill { - if self.rank == 0 { - // TP prefill samples final chunks one row at a time. Offset - // by the final-row index so rows from the same command do - // not reuse the same sampling stream. - let row_seed = sample_seed.wrapping_add(final_row_idx as u64); - let result = self.sample_final_prefill_chunk(chunk, &logits, row_seed)?; - primary_results.push(result); - } - final_row_idx += 1; - self.requests[state_idx].phase = TpRequestPhase::Decoding; - } - } - - Ok(primary_results) - } - - /// Run one batched eager decode step over all rows in command order: a - /// single forward for the whole batch on every rank, then (rank 0 only) - /// one batched sampling pass over the per-row sampling params. Returns one - /// result row per request in command order on rank 0, empty elsewhere. - fn run_decode_batch( - &mut self, - requests: &[TpDecodeStepItem], - sample_seed: u64, - ) -> Result> { - let bs = requests.len(); - if bs == 0 { - return Ok(Vec::new()); - } - if self.graph_state.is_some() { - return self.run_decode_batch_graph(requests, sample_seed); - } - - // Resolve the worker state slot of every row in command order. - // Decode request ids are unique within one command - // (validate_decode_requests), so each slot is borrowed at most once. - let mut row_of_state: Vec> = vec![None; self.requests.len()]; - for (row, request) in requests.iter().enumerate() { - let state_idx = self.request_index(request.request_id).ok_or_else(|| { - anyhow::anyhow!( - "Qwen3.5 TP decode request {} has no worker state", - request.request_id.get() - ) - })?; - anyhow::ensure!( - self.requests[state_idx].phase == TpRequestPhase::Decoding, - "Qwen3.5 TP request {} is not ready for decode", - request.request_id.get() - ); - debug_assert!(row_of_state[state_idx].is_none()); - row_of_state[state_idx] = Some(row); - } - let mut kv_refs: Vec<&mut KvState> = Vec::with_capacity(bs); - let mut recurrent_refs: Vec<&mut RecurrentState> = Vec::with_capacity(bs); - for state in states_in_row_order(&mut self.requests, &row_of_state) { - let TpRequestState { kv, recurrent, .. } = state; - kv_refs.push(kv); - recurrent_refs.push( - recurrent - .as_mut() - .expect("eager TP decode request owns its recurrent state"), - ); - } - - // GDR pointer tables over the full decode batch: allocated once at - // capacity, refilled from the live rows every step (H2D only), so - // swap_remove retirement between steps can never leave a stale row - // addressed and no step pays for device allocations. - self.decode_pointer_tables.refill_from_recurrent_refs( - self.model.device_ctx(), - &mut recurrent_refs, - bs, - "Qwen3.5 TP eager decode", - )?; - let token_ids: Vec = requests.iter().map(|request| request.token_id).collect(); - self.model.batch_decode_eager_logits( - &token_ids, - &mut kv_refs, - &mut recurrent_refs, - &self.decode_pointer_tables, - &mut self.decode_buffers, - )?; - - if self.rank != 0 { - return Ok(Vec::new()); - } - sample_decode_rows( - self.model.device_ctx(), - &self.decode_buffers.logits, - requests, - sample_seed, - &mut self.sample_scratch, - ) - } - - /// CUDA Graph decode step under TP: replay-only (every bucket was recorded - /// by the startup pre-capture sweep), one forward for the whole batch on - /// every rank, then (rank 0 only) the same batched host-side sampling pass - /// as the eager path. - /// - /// Rows must arrive in the scheduler-owned dense slot order - /// (`slot_idx == row`). On a request's first decode row its prefill-owned - /// recurrent state is D2D-copied into `graph_state.slot_states[slot]` and - /// the per-request allocation is dropped; the persistent linear-state - /// pointer tables then keep every replay reading the fixed slot addresses. - fn run_decode_batch_graph( - &mut self, - requests: &[TpDecodeStepItem], - sample_seed: u64, - ) -> Result> { - let bs = requests.len(); - let graph_state = self - .graph_state - .as_mut() - .expect("graph decode arm requires graph state"); - let ctx = self.model.device_ctx(); - - // Resolve the worker state of every row, enforce dense slot order, and - // admit first-decode rows into their slots. Decode request ids are - // unique within one command (validate_decode_requests), so each slot - // is borrowed at most once. - let mut row_of_state: Vec> = vec![None; self.requests.len()]; - for (row, request) in requests.iter().enumerate() { - anyhow::ensure!( - request.slot_idx == Some(row), - "Qwen3.5 TP graph decode row {row} carries slot {:?}; rows must arrive in dense slot order 0..{bs}", - request.slot_idx - ); - let state_idx = self - .requests - .iter() - .position(|state| state.request_id == request.request_id) - .ok_or_else(|| { - anyhow::anyhow!( - "Qwen3.5 TP decode request {} has no worker state", - request.request_id.get() - ) - })?; - anyhow::ensure!( - self.requests[state_idx].phase == TpRequestPhase::Decoding, - "Qwen3.5 TP request {} is not ready for decode", - request.request_id.get() - ); - debug_assert!(row_of_state[state_idx].is_none()); - row_of_state[state_idx] = Some(row); - - if self.slot_map.get(row).copied().flatten() == Some(request.request_id) { - anyhow::ensure!( - self.requests[state_idx].recurrent.is_none(), - "Qwen3.5 TP request {} was admitted to slot {row} but still owns prefill recurrent state", - request.request_id.get() - ); - } else { - slot_admit(&mut self.slot_map, row, request.request_id)?; - let recurrent = self.requests[state_idx].recurrent.take().ok_or_else(|| { - anyhow::anyhow!( - "Qwen3.5 TP request {} lost its prefill recurrent state before slot admission", - request.request_id.get() - ) - })?; - graph_state.copy_state_to_slot(ctx, &recurrent, row)?; - } - } - - // KV refs in row (slot) order; page tables stay per-step H2D via - // sync_paged_meta inside batch_decode_graph. - let mut kv_refs: Vec<&mut KvState> = states_in_row_order(&mut self.requests, &row_of_state) - .into_iter() - .map(|state| &mut state.kv) - .collect(); - let token_ids: Vec = requests.iter().map(|request| request.token_id).collect(); - self.model.batch_decode_graph( - &token_ids, - &mut kv_refs, - graph_state, - DecodeGraphUse::Replay, - )?; - - if self.rank != 0 { - return Ok(Vec::new()); - } - sample_decode_rows( - ctx, - &graph_state.buffers.logits, - requests, - sample_seed, - &mut self.sample_scratch, - ) - } - - fn sample_final_prefill_chunk( - &mut self, - chunk: &TpPrefillChunkItem, - logits: &pegainfer_core::tensor::HiddenStates, - sample_seed: u64, - ) -> Result { - let cpu_logits = - snapshot_requested_logprobs(self.model.device_ctx(), logits, &[chunk.logprobs])?; - let params_refs = [&chunk.sampling_params]; - let tokens = pegainfer_sample::select_batch( - self.model.device_ctx(), - logits, - ¶ms_refs, - &[0], - sample_seed, - &mut self.sample_scratch, - )?; - let first_token = tokens[0]; - let first_token_logprob = cpu_logits[0].as_ref().and_then(|(row, top_k)| { - pegainfer_sample::token_logprob_from_row(row, first_token, *top_k) - }); - Ok(PrefillRequestResult { - request_id: chunk.request_id, - first_token, - first_token_logprob, - }) - } - - fn execute_decode( - &mut self, - requests: &[TpDecodeStepItem], - sample_seed: u64, - ) -> Result { - let requests = self.execute_decode_rows(requests, sample_seed)?; - if self.rank == 0 { - Ok(TpWorkerReply::Decode(DecodeResult { requests })) - } else { - Ok(TpWorkerReply::Ack) - } - } - - fn execute_decode_rows( - &mut self, - requests: &[TpDecodeStepItem], - sample_seed: u64, - ) -> Result> { - anyhow::ensure!( - !requests.is_empty(), - "Qwen3.5 TP decode command requires at least one request" - ); - validate_decode_requests(requests)?; - anyhow::ensure!( - requests.len() <= self.max_batch, - "Qwen3.5 TP decode batch {} exceeds worker capacity {}", - requests.len(), - self.max_batch - ); - - self.run_decode_batch(requests, sample_seed) - } - - fn execute_unified(&mut self, plan: &TpUnifiedPlan) -> Result { - validate_unified_worker_state(self, plan)?; - - // The command order is canonical across ranks. Sampling seeds are - // selected by the scheduler in decode-then-prefill order, independent - // of this forward order. - let prefill_requests = - self.execute_prefill_rows(&plan.prefill, plan.prefill_sample_seed)?; - let decode_requests = self.execute_decode_rows(&plan.decode, plan.decode_sample_seed)?; - - if self.rank == 0 { - Ok(TpWorkerReply::Unified(TpUnifiedResult { - prefill: PrefillResult { - requests: prefill_requests, - }, - decode: DecodeResult { - requests: decode_requests, - }, - })) - } else { - Ok(TpWorkerReply::Ack) - } - } - - fn ensure_prefill_state(&mut self, request_id: RequestId) -> Result { - if let Some(idx) = self.request_index(request_id) { - return Ok(idx); - } - let recurrent = RecurrentState::new( - self.model.device_ctx(), - self.model.config(), - self.model.geometry, - )?; - let state = TpRequestState { - request_id, - phase: TpRequestPhase::Prefilling, - kv: self.model.alloc_kv(), - recurrent: Some(recurrent), - }; - self.requests.push(state); - Ok(self.requests.len() - 1) - } - - fn request_index(&self, request_id: RequestId) -> Option { - self.requests - .iter() - .position(|state| state.request_id == request_id) - } - - /// One phase of the startup pre-capture sweep (graph mode only). - fn precapture_phase(&mut self, phase: PrecapturePhase) -> Result<()> { - match phase { - PrecapturePhase::Warmup => self.model.warmup_tp_collective(), - PrecapturePhase::Capture { bucket_idx } => { - self.precapture_bucket(bucket_idx, DecodeGraphUse::CaptureOnly) - } - PrecapturePhase::Launch { bucket_idx } => { - self.precapture_bucket(bucket_idx, DecodeGraphUse::Replay) - } - PrecapturePhase::Finalize => { - let graph_state = self.graph_state.as_ref().ok_or_else(|| { - anyhow::anyhow!("Qwen3.5 TP pre-capture Finalize without graph state") - })?; - for (bucket_idx, &bucket) in BATCH_BUCKETS.iter().enumerate() { - if bucket > graph_state.slot_states.len() { - break; - } - anyhow::ensure!( - graph_state.graphs[bucket_idx].is_captured(), - "Qwen3.5 TP decode graph pre-capture left bucket {bucket} uncaptured" - ); - } - Ok(()) - } - } - } - - /// Capture or launch one bucket with synthetic rows. Outputs are - /// discarded; the rows exist only to give the recorded kernels valid - /// addresses. One real row (token 0 at position 0 over a freshly - /// allocated one-page KV state) selects nothing — the bucket is passed - /// explicitly — and every other row is padding on the pool's reserved - /// padding page, exactly as when serving. The sweep therefore holds one - /// KV page at a time regardless of pool size or bucket. - fn precapture_bucket(&mut self, bucket_idx: usize, graph_use: DecodeGraphUse) -> Result<()> { - let bucket = BATCH_BUCKETS[bucket_idx]; - let graph_state = self.graph_state.as_mut().ok_or_else(|| { - anyhow::anyhow!("Qwen3.5 TP pre-capture on a worker without graph state") - })?; - anyhow::ensure!( - bucket <= graph_state.slot_states.len(), - "Qwen3.5 TP pre-capture bucket {bucket} exceeds {} slots", - graph_state.slot_states.len() - ); - let mut synthetic_kv = self.model.alloc_kv(); - let mut kv_refs = [&mut synthetic_kv]; - self.model.batch_decode_graph_padded( - &[0u32], - &mut kv_refs, - graph_state, - graph_use, - bucket, - )?; - // Capture acks only after the async cuGraphUpload lands; Launch acks - // only after the collectives drained. - self.model - .device_ctx() - .stream - .synchronize() - .map_err(|e| anyhow::anyhow!("Qwen3.5 TP pre-capture bucket {bucket} sync: {e}"))?; - Ok(()) - } - - /// Retire a request. Graph mode also applies the scheduler's slot - /// compaction (D2D move + occupancy assertions) so the slot layout stays - /// dense; any mismatch between the scheduler's claim and the worker's slot - /// map is a divergence and fails the command (poisoning the executor). - fn drop_request( - &mut self, - request_id: RequestId, - compaction: Option, - ) -> Result { - let Some(idx) = self.request_index(request_id) else { - anyhow::ensure!( - compaction.is_none(), - "Qwen3.5 TP drop of absent request {} carries a slot compaction", - request_id.get() - ); - return Ok(false); - }; - if let Some(graph_state) = self.graph_state.as_mut() { - match compaction { - Some(compaction) => { - let needs_move = slot_compact(&mut self.slot_map, request_id, compaction)?; - if needs_move { - graph_state.move_slot_within( - self.model.device_ctx(), - compaction.from, - compaction.to, - )?; - } - } - None => { - slot_release(&mut self.slot_map, request_id); - } - } - } - self.requests.swap_remove(idx); - Ok(true) - } -} - -/// Admit `request_id` to decode `slot`: the slot must be free (retirement and -/// compaction keep the map dense, so an occupied slot here is a scheduler -/// divergence). -fn slot_admit(owners: &mut [Option], slot: usize, request_id: RequestId) -> Result<()> { - let slot_count = owners.len(); - let owner = owners.get_mut(slot).ok_or_else(|| { - anyhow::anyhow!("Qwen3.5 TP decode slot {slot} exceeds worker slot map {slot_count}") - })?; - anyhow::ensure!( - owner.is_none(), - "Qwen3.5 TP decode slot {slot} still owned by request {} at admission of request {}", - owner.expect("checked").get(), - request_id.get() - ); - *owner = Some(request_id); - Ok(()) -} - -/// Clear `request_id`'s slot if it held one. Requests retired before their -/// first decode row never materialized a slot; that is not an error. -fn slot_release(owners: &mut [Option], request_id: RequestId) -> Option { - let slot = owners.iter().position(|owner| *owner == Some(request_id))?; - owners[slot] = None; - Some(slot) -} - -/// Apply the scheduler's slot compaction to the worker's slot map and report -/// whether a GPU state move is needed. Both requests may legitimately be -/// unmaterialized (retired/compacted before their first decode row), but a -/// materialized slot must hold exactly the request the scheduler claims. -fn slot_compact( - owners: &mut [Option], - dropped: RequestId, - compaction: TpSlotCompaction, -) -> Result { - let TpSlotCompaction { - moved_request_id, - from, - to, - } = compaction; - anyhow::ensure!( - from < owners.len() && to < owners.len(), - "Qwen3.5 TP slot compaction {from} -> {to} exceeds worker slot map {}", - owners.len() - ); - let dropped_owner = owners[to]; - let moved_owner = owners[from]; - if let Some(owner) = dropped_owner { - anyhow::ensure!( - owner == dropped, - "Qwen3.5 TP slot {to} holds request {} where the scheduler dropped request {}", - owner.get(), - dropped.get() - ); - } - if let Some(owner) = moved_owner { - anyhow::ensure!( - owner == moved_request_id, - "Qwen3.5 TP slot {from} holds request {} where the scheduler moved request {}", - owner.get(), - moved_request_id.get() - ); - } - owners[to] = moved_owner; - owners[from] = None; - Ok(moved_owner.is_some()) -} - fn validate_prefill_chunks(chunks: &[TpPrefillChunkItem]) -> Result<()> { let mut seen = HashSet::with_capacity(chunks.len()); for chunk in chunks { @@ -2175,60 +1090,6 @@ fn validate_prefill_chunks(chunks: &[TpPrefillChunkItem]) -> Result<()> { /// Worker request states in decode-row order: `row_of_state[i]` is the row /// `states[i]` occupies in this command, `None` when it is not part of it. -fn states_in_row_order<'a>( - states: &'a mut [TpRequestState], - row_of_state: &[Option], -) -> Vec<&'a mut TpRequestState> { - let mut rows: Vec<(usize, &'a mut TpRequestState)> = states - .iter_mut() - .zip(row_of_state) - .filter_map(|(state, row)| row.map(|row| (row, state))) - .collect(); - rows.sort_unstable_by_key(|(row, _)| *row); - rows.into_iter().map(|(_, state)| state).collect() -} - -/// Rank-0 sampling pass over one decode batch: snapshot the requested logprob -/// rows, select one token per row, and pair each token with its logprob. -fn sample_decode_rows( - ctx: &pegainfer_core::tensor::DeviceContext, - logits: &pegainfer_core::tensor::HiddenStates, - requests: &[TpDecodeStepItem], - sample_seed: u64, - scratch: &mut pegainfer_sample::SampleScratch, -) -> Result> { - let bs = requests.len(); - let requested_logprobs: Vec> = - requests.iter().map(|request| request.logprobs).collect(); - let cpu_logits = snapshot_requested_logprobs(ctx, logits, &requested_logprobs)?; - let params_refs: Vec<&SamplingParams> = requests - .iter() - .map(|request| &request.sampling_params) - .collect(); - let steps = vec![0u64; bs]; - let tokens = - pegainfer_sample::select_batch(ctx, logits, ¶ms_refs, &steps, sample_seed, scratch)?; - anyhow::ensure!( - tokens.len() == bs, - "Qwen3.5 TP decode sampling returned {} tokens for {bs} rows", - tokens.len() - ); - Ok(requests - .iter() - .enumerate() - .map(|(row, request)| { - let logprob = cpu_logits[row].as_ref().and_then(|(logits_row, top_k)| { - pegainfer_sample::token_logprob_from_row(logits_row, tokens[row], *top_k) - }); - DecodeRequestResult { - request_id: request.request_id, - token: tokens[row], - logprob, - } - }) - .collect()) -} - fn validate_decode_requests(requests: &[TpDecodeStepItem]) -> Result<()> { let mut seen = HashSet::with_capacity(requests.len()); for request in requests { @@ -2629,37 +1490,6 @@ fn fatal_tp_abort(message: &str) -> ! { std::process::abort(); } -struct CublasThreadGuard; - -impl Drop for CublasThreadGuard { - fn drop(&mut self) { - unsafe { - crate::ffi::cublas_destroy(); - } - } -} - -fn bind_worker_thread(model: &Qwen35Model) -> Result { - let ctx = model.device_ctx(); - unsafe { - let err = crate::ffi::cuda_set_device(ctx.device_ordinal as i32); - if err != 0 { - return Err(anyhow::anyhow!( - "Failed to set CUDA device {} on Qwen3.5 TP worker thread: cudaError={}", - ctx.device_ordinal, - err - )); - } - } - ctx.ctx.bind_to_thread().map_err(|e| { - anyhow::anyhow!("Failed to bind CUDA context to Qwen3.5 TP worker thread: {e}") - })?; - unsafe { - crate::ffi::cublas_init(); - } - Ok(CublasThreadGuard) -} - #[cfg(test)] mod tests { use super::*; diff --git a/pegainfer-qwen35/src/tp_executor/worker.rs b/pegainfer-qwen35/src/tp_executor/worker.rs new file mode 100644 index 000000000..e2b90f86e --- /dev/null +++ b/pegainfer-qwen35/src/tp_executor/worker.rs @@ -0,0 +1,1197 @@ +//! Tensor-parallel worker runtime: the per-rank worker thread, startup +//! gating, the NCCL startup watchdog, and the per-rank command loop. + +use super::*; + +const TP_NCCL_STARTUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); +const TP_WORKER_SHUTDOWN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); +const TP_RUNTIME_MEMORY_RESERVE_BYTES: usize = 512 * 1024 * 1024; + +pub(super) fn spawn_nccl_startup_watchdog() -> Result<(mpsc::SyncSender<()>, JoinHandle<()>)> { + let (done_tx, done_rx) = mpsc::sync_channel(1); + let watchdog = thread::Builder::new() + .name("qwen35-tp-nccl-startup-watchdog".into()) + .spawn(move || { + if done_rx.recv_timeout(TP_NCCL_STARTUP_TIMEOUT).is_ok() { + return; + } + eprintln!( + "Qwen3.5 TP NCCL startup did not complete within {}s; aborting", + TP_NCCL_STARTUP_TIMEOUT.as_secs() + ); + log::error!( + "Qwen3.5 TP NCCL startup did not complete within {}s; aborting", + TP_NCCL_STARTUP_TIMEOUT.as_secs() + ); + std::process::abort(); + }) + .map_err(|err| anyhow::anyhow!("failed to spawn Qwen3.5 TP NCCL watchdog: {err}"))?; + Ok((done_tx, watchdog)) +} + +#[allow(clippy::needless_pass_by_value)] +pub(super) fn disarm_nccl_startup_watchdog( + done_tx: mpsc::SyncSender<()>, + watchdog: JoinHandle<()>, +) -> Result<()> { + done_tx + .send(()) + .map_err(|_| anyhow::anyhow!("Qwen3.5 TP NCCL watchdog exited unexpectedly"))?; + watchdog + .join() + .map_err(|_| anyhow::anyhow!("Qwen3.5 TP NCCL watchdog panicked")) +} + +pub(super) struct TpWorker { + pub(super) tx: mpsc::Sender, + handle: Option>, + done: mpsc::Receiver<()>, +} + +#[derive(Clone, Copy, Default, PartialEq, Eq)] +enum TpStartupDecision { + #[default] + Pending, + Connect, + Cancel, +} + +#[derive(Default)] +pub(super) struct TpStartupGate { + decision: Mutex, + changed: Condvar, +} + +impl TpStartupGate { + pub(super) fn connect(&self) { + self.set(TpStartupDecision::Connect); + } + + pub(super) fn cancel(&self) { + self.set(TpStartupDecision::Cancel); + } + + pub(super) fn wait(&self) -> bool { + let mut decision = self.decision.lock().unwrap_or_else(PoisonError::into_inner); + while *decision == TpStartupDecision::Pending { + decision = self + .changed + .wait(decision) + .unwrap_or_else(PoisonError::into_inner); + } + *decision == TpStartupDecision::Connect + } + + fn set(&self, next: TpStartupDecision) { + let mut decision = self.decision.lock().unwrap_or_else(PoisonError::into_inner); + if *decision == TpStartupDecision::Pending { + *decision = next; + self.changed.notify_all(); + } + } +} + +impl TpWorker { + #[allow(clippy::too_many_arguments)] + #[allow(clippy::type_complexity)] + pub(super) fn spawn( + rank: usize, + world_size: usize, + model: Qwen35Model, + max_batch: usize, + max_prefill_tokens: usize, + graph_enabled: bool, + nccl_id: cudarc::nccl::safe::Id, + startup_gate: Arc, + effective_max_batch: Arc, + poison: Arc, + ) -> Result<( + Self, + mpsc::Receiver>, + mpsc::Receiver>, + )> { + let (tx, rx) = mpsc::channel(); + let (preflight_tx, preflight_rx) = mpsc::channel(); + let (startup_tx, startup_rx) = mpsc::channel(); + let (done_tx, done_rx) = mpsc::channel(); + let panic_poison = Arc::clone(&poison); + let handle = thread::Builder::new() + .name(format!("qwen35-tp-rank-{rank}")) + .spawn(move || { + let outcome = catch_unwind(AssertUnwindSafe(|| { + let prepared = TpWorkerPrepared::new( + rank, + world_size, + model, + max_batch, + max_prefill_tokens, + graph_enabled, + ); + let prepared = match prepared { + Ok((prepared, rank_max_batch)) => { + let _ = preflight_tx.send(Ok(rank_max_batch)); + prepared + } + Err(err) => { + let _ = preflight_tx.send(Err(err)); + return; + } + }; + if !startup_gate.wait() { + return; + } + let max_batch = effective_max_batch.load(Ordering::Acquire); + match prepared.connect(nccl_id, max_batch, graph_enabled, poison) { + Ok(mut state) => { + let _ = startup_tx.send(Ok(())); + state.run(rx); + } + Err(err) => { + let _ = startup_tx.send(Err(err)); + } + } + })); + if outcome.is_err() { + panic_poison.poison(format!("worker rank {rank} panicked")); + } + let _ = done_tx.send(()); + }) + .map_err(|e| anyhow::anyhow!("failed to spawn Qwen3.5 TP worker {rank}: {e}"))?; + + Ok(( + Self { + tx, + handle: Some(handle), + done: done_rx, + }, + preflight_rx, + startup_rx, + )) + } + + pub(super) fn send(&self, command: TpWorkerCommand) -> Result<()> { + self.tx + .send(command) + .map_err(|_| anyhow::anyhow!("Qwen3.5 TP worker channel closed")) + } + + pub(super) fn join_bounded(&mut self) { + if self.handle.is_none() { + return; + } + if self.done.recv_timeout(TP_WORKER_SHUTDOWN_TIMEOUT).is_err() { + fatal_tp_abort("Qwen3.5 TP worker did not exit during bounded shutdown"); + } + if let Some(handle) = self.handle.take() { + let _ = handle.join(); + } + } +} + +impl Drop for TpWorker { + fn drop(&mut self) { + let _ = self.tx.send(TpWorkerCommand::Shutdown); + self.join_bounded(); + } +} + +pub(super) struct TpWorkerState { + rank: usize, + _world_size: usize, + pub(super) max_batch: usize, + /// Before `model` on purpose: NCCL comm teardown polls until every graph + /// that recorded its collectives is destroyed, so the decode graphs must + /// drop before `model.tp_comm` (qwen3 teardown-hang precedent). + graph_state: Option, + model: Qwen35Model, + pub(super) requests: Vec, + /// Graph-mode slot ownership: `slot_map[i]` is the request whose recurrent + /// state lives in `graph_state.slot_states[i]`. The scheduler owns slot + /// assignment and compaction; the worker only applies and checks them. + /// Empty in eager mode. + slot_map: Vec>, + decode_buffers: BatchDecodeBuffers35, + /// Eager decode GDR pointer tables: allocated once at capacity, refilled + /// with the live rows every step. + decode_pointer_tables: LinearStatePointerTables, + sample_scratch: pegainfer_sample::SampleScratch, + _cublas_guard: CublasThreadGuard, + poison: Arc, +} + +struct TpWorkerPrepared { + rank: usize, + world_size: usize, + max_batch: usize, + model: Qwen35Model, + decode_buffers: BatchDecodeBuffers35, + sample_scratch: pegainfer_sample::SampleScratch, + cublas_guard: CublasThreadGuard, +} + +pub(super) struct TpRequestState { + pub(super) request_id: RequestId, + pub(super) phase: TpRequestPhase, + kv: KvState, + /// Prefill-owned recurrent state. Graph mode moves it into the decode slot + /// on the request's first decode row (`None` afterwards); the eager path + /// keeps it for the request's whole lifetime. + recurrent: Option, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(super) enum TpRequestPhase { + Prefilling, + Decoding, +} + +#[cfg(test)] +#[derive(Clone, Debug, PartialEq, Eq)] +pub(super) struct WorkerStateSnapshot { + pub(super) rank: usize, + pub(super) request_count: usize, + pub(super) requests: Vec<(RequestId, TpRequestPhase)>, +} + +impl TpWorkerPrepared { + fn new( + rank: usize, + world_size: usize, + model: Qwen35Model, + requested_max_batch: usize, + max_prefill_tokens: usize, + graph_enabled: bool, + ) -> Result<(Self, usize)> { + let cublas_guard = bind_worker_thread(&model)?; + let (free_bytes, total_bytes) = model + .device_ctx() + .ctx + .mem_get_info() + .map_err(|err| anyhow::anyhow!("failed to query TP rank {rank} memory: {err}"))?; + // Recurrent state is rank-local, so worker capacity math uses the + // local value-head/qkv sizes. + let recurrent_bytes = RecurrentState::allocation_bytes(model.config(), model.geometry); + let prefill_scratch_tokens = prefill_scratch_tokens(max_prefill_tokens); + let prefill_scratch_bytes = GdrChunkwiseScratch35::estimate_bytes( + model.config(), + model.geometry, + prefill_scratch_tokens, + ); + // Graph mode pre-allocates one fixed-address slot state per decode + // bucket position up front; reserve that before sizing per-request + // (prefill-transient) state capacity. The reserve must track the + // bucket of the *effective* batch, not the requested one: reserving + // for `bucket_for(requested)` can starve a tight-memory rank down to + // zero capacity. Iterate the bucket downward until it stabilises — + // the bucket only shrinks, so this converges — and clamp the fitted + // batch to the reserved bucket so the later `bucket_for(effective)` + // graph allocation never exceeds the reserve. + let max_batch = if graph_enabled { + let mut slot_bucket = bucket_for(requested_max_batch); + loop { + let reserve = slot_bucket * recurrent_bytes; + let candidate = effective_recurrent_capacity( + requested_max_batch, + free_bytes.saturating_sub(reserve), + recurrent_bytes, + TP_RUNTIME_MEMORY_RESERVE_BYTES, + prefill_scratch_bytes, + ); + let fitted = candidate.min(slot_bucket); + let next = bucket_for(fitted); + if next >= slot_bucket { + break fitted; + } + slot_bucket = next; + } + } else { + effective_recurrent_capacity( + requested_max_batch, + free_bytes, + recurrent_bytes, + TP_RUNTIME_MEMORY_RESERVE_BYTES, + prefill_scratch_bytes, + ) + }; + anyhow::ensure!( + max_batch > 0, + "Qwen3.5 TP rank {rank} has {} MiB free after fixed buffers, but one recurrent request needs {} MiB plus {} MiB runtime reserve and {} MiB prefill scratch for {} tokens", + free_bytes / (1024 * 1024), + recurrent_bytes / (1024 * 1024), + TP_RUNTIME_MEMORY_RESERVE_BYTES / (1024 * 1024), + prefill_scratch_bytes / (1024 * 1024), + prefill_scratch_tokens, + ); + log::info!( + "Qwen3.5 TP rank {rank} recurrent capacity: requested={requested_max_batch}, effective={max_batch}, per_request={:.3} MiB, free={:.0} MiB/{:.0} MiB, runtime_reserve={} MiB, prefill_tokens={}, prefill_scratch={:.0} MiB", + recurrent_bytes as f64 / 1024.0 / 1024.0, + free_bytes as f64 / 1024.0 / 1024.0, + total_bytes as f64 / 1024.0 / 1024.0, + TP_RUNTIME_MEMORY_RESERVE_BYTES / (1024 * 1024), + prefill_scratch_tokens, + prefill_scratch_bytes as f64 / 1024.0 / 1024.0, + ); + let decode_buffers = model.create_batch_decode_buffers_with_capacity(max_batch)?; + let sample_scratch = pegainfer_sample::SampleScratch::new( + model.device_ctx(), + model.config().selection_vocab, + max_batch, + )?; + Ok(( + Self { + rank, + world_size, + max_batch, + model, + decode_buffers, + sample_scratch, + cublas_guard, + }, + max_batch, + )) + } + + fn connect( + self, + nccl_id: cudarc::nccl::safe::Id, + effective_max_batch: usize, + graph_enabled: bool, + poison: Arc, + ) -> Result { + let Self { + rank, + world_size, + max_batch, + mut model, + decode_buffers, + sample_scratch, + cublas_guard, + } = self; + anyhow::ensure!( + effective_max_batch > 0 && effective_max_batch <= max_batch, + "Qwen3.5 TP rank {rank} effective max_batch {effective_max_batch} exceeds local capacity {max_batch}" + ); + let comm = cudarc::nccl::safe::Comm::from_rank( + model.device_ctx().stream.clone(), + rank, + world_size, + nccl_id, + ) + .map_err(|e| anyhow::anyhow!("failed to initialize Qwen3.5 TP NCCL rank {rank}: {e:?}"))?; + model.attach_tp_comm(comm); + let decode_pointer_tables = LinearStatePointerTables::with_capacity( + model.device_ctx(), + model.config(), + effective_max_batch, + "Qwen3.5 TP eager decode", + )?; + let (graph_state, slot_map) = if graph_enabled { + // cuBLASLt plans are thread-local: tune the decode bucket GEMMs on + // this worker thread now so plan selection never runs inside + // cuStreamBeginCapture during the pre-capture sweep. + model.tune_decode_gemm_algos()?; + let slots = bucket_for(effective_max_batch); + let graph_state = model.create_batch_decode_graph_state_with_capacity(slots)?; + (Some(graph_state), vec![None; slots]) + } else { + (None, Vec::new()) + }; + Ok(TpWorkerState { + rank, + _world_size: world_size, + max_batch: effective_max_batch, + graph_state, + model, + requests: Vec::new(), + slot_map, + decode_buffers, + decode_pointer_tables, + sample_scratch, + _cublas_guard: cublas_guard, + poison, + }) + } +} + +fn prefill_scratch_tokens(max_prefill_tokens: usize) -> usize { + max_prefill_tokens.min(PREFILL_CHUNK_LEN) +} + +fn effective_recurrent_capacity( + requested_max_batch: usize, + free_bytes: usize, + recurrent_bytes_per_request: usize, + runtime_reserve_bytes: usize, + prefill_scratch_bytes: usize, +) -> usize { + if recurrent_bytes_per_request == 0 { + return requested_max_batch; + } + requested_max_batch.min( + free_bytes + .saturating_sub(runtime_reserve_bytes) + .saturating_sub(prefill_scratch_bytes) + / recurrent_bytes_per_request, + ) +} + +impl TpWorkerState { + #[allow(clippy::needless_pass_by_value)] + fn run(&mut self, rx: mpsc::Receiver) { + while let Ok(command) = rx.recv() { + let fatal = match command { + TpWorkerCommand::Ping { resp } => { + self.respond(resp, "ping", Ok(TpWorkerReply::Ack)) + } + TpWorkerCommand::RunPrefillChunks { + chunks, + sample_seed, + start, + resp, + } => { + if start.wait() == TpCommandDecision::Cancel { + false + } else { + let result = self.execute_prefill_chunks(&chunks, sample_seed); + self.respond(resp, "prefill", result) + } + } + TpWorkerCommand::RunDecodeStep { + requests, + sample_seed, + start, + resp, + } => { + if start.wait() == TpCommandDecision::Cancel { + false + } else { + let result = self.execute_decode(&requests, sample_seed); + self.respond(resp, "decode", result) + } + } + TpWorkerCommand::RunUnifiedStep { plan, start, resp } => { + if start.wait() == TpCommandDecision::Cancel { + false + } else { + let result = self.execute_unified(&plan); + self.respond(resp, "unified step", result) + } + } + TpWorkerCommand::DropRequest { + request_id, + compaction, + start, + resp, + } => { + if start.wait() == TpCommandDecision::Cancel { + false + } else { + let result = self + .drop_request(request_id, compaction) + .map(|existed| TpWorkerReply::DropAck { existed }); + self.respond(resp, "drop request", result) + } + } + TpWorkerCommand::Precapture { phase, start, resp } => { + if start.wait() == TpCommandDecision::Cancel { + false + } else { + let result = self.precapture_phase(phase).map(|()| TpWorkerReply::Ack); + self.respond(resp, "decode graph precapture", result) + } + } + #[cfg(test)] + TpWorkerCommand::SnapshotState { resp } => { + let snapshot = WorkerStateSnapshot { + rank: self.rank, + request_count: self.requests.len(), + requests: self + .requests + .iter() + .map(|state| (state.request_id, state.phase)) + .collect(), + }; + self.respond( + resp, + "snapshot state", + Ok(TpWorkerReply::Snapshot(snapshot)), + ) + } + #[cfg(test)] + TpWorkerCommand::RemoveRequestStateForTest { request_id, resp } => { + let _ = resp.send(self.drop_request(request_id, None).unwrap_or(false)); + false + } + #[cfg(test)] + TpWorkerCommand::DisconnectForTest { ready } => { + let _ = ready.send(()); + break; + } + TpWorkerCommand::Shutdown => break, + }; + if fatal { + break; + } + } + } + + #[allow(clippy::needless_pass_by_value)] + fn respond( + &self, + resp: mpsc::Sender, + operation: &'static str, + result: Result, + ) -> bool { + match result { + Ok(reply) => { + let _ = resp.send(TpWorkerResponse { + rank: self.rank, + result: Ok(reply), + }); + false + } + Err(err) => { + let reason = self.poison.poison(format!( + "rank {} failed during {operation}: {err:#}", + self.rank + )); + let _ = resp.send(TpWorkerResponse { + rank: self.rank, + result: Err(anyhow::anyhow!(reason)), + }); + true + } + } + } + + fn execute_prefill_chunks( + &mut self, + chunks: &[TpPrefillChunkItem], + sample_seed: u64, + ) -> Result { + let requests = self.execute_prefill_rows(chunks, sample_seed)?; + if self.rank == 0 { + Ok(TpWorkerReply::Prefill(PrefillResult { requests })) + } else { + Ok(TpWorkerReply::Ack) + } + } + + fn execute_prefill_rows( + &mut self, + chunks: &[TpPrefillChunkItem], + sample_seed: u64, + ) -> Result> { + anyhow::ensure!( + !chunks.is_empty(), + "Qwen3.5 TP prefill chunk command requires at least one chunk" + ); + validate_prefill_chunks(chunks)?; + let new_requests = chunks + .iter() + .filter(|chunk| self.request_index(chunk.request_id).is_none()) + .count(); + anyhow::ensure!( + self.requests.len() + new_requests <= self.max_batch, + "Qwen3.5 TP prefill chunks would exceed worker capacity {}", + self.max_batch + ); + + let mut primary_results = Vec::new(); + let mut final_row_idx = 0usize; + for chunk in chunks { + let state_idx = self.ensure_prefill_state(chunk.request_id)?; + let state = &mut self.requests[state_idx]; + anyhow::ensure!( + state.phase == TpRequestPhase::Prefilling, + "Qwen3.5 TP request {} is already in decode state", + chunk.request_id.get() + ); + + let prompt = [chunk.prompt_tokens.as_slice()]; + let mut recurrent_refs = vec![ + state + .recurrent + .as_mut() + .expect("prefill-phase TP request owns its recurrent state"), + ]; + let logits = self.model.batch_prefill_logits( + &prompt, + std::slice::from_mut(&mut state.kv), + &mut recurrent_refs, + )?; + + if chunk.finish_prefill { + if self.rank == 0 { + // TP prefill samples final chunks one row at a time. Offset + // by the final-row index so rows from the same command do + // not reuse the same sampling stream. + let row_seed = sample_seed.wrapping_add(final_row_idx as u64); + let result = self.sample_final_prefill_chunk(chunk, &logits, row_seed)?; + primary_results.push(result); + } + final_row_idx += 1; + self.requests[state_idx].phase = TpRequestPhase::Decoding; + } + } + + Ok(primary_results) + } + + /// Run one batched eager decode step over all rows in command order: a + /// single forward for the whole batch on every rank, then (rank 0 only) + /// one batched sampling pass over the per-row sampling params. Returns one + /// result row per request in command order on rank 0, empty elsewhere. + fn run_decode_batch( + &mut self, + requests: &[TpDecodeStepItem], + sample_seed: u64, + ) -> Result> { + let bs = requests.len(); + if bs == 0 { + return Ok(Vec::new()); + } + if self.graph_state.is_some() { + return self.run_decode_batch_graph(requests, sample_seed); + } + + // Resolve the worker state slot of every row in command order. + // Decode request ids are unique within one command + // (validate_decode_requests), so each slot is borrowed at most once. + let mut row_of_state: Vec> = vec![None; self.requests.len()]; + for (row, request) in requests.iter().enumerate() { + let state_idx = self.request_index(request.request_id).ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP decode request {} has no worker state", + request.request_id.get() + ) + })?; + anyhow::ensure!( + self.requests[state_idx].phase == TpRequestPhase::Decoding, + "Qwen3.5 TP request {} is not ready for decode", + request.request_id.get() + ); + debug_assert!(row_of_state[state_idx].is_none()); + row_of_state[state_idx] = Some(row); + } + let mut kv_refs: Vec<&mut KvState> = Vec::with_capacity(bs); + let mut recurrent_refs: Vec<&mut RecurrentState> = Vec::with_capacity(bs); + for state in states_in_row_order(&mut self.requests, &row_of_state) { + let TpRequestState { kv, recurrent, .. } = state; + kv_refs.push(kv); + recurrent_refs.push( + recurrent + .as_mut() + .expect("eager TP decode request owns its recurrent state"), + ); + } + + // GDR pointer tables over the full decode batch: allocated once at + // capacity, refilled from the live rows every step (H2D only), so + // swap_remove retirement between steps can never leave a stale row + // addressed and no step pays for device allocations. + self.decode_pointer_tables.refill_from_recurrent_refs( + self.model.device_ctx(), + &mut recurrent_refs, + bs, + "Qwen3.5 TP eager decode", + )?; + let token_ids: Vec = requests.iter().map(|request| request.token_id).collect(); + self.model.batch_decode_eager_logits( + &token_ids, + &mut kv_refs, + &mut recurrent_refs, + &self.decode_pointer_tables, + &mut self.decode_buffers, + )?; + + if self.rank != 0 { + return Ok(Vec::new()); + } + sample_decode_rows( + self.model.device_ctx(), + &self.decode_buffers.logits, + requests, + sample_seed, + &mut self.sample_scratch, + ) + } + + /// CUDA Graph decode step under TP: replay-only (every bucket was recorded + /// by the startup pre-capture sweep), one forward for the whole batch on + /// every rank, then (rank 0 only) the same batched host-side sampling pass + /// as the eager path. + /// + /// Rows must arrive in the scheduler-owned dense slot order + /// (`slot_idx == row`). On a request's first decode row its prefill-owned + /// recurrent state is D2D-copied into `graph_state.slot_states[slot]` and + /// the per-request allocation is dropped; the persistent linear-state + /// pointer tables then keep every replay reading the fixed slot addresses. + fn run_decode_batch_graph( + &mut self, + requests: &[TpDecodeStepItem], + sample_seed: u64, + ) -> Result> { + let bs = requests.len(); + let graph_state = self + .graph_state + .as_mut() + .expect("graph decode arm requires graph state"); + let ctx = self.model.device_ctx(); + + // Resolve the worker state of every row, enforce dense slot order, and + // admit first-decode rows into their slots. Decode request ids are + // unique within one command (validate_decode_requests), so each slot + // is borrowed at most once. + let mut row_of_state: Vec> = vec![None; self.requests.len()]; + for (row, request) in requests.iter().enumerate() { + anyhow::ensure!( + request.slot_idx == Some(row), + "Qwen3.5 TP graph decode row {row} carries slot {:?}; rows must arrive in dense slot order 0..{bs}", + request.slot_idx + ); + let state_idx = self + .requests + .iter() + .position(|state| state.request_id == request.request_id) + .ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP decode request {} has no worker state", + request.request_id.get() + ) + })?; + anyhow::ensure!( + self.requests[state_idx].phase == TpRequestPhase::Decoding, + "Qwen3.5 TP request {} is not ready for decode", + request.request_id.get() + ); + debug_assert!(row_of_state[state_idx].is_none()); + row_of_state[state_idx] = Some(row); + + if self.slot_map.get(row).copied().flatten() == Some(request.request_id) { + anyhow::ensure!( + self.requests[state_idx].recurrent.is_none(), + "Qwen3.5 TP request {} was admitted to slot {row} but still owns prefill recurrent state", + request.request_id.get() + ); + } else { + slot_admit(&mut self.slot_map, row, request.request_id)?; + let recurrent = self.requests[state_idx].recurrent.take().ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP request {} lost its prefill recurrent state before slot admission", + request.request_id.get() + ) + })?; + graph_state.copy_state_to_slot(ctx, &recurrent, row)?; + } + } + + // KV refs in row (slot) order; page tables stay per-step H2D via + // sync_paged_meta inside batch_decode_graph. + let mut kv_refs: Vec<&mut KvState> = states_in_row_order(&mut self.requests, &row_of_state) + .into_iter() + .map(|state| &mut state.kv) + .collect(); + let token_ids: Vec = requests.iter().map(|request| request.token_id).collect(); + self.model.batch_decode_graph( + &token_ids, + &mut kv_refs, + graph_state, + DecodeGraphUse::Replay, + )?; + + if self.rank != 0 { + return Ok(Vec::new()); + } + sample_decode_rows( + ctx, + &graph_state.buffers.logits, + requests, + sample_seed, + &mut self.sample_scratch, + ) + } + + fn sample_final_prefill_chunk( + &mut self, + chunk: &TpPrefillChunkItem, + logits: &pegainfer_core::tensor::HiddenStates, + sample_seed: u64, + ) -> Result { + let cpu_logits = + snapshot_requested_logprobs(self.model.device_ctx(), logits, &[chunk.logprobs])?; + let params_refs = [&chunk.sampling_params]; + let tokens = pegainfer_sample::select_batch( + self.model.device_ctx(), + logits, + ¶ms_refs, + &[0], + sample_seed, + &mut self.sample_scratch, + )?; + let first_token = tokens[0]; + let first_token_logprob = cpu_logits[0].as_ref().and_then(|(row, top_k)| { + pegainfer_sample::token_logprob_from_row(row, first_token, *top_k) + }); + Ok(PrefillRequestResult { + request_id: chunk.request_id, + first_token, + first_token_logprob, + }) + } + + fn execute_decode( + &mut self, + requests: &[TpDecodeStepItem], + sample_seed: u64, + ) -> Result { + let requests = self.execute_decode_rows(requests, sample_seed)?; + if self.rank == 0 { + Ok(TpWorkerReply::Decode(DecodeResult { requests })) + } else { + Ok(TpWorkerReply::Ack) + } + } + + fn execute_decode_rows( + &mut self, + requests: &[TpDecodeStepItem], + sample_seed: u64, + ) -> Result> { + anyhow::ensure!( + !requests.is_empty(), + "Qwen3.5 TP decode command requires at least one request" + ); + validate_decode_requests(requests)?; + anyhow::ensure!( + requests.len() <= self.max_batch, + "Qwen3.5 TP decode batch {} exceeds worker capacity {}", + requests.len(), + self.max_batch + ); + + self.run_decode_batch(requests, sample_seed) + } + + fn execute_unified(&mut self, plan: &TpUnifiedPlan) -> Result { + validate_unified_worker_state(self, plan)?; + + // The command order is canonical across ranks. Sampling seeds are + // selected by the scheduler in decode-then-prefill order, independent + // of this forward order. + let prefill_requests = + self.execute_prefill_rows(&plan.prefill, plan.prefill_sample_seed)?; + let decode_requests = self.execute_decode_rows(&plan.decode, plan.decode_sample_seed)?; + + if self.rank == 0 { + Ok(TpWorkerReply::Unified(TpUnifiedResult { + prefill: PrefillResult { + requests: prefill_requests, + }, + decode: DecodeResult { + requests: decode_requests, + }, + })) + } else { + Ok(TpWorkerReply::Ack) + } + } + + fn ensure_prefill_state(&mut self, request_id: RequestId) -> Result { + if let Some(idx) = self.request_index(request_id) { + return Ok(idx); + } + let recurrent = RecurrentState::new( + self.model.device_ctx(), + self.model.config(), + self.model.geometry, + )?; + let state = TpRequestState { + request_id, + phase: TpRequestPhase::Prefilling, + kv: self.model.alloc_kv(), + recurrent: Some(recurrent), + }; + self.requests.push(state); + Ok(self.requests.len() - 1) + } + + pub(super) fn request_index(&self, request_id: RequestId) -> Option { + self.requests + .iter() + .position(|state| state.request_id == request_id) + } + + /// One phase of the startup pre-capture sweep (graph mode only). + fn precapture_phase(&mut self, phase: PrecapturePhase) -> Result<()> { + match phase { + PrecapturePhase::Warmup => self.model.warmup_tp_collective(), + PrecapturePhase::Capture { bucket_idx } => { + self.precapture_bucket(bucket_idx, DecodeGraphUse::CaptureOnly) + } + PrecapturePhase::Launch { bucket_idx } => { + self.precapture_bucket(bucket_idx, DecodeGraphUse::Replay) + } + PrecapturePhase::Finalize => { + let graph_state = self.graph_state.as_ref().ok_or_else(|| { + anyhow::anyhow!("Qwen3.5 TP pre-capture Finalize without graph state") + })?; + for (bucket_idx, &bucket) in BATCH_BUCKETS.iter().enumerate() { + if bucket > graph_state.slot_states.len() { + break; + } + anyhow::ensure!( + graph_state.graphs[bucket_idx].is_captured(), + "Qwen3.5 TP decode graph pre-capture left bucket {bucket} uncaptured" + ); + } + Ok(()) + } + } + } + + /// Capture or launch one bucket with synthetic rows. Outputs are + /// discarded; the rows exist only to give the recorded kernels valid + /// addresses. One real row (token 0 at position 0 over a freshly + /// allocated one-page KV state) selects nothing — the bucket is passed + /// explicitly — and every other row is padding on the pool's reserved + /// padding page, exactly as when serving. The sweep therefore holds one + /// KV page at a time regardless of pool size or bucket. + fn precapture_bucket(&mut self, bucket_idx: usize, graph_use: DecodeGraphUse) -> Result<()> { + let bucket = BATCH_BUCKETS[bucket_idx]; + let graph_state = self.graph_state.as_mut().ok_or_else(|| { + anyhow::anyhow!("Qwen3.5 TP pre-capture on a worker without graph state") + })?; + anyhow::ensure!( + bucket <= graph_state.slot_states.len(), + "Qwen3.5 TP pre-capture bucket {bucket} exceeds {} slots", + graph_state.slot_states.len() + ); + let mut synthetic_kv = self.model.alloc_kv(); + let mut kv_refs = [&mut synthetic_kv]; + self.model.batch_decode_graph_padded( + &[0u32], + &mut kv_refs, + graph_state, + graph_use, + bucket, + )?; + // Capture acks only after the async cuGraphUpload lands; Launch acks + // only after the collectives drained. + self.model + .device_ctx() + .stream + .synchronize() + .map_err(|e| anyhow::anyhow!("Qwen3.5 TP pre-capture bucket {bucket} sync: {e}"))?; + Ok(()) + } + + /// Retire a request. Graph mode also applies the scheduler's slot + /// compaction (D2D move + occupancy assertions) so the slot layout stays + /// dense; any mismatch between the scheduler's claim and the worker's slot + /// map is a divergence and fails the command (poisoning the executor). + fn drop_request( + &mut self, + request_id: RequestId, + compaction: Option, + ) -> Result { + let Some(idx) = self.request_index(request_id) else { + anyhow::ensure!( + compaction.is_none(), + "Qwen3.5 TP drop of absent request {} carries a slot compaction", + request_id.get() + ); + return Ok(false); + }; + if let Some(graph_state) = self.graph_state.as_mut() { + match compaction { + Some(compaction) => { + let needs_move = slot_compact(&mut self.slot_map, request_id, compaction)?; + if needs_move { + graph_state.move_slot_within( + self.model.device_ctx(), + compaction.from, + compaction.to, + )?; + } + } + None => { + slot_release(&mut self.slot_map, request_id); + } + } + } + self.requests.swap_remove(idx); + Ok(true) + } +} + +/// Admit `request_id` to decode `slot`: the slot must be free (retirement and +/// compaction keep the map dense, so an occupied slot here is a scheduler +/// divergence). +pub(super) fn slot_admit( + owners: &mut [Option], + slot: usize, + request_id: RequestId, +) -> Result<()> { + let slot_count = owners.len(); + let owner = owners.get_mut(slot).ok_or_else(|| { + anyhow::anyhow!("Qwen3.5 TP decode slot {slot} exceeds worker slot map {slot_count}") + })?; + anyhow::ensure!( + owner.is_none(), + "Qwen3.5 TP decode slot {slot} still owned by request {} at admission of request {}", + owner.expect("checked").get(), + request_id.get() + ); + *owner = Some(request_id); + Ok(()) +} + +/// Clear `request_id`'s slot if it held one. Requests retired before their +/// first decode row never materialized a slot; that is not an error. +pub(super) fn slot_release( + owners: &mut [Option], + request_id: RequestId, +) -> Option { + let slot = owners.iter().position(|owner| *owner == Some(request_id))?; + owners[slot] = None; + Some(slot) +} + +/// Apply the scheduler's slot compaction to the worker's slot map and report +/// whether a GPU state move is needed. Both requests may legitimately be +/// unmaterialized (retired/compacted before their first decode row), but a +/// materialized slot must hold exactly the request the scheduler claims. +pub(super) fn slot_compact( + owners: &mut [Option], + dropped: RequestId, + compaction: TpSlotCompaction, +) -> Result { + let TpSlotCompaction { + moved_request_id, + from, + to, + } = compaction; + anyhow::ensure!( + from < owners.len() && to < owners.len(), + "Qwen3.5 TP slot compaction {from} -> {to} exceeds worker slot map {}", + owners.len() + ); + let dropped_owner = owners[to]; + let moved_owner = owners[from]; + if let Some(owner) = dropped_owner { + anyhow::ensure!( + owner == dropped, + "Qwen3.5 TP slot {to} holds request {} where the scheduler dropped request {}", + owner.get(), + dropped.get() + ); + } + if let Some(owner) = moved_owner { + anyhow::ensure!( + owner == moved_request_id, + "Qwen3.5 TP slot {from} holds request {} where the scheduler moved request {}", + owner.get(), + moved_request_id.get() + ); + } + owners[to] = moved_owner; + owners[from] = None; + Ok(moved_owner.is_some()) +} + +fn states_in_row_order<'a>( + states: &'a mut [TpRequestState], + row_of_state: &[Option], +) -> Vec<&'a mut TpRequestState> { + let mut rows: Vec<(usize, &'a mut TpRequestState)> = states + .iter_mut() + .zip(row_of_state) + .filter_map(|(state, row)| row.map(|row| (row, state))) + .collect(); + rows.sort_unstable_by_key(|(row, _)| *row); + rows.into_iter().map(|(_, state)| state).collect() +} + +/// Rank-0 sampling pass over one decode batch: snapshot the requested logprob +/// rows, select one token per row, and pair each token with its logprob. +fn sample_decode_rows( + ctx: &pegainfer_core::tensor::DeviceContext, + logits: &pegainfer_core::tensor::HiddenStates, + requests: &[TpDecodeStepItem], + sample_seed: u64, + scratch: &mut pegainfer_sample::SampleScratch, +) -> Result> { + let bs = requests.len(); + let requested_logprobs: Vec> = + requests.iter().map(|request| request.logprobs).collect(); + let cpu_logits = snapshot_requested_logprobs(ctx, logits, &requested_logprobs)?; + let params_refs: Vec<&SamplingParams> = requests + .iter() + .map(|request| &request.sampling_params) + .collect(); + let steps = vec![0u64; bs]; + let tokens = + pegainfer_sample::select_batch(ctx, logits, ¶ms_refs, &steps, sample_seed, scratch)?; + anyhow::ensure!( + tokens.len() == bs, + "Qwen3.5 TP decode sampling returned {} tokens for {bs} rows", + tokens.len() + ); + Ok(requests + .iter() + .enumerate() + .map(|(row, request)| { + let logprob = cpu_logits[row].as_ref().and_then(|(logits_row, top_k)| { + pegainfer_sample::token_logprob_from_row(logits_row, tokens[row], *top_k) + }); + DecodeRequestResult { + request_id: request.request_id, + token: tokens[row], + logprob, + } + }) + .collect()) +} + +struct CublasThreadGuard; + +impl Drop for CublasThreadGuard { + fn drop(&mut self) { + unsafe { + crate::ffi::cublas_destroy(); + } + } +} + +fn bind_worker_thread(model: &Qwen35Model) -> Result { + let ctx = model.device_ctx(); + unsafe { + let err = crate::ffi::cuda_set_device(ctx.device_ordinal as i32); + if err != 0 { + return Err(anyhow::anyhow!( + "Failed to set CUDA device {} on Qwen3.5 TP worker thread: cudaError={}", + ctx.device_ordinal, + err + )); + } + } + ctx.ctx.bind_to_thread().map_err(|e| { + anyhow::anyhow!("Failed to bind CUDA context to Qwen3.5 TP worker thread: {e}") + })?; + unsafe { + crate::ffi::cublas_init(); + } + Ok(CublasThreadGuard) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn nccl_startup_watchdog_disarms_after_success() { + let (done_tx, watchdog) = spawn_nccl_startup_watchdog().unwrap(); + disarm_nccl_startup_watchdog(done_tx, watchdog).unwrap(); + } +} From 700933c764ec1f40678e0d2c7798cc563187c0a8 Mon Sep 17 00:00:00 2001 From: CAICAIIs <3360776475@qq.com> Date: Thu, 10 Sep 2026 12:00:35 +0800 Subject: [PATCH 2/5] refactor(qwen35): unify the TP one-shot gates into TpGate The tp_executor entry and the worker module carried two near-identical condvar gates: TpCommandStartGate (Pending/Execute/Cancel, first-write resolve returning whether it won) and TpStartupGate (Pending/Connect/ Cancel, silent set). Both broadcast one go/cancel decision that is resolved exactly once, so they collapse into TpGate + TpGateDecision with the first-write-wins resolve; startup treats Go as connect and asserts on the returned decision instead of a bool. Behavior unchanged. Signed-off-by: CAICAIIs <3360776475@qq.com> --- pegainfer-qwen35/src/tp_executor.rs | 86 +++++++++++----------- pegainfer-qwen35/src/tp_executor/worker.rs | 57 ++------------ 2 files changed, 49 insertions(+), 94 deletions(-) diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index dcc7a0227..e999d392f 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -95,18 +95,18 @@ enum TpWorkerCommand { RunPrefillChunks { chunks: Vec, sample_seed: u64, - start: Arc, + start: Arc, resp: mpsc::Sender, }, RunDecodeStep { requests: Vec, sample_seed: u64, - start: Arc, + start: Arc, resp: mpsc::Sender, }, RunUnifiedStep { plan: TpUnifiedPlan, - start: Arc, + start: Arc, resp: mpsc::Sender, }, DropRequest { @@ -115,14 +115,14 @@ enum TpWorkerCommand { /// `Some` only when the dropped request held a decode slot that a /// still-active request now takes over. Eager workers ignore it. compaction: Option, - start: Arc, + start: Arc, resp: mpsc::Sender, }, /// Startup-only (graph-enabled TP): one phase of the decode-graph /// pre-capture sweep, barriered across ranks by the controller. Precapture { phase: PrecapturePhase, - start: Arc, + start: Arc, resp: mpsc::Sender, }, #[cfg(test)] @@ -167,32 +167,39 @@ pub enum DropExpectation { MustExist, } +/// One-shot go/cancel decision broadcast to every rank's thread. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] -enum TpCommandDecision { +enum TpGateDecision { #[default] Pending, - Execute, + Go, Cancel, } +/// A gate the dispatcher resolves exactly once; waiters block until the +/// decision leaves `Pending`. Used for per-command starts and for startup. #[derive(Default)] -struct TpCommandStartGate { - decision: Mutex, +struct TpGate { + decision: Mutex, changed: Condvar, } -impl TpCommandStartGate { - fn execute(&self) -> bool { - self.resolve(TpCommandDecision::Execute) - } - - fn cancel(&self) -> bool { - self.resolve(TpCommandDecision::Cancel) +impl TpGate { + /// Resolve the decision; returns false if someone resolved it first. + fn resolve(&self, next: TpGateDecision) -> bool { + let mut decision = self.decision.lock().unwrap_or_else(PoisonError::into_inner); + if *decision != TpGateDecision::Pending { + return false; + } + *decision = next; + self.changed.notify_all(); + true } - fn wait(&self) -> TpCommandDecision { + /// Block until the gate is resolved, then return the decision. + fn wait(&self) -> TpGateDecision { let mut decision = self.decision.lock().unwrap_or_else(PoisonError::into_inner); - while *decision == TpCommandDecision::Pending { + while *decision == TpGateDecision::Pending { decision = self .changed .wait(decision) @@ -200,16 +207,6 @@ impl TpCommandStartGate { } *decision } - - fn resolve(&self, next: TpCommandDecision) -> bool { - let mut decision = self.decision.lock().unwrap_or_else(PoisonError::into_inner); - if *decision != TpCommandDecision::Pending { - return false; - } - *decision = next; - self.changed.notify_all(); - true - } } #[derive(Default)] @@ -442,7 +439,7 @@ impl Qwen35TpExecutor { let nccl_id = cudarc::nccl::safe::Id::new() .map_err(|e| anyhow::anyhow!("failed to create Qwen3.5 TP NCCL id: {e:?}"))?; - let startup_gate = Arc::new(TpStartupGate::default()); + let startup_gate = Arc::new(TpGate::default()); let effective_max_batch = Arc::new(AtomicUsize::new(0)); let poison = Arc::new(TpRuntimePoison::default()); let mut workers = Vec::with_capacity(world_size); @@ -467,7 +464,7 @@ impl Qwen35TpExecutor { startups.push(startup); } Err(err) => { - startup_gate.cancel(); + startup_gate.resolve(TpGateDecision::Cancel); return Err(err); } } @@ -479,11 +476,11 @@ impl Qwen35TpExecutor { min_rank_max_batch = min_rank_max_batch.min(rank_max_batch); } Ok(Err(err)) => { - startup_gate.cancel(); + startup_gate.resolve(TpGateDecision::Cancel); return Err(err); } Err(_) => { - startup_gate.cancel(); + startup_gate.resolve(TpGateDecision::Cancel); return Err(anyhow::anyhow!( "Qwen3.5 TP worker {rank} exited during pre-NCCL startup" )); @@ -503,11 +500,11 @@ impl Qwen35TpExecutor { let (watchdog_done, watchdog) = match spawn_nccl_startup_watchdog() { Ok(watchdog) => watchdog, Err(err) => { - startup_gate.cancel(); + startup_gate.resolve(TpGateDecision::Cancel); return Err(err); } }; - startup_gate.connect(); + startup_gate.resolve(TpGateDecision::Go); let startup_result = startups .into_iter() .enumerate() @@ -1013,7 +1010,7 @@ impl Qwen35TpExecutor { fn dispatch_mutating( &self, operation: &'static str, - build: impl Fn(Arc, mpsc::Sender) -> TpWorkerCommand, + build: impl Fn(Arc, mpsc::Sender) -> TpWorkerCommand, ) -> Result> { dispatch_mutating_commands( self.world_size, @@ -1039,15 +1036,15 @@ fn dispatch_mutating_commands( world_size: usize, operation: &'static str, poison: &TpRuntimePoison, - build: impl Fn(Arc, mpsc::Sender) -> TpWorkerCommand, + build: impl Fn(Arc, mpsc::Sender) -> TpWorkerCommand, mut send: impl FnMut(usize, TpWorkerCommand) -> Result<()>, ) -> Result> { - let start = Arc::new(TpCommandStartGate::default()); + let start = Arc::new(TpGate::default()); let (resp_tx, resp_rx) = mpsc::channel(); for rank in 0..world_size { let command = build(Arc::clone(&start), resp_tx.clone()); if let Err(err) = send(rank, command) { - start.cancel(); + start.resolve(TpGateDecision::Cancel); let reason = poison.poison(format!( "failed to dispatch {operation} to TP worker rank {rank}: {err:#}" )); @@ -1055,7 +1052,7 @@ fn dispatch_mutating_commands( } } drop(resp_tx); - let resolved = start.execute(); + let resolved = start.resolve(TpGateDecision::Go); debug_assert!(resolved, "fresh TP command gate resolved more than once"); Ok(resp_rx) } @@ -1496,7 +1493,7 @@ mod tests { #[test] fn startup_gate_cancel_releases_waiting_workers() { - let gate = Arc::new(TpStartupGate::default()); + let gate = Arc::new(TpGate::default()); let worker_gate = Arc::clone(&gate); let (done_tx, done_rx) = mpsc::channel(); let waiter = thread::spawn(move || { @@ -1505,10 +1502,11 @@ mod tests { gate.cancel(); - assert!( - !done_rx + assert_eq!( + done_rx .recv_timeout(std::time::Duration::from_secs(1)) - .expect("cancelled startup gate should release workers within one second") + .expect("cancelled startup gate should release workers within one second"), + TpGateDecision::Cancel, ); waiter.join().unwrap(); } @@ -1634,7 +1632,7 @@ mod tests { let TpWorkerCommand::RunPrefillChunks { start, .. } = rank0_rx.recv().unwrap() else { panic!("expected prefill command") }; - assert_eq!(start.wait(), TpCommandDecision::Cancel); + assert_eq!(start.wait(), TpGateDecision::Cancel); assert!(matches!( rank1_rx.try_recv(), Err(mpsc::TryRecvError::Empty) diff --git a/pegainfer-qwen35/src/tp_executor/worker.rs b/pegainfer-qwen35/src/tp_executor/worker.rs index e2b90f86e..66a3fcbdf 100644 --- a/pegainfer-qwen35/src/tp_executor/worker.rs +++ b/pegainfer-qwen35/src/tp_executor/worker.rs @@ -48,49 +48,6 @@ pub(super) struct TpWorker { done: mpsc::Receiver<()>, } -#[derive(Clone, Copy, Default, PartialEq, Eq)] -enum TpStartupDecision { - #[default] - Pending, - Connect, - Cancel, -} - -#[derive(Default)] -pub(super) struct TpStartupGate { - decision: Mutex, - changed: Condvar, -} - -impl TpStartupGate { - pub(super) fn connect(&self) { - self.set(TpStartupDecision::Connect); - } - - pub(super) fn cancel(&self) { - self.set(TpStartupDecision::Cancel); - } - - pub(super) fn wait(&self) -> bool { - let mut decision = self.decision.lock().unwrap_or_else(PoisonError::into_inner); - while *decision == TpStartupDecision::Pending { - decision = self - .changed - .wait(decision) - .unwrap_or_else(PoisonError::into_inner); - } - *decision == TpStartupDecision::Connect - } - - fn set(&self, next: TpStartupDecision) { - let mut decision = self.decision.lock().unwrap_or_else(PoisonError::into_inner); - if *decision == TpStartupDecision::Pending { - *decision = next; - self.changed.notify_all(); - } - } -} - impl TpWorker { #[allow(clippy::too_many_arguments)] #[allow(clippy::type_complexity)] @@ -102,7 +59,7 @@ impl TpWorker { max_prefill_tokens: usize, graph_enabled: bool, nccl_id: cudarc::nccl::safe::Id, - startup_gate: Arc, + startup_gate: Arc, effective_max_batch: Arc, poison: Arc, ) -> Result<( @@ -137,7 +94,7 @@ impl TpWorker { return; } }; - if !startup_gate.wait() { + if startup_gate.wait() != TpGateDecision::Go { return; } let max_batch = effective_max_batch.load(Ordering::Acquire); @@ -449,7 +406,7 @@ impl TpWorkerState { start, resp, } => { - if start.wait() == TpCommandDecision::Cancel { + if start.wait() == TpGateDecision::Cancel { false } else { let result = self.execute_prefill_chunks(&chunks, sample_seed); @@ -462,7 +419,7 @@ impl TpWorkerState { start, resp, } => { - if start.wait() == TpCommandDecision::Cancel { + if start.wait() == TpGateDecision::Cancel { false } else { let result = self.execute_decode(&requests, sample_seed); @@ -470,7 +427,7 @@ impl TpWorkerState { } } TpWorkerCommand::RunUnifiedStep { plan, start, resp } => { - if start.wait() == TpCommandDecision::Cancel { + if start.wait() == TpGateDecision::Cancel { false } else { let result = self.execute_unified(&plan); @@ -483,7 +440,7 @@ impl TpWorkerState { start, resp, } => { - if start.wait() == TpCommandDecision::Cancel { + if start.wait() == TpGateDecision::Cancel { false } else { let result = self @@ -493,7 +450,7 @@ impl TpWorkerState { } } TpWorkerCommand::Precapture { phase, start, resp } => { - if start.wait() == TpCommandDecision::Cancel { + if start.wait() == TpGateDecision::Cancel { false } else { let result = self.precapture_phase(phase).map(|()| TpWorkerReply::Ack); From 4c087a4998706fa4b4a5ce4ae783b93817e285b7 Mon Sep 17 00:00:00 2001 From: CAICAIIs <3360776475@qq.com> Date: Thu, 10 Sep 2026 12:01:58 +0800 Subject: [PATCH 3/5] refactor(qwen35): dedupe the CUDA thread binding into cublas_thread.rs The scheduler and the TP worker carried byte-identical CUDA/cuBLAS thread bindings (guard + set_device + bind_to_thread + cublas_init); the worker split moved one of the two copies, making the duplication visible. Both now call crate::cublas_thread::bind_model_thread with a role tag (the failure messages are unchanged); the scheduler keeps its tune_decode_gemm_algos step at its own call site. Behavior unchanged. Signed-off-by: CAICAIIs <3360776475@qq.com> --- pegainfer-qwen35/src/cublas_thread.rs | 40 +++++++++++++ pegainfer-qwen35/src/lib.rs | 1 + pegainfer-qwen35/src/scheduler/mod.rs | 66 +++++++--------------- pegainfer-qwen35/src/tp_executor.rs | 2 +- pegainfer-qwen35/src/tp_executor/worker.rs | 34 +---------- 5 files changed, 64 insertions(+), 79 deletions(-) create mode 100644 pegainfer-qwen35/src/cublas_thread.rs diff --git a/pegainfer-qwen35/src/cublas_thread.rs b/pegainfer-qwen35/src/cublas_thread.rs new file mode 100644 index 000000000..4c5cc9b0d --- /dev/null +++ b/pegainfer-qwen35/src/cublas_thread.rs @@ -0,0 +1,40 @@ +//! Bind the calling thread to a model's CUDA device/context and cuBLAS +//! handle; the guard destroys the handle when the binding ends. + +use anyhow::Result; + +use crate::weights::Qwen35Model; + +pub(crate) struct CublasThreadGuard; + +impl Drop for CublasThreadGuard { + fn drop(&mut self) { + unsafe { + crate::ffi::cublas_destroy(); + } + } +} + +/// Bind this thread to `model`'s device and context, then initialize cuBLAS. +/// `role` names the thread kind in failure messages ("scheduler", "TP worker"); +/// every thread that runs model work must hold the returned guard. +pub(crate) fn bind_model_thread(model: &Qwen35Model, role: &str) -> Result { + let ctx = model.device_ctx(); + unsafe { + let err = crate::ffi::cuda_set_device(ctx.device_ordinal as i32); + if err != 0 { + return Err(anyhow::anyhow!( + "Failed to set CUDA device {} on Qwen3.5 {role} thread: cudaError={}", + ctx.device_ordinal, + err + )); + } + } + ctx.ctx.bind_to_thread().map_err(|e| { + anyhow::anyhow!("Failed to bind CUDA context to Qwen3.5 {role} thread: {e}") + })?; + unsafe { + crate::ffi::cublas_init(); + } + Ok(CublasThreadGuard) +} diff --git a/pegainfer-qwen35/src/lib.rs b/pegainfer-qwen35/src/lib.rs index 3da029e66..9aea7c4eb 100644 --- a/pegainfer-qwen35/src/lib.rs +++ b/pegainfer-qwen35/src/lib.rs @@ -9,6 +9,7 @@ mod batch_decode; pub(crate) mod batch_decode_graph; pub(crate) mod config; +mod cublas_thread; mod decode_buffers; mod executor; mod ffi; diff --git a/pegainfer-qwen35/src/scheduler/mod.rs b/pegainfer-qwen35/src/scheduler/mod.rs index 968aacb58..379470a8f 100644 --- a/pegainfer-qwen35/src/scheduler/mod.rs +++ b/pegainfer-qwen35/src/scheduler/mod.rs @@ -379,20 +379,26 @@ pub(crate) fn start_with_capacity_and_policy( let join_handle = thread::Builder::new() .name("scheduler-qwen35".into()) - .spawn(move || match bind_model_thread(backend.model()) { - Ok(_guard) => { - let _ = startup_tx.send(Ok(())); - scheduler_loop( - SchedulerBackend::Single(backend), - submit_rx, - seed, - max_prefill_tokens, - scheduler_policy, - load_tx, - ); - } - Err(err) => { - let _ = startup_tx.send(Err(err)); + .spawn(move || { + match crate::cublas_thread::bind_model_thread(backend.model(), "scheduler") { + Ok(_guard) => { + if let Err(err) = backend.model().tune_decode_gemm_algos() { + let _ = startup_tx.send(Err(err)); + return; + } + let _ = startup_tx.send(Ok(())); + scheduler_loop( + SchedulerBackend::Single(backend), + submit_rx, + seed, + max_prefill_tokens, + scheduler_policy, + load_tx, + ); + } + Err(err) => { + let _ = startup_tx.send(Err(err)); + } } }) .expect("failed to spawn Qwen3.5 scheduler thread"); @@ -491,38 +497,6 @@ fn servable_len(max_context: usize, max_pages: usize, page_size: usize) -> u32 { .unwrap_or(u32::MAX) } -struct CublasThreadGuard; - -impl Drop for CublasThreadGuard { - fn drop(&mut self) { - unsafe { - crate::ffi::cublas_destroy(); - } - } -} - -fn bind_model_thread(model: &Qwen35Model) -> Result { - let ctx = model.device_ctx(); - unsafe { - let err = crate::ffi::cuda_set_device(ctx.device_ordinal as i32); - if err != 0 { - return Err(anyhow::anyhow!( - "Failed to set CUDA device {} on Qwen3.5 scheduler thread: cudaError={}", - ctx.device_ordinal, - err - )); - } - } - ctx.ctx.bind_to_thread().map_err(|e| { - anyhow::anyhow!("Failed to bind CUDA context to Qwen3.5 scheduler thread: {e}") - })?; - unsafe { - crate::ffi::cublas_init(); - } - model.tune_decode_gemm_algos()?; - Ok(CublasThreadGuard) -} - // ── Main loop ─────────────────────────────────────────────────────────── fn publish_load( diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index e999d392f..3e721e50a 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -1500,7 +1500,7 @@ mod tests { let _ = done_tx.send(worker_gate.wait()); }); - gate.cancel(); + gate.resolve(TpGateDecision::Cancel); assert_eq!( done_rx diff --git a/pegainfer-qwen35/src/tp_executor/worker.rs b/pegainfer-qwen35/src/tp_executor/worker.rs index 66a3fcbdf..0058502cc 100644 --- a/pegainfer-qwen35/src/tp_executor/worker.rs +++ b/pegainfer-qwen35/src/tp_executor/worker.rs @@ -2,6 +2,7 @@ //! gating, the NCCL startup watchdog, and the per-rank command loop. use super::*; +use crate::cublas_thread::CublasThreadGuard; const TP_NCCL_STARTUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); const TP_WORKER_SHUTDOWN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30); @@ -219,7 +220,7 @@ impl TpWorkerPrepared { max_prefill_tokens: usize, graph_enabled: bool, ) -> Result<(Self, usize)> { - let cublas_guard = bind_worker_thread(&model)?; + let cublas_guard = crate::cublas_thread::bind_model_thread(&model, "TP worker")?; let (free_bytes, total_bytes) = model .device_ctx() .ctx @@ -1111,37 +1112,6 @@ fn sample_decode_rows( .collect()) } -struct CublasThreadGuard; - -impl Drop for CublasThreadGuard { - fn drop(&mut self) { - unsafe { - crate::ffi::cublas_destroy(); - } - } -} - -fn bind_worker_thread(model: &Qwen35Model) -> Result { - let ctx = model.device_ctx(); - unsafe { - let err = crate::ffi::cuda_set_device(ctx.device_ordinal as i32); - if err != 0 { - return Err(anyhow::anyhow!( - "Failed to set CUDA device {} on Qwen3.5 TP worker thread: cudaError={}", - ctx.device_ordinal, - err - )); - } - } - ctx.ctx.bind_to_thread().map_err(|e| { - anyhow::anyhow!("Failed to bind CUDA context to Qwen3.5 TP worker thread: {e}") - })?; - unsafe { - crate::ffi::cublas_init(); - } - Ok(CublasThreadGuard) -} - #[cfg(test)] mod tests { use super::*; From 676daedf3dd5c5604cced4cf0b4773c256d56f29 Mon Sep 17 00:00:00 2001 From: CAICAIIs <3360776475@qq.com> Date: Tue, 8 Sep 2026 14:30:37 +0800 Subject: [PATCH 4/5] refactor(qwen35): split TP response contract into tp_executor/responses.rs Move the reply envelope (TpWorkerReply/TpWorkerResponse), bounded response collection, per-command response validation, and the step-timeout/abort helpers out of tp_executor.rs; the three response-contract tests move with them. Behavior unchanged, with two behavior-preserving consolidations in the moved code: validate_exact_rank_responses folds its push loop into an iterator pipeline, and the three rank-0 payload validators (prefill / decode / unified) collapse into one validate_primary_responses helper over payload extractors, keeping the exact prior error messages. Rebased over #1004/#1005/#1033/#1038 (the precapture sweep keeps using validate_ack_responses, which stays production-visible). Signed-off-by: CAICAIIs <3360776475@qq.com> --- pegainfer-qwen35/src/tp_executor.rs | 372 +--------------- pegainfer-qwen35/src/tp_executor/responses.rs | 396 ++++++++++++++++++ 2 files changed, 398 insertions(+), 370 deletions(-) create mode 100644 pegainfer-qwen35/src/tp_executor/responses.rs diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index 3e721e50a..443cb6171 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -46,11 +46,12 @@ use crate::recurrent_state::RecurrentState; use crate::weights::ModelRuntimeConfig; use crate::weights::Qwen35Model; +mod responses; mod worker; +use responses::*; use worker::*; -const TP_RUNTIME_STEP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300); /// The pre-capture sweep records every decode bucket per rank; the 60 s NCCL /// startup budget is far too small for that. const TP_PRECAPTURE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(600); @@ -141,25 +142,6 @@ enum TpWorkerCommand { Shutdown, } -#[derive(Debug)] -enum TpWorkerReply { - Ack, - DropAck { - existed: bool, - }, - Prefill(PrefillResult), - Decode(DecodeResult), - Unified(TpUnifiedResult), - #[cfg(test)] - Snapshot(WorkerStateSnapshot), -} - -#[derive(Debug)] -struct TpWorkerResponse { - rank: usize, - result: Result, -} - /// Scheduler-owned lifecycle proof required from every TP rank during cleanup. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum DropExpectation { @@ -1205,288 +1187,6 @@ impl From for TpPrefillChunkItem { } } -fn recv_runtime_responses( - responses: &mpsc::Receiver, - expected: usize, - operation: &'static str, - poison: &TpRuntimePoison, -) -> Result> { - collect_runtime_responses(expected, operation, poison, || { - recv_runtime_response(responses, operation, poison) - }) -} - -fn collect_runtime_responses( - expected: usize, - operation: &'static str, - poison: &TpRuntimePoison, - mut recv_next: impl FnMut() -> Result, -) -> Result> { - let mut collected = Vec::with_capacity(expected); - for _ in 0..expected { - let response = recv_next()?; - if let Err(err) = &response.result { - // A failed rank may leave peers blocked in a collective, so response-set - // completeness is no longer recoverable or useful. - let reason = poison.poison(format!( - "rank {} failed during {operation}: {err:#}", - response.rank - )); - return Err(anyhow::anyhow!(reason)); - } - collected.push(response); - } - Ok(collected) -} - -fn validate_dispatched_responses( - result: Result, - operation: &'static str, - poison: &TpRuntimePoison, -) -> Result { - result.map_err(|err| { - let reason = poison.poison(format!( - "invalid Qwen3.5 TP {operation} response set: {err:#}" - )); - anyhow::anyhow!(reason) - }) -} - -fn validate_exact_rank_responses( - responses: Vec, - world_size: usize, - operation: &'static str, -) -> Result> { - anyhow::ensure!( - responses.len() == world_size, - "{operation} expected {world_size} responses, got {}", - responses.len() - ); - let mut seen_ranks = HashSet::with_capacity(world_size); - let mut replies = Vec::with_capacity(world_size); - for response in responses { - anyhow::ensure!( - response.rank < world_size, - "{operation} returned out-of-range rank {} for world size {world_size}", - response.rank - ); - anyhow::ensure!( - seen_ranks.insert(response.rank), - "{operation} returned duplicate rank {}", - response.rank - ); - replies.push((response.rank, response.result?)); - } - anyhow::ensure!( - (0..world_size).all(|rank| seen_ranks.contains(&rank)), - "{operation} response set did not contain every rank" - ); - replies.sort_unstable_by_key(|(rank, _)| *rank); - Ok(replies) -} - -fn validate_ack_responses( - responses: Vec, - world_size: usize, - operation: &'static str, -) -> Result<()> { - for (rank, reply) in validate_exact_rank_responses(responses, world_size, operation)? { - anyhow::ensure!( - matches!(reply, TpWorkerReply::Ack), - "{operation} rank {rank} returned {} instead of acknowledgement", - reply_name(&reply) - ); - } - Ok(()) -} - -fn validate_drop_responses( - responses: Vec, - world_size: usize, - expectation: DropExpectation, -) -> Result<()> { - let mut existence = Vec::with_capacity(world_size); - for (rank, reply) in validate_exact_rank_responses(responses, world_size, "drop request")? { - let TpWorkerReply::DropAck { existed } = reply else { - anyhow::bail!( - "drop request rank {rank} returned {} instead of drop acknowledgement", - reply_name(&reply) - ); - }; - existence.push((rank, existed)); - } - let expected = expectation == DropExpectation::MustExist; - anyhow::ensure!( - existence.iter().all(|(_, existed)| *existed == expected), - "drop request expected {expectation:?}, got rank existence {existence:?}" - ); - Ok(()) -} - -fn validate_prefill_responses( - responses: Vec, - world_size: usize, -) -> Result { - let mut primary = None; - for (rank, reply) in validate_exact_rank_responses(responses, world_size, "prefill")? { - match (rank, reply) { - (0, TpWorkerReply::Prefill(result)) => primary = Some(result), - (0, reply) => anyhow::bail!( - "prefill rank 0 returned {} instead of primary prefill result", - reply_name(&reply) - ), - (_, TpWorkerReply::Ack) => {} - (rank, reply) => anyhow::bail!( - "prefill non-primary rank {rank} returned {} instead of acknowledgement", - reply_name(&reply) - ), - } - } - primary.ok_or_else(|| anyhow::anyhow!("prefill returned no primary result")) -} - -fn validate_decode_responses( - responses: Vec, - world_size: usize, -) -> Result { - let mut primary = None; - for (rank, reply) in validate_exact_rank_responses(responses, world_size, "decode")? { - match (rank, reply) { - (0, TpWorkerReply::Decode(result)) => primary = Some(result), - (0, reply) => anyhow::bail!( - "decode rank 0 returned {} instead of primary decode result", - reply_name(&reply) - ), - (_, TpWorkerReply::Ack) => {} - (rank, reply) => anyhow::bail!( - "decode non-primary rank {rank} returned {} instead of acknowledgement", - reply_name(&reply) - ), - } - } - primary.ok_or_else(|| anyhow::anyhow!("decode returned no primary result")) -} - -fn validate_unified_responses( - responses: Vec, - world_size: usize, -) -> Result { - let mut primary = None; - for (rank, reply) in validate_exact_rank_responses(responses, world_size, "unified step")? { - match (rank, reply) { - (0, TpWorkerReply::Unified(result)) => primary = Some(result), - (0, reply) => anyhow::bail!( - "unified step rank 0 returned {} instead of primary unified result", - reply_name(&reply) - ), - (_, TpWorkerReply::Ack) => {} - (rank, reply) => anyhow::bail!( - "unified step non-primary rank {rank} returned {} instead of acknowledgement", - reply_name(&reply) - ), - } - } - primary.ok_or_else(|| anyhow::anyhow!("unified step returned no primary result")) -} - -fn reply_name(reply: &TpWorkerReply) -> &'static str { - match reply { - TpWorkerReply::Ack => "acknowledgement", - TpWorkerReply::DropAck { .. } => "drop acknowledgement", - TpWorkerReply::Prefill(_) => "prefill result", - TpWorkerReply::Decode(_) => "decode result", - TpWorkerReply::Unified(_) => "unified result", - #[cfg(test)] - TpWorkerReply::Snapshot(_) => "worker snapshot", - } -} - -#[cfg(test)] -fn wait_for_worker_snapshots( - responses: &mpsc::Receiver, - world_size: usize, - poison: &TpRuntimePoison, -) -> Result> { - let mut seen_ranks = HashSet::with_capacity(world_size); - let mut snapshots = Vec::with_capacity(world_size); - for _ in 0..world_size { - let response = recv_runtime_response(responses, "snapshot state", poison)?; - anyhow::ensure!( - response.rank < world_size, - "Qwen3.5 TP snapshot returned out-of-range rank {} for world size {world_size}", - response.rank - ); - anyhow::ensure!( - seen_ranks.insert(response.rank), - "Qwen3.5 TP snapshot returned duplicate rank {}", - response.rank - ); - match response.result? { - TpWorkerReply::Snapshot(snapshot) => { - anyhow::ensure!( - snapshot.rank == response.rank, - "Qwen3.5 TP snapshot payload rank {} does not match response rank {}", - snapshot.rank, - response.rank - ); - anyhow::ensure!( - snapshot.request_count == snapshot.requests.len(), - "Qwen3.5 TP rank {} snapshot count {} does not match {} request entries", - snapshot.rank, - snapshot.request_count, - snapshot.requests.len() - ); - snapshots.push(snapshot); - } - TpWorkerReply::Ack => { - anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned acknowledgement") - } - TpWorkerReply::DropAck { .. } => { - anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned drop acknowledgement") - } - TpWorkerReply::Prefill(_) => { - anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned prefill result") - } - TpWorkerReply::Decode(_) => { - anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned decode result") - } - TpWorkerReply::Unified(_) => { - anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned unified result") - } - } - } - anyhow::ensure!( - (0..world_size).all(|rank| seen_ranks.contains(&rank)), - "Qwen3.5 TP snapshot response set did not contain every rank" - ); - snapshots.sort_unstable_by_key(|snapshot| snapshot.rank); - Ok(snapshots) -} - -fn recv_runtime_response( - responses: &mpsc::Receiver, - operation: &'static str, - poison: &TpRuntimePoison, -) -> Result { - match responses.recv_timeout(TP_RUNTIME_STEP_TIMEOUT) { - Ok(response) => Ok(response), - Err(mpsc::RecvTimeoutError::Disconnected) => { - let reason = poison.poison(format!("response channel disconnected during {operation}")); - Err(anyhow::anyhow!(reason)) - } - Err(mpsc::RecvTimeoutError::Timeout) => fatal_tp_abort(&format!( - "Qwen3.5 TP {operation} did not complete within {}s", - TP_RUNTIME_STEP_TIMEOUT.as_secs() - )), - } -} - -fn fatal_tp_abort(message: &str) -> ! { - eprintln!("{message}; aborting"); - log::error!("{message}; aborting"); - std::process::abort(); -} - #[cfg(test)] mod tests { use super::*; @@ -1527,74 +1227,6 @@ mod tests { assert!(!err.contains("rank 0 NCCL error")); } - #[test] - fn runtime_response_failure_poisons_executor() { - let poison = TpRuntimePoison::default(); - let responses = vec![ - reply(0, TpWorkerReply::Ack), - TpWorkerResponse { - rank: 1, - result: Err(anyhow::anyhow!("rank 1 failed")), - }, - ]; - - let err = validate_dispatched_responses( - validate_ack_responses(responses, 2, "test"), - "test", - &poison, - ) - .unwrap_err() - .to_string(); - assert!(err.contains("rank 1 failed")); - assert!(poison.ensure_healthy().is_err()); - } - - #[test] - fn runtime_response_collection_fails_fast_when_peer_never_responds() { - let poison = TpRuntimePoison::default(); - let (tx, rx) = mpsc::channel(); - tx.send(TpWorkerResponse { - rank: 0, - result: Err(anyhow::anyhow!("rank 0 failed")), - }) - .unwrap(); - let _keep_peer_channel_connected = tx; - let mut receive_attempts = 0; - - let err = collect_runtime_responses(2, "test", &poison, || { - receive_attempts += 1; - rx.recv_timeout(std::time::Duration::from_millis(50)) - .map_err(|err| anyhow::anyhow!("waited for nonresponding rank: {err}")) - }) - .unwrap_err() - .to_string(); - - assert_eq!(receive_attempts, 1, "collector waited for the missing rank"); - assert!(err.contains("rank 0 failed")); - assert!(!err.contains("waited for nonresponding rank")); - assert!(poison.ensure_healthy().is_err()); - } - - #[test] - fn disconnected_runtime_response_poisons_executor() { - let poison = TpRuntimePoison::default(); - let (tx, rx) = mpsc::channel(); - drop(tx); - - let err = recv_runtime_response(&rx, "test", &poison) - .unwrap_err() - .to_string(); - assert!(err.contains("response channel disconnected during test")); - assert!(poison.ensure_healthy().is_err()); - } - - fn reply(rank: usize, reply: TpWorkerReply) -> TpWorkerResponse { - TpWorkerResponse { - rank, - result: Ok(reply), - } - } - #[test] fn mutating_partial_dispatch_cancels_delivered_prefix_and_poisons() { let poison = TpRuntimePoison::default(); diff --git a/pegainfer-qwen35/src/tp_executor/responses.rs b/pegainfer-qwen35/src/tp_executor/responses.rs new file mode 100644 index 000000000..92aa4b03d --- /dev/null +++ b/pegainfer-qwen35/src/tp_executor/responses.rs @@ -0,0 +1,396 @@ +//! Tensor-parallel response contract: the worker reply envelope, bounded +//! response collection, and per-command response validation. + +use super::*; + +#[derive(Debug)] +pub(super) enum TpWorkerReply { + Ack, + DropAck { + existed: bool, + }, + Prefill(PrefillResult), + Decode(DecodeResult), + Unified(TpUnifiedResult), + #[cfg(test)] + Snapshot(WorkerStateSnapshot), +} + +#[derive(Debug)] +pub(super) struct TpWorkerResponse { + pub(super) rank: usize, + pub(super) result: Result, +} + +const TP_RUNTIME_STEP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300); + +pub(super) fn recv_runtime_responses( + responses: &mpsc::Receiver, + expected: usize, + operation: &'static str, + poison: &TpRuntimePoison, +) -> Result> { + collect_runtime_responses(expected, operation, poison, || { + recv_runtime_response(responses, operation, poison) + }) +} + +fn collect_runtime_responses( + expected: usize, + operation: &'static str, + poison: &TpRuntimePoison, + mut recv_next: impl FnMut() -> Result, +) -> Result> { + let mut collected = Vec::with_capacity(expected); + for _ in 0..expected { + let response = recv_next()?; + if let Err(err) = &response.result { + // A failed rank may leave peers blocked in a collective, so response-set + // completeness is no longer recoverable or useful. + let reason = poison.poison(format!( + "rank {} failed during {operation}: {err:#}", + response.rank + )); + return Err(anyhow::anyhow!(reason)); + } + collected.push(response); + } + Ok(collected) +} + +pub(super) fn validate_dispatched_responses( + result: Result, + operation: &'static str, + poison: &TpRuntimePoison, +) -> Result { + result.map_err(|err| { + let reason = poison.poison(format!( + "invalid Qwen3.5 TP {operation} response set: {err:#}" + )); + anyhow::anyhow!(reason) + }) +} + +fn validate_exact_rank_responses( + responses: Vec, + world_size: usize, + operation: &'static str, +) -> Result> { + anyhow::ensure!( + responses.len() == world_size, + "{operation} expected {world_size} responses, got {}", + responses.len() + ); + let mut seen_ranks = HashSet::with_capacity(world_size); + let mut replies = responses + .into_iter() + .map(|response| { + anyhow::ensure!( + response.rank < world_size, + "{operation} returned out-of-range rank {} for world size {world_size}", + response.rank + ); + anyhow::ensure!( + seen_ranks.insert(response.rank), + "{operation} returned duplicate rank {}", + response.rank + ); + Ok((response.rank, response.result?)) + }) + .collect::>>()?; + anyhow::ensure!( + (0..world_size).all(|rank| seen_ranks.contains(&rank)), + "{operation} response set did not contain every rank" + ); + replies.sort_unstable_by_key(|(rank, _)| *rank); + Ok(replies) +} + +pub(super) fn validate_ack_responses( + responses: Vec, + world_size: usize, + operation: &'static str, +) -> Result<()> { + for (rank, reply) in validate_exact_rank_responses(responses, world_size, operation)? { + anyhow::ensure!( + matches!(reply, TpWorkerReply::Ack), + "{operation} rank {rank} returned {} instead of acknowledgement", + reply_name(&reply) + ); + } + Ok(()) +} + +pub(super) fn validate_drop_responses( + responses: Vec, + world_size: usize, + expectation: DropExpectation, +) -> Result<()> { + let expected = expectation == DropExpectation::MustExist; + let existence = validate_exact_rank_responses(responses, world_size, "drop request")? + .into_iter() + .map(|(rank, reply)| match reply { + TpWorkerReply::DropAck { existed } => Ok((rank, existed)), + reply => anyhow::bail!( + "drop request rank {rank} returned {} instead of drop acknowledgement", + reply_name(&reply) + ), + }) + .collect::>>()?; + anyhow::ensure!( + existence.iter().all(|(_, existed)| *existed == expected), + "drop request expected {expectation:?}, got rank existence {existence:?}" + ); + Ok(()) +} + +/// Reduce one exact-rank reply set to the rank-0 primary payload; every +/// non-primary rank must acknowledge. `operation` names the command in error +/// messages and `noun` names the expected payload. +fn validate_primary_responses( + responses: Vec, + world_size: usize, + operation: &'static str, + noun: &'static str, + unwrap_payload: fn(TpWorkerReply) -> Option, +) -> Result { + let mut primary = None; + for (rank, reply) in validate_exact_rank_responses(responses, world_size, operation)? { + match (rank, reply) { + (0, reply) => { + let wrong_payload = reply_name(&reply); + primary = Some(unwrap_payload(reply).ok_or_else(|| { + anyhow::anyhow!( + "{operation} rank 0 returned {wrong_payload} instead of primary {noun} result", + ) + })?); + } + (_, reply) => anyhow::ensure!( + matches!(reply, TpWorkerReply::Ack), + "{operation} non-primary rank {rank} returned {} instead of acknowledgement", + reply_name(&reply) + ), + } + } + primary.ok_or_else(|| anyhow::anyhow!("{operation} returned no primary result")) +} + +fn prefill_payload(reply: TpWorkerReply) -> Option { + match reply { + TpWorkerReply::Prefill(result) => Some(result), + _ => None, + } +} + +fn decode_payload(reply: TpWorkerReply) -> Option { + match reply { + TpWorkerReply::Decode(result) => Some(result), + _ => None, + } +} + +fn unified_payload(reply: TpWorkerReply) -> Option { + match reply { + TpWorkerReply::Unified(result) => Some(result), + _ => None, + } +} + +pub(super) fn validate_prefill_responses( + responses: Vec, + world_size: usize, +) -> Result { + validate_primary_responses(responses, world_size, "prefill", "prefill", prefill_payload) +} + +pub(super) fn validate_decode_responses( + responses: Vec, + world_size: usize, +) -> Result { + validate_primary_responses(responses, world_size, "decode", "decode", decode_payload) +} + +pub(super) fn validate_unified_responses( + responses: Vec, + world_size: usize, +) -> Result { + validate_primary_responses( + responses, + world_size, + "unified step", + "unified", + unified_payload, + ) +} + +fn reply_name(reply: &TpWorkerReply) -> &'static str { + match reply { + TpWorkerReply::Ack => "acknowledgement", + TpWorkerReply::DropAck { .. } => "drop acknowledgement", + TpWorkerReply::Prefill(_) => "prefill result", + TpWorkerReply::Decode(_) => "decode result", + TpWorkerReply::Unified(_) => "unified result", + #[cfg(test)] + TpWorkerReply::Snapshot(_) => "worker snapshot", + } +} + +#[cfg(test)] +pub(super) fn wait_for_worker_snapshots( + responses: &mpsc::Receiver, + world_size: usize, + poison: &TpRuntimePoison, +) -> Result> { + let mut seen_ranks = HashSet::with_capacity(world_size); + let mut snapshots = Vec::with_capacity(world_size); + for _ in 0..world_size { + let response = recv_runtime_response(responses, "snapshot state", poison)?; + anyhow::ensure!( + response.rank < world_size, + "Qwen3.5 TP snapshot returned out-of-range rank {} for world size {world_size}", + response.rank + ); + anyhow::ensure!( + seen_ranks.insert(response.rank), + "Qwen3.5 TP snapshot returned duplicate rank {}", + response.rank + ); + match response.result? { + TpWorkerReply::Snapshot(snapshot) => { + anyhow::ensure!( + snapshot.rank == response.rank, + "Qwen3.5 TP snapshot payload rank {} does not match response rank {}", + snapshot.rank, + response.rank + ); + anyhow::ensure!( + snapshot.request_count == snapshot.requests.len(), + "Qwen3.5 TP rank {} snapshot count {} does not match {} request entries", + snapshot.rank, + snapshot.request_count, + snapshot.requests.len() + ); + snapshots.push(snapshot); + } + TpWorkerReply::Ack => { + anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned acknowledgement") + } + TpWorkerReply::DropAck { .. } => { + anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned drop acknowledgement") + } + TpWorkerReply::Prefill(_) => { + anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned prefill result") + } + TpWorkerReply::Decode(_) => { + anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned decode result") + } + TpWorkerReply::Unified(_) => { + anyhow::bail!("Qwen3.5 TP snapshot unexpectedly returned unified result") + } + } + } + anyhow::ensure!( + (0..world_size).all(|rank| seen_ranks.contains(&rank)), + "Qwen3.5 TP snapshot response set did not contain every rank" + ); + snapshots.sort_unstable_by_key(|snapshot| snapshot.rank); + Ok(snapshots) +} + +fn recv_runtime_response( + responses: &mpsc::Receiver, + operation: &'static str, + poison: &TpRuntimePoison, +) -> Result { + match responses.recv_timeout(TP_RUNTIME_STEP_TIMEOUT) { + Ok(response) => Ok(response), + Err(mpsc::RecvTimeoutError::Disconnected) => { + let reason = poison.poison(format!("response channel disconnected during {operation}")); + Err(anyhow::anyhow!(reason)) + } + Err(mpsc::RecvTimeoutError::Timeout) => fatal_tp_abort(&format!( + "Qwen3.5 TP {operation} did not complete within {}s", + TP_RUNTIME_STEP_TIMEOUT.as_secs() + )), + } +} + +pub(super) fn fatal_tp_abort(message: &str) -> ! { + eprintln!("{message}; aborting"); + log::error!("{message}; aborting"); + std::process::abort(); +} + +#[cfg(test)] +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn runtime_response_failure_poisons_executor() { + let poison = TpRuntimePoison::default(); + let responses = vec![ + reply(0, TpWorkerReply::Ack), + TpWorkerResponse { + rank: 1, + result: Err(anyhow::anyhow!("rank 1 failed")), + }, + ]; + + let err = validate_dispatched_responses( + validate_ack_responses(responses, 2, "test"), + "test", + &poison, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("rank 1 failed")); + assert!(poison.ensure_healthy().is_err()); + } + + #[test] + fn runtime_response_collection_fails_fast_when_peer_never_responds() { + let poison = TpRuntimePoison::default(); + let (tx, rx) = mpsc::channel(); + tx.send(TpWorkerResponse { + rank: 0, + result: Err(anyhow::anyhow!("rank 0 failed")), + }) + .unwrap(); + let _keep_peer_channel_connected = tx; + let mut receive_attempts = 0; + + let err = collect_runtime_responses(2, "test", &poison, || { + receive_attempts += 1; + rx.recv_timeout(std::time::Duration::from_millis(50)) + .map_err(|err| anyhow::anyhow!("waited for nonresponding rank: {err}")) + }) + .unwrap_err() + .to_string(); + + assert_eq!(receive_attempts, 1, "collector waited for the missing rank"); + assert!(err.contains("rank 0 failed")); + assert!(!err.contains("waited for nonresponding rank")); + assert!(poison.ensure_healthy().is_err()); + } + + #[test] + fn disconnected_runtime_response_poisons_executor() { + let poison = TpRuntimePoison::default(); + let (tx, rx) = mpsc::channel(); + drop(tx); + + let err = recv_runtime_response(&rx, "test", &poison) + .unwrap_err() + .to_string(); + assert!(err.contains("response channel disconnected during test")); + assert!(poison.ensure_healthy().is_err()); + } + + fn reply(rank: usize, reply: TpWorkerReply) -> TpWorkerResponse { + TpWorkerResponse { + rank, + result: Ok(reply), + } + } +} From 1e82b135e966fd257075a1039ef0db718c96fea8 Mon Sep 17 00:00:00 2001 From: CAICAIIs <3360776475@qq.com> Date: Tue, 8 Sep 2026 14:32:13 +0800 Subject: [PATCH 5/5] refactor(qwen35): split scheduler step drivers into scheduler/{steps,emit,telemetry}.rs Continue the scheduler.rs slim-down along the seams backend.rs/tp.rs already established: step drivers (batch prefill, overlap launch, unified step, pure decode, token dispatch/retirement, chunk vocabulary) move to steps.rs, terminal shutdown / closed-request pruning / rejection fan-out move to emit.rs, and the PEGAINFER_ITL_DEBUG telemetry trio moves to telemetry.rs. scheduler/mod.rs keeps the request/artifact types, the entry points (including the #1033 policy-aware start), load publishing, and the main loop. Behavior unchanged; moved items gain pub(super) only where the loop, backend, tp, emit siblings, or scheduler tests name them. Signed-off-by: CAICAIIs <3360776475@qq.com> --- pegainfer-qwen35/src/scheduler/emit.rs | 132 +++ pegainfer-qwen35/src/scheduler/mod.rs | 1037 +------------------ pegainfer-qwen35/src/scheduler/steps.rs | 862 +++++++++++++++ pegainfer-qwen35/src/scheduler/telemetry.rs | 52 + 4 files changed, 1052 insertions(+), 1031 deletions(-) create mode 100644 pegainfer-qwen35/src/scheduler/emit.rs create mode 100644 pegainfer-qwen35/src/scheduler/steps.rs create mode 100644 pegainfer-qwen35/src/scheduler/telemetry.rs diff --git a/pegainfer-qwen35/src/scheduler/emit.rs b/pegainfer-qwen35/src/scheduler/emit.rs new file mode 100644 index 000000000..ae4a8e8f3 --- /dev/null +++ b/pegainfer-qwen35/src/scheduler/emit.rs @@ -0,0 +1,132 @@ +//! Terminal request emission: shutdown fan-out, closed-request pruning, and +//! rejection / error delivery to callers. + +use super::*; + +pub(super) fn terminal_scheduler_shutdown( + submit_rx: &mut mpsc::UnboundedReceiver, + load_tx: &watch::Sender, + kv_total_blocks: u64, + active: Vec, + prefilling: Vec, + pending: Vec, + deferred: Vec, + inflight_prefill: Option, + failure: FatalSchedulerError, +) { + submit_rx.close(); + + let mut requests = failure.transient; + requests.extend(active.into_iter().map(Into::into)); + requests.extend(prefilling.into_iter().map(Into::into)); + requests.extend(pending.into_iter().map(Into::into)); + requests.extend(deferred.into_iter().map(Into::into)); + if let Some(InflightPrefill { output, chunk, .. }) = inflight_prefill { + // The stream must drain before the chunk's KV/recurrent/conv state is + // released or transferred into terminal request ownership. + drop(output); + requests.extend(chunk.reqs.into_iter().map(Into::into)); + } + while let Ok((req, _kv_prefix)) = submit_rx.try_recv() { + requests.push(req.into()); + } + + warn!( + "Qwen3.5 TP scheduler terminating after replica failure: {}", + failure.message + ); + for request in requests { + request.send_error(&failure.message); + } + load_tx.send_replace(SchedulerMetrics { + kv_used_blocks: 0, + kv_total_blocks, + num_running_reqs: 0, + num_waiting_reqs: 0, + spec_decode: None, + }); +} + +pub(super) fn prune_closed_requests( + backend: &mut B, + active: &mut Vec, + prefilling: &mut Vec, + pending: &mut Vec, +) -> std::result::Result<(), FatalSchedulerError> +where + B: DecodeDispatchBackend + PrefillPromoteBackend, +{ + pending.retain(|req| !req.token_tx.is_closed()); + + for idx in (0..active.len()).rev() { + if active[idx].token_tx.is_closed() { + debug!( + "request pruned before scheduling: request_id={:?} phase=decode tokens_generated={}", + active[idx].request_id, active[idx].generated_count + ); + let removed = backend.take_active_request(active, idx); + if let Err(err) = backend.drop_active_state(&removed.backend_state) { + return Err(FatalSchedulerError::new(err.to_string()).with_request(removed)); + } + } + } + + for idx in (0..prefilling.len()).rev() { + if prefilling[idx].req.token_tx.is_closed() { + let removed = prefilling.remove(idx); + debug!( + "request pruned before scheduling: request_id={:?} phase=prefill cursor={}", + removed.req.request_id, removed.cursor + ); + let expectation = if removed.cursor == 0 { + DropExpectation::MustBeAbsent + } else { + DropExpectation::MustExist + }; + if let Err(err) = backend.drop_prefill_state(&removed.backend_state, expectation) { + return Err(FatalSchedulerError::new(err.to_string()).with_request(removed)); + } + } + } + Ok(()) +} + +pub(super) const UNSUPPORTED_PROMPT_LOGPROBS_MESSAGE: &str = + "prompt_logprobs is unsupported by the Qwen3.5 serving contract"; + +pub(super) fn reject_unsupported_prompt_logprobs(pending: &mut Vec) { + pending.retain(|req| { + if req.prompt_logprobs.is_none() { + return true; + } + let _ = req.token_tx.send(TokenEvent::Rejected { + message: UNSUPPORTED_PROMPT_LOGPROBS_MESSAGE.to_string(), + prompt_tokens: req.prompt_tokens.len(), + completion_tokens: 0, + }); + false + }); +} + +pub(super) fn send_rejection(req: &SchedulerRequest, reason: RejectReason) { + let message = match reason { + RejectReason::ContextLength { limit } => format!( + "request exceeds this model's maximum context length of {limit} tokens: requested {} (prompt={} + max_tokens={})", + req.prompt_tokens.len().saturating_add(req.max_tokens), + req.prompt_tokens.len(), + req.max_tokens + ), + RejectReason::KvBudget => { + let max_request_tokens = max_kv_tokens(req.prompt_tokens.len(), req.max_tokens); + format!( + "request requires more KV pages than this model instance can provide: prompt_tokens={}, max_request_tokens={max_request_tokens}", + req.prompt_tokens.len() + ) + } + }; + let _ = req.token_tx.send(TokenEvent::Rejected { + message, + prompt_tokens: req.prompt_tokens.len(), + completion_tokens: 0, + }); +} diff --git a/pegainfer-qwen35/src/scheduler/mod.rs b/pegainfer-qwen35/src/scheduler/mod.rs index 379470a8f..1c128019f 100644 --- a/pegainfer-qwen35/src/scheduler/mod.rs +++ b/pegainfer-qwen35/src/scheduler/mod.rs @@ -5,7 +5,10 @@ //! - `BatchDecodeGraphState` for CUDA Graph batch decode (stable-address slots) mod backend; +mod emit; mod plan; +mod steps; +mod telemetry; mod tp; use std::collections::HashMap; use std::collections::HashSet; @@ -44,6 +47,7 @@ use tokio::sync::mpsc; use tokio::sync::watch; use self::backend::*; +use self::emit::*; use self::plan::ActiveDecodeState; use self::plan::ActiveKvBudget; use self::plan::ExecutionPlan; @@ -57,6 +61,8 @@ use self::plan::max_kv_tokens; use self::plan::plan_prefill_chunks; use self::plan::prefilling_future_pages; use self::plan::slot_for_new_request; +use self::steps::*; +use self::telemetry::*; use self::tp::*; use crate::Qwen35DecodeOverlap; use crate::Qwen35SchedulerPolicy; @@ -282,52 +288,6 @@ impl FatalSchedulerError { pub const DEFAULT_MAX_PREFILL_TOKENS: usize = 1024; -/// Env-gated per-step ITL diagnostics (issue #470). When `PEGAINFER_ITL_DEBUG` -/// is set, the scheduler emits one `ITL_STEP` line per executed step, tagging -/// the plan kind, the *actual* prefill-chunk token count associated with the -/// action, the active decode width, and the CPU wall-time. This lets the -/// mixed-load bench separate serial Unified stalls from overlap launch, -/// decode, completion, and wait actions instead of relying on the coarse -/// `[submit, last-token]` injection window. Off by default: no cost on the -/// normal bench path. -fn itl_debug_enabled() -> bool { - static ENABLED: OnceLock = OnceLock::new(); - *ENABLED.get_or_init(|| std::env::var_os("PEGAINFER_ITL_DEBUG").is_some()) -} - -/// Monotonic microseconds since the first ITL step, so `ITL_STEP` timestamps -/// are correlatable within one process run (paired with wall-clock epoch us). -fn itl_debug_mono_us() -> u128 { - static ORIGIN: OnceLock = OnceLock::new(); - ORIGIN.get_or_init(Instant::now).elapsed().as_micros() -} - -fn log_itl_step( - step_start: Option, - plan: &str, - prefill_tokens: usize, - prefill_reqs: usize, - decode_n: usize, -) { - let Some(step_start) = step_start else { - return; - }; - let dur_us = step_start.elapsed().as_micros(); - let epoch_us = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map_or(0, |d| d.as_micros()); - info!( - "ITL_STEP mono_us={} epoch_us={} plan={} prefill_tok={} prefill_reqs={} decode_n={} dur_us={}", - itl_debug_mono_us(), - epoch_us, - plan, - prefill_tokens, - prefill_reqs, - decode_n, - dur_us - ); -} - // ── Entry point ───────────────────────────────────────────────────────── pub fn start_with_capacity( @@ -536,111 +496,6 @@ fn should_block_on_submit(owned_work_empty: bool, inflight_prefill: bool) -> boo owned_work_empty && !inflight_prefill } -fn terminal_scheduler_shutdown( - submit_rx: &mut mpsc::UnboundedReceiver, - load_tx: &watch::Sender, - kv_total_blocks: u64, - active: Vec, - prefilling: Vec, - pending: Vec, - deferred: Vec, - inflight_prefill: Option, - failure: FatalSchedulerError, -) { - submit_rx.close(); - - let mut requests = failure.transient; - requests.extend(active.into_iter().map(Into::into)); - requests.extend(prefilling.into_iter().map(Into::into)); - requests.extend(pending.into_iter().map(Into::into)); - requests.extend(deferred.into_iter().map(Into::into)); - if let Some(InflightPrefill { output, chunk, .. }) = inflight_prefill { - // The stream must drain before the chunk's KV/recurrent/conv state is - // released or transferred into terminal request ownership. - drop(output); - requests.extend(chunk.reqs.into_iter().map(Into::into)); - } - while let Ok((req, _kv_prefix)) = submit_rx.try_recv() { - requests.push(req.into()); - } - - warn!( - "Qwen3.5 TP scheduler terminating after replica failure: {}", - failure.message - ); - for request in requests { - request.send_error(&failure.message); - } - load_tx.send_replace(SchedulerMetrics { - kv_used_blocks: 0, - kv_total_blocks, - num_running_reqs: 0, - num_waiting_reqs: 0, - spec_decode: None, - }); -} - -fn prune_closed_requests( - backend: &mut B, - active: &mut Vec, - prefilling: &mut Vec, - pending: &mut Vec, -) -> std::result::Result<(), FatalSchedulerError> -where - B: DecodeDispatchBackend + PrefillPromoteBackend, -{ - pending.retain(|req| !req.token_tx.is_closed()); - - for idx in (0..active.len()).rev() { - if active[idx].token_tx.is_closed() { - debug!( - "request pruned before scheduling: request_id={:?} phase=decode tokens_generated={}", - active[idx].request_id, active[idx].generated_count - ); - let removed = backend.take_active_request(active, idx); - if let Err(err) = backend.drop_active_state(&removed.backend_state) { - return Err(FatalSchedulerError::new(err.to_string()).with_request(removed)); - } - } - } - - for idx in (0..prefilling.len()).rev() { - if prefilling[idx].req.token_tx.is_closed() { - let removed = prefilling.remove(idx); - debug!( - "request pruned before scheduling: request_id={:?} phase=prefill cursor={}", - removed.req.request_id, removed.cursor - ); - let expectation = if removed.cursor == 0 { - DropExpectation::MustBeAbsent - } else { - DropExpectation::MustExist - }; - if let Err(err) = backend.drop_prefill_state(&removed.backend_state, expectation) { - return Err(FatalSchedulerError::new(err.to_string()).with_request(removed)); - } - } - } - Ok(()) -} - -const UNSUPPORTED_PROMPT_LOGPROBS_MESSAGE: &str = - "prompt_logprobs is unsupported by the Qwen3.5 serving contract"; - -fn reject_unsupported_prompt_logprobs(pending: &mut Vec) { - pending.retain(|req| { - if req.prompt_logprobs.is_none() { - return true; - } - let _ = req.token_tx.send(TokenEvent::Rejected { - message: UNSUPPORTED_PROMPT_LOGPROBS_MESSAGE.to_string(), - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - }); - false - }); -} - #[allow(clippy::needless_pass_by_value)] fn scheduler_loop( mut backend: SchedulerBackend, @@ -1007,885 +862,5 @@ fn scheduler_loop( } } -fn send_rejection(req: &SchedulerRequest, reason: RejectReason) { - let message = match reason { - RejectReason::ContextLength { limit } => format!( - "request exceeds this model's maximum context length of {limit} tokens: requested {} (prompt={} + max_tokens={})", - req.prompt_tokens.len().saturating_add(req.max_tokens), - req.prompt_tokens.len(), - req.max_tokens - ), - RejectReason::KvBudget => { - let max_request_tokens = max_kv_tokens(req.prompt_tokens.len(), req.max_tokens); - format!( - "request requires more KV pages than this model instance can provide: prompt_tokens={}, max_request_tokens={max_request_tokens}", - req.prompt_tokens.len() - ) - } - }; - let _ = req.token_tx.send(TokenEvent::Rejected { - message, - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - }); -} - -// ── Batch prefill ─────────────────────────────────────────────────────── - -fn prefill_batch( - backend: &mut SchedulerBackend, - active: &mut Vec, - scheduled: Vec, - prefilling: &mut Vec, - rng: &mut StdRng, -) -> std::result::Result<(), FatalSchedulerError> { - let mut chunk = ScheduledChunk::from(scheduled); - let sample_seed = rand::RngExt::random(rng); - let artifacts = match backend { - SchedulerBackend::Single(single) => { - // Scope the borrows of `chunk` to the executor call so the error path can - // move `chunk` into `fail_chunk`. - let logits = match single.batch_prefill_logits(&mut chunk) { - Ok(v) => v, - Err(e) => { - warn!("batch prefill failed: {e}"); - fail_chunk(chunk, &e.to_string()); - return Ok(()); - } - }; - let prefill_sample_seed = rand::RngExt::random(rng); - match single.sample_prefill_logits(&chunk.reqs, &logits, prefill_sample_seed) { - Ok((tokens, logprobs)) => PrefillStepArtifacts::Single { tokens, logprobs }, - Err(e) => { - warn!("prefill sampling failed: {e}"); - fail_chunk(chunk, &e.to_string()); - return Ok(()); - } - } - } - SchedulerBackend::Tp(tp) => match tp.execute_prefill_chunk(&chunk, sample_seed) { - Ok(v) => PrefillStepArtifacts::Tp(v), - Err(e) => { - warn!("TP prefill chunk failed: {e}"); - return Err(FatalSchedulerError::new(e.to_string()).with_requests(chunk.reqs)); - } - }, - }; - - promote_or_requeue(backend, active, prefilling, chunk, &artifacts) -} - -fn launch_overlap_step( - backend: &mut SchedulerBackend, - active: &mut Vec, - scheduled: Vec, - inflight_prefill: &mut Option, - rng: &mut StdRng, -) -> std::result::Result<(), FatalSchedulerError> { - debug_assert!(inflight_prefill.is_none()); - let mut chunk = ScheduledChunk::from(scheduled); - let decode_seed = rand::RngExt::random(rng); - let prefill_seed = rand::RngExt::random(rng); - let output = match backend { - SchedulerBackend::Single(single) => single.launch_async_prefill(&mut chunk), - SchedulerBackend::Tp(_) => unreachable!("Qwen3.5 TP cannot launch async prefill"), - }; - match output { - Ok(output) => { - *inflight_prefill = Some(InflightPrefill { - chunk, - output, - sample_seed: prefill_seed, - }); - } - Err(err) => { - warn!("async prefill launch failed: {err}"); - fail_chunk(chunk, &err.to_string()); - } - } - decode_step_with_seed(backend, active, decode_seed) -} - -fn finish_async_prefill( - backend: &mut SchedulerBackend, - active: &mut Vec, - prefilling: &mut Vec, - inflight: InflightPrefill, -) -> std::result::Result<(), FatalSchedulerError> { - let InflightPrefill { - chunk, - output, - sample_seed, - } = inflight; - let logits = output.into_logits(); - let SchedulerBackend::Single(single) = backend else { - unreachable!("Qwen3.5 TP cannot finish async prefill"); - }; - let (tokens, logprobs) = match single.sample_prefill_logits(&chunk.reqs, &logits, sample_seed) { - Ok(result) => result, - Err(err) => { - warn!("async prefill sampling failed: {err}"); - fail_chunk(chunk, &err.to_string()); - return Ok(()); - } - }; - let artifacts = PrefillStepArtifacts::Single { tokens, logprobs }; - promote_or_requeue(single, active, prefilling, chunk, &artifacts) -} - -// ── Unified step (prefill chunk + decode in one forward pass) ────────────── - -fn unified_step_sched( - backend: &mut SchedulerBackend, - active: &mut Vec, - scheduled: Vec, - prefilling: &mut Vec, - rng: &mut StdRng, -) -> std::result::Result<(), FatalSchedulerError> { - let mut chunk = ScheduledChunk::from(scheduled); - if matches!(backend, SchedulerBackend::Tp(_)) { - // Preserve the established scheduler RNG order: decode seed first, - // prefill seed second. Workers execute the forwards in the opposite - // (prefill-then-decode) order using these preselected seeds. - let decode_sample_seed = rand::RngExt::random(rng); - let prefill_sample_seed = rand::RngExt::random(rng); - let result = { - let SchedulerBackend::Tp(tp) = backend else { - unreachable!() - }; - tp.execute_unified(&chunk, active, decode_sample_seed, prefill_sample_seed) - }; - let artifacts = match result { - Ok(artifacts) => artifacts, - Err(err) => { - warn!("TP unified step failed: {err}"); - return Err(FatalSchedulerError::new(err.to_string()).with_requests(chunk.reqs)); - } - }; - - let (decode_tokens, decode_logprobs) = split_decode_artifacts(&artifacts.decode); - if let Err(failure) = - dispatch_decode_tokens(backend, active, &decode_tokens, &decode_logprobs) - { - return Err(failure.with_requests(chunk.reqs)); - } - - let prefill = PrefillStepArtifacts::Tp(artifacts.prefill); - return promote_or_requeue(backend, active, prefilling, chunk, &prefill); - } - - let SchedulerBackend::Single(backend) = backend else { - unreachable!() - }; - // Scope the borrows of `chunk` / `active` to the executor call so the error - // and decode-processing paths can use them afterwards. - let result = backend.unified_step(&mut chunk, active); - let output = match result { - Ok(v) => v, - Err(e) => { - warn!("unified step failed: {e}"); - let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } - fail_chunk(chunk, &message); - return Ok(()); - } - }; - let decode_seed = rand::RngExt::random(rng); - let prefill_seed = rand::RngExt::random(rng); - - // Process decode results FIRST (it may retire requests and free graph slots - // that promotion then fills densely). - if output.decoded { - process_decode_logits(backend, active, decode_seed)?; - } - - let prefill_logits = output - .prefill_logits - .as_ref() - .expect("scheduled prefill chunk must return prefill logits"); - let (tokens, logprobs) = - match backend.sample_prefill_logits(&chunk.reqs, prefill_logits, prefill_seed) { - Ok(v) => v, - Err(e) => { - warn!("unified prefill sampling failed: {e}"); - fail_chunk(chunk, &e.to_string()); - return Ok(()); - } - }; - let prefill = PrefillStepArtifacts::Single { tokens, logprobs }; - promote_or_requeue(backend, active, prefilling, chunk, &prefill) -} - -// ── Decode step (pure decode, CUDA Graph enabled) ────────────────────── - -fn decode_step( - backend: &mut SchedulerBackend, - active: &mut Vec, - rng: &mut StdRng, -) -> std::result::Result<(), FatalSchedulerError> { - // Preserve the historical scheduler RNG sequence: TP consumes the first - // seed, while single-GPU decode consumed a second seed inside sampling. - let first_seed = rand::RngExt::random(rng); - let sample_seed = if matches!(backend, SchedulerBackend::Single(_)) { - rand::RngExt::random(rng) - } else { - first_seed - }; - decode_step_with_seed(backend, active, sample_seed) -} - -fn decode_step_with_seed( - backend: &mut SchedulerBackend, - active: &mut Vec, - sample_seed: u64, -) -> std::result::Result<(), FatalSchedulerError> { - let (tokens, logprobs_vec) = match backend { - SchedulerBackend::Single(single) => { - if let Err(e) = single.decode_graph(active) { - warn!("batch_decode_graph error: {e}"); - let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } - return Ok(()); - } - // Snapshot logits to CPU BEFORE sampling (sampling may modify bufs.logits) - match single.sample_decode_logits(active, sample_seed) { - Ok(v) => v, - Err(e) => { - warn!("decode sampling/logprobs error: {e}"); - let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } - return Ok(()); - } - } - } - SchedulerBackend::Tp(tp) => match tp.execute_decode(active, sample_seed) { - Ok(v) => split_decode_artifacts(&v), - Err(e) => { - warn!("TP eager decode error: {e}"); - return Err(FatalSchedulerError::new(e.to_string())); - } - }, - }; - - dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec) -} - -/// Process decode logits from unified step: sample, extract logprobs, dispatch. -fn process_decode_logits( - backend: &mut SingleGpuBackend, - active: &mut Vec, - sample_seed: u64, -) -> std::result::Result<(), FatalSchedulerError> { - let (tokens, logprobs_vec) = match backend.sample_decode_logits(active, sample_seed) { - Ok(v) => v, - Err(e) => { - warn!("decode sampling/logprobs error: {e}"); - let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } - return Ok(()); - } - }; - - dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec) -} - -/// Dispatch sampled decode tokens: send events, check EOS/limits, retire finished. -/// -/// `tokens` and `logprobs` are indexed by original position in `active`. -/// Retirements collected first, then compacted in reverse order. -fn dispatch_decode_tokens( - backend: &mut impl DecodeDispatchBackend, - active: &mut Vec, - tokens: &[u32], - logprobs: &[Option], -) -> std::result::Result<(), FatalSchedulerError> { - enum Retirement { - Completion(Vec), - CleanupOnly, - Disconnected, - } - - let n = active.len(); - let mut to_retire = Vec::new(); - - for i in 0..n { - let token = tokens[i]; - let logprob = logprobs[i].clone(); - let req = &mut active[i]; - req.generated_count += 1; - - let is_eos = !req.params.ignore_eos && backend.is_stop_token(token); - let at_limit = req.generated_count >= req.max_tokens; - - if is_eos { - debug!( - "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", - req.request_id, - req.prompt_len, - req.generated_count, - FinishReason::Stop - ); - let event = TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }; - if backend.completion_requires_drop_ack() { - to_retire.push((i, Retirement::Completion(vec![event]))); - } else { - let _ = req.token_tx.send(event); - to_retire.push((i, Retirement::CleanupOnly)); - } - } else if at_limit { - debug!( - "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", - req.request_id, - req.prompt_len, - req.generated_count, - FinishReason::Length - ); - let events = vec![ - TokenEvent::Token { id: token, logprob }, - TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }, - ]; - if backend.completion_requires_drop_ack() { - to_retire.push((i, Retirement::Completion(events))); - } else { - for event in events { - let _ = req.token_tx.send(event); - } - to_retire.push((i, Retirement::CleanupOnly)); - } - } else if req - .token_tx - .send(TokenEvent::Token { id: token, logprob }) - .is_err() - { - debug!( - "request dropped: client disconnected: request_id={:?} tokens_generated={}", - req.request_id, req.generated_count - ); - to_retire.push((i, Retirement::Disconnected)); - } else { - req.last_token = token; - } - } - - // Remove in reverse order so compact_slot indices stay valid - for (i, retirement) in to_retire.into_iter().rev() { - let request = backend.take_active_request(active, i); - match retirement { - Retirement::Completion(final_events) => { - let candidate = CompletionCandidate { - request, - final_events, - }; - if let Err(err) = backend.drop_active_state(&candidate.request.backend_state) { - return Err(FatalSchedulerError::new(err.to_string()) - .with_request(candidate.into_terminal())); - } - candidate.commit(); - } - Retirement::CleanupOnly | Retirement::Disconnected => { - if let Err(err) = backend.drop_active_state(&request.backend_state) { - return Err(FatalSchedulerError::new(err.to_string()).with_request(request)); - } - } - } - } - Ok(()) -} - -trait DecodeDispatchBackend { - fn is_stop_token(&self, token: u32) -> bool; - fn completion_requires_drop_ack(&self) -> bool; - fn take_active_request( - &mut self, - active: &mut Vec, - idx: usize, - ) -> ActiveRequest35; - fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()>; -} - -impl DecodeDispatchBackend for SingleGpuBackend { - fn is_stop_token(&self, token: u32) -> bool { - self.is_stop_token(token) - } - - fn completion_requires_drop_ack(&self) -> bool { - false - } - - fn take_active_request( - &mut self, - active: &mut Vec, - idx: usize, - ) -> ActiveRequest35 { - compact_single_slot(self, active, idx) - } - - fn drop_active_state(&mut self, _state: &ActiveBackendState) -> Result<()> { - Ok(()) - } -} - -impl DecodeDispatchBackend for SchedulerBackend { - fn is_stop_token(&self, token: u32) -> bool { - self.is_stop_token(token) - } - - fn completion_requires_drop_ack(&self) -> bool { - matches!(self, SchedulerBackend::Tp(_)) - } - - fn take_active_request( - &mut self, - active: &mut Vec, - idx: usize, - ) -> ActiveRequest35 { - match self { - SchedulerBackend::Single(backend) => compact_single_slot(backend, active, idx), - SchedulerBackend::Tp(backend) => backend.take_active_request(active, idx), - } - } - - fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { - match (self, state) { - (SchedulerBackend::Single(_), ActiveBackendState::Single { .. }) => Ok(()), - (SchedulerBackend::Tp(backend), ActiveBackendState::Tp { request_id, .. }) => { - let compaction = backend.pending_compaction.take(); - backend.executor.drop_request_with_compaction( - *request_id, - DropExpectation::MustExist, - compaction, - ) - } - _ => anyhow::bail!("mismatched Qwen3.5 scheduler backend state during retirement"), - } - } -} - -/// Remove single-GPU request at `idx` via swap_remove and compact graph slots. -/// -/// After swap_remove, the element that was at `active.len()-1` (before remove) -/// now sits at `idx`. Its graph slot must be copied into the vacated slot so -/// that slots 0..active.len() remain dense. -fn compact_single_slot( - backend: &mut SingleGpuBackend, - active: &mut Vec, - idx: usize, -) -> ActiveRequest35 { - let compaction = compaction_after_retire(active.len(), idx); - let removed = active.swap_remove(idx); - - if let Some(compaction) = compaction { - backend.compact_slot(active, compaction); - } - removed -} - -// ── Chunked-prefill helpers ──────────────────────────────────────────────── - -/// Step's scheduled prefill set -struct ScheduledChunk { - reqs: Vec, - backend_state: ScheduledChunkBackendState, - /// Prompt cursor after this step's chunk - ends: Vec, - /// This step's chunked token slice per request - windows: Vec>, -} - -struct InflightPrefill { - // Fields drop in declaration order. Drain the stream before request state - // can return KV pages or release recurrent/convolution buffers on unwind. - output: AsyncPrefillOutput, - chunk: ScheduledChunk, - sample_seed: u64, -} - -enum ScheduledChunkBackendState { - Single { - kvs: Vec, - recs: Vec, - }, - Tp { - request_ids: Vec, - }, -} - -impl From> for ScheduledChunk { - fn from(scheduled: Vec) -> Self { - let n = scheduled.len(); - let is_tp = scheduled - .first() - .is_some_and(|p| matches!(p.backend_state, PrefillBackendState::Tp { .. })); - let mut chunk = ScheduledChunk { - reqs: Vec::with_capacity(n), - backend_state: if is_tp { - ScheduledChunkBackendState::Tp { - request_ids: Vec::with_capacity(n), - } - } else { - ScheduledChunkBackendState::Single { - kvs: Vec::with_capacity(n), - recs: Vec::with_capacity(n), - } - }, - ends: Vec::with_capacity(n), - windows: Vec::with_capacity(n), - }; - for p in scheduled { - let end = p.cursor + p.step_chunk; - chunk - .windows - .push(p.req.prompt_tokens[p.cursor..end].to_vec()); - chunk.ends.push(end); - chunk.reqs.push(p.req); - match (&mut chunk.backend_state, p.backend_state) { - ( - ScheduledChunkBackendState::Single { kvs, recs }, - PrefillBackendState::Single { kv, rec }, - ) => { - kvs.push(kv); - recs.push(rec); - } - ( - ScheduledChunkBackendState::Tp { request_ids }, - PrefillBackendState::Tp { request_id }, - ) => request_ids.push(request_id), - _ => unreachable!("mixed Qwen3.5 scheduler backend states in one chunk"), - } - } - chunk - } -} - -/// Pull this step's prefill set off the FRONT of `prefilling`, capping the -/// step's total forwarded prompt tokens at `prefill_budget`. -fn take_prefill_chunks( - prefilling: &mut Vec, - prefill_budget: usize, -) -> Vec { - let remaining: Vec = prefilling - .iter() - .map(|p| p.req.prompt_tokens.len() - p.cursor) - .collect(); - let chunks = plan_prefill_chunks(&remaining, prefill_budget); - let mut scheduled: Vec = prefilling.drain(0..chunks.len()).collect(); - for (p, chunk) in scheduled.iter_mut().zip(&chunks) { - p.step_chunk = *chunk; - } - scheduled -} - -/// Report a forward/sampling failure to every request in the failed chunk. -fn fail_chunk(chunk: ScheduledChunk, message: &str) { - for req in chunk.reqs { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.to_string(), - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - }); - } -} - -/// For each request in the just-prefilled chunk: if its prompt is now exhausted, -/// sample its first token, emit events, and move it into the decode batch; -/// otherwise re-queue it (with an advanced cursor) at the FRONT of `prefilling`. -/// `artifacts` are indexed by request order in `chunk`. -fn promote_or_requeue( - backend: &mut impl PrefillPromoteBackend, - active: &mut Vec, - prefilling: &mut Vec, - chunk: ScheduledChunk, - artifacts: &PrefillStepArtifacts, -) -> std::result::Result<(), FatalSchedulerError> { - let ScheduledChunk { - reqs, - backend_state, - ends, - .. - } = chunk; - let mut still_prefilling: Vec = Vec::new(); - let backend_states = split_scheduled_backend_state(backend_state); - let mut entries: VecDeque<_> = reqs - .into_iter() - .zip(backend_states) - .zip(ends) - .enumerate() - .map(|(i, ((req, backend_state), end))| (i, req, backend_state, end)) - .collect(); - - while let Some((i, req, backend_state, end)) = entries.pop_front() { - // Not finished: re-queue with the advanced cursor - if end < req.prompt_tokens.len() { - still_prefilling.push(PrefillingRequest35 { - req, - backend_state, - cursor: end, - step_chunk: 0, - }); - continue; - } - - let prompt_len = req.prompt_tokens.len(); - let artifact = artifacts.final_artifact(i); - let first_token = artifact.token; - let logprob = artifact.logprob; - - if !req.params.ignore_eos && backend.is_stop_token(first_token) { - debug!( - "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", - req.request_id, - prompt_len, - 0, - FinishReason::Stop - ); - let candidate = CompletionCandidate { - request: PrefillCompletionRequest { req, backend_state }, - final_events: vec![TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: prompt_len, - completion_tokens: 0, - }], - }; - if let Err(err) = backend - .drop_prefill_state(&candidate.request.backend_state, DropExpectation::MustExist) - { - return Err(prefill_lifecycle_failure( - err.to_string(), - candidate.into_terminal(), - still_prefilling, - entries, - )); - } - candidate.commit(); - continue; - } - - if req.max_tokens <= 1 { - debug!( - "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", - req.request_id, - prompt_len, - 1, - FinishReason::Length - ); - let candidate = CompletionCandidate { - request: PrefillCompletionRequest { req, backend_state }, - final_events: vec![ - TokenEvent::Token { - id: first_token, - logprob, - }, - TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: prompt_len, - completion_tokens: 1, - }, - ], - }; - if let Err(err) = backend - .drop_prefill_state(&candidate.request.backend_state, DropExpectation::MustExist) - { - return Err(prefill_lifecycle_failure( - err.to_string(), - candidate.into_terminal(), - still_prefilling, - entries, - )); - } - candidate.commit(); - continue; - } - - if req - .token_tx - .send(TokenEvent::Token { - id: first_token, - logprob, - }) - .is_err() - { - debug!( - "request dropped: client disconnected: request_id={:?} tokens_generated={}", - req.request_id, 0 - ); - let removed = PrefillCompletionRequest { req, backend_state }; - if let Err(err) = - backend.drop_prefill_state(&removed.backend_state, DropExpectation::MustExist) - { - return Err(prefill_lifecycle_failure( - err.to_string(), - removed.into_terminal(), - still_prefilling, - entries, - )); - } - continue; - } - - let active_backend_state = backend.promote_prefill_state(active.len(), backend_state); - active.push(ActiveRequest35 { - request_id: req.request_id, - token_tx: req.token_tx, - backend_state: active_backend_state, - last_token: first_token, - generated_count: 1, - max_tokens: req.max_tokens, - prompt_len, - params: req.params, - logprobs: req.logprobs, - }); - } - - prefilling.splice(0..0, still_prefilling); - Ok(()) -} - -fn prefill_lifecycle_failure( - message: String, - current: TerminalRequest, - still_prefilling: Vec, - remaining: VecDeque<(usize, SchedulerRequest, PrefillBackendState, usize)>, -) -> FatalSchedulerError { - FatalSchedulerError::new(message) - .with_request(current) - .with_requests(still_prefilling) - .with_requests(remaining.into_iter().map(|(_, req, _, _)| req)) -} - -trait PrefillPromoteBackend { - fn is_stop_token(&self, token: u32) -> bool; - fn promote_prefill_state( - &mut self, - active_len: usize, - state: PrefillBackendState, - ) -> ActiveBackendState; - fn drop_prefill_state( - &mut self, - state: &PrefillBackendState, - expectation: DropExpectation, - ) -> Result<()>; -} - -impl PrefillPromoteBackend for SingleGpuBackend { - fn is_stop_token(&self, token: u32) -> bool { - self.is_stop_token(token) - } - - fn promote_prefill_state( - &mut self, - active_len: usize, - state: PrefillBackendState, - ) -> ActiveBackendState { - let PrefillBackendState::Single { kv, rec } = state else { - panic!("single-GPU promotion received TP prefill state"); - }; - let slot_idx = slot_for_new_request(active_len, self.max_batch()) - .expect("admission must reserve a graph slot"); - self.copy_recurrent_to_slot(&rec, slot_idx) - .expect("copy recurrent state to slot failed"); - ActiveBackendState::Single { - kv, - graph_slot_idx: slot_idx, - } - } - - fn drop_prefill_state( - &mut self, - _state: &PrefillBackendState, - _expectation: DropExpectation, - ) -> Result<()> { - Ok(()) - } -} - -impl PrefillPromoteBackend for SchedulerBackend { - fn is_stop_token(&self, token: u32) -> bool { - self.is_stop_token(token) - } - - fn promote_prefill_state( - &mut self, - active_len: usize, - state: PrefillBackendState, - ) -> ActiveBackendState { - match (self, state) { - (SchedulerBackend::Single(single), state @ PrefillBackendState::Single { .. }) => { - single.promote_prefill_state(active_len, state) - } - (SchedulerBackend::Tp(backend), PrefillBackendState::Tp { request_id }) => { - let slot_idx = slot_for_new_request(active_len, backend.max_batch()) - .expect("admission must reserve a TP decode slot"); - ActiveBackendState::Tp { - request_id, - slot_idx, - } - } - _ => panic!("mismatched Qwen3.5 scheduler backend state during promotion"), - } - } - - fn drop_prefill_state( - &mut self, - state: &PrefillBackendState, - expectation: DropExpectation, - ) -> Result<()> { - match (self, state) { - (SchedulerBackend::Single(_), PrefillBackendState::Single { .. }) => Ok(()), - (SchedulerBackend::Tp(backend), PrefillBackendState::Tp { request_id }) => { - backend.drop_request(*request_id, expectation) - } - _ => anyhow::bail!("mismatched Qwen3.5 scheduler backend state during prefill drop"), - } - } -} - -fn split_scheduled_backend_state( - backend_state: ScheduledChunkBackendState, -) -> Vec { - match backend_state { - ScheduledChunkBackendState::Single { kvs, recs } => kvs - .into_iter() - .zip(recs) - .map(|(kv, rec)| PrefillBackendState::Single { kv, rec }) - .collect(), - ScheduledChunkBackendState::Tp { request_ids } => request_ids - .into_iter() - .map(|request_id| PrefillBackendState::Tp { request_id }) - .collect(), - } -} - #[cfg(test)] mod tests; diff --git a/pegainfer-qwen35/src/scheduler/steps.rs b/pegainfer-qwen35/src/scheduler/steps.rs new file mode 100644 index 000000000..a76222571 --- /dev/null +++ b/pegainfer-qwen35/src/scheduler/steps.rs @@ -0,0 +1,862 @@ +//! Scheduler step drivers: batch prefill, overlap launch, unified +//! prefill+decode, pure decode, token dispatch/retirement, and the +//! chunked-prefill vocabulary shared with the backend. + +use super::*; + +// ── Batch prefill ─────────────────────────────────────────────────────── + +pub(super) fn prefill_batch( + backend: &mut SchedulerBackend, + active: &mut Vec, + scheduled: Vec, + prefilling: &mut Vec, + rng: &mut StdRng, +) -> std::result::Result<(), FatalSchedulerError> { + let mut chunk = ScheduledChunk::from(scheduled); + let sample_seed = rand::RngExt::random(rng); + let artifacts = match backend { + SchedulerBackend::Single(single) => { + // Scope the borrows of `chunk` to the executor call so the error path can + // move `chunk` into `fail_chunk`. + let logits = match single.batch_prefill_logits(&mut chunk) { + Ok(v) => v, + Err(e) => { + warn!("batch prefill failed: {e}"); + fail_chunk(chunk, &e.to_string()); + return Ok(()); + } + }; + let prefill_sample_seed = rand::RngExt::random(rng); + match single.sample_prefill_logits(&chunk.reqs, &logits, prefill_sample_seed) { + Ok((tokens, logprobs)) => PrefillStepArtifacts::Single { tokens, logprobs }, + Err(e) => { + warn!("prefill sampling failed: {e}"); + fail_chunk(chunk, &e.to_string()); + return Ok(()); + } + } + } + SchedulerBackend::Tp(tp) => match tp.execute_prefill_chunk(&chunk, sample_seed) { + Ok(v) => PrefillStepArtifacts::Tp(v), + Err(e) => { + warn!("TP prefill chunk failed: {e}"); + return Err(FatalSchedulerError::new(e.to_string()).with_requests(chunk.reqs)); + } + }, + }; + + promote_or_requeue(backend, active, prefilling, chunk, &artifacts) +} + +pub(super) fn launch_overlap_step( + backend: &mut SchedulerBackend, + active: &mut Vec, + scheduled: Vec, + inflight_prefill: &mut Option, + rng: &mut StdRng, +) -> std::result::Result<(), FatalSchedulerError> { + debug_assert!(inflight_prefill.is_none()); + let mut chunk = ScheduledChunk::from(scheduled); + let decode_seed = rand::RngExt::random(rng); + let prefill_seed = rand::RngExt::random(rng); + let output = match backend { + SchedulerBackend::Single(single) => single.launch_async_prefill(&mut chunk), + SchedulerBackend::Tp(_) => unreachable!("Qwen3.5 TP cannot launch async prefill"), + }; + match output { + Ok(output) => { + *inflight_prefill = Some(InflightPrefill { + chunk, + output, + sample_seed: prefill_seed, + }); + } + Err(err) => { + warn!("async prefill launch failed: {err}"); + fail_chunk(chunk, &err.to_string()); + } + } + decode_step_with_seed(backend, active, decode_seed) +} + +pub(super) fn finish_async_prefill( + backend: &mut SchedulerBackend, + active: &mut Vec, + prefilling: &mut Vec, + inflight: InflightPrefill, +) -> std::result::Result<(), FatalSchedulerError> { + let InflightPrefill { + chunk, + output, + sample_seed, + } = inflight; + let logits = output.into_logits(); + let SchedulerBackend::Single(single) = backend else { + unreachable!("Qwen3.5 TP cannot finish async prefill"); + }; + let (tokens, logprobs) = match single.sample_prefill_logits(&chunk.reqs, &logits, sample_seed) { + Ok(result) => result, + Err(err) => { + warn!("async prefill sampling failed: {err}"); + fail_chunk(chunk, &err.to_string()); + return Ok(()); + } + }; + let artifacts = PrefillStepArtifacts::Single { tokens, logprobs }; + promote_or_requeue(single, active, prefilling, chunk, &artifacts) +} + +// ── Unified step (prefill chunk + decode in one forward pass) ────────────── + +pub(super) fn unified_step_sched( + backend: &mut SchedulerBackend, + active: &mut Vec, + scheduled: Vec, + prefilling: &mut Vec, + rng: &mut StdRng, +) -> std::result::Result<(), FatalSchedulerError> { + let mut chunk = ScheduledChunk::from(scheduled); + if matches!(backend, SchedulerBackend::Tp(_)) { + // Preserve the established scheduler RNG order: decode seed first, + // prefill seed second. Workers execute the forwards in the opposite + // (prefill-then-decode) order using these preselected seeds. + let decode_sample_seed = rand::RngExt::random(rng); + let prefill_sample_seed = rand::RngExt::random(rng); + let result = { + let SchedulerBackend::Tp(tp) = backend else { + unreachable!() + }; + tp.execute_unified(&chunk, active, decode_sample_seed, prefill_sample_seed) + }; + let artifacts = match result { + Ok(artifacts) => artifacts, + Err(err) => { + warn!("TP unified step failed: {err}"); + return Err(FatalSchedulerError::new(err.to_string()).with_requests(chunk.reqs)); + } + }; + + let (decode_tokens, decode_logprobs) = split_decode_artifacts(&artifacts.decode); + if let Err(failure) = + dispatch_decode_tokens(backend, active, &decode_tokens, &decode_logprobs) + { + return Err(failure.with_requests(chunk.reqs)); + } + + let prefill = PrefillStepArtifacts::Tp(artifacts.prefill); + return promote_or_requeue(backend, active, prefilling, chunk, &prefill); + } + + let SchedulerBackend::Single(backend) = backend else { + unreachable!() + }; + // Scope the borrows of `chunk` / `active` to the executor call so the error + // and decode-processing paths can use them afterwards. + let result = backend.unified_step(&mut chunk, active); + let output = match result { + Ok(v) => v, + Err(e) => { + warn!("unified step failed: {e}"); + let message = e.to_string(); + for req in active.drain(..) { + let _ = req.token_tx.send(TokenEvent::Error { + message: message.clone(), + prompt_tokens: req.prompt_len, + completion_tokens: req.generated_count, + }); + } + fail_chunk(chunk, &message); + return Ok(()); + } + }; + let decode_seed = rand::RngExt::random(rng); + let prefill_seed = rand::RngExt::random(rng); + + // Process decode results FIRST (it may retire requests and free graph slots + // that promotion then fills densely). + if output.decoded { + process_decode_logits(backend, active, decode_seed)?; + } + + let prefill_logits = output + .prefill_logits + .as_ref() + .expect("scheduled prefill chunk must return prefill logits"); + let (tokens, logprobs) = + match backend.sample_prefill_logits(&chunk.reqs, prefill_logits, prefill_seed) { + Ok(v) => v, + Err(e) => { + warn!("unified prefill sampling failed: {e}"); + fail_chunk(chunk, &e.to_string()); + return Ok(()); + } + }; + let prefill = PrefillStepArtifacts::Single { tokens, logprobs }; + promote_or_requeue(backend, active, prefilling, chunk, &prefill) +} + +// ── Decode step (pure decode, CUDA Graph enabled) ────────────────────── + +pub(super) fn decode_step( + backend: &mut SchedulerBackend, + active: &mut Vec, + rng: &mut StdRng, +) -> std::result::Result<(), FatalSchedulerError> { + // Preserve the historical scheduler RNG sequence: TP consumes the first + // seed, while single-GPU decode consumed a second seed inside sampling. + let first_seed = rand::RngExt::random(rng); + let sample_seed = if matches!(backend, SchedulerBackend::Single(_)) { + rand::RngExt::random(rng) + } else { + first_seed + }; + decode_step_with_seed(backend, active, sample_seed) +} + +fn decode_step_with_seed( + backend: &mut SchedulerBackend, + active: &mut Vec, + sample_seed: u64, +) -> std::result::Result<(), FatalSchedulerError> { + let (tokens, logprobs_vec) = match backend { + SchedulerBackend::Single(single) => { + if let Err(e) = single.decode_graph(active) { + warn!("batch_decode_graph error: {e}"); + let message = e.to_string(); + for req in active.drain(..) { + let _ = req.token_tx.send(TokenEvent::Error { + message: message.clone(), + prompt_tokens: req.prompt_len, + completion_tokens: req.generated_count, + }); + } + return Ok(()); + } + // Snapshot logits to CPU BEFORE sampling (sampling may modify bufs.logits) + match single.sample_decode_logits(active, sample_seed) { + Ok(v) => v, + Err(e) => { + warn!("decode sampling/logprobs error: {e}"); + let message = e.to_string(); + for req in active.drain(..) { + let _ = req.token_tx.send(TokenEvent::Error { + message: message.clone(), + prompt_tokens: req.prompt_len, + completion_tokens: req.generated_count, + }); + } + return Ok(()); + } + } + } + SchedulerBackend::Tp(tp) => match tp.execute_decode(active, sample_seed) { + Ok(v) => split_decode_artifacts(&v), + Err(e) => { + warn!("TP eager decode error: {e}"); + return Err(FatalSchedulerError::new(e.to_string())); + } + }, + }; + + dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec) +} + +/// Process decode logits from unified step: sample, extract logprobs, dispatch. +fn process_decode_logits( + backend: &mut SingleGpuBackend, + active: &mut Vec, + sample_seed: u64, +) -> std::result::Result<(), FatalSchedulerError> { + let (tokens, logprobs_vec) = match backend.sample_decode_logits(active, sample_seed) { + Ok(v) => v, + Err(e) => { + warn!("decode sampling/logprobs error: {e}"); + let message = e.to_string(); + for req in active.drain(..) { + let _ = req.token_tx.send(TokenEvent::Error { + message: message.clone(), + prompt_tokens: req.prompt_len, + completion_tokens: req.generated_count, + }); + } + return Ok(()); + } + }; + + dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec) +} + +/// Dispatch sampled decode tokens: send events, check EOS/limits, retire finished. +/// +/// `tokens` and `logprobs` are indexed by original position in `active`. +/// Retirements collected first, then compacted in reverse order. +pub(super) fn dispatch_decode_tokens( + backend: &mut impl DecodeDispatchBackend, + active: &mut Vec, + tokens: &[u32], + logprobs: &[Option], +) -> std::result::Result<(), FatalSchedulerError> { + enum Retirement { + Completion(Vec), + CleanupOnly, + Disconnected, + } + + let n = active.len(); + let mut to_retire = Vec::new(); + + for i in 0..n { + let token = tokens[i]; + let logprob = logprobs[i].clone(); + let req = &mut active[i]; + req.generated_count += 1; + + let is_eos = !req.params.ignore_eos && backend.is_stop_token(token); + let at_limit = req.generated_count >= req.max_tokens; + + if is_eos { + debug!( + "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", + req.request_id, + req.prompt_len, + req.generated_count, + FinishReason::Stop + ); + let event = TokenEvent::Finished { + finish_reason: FinishReason::Stop, + prompt_tokens: req.prompt_len, + completion_tokens: req.generated_count, + }; + if backend.completion_requires_drop_ack() { + to_retire.push((i, Retirement::Completion(vec![event]))); + } else { + let _ = req.token_tx.send(event); + to_retire.push((i, Retirement::CleanupOnly)); + } + } else if at_limit { + debug!( + "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", + req.request_id, + req.prompt_len, + req.generated_count, + FinishReason::Length + ); + let events = vec![ + TokenEvent::Token { id: token, logprob }, + TokenEvent::Finished { + finish_reason: FinishReason::Length, + prompt_tokens: req.prompt_len, + completion_tokens: req.generated_count, + }, + ]; + if backend.completion_requires_drop_ack() { + to_retire.push((i, Retirement::Completion(events))); + } else { + for event in events { + let _ = req.token_tx.send(event); + } + to_retire.push((i, Retirement::CleanupOnly)); + } + } else if req + .token_tx + .send(TokenEvent::Token { id: token, logprob }) + .is_err() + { + debug!( + "request dropped: client disconnected: request_id={:?} tokens_generated={}", + req.request_id, req.generated_count + ); + to_retire.push((i, Retirement::Disconnected)); + } else { + req.last_token = token; + } + } + + // Remove in reverse order so compact_slot indices stay valid + for (i, retirement) in to_retire.into_iter().rev() { + let request = backend.take_active_request(active, i); + match retirement { + Retirement::Completion(final_events) => { + let candidate = CompletionCandidate { + request, + final_events, + }; + if let Err(err) = backend.drop_active_state(&candidate.request.backend_state) { + return Err(FatalSchedulerError::new(err.to_string()) + .with_request(candidate.into_terminal())); + } + candidate.commit(); + } + Retirement::CleanupOnly | Retirement::Disconnected => { + if let Err(err) = backend.drop_active_state(&request.backend_state) { + return Err(FatalSchedulerError::new(err.to_string()).with_request(request)); + } + } + } + } + Ok(()) +} + +pub(super) trait DecodeDispatchBackend { + fn is_stop_token(&self, token: u32) -> bool; + fn completion_requires_drop_ack(&self) -> bool; + fn take_active_request( + &mut self, + active: &mut Vec, + idx: usize, + ) -> ActiveRequest35; + fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()>; +} + +impl DecodeDispatchBackend for SingleGpuBackend { + fn is_stop_token(&self, token: u32) -> bool { + self.is_stop_token(token) + } + + fn completion_requires_drop_ack(&self) -> bool { + false + } + + fn take_active_request( + &mut self, + active: &mut Vec, + idx: usize, + ) -> ActiveRequest35 { + compact_single_slot(self, active, idx) + } + + fn drop_active_state(&mut self, _state: &ActiveBackendState) -> Result<()> { + Ok(()) + } +} + +impl DecodeDispatchBackend for SchedulerBackend { + fn is_stop_token(&self, token: u32) -> bool { + self.is_stop_token(token) + } + + fn completion_requires_drop_ack(&self) -> bool { + matches!(self, SchedulerBackend::Tp(_)) + } + + fn take_active_request( + &mut self, + active: &mut Vec, + idx: usize, + ) -> ActiveRequest35 { + match self { + SchedulerBackend::Single(backend) => compact_single_slot(backend, active, idx), + SchedulerBackend::Tp(backend) => backend.take_active_request(active, idx), + } + } + + fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { + match (self, state) { + (SchedulerBackend::Single(_), ActiveBackendState::Single { .. }) => Ok(()), + (SchedulerBackend::Tp(backend), ActiveBackendState::Tp { request_id, .. }) => { + let compaction = backend.pending_compaction.take(); + backend.executor.drop_request_with_compaction( + *request_id, + DropExpectation::MustExist, + compaction, + ) + } + _ => anyhow::bail!("mismatched Qwen3.5 scheduler backend state during retirement"), + } + } +} + +/// Remove single-GPU request at `idx` via swap_remove and compact graph slots. +/// +/// After swap_remove, the element that was at `active.len()-1` (before remove) +/// now sits at `idx`. Its graph slot must be copied into the vacated slot so +/// that slots 0..active.len() remain dense. +fn compact_single_slot( + backend: &mut SingleGpuBackend, + active: &mut Vec, + idx: usize, +) -> ActiveRequest35 { + let compaction = compaction_after_retire(active.len(), idx); + let removed = active.swap_remove(idx); + + if let Some(compaction) = compaction { + backend.compact_slot(active, compaction); + } + removed +} + +// ── Chunked-prefill helpers ──────────────────────────────────────────────── + +/// Step's scheduled prefill set +pub(super) struct ScheduledChunk { + pub(super) reqs: Vec, + pub(super) backend_state: ScheduledChunkBackendState, + /// Prompt cursor after this step's chunk + pub(super) ends: Vec, + /// This step's chunked token slice per request + pub(super) windows: Vec>, +} + +pub(super) struct InflightPrefill { + // Fields drop in declaration order. Drain the stream before request state + // can return KV pages or release recurrent/convolution buffers on unwind. + pub(super) output: AsyncPrefillOutput, + pub(super) chunk: ScheduledChunk, + pub(super) sample_seed: u64, +} + +pub(super) enum ScheduledChunkBackendState { + Single { + kvs: Vec, + recs: Vec, + }, + Tp { + request_ids: Vec, + }, +} + +impl From> for ScheduledChunk { + fn from(scheduled: Vec) -> Self { + let n = scheduled.len(); + let is_tp = scheduled + .first() + .is_some_and(|p| matches!(p.backend_state, PrefillBackendState::Tp { .. })); + let mut chunk = ScheduledChunk { + reqs: Vec::with_capacity(n), + backend_state: if is_tp { + ScheduledChunkBackendState::Tp { + request_ids: Vec::with_capacity(n), + } + } else { + ScheduledChunkBackendState::Single { + kvs: Vec::with_capacity(n), + recs: Vec::with_capacity(n), + } + }, + ends: Vec::with_capacity(n), + windows: Vec::with_capacity(n), + }; + for p in scheduled { + let end = p.cursor + p.step_chunk; + chunk + .windows + .push(p.req.prompt_tokens[p.cursor..end].to_vec()); + chunk.ends.push(end); + chunk.reqs.push(p.req); + match (&mut chunk.backend_state, p.backend_state) { + ( + ScheduledChunkBackendState::Single { kvs, recs }, + PrefillBackendState::Single { kv, rec }, + ) => { + kvs.push(kv); + recs.push(rec); + } + ( + ScheduledChunkBackendState::Tp { request_ids }, + PrefillBackendState::Tp { request_id }, + ) => request_ids.push(request_id), + _ => unreachable!("mixed Qwen3.5 scheduler backend states in one chunk"), + } + } + chunk + } +} + +/// Pull this step's prefill set off the FRONT of `prefilling`, capping the +/// step's total forwarded prompt tokens at `prefill_budget`. +pub(super) fn take_prefill_chunks( + prefilling: &mut Vec, + prefill_budget: usize, +) -> Vec { + let remaining: Vec = prefilling + .iter() + .map(|p| p.req.prompt_tokens.len() - p.cursor) + .collect(); + let chunks = plan_prefill_chunks(&remaining, prefill_budget); + let mut scheduled: Vec = prefilling.drain(0..chunks.len()).collect(); + for (p, chunk) in scheduled.iter_mut().zip(&chunks) { + p.step_chunk = *chunk; + } + scheduled +} + +/// Report a forward/sampling failure to every request in the failed chunk. +fn fail_chunk(chunk: ScheduledChunk, message: &str) { + for req in chunk.reqs { + let _ = req.token_tx.send(TokenEvent::Error { + message: message.to_string(), + prompt_tokens: req.prompt_tokens.len(), + completion_tokens: 0, + }); + } +} + +/// For each request in the just-prefilled chunk: if its prompt is now exhausted, +/// sample its first token, emit events, and move it into the decode batch; +/// otherwise re-queue it (with an advanced cursor) at the FRONT of `prefilling`. +/// `artifacts` are indexed by request order in `chunk`. +pub(super) fn promote_or_requeue( + backend: &mut impl PrefillPromoteBackend, + active: &mut Vec, + prefilling: &mut Vec, + chunk: ScheduledChunk, + artifacts: &PrefillStepArtifacts, +) -> std::result::Result<(), FatalSchedulerError> { + let ScheduledChunk { + reqs, + backend_state, + ends, + .. + } = chunk; + let mut still_prefilling: Vec = Vec::new(); + let backend_states = split_scheduled_backend_state(backend_state); + let mut entries: VecDeque<_> = reqs + .into_iter() + .zip(backend_states) + .zip(ends) + .enumerate() + .map(|(i, ((req, backend_state), end))| (i, req, backend_state, end)) + .collect(); + + while let Some((i, req, backend_state, end)) = entries.pop_front() { + // Not finished: re-queue with the advanced cursor + if end < req.prompt_tokens.len() { + still_prefilling.push(PrefillingRequest35 { + req, + backend_state, + cursor: end, + step_chunk: 0, + }); + continue; + } + + let prompt_len = req.prompt_tokens.len(); + let artifact = artifacts.final_artifact(i); + let first_token = artifact.token; + let logprob = artifact.logprob; + + if !req.params.ignore_eos && backend.is_stop_token(first_token) { + debug!( + "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", + req.request_id, + prompt_len, + 0, + FinishReason::Stop + ); + let candidate = CompletionCandidate { + request: PrefillCompletionRequest { req, backend_state }, + final_events: vec![TokenEvent::Finished { + finish_reason: FinishReason::Stop, + prompt_tokens: prompt_len, + completion_tokens: 0, + }], + }; + if let Err(err) = backend + .drop_prefill_state(&candidate.request.backend_state, DropExpectation::MustExist) + { + return Err(prefill_lifecycle_failure( + err.to_string(), + candidate.into_terminal(), + still_prefilling, + entries, + )); + } + candidate.commit(); + continue; + } + + if req.max_tokens <= 1 { + debug!( + "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", + req.request_id, + prompt_len, + 1, + FinishReason::Length + ); + let candidate = CompletionCandidate { + request: PrefillCompletionRequest { req, backend_state }, + final_events: vec![ + TokenEvent::Token { + id: first_token, + logprob, + }, + TokenEvent::Finished { + finish_reason: FinishReason::Length, + prompt_tokens: prompt_len, + completion_tokens: 1, + }, + ], + }; + if let Err(err) = backend + .drop_prefill_state(&candidate.request.backend_state, DropExpectation::MustExist) + { + return Err(prefill_lifecycle_failure( + err.to_string(), + candidate.into_terminal(), + still_prefilling, + entries, + )); + } + candidate.commit(); + continue; + } + + if req + .token_tx + .send(TokenEvent::Token { + id: first_token, + logprob, + }) + .is_err() + { + debug!( + "request dropped: client disconnected: request_id={:?} tokens_generated={}", + req.request_id, 0 + ); + let removed = PrefillCompletionRequest { req, backend_state }; + if let Err(err) = + backend.drop_prefill_state(&removed.backend_state, DropExpectation::MustExist) + { + return Err(prefill_lifecycle_failure( + err.to_string(), + removed.into_terminal(), + still_prefilling, + entries, + )); + } + continue; + } + + let active_backend_state = backend.promote_prefill_state(active.len(), backend_state); + active.push(ActiveRequest35 { + request_id: req.request_id, + token_tx: req.token_tx, + backend_state: active_backend_state, + last_token: first_token, + generated_count: 1, + max_tokens: req.max_tokens, + prompt_len, + params: req.params, + logprobs: req.logprobs, + }); + } + + prefilling.splice(0..0, still_prefilling); + Ok(()) +} + +fn prefill_lifecycle_failure( + message: String, + current: TerminalRequest, + still_prefilling: Vec, + remaining: VecDeque<(usize, SchedulerRequest, PrefillBackendState, usize)>, +) -> FatalSchedulerError { + FatalSchedulerError::new(message) + .with_request(current) + .with_requests(still_prefilling) + .with_requests(remaining.into_iter().map(|(_, req, _, _)| req)) +} + +pub(super) trait PrefillPromoteBackend { + fn is_stop_token(&self, token: u32) -> bool; + fn promote_prefill_state( + &mut self, + active_len: usize, + state: PrefillBackendState, + ) -> ActiveBackendState; + fn drop_prefill_state( + &mut self, + state: &PrefillBackendState, + expectation: DropExpectation, + ) -> Result<()>; +} + +impl PrefillPromoteBackend for SingleGpuBackend { + fn is_stop_token(&self, token: u32) -> bool { + self.is_stop_token(token) + } + + fn promote_prefill_state( + &mut self, + active_len: usize, + state: PrefillBackendState, + ) -> ActiveBackendState { + let PrefillBackendState::Single { kv, rec } = state else { + panic!("single-GPU promotion received TP prefill state"); + }; + let slot_idx = slot_for_new_request(active_len, self.max_batch()) + .expect("admission must reserve a graph slot"); + self.copy_recurrent_to_slot(&rec, slot_idx) + .expect("copy recurrent state to slot failed"); + ActiveBackendState::Single { + kv, + graph_slot_idx: slot_idx, + } + } + + fn drop_prefill_state( + &mut self, + _state: &PrefillBackendState, + _expectation: DropExpectation, + ) -> Result<()> { + Ok(()) + } +} + +impl PrefillPromoteBackend for SchedulerBackend { + fn is_stop_token(&self, token: u32) -> bool { + self.is_stop_token(token) + } + + fn promote_prefill_state( + &mut self, + active_len: usize, + state: PrefillBackendState, + ) -> ActiveBackendState { + match (self, state) { + (SchedulerBackend::Single(single), state @ PrefillBackendState::Single { .. }) => { + single.promote_prefill_state(active_len, state) + } + (SchedulerBackend::Tp(backend), PrefillBackendState::Tp { request_id }) => { + let slot_idx = slot_for_new_request(active_len, backend.max_batch()) + .expect("admission must reserve a TP decode slot"); + ActiveBackendState::Tp { + request_id, + slot_idx, + } + } + _ => panic!("mismatched Qwen3.5 scheduler backend state during promotion"), + } + } + + fn drop_prefill_state( + &mut self, + state: &PrefillBackendState, + expectation: DropExpectation, + ) -> Result<()> { + match (self, state) { + (SchedulerBackend::Single(_), PrefillBackendState::Single { .. }) => Ok(()), + (SchedulerBackend::Tp(backend), PrefillBackendState::Tp { request_id }) => { + backend.drop_request(*request_id, expectation) + } + _ => anyhow::bail!("mismatched Qwen3.5 scheduler backend state during prefill drop"), + } + } +} + +fn split_scheduled_backend_state( + backend_state: ScheduledChunkBackendState, +) -> Vec { + match backend_state { + ScheduledChunkBackendState::Single { kvs, recs } => kvs + .into_iter() + .zip(recs) + .map(|(kv, rec)| PrefillBackendState::Single { kv, rec }) + .collect(), + ScheduledChunkBackendState::Tp { request_ids } => request_ids + .into_iter() + .map(|request_id| PrefillBackendState::Tp { request_id }) + .collect(), + } +} diff --git a/pegainfer-qwen35/src/scheduler/telemetry.rs b/pegainfer-qwen35/src/scheduler/telemetry.rs new file mode 100644 index 000000000..68525ca68 --- /dev/null +++ b/pegainfer-qwen35/src/scheduler/telemetry.rs @@ -0,0 +1,52 @@ +//! Env-gated per-step ITL logging (`PEGAINFER_ITL_DEBUG`): one `ITL_STEP` +//! line per executed scheduler step with the plan kind, chunk/decode widths, +//! and CPU wall-time. Off by default; the check is cached so the serving path +//! pays one static lookup. + +use super::*; + +/// Env-gated per-step ITL diagnostics (issue #470). When `PEGAINFER_ITL_DEBUG` +/// is set, the scheduler emits one `ITL_STEP` line per executed step, tagging +/// the plan kind, the *actual* prefill-chunk token count associated with the +/// action, the active decode width, and the CPU wall-time. This lets the +/// mixed-load bench separate serial Unified stalls from overlap launch, +/// decode, completion, and wait actions instead of relying on the coarse +/// `[submit, last-token]` injection window. Off by default: no cost on the +/// normal bench path. +pub(super) fn itl_debug_enabled() -> bool { + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| std::env::var_os("PEGAINFER_ITL_DEBUG").is_some()) +} + +/// Monotonic microseconds since the first ITL step, so `ITL_STEP` timestamps +/// are correlatable within one process run (paired with wall-clock epoch us). +fn itl_debug_mono_us() -> u128 { + static ORIGIN: OnceLock = OnceLock::new(); + ORIGIN.get_or_init(Instant::now).elapsed().as_micros() +} + +pub(super) fn log_itl_step( + step_start: Option, + plan: &str, + prefill_tokens: usize, + prefill_reqs: usize, + decode_n: usize, +) { + let Some(step_start) = step_start else { + return; + }; + let dur_us = step_start.elapsed().as_micros(); + let epoch_us = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |d| d.as_micros()); + info!( + "ITL_STEP mono_us={} epoch_us={} plan={} prefill_tok={} prefill_reqs={} decode_n={} dur_us={}", + itl_debug_mono_us(), + epoch_us, + plan, + prefill_tokens, + prefill_reqs, + decode_n, + dur_us + ); +}