From d9cd1e06667306c7bb76fd6fee115bc29e83a78b Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Thu, 20 Aug 2026 15:06:57 +0000 Subject: [PATCH 01/15] 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 ba12041c39c06b5a50805513e189688987b68131 Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Sun, 30 Aug 2026 14:12:34 +0000 Subject: [PATCH 02/15] fix(qwen35): allow the f32 shard loader's pointer alignment cast The TP rank-sliced 1D f32 loader casts the safetensors byte span to f32 exactly like the whole-tensor loader next to it (which already carries the allow); the missed attribute trips clippy::cast-ptr-alignment under the workspace's -D warnings gates. Signed-off-by: Ziyang Zhang --- pegainfer-core/src/weight_loader.rs | 1 + 1 file changed, 1 insertion(+) 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, From fcdeb5a4a75fe81f01dc16c2af520a10670a61ac Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Sun, 30 Aug 2026 14:20:12 +0000 Subject: [PATCH 03/15] style(qwen35): swap the negated decode plan branch and parse fixture f32s via as_chunks if_not_else and chunks_exact_to_as_chunks (pedantic/default) break the Qwen3.5 clippy gate under -D warnings. Signed-off-by: Ziyang Zhang --- pegainfer-qwen35/src/batch_decode.rs | 6 +++--- pegainfer-qwen35/src/weights.rs | 12 ++++++++---- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/pegainfer-qwen35/src/batch_decode.rs b/pegainfer-qwen35/src/batch_decode.rs index 0c80feaae..cacfe97a8 100644 --- a/pegainfer-qwen35/src/batch_decode.rs +++ b/pegainfer-qwen35/src/batch_decode.rs @@ -296,7 +296,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 +328,6 @@ impl Qwen35Model { ) })?, ) - } else { - None }; let kv_buffer = kv_states[0].buffer(); diff --git a/pegainfer-qwen35/src/weights.rs b/pegainfer-qwen35/src/weights.rs index c4a1d5d54..8d70df598 100644 --- a/pegainfer-qwen35/src/weights.rs +++ b/pegainfer-qwen35/src/weights.rs @@ -1108,8 +1108,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,8 +1120,10 @@ 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; From 44e6dc0d445a4659e08ffed3cd0e497b202501f2 Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Thu, 3 Sep 2026 06:55:57 +0000 Subject: [PATCH 04/15] docs(qwen35): record landed P2B GDR sharding in the TP implementation record Signed-off-by: Ziyang Zhang --- docs/index.md | 2 +- docs/models/qwen35/tp-implementation.md | 37 +++++++++++++++---------- 2 files changed, 24 insertions(+), 15 deletions(-) diff --git a/docs/index.md b/docs/index.md index 63a5950b1..2aae4671f 100644 --- a/docs/index.md +++ b/docs/index.md @@ -54,7 +54,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `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-implementation.md` | Qwen3.5 TP Phase 1, P2A, and P2B GDR state sharding are complete: TP2 has start-gated eager unified prefill+decode, fail-closed lifecycle recovery, and rank-local linear-attention weights/state with a post-`out_proj` hidden all-reduce; batched TP decode (#1004) and TP CUDA Graph (#1005) are next. | | `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-implementation.md b/docs/models/qwen35/tp-implementation.md index d8ebda80d..671dfc888 100644 --- a/docs/models/qwen35/tp-implementation.md +++ b/docs/models/qwen35/tp-implementation.md @@ -1,8 +1,8 @@ # 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 Phase 1, P2A, and the P2B GDR state sharding are complete: TP2 has start-gated eager unified prefill+decode, fail-closed lifecycle recovery, and rank-local linear-attention weights/state with a single post-`out_proj` hidden all-reduce; batched eager TP decode and TP CUDA Graph are next. > -> **Last touched:** 2026-08 +> **Last touched:** 2026-09 ## Scope @@ -139,7 +139,7 @@ Stable test knobs: ## Phase 2 Progress -Phase 2 is locked in `docs/models/qwen35/tp-design.md` as two separate implementation series: P2a is eager mixed unified execution on the replicated Phase 1 GDR path; P2b shards the head-indexed linear-attention/GDR weight and state surface. P2a protocol/lifecycle gates are complete, so P2b can now change loader, kernel, and state shapes while preserving those contracts. +Phase 2 is locked in `docs/models/qwen35/tp-design.md` as two separate implementation series: P2a is eager mixed unified execution on the replicated Phase 1 GDR path; P2b shards the head-indexed linear-attention/GDR weight and state surface. P2a protocol/lifecycle gates are complete, and P2b's core sharding has landed on top of them without weakening the P2A lifecycle and ID contracts (see below). The remaining Phase 2 work is batched eager TP decode (#1004) and TP CUDA Graph (#1005). ### P2a: TP mixed-step unified execution @@ -445,24 +445,33 @@ Why this should be separated from GDR sharding: ### P2b: sharded linear-attention/GDR state -Shard the Qwen3.5 linear-attention/GDR path after P2a establishes the mixed-step and state-lifecycle contract. +Landed as #946 split 1/4 (#1003), after P2a had established the mixed-step and state-lifecycle contract. Each TP rank now owns a rank-local slice of the linear-attention/GDR surface instead of replicating it: -Expected work: +- `LocalGeometry` computes rank-local linear dims (`local_linear_num_key_heads`, `local_linear_num_value_heads`, `local_linear_q_dim`, `local_linear_v_dim`, `local_linear_qkv_dim`, `local_linear_z_dim`) and fails closed with `ConfigError::TpIndivisible` when `linear_num_key_heads` or `linear_num_value_heads` do not divide by `world_size`; there is no silent replication fallback. +- Weight loading shards the head-indexed tensors: the fused QKV projection and the depthwise conv1d are stitched head-locally per segment (`load_linear_in_proj_qkv_shard` / `load_linear_conv1d_shard`; Q/K segments follow key-head ranges, V/conv follow value-head ranges), z/beta/alpha are row shards, `dt_bias`/`A_log` are 1-D shards (A_log stays f32), and linear `out_proj` is column sharded as a row-parallel `[hidden, local_z]` matrix. +- `RecurrentState` (`[local_value_heads, K, V]` f32), conv state (`[local_qkv x (kernel_dim - 1)]` bf16), and all prefill/decode scratch (`GdrChunkwiseScratch35`, prefill/decode buffers) size themselves from the local geometry; worker capacity math uses the same locals. +- The hidden-residual all-reduce happens once after the local linear-attention `out_proj` (`all_reduce_hidden`), on prefill and decode alike; the column-sharded `out_proj` is what makes that reduction point sufficient. +- Full-attention decode-group supportability is evaluated on the rank-local GQA group (`LocalGeometry::local_decode_group_is_compiled`): at TP2 the 27B case leaves `q/kv = 6`, which has no compiled FlashInfer batch-decode kernel, so those layers reroute decode through the eager/paged fallback. The predicate is deterministic per model and identical on every rank, so the reroute adds no collectives. +- TP1 contract is unchanged: at `world_size == 1` every local dim equals the global dim, so kernels, buffers, and fixture behavior are byte-identical to pre-P2b. -- shard linear-attention projection weights -- shard conv state and GDR recurrent state by local value/key heads -- adapt or regenerate GDR kernels for local state shapes -- keep recurrent/conv state rank-local and request-local -- all-reduce only after local linear-attention `out_proj` -- report matched Phase 1 TP2 versus P2b TP2 HBM/latency/throughput data before making a performance claim - -Non-negotiable invariant: +Non-negotiable invariant (still held): - Never all-reduce GDR recurrent state or conv state. These states are owned by rank-local request state. +Acceptance at `fcdeb5a4` (27B TP2 on 2x RTX 4090 48GB, sm_89; fixture-pinned 27B revision `fc05daec`): + +- TP2 short HF logits gate passes: + - sequential eager: `108` positions, mean `0.0210`, p99 `0.0749`, max `0.1240` + - batched eager: `72` positions, mean `0.0201`, p99 `0.0749`, max `0.0803`; the batched leg includes drop -> re-prefill slot cycles +- TP2 long HF logits gate passes with prompts `4097` and `8192`: sequential eager, `18` positions, mean `0.0177`, p99 `0.0660` +- TP2 scheduler E2E and TP2 HTTP serving gate pass. +- Peak per-rank HBM is `35,988` / `36,822` MiB of `49,140` MiB: 27B TP2 now fits the 2x48GB pair that Phase-1 replicated state OOMed, and memory fully releases between test processes. + +Not in this step: batching the TP decode loop across rows, TP CUDA Graph capture, and the matched Phase-1-vs-P2b HBM/latency/throughput A/B promised in #1001; no performance claim is made until that rerun lands on the merged stack. + ## Follow-Ups -- Design and implement P2B sharded linear-attention/GDR state without weakening the completed P2A lifecycle and ID contracts. +- Land batched eager TP decode (#1004) on top of the P2B state sharding, then TP CUDA Graph (#1005); rerun the #946 throughput A/B on the merged stack before any performance claim. - 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. From e6629f93e3ed238ec0c071c829f2cda07e66510b Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Fri, 4 Sep 2026 05:40:43 +0000 Subject: [PATCH 05/15] fix(qwen35): validate dtype and alignment in the 1D F32 weight loaders Signed-off-by: Ziyang Zhang --- pegainfer-core/src/weight_loader.rs | 114 +++++++++++++++++++++------- 1 file changed, 88 insertions(+), 26 deletions(-) diff --git a/pegainfer-core/src/weight_loader.rs b/pegainfer-core/src/weight_loader.rs index 3261b3671..1e502a658 100644 --- a/pegainfer-core/src/weight_loader.rs +++ b/pegainfer-core/src/weight_loader.rs @@ -371,6 +371,49 @@ fn tensor_bf16_cow<'d>( } } +/// Typed F32 payload with dtype and 1D-shape validation. Aligned payloads +/// borrow zero-copy; misaligned ones (legal in safetensors) decode +/// little-endian into an owned buffer, since a misaligned f32 view is UB. +#[allow(clippy::cast_ptr_alignment)] +fn tensor_f32_cow<'d>( + tensor: &safetensors::tensor::TensorView<'d>, + name: &str, +) -> Result> { + anyhow::ensure!( + tensor.dtype() == Dtype::F32, + "Tensor '{name}': expected dtype F32, got {:?}", + tensor.dtype() + ); + anyhow::ensure!( + tensor.shape().len() == 1, + "Tensor '{name}': expected 1D shape, got {:?}", + tensor.shape() + ); + let data = tensor.data(); + anyhow::ensure!( + data.len().is_multiple_of(std::mem::size_of::()), + "Tensor '{name}': {} bytes is not a whole number of f32 elements", + data.len() + ); + if (data.as_ptr() as usize).is_multiple_of(std::mem::align_of::()) { + // SAFETY: alignment checked; any bit pattern is a valid f32. + Ok(Cow::Borrowed(unsafe { + std::slice::from_raw_parts( + data.as_ptr().cast::(), + data.len() / std::mem::size_of::(), + ) + })) + } else { + Ok(Cow::Owned( + data.as_chunks::<4>() + .0 + .iter() + .map(|&b| f32::from_le_bytes(b)) + .collect(), + )) + } +} + /// One row-consecutive part of a fused matrix: `rows` rows starting at /// `row_offset` of a source tensor that must have exactly `src_rows` rows. pub struct FusedPart<'a> { @@ -868,7 +911,6 @@ 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, @@ -879,33 +921,23 @@ pub fn load_tensor_1d_f32_shard( 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 { + let elems = tensor_f32_cow(&tensor, name)?; + if offset + len > elems.len() { return Err(anyhow::anyhow!( "F32 1D shard out of bounds for '{}': offset={} len={} total_len={}", name, offset, len, - total + elems.len() )); } - let slice = unsafe { std::slice::from_raw_parts(data.as_ptr().cast::(), total) }; let gpu_data = ctx .stream - .clone_htod(&slice[offset..offset + len]) + .clone_htod(&elems[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). pub fn load_tensor_1d_f32( @@ -915,19 +947,10 @@ pub fn load_tensor_1d_f32( name: &str, ) -> 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 len = data.len() / 4; - let slice = unsafe { std::slice::from_raw_parts(data.as_ptr().cast::(), len) }; + let elems = tensor_f32_cow(&tensor, name)?; let gpu_data = ctx .stream - .clone_htod(slice) + .clone_htod(elems.as_ref()) .map_err(|e| anyhow::anyhow!("H2D copy failed for '{}': {}", name, e))?; Ok(gpu_data) } @@ -1045,6 +1068,45 @@ mod tests { use safetensors::tensor::TensorView; use super::tensor_bf16_cow; + use super::tensor_f32_cow; + + #[test] + fn tensor_f32_cow_borrows_aligned_and_decodes_unaligned() { + let vals: [u32; 4] = [0x3f80_0000, 0x0000_0001, 0xbf12_3456, 0x7f80_0001]; + let mut bytes = vec![0u8; vals.len() * 4 + 3]; + // A Vec base has no alignment guarantee; derive both offsets from + // the actual address so each branch is forced deterministically. + let base = bytes.as_ptr() as usize; + let aligned_off = base.next_multiple_of(4) - base; + for (off, expect_borrowed) in [(aligned_off, true), (aligned_off + 1, false)] { + for (i, v) in vals.iter().enumerate() { + bytes[off + i * 4..off + i * 4 + 4].copy_from_slice(&v.to_le_bytes()); + } + let view = TensorView::new( + Dtype::F32, + vec![vals.len()], + &bytes[off..off + vals.len() * 4], + ) + .unwrap(); + let cow = tensor_f32_cow(&view, "w").unwrap(); + assert_eq!( + matches!(cow, Cow::Borrowed(_)), + expect_borrowed, + "off={off}" + ); + let got: Vec = cow.iter().map(|f| f.to_bits()).collect(); + assert_eq!(got, vals, "off={off}"); + } + } + + #[test] + fn tensor_f32_cow_rejects_wrong_dtype_and_rank() { + let bytes = vec![0u8; 8]; + let bf16_view = TensorView::new(Dtype::BF16, vec![4], &bytes).unwrap(); + assert!(tensor_f32_cow(&bf16_view, "w").is_err()); + let f32_2d_view = TensorView::new(Dtype::F32, vec![2, 1], &bytes).unwrap(); + assert!(tensor_f32_cow(&f32_2d_view, "w").is_err()); + } #[test] fn tensor_bf16_cow_borrows_aligned_and_decodes_unaligned() { From 68671114dba8b0c83c7c764fd7e81bbdddf7e878 Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Sat, 5 Sep 2026 09:57:39 +0000 Subject: [PATCH 06/15] refactor(qwen35): drop derived TP guards, dedupe decode-group predicate, prune fixture tests Address the P2B review on the GDR state-sharding split: - config/tp.rs drops the linear_num_value_heads % world_size guard: Config35::try_from already validates value % key == 0, so the key-head guard implies value divisibility; one source of truth. The surviving test uses an internally consistent key/value = 17/34 fixture. - Drop the test-only local_linear_q_dim field; the constructor keeps q_dim as a local to derive local_linear_qkv_dim. Geometry tests re-derive q from qkv/z dims. - batch_decode_eager_logits derives start_positions only inside the paged-prefill fallback branch instead of allocating per rank per token on the compiled path, and routes on Config35::decode_group_is_compiled: head sharding leaves the q-per-kv group unchanged, so the rank-local predicate duplicated the config one. The one-token paged plan is now shared with the hybrid fallback via one_token_paged_plan(). - weights.rs deletes the F32 safetensors fixture test and the TP1 segment test: they re-asserted the segment tuples they built expectations from without exercising the production BF16 stitch loaders. Segment coverage stays on the direct QKV/conv tests plus the HF golden gates. - Trim the repeated Phase-2b ownership explanation across recurrent_state, prefill, prefill_buffers, batch_decode, tp_executor, and config/tp to one module-level invariant plus the non-obvious fused-QKV stitching and post-out_proj all-reduce notes. - tp-implementation.md P2B record updated for the dropped value-head guard, the deleted local predicate, and the removed local_linear_q_dim dim. Evidence (2x RTX 4090, sm_89): cargo check/clippy --release --all-targets -D warnings clean; qwen35 lib tests 104 passed / 0 failed; cargo fmt clean. TP2 gpu gates not rerun (no local weights). Signed-off-by: Ziyang Zhang --- docs/models/qwen35/tp-implementation.md | 4 +- pegainfer-qwen35/src/batch_decode.rs | 116 +++++++++++------------ pegainfer-qwen35/src/config/tp.rs | 58 +++--------- pegainfer-qwen35/src/prefill.rs | 6 +- pegainfer-qwen35/src/prefill_buffers.rs | 3 +- pegainfer-qwen35/src/recurrent_state.rs | 15 ++- pegainfer-qwen35/src/tp_executor.rs | 8 +- pegainfer-qwen35/src/weights.rs | 120 ------------------------ 8 files changed, 87 insertions(+), 243 deletions(-) diff --git a/docs/models/qwen35/tp-implementation.md b/docs/models/qwen35/tp-implementation.md index 671dfc888..8d8209ac5 100644 --- a/docs/models/qwen35/tp-implementation.md +++ b/docs/models/qwen35/tp-implementation.md @@ -447,11 +447,11 @@ Why this should be separated from GDR sharding: Landed as #946 split 1/4 (#1003), after P2a had established the mixed-step and state-lifecycle contract. Each TP rank now owns a rank-local slice of the linear-attention/GDR surface instead of replicating it: -- `LocalGeometry` computes rank-local linear dims (`local_linear_num_key_heads`, `local_linear_num_value_heads`, `local_linear_q_dim`, `local_linear_v_dim`, `local_linear_qkv_dim`, `local_linear_z_dim`) and fails closed with `ConfigError::TpIndivisible` when `linear_num_key_heads` or `linear_num_value_heads` do not divide by `world_size`; there is no silent replication fallback. +- `LocalGeometry` computes rank-local linear dims (`local_linear_num_key_heads`, `local_linear_num_value_heads`, `local_linear_v_dim`, `local_linear_qkv_dim`, `local_linear_z_dim`) and fails closed with `ConfigError::TpIndivisible` when `linear_num_key_heads` does not divide by `world_size`; there is no silent replication fallback. Value-head divisibility needs no second guard: `Config35` already validates `linear_num_value_heads % linear_num_key_heads == 0`. - Weight loading shards the head-indexed tensors: the fused QKV projection and the depthwise conv1d are stitched head-locally per segment (`load_linear_in_proj_qkv_shard` / `load_linear_conv1d_shard`; Q/K segments follow key-head ranges, V/conv follow value-head ranges), z/beta/alpha are row shards, `dt_bias`/`A_log` are 1-D shards (A_log stays f32), and linear `out_proj` is column sharded as a row-parallel `[hidden, local_z]` matrix. - `RecurrentState` (`[local_value_heads, K, V]` f32), conv state (`[local_qkv x (kernel_dim - 1)]` bf16), and all prefill/decode scratch (`GdrChunkwiseScratch35`, prefill/decode buffers) size themselves from the local geometry; worker capacity math uses the same locals. - The hidden-residual all-reduce happens once after the local linear-attention `out_proj` (`all_reduce_hidden`), on prefill and decode alike; the column-sharded `out_proj` is what makes that reduction point sufficient. -- Full-attention decode-group supportability is evaluated on the rank-local GQA group (`LocalGeometry::local_decode_group_is_compiled`): at TP2 the 27B case leaves `q/kv = 6`, which has no compiled FlashInfer batch-decode kernel, so those layers reroute decode through the eager/paged fallback. The predicate is deterministic per model and identical on every rank, so the reroute adds no collectives. +- Full-attention decode-group supportability uses the config-level GQA group (`Config35::decode_group_is_compiled`): head sharding leaves the q-per-kv group size unchanged, so the predicate is identical on every rank and the reroute adds no collectives. The 27B case leaves `q/kv = 6`, which has no compiled FlashInfer batch-decode kernel, so those layers reroute decode through the eager/paged fallback. - TP1 contract is unchanged: at `world_size == 1` every local dim equals the global dim, so kernels, buffers, and fixture behavior are byte-identical to pre-P2b. Non-negotiable invariant (still held): diff --git a/pegainfer-qwen35/src/batch_decode.rs b/pegainfer-qwen35/src/batch_decode.rs index cacfe97a8..204a3e080 100644 --- a/pegainfer-qwen35/src/batch_decode.rs +++ b/pegainfer-qwen35/src/batch_decode.rs @@ -271,7 +271,6 @@ 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)?; @@ -279,7 +278,6 @@ 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); @@ -293,41 +291,22 @@ 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 + // When this GQA 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() { + // Head sharding leaves the q-per-kv group size unchanged, so the + // config-level predicate decides the per-rank route identically on + // every rank; the reroute adds no collectives. + let prefill_attn_plan = if self.config.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 = - 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 - ) - })?, - ) + let start_positions: Vec = positions.iter().map(|&p| p as usize).collect(); + Some(self.one_token_paged_plan( + &kv_refs, + &start_positions, + self.geometry.local_num_attention_heads(), + self.geometry.local_num_key_value_heads(), + "eager decode", + )?) }; let kv_buffer = kv_states[0].buffer(); @@ -494,31 +473,14 @@ impl Qwen35Model { ) })?; - 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; the hd256 FFI takes no override. - let plan = ops::PrefillPagedPlan::from_raw_batch_with_cta_tile_q( - &self.ctx, - &page_indices, - &last_page_lens, + let kv_refs: Vec<&KvState> = kv_states.iter().map(|s| &**s).collect(); + let plan = self.one_token_paged_plan( + &kv_refs, &start_positions, - &seq_lens, self.config.num_attention_heads, self.config.num_key_value_heads, - self.config.head_dim, - 0, - ) - .with_context(|| { - format!( - "hybrid decode build PrefillPagedPlan bs={bs}, pages={}, heads={}/{}, head_dim={}", - page_indices.iter().map(Vec::len).sum::(), - self.config.num_attention_heads, - self.config.num_key_value_heads, - self.config.head_dim - ) - })?; + "hybrid decode", + )?; let kv_buffer = kv_states[0].buffer(); let layout = *kv_states[0].layout(); @@ -542,6 +504,42 @@ impl Qwen35Model { ) } + /// Paged-prefill plan that runs one decode row per request through the + /// prefill attention kernel; used when the GQA group has no compiled + /// batch-decode kernel. `cta_tile_q` 0 = the kernel's own FA2 derivation; + /// the hd256 FFI takes no override. + fn one_token_paged_plan( + &self, + kv_refs: &[&KvState], + start_positions: &[usize], + num_q_heads: usize, + num_kv_heads: usize, + label: &str, + ) -> Result { + let bs = kv_refs.len(); + let page_indices: Vec> = kv_refs.iter().map(|kv| kv.page_indices_i32()).collect(); + let last_page_lens: Vec = kv_refs.iter().map(|kv| kv.last_page_len()).collect(); + let seq_lens = vec![1usize; bs]; + ops::PrefillPagedPlan::from_raw_batch_with_cta_tile_q( + &self.ctx, + &page_indices, + &last_page_lens, + start_positions, + &seq_lens, + num_q_heads, + num_kv_heads, + self.config.head_dim, + 0, + ) + .with_context(|| { + format!( + "{label} build PrefillPagedPlan bs={bs}, pages={}, heads={num_q_heads}/{num_kv_heads}, head_dim={}", + page_indices.iter().map(Vec::len).sum::(), + self.config.head_dim + ) + }) + } + fn batch_decode_kernels_graph( &self, kv_buffer: &cudarc::driver::CudaSlice, @@ -773,10 +771,8 @@ impl Qwen35Model { /// (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). + /// `out_proj` is column-sharded, so its partial hidden sum is the one + /// linear-attention output all-reduced under TP (no-op at world_size 1). fn batch_decode_linear_attention_slots( &self, attn: &LinearAttentionLayer, diff --git a/pegainfer-qwen35/src/config/tp.rs b/pegainfer-qwen35/src/config/tp.rs index 5506067ac..7a8b39347 100644 --- a/pegainfer-qwen35/src/config/tp.rs +++ b/pegainfer-qwen35/src/config/tp.rs @@ -85,7 +85,6 @@ pub(crate) struct LocalGeometry { 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, } @@ -97,7 +96,8 @@ 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 heads included: Phase 2b shards them per rank); + /// (linear-attention key heads included; the value-head count follows + /// from the `Config35` key/value-head invariant); /// - `rank < world_size` and `world_size >= 1` are guaranteed by /// `TensorParallelConfig::try_from`. pub(crate) fn try_new( @@ -131,8 +131,9 @@ 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. + // Fail closed on an indivisible key-head count rather than falling + // back to replication; value-head divisibility follows from the + // Config35 value % key invariant and needs no second guard. if !config.linear_num_key_heads.is_multiple_of(tp.world_size()) { return Err(ConfigError::TpIndivisible { field: "linear_num_key_heads", @@ -140,16 +141,6 @@ impl LocalGeometry { 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(); @@ -173,7 +164,6 @@ impl LocalGeometry { 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, @@ -221,7 +211,7 @@ impl LocalGeometry { self.local_full_attn_gated_q_dim } - // ── Linear-attention local dims (Phase 2b TP sharding) ──────────────── + // ── Linear-attention local dims ─────────────────────────────────────── // 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. @@ -242,16 +232,6 @@ impl LocalGeometry { 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)] @@ -376,10 +356,13 @@ mod tests { } #[test] - fn requires_linear_attention_head_divisibility() { + fn requires_linear_attention_key_head_divisibility() { let tp = TensorParallelConfig::try_from((1, 2)).unwrap(); let mut broken = config(); + // Keep the Config35 value % key invariant intact so the failing + // branch is the key-head TP guard itself. broken.linear_num_key_heads = 17; + broken.linear_num_value_heads = 34; let err = LocalGeometry::try_new(&broken, tp, false).unwrap_err(); assert_eq!( err, @@ -389,18 +372,6 @@ mod tests { 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] @@ -410,10 +381,13 @@ mod tests { 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); + // q/k rows derive as key_heads * key_head_dim; qkv = 2q + v, z = v. assert_eq!(geom.local_linear_qkv_dim(), 4096); assert_eq!(geom.local_linear_z_dim(), 2048); + assert_eq!( + (geom.local_linear_qkv_dim() - geom.local_linear_z_dim()) / 2, + 8 * cfg.linear_key_head_dim + ); } #[test] @@ -424,8 +398,6 @@ mod tests { 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); diff --git a/pegainfer-qwen35/src/prefill.rs b/pegainfer-qwen35/src/prefill.rs index 6644764b1..043dfa4dc 100644 --- a/pegainfer-qwen35/src/prefill.rs +++ b/pegainfer-qwen35/src/prefill.rs @@ -436,10 +436,8 @@ 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). + /// `out_proj` is column-sharded, so its partial hidden sum is the one + /// linear-attention output all-reduced under TP (no-op at world_size 1). fn prefill_linear_attention( &self, attn: &LinearAttentionLayer, diff --git a/pegainfer-qwen35/src/prefill_buffers.rs b/pegainfer-qwen35/src/prefill_buffers.rs index 027cf0e39..b4f4b023c 100644 --- a/pegainfer-qwen35/src/prefill_buffers.rs +++ b/pegainfer-qwen35/src/prefill_buffers.rs @@ -57,8 +57,7 @@ impl GdrChunkwiseScratch35 { geometry: LocalGeometry, seq_len: usize, ) -> Result { - // GDR scratch is rank-local under TP: local value heads, global dims - // at world_size 1. + // GDR scratch sizes follow the rank's local value-head geometry. Self::from_dims( ctx, geometry.local_linear_num_value_heads(), diff --git a/pegainfer-qwen35/src/recurrent_state.rs b/pegainfer-qwen35/src/recurrent_state.rs index 3d868332a..d35344343 100644 --- a/pegainfer-qwen35/src/recurrent_state.rs +++ b/pegainfer-qwen35/src/recurrent_state.rs @@ -4,8 +4,9 @@ //! - 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. +//! Under TP, value heads (and fused qkv channels) are sharded across ranks, +//! so every rank owns its own recurrent/conv state; these states are never +//! all-reduced. use anyhow::Result; use cudarc::driver::CudaSlice; @@ -19,11 +20,10 @@ use super::config::LocalGeometry; /// Per-layer recurrent state for a single linear attention layer. pub(crate) struct LayerRecurrentState { /// 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. + /// Stored as f32 per mamba_ssm_dtype="float32" in config. pub(crate) state: CudaSlice, /// Conv1d state buffer: [local_linear_qkv_dim * (conv_kernel_dim - 1)] bf16 - /// Stores the last (kernel_dim - 1) inputs for causal conv1d. Rank-local. + /// Stores the last (kernel_dim - 1) inputs for causal conv1d. pub(crate) conv_state: DeviceVec, } @@ -47,8 +47,7 @@ pub(crate) struct LinearStatePointerTables { } /// Per-layer element counts shared by allocation and reservation: -/// (linear layers, f32 state elements, bf16 conv elements). Sizes are -/// rank-local under TP; at world_size 1 they equal the global dims. +/// (linear layers, f32 state elements, bf16 conv elements). 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 = geometry.local_linear_num_value_heads() @@ -155,7 +154,7 @@ impl LinearStatePointerTables { } } -/// Device bytes of one request's (rank-local) recurrent state. +/// Device bytes of one request's 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 diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index 830b3c231..16d245b3a 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. Phase 2b shards the -//! linear-attention/GDR weight and state surface per rank. +//! One canonical eager unified command per step. Linear-attention/GDR weights +//! and state are sharded per rank. use std::collections::HashSet; use std::panic::AssertUnwindSafe; @@ -1037,8 +1037,8 @@ impl TpWorkerPrepared { .ctx .mem_get_info() .map_err(|err| anyhow::anyhow!("failed to query TP rank {rank} memory: {err}"))?; - // Phase 2b: per-request recurrent state is rank-local, so worker - // capacity math uses the local value-head/qkv sizes. + // 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( diff --git a/pegainfer-qwen35/src/weights.rs b/pegainfer-qwen35/src/weights.rs index 8d70df598..ed90e411d 100644 --- a/pegainfer-qwen35/src/weights.rs +++ b/pegainfer-qwen35/src/weights.rs @@ -1068,124 +1068,4 @@ mod tests { // 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() - .as_chunks::<4>() - .0 - .iter() - .map(|b| f32::from_le_bytes(*b)) - .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() - .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 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 aa8be2a9b268e00e85885cb09809e77f06858e6a Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Thu, 20 Aug 2026 15:07:11 +0000 Subject: [PATCH 07/15] 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 | 199 ++++++++++++++++++---------- 1 file changed, 131 insertions(+), 68 deletions(-) diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index 16d245b3a..568c98edc 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -1,7 +1,8 @@ //! Tensor-parallel worker runtime for Qwen3.5. //! //! One canonical eager unified command per step. Linear-attention/GDR weights -//! and state are sharded per rank. +//! and state are sharded per rank, and decode rows run as one batched forward +//! 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 +1007,6 @@ struct TpRequestState { phase: TpRequestPhase, kv: KvState, recurrent: RecurrentState, - linear_pointer_tables: LinearStatePointerTables, } #[derive(Clone, Copy, Debug, PartialEq, Eq)] @@ -1341,6 +1341,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 +1526,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 +1559,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 d2242de35b01ed95e92aee510f41096b17760867 Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Thu, 20 Aug 2026 17:50:18 +0000 Subject: [PATCH 08/15] 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-qwen35/src/batch_decode.rs | 61 +- pegainfer-qwen35/src/batch_decode_graph.rs | 48 ++ pegainfer-qwen35/src/config/error.rs | 4 - pegainfer-qwen35/src/config/tp.rs | 38 +- 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 | 34 +- pegainfer-qwen35/tests/e2e_scheduler.rs | 32 +- pegainfer-qwen35/tests/hf_golden_gate.rs | 170 +++++ pegainfer-qwen35/tests/serving_tp2.rs | 34 +- 17 files changed, 1454 insertions(+), 159 deletions(-) diff --git a/docs/index.md b/docs/index.md index 2aae4671f..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, P2A, and P2B GDR state sharding are complete: TP2 has start-gated eager unified prefill+decode, fail-closed lifecycle recovery, and rank-local linear-attention weights/state with a post-`out_proj` hidden all-reduce; batched TP decode (#1004) and TP CUDA Graph (#1005) are 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 8d8209ac5..50bfa966b 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, P2A, and the P2B GDR state sharding are complete: TP2 has start-gated eager unified prefill+decode, fail-closed lifecycle recovery, and rank-local linear-attention weights/state with a single post-`out_proj` hidden all-reduce; batched eager TP decode and TP CUDA Graph are 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-09 @@ -469,9 +469,160 @@ Acceptance at `fcdeb5a4` (27B TP2 on 2x RTX 4090 48GB, sm_89; fixture-pinned 27B Not in this step: batching the TP decode loop across rows, TP CUDA Graph capture, and the matched Phase-1-vs-P2b HBM/latency/throughput A/B promised in #1001; no performance claim is made until that rerun lands on the merged stack. +## 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 -- Land batched eager TP decode (#1004) on top of the P2B state sharding, then TP CUDA Graph (#1005); rerun the #946 throughput A/B on the merged stack before any performance claim. +- 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-qwen35/src/batch_decode.rs b/pegainfer-qwen35/src/batch_decode.rs index 204a3e080..e3da74d25 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, @@ -340,6 +356,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"); @@ -351,6 +368,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!( @@ -411,17 +432,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 7a8b39347..36c70a8a0 100644 --- a/pegainfer-qwen35/src/config/tp.rs +++ b/pegainfer-qwen35/src/config/tp.rs @@ -90,26 +90,22 @@ 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 key heads included; the value-head count follows /// from the `Config35` key/value-head invariant); /// - `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", @@ -279,7 +275,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); @@ -309,7 +305,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 { @@ -322,7 +318,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 { @@ -334,7 +330,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 { @@ -345,16 +341,6 @@ 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_key_head_divisibility() { let tp = TensorParallelConfig::try_from((1, 2)).unwrap(); @@ -363,7 +349,7 @@ mod tests { // branch is the key-head TP guard itself. broken.linear_num_key_heads = 17; broken.linear_num_value_heads = 34; - let err = LocalGeometry::try_new(&broken, tp, false).unwrap_err(); + let err = LocalGeometry::try_new(&broken, tp).unwrap_err(); assert_eq!( err, ConfigError::TpIndivisible { @@ -378,7 +364,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); // q/k rows derive as key_heads * key_head_dim; qkv = 2q + v, z = v. @@ -395,7 +381,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); let global_qkv = 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 568c98edc..9d98dd2c1 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -16,16 +16,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; @@ -43,9 +49,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 { @@ -70,6 +110,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, }, @@ -195,6 +246,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)] @@ -245,6 +308,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 { @@ -259,6 +327,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) } } } @@ -306,10 +388,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" @@ -321,7 +399,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, }, @@ -330,6 +408,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() { @@ -359,6 +454,7 @@ impl Qwen35TpExecutor { model, max_batch, max_prefill_tokens, + graph_enabled, nccl_id, Arc::clone(&startup_gate), Arc::clone(&effective_max_batch), @@ -426,7 +522,7 @@ impl Qwen35TpExecutor { } disarm_nccl_startup_watchdog(watchdog_done, watchdog)?; - Ok(Self { + let executor = Self { workers, poison, world_size, @@ -435,7 +531,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)] @@ -463,6 +583,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()?; @@ -495,7 +699,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 { @@ -536,18 +753,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) } @@ -611,10 +859,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, })?; @@ -627,6 +889,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()?; @@ -880,6 +1165,7 @@ impl TpStartupGate { } impl TpWorker { + #[allow(clippy::too_many_arguments)] #[allow(clippy::type_complexity)] fn spawn( rank: usize, @@ -887,6 +1173,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, @@ -911,6 +1198,7 @@ impl TpWorker { model, max_batch, max_prefill_tokens, + graph_enabled, ); let prepared = match prepared { Ok((prepared, rank_max_batch)) => { @@ -926,7 +1214,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); @@ -984,8 +1272,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, @@ -1006,7 +1303,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)] @@ -1030,6 +1330,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 @@ -1046,9 +1347,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, @@ -1095,6 +1404,7 @@ impl TpWorkerPrepared { self, nccl_id: cudarc::nccl::safe::Id, effective_max_batch: usize, + graph_enabled: bool, poison: Arc, ) -> Result { let Self { @@ -1118,12 +1428,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, @@ -1198,14 +1521,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)] @@ -1227,7 +1561,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)] @@ -1317,7 +1651,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), @@ -1361,6 +1700,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 @@ -1388,7 +1730,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); @@ -1468,6 +1815,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, @@ -1568,7 +2053,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) @@ -1580,14 +2065,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<()> { @@ -1725,17 +2363,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, @@ -1816,7 +2443,6 @@ fn validate_exact_rank_responses( Ok(replies) } -#[cfg(test)] fn validate_ack_responses( responses: Vec, world_size: usize, @@ -2221,12 +2847,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 ed90e411d..0585fa84e 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] 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 ab02c2334ab0d6bdfd301ae824e4345220c224bb Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Sun, 30 Aug 2026 13:33:09 +0000 Subject: [PATCH 09/15] 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 9d98dd2c1..05b179f83 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -1349,19 +1349,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", From 8be07c93e7037ecd2983844b8ff4e4398489ec3e Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Sat, 5 Sep 2026 10:01:17 +0000 Subject: [PATCH 10/15] refactor(qwen35): point the TP graph gate at the config-level decode-group predicate The rank-local GQA group equals the global group under validated head sharding, so the P2B review removed LocalGeometry::local_decode_group_is_compiled. Route the P2c graph gate through Config35::decode_group_is_compiled instead of re-adding the local duplicate; the gate decision was already identical on every rank. Evidence (2x RTX 4090, sm_89): cargo check/clippy --release --all-targets -D warnings clean; qwen35 lib tests 107 passed / 0 failed; cargo fmt clean. Signed-off-by: Ziyang Zhang --- docs/models/qwen35/tp-design.md | 4 ++-- docs/models/qwen35/tp-implementation.md | 2 +- pegainfer-qwen35/src/config/tp.rs | 2 +- pegainfer-qwen35/src/tp_executor.rs | 2 +- 4 files changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/models/qwen35/tp-design.md b/docs/models/qwen35/tp-design.md index 782fa38cc..c20682664 100644 --- a/docs/models/qwen35/tp-design.md +++ b/docs/models/qwen35/tp-design.md @@ -272,12 +272,12 @@ Validation scope: ## 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; +`Config35::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. +**Gate**: graph mode active iff `enable_cuda_graph && config.decode_group_is_compiled()`. 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. diff --git a/docs/models/qwen35/tp-implementation.md b/docs/models/qwen35/tp-implementation.md index 50bfa966b..5a1f668e6 100644 --- a/docs/models/qwen35/tp-implementation.md +++ b/docs/models/qwen35/tp-implementation.md @@ -525,7 +525,7 @@ 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, +config.decode_group_is_compiled()`. 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, diff --git a/pegainfer-qwen35/src/config/tp.rs b/pegainfer-qwen35/src/config/tp.rs index 36c70a8a0..4bdf9fba4 100644 --- a/pegainfer-qwen35/src/config/tp.rs +++ b/pegainfer-qwen35/src/config/tp.rs @@ -100,7 +100,7 @@ impl LocalGeometry { /// `TensorParallelConfig::try_from`. /// /// CUDA Graph under TP is gated at executor startup on - /// [`LocalGeometry::local_decode_group_is_compiled`] (P2c): uncompiled GQA + /// [`Config35::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, diff --git a/pegainfer-qwen35/src/tp_executor.rs b/pegainfer-qwen35/src/tp_executor.rs index 05b179f83..47def95d8 100644 --- a/pegainfer-qwen35/src/tp_executor.rs +++ b/pegainfer-qwen35/src/tp_executor.rs @@ -413,7 +413,7 @@ impl Qwen35TpExecutor { // 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(); + let graph_enabled = enable_cuda_graph && first.config().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(|| { From 5e80c5134647e969a0986d8ae28921e561184a0d Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Thu, 20 Aug 2026 15:19:42 +0000 Subject: [PATCH 11/15] docs(qwen35): record the TP rebase onto #870 Signed-off-by: Ziyang Zhang --- docs/models/qwen35/tp-implementation.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/models/qwen35/tp-implementation.md b/docs/models/qwen35/tp-implementation.md index 5a1f668e6..5a2ba8bef 100644 --- a/docs/models/qwen35/tp-implementation.md +++ b/docs/models/qwen35/tp-implementation.md @@ -506,7 +506,7 @@ scheduler planner-gate/test updates — ours' `scheduler.rs`, upstream, so those files resolved to #870's versions except the `alloc_recurrent` signature change. -Validation on 2× RTX 4090: +Validation on 2× RTX 4090 (venv NCCL on `LD_LIBRARY_PATH`): - `cargo check --release -p pegainfer-qwen35 --features qwen35` clean; `cargo fmt --check -p pegainfer-qwen35` clean. @@ -622,6 +622,8 @@ the *eager* test while the graph test ran concurrently. cross-harness band, no TP-induced accuracy regression. MMLU-Pro / SuperGPQA sampled runs remain outstanding; rerun on this rebased branch before citing parity. +## Follow-Ups + - 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`. From b6c39cc9ca7fc2fa831f6dc1d439ce53783853e7 Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Thu, 20 Aug 2026 11:19:40 +0000 Subject: [PATCH 12/15] docs(qwen35): TP implementation record + 27B TP2 knowledge benchmark eval - tp-implementation.md: Phase 2a/2b + Step 3 batched-decode landing record - benchmarks/qwen35-27b-tp2-knowledge-eval.md: MMLU-Redux 94.09 vs 93.2, C-Eval 88.11 vs 90.5 (full runs, truncation rerun-merged); in band - scripts/eval_mc.py + eval_rerun_truncated.py: benchmark runner (OpenCompass/lm-eval recipe replicas over the chat completions API) - developer-onboarding: NCCL libnccl.so dlopen note Signed-off-by: Ziyang Zhang --- .gitignore | 1 + .../qwen35-27b-tp2-knowledge-eval.md | 52 +++ docs/index.md | 1 + docs/models/qwen35/roadmap.md | 2 +- docs/models/qwen35/tp-design.md | 2 +- docs/models/qwen35/tp-implementation.md | 7 + docs/playbooks/developer-onboarding.md | 2 + scripts/eval_mc.py | 393 ++++++++++++++++++ scripts/eval_rerun_truncated.py | 123 ++++++ 9 files changed, 581 insertions(+), 2 deletions(-) create mode 100644 docs/benchmarks/qwen35-27b-tp2-knowledge-eval.md create mode 100644 scripts/eval_mc.py create mode 100644 scripts/eval_rerun_truncated.py diff --git a/.gitignore b/.gitignore index 30b60a11e..89f8d4332 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ __pycache__/ /docs/private/ /profile/ .claude/scheduled_tasks.lock +/results/ diff --git a/docs/benchmarks/qwen35-27b-tp2-knowledge-eval.md b/docs/benchmarks/qwen35-27b-tp2-knowledge-eval.md new file mode 100644 index 000000000..310e19546 --- /dev/null +++ b/docs/benchmarks/qwen35-27b-tp2-knowledge-eval.md @@ -0,0 +1,52 @@ +# Qwen3.5-27B TP2 知识基准评测(官方分对比) + +> TL;DR:Qwen3.5-27B 在 pegainfer TP2(2× RTX 4090,batched eager decode)上跑知识基准,C-Eval 88.11(官方 90.5)、MMLU-Redux 94.09(官方 93.2),均在跨 harness 正常带内;MMLU-Pro / SuperGPQA 因运行时长原因仅完成抽样冒烟,未出最终分(见下文)。模型数值无 TP 引入的精度问题。 +> +> 注:分数实测于 rebase 前的 f4c66780 分支(自研 Phase 1/2a 线);rebase 到 #870 后 logits golden gate 两侧一致通过,数值可迁移,但若正式引用请在本 PR 分支上复跑确认。 + +## 环境 + +- GPU:2× RTX 4090(48 GB 版本),`--tp-size 2 --cuda-graph false`(TP+CUDA Graph 仍 fail-closed) +- 模型:Qwen/Qwen3.5-27B `fc05daec`,BF16,served-model-name `qwen35-27b-tp2` +- 采样:temperature=0.0,top_p=1.0,chat completions(thinking 模式,即模板默认行为) +- 评测器:`scripts/eval_mc.py`(自研,统一 `/v1/chat/completions` 并发 48),配方逐项复刻官方 harness: + - **C-Eval** = OpenCompass `ceval_gen`:52 学科 val split 全量 1346 题,dev split 5-shot,"答案: " 续写,首大写字母抽取 + - **MMLU-Redux** = lm-eval `mmlu_redux_generative`:`fxmarty/mmlu-redux-2.0-ok` 57 学科 test 全量 5330 题,0-shot,首个 `[ABCD]` 抽取 + - **MMLU-Pro** = lm-eval `mmlu_pro`:TIGER-Lab/MMLU-Pro test,validation split 5-shot CoT,`answer is (X)` 抽取 + - **SuperGPQA** = OpenCompass `supergpqa_gen`:`m-a-p/SuperGPQA` train 26529 题,0-shot,"Answer: X" 字母/内容两层抽取 +- 启动命令:`LD_LIBRARY_PATH=<.venv>/nvidia/nccl/lib ./target/release/pegainfer --model-path <27B> --served-model-name qwen35-27b-tp2 --tp-size 2 --cuda-graph false --port 18082` +- 结果原始数据:`results/qwen35-27b-tp2-eval/{ceval,mmlu_redux}_samples*.json`(本地,未入库) + +## 结果(截至 2026-08-20,评测按需要提前终止) + +| 基准 | 官方 | 实测 | n | 口径 | Δ 判定 | +|---|---|---|---|---|---| +| MMLU-Redux | 93.2 | **94.09** | 5330 全量 | 8192 cap + 截断重跑合并(32 条重跑,2 条仍截断) | **同带** | +| C-Eval | 90.5 | **88.11** | 1346 全量 | 8192 cap + 截断重跑合并(48 条重跑,0 条残留) | **同带边缘**(-2.4pp,CI95 ±1.7pp) | +| MMLU-Pro | 86.1 | — | 600/2000 中止 | 抽样 n=2000(cap 24576)跑到 30% 人工终止;100 题冒烟在 4096 cap 下 51% 截断 | 无最终分 | +| SuperGPQA | 65.6 | — | 未正式跑 | 100 题冒烟:可完成子集 43 题对金标准确 27/43≈63% | 无最终分 | + +## 关键观察 + +- **没有 TP 精度问题**:27B TP2 HF logits golden gate 全绿;C-Eval 非截断子集(1290/1346)准确率 90.2% ≈ 官方 90.5。C-Eval 的差距全部来自 thinking 长度上限被掐断的最难题,而非模型错算。 +- **MMLU-Redux 略高于官方**(+0.9pp):同带,说明 prompt/抽取/数值链路都对。 +- **thinking 长度是最大的系统变量**:thinking 模型在 C-Eval 上 ~4% 题需要 >8192 token,MMLU-Pro 上 >1/3 题在 4096 内收不住。官方 harness 的 max_tokens 未知(推测 ≥32k);本评测用 8192 首轮 + 32768 重跑合并来逼近。跨 harness ±1–2pp 属正常。 +- **吞吐前置条件**:此评测可行完全依赖 Step 3 的 batched eager TP decode 修复(此前 16 并发聚合仅 ~25 tok/s,全量不可行;修复后 ~450 tok/s @48 并发)。 + +## 复现命令 + +```bash +# server(见上);评测(hf 镜像): +HF_ENDPOINT=https://hf-mirror.com .venv/bin/python -u scripts/eval_mc.py ceval --max-tokens 8192 --concurrency 48 --out-dir results/qwen35-27b-tp2-eval +HF_ENDPOINT=https://hf-mirror.com .venv/bin/python -u scripts/eval_mc.py mmlu_redux --max-tokens 8192 --concurrency 48 --out-dir results/qwen35-27b-tp2-eval +# 截断样本重跑合并: +.venv/bin/python -u scripts/eval_rerun_truncated.py ceval --max-tokens 32768 --concurrency 16 +# 抽样(--sample 按学科比例分层,seed 1337,n=2000 时 CI95 半宽 ±1.3pp): +.venv/bin/python -u scripts/eval_mc.py mmlu_pro --sample 2000 --max-tokens 24576 --concurrency 48 +.venv/bin/python -u scripts/eval_mc.py supergpqa --sample 2000 --max-tokens 24576 --concurrency 48 +``` + +## 下一步 + +- 补齐 MMLU-Pro / SuperGPQA 抽样全量(各 2000 题,预估合计 ~10h,吞吐 ~450 tok/s 前提)。 +- 若要把 C-Eval 收敛到官方 ±1pp:换更大的 thinking 预算复跑全量(无合并),并确认 Qwen 官方 harness 的 prompt 模板与本评测是否一致。 diff --git a/docs/index.md b/docs/index.md index 322816bf3..3d6a7a7d3 100644 --- a/docs/index.md +++ b/docs/index.md @@ -234,6 +234,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `benchmarks/bs1-4k64-vllm-pegainfer.md` | RTX 5090 single-concurrency probe: `input_len=4096`, `output_len=64`, no vLLM prefix cache. PegaInfer TTFT median `177ms` vs vLLM `198ms`; TPOT median `6.47ms` vs `6.36ms`; corrected output throughput `+6%` for PegaInfer. | | `benchmarks/mixed-load-itl.md` | Qwen3-4B + Qwen3.5 mixed-load ITL (#244, #375): chunking-off sweeps via `bench_serving mixed`. Both freeze active decode for the full prefill. Qwen3 p99 blows up with prompt/QPS; the old Qwen3.5 “p99-immune” table is a **measurement artifact** (primary: hardcoded `max_batch=4` slot starvation — see #470 / `models/qwen35/mixed-load-itl-470.md`; secondary: short `bg_output_len`). Prefix reuse defeats it on Qwen3. | | `benchmarks/accuracy-eval-results.md` | Phase 1 GSM8K: Qwen3-4B PASS (pegainfer 85.37% vs HF 85.82%, delta -0.45 pp). Qwen3.5-4B historical FAIL recovered by #250 (strict 79.38%, flexible 79.30% vs HF 79.45%). | +| `benchmarks/qwen35-27b-tp2-knowledge-eval.md` | Qwen3.5-27B TP2 (2× RTX 4090, batched eager decode) knowledge benchmarks vs official: MMLU-Redux 94.09 vs 93.2 (full 5330), C-Eval 88.11 vs 90.5 (full 1346, thinking-cap truncation rerun-merged) — both in cross-harness band; MMLU-Pro/SuperGPQA partial smokes only. | | `benchmarks/qwen3-8b-pd-vs-mix-h200.md` | Qwen3-8B 多轮负载三方 A/B（2×H200）：P/D 1P+1D vs mixed×2（会话亲和 LB）vs mixed×1。吞吐持平（47.8k vs 47.0k tok/s），P/D 赢在 decode 稳定性（TPOT p99 10.08 vs 12.77ms，turn2+ TTFT 恒定 ~107ms vs 爬升 71→132ms），冷 turn1 多付 ~200ms（M3 目标）。含 vllm-bench 命令与 `max_completion_tokens` 坑。 | ## conventions diff --git a/docs/models/qwen35/roadmap.md b/docs/models/qwen35/roadmap.md index 4292afafe..ad302a3db 100644 --- a/docs/models/qwen35/roadmap.md +++ b/docs/models/qwen35/roadmap.md @@ -47,7 +47,7 @@ out: | Fault isolation | Open risk: batch-level execution errors can still fail multiple active requests | #654 | | Prefix reuse | Open: bounded joint KV/recurrent/conv snapshot design and implementation | #257 | | DFlash | In flight and opt-in: correctness-first work must stay default-off until gates pass | #434, PR #626, #654 | -| Tensor parallel | Phase 1 complete: eager dense TP2 worker/scheduler execution; Phase 2 still needs mixed-step execution and sharded linear-attention/GDR state. | `docs/models/qwen35/tp-implementation.md`, #446 | +| Tensor parallel | Done through Phase 2b: eager dense TP2, mixed-step unified execution, and sharded linear/GDR state, verified on 2× RTX 4090 for 9B and 27B; TP CUDA Graph + perf gates remain | `docs/models/qwen35/tp-implementation.md`, #446 | ## Active Contract diff --git a/docs/models/qwen35/tp-design.md b/docs/models/qwen35/tp-design.md index c20682664..f6b2b960d 100644 --- a/docs/models/qwen35/tp-design.md +++ b/docs/models/qwen35/tp-design.md @@ -1,6 +1,6 @@ # Qwen3.5 Tensor Parallelism Design -> **TL;DR:** Qwen3.5 TP Phase 2 is two separately delivered correctness milestones: P2a adds eager `RunUnifiedStep` with a shared ordered `RequestId` plan while retaining Phase 1 replicated GDR; P2b shards the head-indexed linear-attention/GDR surface and adds only the hidden all-reduce after local `out_proj`. +> **TL;DR:** Qwen3.5 tensor parallelism should reuse Qwen3's controller/worker TP runtime and stay degree-parametric. Phases 1, 2a, and 2b are implemented (see `tp-implementation.md` for the landing record, including the rebase onto #870): eager dense TP, TP mixed-step unified execution, and sharded linear-attention/GDR state. Remaining design work: TP CUDA Graph capture. > > **Last touched:** 2026-08 diff --git a/docs/models/qwen35/tp-implementation.md b/docs/models/qwen35/tp-implementation.md index 5a2ba8bef..a04b0a485 100644 --- a/docs/models/qwen35/tp-implementation.md +++ b/docs/models/qwen35/tp-implementation.md @@ -624,6 +624,13 @@ the *eager* test while the graph test ran concurrently. before citing parity. ## Follow-Ups +- 27B TP2 knowledge-benchmark parity (2026-08-20, validated pre-rebase on + the f4c66780 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`. diff --git a/docs/playbooks/developer-onboarding.md b/docs/playbooks/developer-onboarding.md index 8dd69c8b9..6cba19719 100644 --- a/docs/playbooks/developer-onboarding.md +++ b/docs/playbooks/developer-onboarding.md @@ -38,6 +38,8 @@ Verify: ``` > build.rs auto-detects `.venv/bin/python` for Triton AOT compilation. Override with `PEGAINFER_TRITON_PYTHON` if needed. +> +> Multi-GPU (TP/NCCL) runtime note: cudarc dlopens `libnccl.so` (its search list has no plain `libnccl.so.2`), while the `nvidia-nccl-cu13` wheel ships only `libnccl.so.2`. Add `libnccl.so -> libnccl.so.2` inside `.venv/lib/python3.*/site-packages/nvidia/nccl/lib/` and put that dir on `LD_LIBRARY_PATH` when running TP. ## 3. Build diff --git a/scripts/eval_mc.py b/scripts/eval_mc.py new file mode 100644 index 000000000..c879751b0 --- /dev/null +++ b/scripts/eval_mc.py @@ -0,0 +1,393 @@ +#!/usr/bin/env python3 +"""Multiple-choice benchmark runner for pegainfer's OpenAI-compatible chat API. + +Recipes copied from the canonical harnesses so scores stay comparable: + * C-Eval -> opencompass ceval_gen (5-shot from dev split, "答案: " + completion, first-capital extraction) + * SuperGPQA -> opencompass supergpqa_gen (zero-shot prompt_format yaml, + extract_option_labels -> (A..J), content fallback) + * MMLU-Pro -> lm-eval mmlu_pro (5-shot CoT from validation split, + 'answer is \\(?(X)\\)?' extraction) + * MMLU-Redux -> lm-eval mmlu_redux_generative (0-shot, first [ABCD] extraction) + +All four go through /v1/chat/completions with one user message; Qwen3.5 is a +thinking model so reasoning lands in `reasoning`, the final answer text in +`content` — extraction runs on `content` only. + +Usage: + python eval_mc.py ceval|supergpqa|mmlu_pro|mmlu_redux \ + --base-url http://127.0.0.1:18082/v1 --model qwen35-27b-tp2 +""" +import argparse +import asyncio +import json +import re +import sys +import time +from pathlib import Path + +import datasets +import httpx + +# ---------------------------------------------------------------- ceval +CEVAL_SUBJECT_CN = { + 'computer_network': '计算机网络', 'operating_system': '操作系统', + 'computer_architecture': '计算机组成', 'college_programming': '大学编程', + 'college_physics': '大学物理', 'college_chemistry': '大学化学', + 'advanced_mathematics': '高等数学', 'probability_and_statistics': '概率统计', + 'discrete_mathematics': '离散数学', 'electrical_engineer': '注册电气工程师', + 'metrology_engineer': '注册计量师', 'high_school_mathematics': '高中数学', + 'high_school_physics': '高中物理', 'high_school_chemistry': '高中化学', + 'high_school_biology': '高中生物', 'middle_school_mathematics': '初中数学', + 'middle_school_biology': '初中生物', 'middle_school_physics': '初中物理', + 'middle_school_chemistry': '初中化学', 'veterinary_medicine': '兽医学', + 'college_economics': '大学经济学', 'business_administration': '工商管理', + 'marxism': '马克思主义基本原理', + 'mao_zedong_thought': '毛泽东思想和中国特色社会主义理论体系概论', + 'education_science': '教育学', 'teacher_qualification': '教师资格', + 'high_school_politics': '高中政治', 'high_school_geography': '高中地理', + 'middle_school_politics': '初中政治', 'middle_school_geography': '初中地理', + 'modern_chinese_history': '近代史纲要', + 'ideological_and_moral_cultivation': '思想道德修养与法律基础', + 'logic': '逻辑学', 'law': '法学', + 'chinese_language_and_literature': '中国语言文学', 'art_studies': '艺术学', + 'professional_tour_guide': '导游资格', 'legal_professional': '法律职业资格', + 'high_school_chinese': '高中语文', 'high_school_history': '高中历史', + 'middle_school_history': '初中历史', 'civil_servant': '公务员', + 'sports_science': '体育学', 'plant_protection': '植物保护', + 'basic_medicine': '基础医学', 'clinical_medicine': '临床医学', + 'urban_and_rural_planner': '注册城乡规划师', 'accountant': '注册会计师', + 'fire_engineer': '注册消防工程师', + 'environmental_impact_assessment_engineer': '环境影响评价工程师', + 'tax_accountant': '税务师', 'physician': '医师资格', +} + + +def ceval_prompt(ch_name, q, shots): + head = f'以下是中国关于{ch_name}考试的单项选择题，请选出其中的正确答案。\n' + + def block(item): + return (f"{item['question']}\nA. {item['A']}\nB. {item['B']}\n" + f"C. {item['C']}\nD. {item['D']}") + + examples = ''.join(f"{block(s)}\n答案: {s['answer']}\n" for s in shots) + return head + examples + block(q) + '\n答案: ' + + +def first_capital(text): + for ch in text: + if ch.isupper(): + return ch + return '' + + +# ---------------------------------------------------------------- mmlu_pro +MMLUPRO_LETTERS = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J'] +MMLUPRO_RE = re.compile(r'answer is \(?([ABCDEFGHIJ])\)?') + + +def mmlupro_format(example, include_answer): + prompt = 'Question:\n' + example['question'] + '\nOptions:\n' + for i, opt in enumerate(example['options'][:len(MMLUPRO_LETTERS)]): + prompt += f"{MMLUPRO_LETTERS[i]}. {opt.strip()}\n" + if include_answer: + cot = example['cot_content'].replace("A: Let's think step by step.", + "Answer: Let's think step by step.") + return prompt + cot + '\n\n' + return prompt + "Answer: Let's think step by step." + + +def mmlupro_extract(text): + m = MMLUPRO_RE.search(text) + return m.group(1) if m else '' + + +# ---------------------------------------------------------------- mmlu-redux +MMLUREDUX_RE = re.compile(r'([ABCD])') + + +# ---------------------------------------------------------------- supergpqa +SG_PROMPT = ("Answer the following multiple choice question. There is only one " + "correct answer. The last line of your response should be in the " + "format 'Answer: $LETTER' (without quotes), where LETTER is one of " + "A, B, C, D, E, F, G, H, I, or J.\n\n{}") + + +def sg_build_questions(item): + opts = '\n'.join(f'{chr(65 + i)}) {o}' for i, o in enumerate(item['options'])) + return item['question'] + '\n' + opts + + +def sg_label_patterns(letters): + return [ + rf'[Tt]he\s+(?:\w+\s+)?(?:answer|option)(?:\w+\s+)?\s+is?:?\s*(?:[\*\$\\{{\(\[\\]*?(?:(?:\\boxed|\\mathbf|\\mathrm|\\text){{)?)*\s*([{letters}])(?:\\?\}}?\$?\)?\]?\}}?)*(?:[\s:\.\*)]|$)', + rf'(?i:Answer)[\*\s]*:\s*(?:[\*\$\\{{\(\[\\]*?(?:(?:\\boxed|\\mathbf|\\mathrm|\\text){{)?)*\s*([{letters}])(?:\\?\}}?\$?\)?\]?\}}?)*(?:[\s:\.\*)]|$)', + rf'^[^\w\r\n]*(?:[\*\$\\{{\(\[\\]*?(?:(?:\\boxed|\\mathbf|\\mathrm|\\text){{)?)*\s*([{letters}])(?:\\?\}}?\$?\)?\]?\}}?)*(?:[\s:\.\*)]|$)', + ] + + +def sg_extract_labels(text, letters='ABCDEFGHIJ'): + if not isinstance(text, str): + return None + text = text.rstrip() + last_line = text.split('\n')[-1] + pats = sg_label_patterns(letters) + for src in (last_line, text): + for p in pats: + try: + m = re.search(p, src, re.IGNORECASE) + except Exception: + m = None + if m: + return m.group(1) + return None + + +def sg_extract_content(text, options_content): + if not isinstance(text, str) or not isinstance(options_content, list): + return None + esc = [re.escape(o) for o in options_content] + alt = '|'.join(esc) + pats = [ + rf'[Tt]he\s+(?:\w+\s+)?(?:answer|option)(?:\w+\s+)?\s+is:?\s*(?:[\*\$\\{{\(\[\\]*?(?:(?:\\boxed|\\mathbf|\\mathrm|\\text){{)?)*\s*({alt})(?:\\?\}}?\$?\)?\]?\}}?)*(?:[\s:\.\*)]|$)', + rf'(?i:Answer)\s*(?:[\*\$\\{{\(\[\\]*?(?:(?:\\boxed|\\mathbf|\\mathrm|\\text){{)?)*\s*({alt})(?:\\?\}}?\$?\)?\]?\}}?)*(?:[\s:\.\*)]|$)', + rf'^[^\w\r\n]*(?:[\*\$\\{{\(\[\\]*?(?:(?:\\boxed|\\mathbf|\\mathrm|\\text){{)?)*\s*({alt})(?:\\?\}}?\$?\)?\]?\}}?)*(?:[\s:\.\*)]|$)', + ] + text = text.rstrip() + last_line = text.split('\n')[-1] + for src in (last_line, text): + for p in pats: + try: + m = re.search(p, src) + except Exception: + m = None + if m: + hit = m.group(1) + if hit in esc: + return options_content[esc.index(hit)] + return hit + return None + + +# ---------------------------------------------------------------- runner +STOP_MAP = {'mmlu_pro': ['Question:']} + + +async def run_completion(client, base_url, model, prompt, max_tokens, temperature, name): + payload = { + 'model': model, 'max_tokens': max_tokens, + 'temperature': temperature, 'top_p': 1.0, + 'messages': [{'role': 'user', 'content': prompt}], + } + if STOP_MAP.get(name): + payload['stop'] = STOP_MAP[name] + for attempt in range(6): + try: + r = await client.post(f'{base_url}/chat/completions', json=payload) + if r.status_code == 200: + data = r.json() + msg = data['choices'][0]['message'] + usage = data.get('usage') or {} + return { + 'content': msg.get('content') or '', + 'reasoning': msg.get('reasoning') or '', + 'prompt_tokens': usage.get('prompt_tokens') or 0, + 'completion_tokens': usage.get('completion_tokens') or 0, + 'finish_reason': data['choices'][0].get('finish_reason') or '', + } + body = f'HTTP {r.status_code} {r.text[:200]!r}' + except Exception as e: # noqa: BLE001 + body = f'{type(e).__name__}: {e!s}' or repr(e) + await asyncio.sleep(min(2 ** attempt, 20)) + if attempt == 5: + return {'content': '', 'reasoning': f'__ERROR__ {body}', + 'prompt_tokens': 0, 'completion_tokens': 0, 'finish_reason': 'error'} + return {'content': '', 'reasoning': '__ERROR__'} + + +async def evaluate(name, items, prompts, golds, args, out_dir): + sem = asyncio.Semaphore(args.concurrency) + limits = httpx.Limits(max_connections=args.concurrency) + async with httpx.AsyncClient(timeout=httpx.Timeout(args.timeout), limits=limits) as client: + async def one(i, prompt): + async with sem: + t0 = time.time() + out = await run_completion(client, args.base_url, args.model, + prompt, args.max_tokens, args.temperature, name) + return i, out, round(time.time() - t0, 2) + + t0 = time.time() + results = {} + tasks = [one(i, p) for i, p in enumerate(prompts)] + done = 0 + for fut in asyncio.as_completed(tasks): + i, out, dt = await fut + results[i] = out + done += 1 + if done % 200 == 0 or done == len(prompts): + print(f'[{name}] {done}/{len(prompts)} ' + f'({(time.time() - t0) / 60:.1f} min)', flush=True) + + preds, fails, trunc = [], 0, 0 + records = [] + for i in range(len(prompts)): + out = results[i] + text = out['content'] + if not text and out['reasoning'] and not out['reasoning'].startswith('__ERROR__'): + trunc += 1 # hit max_tokens mid-thinking; no final answer produced + if name == 'ceval': + pred = first_capital(text) + elif name == 'mmlu_pro': + pred = mmlupro_extract(text).lower() + elif name == 'mmlu_redux': + m = MMLUREDUX_RE.search(text) + pred = m.group(1) if m else '' + else: + pred = sg_extract_labels(text) + if pred is None: + content = sg_extract_content(text, items[i]['options']) + if content is not None: + try: + pred = chr(items[i]['options'].index(content) + 65) + except ValueError: + pred = None + if pred is None: + pred = '' + correct = pred.lower() == golds[i].lower() + if out['reasoning'].startswith('__ERROR__'): + fails += 1 + extra = {'subject': items[i].get('subject')} + if name == 'supergpqa': + extra.update({'discipline': items[i].get('discipline'), + 'field': items[i].get('field'), + 'difficulty': items[i].get('difficulty')}) + records.append({ + 'idx': i, 'prompt': prompts[i], 'reasoning': out['reasoning'], + 'output': text, 'gold': golds[i], 'pred': pred, 'correct': correct, + 'completion_tokens': out.get('completion_tokens', 0), + 'finish_reason': out.get('finish_reason', ''), + **extra, + }) + preds.append(correct) + acc = sum(preds) / max(len(preds), 1) + wall_min = (time.time() - t0) / 60.0 + tot_completion = sum(results[i].get('completion_tokens', 0) for i in results) + out_path = Path(out_dir) + out_path.mkdir(parents=True, exist_ok=True) + (out_path / f'{name}_samples.json').write_text(json.dumps(records, ensure_ascii=False, indent=1)) + summary = {'benchmark': name, 'model': args.model, 'n': len(preds), + 'acc': round(acc * 100, 2), 'api_errors': fails, + 'truncated_thinking': trunc, + 'max_tokens': args.max_tokens, 'temperature': args.temperature, + 'concurrency': args.concurrency, + 'sample': args.sample or None, 'seed': args.seed, + 'wall_min': round(wall_min, 2), + 'completion_tokens_total': tot_completion, + 'agg_toks_per_s': round(tot_completion / (wall_min * 60), 1) if wall_min else 0} + (out_path / f'{name}_summary.json').write_text(json.dumps(summary, ensure_ascii=False, indent=1)) + print(json.dumps(summary, ensure_ascii=False)) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('benchmark', choices=['ceval', 'supergpqa', 'mmlu_pro', 'mmlu_redux']) + ap.add_argument('--base-url', default='http://127.0.0.1:18082/v1') + ap.add_argument('--model', default='qwen35-27b-tp2') + ap.add_argument('--split') + ap.add_argument('--limit', type=int, default=0, + help='take first N items per benchmark (smoke only)') + ap.add_argument('--sample', type=int, default=0, + help='stratified random subsample of N items ' + '(strata: subject/category/discipline)') + ap.add_argument('--seed', type=int, default=1337) + ap.add_argument('--max-tokens', type=int, default=4096) + ap.add_argument('--temperature', type=float, default=0.0) + ap.add_argument('--concurrency', type=int, default=16) + ap.add_argument('--timeout', type=float, default=3600.0) + ap.add_argument('--out-dir', default='results/qwen35-27b-tp2-eval') + args = ap.parse_args() + + if args.benchmark == 'ceval': + split = args.split or 'val' + items, prompts, golds = [], [], [] + for subject, cn in CEVAL_SUBJECT_CN.items(): + dev = datasets.load_dataset('ceval/ceval-exam', subject, split='dev') + split_ds = datasets.load_dataset('ceval/ceval-exam', subject, split=split) + shots = list(dev)[:5] + for item in split_ds: + item = dict(item) + item['subject'] = subject + items.append(item) + prompts.append(ceval_prompt(cn, item, shots)) + golds.append(item.get('answer', '')) + elif args.benchmark == 'mmlu_pro': + test = datasets.load_dataset('TIGER-Lab/MMLU-Pro', split='test') + shots_by_cat = {} + for eg in datasets.load_dataset('TIGER-Lab/MMLU-Pro', split='validation'): + shots_by_cat.setdefault(eg['category'], []).append(eg) + items, prompts, golds = [], [], [] + for item in test: + cat = item['category'] + head = ('The following are multiple choice questions (with answers) ' + f'about {cat}. Think step by step and then finish your answer ' + 'with "the answer is (X)" where X is the correct letter choice.\n') + shots = ''.join(mmlupro_format(s, True) for s in shots_by_cat.get(cat, [])[:5]) + prompts.append(head + '\n' + shots + mmlupro_format(item, False)) + items.append({'subject': cat}) + golds.append(MMLUPRO_LETTERS[item['answer_index']] if isinstance( + item['answer_index'], int) else item['answer']) + elif args.benchmark == 'mmlu_redux': + split = args.split or 'test' + subjects = datasets.get_dataset_config_names('fxmarty/mmlu-redux-2.0-ok') + items, prompts, golds = [], [], [] + for subj in subjects: + ds = datasets.load_dataset('fxmarty/mmlu-redux-2.0-ok', subj, split=split) + desc = ('The following are multiple choice questions (with answers) ' + f"about {subj.replace('_', ' ')}.\n\n") + for item in ds: + prompt = (desc + item['question'].strip() + + f"\nA. {item['choices'][0]}\nB. {item['choices'][1]}" + f"\nC. {item['choices'][2]}\nD. {item['choices'][3]}" + '\nPlease respond with the correct letter (A, B, C or D) ' + 'without any additional comments, only the correct letter:') + prompts.append(prompt) + items.append({'subject': subj}) + golds.append('ABCD'[item['answer']]) + else: + ds = datasets.load_dataset('m-a-p/SuperGPQA', split='train') + items, prompts, golds = [], [], [] + for item in ds: + items.append(item) + prompts.append(SG_PROMPT.format(sg_build_questions(item))) + golds.append(item['answer_letter']) + + if args.limit or args.sample: + import random as rnd + indices = list(range(len(items))) + if args.sample and args.sample < len(items): + strata = {} + for idx in indices: + key = items[idx].get('subject') or items[idx].get('discipline') or '_' + strata.setdefault(key, []).append(idx) + rng = rnd.Random(args.seed) + for group in strata.values(): + rng.shuffle(group) + picked = [] + for key in sorted(strata): + group = strata[key] + take = max(1, round(args.sample * len(group) / len(items))) + take = min(take, len(group)) + picked.extend(sorted(group[:take])) + indices = sorted(picked) + print(f'strata: {len(strata)}, picked {len(indices)} of {len(items)}', flush=True) + else: + indices = indices[:args.limit] + items = [items[i] for i in indices] + prompts = [prompts[i] for i in indices] + golds = [golds[i] for i in indices] + print(f'{args.benchmark}: {len(prompts)} samples', flush=True) + asyncio.run(evaluate(args.benchmark, items, prompts, golds, args, args.out_dir)) + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/scripts/eval_rerun_truncated.py b/scripts/eval_rerun_truncated.py new file mode 100644 index 000000000..44736291b --- /dev/null +++ b/scripts/eval_rerun_truncated.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +"""Re-run items truncated by the thinking-token budget with a larger cap. + +Reads /_samples.json, finds rows with empty final output +(finish_reason == 'length' / 'stop' with empty content — mid-thinking ends), +re-generates them against the server with a bigger --max-tokens, merges the +new outputs back, and writes: + /_samples_merged.json + /_summary.json (updated acc, with truncation stats) +""" +import argparse +import asyncio +import json +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent)) +from eval_mc import (first_capital, mmlupro_extract, MMLUREDUX_RE, + sg_extract_labels, sg_extract_content, STOP_MAP) +import httpx + + +def extract(bench, item, text): + if bench == 'ceval': + return first_capital(text) + if bench == 'mmlu_pro': + return mmlupro_extract(text).lower() + if bench == 'mmlu_redux': + m = MMLUREDUX_RE.search(text) + return m.group(1) if m else '' + pred = sg_extract_labels(text) + if pred is None and item.get('options'): + content = sg_extract_content(text, item['options']) + if content is not None: + try: + pred = chr(item['options'].index(content) + 65) + except ValueError: + pred = None + return pred or '' + + +async def run(base_url, model, rows, bench, max_tokens, concurrency, timeout): + sem = asyncio.Semaphore(concurrency) + async with httpx.AsyncClient(timeout=httpx.Timeout(timeout)) as client: + async def one(row): + async with sem: + payload = {'model': model, 'max_tokens': max_tokens, + 'temperature': 0.0, 'top_p': 1.0, + 'messages': [{'role': 'user', 'content': row['prompt']}]} + if STOP_MAP.get(bench): + payload['stop'] = STOP_MAP[bench] + for attempt in range(6): + try: + r = await client.post(f'{base_url}/chat/completions', json=payload) + if r.status_code == 200: + data = r.json() + msg = data['choices'][0]['message'] + usage = data.get('usage') or {} + row = dict(row) + row['output'] = msg.get('content') or '' + row['reasoning'] = msg.get('reasoning') or '' + row['completion_tokens'] = usage.get('completion_tokens') or 0 + row['finish_reason'] = data['choices'][0].get('finish_reason') or '' + return row + err = f'HTTP {r.status_code} {r.text[:200]!r}' + except Exception as e: # noqa: BLE001 + err = f'{type(e).__name__}: {e!s}' + await asyncio.sleep(min(2 ** attempt, 30)) + raise RuntimeError(f"rerun failed: {err}") + + return await asyncio.gather(*[one(r) for r in rows]) + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument('benchmark') + ap.add_argument('--base-url', default='http://127.0.0.1:18082/v1') + ap.add_argument('--model', default='qwen35-27b-tp2') + ap.add_argument('--out-dir', default='results/qwen35-27b-tp2-eval') + ap.add_argument('--max-tokens', type=int, default=32768) + ap.add_argument('--concurrency', type=int, default=48) + ap.add_argument('--timeout', type=float, default=7200.0) + args = ap.parse_args() + + out = Path(args.out_dir) + samples = json.loads((out / f'{args.benchmark}_samples.json').read_text()) + summary = json.loads((out / f'{args.benchmark}_summary.json').read_text()) + + bad = [s for s in samples if not s['output']] + print(f"{args.benchmark}: {len(samples)} total, {len(bad)} to re-run with " + f"max_tokens={args.max_tokens}", flush=True) + if not bad: + return + fixed = asyncio.run(run(args.base_url, args.model, bad, args.benchmark, + args.max_tokens, args.concurrency, args.timeout)) + fixed_by_idx = {f['idx']: f for f in fixed} + merged = [] + for s in samples: + m = fixed_by_idx.get(s['idx'], s) + m['pred'] = extract(args.benchmark, m, m['output']) + m['correct'] = m['pred'].lower() == m['gold'].lower() + merged.append(m) + n_ok = sum(1 for m in merged if m['correct']) + n_trunc_left = sum(1 for m in merged if not m['output']) + acc = round(100.0 * n_ok / len(merged), 2) + summary.update({ + 'acc_merged': acc, + 'rerun_max_tokens': args.max_tokens, + 'rerun_n': len(bad), + 'still_truncated': n_trunc_left, + 'completion_tokens_total_merged': sum(m.get('completion_tokens', 0) for m in merged), + }) + (out / f'{args.benchmark}_samples_merged.json').write_text( + json.dumps(merged, ensure_ascii=False, indent=1)) + (out / f'{args.benchmark}_summary.json').write_text( + json.dumps(summary, ensure_ascii=False, indent=1)) + print(json.dumps({k: summary[k] for k in ('benchmark', 'acc', 'acc_merged', + 'rerun_n', 'still_truncated')}, + ensure_ascii=False)) + + +if __name__ == '__main__': + main() From 2d33c9bf7380b989f10376d1c93dac3d6cdaf8c3 Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Sun, 30 Aug 2026 12:55:12 +0000 Subject: [PATCH 13/15] fix(bench): mmlu-redux extracts answers by markers, not first capital The MMLU-Redux extractor searched the whole completion for the first [ABCD], so any answer prefixed with "Answer: B" scored as the "A" in "Answer", silently marking correct B/C/D responses wrong (codex review on #946). Reuse the SuperGPQA marker-preferring patterns restricted to A-D, and drop the now-dead MMLUREDUX_RE. Redo the MMLU-Redux numbers before citing them. Signed-off-by: Ziyang Zhang --- scripts/eval_mc.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/scripts/eval_mc.py b/scripts/eval_mc.py index c879751b0..51563085b 100644 --- a/scripts/eval_mc.py +++ b/scripts/eval_mc.py @@ -103,7 +103,9 @@ def mmlupro_extract(text): # ---------------------------------------------------------------- mmlu-redux -MMLUREDUX_RE = re.compile(r'([ABCD])') +# Extraction reuses the SuperGPQA marker-preferring patterns with the MMLU +# letter range: a first-anywhere `[ABCD]` search scores "Answer: B" as the +# "A" of "Answer", silently flipping correct answers (codex review on #946). # ---------------------------------------------------------------- supergpqa @@ -240,8 +242,7 @@ async def one(i, prompt): elif name == 'mmlu_pro': pred = mmlupro_extract(text).lower() elif name == 'mmlu_redux': - m = MMLUREDUX_RE.search(text) - pred = m.group(1) if m else '' + pred = sg_extract_labels(text, 'ABCD') or '' else: pred = sg_extract_labels(text) if pred is None: From e404290c002dae1a9da1465f56c970e4215addd9 Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Sun, 30 Aug 2026 12:55:30 +0000 Subject: [PATCH 14/15] docs(bench): flag the pre-fix MMLU-Redux figure as rerun-pending The 94.09 snapshot predates the answer-extraction fix; mark it as not citable until rerun with the marker-preferring extractor. Signed-off-by: Ziyang Zhang --- docs/benchmarks/qwen35-27b-tp2-knowledge-eval.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/docs/benchmarks/qwen35-27b-tp2-knowledge-eval.md b/docs/benchmarks/qwen35-27b-tp2-knowledge-eval.md index 310e19546..cffd0d932 100644 --- a/docs/benchmarks/qwen35-27b-tp2-knowledge-eval.md +++ b/docs/benchmarks/qwen35-27b-tp2-knowledge-eval.md @@ -3,6 +3,8 @@ > TL;DR:Qwen3.5-27B 在 pegainfer TP2(2× RTX 4090,batched eager decode)上跑知识基准,C-Eval 88.11(官方 90.5)、MMLU-Redux 94.09(官方 93.2),均在跨 harness 正常带内;MMLU-Pro / SuperGPQA 因运行时长原因仅完成抽样冒烟,未出最终分(见下文)。模型数值无 TP 引入的精度问题。 > > 注:分数实测于 rebase 前的 f4c66780 分支(自研 Phase 1/2a 线);rebase 到 #870 后 logits golden gate 两侧一致通过,数值可迁移,但若正式引用请在本 PR 分支上复跑确认。 +> +> 另注:**MMLU-Redux 94.09 是在旧抽取器下实测的,在复跑前不要引用**。旧抽取取全文首个 `[ABCD]`,"Answer: B" 会被算成 "Answer" 里的 A(codex review 发现);本 PR 已修复为 marker 优先 + 独立字母抽取,需以修复后的 `scripts/eval_mc.py` 在 exact head 上重测。 ## 环境 @@ -11,7 +13,7 @@ - 采样:temperature=0.0,top_p=1.0,chat completions(thinking 模式,即模板默认行为) - 评测器:`scripts/eval_mc.py`(自研,统一 `/v1/chat/completions` 并发 48),配方逐项复刻官方 harness: - **C-Eval** = OpenCompass `ceval_gen`:52 学科 val split 全量 1346 题,dev split 5-shot,"答案: " 续写,首大写字母抽取 - - **MMLU-Redux** = lm-eval `mmlu_redux_generative`:`fxmarty/mmlu-redux-2.0-ok` 57 学科 test 全量 5330 题,0-shot,首个 `[ABCD]` 抽取 + - **MMLU-Redux** = lm-eval `mmlu_redux_generative`:`fxmarty/mmlu-redux-2.0-ok` 57 学科 test 全量 5330 题,0-shot,marker/独立字母抽取(早期版本为全文首个 `[ABCD]`,会误吸 "Answer: X" 前缀,已修复待复跑) - **MMLU-Pro** = lm-eval `mmlu_pro`:TIGER-Lab/MMLU-Pro test,validation split 5-shot CoT,`answer is (X)` 抽取 - **SuperGPQA** = OpenCompass `supergpqa_gen`:`m-a-p/SuperGPQA` train 26529 题,0-shot,"Answer: X" 字母/内容两层抽取 - 启动命令:`LD_LIBRARY_PATH=<.venv>/nvidia/nccl/lib ./target/release/pegainfer --model-path <27B> --served-model-name qwen35-27b-tp2 --tp-size 2 --cuda-graph false --port 18082` From 7a3f54380191933919f9003a0b31898856ee8ea1 Mon Sep 17 00:00:00 2001 From: Ziyang Zhang Date: Sat, 5 Sep 2026 10:16:22 +0000 Subject: [PATCH 15/15] docs(qwen35): derive value-head TP divisibility from the key-head guard The second LocalGeometry guard was dropped in review: Config35 already validates linear_num_value_heads % linear_num_key_heads == 0, so the key-head % tp guard implies value-head divisibility. Align the P2/P2b requirement wording with the single source of truth. Signed-off-by: Ziyang Zhang --- docs/models/qwen35/tp-design.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/models/qwen35/tp-design.md b/docs/models/qwen35/tp-design.md index f6b2b960d..906769258 100644 --- a/docs/models/qwen35/tp-design.md +++ b/docs/models/qwen35/tp-design.md @@ -110,7 +110,7 @@ For any candidate `tp`, require: - `num_attention_heads % tp == 0` - `num_key_value_heads % tp == 0` - `intermediate_size % tp == 0` -- Phase 2 additionally requires `linear_num_key_heads % tp == 0` and `linear_num_value_heads % tp == 0` +- Phase 2 requires `linear_num_key_heads % tp == 0`; `linear_num_value_heads % tp == 0` then follows from the checkpoint invariant `linear_num_value_heads % linear_num_key_heads == 0`, so no second runtime guard Full attention local dimensions: @@ -233,7 +233,7 @@ Lifecycle observability, cancellation ordering, fail-closed cleanup, and unified ## P2b: Local-Head Linear Attention / GDR -P2b converts the 24 linear-attention layers from replicated execution to true TP execution. It additionally requires `linear_num_key_heads % tp == 0` and `linear_num_value_heads % tp == 0`; unsupported degrees and unsupported local kernel shapes fail before model loading. +P2b converts the 24 linear-attention layers from replicated execution to true TP execution. It requires `linear_num_key_heads % tp == 0` (value-head divisibility follows from the checkpoint invariant `linear_num_value_heads % linear_num_key_heads == 0`); unsupported degrees and unsupported local kernel shapes fail before model loading. Shard every head-indexed linear-attention/GDR surface by the local key/value-head ranges: