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/4] 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/4] 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/4] 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/4] 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), + } + } +}