From d9cd1e06667306c7bb76fd6fee115bc29e83a78b Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Thu, 20 Aug 2026 15:06:57 +0000 Subject: [PATCH 1/4] feat(qwen35): shard linear-attention/GDR state per TP rank Phase 2b ported onto #870: recurrent/conv state, GDR scratch, and the linear-attention weight surface are allocated and addressed per rank instead of replicated, which is what makes 27B TP2 fit on 2x48 GB cards. - weight_loader: additive stitch/shard loaders (2D row stitch for the fused qkv [q|k|v] segments, 1D element stitch for conv1d channels, bf16/f32 1D shards for dt_bias/A_log) - config: local_linear_* accessors mirroring the full-attn TP style; linear head divisibility fails closed in TensorParallelConfig - weights: per-rank stitched shards for in_proj_qkv/conv1d (head-local slices per segment), row shards for z/b/a, col shard for out_proj, dt_bias/A_log sliced, norm_weight kept replicated (head-shared); loader reserve uses TP-aware estimates - recurrent_state/decode_buffers/prefill_buffers: state and GDR scratch at local value-head/qkv sizes; capacity math derives from local allocation_bytes - batch_decode/prefill: local head counts into the GDR decode/conv/ Triton-AOT prefill chains, gated RMSNorm at local v heads, all-reduce after linear out_proj; batch_decode_full_attention_via_prefill is now TP-local so eager decode routes 27B TP2 group-6 full attention through prefill (was FlashInfer Unsupported group_size: 6) - tp_executor: worker capacity math and per-request state use the rank-local sizes; decode rows still run as a per-request bs=1 loop (batched in a follow-up) Recurrent/conv state is never all-reduced. Signed-off-by: Ziyang Zhang --- pegainfer-core/src/weight_loader.rs | 121 ++++++++ pegainfer-qwen35/src/batch_decode.rs | 88 +++++- pegainfer-qwen35/src/batch_decode_graph.rs | 2 +- pegainfer-qwen35/src/config/model.rs | 8 - pegainfer-qwen35/src/config/tp.rs | 128 +++++++- pegainfer-qwen35/src/decode_buffers.rs | 6 +- pegainfer-qwen35/src/executor.rs | 8 +- pegainfer-qwen35/src/prefill.rs | 26 +- pegainfer-qwen35/src/prefill_buffers.rs | 28 +- pegainfer-qwen35/src/recurrent_state.rs | 47 +-- pegainfer-qwen35/src/scheduler.rs | 6 +- pegainfer-qwen35/src/tp_executor.rs | 21 +- pegainfer-qwen35/src/unified_forward.rs | 8 +- pegainfer-qwen35/src/weights.rs | 340 ++++++++++++++++++--- 14 files changed, 725 insertions(+), 112 deletions(-) diff --git a/pegainfer-core/src/weight_loader.rs b/pegainfer-core/src/weight_loader.rs index 6c6a263c3..c2282d322 100644 --- a/pegainfer-core/src/weight_loader.rs +++ b/pegainfer-core/src/weight_loader.rs @@ -783,6 +783,127 @@ pub fn load_tensor_2d_col_shard( DeviceMatrix::from_host(ctx, &host, rows, cols) } +/// Load a 2D tensor assembled from multiple row ranges of one source tensor, +/// stitched in `ranges` order: each entry is (row_offset, rows). +pub fn load_tensor_2d_row_stitch( + ctx: &DeviceContext, + shards: &[SafeTensors], + weight_map: &HashMap, + name: &str, + ranges: &[(usize, usize)], +) -> Result { + let tensor = find_tensor(shards, weight_map, name)?; + let shape = tensor.shape(); + if shape.len() != 2 { + return Err(anyhow::anyhow!( + "Tensor '{}' expected 2D, got shape {:?}", + name, + shape + )); + } + let total_rows = shape[0]; + let cols = shape[1]; + let mut total = 0usize; + for &(row_offset, rows) in ranges { + if row_offset + rows > total_rows { + return Err(anyhow::anyhow!( + "2D row stitch out of bounds for '{}': row_offset={} rows={} total_rows={}", + name, + row_offset, + rows, + total_rows + )); + } + total += rows; + } + let elems = tensor_bf16_cow(&tensor, name)?; + let mut host = Vec::with_capacity(total * cols); + for &(row_offset, rows) in ranges { + let start = row_offset * cols; + host.extend_from_slice(&elems[start..start + rows * cols]); + } + DeviceMatrix::from_host(ctx, &host, total, cols) +} + +/// Load a 1D BF16 tensor assembled from multiple element ranges of one source +/// tensor, stitched in `ranges` order: each entry is (offset, len). +pub fn load_tensor_1d_stitch( + ctx: &DeviceContext, + shards: &[SafeTensors], + weight_map: &HashMap, + name: &str, + ranges: &[(usize, usize)], +) -> Result { + let tensor = find_tensor(shards, weight_map, name)?; + let elems = tensor_bf16_cow(&tensor, name)?; + let mut total = 0usize; + for &(offset, len) in ranges { + if offset + len > elems.len() { + return Err(anyhow::anyhow!( + "1D stitch out of bounds for '{}': offset={} len={} total_len={}", + name, + offset, + len, + elems.len() + )); + } + total += len; + } + let mut host = Vec::with_capacity(total); + for &(offset, len) in ranges { + host.extend_from_slice(&elems[offset..offset + len]); + } + DeviceVec::from_host(ctx, &host) +} + +/// Load a 1D BF16 element range to GPU (tensor-parallel shard of a 1D weight). +pub fn load_tensor_1d_shard( + ctx: &DeviceContext, + shards: &[SafeTensors], + weight_map: &HashMap, + name: &str, + offset: usize, + len: usize, +) -> Result { + load_tensor_1d_stitch(ctx, shards, weight_map, name, &[(offset, len)]) +} + +/// Load a 1D F32 element range to GPU (tensor-parallel shard of a 1D weight). +pub fn load_tensor_1d_f32_shard( + ctx: &DeviceContext, + shards: &[SafeTensors], + weight_map: &HashMap, + name: &str, + offset: usize, + len: usize, +) -> Result> { + let tensor = find_tensor(shards, weight_map, name)?; + let data = tensor.data(); + if data.len() % 4 != 0 { + return Err(anyhow::anyhow!( + "F32 tensor '{}': data length {} not multiple of 4", + name, + data.len() + )); + } + let total = data.len() / 4; + if offset + len > total { + return Err(anyhow::anyhow!( + "F32 1D shard out of bounds for '{}': offset={} len={} total_len={}", + name, + offset, + len, + total + )); + } + let slice = unsafe { std::slice::from_raw_parts(data.as_ptr().cast::(), total) }; + let gpu_data = ctx + .stream + .clone_htod(&slice[offset..offset + len]) + .map_err(|e| anyhow::anyhow!("H2D copy failed for '{}': {}", name, e))?; + Ok(gpu_data) +} + #[allow(clippy::cast_ptr_alignment)] /// Load a 1D F32 tensor to GPU as CudaSlice. /// For weights stored in float32 (e.g., A_log, norm.weight in linear attention). diff --git a/pegainfer-qwen35/src/batch_decode.rs b/pegainfer-qwen35/src/batch_decode.rs index 8d1c373c6..0c80feaae 100644 --- a/pegainfer-qwen35/src/batch_decode.rs +++ b/pegainfer-qwen35/src/batch_decode.rs @@ -179,6 +179,9 @@ impl Qwen35Model { bufs: &mut BatchDecodeBuffers35, ) -> Result<()> { let eps = self.config.rms_norm_eps; + let geom = self.geometry; + let num_attention_heads = geom.local_num_attention_heads(); + let num_key_value_heads = geom.local_num_key_value_heads(); ops::gemm_into(&self.ctx, &attn.q_proj, &bufs.normed, &mut bufs.q_full); ops::gemm_into(&self.ctx, &attn.k_proj, &bufs.normed, &mut bufs.k_attn); @@ -194,8 +197,8 @@ impl Qwen35Model { &self.cos_cache, &self.sin_cache, &bufs.positions_d, - self.config.num_attention_heads, - self.config.num_key_value_heads, + num_attention_heads, + num_key_value_heads, self.config.rotary_dim, eps, ); @@ -211,7 +214,7 @@ impl Qwen35Model { plan, &bufs.positions_d, &mut bufs.attn_out_full, - self.config.num_attention_heads, + num_attention_heads, bs, )?; @@ -221,7 +224,7 @@ impl Qwen35Model { crate::ffi::attention_gate_batch_hd256_cuda( qf_ptr as *const crate::ffi::Half, out_ptr as *mut crate::ffi::Half, - self.config.num_attention_heads as i32, + num_attention_heads as i32, bs as i32, self.ctx.stream.cu_stream(), ); @@ -268,6 +271,7 @@ impl Qwen35Model { linear_pointer_tables.validate_for(&self.config, bs, "Qwen3.5 eager decode")?; let mut positions = Vec::with_capacity(bs); + let mut start_positions = Vec::with_capacity(bs); for (i, kv) in kv_states.iter_mut().enumerate() { let pos = kv.seq_len(); self.ensure_rope_cache_covers(pos + 1)?; @@ -275,6 +279,7 @@ impl Qwen35Model { kv.advance(1); recurrent_states[i].seq_len += 1; positions.push(pos as i32); + start_positions.push(pos); } bufs.set_batch_size(bs); @@ -288,12 +293,50 @@ impl Qwen35Model { let kv_refs: Vec<&KvState> = kv_states.iter().map(|s| &**s).collect(); bufs.sync_paged_meta(&self.ctx, &kv_refs, bs)?; + // Route by TP-local GQA group supportability: when the rank-local + // q-per-kv group has no compiled batch-decode kernel, run full + // attention through the paged-prefill kernel with a per-step plan. + let prefill_attn_plan = if !self.geometry.local_decode_group_is_compiled() { + let page_indices: Vec> = + kv_states.iter().map(|kv| kv.page_indices_i32()).collect(); + let last_page_lens: Vec = + kv_states.iter().map(|kv| kv.last_page_len()).collect(); + let seq_lens = vec![1usize; bs]; + // cta_tile_q 0 = the kernel's own FA2 derivation; matches the + // hybrid fallback path in batch_decode_batched_hybrid. + Some( + ops::PrefillPagedPlan::from_raw_batch_with_cta_tile_q( + &self.ctx, + &page_indices, + &last_page_lens, + &start_positions, + &seq_lens, + self.geometry.local_num_attention_heads(), + self.geometry.local_num_key_value_heads(), + self.config.head_dim, + 0, + ) + .with_context(|| { + format!( + "eager decode build PrefillPagedPlan bs={bs}, pages={}, local heads={}/{}, head_dim={}", + page_indices.iter().map(Vec::len).sum::(), + self.geometry.local_num_attention_heads(), + self.geometry.local_num_key_value_heads(), + self.config.head_dim + ) + })?, + ) + } else { + None + }; + let kv_buffer = kv_states[0].buffer(); let layout = *kv_states[0].layout(); self.batch_decode_kernels_graph( kv_buffer, &layout, bs, + prefill_attn_plan.as_ref(), &linear_pointer_tables.state_ptrs, &linear_pointer_tables.conv_state_ptrs, bufs, @@ -394,6 +437,7 @@ impl Qwen35Model { kv_buffer, &layout, padded_bs, + None, linear_state_ptrs, linear_conv_state_ptrs, &mut graph_state.buffers, @@ -503,6 +547,7 @@ impl Qwen35Model { kv_buffer: &cudarc::driver::CudaSlice, layout: &KvLayout, padded_bs: usize, + prefill_attn_plan: Option<&ops::PrefillPagedPlan>, linear_state_ptrs: &[CudaSlice], linear_conv_state_ptrs: &[CudaSlice], bufs: &mut BatchDecodeBuffers35, @@ -529,9 +574,17 @@ impl Qwen35Model { match &layer.attn { LayerKind::FullAttention(attn) => { - self.batch_decode_full_attention( - attn, kv_buffer, layout, full_idx, padded_bs, bufs, - )?; + // The eager TP path passes a per-step prefill plan when the + // TP-local GQA group has no compiled batch-decode kernel; + // graph capture always passes None (rerouted earlier). + match prefill_attn_plan { + Some(plan) => self.batch_decode_full_attention_via_prefill( + attn, kv_buffer, layout, plan, full_idx, padded_bs, bufs, + )?, + None => self.batch_decode_full_attention( + attn, kv_buffer, layout, full_idx, padded_bs, bufs, + )?, + } full_idx += 1; } LayerKind::LinearAttention(attn) => { @@ -541,7 +594,7 @@ impl Qwen35Model { &linear_conv_state_ptrs[linear_idx], padded_bs, bufs, - ); + )?; linear_idx += 1; } } @@ -646,7 +699,7 @@ impl Qwen35Model { &linear_conv_state_ptrs[linear_idx], bs, bufs, - ); + )?; linear_idx += 1; } } @@ -719,6 +772,11 @@ impl Qwen35Model { /// Iterates 0..`padded_bs`. Real requests are in 0..real_bs; padding slots /// (real_bs..padded_bs) run but their output columns are ignored by the caller. /// All GPU addresses are stable per slot index, making this CUDA Graph safe. + /// + /// Phase 2b: every rank computes only its local value heads against + /// rank-local (never all-reduced) recurrent/conv state; the col-sharded + /// out_proj yields a partial hidden sum that is all-reduced under TP + /// (no-op at world_size 1). fn batch_decode_linear_attention_slots( &self, attn: &LinearAttentionLayer, @@ -726,7 +784,9 @@ impl Qwen35Model { conv_state_ptrs: &CudaSlice, padded_bs: usize, bufs: &mut BatchDecodeBuffers35, - ) { + ) -> Result<()> { + let geom = self.geometry; + ops::gemm_into(&self.ctx, &attn.in_proj_qkv, &bufs.normed, &mut bufs.qkv); ops::gemm_into(&self.ctx, &attn.in_proj_z, &bufs.normed, &mut bufs.z); ops::gemm_into(&self.ctx, &attn.in_proj_b, &bufs.normed, &mut bufs.b_proj); @@ -750,8 +810,8 @@ impl Qwen35Model { state_ptrs, &mut bufs.gdr_out, padded_bs, - self.config.linear_num_key_heads, - self.config.linear_num_value_heads, + geom.local_linear_num_key_heads(), + geom.local_linear_num_value_heads(), self.config.linear_key_head_dim, self.config.linear_value_head_dim, ); @@ -762,7 +822,7 @@ impl Qwen35Model { &attn.norm_weight, &bufs.z, &mut bufs.normed_gated, - self.config.linear_num_value_heads, + geom.local_linear_num_value_heads(), self.config.linear_value_head_dim, self.config.rms_norm_eps, ); @@ -772,5 +832,7 @@ impl Qwen35Model { &bufs.normed_gated, &mut bufs.attn_results, ); + self.all_reduce_hidden(&mut bufs.attn_results)?; + Ok(()) } } diff --git a/pegainfer-qwen35/src/batch_decode_graph.rs b/pegainfer-qwen35/src/batch_decode_graph.rs index eb88ba691..59d5225d3 100644 --- a/pegainfer-qwen35/src/batch_decode_graph.rs +++ b/pegainfer-qwen35/src/batch_decode_graph.rs @@ -79,7 +79,7 @@ impl BatchDecodeGraphState { let mut slot_states = Vec::with_capacity(max_batch); for _ in 0..max_batch { - slot_states.push(RecurrentState::new(ctx, config)?); + slot_states.push(RecurrentState::new(ctx, config, geometry)?); } let linear_pointer_tables = { let mut slot_refs: Vec<&mut RecurrentState> = slot_states.iter_mut().collect(); diff --git a/pegainfer-qwen35/src/config/model.rs b/pegainfer-qwen35/src/config/model.rs index 64c7fdecd..566c25459 100644 --- a/pegainfer-qwen35/src/config/model.rs +++ b/pegainfer-qwen35/src/config/model.rs @@ -141,14 +141,6 @@ impl Config35 { .contains(&(self.num_attention_heads / self.num_key_value_heads)) } - /// QKV projection output dimension for linear attention. - pub(crate) fn linear_attn_qkv_dim(&self) -> usize { - let q_dim = self.linear_num_key_heads * self.linear_key_head_dim; - let k_dim = q_dim; - let v_dim = self.linear_num_value_heads * self.linear_value_head_dim; - q_dim + k_dim + v_dim - } - /// Z projection output dimension for linear attention. pub(crate) fn linear_attn_z_dim(&self) -> usize { self.linear_num_value_heads * self.linear_value_head_dim diff --git a/pegainfer-qwen35/src/config/tp.rs b/pegainfer-qwen35/src/config/tp.rs index 9f2a97e1a..5506067ac 100644 --- a/pegainfer-qwen35/src/config/tp.rs +++ b/pegainfer-qwen35/src/config/tp.rs @@ -83,6 +83,11 @@ pub(crate) struct LocalGeometry { local_full_attn_q_dim: usize, local_full_attn_kv_dim: usize, local_full_attn_gated_q_dim: usize, + local_linear_num_key_heads: usize, + local_linear_num_value_heads: usize, + local_linear_q_dim: usize, + local_linear_v_dim: usize, + local_linear_qkv_dim: usize, } impl LocalGeometry { @@ -92,7 +97,7 @@ impl LocalGeometry { /// Fails on unsupported combinations before expensive loading: /// - sharded TP demands eager execution (`enable_cuda_graph` off); /// - every sharded model dimension must divide evenly by `world_size` - /// (linear-attention head counts are intentionally exempt); + /// (linear-attention heads included: Phase 2b shards them per rank); /// - `rank < world_size` and `world_size >= 1` are guaranteed by /// `TensorParallelConfig::try_from`. pub(crate) fn try_new( @@ -126,12 +131,37 @@ impl LocalGeometry { world_size: tp.world_size(), }); } + // Phase 2b shards linear attention/GDR heads per rank; fail closed on + // indivisible head counts rather than falling back to replication. + if !config.linear_num_key_heads.is_multiple_of(tp.world_size()) { + return Err(ConfigError::TpIndivisible { + field: "linear_num_key_heads", + value: config.linear_num_key_heads, + world_size: tp.world_size(), + }); + } + if !config + .linear_num_value_heads + .is_multiple_of(tp.world_size()) + { + return Err(ConfigError::TpIndivisible { + field: "linear_num_value_heads", + value: config.linear_num_value_heads, + world_size: tp.world_size(), + }); + } let local_num_attention_heads = config.num_attention_heads / tp.world_size(); let local_num_key_value_heads = config.num_key_value_heads / tp.world_size(); let local_intermediate_size = config.intermediate_size / tp.world_size(); let local_full_attn_q_dim = local_num_attention_heads * config.head_dim; let local_full_attn_kv_dim = local_num_key_value_heads * config.head_dim; + let local_linear_num_key_heads = config.linear_num_key_heads / tp.world_size(); + let local_linear_num_value_heads = config.linear_num_value_heads / tp.world_size(); + // Local q/k segment rows of the fused linear-attention qkv projection; + // q is keyed by key heads (one key head per value-head group). + let local_linear_q_dim = local_linear_num_key_heads * config.linear_key_head_dim; + let local_linear_v_dim = local_linear_num_value_heads * config.linear_value_head_dim; Ok(Self { tp, @@ -141,6 +171,12 @@ impl LocalGeometry { local_full_attn_q_dim, local_full_attn_kv_dim, local_full_attn_gated_q_dim: local_full_attn_q_dim * 2, + local_linear_num_key_heads, + local_linear_num_value_heads, + local_linear_q_dim, + local_linear_v_dim, + // [q_local | k_local | v_local] in storage order; k == q. + local_linear_qkv_dim: local_linear_q_dim * 2 + local_linear_v_dim, }) } @@ -184,6 +220,38 @@ impl LocalGeometry { pub(crate) fn local_full_attn_gated_q_dim(&self) -> usize { self.local_full_attn_gated_q_dim } + + // ── Linear-attention local dims (Phase 2b TP sharding) ──────────────── + // TP1 contract: at world_size 1 every local dim equals the global dim, so + // all linear-attention kernels/buffers/state keep their pre-TP shapes. + + pub(crate) fn local_linear_num_key_heads(&self) -> usize { + self.local_linear_num_key_heads + } + + pub(crate) fn local_linear_num_value_heads(&self) -> usize { + self.local_linear_num_value_heads + } + + /// Local fused qkv rows: [q_local | k_local | v_local] in storage order. + pub(crate) fn local_linear_qkv_dim(&self) -> usize { + self.local_linear_qkv_dim + } + + /// Local z projection output dimension (equals local v dim). + pub(crate) fn local_linear_z_dim(&self) -> usize { + self.local_linear_v_dim + } + + /// TP-local decode GQA group supportability for the eager decode path: + /// whether the FlashInfer batch-decode kernel supports the rank-local + /// q-per-kv group. Deterministic per model and identical on every rank, + /// so both arms are collective-safe (the reroute adds no collectives). + /// At world_size 1 this equals `Config35::decode_group_is_compiled`. + pub(crate) fn local_decode_group_is_compiled(&self) -> bool { + pegainfer_core::ops::SUPPORTED_GQA_GROUP_SIZES + .contains(&(self.local_num_attention_heads / self.local_num_key_value_heads)) + } } #[cfg(test)] @@ -308,11 +376,59 @@ mod tests { } #[test] - fn linear_attention_heads_need_not_divide_world_size() { - let mut cfg = config(); - cfg.linear_num_key_heads = 17; - cfg.linear_num_value_heads = 31; + fn requires_linear_attention_head_divisibility() { + let tp = TensorParallelConfig::try_from((1, 2)).unwrap(); + let mut broken = config(); + broken.linear_num_key_heads = 17; + let err = LocalGeometry::try_new(&broken, tp, false).unwrap_err(); + assert_eq!( + err, + ConfigError::TpIndivisible { + field: "linear_num_key_heads", + value: 17, + world_size: 2, + } + ); + + let mut broken = config(); + broken.linear_num_value_heads = 31; + let err = LocalGeometry::try_new(&broken, tp, false).unwrap_err(); + assert_eq!( + err, + ConfigError::TpIndivisible { + field: "linear_num_value_heads", + value: 31, + world_size: 2, + } + ); + } + + #[test] + fn computes_tp2_linear_attention_local_dimensions() { + let cfg = config(); let tp = TensorParallelConfig::try_from((1, 2)).unwrap(); - LocalGeometry::try_new(&cfg, tp, false).unwrap(); + let geom = LocalGeometry::try_new(&cfg, tp, false).unwrap(); + assert_eq!(geom.local_linear_num_key_heads(), 8); + assert_eq!(geom.local_linear_num_value_heads(), 16); + assert_eq!(geom.local_linear_q_dim, 1024); + assert_eq!(geom.local_linear_v_dim, 2048); + assert_eq!(geom.local_linear_qkv_dim(), 4096); + assert_eq!(geom.local_linear_z_dim(), 2048); + } + + #[test] + fn tp1_linear_attention_local_dimensions_equal_global() { + // TP1 invariant: every local dim equals the global dim, keeping TP1 + // numerics byte-identical to pre-sharding execution. + let cfg = config(); + let geom = LocalGeometry::try_new(&cfg, TensorParallelConfig::default(), false).unwrap(); + assert_eq!(geom.local_linear_num_key_heads(), 16); + assert_eq!(geom.local_linear_num_value_heads(), 32); + assert_eq!(geom.local_linear_q_dim, 2048); + assert_eq!(geom.local_linear_v_dim, 4096); + let global_qkv = + 2 * (cfg.linear_num_key_heads * cfg.linear_key_head_dim) + cfg.linear_attn_z_dim(); + assert_eq!(geom.local_linear_qkv_dim(), global_qkv); + assert_eq!(geom.local_linear_z_dim(), cfg.linear_attn_z_dim()); } } diff --git a/pegainfer-qwen35/src/decode_buffers.rs b/pegainfer-qwen35/src/decode_buffers.rs index 2275d2167..cea677445 100644 --- a/pegainfer-qwen35/src/decode_buffers.rs +++ b/pegainfer-qwen35/src/decode_buffers.rs @@ -75,9 +75,9 @@ impl BatchDecodeBuffers35 { let q_proj_dim = geometry.local_full_attn_gated_q_dim(); let q_dim = geometry.local_full_attn_q_dim(); let kv_dim = geometry.local_full_attn_kv_dim(); - let qkv_dim = config.linear_attn_qkv_dim(); - let z_dim = config.linear_attn_z_dim(); - let b_dim = config.linear_num_value_heads; + let qkv_dim = geometry.local_linear_qkv_dim(); + let z_dim = geometry.local_linear_z_dim(); + let b_dim = geometry.local_linear_num_value_heads(); let a_dim = b_dim; let intermediate = geometry.local_intermediate_size(); diff --git a/pegainfer-qwen35/src/executor.rs b/pegainfer-qwen35/src/executor.rs index 60e021330..956d0f9ab 100644 --- a/pegainfer-qwen35/src/executor.rs +++ b/pegainfer-qwen35/src/executor.rs @@ -167,7 +167,13 @@ impl Qwen35Executor { let mut recurrent_states: Vec = plan .requests .iter() - .map(|_| RecurrentState::new(self.model.device_ctx(), self.model.config())) + .map(|_| { + RecurrentState::new( + self.model.device_ctx(), + self.model.config(), + self.model.geometry, + ) + }) .collect::>()?; let mut recurrent_refs: Vec<&mut RecurrentState> = recurrent_states.iter_mut().collect(); let logits = diff --git a/pegainfer-qwen35/src/prefill.rs b/pegainfer-qwen35/src/prefill.rs index 6b65d9459..6644764b1 100644 --- a/pegainfer-qwen35/src/prefill.rs +++ b/pegainfer-qwen35/src/prefill.rs @@ -170,7 +170,8 @@ impl Qwen35Model { // Allocate the chunk scratch before advancing the KV state. It is the // largest, most allocation-prone buffer here, so failing first leaves // `kv_state` untouched and the request can be rejected cleanly. - let mut gdr_chunkwise_scratch = GdrChunkwiseScratch35::new(&self.ctx, c, seq_len)?; + let mut gdr_chunkwise_scratch = + GdrChunkwiseScratch35::new(&self.ctx, c, self.geometry, seq_len)?; // Advance paged KV state and build this chunk's prefill plan. kv_state.ensure_capacity(end_pos)?; @@ -240,7 +241,7 @@ impl Qwen35Model { let geom = self.geometry; let attn_out_dim = match &layer.attn { LayerKind::FullAttention(_) => geom.local_full_attn_q_dim(), - LayerKind::LinearAttention(_) => c.linear_attn_z_dim(), + LayerKind::LinearAttention(_) => geom.local_linear_z_dim(), }; // Batch project, then per-token attention/recurrent @@ -435,6 +436,10 @@ impl Qwen35Model { Ok(projected) } + /// Phase 2b: every rank prefills only its local value heads against + /// rank-local (never all-reduced) recurrent/conv state; the col-sharded + /// out_proj yields a partial hidden sum that is all-reduced under TP + /// (no-op at world_size 1). fn prefill_linear_attention( &self, attn: &LinearAttentionLayer, @@ -445,6 +450,7 @@ impl Qwen35Model { seq_len: usize, ) -> Result { let c = &self.config; + let geom = self.geometry; // Batch projections let qkv_batch = ops::gemm(&self.ctx, &attn.in_proj_qkv, normed_batch)?; @@ -452,8 +458,8 @@ impl Qwen35Model { let b_batch = ops::gemm(&self.ctx, &attn.in_proj_b, normed_batch)?; let a_batch = ops::gemm(&self.ctx, &attn.in_proj_a, normed_batch)?; - let qkv_dim = c.linear_attn_qkv_dim(); - let z_dim = c.linear_attn_z_dim(); + let qkv_dim = geom.local_linear_qkv_dim(); + let z_dim = geom.local_linear_z_dim(); let layer_state = &mut recurrent.layers[*linear_idx]; let mut qkv_conv_batch = HiddenStates::zeros(&self.ctx, qkv_dim, seq_len)?; @@ -477,8 +483,8 @@ impl Qwen35Model { &mut layer_state.state, gdr_chunkwise_scratch, &mut gdr_out_batch, - c.linear_num_key_heads, - c.linear_num_value_heads, + geom.local_linear_num_key_heads(), + geom.local_linear_num_value_heads(), c.linear_key_head_dim, c.linear_value_head_dim, )?; @@ -490,15 +496,17 @@ impl Qwen35Model { &attn.norm_weight, &z_batch, &mut normed_out_batch, - c.linear_num_value_heads, + geom.local_linear_num_value_heads(), c.linear_value_head_dim, c.rms_norm_eps, ); *linear_idx += 1; - // Output projection (batched) - ops::gemm(&self.ctx, &attn.out_proj, &normed_out_batch) + // Output projection (batched), then all-reduce the partial hidden sum. + let mut projected = ops::gemm(&self.ctx, &attn.out_proj, &normed_out_batch)?; + self.all_reduce_hidden(&mut projected)?; + Ok(projected) } fn batched_rms_norm_offset( diff --git a/pegainfer-qwen35/src/prefill_buffers.rs b/pegainfer-qwen35/src/prefill_buffers.rs index 92b38eb8a..027cf0e39 100644 --- a/pegainfer-qwen35/src/prefill_buffers.rs +++ b/pegainfer-qwen35/src/prefill_buffers.rs @@ -7,6 +7,7 @@ use pegainfer_core::tensor::DeviceContext; use pegainfer_core::tensor::HiddenStates; use super::config::Config35; +use super::config::LocalGeometry; /// Scratch buffers for a single Qwen3.5 linear-attention chunk-wise GDR prefill call. /// @@ -50,10 +51,17 @@ pub struct GdrChunkwiseScratch35 { impl GdrChunkwiseScratch35 { pub(crate) const CHUNK_SIZE: usize = 64; - pub(crate) fn new(ctx: &DeviceContext, config: &Config35, seq_len: usize) -> Result { + pub(crate) fn new( + ctx: &DeviceContext, + config: &Config35, + geometry: LocalGeometry, + seq_len: usize, + ) -> Result { + // GDR scratch is rank-local under TP: local value heads, global dims + // at world_size 1. Self::from_dims( ctx, - config.linear_num_value_heads, + geometry.local_linear_num_value_heads(), config.linear_key_head_dim, config.linear_value_head_dim, seq_len, @@ -120,8 +128,12 @@ impl GdrChunkwiseScratch35 { /// /// Direct-paged prefill writes full-attention K/V into the paged pool, so /// HND KVCache staging buffers are no longer part of the prefill scratch. - pub(crate) fn estimate_bytes(config: &Config35, max_seq_len: usize) -> usize { - let num_vh = config.linear_num_value_heads; + pub(crate) fn estimate_bytes( + config: &Config35, + geometry: LocalGeometry, + max_seq_len: usize, + ) -> usize { + let num_vh = geometry.local_linear_num_value_heads(); let key_dim = config.linear_key_head_dim; let val_dim = config.linear_value_head_dim; let chunk_sz = Self::CHUNK_SIZE; @@ -150,15 +162,15 @@ impl GdrChunkwiseScratch35 { // 2. Per-layer transient peak (all bf16 = 2 bytes). // Attention and MLP temps don't coexist — MLP runs after attention. let hidden_dim = config.hidden_size; - let intermediate = config.intermediate_size; + let intermediate = geometry.local_intermediate_size(); // Shared: hidden_batch + normed + hidden_plus_attn + normed_for_mlp let shared_layer = hidden_dim * seq * 4; // Full attention: q_full(with gate) + k + v + attn_out + q_prepped - let full_qkv = config.num_attention_heads * config.head_dim * 2; - let full_kv = config.num_key_value_heads * config.head_dim; - let full_out = config.num_attention_heads * config.head_dim; + let full_qkv = geometry.local_full_attn_gated_q_dim(); + let full_kv = geometry.local_full_attn_kv_dim(); + let full_out = geometry.local_full_attn_q_dim(); let full_attn_temps = (full_qkv + full_kv * 2 + full_out * 2) * seq; // MLP: gate_up_out + act_out (same peak footprint as separate gate/up) diff --git a/pegainfer-qwen35/src/recurrent_state.rs b/pegainfer-qwen35/src/recurrent_state.rs index f7279ec12..3d868332a 100644 --- a/pegainfer-qwen35/src/recurrent_state.rs +++ b/pegainfer-qwen35/src/recurrent_state.rs @@ -1,8 +1,11 @@ //! Recurrent state for Qwen3.5 linear attention layers. //! //! Each linear attention layer maintains: -//! - Recurrent state: [num_value_heads, key_head_dim, value_head_dim] f32, V contiguous ([H,K,V]) -//! - Conv state: [qkv_dim × (conv_kernel_dim - 1)] bf16 +//! - Recurrent state: [local_value_heads, key_head_dim, value_head_dim] f32, V contiguous ([H,K,V]) +//! - Conv state: [local_qkv_dim × (conv_kernel_dim - 1)] bf16 +//! +//! Under TP both are rank-local: Phase 2b shards value heads (and fused qkv +//! channels) across ranks. Recurrent and conv state are NEVER all-reduced. use anyhow::Result; use cudarc::driver::CudaSlice; @@ -11,14 +14,16 @@ use pegainfer_core::tensor::DeviceContext; use pegainfer_core::tensor::DeviceVec; use super::config::Config35; +use super::config::LocalGeometry; /// Per-layer recurrent state for a single linear attention layer. pub(crate) struct LayerRecurrentState { - /// Recurrent state matrix: [num_value_heads * key_head_dim * value_head_dim] f32 - /// Stored as f32 per mamba_ssm_dtype="float32" in config. + /// Recurrent state matrix: [local_value_heads * key_head_dim * value_head_dim] f32 + /// Stored as f32 per mamba_ssm_dtype="float32" in config. Rank-local under + /// TP (local dims equal global dims at world_size 1). Never all-reduced. pub(crate) state: CudaSlice, - /// Conv1d state buffer: [qkv_dim * (conv_kernel_dim - 1)] bf16 - /// Stores the last (kernel_dim - 1) inputs for causal conv1d. + /// Conv1d state buffer: [local_linear_qkv_dim * (conv_kernel_dim - 1)] bf16 + /// Stores the last (kernel_dim - 1) inputs for causal conv1d. Rank-local. pub(crate) conv_state: DeviceVec, } @@ -42,19 +47,25 @@ pub(crate) struct LinearStatePointerTables { } /// Per-layer element counts shared by allocation and reservation: -/// (linear layers, f32 state elements, bf16 conv elements). -fn per_layer_dims(config: &Config35) -> (usize, usize, usize) { +/// (linear layers, f32 state elements, bf16 conv elements). Sizes are +/// rank-local under TP; at world_size 1 they equal the global dims. +fn per_layer_dims(config: &Config35, geometry: LocalGeometry) -> (usize, usize, usize) { let num_linear_layers = config.num_hidden_layers - config.num_full_attention_layers(); - let state_size = - config.linear_num_value_heads * config.linear_key_head_dim * config.linear_value_head_dim; - let conv_state_size = config.linear_attn_qkv_dim() * (config.linear_conv_kernel_dim - 1); + let state_size = geometry.local_linear_num_value_heads() + * config.linear_key_head_dim + * config.linear_value_head_dim; + let conv_state_size = geometry.local_linear_qkv_dim() * (config.linear_conv_kernel_dim - 1); (num_linear_layers, state_size, conv_state_size) } impl RecurrentState { /// Allocate zeroed recurrent state for all linear attention layers. - pub(crate) fn new(ctx: &DeviceContext, config: &Config35) -> Result { - let (num_linear_layers, state_size, conv_state_size) = per_layer_dims(config); + pub(crate) fn new( + ctx: &DeviceContext, + config: &Config35, + geometry: LocalGeometry, + ) -> Result { + let (num_linear_layers, state_size, conv_state_size) = per_layer_dims(config, geometry); let mut layers = Vec::with_capacity(num_linear_layers); for _ in 0..num_linear_layers { @@ -144,17 +155,17 @@ impl LinearStatePointerTables { } } -/// Device bytes of one request's recurrent state. -pub(crate) fn bytes_per_request(config: &Config35) -> usize { - let (num_linear_layers, state_size, conv_state_size) = per_layer_dims(config); +/// Device bytes of one request's (rank-local) recurrent state. +pub(crate) fn bytes_per_request(config: &Config35, geometry: LocalGeometry) -> usize { + let (num_linear_layers, state_size, conv_state_size) = per_layer_dims(config, geometry); num_linear_layers * (state_size * std::mem::size_of::() + conv_state_size * std::mem::size_of::()) } impl RecurrentState { - pub(crate) fn allocation_bytes(config: &Config35) -> usize { - bytes_per_request(config) + pub(crate) fn allocation_bytes(config: &Config35, geometry: LocalGeometry) -> usize { + bytes_per_request(config, geometry) } } diff --git a/pegainfer-qwen35/src/scheduler.rs b/pegainfer-qwen35/src/scheduler.rs index c2552c23c..bf4259c47 100644 --- a/pegainfer-qwen35/src/scheduler.rs +++ b/pegainfer-qwen35/src/scheduler.rs @@ -592,7 +592,11 @@ impl SingleGpuBackend { } fn alloc_recurrent(&self) -> Result { - RecurrentState::new(self.model.device_ctx(), self.model.config()) + RecurrentState::new( + self.model.device_ctx(), + self.model.config(), + self.model.geometry, + ) } fn batch_prefill_logits(&self, chunk: &mut ScheduledChunk) -> Result { diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index e35a05d36..830b3c231 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -1,7 +1,7 @@ //! Tensor-parallel worker runtime for Qwen3.5. //! -//! Phase 2A adds one canonical eager unified command while retaining the -//! replicated linear-attention state layout from Phase 1. +//! Phase 2A adds one canonical eager unified command. Phase 2b shards the +//! linear-attention/GDR weight and state surface per rank. use std::collections::HashSet; use std::panic::AssertUnwindSafe; @@ -1037,10 +1037,15 @@ impl TpWorkerPrepared { .ctx .mem_get_info() .map_err(|err| anyhow::anyhow!("failed to query TP rank {rank} memory: {err}"))?; - let recurrent_bytes = RecurrentState::allocation_bytes(model.config()); + // Phase 2b: per-request 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(), prefill_scratch_tokens); + let prefill_scratch_bytes = GdrChunkwiseScratch35::estimate_bytes( + model.config(), + model.geometry, + prefill_scratch_tokens, + ); let max_batch = effective_recurrent_capacity( requested_max_batch, free_bytes, @@ -1480,7 +1485,11 @@ impl TpWorkerState { if let Some(idx) = self.request_index(request_id) { return Ok(idx); } - let mut recurrent = RecurrentState::new(self.model.device_ctx(), self.model.config())?; + let mut recurrent = RecurrentState::new( + self.model.device_ctx(), + self.model.config(), + self.model.geometry, + )?; let linear_pointer_tables = { let mut recurrent_refs = [&mut recurrent]; LinearStatePointerTables::from_recurrent_refs( diff --git a/pegainfer-qwen35/src/unified_forward.rs b/pegainfer-qwen35/src/unified_forward.rs index 6526225bc..32f216672 100644 --- a/pegainfer-qwen35/src/unified_forward.rs +++ b/pegainfer-qwen35/src/unified_forward.rs @@ -174,8 +174,8 @@ mod tests { let prompts_ref: Vec<&[u32]> = vec![&prompt_a, &prompt_b]; let mut kv_states: Vec = vec![model.alloc_kv(), model.alloc_kv()]; let mut rec_states: Vec = vec![ - RecurrentState::new(&model.ctx, &model.config).unwrap(), - RecurrentState::new(&model.ctx, &model.config).unwrap(), + RecurrentState::new(&model.ctx, &model.config, model.geometry).unwrap(), + RecurrentState::new(&model.ctx, &model.config, model.geometry).unwrap(), ]; let mut rec_refs: Vec<&mut RecurrentState> = rec_states.iter_mut().collect(); let first_logits = model @@ -213,8 +213,8 @@ mod tests { let prompts_ref: Vec<&[u32]> = vec![&prompt_a, &prompt_b]; let mut kv_states: Vec = vec![model.alloc_kv(), model.alloc_kv()]; let mut rec_states: Vec = vec![ - RecurrentState::new(&model.ctx, &model.config).unwrap(), - RecurrentState::new(&model.ctx, &model.config).unwrap(), + RecurrentState::new(&model.ctx, &model.config, model.geometry).unwrap(), + RecurrentState::new(&model.ctx, &model.config, model.geometry).unwrap(), ]; let mut rec_refs: Vec<&mut RecurrentState> = rec_states.iter_mut().collect(); diff --git a/pegainfer-qwen35/src/weights.rs b/pegainfer-qwen35/src/weights.rs index 70cb6bb80..c4a1d5d54 100644 --- a/pegainfer-qwen35/src/weights.rs +++ b/pegainfer-qwen35/src/weights.rs @@ -18,9 +18,13 @@ use pegainfer_core::weight_loader::deserialize_shards; use pegainfer_core::weight_loader::load_shard_info_fixed; use pegainfer_core::weight_loader::load_tensor_1d; use pegainfer_core::weight_loader::load_tensor_1d_f32; +use pegainfer_core::weight_loader::load_tensor_1d_f32_shard; +use pegainfer_core::weight_loader::load_tensor_1d_shard; +use pegainfer_core::weight_loader::load_tensor_1d_stitch; use pegainfer_core::weight_loader::load_tensor_2d; use pegainfer_core::weight_loader::load_tensor_2d_col_shard; use pegainfer_core::weight_loader::load_tensor_2d_row_shard; +use pegainfer_core::weight_loader::load_tensor_2d_row_stitch; use pegainfer_core::weight_loader::mmap_shards; use safetensors::SafeTensors; @@ -46,23 +50,28 @@ pub(super) struct FullAttentionLayer { /// Linear attention layer weights (24 layers in Qwen3.5-4B). pub(super) struct LinearAttentionLayer { - /// Fused QKV projection: [q_dim + k_dim + v_dim, hidden_size] + /// Fused QKV projection: [local_linear_qkv_dim, hidden_size] — rows keep + /// the global [q | k | v] segment layout, with each segment restricted to + /// this rank's head-local slice (see `linear_qkv_shard_segments`). pub(super) in_proj_qkv: DeviceMatrix, - /// Z projection (for output gating): [z_dim, hidden_size] + /// Z projection (for output gating): [local_linear_z_dim, hidden_size] pub(super) in_proj_z: DeviceMatrix, - /// Beta projection: [num_value_heads, hidden_size] + /// Beta projection: [local_linear_num_value_heads, hidden_size] pub(super) in_proj_b: DeviceMatrix, - /// Alpha projection: [num_value_heads, hidden_size] + /// Alpha projection: [local_linear_num_value_heads, hidden_size] pub(super) in_proj_a: DeviceMatrix, - /// Depthwise conv1d weight: [qkv_dim * conv_kernel_dim] (flattened from [qkv_dim, 1, 4]) + /// Depthwise conv1d weight: [local_linear_qkv_dim * conv_kernel_dim] + /// (flattened from [qkv_dim, 1, 4]); channel layout mirrors in_proj_qkv. pub(super) conv1d_weight: DeviceVec, - /// dt_bias: [num_value_heads] bf16 + /// dt_bias: [local_linear_num_value_heads] bf16 pub(super) dt_bias: DeviceVec, - /// A_log: [num_value_heads] f32 + /// A_log: [local_linear_num_value_heads] f32 pub(super) a_log: CudaSlice, - /// RMSNorm weight for output normalization: [value_head_dim] f32 + /// RMSNorm weight for output normalization: [value_head_dim] f32 — + /// head-shared, so replicated on every rank. pub(super) norm_weight: CudaSlice, - /// Output projection: [hidden_size, z_dim] + /// Output projection: [hidden_size, local_linear_z_dim] (row-parallel; + /// the layer all-reduces the partial hidden sum under TP). pub(super) out_proj: DeviceMatrix, } @@ -338,60 +347,110 @@ impl Qwen35Model { } LayerType::LinearAttention => { let attn_prefix = format!("{}.linear_attn", prefix); + // Phase 2b: shard linear attention over TP ranks. Value-head + // unit drives z/b/a/dt_bias/a_log rows; the fused qkv weight + // and conv need per-segment head-local stitching. + let (vh_offset, vh_rows) = + tensor_parallel.shard_range(config.linear_num_value_heads); + let (z_row_offset, z_rows) = + tensor_parallel.shard_range(config.linear_attn_z_dim()); + let (z_col_offset, z_cols) = (z_row_offset, z_rows); LayerKind::LinearAttention(LinearAttentionLayer { - in_proj_qkv: load_tensor_2d( + in_proj_qkv: load_linear_in_proj_qkv_shard( &ctx, &shards, &weight_map, &format!("{}.in_proj_qkv.weight", attn_prefix), + &config, + tensor_parallel, )?, - in_proj_z: load_tensor_2d( + in_proj_z: load_tensor_2d_row_shard_if_needed( &ctx, &shards, &weight_map, &format!("{}.in_proj_z.weight", attn_prefix), + geometry, + z_row_offset, + z_rows, )?, - in_proj_b: load_tensor_2d( + in_proj_b: load_tensor_2d_row_shard_if_needed( &ctx, &shards, &weight_map, &format!("{}.in_proj_b.weight", attn_prefix), + geometry, + vh_offset, + vh_rows, )?, - in_proj_a: load_tensor_2d( + in_proj_a: load_tensor_2d_row_shard_if_needed( &ctx, &shards, &weight_map, &format!("{}.in_proj_a.weight", attn_prefix), + geometry, + vh_offset, + vh_rows, )?, - conv1d_weight: load_tensor_1d( + conv1d_weight: load_linear_conv1d_shard( &ctx, &shards, &weight_map, &format!("{}.conv1d.weight", attn_prefix), + &config, + tensor_parallel, )?, - dt_bias: load_tensor_1d( - &ctx, - &shards, - &weight_map, - &format!("{}.dt_bias", attn_prefix), - )?, - a_log: load_tensor_1d_f32( - &ctx, - &shards, - &weight_map, - &format!("{}.A_log", attn_prefix), - )?, + dt_bias: if tensor_parallel.is_sharded() { + load_tensor_1d_shard( + &ctx, + &shards, + &weight_map, + &format!("{}.dt_bias", attn_prefix), + vh_offset, + vh_rows, + )? + } else { + load_tensor_1d( + &ctx, + &shards, + &weight_map, + &format!("{}.dt_bias", attn_prefix), + )? + }, + a_log: if tensor_parallel.is_sharded() { + load_tensor_1d_f32_shard( + &ctx, + &shards, + &weight_map, + &format!("{}.A_log", attn_prefix), + vh_offset, + vh_rows, + )? + } else { + load_tensor_1d_f32( + &ctx, + &shards, + &weight_map, + &format!("{}.A_log", attn_prefix), + )? + }, + // Gated RMSNorm weight is per value-head dim (128) and + // shared by every head: replicated, never sharded. norm_weight: load_tensor_1d_f32( &ctx, &shards, &weight_map, &format!("{}.norm.weight", attn_prefix), )?, - out_proj: load_tensor_2d( + // Row-parallel out_proj: shard input columns to the + // local z dim; the layer all-reduces the partial sum. + out_proj: load_tensor_2d_col_shard_if_needed( &ctx, &shards, &weight_map, &format!("{}.out_proj.weight", attn_prefix), + geometry, + z_col_offset, + z_cols, )?, }) } @@ -499,10 +558,14 @@ impl Qwen35Model { // Reserve space for prefill scratch (GDR chunkwise + per-layer transients) // before allocating KV pool, so prefill doesn't OOM. let max_prefill_len = super::prefill::SCRATCH_ESTIMATE_SEQ; - let scratch_reserve = - super::prefill_buffers::GdrChunkwiseScratch35::estimate_bytes(&config, max_prefill_len); - let recurrent_reserve = - STATES_PER_DECODE_SLOT * max_batch * super::recurrent_state::bytes_per_request(&config); + let scratch_reserve = super::prefill_buffers::GdrChunkwiseScratch35::estimate_bytes( + &config, + geometry, + max_prefill_len, + ); + let recurrent_reserve = STATES_PER_DECODE_SLOT + * max_batch + * super::recurrent_state::bytes_per_request(&config, geometry); let min_kv_bytes = MIN_KV_PAGES * bytes_per_page; anyhow::ensure!( free_bytes >= scratch_reserve + recurrent_reserve + min_kv_bytes, @@ -610,9 +673,9 @@ impl Qwen35Model { let geom = self.geometry; let full_q = geom.local_full_attn_gated_q_dim(); let full_kv = geom.local_full_attn_kv_dim(); - let linear_qkv = self.config.linear_attn_qkv_dim(); - let linear_z = self.config.linear_attn_z_dim(); - let linear_ba = self.config.linear_num_value_heads; + let linear_qkv = geom.local_linear_qkv_dim(); + let linear_z = geom.local_linear_z_dim(); + let linear_ba = geom.local_linear_num_value_heads(); let intermediate = geom.local_intermediate_size(); let full_q_samples: Vec<_> = self @@ -802,6 +865,70 @@ fn load_full_attention_gated_q_proj( load_tensor_2d_row_shard(ctx, shards, weight_map, name, range.row_offset, range.rows) } +/// Row ranges this rank owns inside the fused global linear-attention qkv +/// projection. The checkpoint stores [all q rows | all k rows | all v rows]; +/// each segment contributes its head-local slice so the rank's stitched rows +/// stay [q_local | k_local | v_local]. Never reblock across segments — q rows +/// key on key heads, v rows on value heads (the gated-q lesson). +fn linear_qkv_shard_segments( + config: &Config35, + tensor_parallel: TensorParallelConfig, +) -> [(usize, usize); 3] { + let global_q = config.linear_num_key_heads * config.linear_key_head_dim; + let global_k = global_q; + let global_v = config.linear_attn_z_dim(); + let (q_rel, q_rows) = tensor_parallel.shard_range(global_q); + let (k_rel, k_rows) = tensor_parallel.shard_range(global_k); + let (v_rel, v_rows) = tensor_parallel.shard_range(global_v); + [ + (q_rel, q_rows), + (global_q + k_rel, k_rows), + (global_q + global_k + v_rel, v_rows), + ] +} + +fn load_linear_in_proj_qkv_shard( + ctx: &DeviceContext, + shards: &[SafeTensors], + weight_map: &HashMap, + name: &str, + config: &Config35, + tensor_parallel: TensorParallelConfig, +) -> Result { + if !tensor_parallel.is_sharded() { + return load_tensor_2d(ctx, shards, weight_map, name); + } + let segments = linear_qkv_shard_segments(config, tensor_parallel); + load_tensor_2d_row_stitch(ctx, shards, weight_map, name, &segments) +} + +/// The flattened conv1d weight keeps each channel's kernel taps contiguous +/// ([channel, 1, kernel_dim]); its channel layout mirrors the fused qkv rows, +/// so shard it with the same per-segment ranges scaled by the kernel dim. +fn linear_conv1d_shard_segments( + config: &Config35, + tensor_parallel: TensorParallelConfig, +) -> [(usize, usize); 3] { + let kernel_dim = config.linear_conv_kernel_dim; + linear_qkv_shard_segments(config, tensor_parallel) + .map(|(offset, len)| (offset * kernel_dim, len * kernel_dim)) +} + +fn load_linear_conv1d_shard( + ctx: &DeviceContext, + shards: &[SafeTensors], + weight_map: &HashMap, + name: &str, + config: &Config35, + tensor_parallel: TensorParallelConfig, +) -> Result { + if !tensor_parallel.is_sharded() { + return load_tensor_1d(ctx, shards, weight_map, name); + } + let segments = linear_conv1d_shard_segments(config, tensor_parallel); + load_tensor_1d_stitch(ctx, shards, weight_map, name, &segments) +} + fn load_tensor_2d_row_shard_if_needed( ctx: &DeviceContext, shards: &[SafeTensors], @@ -912,4 +1039,149 @@ mod tests { assert_eq!(local_gate_up_rows, 9216); assert_eq!(local_down_cols, 4608); } + + #[test] + fn linear_qkv_shard_segments_stitch_head_local_slices() { + // test_config: k heads 16, v heads 32, head dim 128 → q=k=2048, v=4096. + let config = test_config(); + let rank0 = + linear_qkv_shard_segments(&config, TensorParallelConfig::try_from((0, 2)).unwrap()); + assert_eq!(rank0, [(0, 1024), (2048, 1024), (4096, 2048)]); + + let rank1 = + linear_qkv_shard_segments(&config, TensorParallelConfig::try_from((1, 2)).unwrap()); + assert_eq!(rank1, [(1024, 1024), (3072, 1024), (6144, 2048)]); + + // Every rank's stitched rows tile [0, qkv) with no overlap: each + // segment's local slices across ranks are contiguous and complete. + for (r0, r1) in rank0.iter().zip(rank1.iter()) { + assert_eq!(r0.1, r1.1); + assert_eq!(r1.0, r0.0 + r0.1); + } + } + + #[test] + fn linear_conv1d_shard_segments_scale_by_kernel_dim() { + let config = test_config(); + let rank1 = + linear_conv1d_shard_segments(&config, TensorParallelConfig::try_from((1, 2)).unwrap()); + // conv1d.weight is [qkv * 4]: same ranges as qkv, scaled by 4. + assert_eq!(rank1, [(4096, 4096), (12288, 4096), (24576, 8192)]); + } + + #[test] + fn tp1_linear_qkv_shard_segments_cover_full_segments() { + // TP1 identity: the segments used by the sharded loader would reproduce + // the full tensor (TP1 itself never stitches — it uses load_tensor_2d). + let config = test_config(); + let segments = linear_qkv_shard_segments(&config, TensorParallelConfig::default()); + assert_eq!(segments, [(0, 2048), (2048, 2048), (4096, 4096)]); + } + + /// In-memory safetensors fixture: one F32 tensor whose element `i` + /// carries value `i` (exact in f32), so any slice maps back to its source + /// offset. The row/col range math is dtype-agnostic (element units), so + /// an f32 blob exercises the same layout contract as the bf16 loaders. + fn safetensors_fixture_f32(name: &str, shape: &[usize]) -> Vec { + let len: usize = shape.iter().product(); + let data: Vec = (0..len).flat_map(|i| (i as f32).to_le_bytes()).collect(); + let view = + safetensors::tensor::TensorView::new(safetensors::Dtype::F32, shape.to_vec(), &data) + .unwrap(); + safetensors::serialize([(name.to_string(), view)], None).unwrap() + } + + #[test] + fn linear_attention_tp2_slices_match_synthetic_checkpoint() { + // CPU-only layout contract test: parse a synthetic safetensors blob and + // verify each rank's stitched slices land on the expected source + // offsets, including per-segment head-local contiguity. + let config = test_config(); + let qkv_rows = test_geometry(0, 1).local_linear_qkv_dim(); // 8192 + let kernel_dim = config.linear_conv_kernel_dim; // 4 + // Column count is irrelevant to the row range math; keep it tiny. + let cols = 8; + + let qkv_blob = safetensors_fixture_f32("w", &[qkv_rows, cols]); + let qkv = safetensors::SafeTensors::deserialize(&qkv_blob).unwrap(); + let qkv_view = qkv.tensor("w").unwrap(); + assert_eq!(qkv_view.shape(), [qkv_rows, cols]); + let qkv_elems: Vec = qkv_view + .data() + .chunks_exact(4) + .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])) + .collect(); + + // conv1d.weight fixture: [qkv * kernel_dim] flattened channels. + let conv_blob = safetensors_fixture_f32("c", &[qkv_rows * kernel_dim]); + let conv = safetensors::SafeTensors::deserialize(&conv_blob).unwrap(); + let conv_view = conv.tensor("c").unwrap(); + let conv_elems: Vec = conv_view + .data() + .chunks_exact(4) + .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])) + .collect(); + + let global_q = config.linear_num_key_heads * config.linear_key_head_dim; + for rank in 0..2usize { + let tp = TensorParallelConfig::try_from((rank, 2)).unwrap(); + let geom = LocalGeometry::try_new(&config, tp, false).unwrap(); + let segments = linear_qkv_shard_segments(&config, tp); + + // Stitched matrix = per-segment head-local row slices, in storage + // order; assert full row content of every stitched row. + let mut stitched = Vec::new(); + for &(offset, rows) in &segments { + stitched.extend_from_slice(&qkv_elems[offset * cols..(offset + rows) * cols]); + } + assert_eq!(stitched.len(), geom.local_linear_qkv_dim() * cols); + let mut expected = Vec::new(); + for row in (0..segments[0].1) // q: rank-local key-head slice + .map(|r| segments[0].0 + r) + .chain((0..segments[1].1).map(|r| segments[1].0 + r)) + .chain((0..segments[2].1).map(|r| segments[2].0 + r)) + { + expected.extend_from_slice(&qkv_elems[row * cols..(row + 1) * cols]); + } + assert_eq!(stitched, expected, "rank {rank} fused qkv layout"); + + // Head-locality: q segment stays inside the rank's key-head range, + // v segment starts after all global q+k rows plus the rank offset. + // Local segment dims derive from the fixture config (TP2). + let lq = config.linear_num_key_heads / 2 * config.linear_key_head_dim; + let lk = lq; + let lv = config.linear_num_value_heads / 2 * config.linear_value_head_dim; + assert_eq!(segments[0].0, rank * lq); + assert_eq!(segments[2].0, 2 * global_q + rank * lv); + + // conv1d channels mirror qkv rows, each scaled by kernel_dim. + // The expectation is built from first principles (rank-local + // head-dim channel windows), not from the segment tuples. + let conv_segments = linear_conv1d_shard_segments(&config, tp); + let mut conv_stitched = Vec::new(); + for &(conv_off, conv_len) in &conv_segments { + conv_stitched.extend_from_slice(&conv_elems[conv_off..conv_off + conv_len]); + } + let channels = (rank * lq..(rank + 1) * lq) + .chain(global_q + rank * lk..global_q + (rank + 1) * lk) + .chain(2 * global_q + rank * lv..2 * global_q + (rank + 1) * lv); + let mut conv_expected = Vec::new(); + for c in channels { + conv_expected.extend_from_slice(&conv_elems[c * kernel_dim..(c + 1) * kernel_dim]); + } + assert_eq!( + conv_stitched.len(), + geom.local_linear_qkv_dim() * kernel_dim + ); + assert_eq!(conv_stitched, conv_expected, "rank {rank} conv1d layout"); + + // Value-head unit drives in_proj_z rows, in_proj_b/a rows, dt_bias + // and a_log; out_proj takes the same range as columns. + let (vh_offset, vh_rows) = tp.shard_range(config.linear_num_value_heads); + assert_eq!((vh_offset, vh_rows), (rank * 16, 16)); + let (z_offset, z_rows) = tp.shard_range(config.linear_attn_z_dim()); + assert_eq!((z_offset, z_rows), (rank * 2048, 2048)); + assert_eq!(z_rows, geom.local_linear_z_dim()); + } + } } From 1a6a86214725a138ccbbb631fa35a9778fd375ec Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Thu, 20 Aug 2026 15:07:11 +0000 Subject: [PATCH 2/4] perf(qwen35): batch eager decode rows under TP Port of the batched eager TP decode step onto #870's worker structure: decode rows in one command now run as ONE batched forward per step on every rank plus one batched rank-0 sampling pass, instead of a per-request bs=1 loop. - run_decode_batch resolves every decode row's worker state in command order, builds a step-scoped LinearStatePointerTables over the whole batch (from_recurrent_refs(..., bs, ...)), runs one batch_decode_eager_logits forward, then rank 0 snapshots all requested logprob rows before one batched select_batch over per-row params - execute_decode_rows (used by both decode-only and unified steps) calls run_decode_batch once; per-row results fan out in command order - TpRequestState.linear_pointer_tables (capacity-1, decode-only) removed; ensure_prefill_state no longer builds it. The step-scoped table is rebuilt every step, so swap_remove retirement can't stale it Seeded rows keep per-row semantics: select_batch isolates each seeded row into its own single-row philox call keyed on (request seed, step 0), so seeded output stays independent of batch composition. Unseeded rows decorrelate via the per-step command seed, same as the single-GPU batched path. Reference (27B TP2, 2x RTX 4090, eager): 16 concurrent 256-token completions aggregate 24.9 -> 292.3 tok/s; single-request unchanged. Signed-off-by: Ziyang Zhang --- pegainfer-qwen35/src/tp_executor.rs | 200 ++++++++++++++++++---------- 1 file changed, 132 insertions(+), 68 deletions(-) diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index 830b3c231..885a5d844 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -1,7 +1,9 @@ //! Tensor-parallel worker runtime for Qwen3.5. //! //! Phase 2A adds one canonical eager unified command. Phase 2b shards the -//! linear-attention/GDR weight and state surface per rank. +//! linear-attention/GDR weight and state surface per rank, and decode rows +//! run as one batched forward per step plus one batched rank-0 sampling pass +//! instead of a per-request bs=1 loop. use std::collections::HashSet; use std::panic::AssertUnwindSafe; @@ -1006,7 +1008,6 @@ struct TpRequestState { phase: TpRequestPhase, kv: KvState, recurrent: RecurrentState, - linear_pointer_tables: LinearStatePointerTables, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -1341,6 +1342,133 @@ impl TpWorkerState { 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. + /// + /// Seeded rows keep their former per-row semantics: they pass step 0 and + /// `select_batch` isolates each seeded row into its own single-row philox + /// call keyed on (request seed, step), so seeded output stays independent + /// of batch composition. Unseeded rows decorrelate inside the batched + /// call through the per-step command seed, same as the single-GPU batched + /// path. 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()); + } + + // 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_slots: Vec> = + std::iter::repeat_with(|| None).take(bs).collect(); + let mut recurrent_slots: Vec> = + std::iter::repeat_with(|| None).take(bs).collect(); + for (state_idx, state) in self.requests.iter_mut().enumerate() { + if let Some(row) = row_of_state[state_idx] { + kv_slots[row] = Some(&mut state.kv); + recurrent_slots[row] = Some(&mut state.recurrent); + } + } + let mut kv_refs: Vec<&mut KvState> = Vec::with_capacity(bs); + let mut recurrent_refs: Vec<&mut RecurrentState> = Vec::with_capacity(bs); + for (kv, recurrent) in kv_slots.into_iter().zip(recurrent_slots) { + kv_refs.push(kv.expect("decode row state resolved above")); + recurrent_refs.push(recurrent.expect("decode row state resolved above")); + } + + // Step-scoped GDR pointer tables over the full decode batch. They are + // rebuilt every step, so swap_remove retirement between steps can + // never leave a stale table behind. + let linear_pointer_tables = LinearStatePointerTables::from_recurrent_refs( + self.model.device_ctx(), + self.model.config(), + &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, + &linear_pointer_tables, + &mut self.decode_buffers, + )?; + + if self.rank != 0 { + return Ok(Vec::new()); + } + + // Snapshot requested logits rows BEFORE sampling: the sampler may + // modify bufs.logits in place. + let requested_logprobs: Vec = + requests.iter().map(|request| request.logprobs).collect(); + let cpu_logits = snapshot_requested_logprobs( + self.model.device_ctx(), + &self.decode_buffers.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( + self.model.device_ctx(), + &self.decode_buffers.logits, + ¶ms_refs, + &steps, + sample_seed, + &mut self.sample_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| { + pegainfer_sample::token_logprob_from_row( + logits_row, + tokens[row], + request.logprobs, + ) + }); + DecodeRequestResult { + request_id: request.request_id, + token: tokens[row], + logprob, + } + }) + .collect()) + } + fn sample_final_prefill_chunk( &mut self, chunk: &TpPrefillChunkItem, @@ -1399,60 +1527,7 @@ impl TpWorkerState { self.max_batch ); - let mut primary_results = - Vec::with_capacity(if self.rank == 0 { requests.len() } else { 0 }); - for (row_idx, 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() - ); - - { - let state = &mut self.requests[state_idx]; - let mut kv_refs = [&mut state.kv]; - let mut recurrent_refs = [&mut state.recurrent]; - self.model.batch_decode_eager_logits( - &[request.token_id], - &mut kv_refs, - &mut recurrent_refs, - &state.linear_pointer_tables, - &mut self.decode_buffers, - )?; - } - - if self.rank == 0 { - let cpu_logits = snapshot_requested_logprobs( - self.model.device_ctx(), - &self.decode_buffers.logits, - &[request.logprobs], - )?; - let params_refs = [&request.sampling_params]; - let tokens = pegainfer_sample::select_batch( - self.model.device_ctx(), - &self.decode_buffers.logits, - ¶ms_refs, - &[0], - sample_seed.wrapping_add(row_idx as u64), - &mut self.sample_scratch, - )?; - let token = tokens[0]; - let logprob = cpu_logits[0].as_ref().and_then(|row| { - pegainfer_sample::token_logprob_from_row(row, token, request.logprobs) - }); - primary_results.push(DecodeRequestResult { - request_id: request.request_id, - token, - logprob, - }); - } - } + let primary_results = self.run_decode_batch(requests, sample_seed)?; Ok(primary_results) } @@ -1485,27 +1560,16 @@ impl TpWorkerState { if let Some(idx) = self.request_index(request_id) { return Ok(idx); } - let mut recurrent = RecurrentState::new( + let recurrent = RecurrentState::new( self.model.device_ctx(), self.model.config(), self.model.geometry, )?; - let linear_pointer_tables = { - let mut recurrent_refs = [&mut recurrent]; - LinearStatePointerTables::from_recurrent_refs( - self.model.device_ctx(), - self.model.config(), - &mut recurrent_refs, - 1, - "Qwen3.5 TP eager", - )? - }; let state = TpRequestState { request_id, phase: TpRequestPhase::Prefilling, kv: self.model.alloc_kv(), recurrent, - linear_pointer_tables, }; self.requests.push(state); Ok(self.requests.len() - 1) From 1fdcc0a4977df8a14d25d954318c843a1e84f84b Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Thu, 20 Aug 2026 17:50:18 +0000 Subject: [PATCH 3/4] feat(qwen35): cuda graph decode under TP P2c from docs/models/qwen35/tp-design.md. TP decode captures/replays per-bucket CUDA Graphs when --cuda-graph is set AND the TP-local decode GQA group has a compiled kernel (4B/9B TP2); 27B TP2 (group 6) keeps the batched eager path byte-for-byte under the gate. - Gate: drop the fail-closed TP+graph rejections in config.rs/lib.rs/ tp_executor.rs; log once when graph was requested but the group gate keeps decode eager. - State: scheduler owns dense decode slots (slot_idx on TP decode rows, slot_for_new_request at promote, compaction_after_retire on retire); workers hold a fixed-address BatchDecodeGraphState plus slot_map, D2D copy prefill state into the slot on the first decode row, and apply DropRequest compactions via move_slot_within with occupancy assertions (poison on mismatch). - Capture/replay: startup pre-capture sweep ported from qwen3 (Warmup via new warmup_tp_collective, Capture/Launch per bucket, Finalize) behind a 600 s abort watchdog; batch_decode_graph gains DecodeGraphUse (Serve/CaptureOnly/Replay); TP serving is replay-only; graphs drop before the NCCL comm. - Convenience executor API (model-local tests) keeps a slot tracker mirroring Qwen35Executor; scheduler flows pass explicit slots. Validation (2x RTX 4090, venv NCCL): lib 105/105; 9B TP2 HF gates eager + graph (sequential replay, bucket-straddling, post-compaction) pass within existing tolerances; 9B TP2 scheduler e2e eager+graph pass; serving_tp2 now launches with graph on; 27B TP2 HF+e2e pass unchanged (group-6 gate keeps eager, graph variant self-skips). Serving benchmark 9B TP2, 16 concurrent x 256 out: 767.15 tok/s graph vs 705.86 eager (+8.7% steady output, TPOT 20.04 vs 21.99 ms). Signed-off-by: Ziyang Zhang --- docs/index.md | 4 +- docs/models/qwen35/tp-design.md | 24 + docs/models/qwen35/tp-implementation.md | 155 +++- pegainfer-core/src/weight_loader.rs | 1 + pegainfer-qwen35/src/batch_decode.rs | 67 +- pegainfer-qwen35/src/batch_decode_graph.rs | 48 ++ pegainfer-qwen35/src/config/error.rs | 4 - pegainfer-qwen35/src/config/tp.rs | 40 +- pegainfer-qwen35/src/executor.rs | 8 +- pegainfer-qwen35/src/lib.rs | 37 +- pegainfer-qwen35/src/scheduler.rs | 102 ++- pegainfer-qwen35/src/scheduler/tests.rs | 16 +- pegainfer-qwen35/src/tp_executor.rs | 836 +++++++++++++++++++-- pegainfer-qwen35/src/unified_forward.rs | 10 +- pegainfer-qwen35/src/weights.rs | 48 +- pegainfer-qwen35/tests/e2e_scheduler.rs | 32 +- pegainfer-qwen35/tests/hf_golden_gate.rs | 170 +++++ pegainfer-qwen35/tests/serving_tp2.rs | 34 +- 18 files changed, 1468 insertions(+), 168 deletions(-) diff --git a/docs/index.md b/docs/index.md index 63a5950b1..322816bf3 100644 --- a/docs/index.md +++ b/docs/index.md @@ -53,8 +53,8 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `models/qwen35/accuracy.md` | Qwen3.5 HF bf16 logits goldens, size-keyed (0.8b/2b/4b/9b/27b all committed), through `past_key_values`: short replay covers sequential graph, bucket-straddling batched graph, and slot-compaction; long replay covers 4097/8192-token prompts; full GSM8K 8-shot now matches the HF baseline within 0.15 percentage points. | | `models/qwen35/model-crate.md` | `pegainfer-qwen35` owns Qwen3.5 model/scheduler/recurrent ops/tests/benches; feature-gated behind `qwen35` (Triton AOT is the only Python build dependency); root loads it through `EngineHandle`. Build/check/clippy, root bench sanity check, historical Qwen3.5 e2e, and scheduler e2e records live here. | | `models/qwen35/batched-step-tail.md` | Qwen3.5 issue #353 implementation record: final prefill tail is batched, decode/unified sample from batched logits, host full-vocab copies are logprobs-only, HF + scheduler e2e pass, and final serving A/B supports only the first-token/short-output TTFT claim. | -| `models/qwen35/tp-design.md` | Qwen3.5 TP design: Phase 1 is eager dense TP on Qwen3's controller/worker runtime; validate TP2 first, fail closed for indivisible degrees and TP+CUDA Graph, shard dense full-attention/MLP, and leave sharded linear/GDR state to follow-up. | -| `models/qwen35/tp-implementation.md` | Qwen3.5 TP Phase 1 and P2A are complete: TP2 has start-gated eager unified prefill+decode, strict ID-aligned artifacts, fail-closed lifecycle recovery, and pre-load ordinal validation; P2B GDR state sharding is next. | +| `models/qwen35/tp-design.md` | Qwen3.5 TP design: Phase 1 eager dense TP on Qwen3's controller/worker runtime, P2a mixed-step protocol, P2b rank-local GDR sharding, P2c CUDA Graph under TP gated on the compiled decode GQA group (27B group-6 stays eager). | +| `models/qwen35/tp-implementation.md` | Qwen3.5 TP landed through P2c on #870 (2026-08, 2× RTX 4090): Phase 1/P2A lifecycle and ID contracts kept; GDR state sharded per rank (27B TP2 fits 48 GB pairs); batched eager decode (27B: 292 tok/s ×16); TP decode CUDA Graphs for 4B/9B (9B: 767 vs 706 tok/s ×16 eager). 9B/27B TP2 HF + e2e gates pass. | | `models/qwen35/mixed-load-itl-470.md` | Issue #470: full cold `--max-batch 8/bg=4` matrix on RTX 4090 (24/24 valid) + starvation negative control. Qwen3.5 is not immune; chunking bounds max/per-step stall but raises p99 at low QPS (~14→~80–92ms) and pulls p99/max back from the prefill wall to the chunk wall at high load; `qps·prefill_s≳1` is a throughput wall (chunking can't fix it, and ON's +15% TTFT can trip it earlier). The old "p99 immunity" was a slot-starvation artifact. | | `models/qwen35/adaptive-scheduler-policy.md` | Issue #727 adaptive scheduler policy record: default `off`, opt-in `auto`, hard `--max-prefill-tokens` cap, TP `auto` rejection, and pre-review whole-prefill benchmark tradeoff retained as non-default evidence. | | `models/qwen35/unified-prefill-overlap.md` | Issue #715 implementation record: opt-in single-GPU shared-SM overlap keeps one prefill chunk in flight while active decode continues; default serial policy and unsupported-combination guards remain explicit. | diff --git a/docs/models/qwen35/tp-design.md b/docs/models/qwen35/tp-design.md index 27d90b170..782fa38cc 100644 --- a/docs/models/qwen35/tp-design.md +++ b/docs/models/qwen35/tp-design.md @@ -269,6 +269,30 @@ Validation scope: - recurrent-state cleanup on finish/drop/cancellation - no stale local recurrent state after a new `RequestId` is admitted +## P2c: CUDA Graph under TP + +Status: landed (2026-08-20) on `feat/qwen35-tp2-rebased`, gated on +`local_decode_group_is_compiled` — 4B/9B TP2 capture and replay decode graphs; +27B TP2 (group 6) stays on the batched eager path byte-for-byte until group-6 +batch-decode kernels are compiled. Execution record: `tp-implementation.md` +section "P2c — CUDA Graph under TP". + +**Gate**: graph mode active iff `enable_cuda_graph && config.local_decode_group_is_compiled(tp)`. 27B TP2 is group-6 (`SUPPORTED_GQA_GROUP_SIZES = [1,2,3,4,8]`, group ratio is TP-invariant), so 27B TP2 keeps the batched eager path byte-for-byte until group-6 batch-decode kernels are compiled; 4B/9B TP2 capture graphs. Startup logs once when graph was requested but the group gate keeps decode eager. + +**State model**: scheduler owns slot semantics (TP1 mirror); workers execute slot copies on command, never infer slots worker-side. + +- KV paged state unchanged (pool stable; page tables are per-step H2D via `sync_paged_meta`). +- Per rank: `BatchDecodeGraphState`-equivalent at `bucket_for(effective_max_batch)` slots — fixed-address `slot_states: Vec` + one persistent `LinearStatePointerTables` built once over slots (contents stable → replay-safe). +- Admission: decode command rows carry explicit `slot_idx` (`slot_for_new_request`); first decode row D2D-copies prefill `RecurrentState` into the slot (`copy_state_to_slot`), drops the per-request allocation. +- Retirement: `DropRequest` gains `compaction: Option<(RequestId, from, to)>`; worker D2D-moves slot state (`move_slot_within`), asserts occupancy, poisons on mismatch. +- Decode rows arrive dense slot order `0..bs`; padding rows clobber free slots (benign — admission overwrites). + +**Capture/replay**: startup pre-capture sweep ported from qwen3 (`executor.rs:1424`): `Warmup` (port `warmup_tp_collective`, one all-reduce per bucket message size — lazy NCCL connect inside capture wedges), `Capture`/`Launch` per bucket `[1,2,4,8,16,32,64]` with synthetic rows, `Finalize` asserts all captured; dedicated 600 s abort watchdog (60 s startup timeout too small). New `TpWorkerCommand::Precapture { phase }` via existing exact-rank dispatch. Serve time: replay-only (`ensure is_captured` + `launch_captured`), never capture mid-serving. Sampling/logprobs stay rank-0 host-side outside the graph. Mixed ticks: prefill eager + decode replay; collective order canonical per plan. `TpWorkerState` declares graph state before `model` so graphs drop before the NCCL comm (teardown hang precedent qwen3 `executor.rs:3076`). + +**Memory** (27B TP2/rank): weights ~17.5 GB + KV pool ~5.9 GiB + slot state reserve ~6.1 GiB + buffers/graphs ~0.3 + scratch/NCCL ~2.5 ≈ 32 GiB → fits 48 GB. 9B TP2 slot state ~1.6 GiB. Loader already reserves `2 × max_batch × bytes_per_request` before sizing KV. + +**Validation ladder**: CPU lib suite → TP2 graph HF gate (9B: sequential + bucket-straddling + post-compaction replay vs eager stats) → e2e scheduler graph variant → serving_tp2 graph smoke → 27B TP2 regression unchanged (group-6 stays eager) → per-bucket eager-vs-graph decode benchmark recorded in `bench_snapshots/`. + ## References - `docs/models/qwen3/tp-design.md` diff --git a/docs/models/qwen35/tp-implementation.md b/docs/models/qwen35/tp-implementation.md index d8ebda80d..0bcfb1c26 100644 --- a/docs/models/qwen35/tp-implementation.md +++ b/docs/models/qwen35/tp-implementation.md @@ -1,6 +1,6 @@ # Qwen3.5 TP Implementation Record -> **TL;DR:** Qwen3.5 TP Phase 1 and P2A are complete: TP2 now supports start-gated eager unified prefill+decode with strict ID-aligned artifacts, fail-closed lifecycle recovery, and pre-load CUDA ordinal validation; P2B GDR state sharding is next. +> **TL;DR:** Qwen3.5 TP is complete through Phase 2b plus batched eager TP decode and P2c CUDA Graph under TP: TP2 supports start-gated eager unified prefill+decode with strict ID-aligned artifacts, fail-closed lifecycle recovery, and pre-load CUDA ordinal validation (#870); linear-attention/GDR state is sharded per rank (27B TP2 fits on 2×48 GB); TP decode rows run as one batched forward per step; and 4B/9B TP2 decode replays pre-captured CUDA Graphs (27B group-6 stays eager by gate). Remaining TP work: group-6 batch-decode kernels for 27B graphs, perf gates. > > **Last touched:** 2026-08 @@ -460,9 +460,160 @@ Non-negotiable invariant: - Never all-reduce GDR recurrent state or conv state. These states are owned by rank-local request state. +## Rebase onto #870 (2026-08-20) + +#870 landed its own Phase 1/2a upstream while the parallel line on +`feat/qwen35-tp2-batched-decode` had implemented its own Phase 1/2a/2b +plus a batched-decode fix on the old main. The rebase onto #870 +ported only the two deltas #870 lacks, with #870's code as the base: + +1. **Phase 2b — sharded linear-attention/GDR state** (commit + `feat(qwen35): shard linear-attention/GDR state per TP rank`). #870's + `recurrent_state.rs` had no rank sharding and its linear-attention + weights loaded replicated; the port adds rank-local slices end to end + (`weight_loader` stitch/shard loaders, `config.rs` `local_linear_*` + accessors, `weights.rs` per-rank stitched qkv/conv1d + row/col shards, + `recurrent_state`/`decode_buffers`/`prefill_buffers` at local sizes, + `batch_decode`/`prefill` local head counts + all-reduce after linear + `out_proj`, TP-local `batch_decode_full_attention_via_prefill` so 27B + TP2 group-6 eager decode routes through prefill). `tp_executor.rs` + only took the capacity-math and `RecurrentState::new` signature + changes; #870's worker protocol untouched. +2. **Step 3 — batched eager TP decode** (commit + `perf(qwen35): batch eager decode rows under TP`). #870's + `execute_decode_rows` looped per request with bs=1 forwards and + capacity-1 per-request pointer tables. The port adds `run_decode_batch` + (one `batch_decode_eager_logits` over all decode rows, step-scoped + `LinearStatePointerTables::from_recurrent_refs(..., bs, ...)`, one + batched rank-0 `select_batch`, per-row fan-out in command order) inside + #870's `execute_decode_rows`, keeping its validation and response + contracts; `TpRequestState.linear_pointer_tables` removed. + +What #870 already covered (not ported): Phase 1 dense TP, the Phase 2a +unified command/scheduler surface (`TpUnifiedPlan`, command start gates, +dispatch/response validators, drop-expectation lifecycle proofs), and the +scheduler planner-gate/test updates — ours' `scheduler.rs`, +`scheduler/tests.rs`, and `e2e_scheduler.rs` deltas were subsumed +upstream, so those files resolved to #870's versions except the +`alloc_recurrent` signature change. + +Validation on 2× RTX 4090: + +- `cargo check --release -p pegainfer-qwen35 --features qwen35` clean; + `cargo fmt --check -p pegainfer-qwen35` clean. +- Lib unit suite 101/101 (includes the four sharding layout tests: + segment tables, conv kernel-dim scaling, TP1 identity, synthetic + safetensors stitch contract). +- 9B TP2 HF short+long gates PASS (24.7 s); 9B TP2 scheduler e2e + (`test_e2e_qwen35_scheduler_tp2`) PASS (27.1 s). +- 27B TP2 HF short+long gates PASS (64.7 s) — 27B TP2 fits on 2×48 GB + only because of the Phase-2b sharding (the acceptance criterion for the + port). 27B TP2 scheduler e2e PASS (68.8 s). + +## P2c — CUDA Graph under TP (2026-08-20) + +Implemented the locked P2c design from `tp-design.md`: decode CUDA Graphs +under TP, gated on the TP-local decode GQA group. + +**Gate.** Graph mode is active iff `enable_cuda_graph && +config.local_decode_group_is_compiled(tp)`. The group ratio is TP-invariant, +so 27B TP2 (group 6, not in `SUPPORTED_GQA_GROUP_SIZES`) keeps the batched +eager path byte-for-byte while 4B/9B TP2 capture. The old fail-closed +rejections (`config.rs` `validate_for`, `tp_executor.rs` startup ensure, +`lib.rs` TP branch) were replaced by the gate; startup logs once when graph +was requested but the group gate keeps decode eager. + +**State model.** Scheduler owns slots, workers execute: + +- `ActiveBackendState::Tp` gains `slot_idx` (dense `active` position, assigned + at promote via `slot_for_new_request`, updated on compaction); + `tp_decode_items` emits rows with explicit `slot_idx` and workers assert + `slot_idx == row`. +- Graph workers hold a `BatchDecodeGraphState` at + `bucket_for(effective_max_batch)` fixed-address slots plus a + `slot_map: Vec>`. On a request's first decode row the + worker D2D-copies its prefill recurrent state into the slot + (`copy_state_to_slot`) and drops the per-request allocation. +- `DropRequest` carries `compaction: Option`; the worker + validates occupancy against `slot_map` (`slot_compact`), applies the D2D + move (`BatchDecodeGraphState::move_slot_within`), and poisons on mismatch. + Requests retired between promotion and their first decode row legitimately + have no materialized slot; `slot_compact` tolerates exactly that case and + skips the GPU move. +- Convenience `execute_prefill`/`execute_decode`/`drop_request` (model-local + tests) keep a Mutex-guarded slot tracker mirroring the single-GPU + `Qwen35Executor`; scheduler flows pass explicit slots via + `execute_decode_items`/`drop_request_with_compaction` and never touch the + tracker. + +**Capture/replay.** Startup pre-capture sweep ported from qwen3: +`TpWorkerCommand::Precapture { phase }` over Warmup (new +`Qwen35Model::warmup_tp_collective` — one all-reduce per bucket message size; +without it the lazy NCCL connect wedges inside capture), Capture + Launch per +bucket `[1,2,4,8,16,32,64]` up to `bucket_for(max_batch)` with synthetic rows, +Finalize asserting all buckets captured. Dedicated 600 s abort watchdog (the +60 s NCCL startup timeout is too small for the sweep). +`batch_decode_graph` gained `DecodeGraphUse` (Serve lazy / CaptureOnly / +Replay); TP serve time is Replay-only. Workers tune decode GEMM algos on the +worker thread before capture (cuBLASLt plans are thread-local). Graph state +is declared before `model` in `TpWorkerState` so graphs drop before the NCCL +comm. Sampling/logprobs stay rank-0 host-side outside the graph. Eager +workers ignore `slot_idx`/`compaction`, keeping 27B TP2 byte-identical. + +**Gotcha fixed during validation:** `track_retired_slot` used +`bool::then_some`, which evaluates eagerly and indexed past the tail when the +retired request was the last slot — use `then` for the lazy closure. + +**Validation (2× RTX 4090):** + +- `cargo check --release -p pegainfer-qwen35 --features qwen35` clean; lib + suite 105/105 (new: group-gate acceptance incl. group-6 stays eager, slot + map admit/compact/mismatch CPU tests); `cargo fmt --check` clean. +- 9B TP2 HF gates (`--test-threads=1`): eager sequential+batched PASS; + graph sequential replay (identical fingerprints across reruns), + bucket-straddling batched replay (5→bucket 8, 3→bucket 4), and + post-compaction replay after a mid-batch drop all within the existing TP2 + tolerances (worst graph arm: mean 0.0228, p99 0.1002 against MEAN_TOL 0.06 + / P99_TOL 0.20; eager arm mean 0.0227/0.0230). +- 9B TP2 scheduler e2e eager + graph variants PASS; 9B TP2 serving smoke now + launches with `--cuda-graph true` (graph acceptance replaced the old + rejection assertion). +- 27B TP2 HF short+long and scheduler e2e PASS unchanged — the gate log + confirms group 6 keeps decode eager (the graph test variant skips itself + via `graph_enabled()`). +- 9B TP2 serving benchmark (`pegainfer-server --tp-size 2 --port 18093`, + vllm-bench `openai` backend, random 128-in/256-out, 64 prompts at + concurrency 16, greedy, ignore_eos, seed 42): + + | arm | steady output tok/s | mean TPOT (ms) | total tok/s | + |-----|--------------------:|---------------:|------------:| + | CUDA Graph on | 767.15 | 20.04 | 1146.65 | + | CUDA Graph off | 705.86 | 21.99 | 1054.56 | + + Graph decode is +8.7% steady output tok/s (-8.8% TPOT) at 16 concurrent. + The design's "record in `bench_snapshots/`" step was skipped: the + in-process snapshot gate is retired (`docs/conventions/bench-regression.md`), + so the HTTP bench numbers live here instead. + +**Test-isolation note:** TP2 GPU tests must run with `--test-threads=1`. Two +TP executors sharing the GPUs perturb cuBLASLt algorithm selection (workspace +pressure), which once flipped a sequential-replay fingerprint comparison in +the *eager* test while the graph test ran concurrently. + ## Follow-Ups -- Design and implement P2B sharded linear-attention/GDR state without weakening the completed P2A lifecycle and ID contracts. +- P2c CUDA Graph under TP landed for compiled decode GQA groups (4B/9B); 27B + TP2 graphs stay gated off until group-6 batch-decode kernels are compiled + (`SUPPORTED_GQA_GROUP_SIZES`). The eager path is the 27B fallback and must + not regress. +- 27B TP2 knowledge-benchmark parity (2026-08-20, validated pre-rebase on + the parallel TP line; `docs/benchmarks/qwen35-27b-tp2-knowledge-eval.md`): + MMLU-Redux 94.09 vs official 93.2 (full 5330), C-Eval 88.11 vs 90.5 + (full 1346, thinking-cap truncation rerun-merged) — inside the + cross-harness band, no TP-induced accuracy regression. MMLU-Pro / + SuperGPQA sampled runs remain outstanding; rerun on this rebased branch + before citing parity. +- P2B sharded linear-attention/GDR state landed (see "Rebase onto #870"); keep the completed P2A lifecycle and ID contracts unweakened. - Promote any stable contract changes discovered here back into `tp-design.md` through the design-doc branch. - Decide whether Qwen3.5 server CLI should accept arbitrary TP device ordinals instead of only `0..tp_size`. - Consider lifting the per-device Triton AOT handle lesson into a kernels or runtime subsystem doc if another model hits the same issue. diff --git a/pegainfer-core/src/weight_loader.rs b/pegainfer-core/src/weight_loader.rs index c2282d322..3261b3671 100644 --- a/pegainfer-core/src/weight_loader.rs +++ b/pegainfer-core/src/weight_loader.rs @@ -868,6 +868,7 @@ pub fn load_tensor_1d_shard( load_tensor_1d_stitch(ctx, shards, weight_map, name, &[(offset, len)]) } +#[allow(clippy::cast_ptr_alignment)] /// Load a 1D F32 element range to GPU (tensor-parallel shard of a 1D weight). pub fn load_tensor_1d_f32_shard( ctx: &DeviceContext, diff --git a/pegainfer-qwen35/src/batch_decode.rs b/pegainfer-qwen35/src/batch_decode.rs index 0c80feaae..24c25d936 100644 --- a/pegainfer-qwen35/src/batch_decode.rs +++ b/pegainfer-qwen35/src/batch_decode.rs @@ -26,6 +26,22 @@ use crate::ops; static LOG_UNCOMPILED_DECODE_ROUTE: std::sync::Once = std::sync::Once::new(); +/// How a `batch_decode_graph` call interacts with the per-bucket CUDA graphs. +/// +/// TP serving never captures lazily: a mid-serving capture on one rank while a +/// peer replays desyncs the recorded NCCL collectives, so tensor-parallel +/// decodes are replay-only after the startup pre-capture sweep. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub(crate) enum DecodeGraphUse { + /// Replay if captured, lazily capture otherwise (single-GPU serving). + Serve, + /// Record + instantiate + upload, no launch (the TP sweep's Capture phase). + CaptureOnly, + /// Replay only; error if never captured (TP serving, and the TP sweep's + /// Launch phase that drains the captured collectives across ranks). + Replay, +} + impl Qwen35Model { pub(crate) fn select_tokens_from_logits_varied( &self, @@ -296,7 +312,9 @@ impl Qwen35Model { // Route by TP-local GQA group supportability: when the rank-local // q-per-kv group has no compiled batch-decode kernel, run full // attention through the paged-prefill kernel with a per-step plan. - let prefill_attn_plan = if !self.geometry.local_decode_group_is_compiled() { + let prefill_attn_plan = if self.geometry.local_decode_group_is_compiled() { + None + } else { let page_indices: Vec> = kv_states.iter().map(|kv| kv.page_indices_i32()).collect(); let last_page_lens: Vec = @@ -326,8 +344,6 @@ impl Qwen35Model { ) })?, ) - } else { - None }; let kv_buffer = kv_states[0].buffer(); @@ -361,6 +377,7 @@ impl Qwen35Model { token_ids: &[u32], kv_states: &mut [&mut KvState], graph_state: &mut BatchDecodeGraphState, + graph_use: DecodeGraphUse, ) -> Result<()> { let bs = token_ids.len(); anyhow::ensure!(bs > 0, "batch_decode_graph requires at least one request"); @@ -372,6 +389,10 @@ impl Qwen35Model { ); if !self.config.decode_group_is_compiled() { + anyhow::ensure!( + graph_use == DecodeGraphUse::Serve, + "Qwen3.5 batched hybrid eager fallback only supports lazy serve-mode decode, got {graph_use:?}" + ); LOG_UNCOMPILED_DECODE_ROUTE.call_once(|| { let group = self.config.num_attention_heads / self.config.num_key_value_heads; log::info!( @@ -432,17 +453,35 @@ impl Qwen35Model { let mut graphs = std::mem::take(&mut graph_state.graphs); let linear_state_ptrs = &graph_state.linear_pointer_tables.state_ptrs; let linear_conv_state_ptrs = &graph_state.linear_pointer_tables.conv_state_ptrs; - let result = graphs[bucket_idx].run_or_capture(&self.ctx, || { - self.batch_decode_kernels_graph( - kv_buffer, - &layout, - padded_bs, - None, - linear_state_ptrs, - linear_conv_state_ptrs, - &mut graph_state.buffers, - ) - }); + let result = match graph_use { + DecodeGraphUse::Serve => graphs[bucket_idx].run_or_capture(&self.ctx, || { + self.batch_decode_kernels_graph( + kv_buffer, + &layout, + padded_bs, + None, + linear_state_ptrs, + linear_conv_state_ptrs, + &mut graph_state.buffers, + ) + }), + DecodeGraphUse::CaptureOnly => graphs[bucket_idx].capture_only(&self.ctx, || { + self.batch_decode_kernels_graph( + kv_buffer, + &layout, + padded_bs, + None, + linear_state_ptrs, + linear_conv_state_ptrs, + &mut graph_state.buffers, + ) + }), + // Replay is a pure enqueue: every bucket was recorded by the + // startup pre-capture sweep, so a missing graph here means the + // sweep was skipped or incomplete — fail loudly, never capture + // mid-serving (a one-sided capture desyncs TP collectives). + DecodeGraphUse::Replay => graphs[bucket_idx].launch_captured(&self.ctx), + }; graph_state.graphs = graphs; result } diff --git a/pegainfer-qwen35/src/batch_decode_graph.rs b/pegainfer-qwen35/src/batch_decode_graph.rs index 59d5225d3..e68348ea1 100644 --- a/pegainfer-qwen35/src/batch_decode_graph.rs +++ b/pegainfer-qwen35/src/batch_decode_graph.rs @@ -128,4 +128,52 @@ impl BatchDecodeGraphState { dst.seq_len = src.seq_len; Ok(()) } + + /// D2D move slot `from` into slot `to`, leaving `to` as the canonical + /// state. Used by TP slot compaction when a mid-batch request retires: + /// the last occupied slot moves into the vacated slot so decode rows stay + /// dense (`0..bs`). Slot `from` keeps stale bytes afterwards; admission + /// overwrites it (`copy_state_to_slot`), and padding rows may clobber it — + /// both benign by design. + pub(crate) fn move_slot_within( + &mut self, + ctx: &DeviceContext, + from: usize, + to: usize, + ) -> Result<()> { + anyhow::ensure!( + from != to, + "TP slot compaction move {from} -> {to} is a no-op" + ); + anyhow::ensure!( + from < self.slot_states.len() && to < self.slot_states.len(), + "TP slot compaction move {from} -> {to} exceeds {} slots", + self.slot_states.len() + ); + let (lo, hi) = (from.min(to), from.max(to)); + let (left, right) = self.slot_states.split_at_mut(hi); + let (src, dst) = if from < to { + (&left[from], &mut right[0]) + } else { + (&right[0], &mut left[lo]) + }; + anyhow::ensure!( + src.layers.len() == dst.layers.len(), + "TP slot compaction layer count mismatch: {} vs {}", + src.layers.len(), + dst.layers.len() + ); + for (src_layer, dst_layer) in src.layers.iter().zip(dst.layers.iter_mut()) { + ctx.stream + .memcpy_dtod(&src_layer.state, &mut dst_layer.state) + .map_err(|e| anyhow::anyhow!("move recurrent state slot {from} -> {to}: {e}"))?; + ctx.stream + .memcpy_dtod(&src_layer.conv_state.data, &mut dst_layer.conv_state.data) + .map_err(|e| anyhow::anyhow!("move conv state slot {from} -> {to}: {e}"))?; + } + let seq_len = src.seq_len; + let dst = &mut self.slot_states[to]; + dst.seq_len = seq_len; + Ok(()) + } } diff --git a/pegainfer-qwen35/src/config/error.rs b/pegainfer-qwen35/src/config/error.rs index b44a5b996..3b8dbfbb1 100644 --- a/pegainfer-qwen35/src/config/error.rs +++ b/pegainfer-qwen35/src/config/error.rs @@ -68,10 +68,6 @@ pub(crate) enum ConfigError { TpZeroWorldSize, #[error("tensor_parallel.rank {rank} must be < world_size {world_size}")] TpRankOutOfRange { rank: usize, world_size: usize }, - #[error( - "Qwen3.5 tensor parallelism is eager-only; disable CUDA Graph for tp world_size={world_size}" - )] - TpRequiresEager { world_size: usize }, #[error("{field}={value} not divisible by tp world_size={world_size}")] TpIndivisible { field: &'static str, diff --git a/pegainfer-qwen35/src/config/tp.rs b/pegainfer-qwen35/src/config/tp.rs index 5506067ac..356de76fc 100644 --- a/pegainfer-qwen35/src/config/tp.rs +++ b/pegainfer-qwen35/src/config/tp.rs @@ -91,25 +91,21 @@ pub(crate) struct LocalGeometry { } impl LocalGeometry { - /// Validate `config` against `tp` and the runtime execution mode, then derive - /// this rank's local dimensions. + /// Validate `config` against `tp`, then derive this rank's local dimensions. /// /// Fails on unsupported combinations before expensive loading: - /// - sharded TP demands eager execution (`enable_cuda_graph` off); /// - every sharded model dimension must divide evenly by `world_size` /// (linear-attention heads included: Phase 2b shards them per rank); /// - `rank < world_size` and `world_size >= 1` are guaranteed by /// `TensorParallelConfig::try_from`. + /// + /// CUDA Graph under TP is gated at executor startup on + /// [`LocalGeometry::local_decode_group_is_compiled`] (P2c): uncompiled GQA + /// groups keep the batched eager path instead of failing validation here. pub(crate) fn try_new( config: &Config35, tp: TensorParallelConfig, - enable_cuda_graph: bool, ) -> Result { - if tp.is_sharded() && enable_cuda_graph { - return Err(ConfigError::TpRequiresEager { - world_size: tp.world_size(), - }); - } if !config.num_attention_heads.is_multiple_of(tp.world_size()) { return Err(ConfigError::TpIndivisible { field: "num_attention_heads", @@ -299,7 +295,7 @@ mod tests { fn tp2_local_geometry_matches_dense_dims() { let cfg = config(); let tp = TensorParallelConfig::try_from((1, 2)).unwrap(); - let geom = LocalGeometry::try_new(&cfg, tp, false).unwrap(); + let geom = LocalGeometry::try_new(&cfg, tp).unwrap(); assert!(geom.is_sharded()); assert_eq!(geom.shard_range(4096), (2048, 2048)); assert_eq!(geom.local_num_attention_heads(), 8); @@ -329,7 +325,7 @@ mod tests { fn rejects_indivisible_dense_dimensions() { let tp = TensorParallelConfig::try_from((0, 3)).unwrap(); let cfg = config(); - let mut err = LocalGeometry::try_new(&cfg, tp, false).unwrap_err(); + let mut err = LocalGeometry::try_new(&cfg, tp).unwrap_err(); assert_eq!( err, ConfigError::TpIndivisible { @@ -342,7 +338,7 @@ mod tests { let mut broken = cfg; broken.num_attention_heads = 15; broken.num_key_value_heads = 4; - err = LocalGeometry::try_new(&broken, tp, false).unwrap_err(); + err = LocalGeometry::try_new(&broken, tp).unwrap_err(); assert_eq!( err, ConfigError::TpIndivisible { @@ -354,7 +350,7 @@ mod tests { broken.num_key_value_heads = 3; broken.intermediate_size = 9217; - err = LocalGeometry::try_new(&broken, tp, false).unwrap_err(); + err = LocalGeometry::try_new(&broken, tp).unwrap_err(); assert_eq!( err, ConfigError::TpIndivisible { @@ -365,22 +361,12 @@ mod tests { ); } - #[test] - fn rejects_tensor_parallel_with_cuda_graph() { - let cfg = config(); - let tp = TensorParallelConfig::try_from((0, 2)).unwrap(); - assert_eq!( - LocalGeometry::try_new(&cfg, tp, true), - Err(ConfigError::TpRequiresEager { world_size: 2 }) - ); - } - #[test] fn requires_linear_attention_head_divisibility() { let tp = TensorParallelConfig::try_from((1, 2)).unwrap(); let mut broken = config(); broken.linear_num_key_heads = 17; - let err = LocalGeometry::try_new(&broken, tp, false).unwrap_err(); + let err = LocalGeometry::try_new(&broken, tp).unwrap_err(); assert_eq!( err, ConfigError::TpIndivisible { @@ -392,7 +378,7 @@ mod tests { let mut broken = config(); broken.linear_num_value_heads = 31; - let err = LocalGeometry::try_new(&broken, tp, false).unwrap_err(); + let err = LocalGeometry::try_new(&broken, tp).unwrap_err(); assert_eq!( err, ConfigError::TpIndivisible { @@ -407,7 +393,7 @@ mod tests { fn computes_tp2_linear_attention_local_dimensions() { let cfg = config(); let tp = TensorParallelConfig::try_from((1, 2)).unwrap(); - let geom = LocalGeometry::try_new(&cfg, tp, false).unwrap(); + let geom = LocalGeometry::try_new(&cfg, tp).unwrap(); assert_eq!(geom.local_linear_num_key_heads(), 8); assert_eq!(geom.local_linear_num_value_heads(), 16); assert_eq!(geom.local_linear_q_dim, 1024); @@ -421,7 +407,7 @@ mod tests { // TP1 invariant: every local dim equals the global dim, keeping TP1 // numerics byte-identical to pre-sharding execution. let cfg = config(); - let geom = LocalGeometry::try_new(&cfg, TensorParallelConfig::default(), false).unwrap(); + let geom = LocalGeometry::try_new(&cfg, TensorParallelConfig::default()).unwrap(); assert_eq!(geom.local_linear_num_key_heads(), 16); assert_eq!(geom.local_linear_num_value_heads(), 32); assert_eq!(geom.local_linear_q_dim, 2048); diff --git a/pegainfer-qwen35/src/executor.rs b/pegainfer-qwen35/src/executor.rs index 956d0f9ab..67bf5568e 100644 --- a/pegainfer-qwen35/src/executor.rs +++ b/pegainfer-qwen35/src/executor.rs @@ -232,8 +232,12 @@ impl Qwen35Executor { let token_ids: Vec = plan.requests.iter().map(|req| req.token_id).collect(); let mut kv_refs: Vec<&mut KvState> = self.active.iter_mut().map(|req| &mut req.kv).collect(); - self.model - .batch_decode_graph(&token_ids, &mut kv_refs, &mut self.graph_state)?; + self.model.batch_decode_graph( + &token_ids, + &mut kv_refs, + &mut self.graph_state, + crate::batch_decode::DecodeGraphUse::Serve, + )?; let requested_logprobs: Vec = plan.requests.iter().map(|req| req.logprobs).collect(); let cpu_logits = snapshot_requested_logprobs( diff --git a/pegainfer-qwen35/src/lib.rs b/pegainfer-qwen35/src/lib.rs index 0638d222c..ec1133871 100644 --- a/pegainfer-qwen35/src/lib.rs +++ b/pegainfer-qwen35/src/lib.rs @@ -96,6 +96,9 @@ pub enum Qwen35DecodeOverlap { SharedSm, } +/// TP decode runs CUDA Graphs when `cuda_graph` is set AND the rank-local +/// decode GQA group has a compiled kernel (`tp-design.md` P2c gate); otherwise +/// TP keeps the batched eager path. pub fn start_engine( model_path: &Path, options: EngineLoadOptions, @@ -117,7 +120,8 @@ pub struct Qwen35LaunchOptions { device_ordinal: usize, /// Tensor-parallel world size; `> 1` uses devices `0..tp_size`. tp_size: usize, - /// TP Phase 1 supports eager-only multi-GPU execution. + /// TP decode captures CUDA Graphs when the rank-local decode GQA group has + /// a compiled kernel; uncompiled groups (27B) stay on the eager path. cuda_graph: bool, max_batch: usize, max_prefill_tokens: usize, @@ -229,11 +233,6 @@ pub fn start_engine_with_capacity_policy_and_overlap( "Qwen3.5 TP uses the fixed off scheduler policy; --qwen35-scheduler-policy=auto is single-GPU only" )); } - if enable_cuda_graph { - return Err(anyhow!( - "Qwen3.5 TP Phase 1 supports eager execution only; disable CUDA Graph" - )); - } let model_path = model_path .to_str() .ok_or_else(|| anyhow!("model path must be valid UTF-8"))?; @@ -243,6 +242,7 @@ pub fn start_engine_with_capacity_policy_and_overlap( &device_ordinals, max_batch, max_prefill_tokens, + enable_cuda_graph, ); } @@ -307,6 +307,31 @@ mod tests { assert_eq!(Qwen35SchedulerPolicy::default(), Qwen35SchedulerPolicy::Off); } + #[test] + fn tp_cuda_graph_is_no_longer_rejected_before_model_load() { + // P2c: the graph/eager decision is gated on the model's TP-local decode + // GQA group, which needs the loaded config — so TP + CUDA Graph passes + // pre-load validation and fails here only because the path is bogus. + let err = start_engine_with_capacity_and_policy( + Path::new("unused-model-path"), + EngineLoadOptions { + enable_cuda_graph: true, + device_ordinals: vec![0, 1], + parallel_config: None, + ep_backend: EpBackend::Nccl, + seed: 42, + }, + 1, + 1, + Qwen35SchedulerPolicy::Off, + ) + .err() + .expect("nonexistent model path should fail at load") + .to_string(); + + assert!(!err.contains("eager execution only")); + } + #[test] fn tp_rejects_auto_scheduler_policy_before_loading_model() { let err = start_engine_with_capacity_and_policy( diff --git a/pegainfer-qwen35/src/scheduler.rs b/pegainfer-qwen35/src/scheduler.rs index bf4259c47..3a276a1db 100644 --- a/pegainfer-qwen35/src/scheduler.rs +++ b/pegainfer-qwen35/src/scheduler.rs @@ -69,6 +69,7 @@ use crate::tp_executor::DropExpectation; use crate::tp_executor::Qwen35TpExecutor; use crate::tp_executor::TpDecodeStepItem; use crate::tp_executor::TpPrefillChunkItem; +use crate::tp_executor::TpSlotCompaction; use crate::tp_executor::TpUnifiedPlan; use crate::weights::Qwen35Model; @@ -109,6 +110,9 @@ enum ActiveBackendState { }, Tp { request_id: RequestId, + /// Dense decode slot (`active` position). Graph-mode workers assert + /// `slot_idx == row` on every decode command; eager workers ignore it. + slot_idx: usize, }, } @@ -418,13 +422,19 @@ pub(crate) fn start_tp_with_capacity( device_ordinals: &[usize], max_batch: usize, max_prefill_tokens: usize, + enable_cuda_graph: bool, ) -> Result { assert!( max_prefill_tokens > 0, "max_prefill_tokens must be positive: a zero budget can never schedule a prefill chunk" ); - let backend = - TpSchedulerBackend::new(model_path, device_ordinals, max_batch, max_prefill_tokens)?; + let backend = TpSchedulerBackend::new( + model_path, + device_ordinals, + max_batch, + max_prefill_tokens, + enable_cuda_graph, + )?; let servable = servable_len( backend.max_position_embeddings(), backend.capacity_pages_for_requests(), @@ -530,6 +540,9 @@ fn fatal_cuda_lifecycle(message: &str) -> ! { struct TpSchedulerBackend { executor: Qwen35TpExecutor, next_request_id: u64, + /// Slot move derived by the in-flight `take_active_request`; consumed by + /// the paired `drop_active_state` so the workers apply the same move. + pending_compaction: Option, } impl SingleGpuBackend { @@ -706,8 +719,12 @@ impl SingleGpuBackend { } }) .collect(); - self.model - .batch_decode_graph(&token_ids, &mut kv_refs, &mut self.graph_state) + self.model.batch_decode_graph( + &token_ids, + &mut kv_refs, + &mut self.graph_state, + crate::batch_decode::DecodeGraphUse::Serve, + ) } fn sample_prefill_logits( @@ -847,10 +864,11 @@ impl TpSchedulerBackend { device_ordinals: &[usize], max_batch: usize, max_prefill_tokens: usize, + enable_cuda_graph: bool, ) -> Result { let executor = Qwen35TpExecutor::from_runtime_with_limits( model_path, - false, + enable_cuda_graph, device_ordinals, max_batch, max_prefill_tokens, @@ -858,6 +876,7 @@ impl TpSchedulerBackend { Ok(Self { executor, next_request_id: 1, + pending_compaction: None, }) } @@ -955,7 +974,40 @@ impl TpSchedulerBackend { } fn drop_request(&self, request_id: RequestId, expectation: DropExpectation) -> Result<()> { - self.executor.drop_request(request_id, expectation) + self.executor + .drop_request_with_compaction(request_id, expectation, None) + } + + /// Remove the TP request at `idx` via swap_remove and stash the resulting + /// slot compaction for the paired `drop_active_state`. Mirrors + /// `compact_single_slot`: after the swap, slots `0..active.len()` stay + /// dense because the moved request's slot follows it. + fn take_active_request( + &mut self, + active: &mut Vec, + idx: usize, + ) -> ActiveRequest35 { + let compaction = compaction_after_retire(active.len(), idx); + let removed = active.swap_remove(idx); + + self.pending_compaction = compaction.map(|compaction| { + let moved = &mut active[idx]; + let ActiveBackendState::Tp { + request_id, + slot_idx, + } = &mut moved.backend_state + else { + panic!("TP scheduler received single-GPU active state") + }; + debug_assert_eq!(*slot_idx, compaction.moved_from); + *slot_idx = compaction.moved_to; + TpSlotCompaction { + moved_request_id: *request_id, + from: compaction.moved_from, + to: compaction.moved_to, + } + }); + removed } } @@ -1059,15 +1111,25 @@ fn tp_prefill_items(chunk: &ScheduledChunk) -> Result> { fn tp_decode_items(active: &[ActiveRequest35]) -> Result> { active .iter() - .map(|req| { - let ActiveBackendState::Tp { request_id } = &req.backend_state else { + .enumerate() + .map(|(row, req)| { + let ActiveBackendState::Tp { + request_id, + slot_idx, + } = &req.backend_state + else { anyhow::bail!("TP decode received single-GPU active state"); }; - Ok(TpDecodeStepItem::new( + debug_assert_eq!( + *slot_idx, row, + "TP decode slots must stay dense in active order" + ); + Ok(TpDecodeStepItem::new_with_slot( *request_id, req.last_token, req.logprobs, req.params, + *slot_idx, )) }) .collect() @@ -1145,7 +1207,7 @@ fn align_decode_results( let expected: Vec = active .iter() .map(|active_req| { - let ActiveBackendState::Tp { request_id } = active_req.backend_state else { + let ActiveBackendState::Tp { request_id, .. } = active_req.backend_state else { anyhow::bail!("align_decode_results requires TP active state"); }; Ok(request_id) @@ -2212,15 +2274,20 @@ impl DecodeDispatchBackend for SchedulerBackend { ) -> ActiveRequest35 { match self { SchedulerBackend::Single(backend) => compact_single_slot(backend, active, idx), - SchedulerBackend::Tp(_) => active.swap_remove(idx), + SchedulerBackend::Tp(backend) => backend.take_active_request(active, idx), } } fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { match (self, state) { (SchedulerBackend::Single(_), ActiveBackendState::Single { .. }) => Ok(()), - (SchedulerBackend::Tp(backend), ActiveBackendState::Tp { request_id }) => { - backend.drop_request(*request_id, DropExpectation::MustExist) + (SchedulerBackend::Tp(backend), ActiveBackendState::Tp { request_id, .. }) => { + let compaction = backend.pending_compaction.take(); + backend.executor.drop_request_with_compaction( + *request_id, + DropExpectation::MustExist, + compaction, + ) } _ => anyhow::bail!("mismatched Qwen3.5 scheduler backend state during retirement"), } @@ -2586,8 +2653,13 @@ impl PrefillPromoteBackend for SchedulerBackend { graph_slot_idx: slot_idx, } } - (SchedulerBackend::Tp(_), PrefillBackendState::Tp { request_id }) => { - ActiveBackendState::Tp { request_id } + (SchedulerBackend::Tp(backend), PrefillBackendState::Tp { request_id }) => { + let slot_idx = slot_for_new_request(active_len, backend.max_batch()) + .expect("admission must reserve a TP decode slot"); + ActiveBackendState::Tp { + request_id, + slot_idx, + } } _ => panic!("mismatched Qwen3.5 scheduler backend state during promotion"), } diff --git a/pegainfer-qwen35/src/scheduler/tests.rs b/pegainfer-qwen35/src/scheduler/tests.rs index 39a521517..1775a1fc9 100644 --- a/pegainfer-qwen35/src/scheduler/tests.rs +++ b/pegainfer-qwen35/src/scheduler/tests.rs @@ -44,6 +44,7 @@ fn active_request(request_id: u64, label: &str, token_tx: TokenSink) -> ActiveRe token_tx, backend_state: ActiveBackendState::Tp { request_id: RequestId::new(request_id), + slot_idx: 0, }, last_token: 1, generated_count: 1, @@ -89,7 +90,7 @@ impl DecodeDispatchBackend for PruneTestBackend { } fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { - let ActiveBackendState::Tp { request_id } = state else { + let ActiveBackendState::Tp { request_id, .. } = state else { panic!("prune test expected TP active state"); }; self.retired_active.push(*request_id); @@ -180,7 +181,7 @@ impl DecodeDispatchBackend for LifecycleTestBackend { } fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { - let ActiveBackendState::Tp { request_id } = state else { + let ActiveBackendState::Tp { request_id, .. } = state else { panic!("lifecycle test expected TP active state"); }; if self.active_completion_requires_drop_ack { @@ -879,7 +880,10 @@ fn inflight_prefill_waits_instead_of_parking_after_last_decode_retires() { } #[test] -fn tp_engine_rejects_cuda_graph_before_model_load() { +fn tp_engine_cuda_graph_passes_preload_validation() { + // P2c: TP + CUDA Graph is gated on the model's TP-local decode GQA group + // after load, so startup with a bogus path fails at load, not at the old + // eager-only rejection. let err = match crate::start_engine_with_capacity( Path::new("unused"), EngineLoadOptions { @@ -892,10 +896,10 @@ fn tp_engine_rejects_cuda_graph_before_model_load() { 1, 1, ) { - Ok(_) => panic!("TP CUDA Graph startup should fail"), + Ok(_) => panic!("TP CUDA Graph startup with a nonexistent path should fail at load"), Err(err) => err.to_string(), }; - assert!(err.contains("eager execution only")); + assert!(!err.contains("eager execution only")); } #[test] @@ -907,7 +911,7 @@ fn tp2_scheduler_runs_forced_mixed_steps() { return; }; let handle = - start_tp_with_capacity(&model_path, 42, &[0, 1], 2, 1).expect("start TP2 scheduler"); + start_tp_with_capacity(&model_path, 42, &[0, 1], 2, 1, false).expect("start TP2 scheduler"); let (decode_tx, mut decode_rx) = TokenSink::standalone(); let (prefill_tx, mut prefill_rx) = TokenSink::standalone(); diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index 885a5d844..a3c2fee3a 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -17,16 +17,22 @@ use std::sync::atomic::Ordering; use std::sync::mpsc; use std::thread::JoinHandle; use std::thread::{self}; +use std::time::Instant; use anyhow::Result; use pegainfer_core::kv_pool::KvState; use pegainfer_frontend::sampler::SamplingParams; +use crate::batch_decode::DecodeGraphUse; +use crate::batch_decode_graph::BATCH_BUCKETS; +use crate::batch_decode_graph::BatchDecodeGraphState; +use crate::batch_decode_graph::bucket_for; use crate::config::TensorParallelConfig; use crate::decode_buffers::BatchDecodeBuffers35; use crate::executor::DecodePlan; use crate::executor::DecodeRequestResult; use crate::executor::DecodeResult; +#[cfg(test)] use crate::executor::DecodeStepItem; use crate::executor::PrefillPlan; use crate::executor::PrefillRequestResult; @@ -44,9 +50,43 @@ use crate::weights::Qwen35Model; const TP_NCCL_STARTUP_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60); 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 (qwen3 uses the same 600 s). +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. +/// +/// Capture and launch are separate phases because a captured collective's +/// first launch blocks on its peers: overlapping that with a peer still in +/// capture/instantiate/upload (which contend driver locks and allocate device +/// memory) deadlocks the driver. So every rank finishes capturing a bucket +/// before any rank launches it. (Ported from qwen3's TP sweep.) +#[derive(Clone, Copy, Debug)] +enum PrecapturePhase { + /// One eager all-reduce per bucket message size, so the size-selected NCCL + /// algorithm connects before any `cuStreamBeginCapture` records it. + Warmup, + /// Record + instantiate + upload one bucket; no launch, no cross-rank dependency. + Capture { bucket_idx: usize }, + /// Launch one bucket (pure enqueue after `Capture`) + sync; collectives pair across ranks. + Launch { bucket_idx: usize }, + /// Verify every reachable bucket captured. + Finalize, +} + +/// Scheduler-owned slot move for TP graph decode: when the request at slot +/// `to` retires mid-batch, the request at slot `from` (the last occupied slot) +/// takes over slot `to` so decode rows stay dense. Workers apply the move and +/// fail (poisoning the executor) if slot occupancy does not match. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub(crate) struct TpSlotCompaction { + pub(crate) moved_request_id: RequestId, + pub(crate) from: usize, + pub(crate) to: usize, +} + #[allow(dead_code)] enum TpWorkerCommand { Ping { @@ -71,6 +111,17 @@ enum TpWorkerCommand { }, DropRequest { request_id: RequestId, + /// Slot move the scheduler already applied to its own bookkeeping; + /// `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, + 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, resp: mpsc::Sender, }, @@ -196,6 +247,18 @@ pub struct Qwen35TpExecutor { capacity_pages_for_requests: usize, max_position_embeddings: usize, eos_token_id: u32, + /// Whether decode steps replay pre-captured CUDA Graphs that record NCCL + /// collectives (`enable_cuda_graph` AND a compiled TP-local decode GQA + /// group; see the P2c gate in `tp-design.md`). + graph_enabled: bool, + /// Slot tracker for the convenience `execute_prefill`/`execute_decode`/ + /// `drop_request` API (model-local tests), mirroring the single-GPU + /// `Qwen35Executor`: prefill completion appends, decode plans must cover + /// every tracked request in slot order, drop swap-removes and derives the + /// slot compaction. Scheduler-driven flows bypass it entirely — they pass + /// explicit slots via `execute_decode_items` and + /// `drop_request_with_compaction`. Never mix the two flows on one executor. + active_slots: Mutex>, } #[derive(Clone)] @@ -246,6 +309,11 @@ pub(crate) struct TpDecodeStepItem { token_id: u32, logprobs: usize, sampling_params: SamplingParams, + /// Scheduler-assigned decode slot under CUDA Graph TP. Rows must arrive in + /// dense slot order (`slot_idx == row`); on the request's first decode row + /// the worker D2D-copies its prefill recurrent state into the slot and + /// drops the per-request allocation. `None` on the slot-free eager path. + slot_idx: Option, } impl TpDecodeStepItem { @@ -260,6 +328,20 @@ impl TpDecodeStepItem { token_id, logprobs, sampling_params, + slot_idx: None, + } + } + + pub(crate) fn new_with_slot( + request_id: RequestId, + token_id: u32, + logprobs: usize, + sampling_params: SamplingParams, + slot_idx: usize, + ) -> Self { + Self { + slot_idx: Some(slot_idx), + ..Self::new(request_id, token_id, logprobs, sampling_params) } } } @@ -307,10 +389,6 @@ impl Qwen35TpExecutor { "Qwen3.5 TP executor requires at least two CUDA devices, got {}", device_ordinals.len() ); - anyhow::ensure!( - !enable_cuda_graph, - "Qwen3.5 TP Phase 1 supports eager execution only; disable CUDA Graph" - ); anyhow::ensure!( max_prefill_tokens > 0, "Qwen3.5 TP max_prefill_tokens must be positive" @@ -322,7 +400,7 @@ impl Qwen35TpExecutor { models.push(Qwen35Model::from_safetensors_with_runtime( model_path, ModelRuntimeConfig { - enable_cuda_graph: false, + enable_cuda_graph, tensor_parallel: Some(TensorParallelConfig::try_from((rank, world_size))?), device_ordinal, }, @@ -331,6 +409,23 @@ impl Qwen35TpExecutor { let first = models .first() .ok_or_else(|| anyhow::anyhow!("Qwen3.5 TP executor loaded no models"))?; + // P2c gate: graph decode under TP requires a compiled batch-decode + // kernel for the rank-local GQA group. The group ratio is TP-invariant, + // so every rank decides identically; an uncompiled group (e.g. 27B's + // group 6) keeps the batched eager path byte-for-byte. + let geometry = first.geometry; + let graph_enabled = enable_cuda_graph && geometry.local_decode_group_is_compiled(); + if enable_cuda_graph && !graph_enabled { + static LOG_GRAPH_GATE: std::sync::Once = std::sync::Once::new(); + LOG_GRAPH_GATE.call_once(|| { + log::info!( + "Qwen3.5 TP decode GQA group {} ({} q heads / {} kv heads per rank) has no compiled batch-decode kernel; CUDA Graph requested but decode stays on the batched eager path", + geometry.local_num_attention_heads() / geometry.local_num_key_value_heads(), + geometry.local_num_attention_heads(), + geometry.local_num_key_value_heads(), + ); + }); + } let page_size = first.kv_pool().layout().page_size; let mut min_capacity_pages = usize::MAX; for (rank, model) in models.iter().enumerate() { @@ -360,6 +455,7 @@ impl Qwen35TpExecutor { model, max_batch, max_prefill_tokens, + graph_enabled, nccl_id, Arc::clone(&startup_gate), Arc::clone(&effective_max_batch), @@ -427,7 +523,7 @@ impl Qwen35TpExecutor { } disarm_nccl_startup_watchdog(watchdog_done, watchdog)?; - Ok(Self { + let executor = Self { workers, poison, world_size, @@ -436,7 +532,31 @@ impl Qwen35TpExecutor { capacity_pages_for_requests, max_position_embeddings, eos_token_id, - }) + graph_enabled, + active_slots: Mutex::new(Vec::new()), + }; + // Pre-capture every reachable decode graph now: after NCCL connect, + // exactly once, before serving. A mid-serving capture on one rank while + // a peer replays would desync the recorded collectives, so serve time + // is replay-only. + if graph_enabled { + executor.run_decode_graph_precapture_sweep()?; + log::info!( + "Qwen3.5 TP decode CUDA Graph enabled: {} bucket(s) up to batch {} captured per rank", + BATCH_BUCKETS + .iter() + .take_while(|&&b| b <= bucket_for(executor.max_batch)) + .count(), + bucket_for(executor.max_batch), + ); + } + Ok(executor) + } + + /// Whether decode replays pre-captured CUDA Graphs (P2c gate: requested + /// AND the TP-local decode GQA group has a compiled kernel). + pub fn graph_enabled(&self) -> bool { + self.graph_enabled } #[cfg(test)] @@ -464,6 +584,90 @@ impl Qwen35TpExecutor { token_id == self.eos_token_id } + /// Pre-capture every reachable decode bucket on every rank, phase-by-phase + /// and barriered by the controller so a captured collective's first launch + /// never overlaps a peer's capture (qwen3 sweep precedent). + fn run_decode_graph_precapture_sweep(&self) -> Result<()> { + // NCCL has no device timeout, so a desynced sweep wedges forever; this + // watchdog aborts on the deadline. abort() not exit() — exit's cudart + // atexit teardown takes the same wedged driver lock — and it disarms + // only on the explicit success send (drop-on-error stays armed). + let (sweep_done_tx, sweep_done_rx) = mpsc::sync_channel::<()>(1); + let deadline = Instant::now() + TP_PRECAPTURE_TIMEOUT; + let watchdog = thread::Builder::new() + .name("qwen35-tp-precapture-watchdog".into()) + .spawn(move || { + // Disarmed only by the explicit success send. A sender drop + // (error path) also returns Err here; stay armed to the + // deadline before deciding startup is wedged. + if sweep_done_rx.recv_timeout(TP_PRECAPTURE_TIMEOUT).is_ok() { + return; + } + std::thread::sleep(deadline.saturating_duration_since(Instant::now())); + eprintln!( + "Qwen3.5 TP decode graph pre-capture did not complete within {}s — NCCL wedge suspected, aborting", + TP_PRECAPTURE_TIMEOUT.as_secs() + ); + log::error!( + "Qwen3.5 TP decode graph pre-capture did not complete within {}s — NCCL wedge suspected, aborting", + TP_PRECAPTURE_TIMEOUT.as_secs() + ); + std::process::abort(); + }) + .map_err(|e| anyhow::anyhow!("failed to spawn Qwen3.5 TP pre-capture watchdog: {e}"))?; + + let started = Instant::now(); + let max_bucket = bucket_for(self.max_batch); + let sweep = (|| { + self.run_precapture_phase(PrecapturePhase::Warmup)?; + for (bucket_idx, &bucket) in BATCH_BUCKETS.iter().enumerate() { + if bucket > max_bucket { + break; + } + self.run_precapture_phase(PrecapturePhase::Capture { bucket_idx })?; + self.run_precapture_phase(PrecapturePhase::Launch { bucket_idx })?; + } + self.run_precapture_phase(PrecapturePhase::Finalize) + })(); + match sweep { + Ok(()) => { + // Disarm: only the explicit success send stops the watchdog. + sweep_done_tx + .send(()) + .map_err(|_| anyhow::anyhow!("Qwen3.5 TP pre-capture watchdog exited"))?; + watchdog + .join() + .map_err(|_| anyhow::anyhow!("Qwen3.5 TP pre-capture watchdog panicked"))?; + log::info!( + "Qwen3.5 TP decode graph pre-capture: buckets up to {max_bucket} captured per rank in {:.2}s", + started.elapsed().as_secs_f64() + ); + Ok(()) + } + // On error the watchdog stays armed: peers may be wedged in + // unpaired collectives, and the abort makes the wedge attributable. + Err(err) => Err(err), + } + } + + fn run_precapture_phase(&self, phase: PrecapturePhase) -> Result<()> { + self.poison.ensure_healthy()?; + let resp_rx = self.dispatch_mutating("decode graph precapture", |start, resp| { + TpWorkerCommand::Precapture { phase, start, resp } + })?; + let responses = recv_runtime_responses( + &resp_rx, + self.world_size, + "decode graph precapture", + &self.poison, + )?; + validate_dispatched_responses( + validate_ack_responses(responses, self.world_size, "decode graph precapture"), + "decode graph precapture", + &self.poison, + ) + } + #[cfg(test)] fn ping_all(&self) -> Result<()> { self.poison.ensure_healthy()?; @@ -496,7 +700,20 @@ impl Qwen35TpExecutor { .cloned() .map(TpPrefillChunkItem::from) .collect(); - self.execute_prefill_chunks(&chunks) + let result = self.execute_prefill_chunks(&chunks)?; + if self.graph_enabled { + // Convenience-API slot tracking: every prefill plan item finishes + // prefill (TpPrefillChunkItem::from sets finish_prefill), so each + // request takes the next dense decode slot. + let mut active = self + .active_slots + .lock() + .unwrap_or_else(PoisonError::into_inner); + for chunk in &chunks { + active.push(chunk.request_id); + } + } + Ok(result) } fn execute_prefill_chunks(&self, chunks: &[TpPrefillChunkItem]) -> Result { @@ -537,18 +754,49 @@ impl Qwen35TpExecutor { !plan.requests.is_empty(), "Qwen3.5 TP decode plan requires at least one request" ); - let requests: Vec = plan - .requests - .iter() - .map(|request| { - TpDecodeStepItem::new( - request.request_id, - request.token_id, - request.logprobs, - SamplingParams::default(), - ) - }) - .collect(); + let requests: Vec = if self.graph_enabled { + let active = self + .active_slots + .lock() + .unwrap_or_else(PoisonError::into_inner); + anyhow::ensure!( + plan.requests.len() == active.len(), + "Qwen3.5 TP graph decode must cover all {} active requests in slot order, got {}", + active.len(), + plan.requests.len() + ); + plan.requests + .iter() + .enumerate() + .map(|(slot, request)| { + anyhow::ensure!( + active[slot] == request.request_id, + "Qwen3.5 TP graph decode slot {slot} holds request {} but the plan carries {}", + active[slot].get(), + request.request_id.get() + ); + Ok(TpDecodeStepItem::new_with_slot( + request.request_id, + request.token_id, + request.logprobs, + SamplingParams::default(), + slot, + )) + }) + .collect::>()? + } else { + plan.requests + .iter() + .map(|request| { + TpDecodeStepItem::new( + request.request_id, + request.token_id, + request.logprobs, + SamplingParams::default(), + ) + }) + .collect() + }; self.execute_decode_items(&requests, 0) } @@ -612,10 +860,24 @@ impl Qwen35TpExecutor { } pub fn drop_request(&self, request_id: RequestId, expectation: DropExpectation) -> Result<()> { + let compaction = self.track_retired_slot(request_id); + self.drop_request_with_compaction(request_id, expectation, compaction) + } + + /// Retire a request, attaching the slot compaction the caller (scheduler) + /// already applied to its own dense-slot bookkeeping. Workers apply the + /// move and poison on occupancy mismatch; eager workers ignore it. + pub(crate) fn drop_request_with_compaction( + &self, + request_id: RequestId, + expectation: DropExpectation, + compaction: Option, + ) -> Result<()> { self.poison.ensure_healthy()?; let resp_rx = self.dispatch_mutating("drop request", |start, resp| TpWorkerCommand::DropRequest { request_id, + compaction, start, resp, })?; @@ -628,6 +890,29 @@ impl Qwen35TpExecutor { ) } + /// Convenience-API tracker: swap-remove the retired request and derive the + /// slot compaction (last occupied slot moves into the vacated one). + /// Returns `None` on the eager path and for untracked requests. + fn track_retired_slot(&self, request_id: RequestId) -> Option { + if !self.graph_enabled { + return None; + } + let mut active = self + .active_slots + .lock() + .unwrap_or_else(PoisonError::into_inner); + let idx = active.iter().position(|&id| id == request_id)?; + let last = active.len() - 1; + active.swap_remove(idx); + // `then`, not `then_some`: the moved request only exists when the + // retired request was not the tail slot. + (idx < active.len()).then(|| TpSlotCompaction { + moved_request_id: active[idx], + from: last, + to: idx, + }) + } + #[cfg(test)] fn snapshot_workers(&self) -> Result> { self.poison.ensure_healthy()?; @@ -881,6 +1166,7 @@ impl TpStartupGate { } impl TpWorker { + #[allow(clippy::too_many_arguments)] #[allow(clippy::type_complexity)] fn spawn( rank: usize, @@ -888,6 +1174,7 @@ impl TpWorker { 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, @@ -912,6 +1199,7 @@ impl TpWorker { model, max_batch, max_prefill_tokens, + graph_enabled, ); let prepared = match prepared { Ok((prepared, rank_max_batch)) => { @@ -927,7 +1215,7 @@ impl TpWorker { return; } let max_batch = effective_max_batch.load(Ordering::Acquire); - match prepared.connect(nccl_id, max_batch, poison) { + match prepared.connect(nccl_id, max_batch, graph_enabled, poison) { Ok(mut state) => { let _ = startup_tx.send(Ok(())); state.run(rx); @@ -985,8 +1273,17 @@ 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, sample_scratch: pegainfer_sample::SampleScratch, _cublas_guard: CublasThreadGuard, @@ -1007,7 +1304,10 @@ struct TpRequestState { request_id: RequestId, phase: TpRequestPhase, kv: KvState, - recurrent: RecurrentState, + /// 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)] @@ -1031,6 +1331,7 @@ impl TpWorkerPrepared { 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 @@ -1047,9 +1348,17 @@ impl TpWorkerPrepared { 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. + let graph_slot_reserve = if graph_enabled { + bucket_for(requested_max_batch) * recurrent_bytes + } else { + 0 + }; let max_batch = effective_recurrent_capacity( requested_max_batch, - free_bytes, + free_bytes.saturating_sub(graph_slot_reserve), recurrent_bytes, TP_RUNTIME_MEMORY_RESERVE_BYTES, prefill_scratch_bytes, @@ -1096,6 +1405,7 @@ impl TpWorkerPrepared { self, nccl_id: cudarc::nccl::safe::Id, effective_max_batch: usize, + graph_enabled: bool, poison: Arc, ) -> Result { let Self { @@ -1119,12 +1429,25 @@ impl TpWorkerPrepared { ) .map_err(|e| anyhow::anyhow!("failed to initialize Qwen3.5 TP NCCL rank {rank}: {e:?}"))?; model.attach_tp_comm(comm); + 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, sample_scratch, _cublas_guard: cublas_guard, @@ -1199,14 +1522,25 @@ impl TpWorkerState { } TpWorkerCommand::DropRequest { request_id, + compaction, start, resp, } => { if start.wait() == TpCommandDecision::Cancel { false } else { - let existed = self.drop_request(request_id); - self.respond(resp, "drop request", Ok(TpWorkerReply::DropAck { existed })) + 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)] @@ -1228,7 +1562,7 @@ impl TpWorkerState { } #[cfg(test)] TpWorkerCommand::RemoveRequestStateForTest { request_id, resp } => { - let _ = resp.send(self.drop_request(request_id)); + let _ = resp.send(self.drop_request(request_id, None).unwrap_or(false)); false } #[cfg(test)] @@ -1318,7 +1652,12 @@ impl TpWorkerState { ); let prompt = [chunk.prompt_tokens.as_slice()]; - let mut recurrent_refs = vec![&mut state.recurrent]; + 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), @@ -1362,6 +1701,9 @@ impl TpWorkerState { 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 @@ -1389,7 +1731,12 @@ impl TpWorkerState { for (state_idx, state) in self.requests.iter_mut().enumerate() { if let Some(row) = row_of_state[state_idx] { kv_slots[row] = Some(&mut state.kv); - recurrent_slots[row] = Some(&mut state.recurrent); + recurrent_slots[row] = Some( + state + .recurrent + .as_mut() + .expect("eager TP decode request owns its recurrent state"), + ); } } let mut kv_refs: Vec<&mut KvState> = Vec::with_capacity(bs); @@ -1469,6 +1816,144 @@ impl TpWorkerState { .collect()) } + /// 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_slots: Vec> = + std::iter::repeat_with(|| None).take(bs).collect(); + for (state_idx, state) in self.requests.iter_mut().enumerate() { + if let Some(row) = row_of_state[state_idx] { + kv_slots[row] = Some(&mut state.kv); + } + } + let mut kv_refs: Vec<&mut KvState> = Vec::with_capacity(bs); + for kv in kv_slots { + kv_refs.push(kv.expect("decode row state resolved above")); + } + 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()); + } + + // Snapshot requested logits rows BEFORE sampling: the sampler may + // modify bufs.logits in place. + let requested_logprobs: Vec = + requests.iter().map(|request| request.logprobs).collect(); + let cpu_logits = + snapshot_requested_logprobs(ctx, &graph_state.buffers.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, + &graph_state.buffers.logits, + ¶ms_refs, + &steps, + sample_seed, + &mut self.sample_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| { + pegainfer_sample::token_logprob_from_row( + logits_row, + tokens[row], + request.logprobs, + ) + }); + DecodeRequestResult { + request_id: request.request_id, + token: tokens[row], + logprob, + } + }) + .collect()) + } + fn sample_final_prefill_chunk( &mut self, chunk: &TpPrefillChunkItem, @@ -1569,7 +2054,7 @@ impl TpWorkerState { request_id, phase: TpRequestPhase::Prefilling, kv: self.model.alloc_kv(), - recurrent, + recurrent: Some(recurrent), }; self.requests.push(state); Ok(self.requests.len() - 1) @@ -1581,14 +2066,167 @@ impl TpWorkerState { .position(|state| state.request_id == request_id) } - fn drop_request(&mut self, request_id: RequestId) -> bool { - if let Some(idx) = self.request_index(request_id) { - self.requests.swap_remove(idx); - true - } else { - false + /// 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: token 0 at position 0 + /// over freshly allocated one-page KV states. Outputs are discarded; the + /// rows exist only to give the recorded kernels valid addresses. + 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: Vec = (0..bucket).map(|_| self.model.alloc_kv()).collect(); + let mut kv_refs: Vec<&mut KvState> = synthetic_kv.iter_mut().collect(); + let token_ids = vec![0u32; bucket]; + self.model + .batch_decode_graph(&token_ids, &mut kv_refs, graph_state, graph_use)?; + // 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<()> { @@ -1726,17 +2364,6 @@ impl From for TpPrefillChunkItem { } } -impl From for TpDecodeStepItem { - fn from(request: DecodeStepItem) -> Self { - Self::new( - request.request_id, - request.token_id, - request.logprobs, - SamplingParams::default(), - ) - } -} - fn recv_runtime_responses( responses: &mpsc::Receiver, expected: usize, @@ -1817,7 +2444,6 @@ fn validate_exact_rank_responses( Ok(replies) } -#[cfg(test)] fn validate_ack_responses( responses: Vec, world_size: usize, @@ -2222,12 +2848,122 @@ mod tests { } #[test] - fn rejects_tensor_parallel_cuda_graph() { + fn tensor_parallel_cuda_graph_gate_defers_to_model_load() { + // P2c: TP + CUDA Graph is no longer rejected up front; the graph/eager + // decision needs the model config, so a nonexistent path fails at load. let err = match Qwen35TpExecutor::from_runtime_with_capacity("unused", true, &[0, 1], 1) { - Ok(_) => panic!("TP CUDA Graph should fail"), + Ok(_) => panic!("TP CUDA Graph with a nonexistent model path should fail at load"), Err(err) => err.to_string(), }; - assert!(err.contains("eager execution only")); + assert!(!err.contains("eager execution only")); + } + + #[test] + fn slot_map_admit_release_and_compact() { + let id = |value: u64| RequestId::new(value); + let mut owners = vec![None, None, None, None]; + + slot_admit(&mut owners, 0, id(1)).unwrap(); + slot_admit(&mut owners, 1, id(2)).unwrap(); + slot_admit(&mut owners, 2, id(3)).unwrap(); + + let err = slot_admit(&mut owners, 1, id(9)).unwrap_err().to_string(); + assert!(err.contains("still owned by request 2")); + + // Retire slot 1: last occupied slot (2, request 3) moves into it. + let needs_move = slot_compact( + &mut owners, + id(2), + TpSlotCompaction { + moved_request_id: id(3), + from: 2, + to: 1, + }, + ) + .unwrap(); + assert!(needs_move, "materialized moved request needs the GPU move"); + assert_eq!(owners, vec![Some(id(1)), Some(id(3)), None, None]); + + // Retire the tail slot: release without compaction. + assert_eq!(slot_release(&mut owners, id(3)), Some(1)); + assert_eq!(owners, vec![Some(id(1)), None, None, None]); + + // Releasing a request that never materialized a slot is not an error. + assert_eq!(slot_release(&mut owners, id(77)), None); + } + + #[test] + fn slot_map_compact_tolerates_unmaterialized_requests() { + let id = |value: u64| RequestId::new(value); + let mut owners = vec![None, None, None]; + + // Dropped request materialized, moved request not yet admitted to its + // slot (retired between promotion and its first decode row): clear + // only, no GPU move. + slot_admit(&mut owners, 0, id(1)).unwrap(); + let needs_move = slot_compact( + &mut owners, + id(1), + TpSlotCompaction { + moved_request_id: id(2), + from: 2, + to: 0, + }, + ) + .unwrap(); + assert!(!needs_move); + assert_eq!(owners, vec![None, None, None]); + + // Moved request materialized, dropped request not: the move is needed + // and the moved request takes over the vacated slot. + slot_admit(&mut owners, 2, id(3)).unwrap(); + let needs_move = slot_compact( + &mut owners, + id(4), + TpSlotCompaction { + moved_request_id: id(3), + from: 2, + to: 0, + }, + ) + .unwrap(); + assert!(needs_move); + assert_eq!(owners, vec![Some(id(3)), None, None]); + } + + #[test] + fn slot_map_compact_poisons_on_occupancy_mismatch() { + let id = |value: u64| RequestId::new(value); + let mut owners = vec![Some(id(1)), Some(id(2))]; + + let err = slot_compact( + &mut owners, + id(9), + TpSlotCompaction { + moved_request_id: id(2), + from: 1, + to: 0, + }, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("slot 0 holds request 1")); + + let err = slot_compact( + &mut owners, + id(1), + TpSlotCompaction { + moved_request_id: id(9), + from: 1, + to: 0, + }, + ) + .unwrap_err() + .to_string(); + assert!(err.contains("slot 1 holds request 2")); + + let err = slot_admit(&mut owners, 5, id(1)).unwrap_err().to_string(); + assert!(err.contains("exceeds worker slot map")); } #[test] diff --git a/pegainfer-qwen35/src/unified_forward.rs b/pegainfer-qwen35/src/unified_forward.rs index 32f216672..e4dfc4560 100644 --- a/pegainfer-qwen35/src/unified_forward.rs +++ b/pegainfer-qwen35/src/unified_forward.rs @@ -21,6 +21,7 @@ use pegainfer_kernels::tensor::StreamOverrideGuard; use super::batch_decode_graph::BatchDecodeGraphState; use super::recurrent_state::RecurrentState; use super::weights::Qwen35Model; +use crate::batch_decode::DecodeGraphUse; pub(crate) struct UnifiedStepOutput { pub(crate) prefill_logits: Option, @@ -125,7 +126,12 @@ impl Qwen35Model { let decoded = if decode_tokens.is_empty() { false } else { - self.batch_decode_graph(decode_tokens, decode_kv_states, graph_state)?; + self.batch_decode_graph( + decode_tokens, + decode_kv_states, + graph_state, + DecodeGraphUse::Serve, + )?; true }; @@ -199,7 +205,7 @@ mod tests { for _ in 1..num_steps { let tids = [*tokens_a.last().unwrap(), *tokens_b.last().unwrap()]; model - .batch_decode_graph(&tids, &mut kv_refs, &mut gs) + .batch_decode_graph(&tids, &mut kv_refs, &mut gs, DecodeGraphUse::Serve) .unwrap(); let next = greedy_sample_batch(&model, &gs.buffers.logits, 2); tokens_a.push(next[0]); diff --git a/pegainfer-qwen35/src/weights.rs b/pegainfer-qwen35/src/weights.rs index c4a1d5d54..dd8e40f8a 100644 --- a/pegainfer-qwen35/src/weights.rs +++ b/pegainfer-qwen35/src/weights.rs @@ -218,8 +218,8 @@ impl Qwen35Model { let mut config = Config35::from_file(model_path)?; let tensor_parallel = runtime.tensor_parallel.unwrap_or_default(); - let geometry = LocalGeometry::try_new(&config, tensor_parallel, runtime.enable_cuda_graph) - .map_err(anyhow::Error::from)?; + let geometry = + LocalGeometry::try_new(&config, tensor_parallel).map_err(anyhow::Error::from)?; debug!( "Config: hidden_size={}, num_layers={}, full_attn={}, linear_attn={}, max_position_embeddings={}, tp_rank={}, tp_world_size={}", config.hidden_size, @@ -649,6 +649,34 @@ impl Qwen35Model { self.tp_comm = Some(comm); } + /// Force NCCL connect before any CUDA Graph capture records a collective + /// (lazy connect inside `cuStreamBeginCapture` wedges the capture). NCCL + /// 2.22+ connects per size-selected algorithm, so warm one all-reduce at + /// every decode bucket's message size. No-op without a TP communicator. + pub(crate) fn warmup_tp_collective(&self) -> Result<()> { + if let Some(comm) = &self.tp_comm { + let buckets = super::batch_decode_graph::BATCH_BUCKETS; + let max_elems = buckets.last().unwrap() * self.config.hidden_size; + let mut scratch = self + .ctx + .stream + .alloc_zeros::(max_elems) + .map_err(|e| anyhow::anyhow!("alloc NCCL warm-up scratch: {e}"))?; + for &bucket in buckets { + let mut view = scratch.slice_mut(0..bucket * self.config.hidden_size); + comm.all_reduce_in_place(&mut view, &ReduceOp::Sum) + .map_err(|e| { + anyhow::anyhow!("Qwen3.5 NCCL warm-up all-reduce failed: {e:?}") + })?; + } + self.ctx + .stream + .synchronize() + .map_err(|e| anyhow::anyhow!("Qwen3.5 NCCL warm-up sync failed: {e}"))?; + } + Ok(()) + } + pub(crate) fn all_reduce_hidden(&self, hidden: &mut HiddenStates) -> Result<()> { self.all_reduce_hidden_untraced(hidden) } @@ -999,7 +1027,7 @@ mod tests { fn test_geometry(rank: usize, world_size: usize) -> LocalGeometry { let config = test_config(); let tp = TensorParallelConfig::try_from((rank, world_size)).unwrap(); - LocalGeometry::try_new(&config, tp, false).unwrap() + LocalGeometry::try_new(&config, tp).unwrap() } #[test] @@ -1108,8 +1136,10 @@ mod tests { assert_eq!(qkv_view.shape(), [qkv_rows, cols]); let qkv_elems: Vec = qkv_view .data() - .chunks_exact(4) - .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])) + .as_chunks::<4>() + .0 + .iter() + .map(|b| f32::from_le_bytes(*b)) .collect(); // conv1d.weight fixture: [qkv * kernel_dim] flattened channels. @@ -1118,14 +1148,16 @@ mod tests { let conv_view = conv.tensor("c").unwrap(); let conv_elems: Vec = conv_view .data() - .chunks_exact(4) - .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])) + .as_chunks::<4>() + .0 + .iter() + .map(|b| f32::from_le_bytes(*b)) .collect(); let global_q = config.linear_num_key_heads * config.linear_key_head_dim; for rank in 0..2usize { let tp = TensorParallelConfig::try_from((rank, 2)).unwrap(); - let geom = LocalGeometry::try_new(&config, tp, false).unwrap(); + let geom = LocalGeometry::try_new(&config, tp).unwrap(); let segments = linear_qkv_shard_segments(&config, tp); // Stitched matrix = per-segment head-local row slices, in storage diff --git a/pegainfer-qwen35/tests/e2e_scheduler.rs b/pegainfer-qwen35/tests/e2e_scheduler.rs index 9886399bb..9933d2e22 100644 --- a/pegainfer-qwen35/tests/e2e_scheduler.rs +++ b/pegainfer-qwen35/tests/e2e_scheduler.rs @@ -829,7 +829,6 @@ fn test_e2e_qwen35_scheduler_tp2() { info!("Loading Qwen3.5 TP2 model for scheduler test..."); let start = Instant::now(); let tokenizer = common::load_tokenizer(&model_path); - // TP Phase 1 is eager-only; CUDA Graph must stay disabled for multi-device startup. let handle = pegainfer_qwen35::start_engine_with_capacity( Path::new(&model_path), EngineLoadOptions { @@ -847,3 +846,34 @@ fn test_e2e_qwen35_scheduler_tp2() { let max_context_tokens = max_position_embeddings(&model_path); run_full_scheduler_e2e(&handle, &tokenizer, max_context_tokens, "TP2"); } + +#[test] +#[ignore = "requires two CUDA devices, NCCL, and Qwen3.5 weights"] +fn test_e2e_qwen35_scheduler_tp2_graph() { + let Some(model_path) = common::model_path_or_skip("test_e2e_qwen35_scheduler_tp2_graph") else { + return; + }; + + info!("Loading Qwen3.5 TP2 model with CUDA Graph for scheduler test..."); + let start = Instant::now(); + let tokenizer = common::load_tokenizer(&model_path); + // P2c: decode replays pre-captured CUDA Graphs when the TP-local decode + // GQA group has a compiled kernel (4B/9B); uncompiled groups (27B group 6) + // keep the batched eager path under the same request flow. + let handle = pegainfer_qwen35::start_engine_with_capacity( + Path::new(&model_path), + EngineLoadOptions { + enable_cuda_graph: true, + device_ordinals: common::tp2_device_ordinals(), + seed: 42, + ..EngineLoadOptions::default() + }, + 8, + pegainfer_qwen35::DEFAULT_MAX_PREFILL_TOKENS, + ) + .expect("Failed to start Qwen3.5 TP2 graph scheduler"); + info!("TP2 graph scheduler loaded in {:.2?}", start.elapsed()); + + let max_context_tokens = max_position_embeddings(&model_path); + run_full_scheduler_e2e(&handle, &tokenizer, max_context_tokens, "TP2 graph"); +} diff --git a/pegainfer-qwen35/tests/hf_golden_gate.rs b/pegainfer-qwen35/tests/hf_golden_gate.rs index 558866a6f..a2e3220b5 100644 --- a/pegainfer-qwen35/tests/hf_golden_gate.rs +++ b/pegainfer-qwen35/tests/hf_golden_gate.rs @@ -747,6 +747,117 @@ fn build_tp2_executor(model_path: &str) -> Qwen35TpExecutor { .expect("build Qwen3.5 TP2 logits executor") } +/// TP2 executor with CUDA Graph requested. Returns `None` when the loaded +/// model's TP-local decode GQA group has no compiled kernel (27B group 6) — +/// the P2c gate keeps that path eager, so there is no graph to gate on. +fn build_tp2_graph_executor(model_path: &str, label: &str) -> Option { + let devices = common::tp2_device_ordinals(); + let ex = Qwen35TpExecutor::from_runtime_with_capacity( + model_path, + true, + &devices, + MAX_EXECUTOR_BATCH, + ) + .expect("build Qwen3.5 TP2 graph logits executor"); + if !ex.graph_enabled() { + eprintln!( + "qwen35 hf_golden_gate [{label}]: CUDA Graph gated off (uncompiled TP-local decode GQA group); skipping" + ); + return None; + } + Some(ex) +} + +/// Mid-batch drop with slot compaction under TP: prefill `seqs`, decode one +/// step, retire `SLOT_COMPACTION_DROP_INDEX` (the last slot moves into the +/// gap), then keep decoding the survivors in their new dense slot order. +fn run_tp_with_slot_compaction( + g: &Golden, + ex: &Qwen35TpExecutor, + seqs: &[usize], +) -> (Stats, Vec) { + assert!( + seqs.len() > SLOT_COMPACTION_DROP_INDEX + 1, + "TP slot-compaction replay needs a non-tail request to drop" + ); + assert!( + g.decode_len >= 2, + "TP slot-compaction replay needs at least two decode tokens" + ); + + let mut stats = Stats::default(); + let mut fingerprint = Vec::new(); + let mut fold = |stats: &mut Stats, seq, pos, pega: &[(u32, f32)]| { + fingerprint.push(pega[0].1); + check_position(stats, seq, pos, pega, &g.topk(seq, pos)); + }; + + let mut live: Vec<(usize, RequestId)> = seqs + .iter() + .map(|&seq| (seq, RequestId::new(30_000 + seq as u64))) + .collect(); + let items: Vec = live + .iter() + .map(|&(seq, id)| prefill_item(id, g.prompt(seq))) + .collect(); + let pr = ex + .execute_prefill(PrefillPlan { requests: &items }) + .expect("TP2 prefill"); + for (i, &(seq, _)) in live.iter().enumerate() { + fold( + &mut stats, + seq, + 0, + &top_logprobs(pr.requests[i].first_token_logprob.as_ref()), + ); + } + + let step0: Vec = live + .iter() + .map(|&(seq, id)| decode_item(id, g.decode(seq, 0))) + .collect(); + let dr = ex + .execute_decode(DecodePlan { requests: &step0 }) + .expect("TP2 decode before compaction"); + for (i, &(seq, _)) in live.iter().enumerate() { + fold( + &mut stats, + seq, + 1, + &top_logprobs(dr.requests[i].logprob.as_ref()), + ); + } + + let (_, dropped_id) = live[SLOT_COMPACTION_DROP_INDEX]; + ex.drop_request(dropped_id, DropExpectation::MustExist) + .expect("TP2 drop request"); + live.swap_remove(SLOT_COMPACTION_DROP_INDEX); + + for step in 1..g.decode_len { + let items: Vec = live + .iter() + .map(|&(seq, id)| decode_item(id, g.decode(seq, step))) + .collect(); + let dr = ex + .execute_decode(DecodePlan { requests: &items }) + .expect("TP2 decode after compaction"); + for (i, &(seq, _)) in live.iter().enumerate() { + fold( + &mut stats, + seq, + step + 1, + &top_logprobs(dr.requests[i].logprob.as_ref()), + ); + } + } + + for (_, id) in live { + ex.drop_request(id, DropExpectation::MustExist) + .expect("TP2 drop request"); + } + (stats, fingerprint) +} + #[test] fn pega_logprobs_match_hf_golden_within_qwen35_tolerance() { let Some(model_path) = common::model_path_or_skip("pega_logprobs_match_hf_golden") else { @@ -888,3 +999,62 @@ fn pega_logprobs_match_hf_long_golden_within_qwen35_tolerance_tp2() { "TP2 long sequential Qwen3.5 replay must reproduce identical logprobs" ); } + +/// P2c TP2 CUDA Graph gate: sequential replay, bucket-straddling batched +/// replay, and post-compaction replay after a mid-batch drop, all compared +/// against the same HF golden within the existing TP2 tolerances. +#[test] +#[ignore = "requires two CUDA devices, NCCL, and Qwen3.5 weights"] +fn pega_logprobs_match_hf_golden_within_qwen35_tolerance_tp2_graph() { + let Some(model_path) = common::model_path_or_skip("pega_logprobs_match_hf_golden_tp2_graph") + else { + return; + }; + let Some(golden) = Golden::load_for(&model_path, false) else { + return; + }; + if !check_fixture_metadata(&model_path, &golden) { + return; + } + report_fixture_shape(&golden); + let all: Vec = (0..golden.num_seqs).collect(); + + let Some(ex) = build_tp2_graph_executor(&model_path, "TP2 graph") else { + return; + }; + let (stats, fp1) = run_tp(&golden, &ex, &all, false); + report_and_assert("TP2 sequential graph", &stats); + let (_, fp2) = run_tp(&golden, &ex, &all, false); + assert_eq!( + fp1, fp2, + "TP2 sequential Qwen3.5 graph replay must reproduce identical logprobs" + ); + + for n in BUCKET_STRADDLES { + if all.len() >= n { + let (batched, _) = run_tp(&golden, &ex, &all[..n], true); + report_and_assert(&format!("TP2 batched graph ({n} padded)"), &batched); + } else { + eprintln!( + "qwen35 hf_golden_gate: skipping TP2 batched graph ({n} padded); fixture has only {} sequence(s)", + all.len() + ); + } + } + + if golden.num_seqs >= SLOT_COMPACTION_BATCH && golden.decode_len >= 2 { + let (compacted, fp1) = + run_tp_with_slot_compaction(&golden, &ex, &all[..SLOT_COMPACTION_BATCH]); + report_and_assert("TP2 slot-compaction graph", &compacted); + let (_, fp2) = run_tp_with_slot_compaction(&golden, &ex, &all[..SLOT_COMPACTION_BATCH]); + assert_eq!( + fp1, fp2, + "TP2 slot-compaction Qwen3.5 graph replay must reproduce identical logprobs" + ); + } else { + eprintln!( + "qwen35 hf_golden_gate: skipping TP2 slot-compaction graph; fixture has {} sequence(s), decode_len {}", + golden.num_seqs, golden.decode_len + ); + } +} diff --git a/pegainfer-qwen35/tests/serving_tp2.rs b/pegainfer-qwen35/tests/serving_tp2.rs index 2ed223e49..7340e0d79 100644 --- a/pegainfer-qwen35/tests/serving_tp2.rs +++ b/pegainfer-qwen35/tests/serving_tp2.rs @@ -1,5 +1,4 @@ use std::net::TcpListener; -use std::path::Path; use std::path::PathBuf; use std::time::Duration; @@ -51,7 +50,10 @@ async fn qwen35_tp2_serves_openai_completions_over_http() -> Result<()> { return Ok(()); }; let frontend_model_path = PathBuf::from(frontend_model_path); - let invalid_graph_model_path = engine_model_path.clone(); + // P2c graph acceptance smoke: CUDA Graph requested at TP2. Models with a + // compiled TP-local decode GQA group (4B/9B) replay pre-captured decode + // graphs; uncompiled groups (27B group 6) stay on the batched eager path + // under the same serving flow. let server = spawn_ready_server(engine_model_path, frontend_model_path, 1).await?; let client = test_client()?; @@ -59,11 +61,6 @@ async fn qwen35_tp2_serves_openai_completions_over_http() -> Result<()> { assert_non_streaming_completion(&client, &server.base_url).await?; assert_streaming_completion(&client, &server.base_url).await?; assert_concurrent_completions(&client, &server.base_url).await?; - assert_invalid_cuda_graph_tp_startup_fails( - invalid_graph_model_path - .to_str() - .context("Qwen3.5 engine fixture path is not valid UTF-8")?, - )?; server.shutdown().await } @@ -78,7 +75,7 @@ async fn spawn_ready_server( pegainfer_qwen35::start_engine_with_capacity( &engine_model_path, EngineLoadOptions { - enable_cuda_graph: false, + enable_cuda_graph: true, device_ordinals, seed: 42, ..EngineLoadOptions::default() @@ -199,27 +196,6 @@ async fn assert_concurrent_completions(client: &Client, base_url: &str) -> Resul Ok(()) } -fn assert_invalid_cuda_graph_tp_startup_fails(model_path: &str) -> Result<()> { - let Err(error) = pegainfer_qwen35::start_engine_with_capacity( - Path::new(model_path), - EngineLoadOptions { - enable_cuda_graph: true, - device_ordinals: common::tp2_device_ordinals(), - seed: 42, - ..EngineLoadOptions::default() - }, - 8, - 1, - ) else { - bail!("TP2 + CUDA Graph must fail before serving requests"); - }; - let message = error.to_string(); - if !message.contains("eager execution only") { - bail!("unexpected TP2 + CUDA Graph startup error: {message}"); - } - Ok(()) -} - async fn post_completion(client: &Client, base_url: &str, body: Value) -> Result { client .post(format!("{base_url}/v1/completions")) From 35880118483634f05b0d442a85f042cb7422b31b Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Sun, 30 Aug 2026 13:33:09 +0000 Subject: [PATCH 4/4] fix(qwen35): size TP graph slot reserve by effective batch bucket The graph-slot pre-reserve used bucket_for(requested_max_batch), but the graph state is later allocated as bucket_for(effective_max_batch). On a tight-memory rank the oversized reserve could starve the effective recurrent capacity to zero and refuse startup. Iterate the reserve bucket downward against the fitted capacity until it stabilises (the bucket only shrinks, so it converges), and clamp the fitted batch to the reserved bucket so the later bucket_for(effective) allocation never exceeds the reserve. Absorbs the codex review comment on Ma1oneZhang/pegainfer PR #946 (tp_executor.rs graph_slot_reserve). Signed-off-by: Ziyang Zhang --- pegainfer-qwen35/src/tp_executor.rs | 43 +++++++++++++++++++++-------- 1 file changed, 32 insertions(+), 11 deletions(-) diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index a3c2fee3a..ed9376048 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -1350,19 +1350,40 @@ impl TpWorkerPrepared { ); // 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. - let graph_slot_reserve = if graph_enabled { - bucket_for(requested_max_batch) * recurrent_bytes + // (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 { - 0 + effective_recurrent_capacity( + requested_max_batch, + free_bytes, + recurrent_bytes, + TP_RUNTIME_MEMORY_RESERVE_BYTES, + prefill_scratch_bytes, + ) }; - let max_batch = effective_recurrent_capacity( - requested_max_batch, - free_bytes.saturating_sub(graph_slot_reserve), - 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",