From 4c9e04614d2c4ba5d1444b75b91699fdceb68b6f Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Thu, 2 Jul 2026 23:57:36 -0400 Subject: [PATCH 01/15] =?UTF-8?q?wip(gdn-tp):=20GDN=20HeadParallel=20MVP?= =?UTF-8?q?=20=E2=80=94=20config-localize=20+=20segmented=20slicers=20+=20?= =?UTF-8?q?post-out=5Fproj=20all-reduce=20(cargo=20check=20green)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../src/layers/qwen3_ssm/ssm_forward.rs | 8 + .../layers/qwen3_ssm/trait_decode_batched.rs | 4 + .../trait_decode_multi_seq/ssm_batched.rs | 4 + .../src/layers/qwen3_ssm/trait_prefill.rs | 3 + .../layers/qwen3_ssm/trait_prefill_helper.rs | 31 ++ .../layers/qwen3_ssm/trait_prefill_phase3.rs | 3 + crates/spark-model/src/tp_shard.rs | 347 ++++++++++++++++++ crates/spark-model/src/tp_shard/tests.rs | 194 ++++++++++ .../spark-model/src/weight_loader/qwen35.rs | 14 +- .../qwen35/load_layers/linear_attn_arms.rs | 167 ++++++--- .../src/main_modules/serve_phases/topology.rs | 28 +- scripts/gdn_tp_golden.py | 294 +++++++++++++++ 12 files changed, 1048 insertions(+), 49 deletions(-) create mode 100755 scripts/gdn_tp_golden.py diff --git a/crates/spark-model/src/layers/qwen3_ssm/ssm_forward.rs b/crates/spark-model/src/layers/qwen3_ssm/ssm_forward.rs index d998a3e32..106c357ea 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/ssm_forward.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/ssm_forward.rs @@ -440,6 +440,14 @@ impl Qwen3SsmLayer { Self::debug_bf16(ctx.gpu, "out-proj", out, 4); } + // GDN HeadParallel: `out` is this rank's PARTIAL row-parallel out_proj + // over its local value heads. Reduce across TP ranks to the complete + // SSM output before the caller's residual add. Single-token path + // (dense_gemv / w8a16_gemv / w4a16_gemv above → one position), so + // num_tokens = 1. No-op at tp=1. Covers single-token decode + // (trait_decode) and per-sequence multi-seq decode (trait_decode_multi_seq). + self.ssm_tp_all_reduce(out, 1, ctx, stream)?; + Ok(out) } } diff --git a/crates/spark-model/src/layers/qwen3_ssm/trait_decode_batched.rs b/crates/spark-model/src/layers/qwen3_ssm/trait_decode_batched.rs index acc82364c..8d3b2c658 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/trait_decode_batched.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/trait_decode_batched.rs @@ -533,6 +533,10 @@ impl Qwen3SsmLayer { )?; } + // GDN HeadParallel: reduce the row-parallel partial out_proj across TP + // ranks (num_tokens × h BF16) before the residual add. No-op at tp=1. + self.ssm_tp_all_reduce(out_proj_buf, num_tokens, ctx, stream)?; + // ── 10. Batched residual + post-norm, then MoE + residual ── // residual_add_rms_norm supports multi-token (grid.x = num_tokens) let normed2_base = ctx.buffers.norm_output(); diff --git a/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched.rs b/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched.rs index 3d3be634c..201985dc8 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched.rs @@ -307,6 +307,10 @@ impl Qwen3SsmLayer { } detail_step!("out_proj"); + // GDN HeadParallel: reduce the row-parallel partial out_proj across TP + // ranks (n × h BF16) before the residual add. No-op at tp=1. + self.ssm_tp_all_reduce(ssm_out_base, n, ctx, stream)?; + // ── 5. Batched residual add + post-attn RMS norm → norm_output[0..n] ── ops::residual_add_rms_norm( ctx.gpu, diff --git a/crates/spark-model/src/layers/qwen3_ssm/trait_prefill.rs b/crates/spark-model/src/layers/qwen3_ssm/trait_prefill.rs index 420a9d351..91dae483b 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/trait_prefill.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/trait_prefill.rs @@ -376,6 +376,9 @@ impl Qwen3SsmLayer { // ── 10. Output projection GEMM: [N, 4096] × [4096, 2048] → [N, 2048] ── let out_proj_buf = ctx.buffers.moe_output(); self.prefill_out_proj_dispatch(ctx, normed_out_buf, out_proj_buf, k, h, value_dim, stream)?; + // GDN HeadParallel: reduce the row-parallel partial out_proj across TP + // ranks (num_tokens × h BF16) before the residual add. No-op at tp=1. + self.ssm_tp_all_reduce(out_proj_buf, num_tokens, ctx, stream)?; // ATLAS_GDN_DUMP hook: SSM out_proj output — drift attribution. super::debug::maybe_dump_gdn_buf( ctx.gpu, diff --git a/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_helper.rs b/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_helper.rs index 1d37f3683..41f45c542 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_helper.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_helper.rs @@ -16,6 +16,37 @@ use crate::layer::ForwardContext; use crate::layers::ops; impl Qwen3SsmLayer { + /// GDN HeadParallel tensor-parallel all-reduce of the row-parallel + /// `out_proj` output. + /// + /// Each TP rank ran the GDN scan over its LOCAL value-head slice and + /// projected with the row-parallel `out_proj` (columns = local value_dim), + /// so its `[num_tokens, h]` BF16 buffer holds a PARTIAL sum over the full + /// hidden dim. Summing across ranks reconstructs the complete SSM output, + /// exactly mirroring attention's post-`o_proj` reduce + /// (`qwen3_attention/trait_impl/decode_inner.rs`). Must run BEFORE the + /// residual add / post-norm that consumes `out_proj_buf`. + /// + /// No-op when `tp_world_size == 1` or no communicator is present (single + /// GPU, or a path that already holds the complete output). + pub(super) fn ssm_tp_all_reduce( + &self, + out_proj_buf: DevicePtr, + num_tokens: usize, + ctx: &ForwardContext, + stream: u64, + ) -> Result<()> { + if ctx.config.tp_world_size > 1 + && let Some(comm) = ctx.comm + { + // BF16 [num_tokens, hidden_size] — same byte count attention + // reduces after o_proj. + let bytes = num_tokens * ctx.config.hidden_size * 2; + comm.all_reduce_async(out_proj_buf.0, bytes, stream)?; + } + Ok(()) + } + pub(super) fn prefill_out_proj_dispatch( &self, ctx: &ForwardContext, diff --git a/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_phase3.rs b/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_phase3.rs index 9317d5322..ce6efa755 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_phase3.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/trait_prefill_phase3.rs @@ -57,6 +57,9 @@ impl Qwen3SsmLayer { // that on co-dispatched requests while single-stream (which already used // this helper) ran NVFP4. Routing through it equalises the two paths. self.prefill_out_proj_dispatch(ctx, normed_out_buf, out_proj_buf, k, h, value_dim, stream)?; + // GDN HeadParallel: reduce the row-parallel partial out_proj across TP + // ranks (num_tokens × h BF16) before the residual add. No-op at tp=1. + self.ssm_tp_all_reduce(out_proj_buf, num_tokens, ctx, stream)?; // ── 11. Batched residual + post-norm + MoE ── ops::residual_add_rms_norm( diff --git a/crates/spark-model/src/tp_shard.rs b/crates/spark-model/src/tp_shard.rs index 57141e3ef..e8537a38e 100644 --- a/crates/spark-model/src/tp_shard.rs +++ b/crates/spark-model/src/tp_shard.rs @@ -327,6 +327,353 @@ impl TpMoeDims { } } +// ════════════════════════════════════════════════════════════════════ +// GDN HeadParallel — tensor-parallel sharding of the Gated-DeltaNet +// (linear_attention / SSM) layers for Qwen3.5 / 3.6. +// +// The GDN recurrence is embarrassingly parallel across *value-head groups*: +// each TP rank owns a contiguous range of key/value heads, runs the whole +// scan locally with LOCAL nk/nv/conv_dim, and the ranks reconcile with a +// single all-reduce after `out_proj` (row-parallel, exactly like attention +// after `o_proj`). No cross-rank comm inside the scan. +// +// CRITICAL LAYOUT FACT — the in-projection is stored as *segmented* +// contiguous blocks, NOT one flat matrix: +// +// in_proj_qkv : [Q | K | V] rows = nk·kd + nk·kd + nv·vd (= conv_dim) +// in_proj_z : [Z] rows = nv·vd +// → gpu_concat_rows → QKVZ [Q | K | V | Z] +// +// A naive "first out_dim/tp rows" slice is WRONG: it would give rank 0 the +// whole Q block plus part of K. Each segment (Q, K, V, Z) must be sliced by +// the LOCAL head range *independently*, then the local slices re-concatenated +// in the same [Q|K|V|Z] order. The `segment_copy_plan` below encodes exactly +// that: one contiguous copy per segment, packed back-to-back into the local +// buffer. +// +// The depthwise `conv1d` weight `[conv_dim, d_conv]` is sharded with the SAME +// segment pattern as QKV (its channels ARE the QKV channels — one filter per +// channel), NOT replicated. `a_log`/`dt_bias` (`[nv]` FP32), `norm` +// (`[nv·vd]` BF16) and `out_proj` (`[h, nv·vd]`, row-parallel) shard on the +// value-head axis. The BA gate buffer is per-group interleaved but the rank +// boundary always lands on a group boundary, so it slices contiguously. +// ════════════════════════════════════════════════════════════════════ + +/// Pre-TP-shard GDN (linear-attention / SSM) dimensions reconstructed from +/// `config`. +/// +/// Mirrors [`TpAttentionDims`]: `topology.rs` divides +/// `linear_num_key_heads` / `linear_num_value_heads` by `tp_world_size` at +/// startup, so by the time a loader runs `config` holds **per-rank-local** +/// head counts. The `full_*` fields multiply back up to the pre-shard sizes +/// that the segment slicers expect. Head *dims* (`kd`, `vd`) and the hidden +/// size `h` are never sharded. +#[derive(Debug, Clone, Copy)] +pub struct TpGdnDims { + pub tp_rank: usize, + /// `tp_world_size` clamped to `>= 1`. Loaders treat `tp_size == 1` as the + /// no-shard fast path. + pub tp_size: usize, + /// Hidden size (model embed dim) — never sharded. + pub h: usize, + /// Key head dim (`linear_key_head_dim`) — never sharded. + pub kd: usize, + /// Value head dim (`linear_value_head_dim`) — never sharded. + pub vd: usize, + /// Per-rank key heads (Q and K share this count). + pub local_nk: usize, + /// Full pre-shard key heads = `local_nk * tp_size`. + pub full_nk: usize, + /// Per-rank value heads. + pub local_nv: usize, + /// Full pre-shard value heads = `local_nv * tp_size`. + pub full_nv: usize, +} + +impl TpGdnDims { + pub fn from_config(config: &ModelConfig) -> Self { + let tp_size = config.tp_world_size.max(1); + let local_nk = config.linear_num_key_heads; + let local_nv = config.linear_num_value_heads; + Self { + tp_rank: config.tp_rank, + tp_size, + h: config.hidden_size, + kd: config.linear_key_head_dim, + vd: config.linear_value_head_dim, + local_nk, + full_nk: local_nk * tp_size, + local_nv, + full_nv: local_nv * tp_size, + } + } + + /// Full (pre-shard) key projection width: `full_nk * kd`. + pub fn full_key_dim(&self) -> usize { + self.full_nk * self.kd + } + /// Local key projection width: `local_nk * kd`. + pub fn local_key_dim(&self) -> usize { + self.local_nk * self.kd + } + /// Full (pre-shard) value projection width: `full_nv * vd`. + pub fn full_value_dim(&self) -> usize { + self.full_nv * self.vd + } + /// Local value projection width: `local_nv * vd`. + pub fn local_value_dim(&self) -> usize { + self.local_nv * self.vd + } + /// Full conv / QKV width: `2*full_nk*kd + full_nv*vd`. + pub fn full_conv_dim(&self) -> usize { + 2 * self.full_key_dim() + self.full_value_dim() + } + /// Local conv / QKV width: `2*local_nk*kd + local_nv*vd`. + pub fn local_conv_dim(&self) -> usize { + 2 * self.local_key_dim() + self.local_value_dim() + } + /// Full QKVZ out dim: `2*full_nk*kd + 2*full_nv*vd`. + pub fn full_qkvz_out(&self) -> usize { + self.full_conv_dim() + self.full_value_dim() + } + /// Local QKVZ out dim: `2*local_nk*kd + 2*local_nv*vd`. + pub fn local_qkvz_out(&self) -> usize { + self.local_conv_dim() + self.local_value_dim() + } + + /// Full-row segment list for the `[Q|K|V]` in-projection. + fn qkv_segments(&self) -> [usize; 3] { + [self.full_key_dim(), self.full_key_dim(), self.full_value_dim()] + } + /// Full-row segment list for the concatenated `[Q|K|V|Z]` in-projection. + fn qkvz_segments(&self) -> [usize; 4] { + [ + self.full_key_dim(), + self.full_key_dim(), + self.full_value_dim(), + self.full_value_dim(), + ] + } +} + +/// A single device-to-device copy in a segmented-slice plan. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct CopyOp { + src_off: usize, + dst_off: usize, + len: usize, +} + +/// Build the copy plan for a SEGMENTED row-slice. +/// +/// `segments` lists the full (pre-shard) row count of each contiguous block +/// (Q, K, V[, Z] for QKVZ). Each block is sliced *independently* to the local +/// rank's head range `[tp_rank * seg/tp, (tp_rank+1) * seg/tp)` and the local +/// slices are packed back-to-back into the output buffer, preserving segment +/// order. `row_bytes` is the byte width of one row (`in_dim * elem_bytes`). +/// +/// Returns `(ops, local_total_rows)`. Every segment must be divisible by +/// `tp_size` — the caller has already reconstructed `full_*` as +/// `local_* * tp_size`, so this holds by construction, but it is checked to +/// fail loudly on a mis-wired config rather than silently corrupt heads. +fn segment_copy_plan( + segments: &[usize], + row_bytes: usize, + tp_rank: usize, + tp_size: usize, +) -> Result<(Vec, usize)> { + ensure!(tp_rank < tp_size, "tp_rank {tp_rank} >= tp_size {tp_size}"); + let mut ops = Vec::with_capacity(segments.len()); + let mut src_rows = 0usize; // running offset into the source (in rows) + let mut dst_rows = 0usize; // running offset into the packed dst (in rows) + for (i, &seg) in segments.iter().enumerate() { + ensure!( + seg.is_multiple_of(tp_size), + "segment {i} ({seg} rows) not divisible by tp_size {tp_size}", + ); + let local = seg / tp_size; + ops.push(CopyOp { + src_off: (src_rows + tp_rank * local) * row_bytes, + dst_off: dst_rows * row_bytes, + len: local * row_bytes, + }); + src_rows += seg; + dst_rows += local; + } + Ok((ops, dst_rows)) +} + +/// Execute a segmented row-slice on the GPU. `row_elems` is the number of +/// elements per row (`in_dim`); `elem_bytes` its size (2 = BF16, 4 = FP32). +/// Returns `(local_ptr, local_total_rows)`. For `tp_size <= 1` returns the +/// source untouched (no allocation, caller must not double-free). +fn slice_segments( + src: DevicePtr, + segments: &[usize], + row_elems: usize, + elem_bytes: usize, + tp_rank: usize, + tp_size: usize, + gpu: &dyn GpuBackend, +) -> Result<(DevicePtr, usize)> { + let full_rows: usize = segments.iter().sum(); + if tp_size <= 1 { + return Ok((src, full_rows)); + } + let row_bytes = row_elems * elem_bytes; + let (ops, local_rows) = segment_copy_plan(segments, row_bytes, tp_rank, tp_size)?; + let dst = gpu.alloc(local_rows * row_bytes)?; + for op in &ops { + gpu.copy_d2d(src.offset(op.src_off), dst.offset(op.dst_off), op.len)?; + } + Ok((dst, local_rows)) +} + +/// Shard the `[Q|K|V]` (`in_proj_qkv`) BF16 weight `[full_conv_dim, h]` to the +/// local rank's `[local_conv_dim, h]`, slicing Q, K and V independently by the +/// local head range. Returns `(ptr, local_rows, h)`. +pub fn shard_gdn_qkv_rows( + src: DevicePtr, + dims: &TpGdnDims, + gpu: &dyn GpuBackend, +) -> Result<(DevicePtr, usize, usize)> { + let (ptr, rows) = slice_segments( + src, + &dims.qkv_segments(), + dims.h, + BF16_BYTES, + dims.tp_rank, + dims.tp_size, + gpu, + )?; + Ok((ptr, rows, dims.h)) +} + +/// Shard the concatenated `[Q|K|V|Z]` (`in_proj_qkvz`) BF16 weight +/// `[full_qkvz_out, h]` to the local rank's `[local_qkvz_out, h]`, slicing all +/// four segments independently. Returns `(ptr, local_rows, h)`. +pub fn shard_gdn_qkvz_rows( + src: DevicePtr, + dims: &TpGdnDims, + gpu: &dyn GpuBackend, +) -> Result<(DevicePtr, usize, usize)> { + let (ptr, rows) = slice_segments( + src, + &dims.qkvz_segments(), + dims.h, + BF16_BYTES, + dims.tp_rank, + dims.tp_size, + gpu, + )?; + Ok((ptr, rows, dims.h)) +} + +/// Shard the BA gate BF16 weight `[2*full_nv, h]` to `[2*local_nv, h]`. +/// +/// The interleave is per key-head group (`[β₀..β_{vpg-1}, α₀..α_{vpg-1}]` per +/// group, `vpg = nv/nk`), but rank `r` owns key-head groups +/// `[r*local_nk, (r+1)*local_nk)` which map to the contiguous row range +/// `[r*2*local_nv, (r+1)*2*local_nv)` — the rank boundary always lands on a +/// group boundary, so a single contiguous slice preserves the interleave. +pub fn shard_gdn_ba_rows( + src: DevicePtr, + dims: &TpGdnDims, + gpu: &dyn GpuBackend, +) -> Result<(DevicePtr, usize, usize)> { + // Group-boundary alignment guarantee: full_nk divisible by tp_size ⇒ + // each rank gets whole groups. + ensure!( + dims.full_nk.is_multiple_of(dims.tp_size), + "BA: full_nk {} not divisible by tp_size {}", + dims.full_nk, + dims.tp_size, + ); + let (ptr, rows) = slice_segments( + src, + &[2 * dims.full_nv], + dims.h, + BF16_BYTES, + dims.tp_rank, + dims.tp_size, + gpu, + )?; + Ok((ptr, rows, dims.h)) +} + +/// Shard the depthwise `conv1d` BF16 weight `[full_conv_dim, d_conv]` to +/// `[local_conv_dim, d_conv]`. Channels ARE the QKV channels (one filter per +/// channel), so this uses the SAME `[Q|K|V]` segment pattern as the QKV +/// in-projection — the conv is NOT replicated across ranks. +pub fn shard_gdn_conv_rows( + src: DevicePtr, + dims: &TpGdnDims, + d_conv: usize, + gpu: &dyn GpuBackend, +) -> Result<(DevicePtr, usize, usize)> { + let (ptr, rows) = slice_segments( + src, + &dims.qkv_segments(), + d_conv, + BF16_BYTES, + dims.tp_rank, + dims.tp_size, + gpu, + )?; + Ok((ptr, rows, d_conv)) +} + +/// Shard a per-value-head 1D vector on the value-head axis. Handles BF16 +/// (`norm`, `[full_nv*vd]` → `[local_nv*vd]` with `elem_bytes = 2`, +/// `unit = vd`) and FP32 (`a_log` / `dt_bias`, `[full_nv]` → `[local_nv]` with +/// `elem_bytes = 4`, `unit = 1`). `unit` is the number of elements per value +/// head. Returns `(ptr, local_len_elems)`. +pub fn shard_gdn_value_vector( + src: DevicePtr, + dims: &TpGdnDims, + unit: usize, + elem_bytes: usize, + gpu: &dyn GpuBackend, +) -> Result<(DevicePtr, usize)> { + let full_len = dims.full_nv * unit; + if dims.tp_size <= 1 { + return Ok((src, full_len)); + } + ensure!( + dims.tp_rank < dims.tp_size, + "tp_rank {} >= tp_size {}", + dims.tp_rank, + dims.tp_size, + ); + let local_len = dims.local_nv * unit; + let local_bytes = local_len * elem_bytes; + let dst = gpu.alloc(local_bytes)?; + let src_off = dims.tp_rank * local_bytes; + gpu.copy_d2d(src.offset(src_off), dst, local_bytes)?; + Ok((dst, local_len)) +} + +/// Shard the `out_proj` BF16 weight `[h, full_value_dim]` row-parallel on its +/// input dim (value_dim). Rank `r` keeps columns +/// `[r*local_value_dim, (r+1)*local_value_dim)` of every output row; the +/// partial products are summed with an all-reduce after the GEMM (mirrors +/// attention `o_proj`). Returns `(ptr, h, local_value_dim)`. +pub fn shard_gdn_out_proj_row_parallel( + src: DevicePtr, + dims: &TpGdnDims, + gpu: &dyn GpuBackend, +) -> Result<(DevicePtr, usize, usize)> { + shard_dense_bf16( + src, + dims.h, + dims.full_value_dim(), + TpShardKind::RowParallel, + dims.tp_rank, + dims.tp_size, + gpu, + ) +} + mod quant_shard; pub use quant_shard::{shard_fp8_block_scaled, shard_quantized_nvfp4}; diff --git a/crates/spark-model/src/tp_shard/tests.rs b/crates/spark-model/src/tp_shard/tests.rs index ecc799b6f..cf748c6ea 100644 --- a/crates/spark-model/src/tp_shard/tests.rs +++ b/crates/spark-model/src/tp_shard/tests.rs @@ -55,3 +55,197 @@ fn divisibility_check() { let tp_size = 4usize; assert_ne!(out_dim % tp_size, 0); } + +// ════════════════════════════════════════════════════════════════════ +// GDN HeadParallel — segmented-slice math (no GPU; validates the copy plan +// the device slicers execute, plus a CPU reference re-concat). +// ════════════════════════════════════════════════════════════════════ + +/// Synthetic GDN dims: full_nk=4, kd=8, full_nv=8, vd=16, tp=2 → local +/// nk=2, nv=4. Q/K width = 32, V/Z width = 128. conv_dim = 2*32+128 = 192. +fn synth_dims(tp_rank: usize) -> TpGdnDims { + TpGdnDims { + tp_rank, + tp_size: 2, + h: 64, + kd: 8, + vd: 16, + local_nk: 2, + full_nk: 4, + local_nv: 4, + full_nv: 8, + } +} + +#[test] +fn gdn_dims_derived() { + let d = synth_dims(0); + assert_eq!(d.full_key_dim(), 32); + assert_eq!(d.local_key_dim(), 16); + assert_eq!(d.full_value_dim(), 128); + assert_eq!(d.local_value_dim(), 64); + assert_eq!(d.full_conv_dim(), 2 * 32 + 128); // 192 + assert_eq!(d.local_conv_dim(), 2 * 16 + 64); // 96 + assert_eq!(d.full_qkvz_out(), 192 + 128); // 320 + assert_eq!(d.local_qkvz_out(), 96 + 64); // 160 + assert_eq!(d.qkv_segments(), [32, 32, 128]); + assert_eq!(d.qkvz_segments(), [32, 32, 128, 128]); +} + +/// The QKVZ segmented plan must slice Q, K, V, Z each by the local head range +/// and pack them back-to-back — NOT take the first `out/tp` contiguous rows. +#[test] +fn qkvz_segment_plan_is_segmented_not_contiguous() { + let d = synth_dims(1); // rank 1 of 2 + let row_bytes = d.h * BF16_BYTES; // 64 * 2 = 128 + let (ops, local_rows) = + segment_copy_plan(&d.qkvz_segments(), row_bytes, d.tp_rank, d.tp_size).unwrap(); + assert_eq!(local_rows, d.local_qkvz_out()); // 160 + + // Full segment starts (rows): Q@0, K@32, V@64, Z@192. + // Rank-1 local halves: Q[16..32], K[48..64], V[128..192], Z[256..320]. + // Packed dst (rows): 0, 16, 32, 96. + let want = [ + CopyOp { src_off: 16 * row_bytes, dst_off: 0, len: 16 * row_bytes }, + CopyOp { src_off: 48 * row_bytes, dst_off: 16 * row_bytes, len: 16 * row_bytes }, + CopyOp { src_off: 128 * row_bytes, dst_off: 32 * row_bytes, len: 64 * row_bytes }, + CopyOp { src_off: 256 * row_bytes, dst_off: 96 * row_bytes, len: 64 * row_bytes }, + ]; + assert_eq!(ops, want); + + // A naive contiguous "first half for rank 0 / second half for rank 1" + // would put rank 1's Q source at row 160, NOT 16 — prove we differ. + assert_ne!(ops[0].src_off, 160 * row_bytes); +} + +/// Rank 0 + rank 1 slices must exactly tile the full buffer with no overlap +/// and no gap, per segment. Reference re-concat on a synthetic u16 buffer. +#[test] +fn qkvz_two_rank_reconcat_tiles_full() { + let d0 = synth_dims(0); + let d1 = synth_dims(1); + let segs = d0.qkvz_segments(); + let full_rows: usize = segs.iter().sum(); // 320 + let h = d0.h; + + // Synthetic full weight: row r filled with value r (u16), h cols each. + let full: Vec = (0..full_rows) + .flat_map(|r| std::iter::repeat(r as u16).take(h)) + .collect(); + + // CPU reference slice using the plan (byte offsets → row indices). + let cpu_slice = |d: &TpGdnDims| -> Vec { + let row_bytes = h * BF16_BYTES; + let (ops, local_rows) = + segment_copy_plan(&segs, row_bytes, d.tp_rank, d.tp_size).unwrap(); + let mut out = vec![0u16; local_rows * h]; + for op in &ops { + let src_row = op.src_off / row_bytes; + let dst_row = op.dst_off / row_bytes; + let nrows = op.len / row_bytes; + for i in 0..nrows { + let s = (src_row + i) * h; + let dd = (dst_row + i) * h; + out[dd..dd + h].copy_from_slice(&full[s..s + h]); + } + } + out + }; + + let r0 = cpu_slice(&d0); + let r1 = cpu_slice(&d1); + + // Expected per-segment source rows for each rank. + // Q: r0=[0..16] r1=[16..32] + // K: r0=[32..48] r1=[48..64] + // V: r0=[64..128] r1=[128..192] + // Z: r0=[192..256] r1=[256..320] + let expect_rows_r0: Vec = (0..16) + .chain(32..48) + .chain(64..128) + .chain(192..256) + .map(|r| r as u16) + .collect(); + let expect_rows_r1: Vec = (16..32) + .chain(48..64) + .chain(128..192) + .chain(256..320) + .map(|r| r as u16) + .collect(); + + let rows_of = |v: &[u16]| -> Vec { v.iter().step_by(h).copied().collect() }; + assert_eq!(rows_of(&r0), expect_rows_r0); + assert_eq!(rows_of(&r1), expect_rows_r1); + + // Union of both ranks (per segment) == the full set of rows: no + // overlap, no gap. + let mut union: Vec = rows_of(&r0); + union.extend(rows_of(&r1)); + union.sort_unstable(); + let all: Vec = (0..full_rows as u16).collect(); + assert_eq!(union, all); +} + +/// conv1d uses the SAME [Q|K|V] 3-segment split but with row width = d_conv. +#[test] +fn conv_segment_plan_uses_qkv_segments() { + let d = synth_dims(0); + let d_conv = 4usize; + let row_bytes = d_conv * BF16_BYTES; // 8 + let (ops, local_rows) = + segment_copy_plan(&d.qkv_segments(), row_bytes, d.tp_rank, d.tp_size).unwrap(); + assert_eq!(local_rows, d.local_conv_dim()); // 96 + // Rank 0: Q[0..16], K[32..48], V[64..128]; packed at 0,16,32. + let want = [ + CopyOp { src_off: 0, dst_off: 0, len: 16 * row_bytes }, + CopyOp { src_off: 32 * row_bytes, dst_off: 16 * row_bytes, len: 16 * row_bytes }, + CopyOp { src_off: 64 * row_bytes, dst_off: 32 * row_bytes, len: 64 * row_bytes }, + ]; + assert_eq!(ops, want); +} + +/// BA is per-group interleaved but the rank boundary lands on a group +/// boundary → a single contiguous slice. rank r → rows [r*2*local_nv, ...). +#[test] +fn ba_single_segment_group_aligned() { + let d = synth_dims(1); + let row_bytes = d.h * BF16_BYTES; + let (ops, local_rows) = + segment_copy_plan(&[2 * d.full_nv], row_bytes, d.tp_rank, d.tp_size).unwrap(); + assert_eq!(local_rows, 2 * d.local_nv); // 8 + // 2*full_nv = 16 rows; rank 1 → rows [8..16). + assert_eq!(ops.len(), 1); + assert_eq!(ops[0].src_off, 8 * row_bytes); + assert_eq!(ops[0].dst_off, 0); + assert_eq!(ops[0].len, 8 * row_bytes); + // Group size = 2*vpg = 2*(nv/nk) = 2*2 = 4 rows; 8 is a multiple → the + // slice starts on a group boundary. + let vpg = d.full_nv / d.full_nk; + assert_eq!((8) % (2 * vpg), 0); +} + +/// Value-vector 1D shard: norm [full_nv*vd] BF16 and a_log/dt_bias [full_nv] +/// FP32, sliced on the value-head axis. +#[test] +fn value_vector_offsets() { + let d = synth_dims(1); + // norm: unit = vd = 16, bf16. full = 8*16 = 128, local = 4*16 = 64. + let full_norm = d.full_nv * d.vd; + let local_norm = d.local_nv * d.vd; + assert_eq!(full_norm, 128); + assert_eq!(local_norm, 64); + let norm_src_off = d.tp_rank * local_norm * BF16_BYTES; // rank1 = 64*2 + assert_eq!(norm_src_off, 128); + // a_log: unit = 1, fp32. full = 8, local = 4. + let f32_bytes = 4usize; + let alog_src_off = d.tp_rank * d.local_nv * f32_bytes; // rank1 = 4*4 + assert_eq!(alog_src_off, 16); +} + +/// A non-divisible segment must be rejected loudly, not silently corrupt. +#[test] +fn segment_plan_rejects_indivisible() { + // 33 rows can't split evenly across tp=2. + let r = segment_copy_plan(&[32, 33], 128, 0, 2); + assert!(r.is_err()); +} diff --git a/crates/spark-model/src/weight_loader/qwen35.rs b/crates/spark-model/src/weight_loader/qwen35.rs index 4ad4d9140..103489416 100644 --- a/crates/spark-model/src/weight_loader/qwen35.rs +++ b/crates/spark-model/src/weight_loader/qwen35.rs @@ -47,13 +47,13 @@ impl ModelWeightLoader for Qwen35WeightLoader { fn supports_tp(&self) -> bool { // FullAttention layers are TP-sharded across all 3 quant paths // (FP8 native, NVFP4-from-disk, BF16 → NVFP4). LinearAttention - // (GDN SSM) layers run full-replica per rank — functionally - // correct (same hidden in → same SSM out across ranks) but - // wastes SSM weight memory. Acceptable trade-off: SSM weights - // are a small fraction of total model size for Qwen3.5-A3B - // (most parameters are in routed MoE experts which are - // EP-sharded). Future work: GDN HeadParallel sharding would - // recover the per-rank SSM memory. + // (GDN SSM) layers are now TP-sharded head-parallel for the BF16 + // and NVFP4 paths (GDN HeadParallel): linear_num_key/value_heads + // are divided per rank in topology.rs, each rank owns a contiguous + // head range, and out_proj is row-parallel with one all-reduce. + // Native block-scaled FP8 SSM still requires TP=1 (per-128-row + // scale slicing deferred) — build_linear_attention_fp8 errors + // clearly when tp_size > 1. true } diff --git a/crates/spark-model/src/weight_loader/qwen35/load_layers/linear_attn_arms.rs b/crates/spark-model/src/weight_loader/qwen35/load_layers/linear_attn_arms.rs index 5326c0bda..836a1c875 100644 --- a/crates/spark-model/src/weight_loader/qwen35/load_layers/linear_attn_arms.rs +++ b/crates/spark-model/src/weight_loader/qwen35/load_layers/linear_attn_arms.rs @@ -11,6 +11,10 @@ use spark_runtime::weights::WeightStore; use crate::layer::TransformerLayer; use crate::layers::{FfnComponent, Qwen3SsmLayer}; +use crate::tp_shard::{ + TpGdnDims, shard_gdn_ba_rows, shard_gdn_conv_rows, shard_gdn_out_proj_row_parallel, + shard_gdn_qkvz_rows, shard_gdn_value_vector, +}; use crate::weight_map::{ DenseWeight, Fp8Weight, Nvfp4Variant, QuantizedWeight, SsmWeights, WeightQuantFormat, dense_auto, dense_f32_safe, dense_keep_f32, gpu_concat_rows, interleave_ba, @@ -130,6 +134,18 @@ pub(super) fn build_linear_attention_fp8( post_attn_norm: DenseWeight, ffn: FfnComponent, ) -> Result> { + // GDN HeadParallel MVP: BF16 + NVFP4 only. Native block-scaled FP8 SSM + // slicing (per-128-row block scales split on the head axis) is deferred — + // slicing rows mid-block would corrupt the scale association. Fail loudly + // rather than ship wrong FP8 scale slicing. + ensure!( + config.tp_world_size.max(1) == 1, + "Native block-scaled FP8 SSM (linear_attn) supports TP=1 only (got tp={}); \ + GDN HeadParallel FP8 scale slicing is deferred. Use the NVFP4 decode path \ + (ATLAS_HOLO_FP4_PROJ_DECODE=1) or run --tp-size 1 for FP8.", + config.tp_world_size, + ); + let p = format!("{lp}.linear_attn"); tracing::info!("Layer {layer_idx}: loading SSM FP8 native (block-scaled decode + prefill)"); @@ -203,22 +219,30 @@ pub(super) fn build_linear_attention_dense_bf16( post_attn_norm: DenseWeight, ffn: FfnComponent, ) -> Result> { - ensure!( - config.tp_world_size.max(1) == 1, - "BF16 dense SSM supports TP=1 only (got tp={})", - config.tp_world_size, + // GDN HeadParallel: `config` already holds per-rank-LOCAL linear head + // counts (topology.rs divided them by tp_size). `TpGdnDims::from_config` + // multiplies back up to the full pre-shard sizes the on-disk weights use; + // the slicers below cut each rank's contiguous head range. For tp=1 every + // slicer returns the source pointer untouched → byte-identical fast path. + let tp_size = config.tp_world_size.max(1); + let dims = TpGdnDims::from_config(config); + tracing::info!( + "Layer {layer_idx}: loading SSM FP8 projections as BF16 dense \ + (tp={tp_size}, local_nk={}, local_nv={})", + dims.local_nk, + dims.local_nv, ); - tracing::info!("Layer {layer_idx}: loading SSM FP8 projections as BF16 dense"); let ssm35 = load_ssm_qwen35(store, lp, gpu, variant)?; - let qkv_rows = config.ssm_qkv_size(); - let z_rows = config.ssm_z_size(); - let qkvz_dense = gpu_concat_rows( + // Concat FULL [Q|K|V] || [Z] (on-disk sizes) then SEGMENT-slice to this + // rank's heads (Q/K/V/Z sliced independently, re-packed local — a naive + // "first half of QKVZ" split is WRONG). + let qkvz_full = gpu_concat_rows( &ssm35.in_proj_qkv, - qkv_rows, + dims.full_conv_dim(), &ssm35.in_proj_z, - z_rows, + dims.full_value_dim(), h, gpu, )?; @@ -229,29 +253,48 @@ pub(super) fn build_linear_attention_dense_bf16( // ~30 GDN layers ≈ 1.5 GB. Free them here — identical numerics. let _ = gpu.free(ssm35.in_proj_qkv.weight); let _ = gpu.free(ssm35.in_proj_z.weight); + let (qkvz_ptr, _, _) = shard_gdn_qkvz_rows(qkvz_full.weight, &dims, gpu)?; + if tp_size > 1 { + let _ = gpu.free(qkvz_full.weight); + } + let qkvz_dense = DenseWeight { weight: qkvz_ptr }; - let nv = config.linear_num_value_heads; - let nk = config.linear_num_key_heads; - let ba_dense = interleave_ba( + // BA: interleave FULL heads (per-group β/α) then slice to local heads — + // the rank boundary always lands on a key-head group boundary. + let ba_full = interleave_ba( &DenseWeight { weight: ssm35.in_proj_a.weight, }, &DenseWeight { weight: ssm35.in_proj_b.weight, }, - nv, - nk, + dims.full_nv, + dims.full_nk, h, gpu, )?; + let (ba_ptr, _, _) = shard_gdn_ba_rows(ba_full.weight, &dims, gpu)?; + if tp_size > 1 { + let _ = gpu.free(ba_full.weight); + } + let ba_dense = DenseWeight { weight: ba_ptr }; + + // conv1d (per-QKV-channel filter), a_log/dt_bias (per value head, FP32), + // norm (per value_dim, BF16), out_proj (row-parallel on value_dim). + let d_conv = config.linear_conv_kernel_dim; + let (conv_ptr, _, _) = shard_gdn_conv_rows(ssm35.conv1d.weight, &dims, d_conv, gpu)?; + let (a_log_ptr, _) = shard_gdn_value_vector(ssm35.a_log.weight, &dims, 1, 4, gpu)?; + let (dt_bias_ptr, _) = shard_gdn_value_vector(ssm35.dt_bias.weight, &dims, 1, 4, gpu)?; + let (norm_ptr, _) = shard_gdn_value_vector(ssm35.norm.weight, &dims, dims.vd, 2, gpu)?; + let (out_proj_ptr, _, _) = shard_gdn_out_proj_row_parallel(ssm35.out_proj.weight, &dims, gpu)?; let ssm = SsmWeights { in_proj_qkvz: qkvz_dense, in_proj_ba: ba_dense, - conv1d: ssm35.conv1d, - a_log: ssm35.a_log, - dt_bias: ssm35.dt_bias, - norm: ssm35.norm, + conv1d: DenseWeight { weight: conv_ptr }, + a_log: DenseWeight { weight: a_log_ptr }, + dt_bias: DenseWeight { weight: dt_bias_ptr }, + norm: DenseWeight { weight: norm_ptr }, out_proj: QuantizedWeight::null(), }; @@ -266,14 +309,19 @@ pub(super) fn build_linear_attention_dense_bf16( config, gpu, )?; - layer.out_proj_dense = Some(ssm35.out_proj); + layer.out_proj_dense = Some(DenseWeight { + weight: out_proj_ptr, + }); // Decode-only FP8 SSM overlay (ATLAS_HOLO_FP8_SSM_DECODE=1): install the // on-disk block-scaled FP8 QKVZ/out_proj so DECODE runs through // w8a16_gemv / w8a16_gemm (half the BF16 weight bandwidth — SSM weights // are the bulk of the per-step fixed decode cost), while PREFILL keeps the // stable BF16 dense path (sidesteps the native-FP8 FLA-prefill crash at // layer 36). Costs ~25 MB/GDN layer extra (BF16 kept for prefill). - if std::env::var("ATLAS_HOLO_FP8_SSM_DECODE").ok().as_deref() == Some("1") { + // The FP8 decode overlay loads FULL (unsliced) block-scaled FP8 weights; + // its per-128-row scale slicing is deferred (same reason as the native-FP8 + // path). Skip under TP>1 so the sharded BF16 path stays correct. + if tp_size == 1 && std::env::var("ATLAS_HOLO_FP8_SSM_DECODE").ok().as_deref() == Some("1") { let p = format!("{lp}.linear_attn"); let (qkvz_fp8, out_fp8) = load_ssm_fp8_decode_weights(layer_idx, store, &p, gpu, h)?; layer.set_fp8_decode_weights(Some(qkvz_fp8), Some(out_fp8)); @@ -297,34 +345,70 @@ pub(super) fn build_linear_attention_nvfp4( post_attn_norm: DenseWeight, ffn: FfnComponent, ) -> Result> { + // GDN HeadParallel: `config` holds per-rank-LOCAL linear head counts. + // Slice each SSM projection to this rank's head range on the dense/BF16 + // intermediate BEFORE quantizing to NVFP4 (dequant→slice→requant is the + // safe path — no NVFP4 packed-buffer surgery). For tp=1 the slicers return + // the source pointer untouched → byte-identical fast path. + let tp_size = config.tp_world_size.max(1); + let dims = TpGdnDims::from_config(config); + let ssm35 = load_ssm_qwen35(store, lp, gpu, variant)?; - let qkv_rows = config.ssm_qkv_size(); - let z_rows = config.ssm_z_size(); - let qkvz_dense = gpu_concat_rows( + // Concat FULL [Q|K|V] || [Z] then SEGMENT-slice to local heads. + let qkvz_full = gpu_concat_rows( &ssm35.in_proj_qkv, - qkv_rows, + dims.full_conv_dim(), &ssm35.in_proj_z, - z_rows, + dims.full_value_dim(), h, gpu, )?; + let (qkvz_ptr, _, _) = shard_gdn_qkvz_rows(qkvz_full.weight, &dims, gpu)?; + if tp_size > 1 { + let _ = gpu.free(qkvz_full.weight); + } + let qkvz_dense = DenseWeight { weight: qkvz_ptr }; - let nv = config.linear_num_value_heads; - let nk = config.linear_num_key_heads; - let ba_dense = interleave_ba( + // BA: interleave FULL heads then slice to local (group-aligned). + let ba_full = interleave_ba( &DenseWeight { weight: ssm35.in_proj_a.weight, }, &DenseWeight { weight: ssm35.in_proj_b.weight, }, - nv, - nk, + dims.full_nv, + dims.full_nk, h, gpu, )?; + let (ba_ptr, _, _) = shard_gdn_ba_rows(ba_full.weight, &dims, gpu)?; + if tp_size > 1 { + let _ = gpu.free(ba_full.weight); + } + let ba_dense = DenseWeight { weight: ba_ptr }; + + // conv1d / a_log / dt_bias / norm sliced to local heads (stay dense/FP32). + let d_conv = config.linear_conv_kernel_dim; + let (conv_ptr, _, _) = shard_gdn_conv_rows(ssm35.conv1d.weight, &dims, d_conv, gpu)?; + let (a_log_ptr, _) = shard_gdn_value_vector(ssm35.a_log.weight, &dims, 1, 4, gpu)?; + let (dt_bias_ptr, _) = shard_gdn_value_vector(ssm35.dt_bias.weight, &dims, 1, 4, gpu)?; + let (norm_ptr, _) = shard_gdn_value_vector(ssm35.norm.weight, &dims, dims.vd, 2, gpu)?; + let conv1d_local = DenseWeight { weight: conv_ptr }; + let a_log_local = DenseWeight { weight: a_log_ptr }; + let dt_bias_local = DenseWeight { weight: dt_bias_ptr }; + let norm_local = DenseWeight { weight: norm_ptr }; + // out_proj is row-parallel: slice its input (value_dim) to local, then + // quantize the LOCAL [h, local_value_dim] weight. + let (out_proj_ptr, _, _) = shard_gdn_out_proj_row_parallel(ssm35.out_proj.weight, &dims, gpu)?; + let out_proj_local = DenseWeight { + weight: out_proj_ptr, + }; + + // All sizes below are LOCAL (config was TP-divided at load). + let nv = config.linear_num_value_heads; let qkvz_size = config.ssm_qkvz_size(); let qkvz_nvfp4 = quantize_to_nvfp4(&qkvz_dense, qkvz_size, h, gpu, absmax_k, quantize_k, stream)?; @@ -333,7 +417,7 @@ pub(super) fn build_linear_attention_nvfp4( let value_dim = nv * config.linear_value_head_dim; let out_proj_nvfp4 = quantize_to_nvfp4( - &ssm35.out_proj, + &out_proj_local, h, value_dim, gpu, @@ -379,7 +463,7 @@ pub(super) fn build_linear_attention_nvfp4( crate::layers::ops::bf16_to_fp8( gpu, b2f_k, - ssm35.out_proj.weight, + out_proj_local.weight, out_fp8, out_total, stream, @@ -393,10 +477,10 @@ pub(super) fn build_linear_attention_nvfp4( let ssm = SsmWeights { in_proj_qkvz: qkvz_dense, in_proj_ba: ba_dense, - conv1d: ssm35.conv1d, - a_log: ssm35.a_log, - dt_bias: ssm35.dt_bias, - norm: ssm35.norm, + conv1d: conv1d_local, + a_log: a_log_local, + dt_bias: dt_bias_local, + norm: norm_local, out_proj: out_proj_nvfp4, }; @@ -439,10 +523,11 @@ pub(super) fn build_linear_attention_nvfp4( std::env::var("ATLAS_GDN_BF16_WEIGHTS").ok().as_deref(), Some("1") ) { - // ssm35.out_proj weight is BF16 on GPU (from load_ssm_qwen35 → - // dense_auto on Fp8Dequanted variant). It's a separate buffer - // from out_proj_nvfp4 / out_proj_fp8_prefill. Set as dense path. - layer.out_proj_dense = Some(ssm35.out_proj); + // out_proj_local weight is BF16 on GPU (from load_ssm_qwen35 → + // dense_auto on Fp8Dequanted variant, sliced to this rank's value + // heads). It's a separate buffer from out_proj_nvfp4 / + // out_proj_fp8_prefill. Set as dense path. + layer.out_proj_dense = Some(out_proj_local); tracing::info!( "SSM[{lp}] ATLAS_GDN_BF16_WEIGHTS: out_proj routed through BF16 dense_gemm (overrides FP8/NVFP4)" ); diff --git a/crates/spark-server/src/main_modules/serve_phases/topology.rs b/crates/spark-server/src/main_modules/serve_phases/topology.rs index 4386f0e08..ecfe976e2 100644 --- a/crates/spark-server/src/main_modules/serve_phases/topology.rs +++ b/crates/spark-server/src/main_modules/serve_phases/topology.rs @@ -96,10 +96,36 @@ pub(crate) fn resolve_topology( } config.num_attention_heads /= tp_size; config.num_key_value_heads /= tp_size; + // GDN HeadParallel: linear-attention (SSM) key/value head counts are + // sharded exactly like attention heads — each rank owns a contiguous + // head range; the recurrence is head-parallel with one all-reduce after + // out_proj. Only relevant for SSM-hybrid models (linear_*_heads > 0); + // pure-attention configs leave these at 0 and skip the divide. + if config.linear_num_key_heads > 0 || config.linear_num_value_heads > 0 { + if !config.linear_num_key_heads.is_multiple_of(tp_size) { + anyhow::bail!( + "TP requires linear_num_key_heads ({}) divisible by tp_size ({})", + config.linear_num_key_heads, + tp_size, + ); + } + if !config.linear_num_value_heads.is_multiple_of(tp_size) { + anyhow::bail!( + "TP requires linear_num_value_heads ({}) divisible by tp_size ({})", + config.linear_num_value_heads, + tp_size, + ); + } + config.linear_num_key_heads /= tp_size; + config.linear_num_value_heads /= tp_size; + } tracing::info!( - "TP-local head counts: num_attention_heads={}, num_key_value_heads={}", + "TP-local head counts: num_attention_heads={}, num_key_value_heads={}, \ + linear_num_key_heads={}, linear_num_value_heads={}", config.num_attention_heads, config.num_key_value_heads, + config.linear_num_key_heads, + config.linear_num_value_heads, ); } if world_size > 1 { diff --git a/scripts/gdn_tp_golden.py b/scripts/gdn_tp_golden.py new file mode 100755 index 000000000..5db0622ca --- /dev/null +++ b/scripts/gdn_tp_golden.py @@ -0,0 +1,294 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: AGPL-3.0-only +""" +GDN HeadParallel golden gate — TP=2 (loopback) vs TP=1 long-sequence parity. + +WHY THIS EXISTS +--------------- +Gated-DeltaNet (linear_attention / SSM) layers are STATEFUL: the recurrence +carries `h_state` + `conv_state` across every token. A layout / head-slicing / +all-reduce bug in the GDN HeadParallel tensor-parallel path does NOT show up on +a short prompt — the state has to accumulate over hundreds of tokens before a +mis-sliced value head or a missing out_proj all-reduce visibly diverges. So the +gate here is a LONG prompt (>~1k tokens) plus a long greedy generation, compared +between a TP=1 reference and a TP=2 (two ranks, one GPU, NCCL loopback) run. + +Correct GDN HeadParallel ⇒ the two runs produce the IDENTICAL greedy token +sequence, and near-identical top-token logprobs (the out_proj all-reduce sums +partial products in a different float order per rank, so expect ~1e-2, not bit +exact). + +WHAT IT CHECKS +-------------- + 1. Greedy (temperature=0) generated token sequence is identical (HARD gate). + 2. Per-position top-1 logprob max-abs-diff < --logprob-tol (soft numeric gate). + +TWO WAYS TO RUN +--------------- +(A) Point it at two servers you launched yourself (recommended on a shared box): + + python3 scripts/gdn_tp_golden.py \ + --url-tp1 http://127.0.0.1:8800 \ + --url-tp2 http://127.0.0.1:8801 + +(B) Let it launch both servers on ONE GB10 (TP=2 = two ranks pinned to the same + GPU → NCCL loopback). Build the server first: + + export LIBRARY_PATH=/home/ms/nccl/build/lib CUDARC_CUDA_VERSION=12000 \ + ATLAS_TARGET_HW=gb10 ATLAS_TARGET_MODEL=qwen3.6-27b \ + ATLAS_TARGET_QUANT=nvfp4 + cargo build --release -p spark-server --bin spark + + python3 scripts/gdn_tp_golden.py --launch \ + --model \ + --bin target/release/spark + + NOTE (single-GB10 memory): TP shards attention + GDN weights, but with + ep_size=1 the MoE experts are REPLICATED per rank, so two ranks on one GPU + roughly double MoE-resident memory. Use a model that fits 2× on the GB10, + or lower --gpu-mem-util, or run the two ranks on two GPUs. FP8-native SSM + checkpoints force TP=1 by design (block-scale slicing deferred) — use a BF16 + or NVFP4 SSM checkpoint for the TP=2 side. + +EXIT CODE: 0 = parity PASS, 1 = FAIL (token mismatch or logprob drift). +""" + +import argparse +import json +import os +import subprocess +import sys +import time +import urllib.request + + +# A deterministic, self-contained long prompt. Repeated to force a multi-chunk +# prefill (>1k tokens) so the GDN state genuinely accumulates before we diff. +_PARAGRAPH = ( + "The recurrence in a gated delta network threads a hidden state and a " + "convolution state through every position, so any tensor-parallel head " + "slicing bug compounds silently over a long context rather than failing " + "loudly on the first token. Consider the counting sequence carefully: " + "one, two, three, four, five, six, seven, eight, nine, ten. " +) + + +def build_prompt(repeat: int) -> str: + body = _PARAGRAPH * repeat + return ( + body + + "\n\nContinue the numeric analysis precisely and deterministically, " + "step by step, without repeating yourself:\n" + ) + + +def http_json(url: str, payload=None, timeout=600): + data = json.dumps(payload).encode() if payload is not None else None + req = urllib.request.Request( + url, + data=data, + headers={"Content-Type": "application/json"}, + method="POST" if data is not None else "GET", + ) + with urllib.request.urlopen(req, timeout=timeout) as resp: + return json.loads(resp.read().decode()) + + +def wait_healthy(base_url: str, timeout_s: int) -> str: + """Poll /v1/models until the server answers; return the model id.""" + deadline = time.time() + timeout_s + last_err = None + while time.time() < deadline: + try: + r = http_json(f"{base_url}/v1/models", timeout=10) + return r["data"][0]["id"] + except Exception as e: # noqa: BLE001 + last_err = e + time.sleep(2.0) + raise RuntimeError(f"server at {base_url} not healthy after {timeout_s}s: {last_err}") + + +def sample(base_url: str, model: str, prompt: str, max_tokens: int): + """Greedy completion with per-token top-1 logprobs.""" + payload = { + "model": model, + "prompt": prompt, + "max_tokens": max_tokens, + "temperature": 0.0, + "top_p": 1.0, + "seed": 0, + "logprobs": 1, # OpenAI completions: top-1 logprob per generated token + "stream": False, + } + r = http_json(f"{base_url}/v1/completions", payload) + choice = r["choices"][0] + text = choice.get("text", "") + lp = choice.get("logprobs") or {} + tokens = lp.get("tokens") or [] + token_logprobs = lp.get("token_logprobs") or [] + return text, tokens, token_logprobs + + +def compare(a, b, logprob_tol: float) -> bool: + text_a, tok_a, lpa = a + text_b, tok_b, lpb = b + + print("\n=== GDN HeadParallel TP golden gate ===") + print(f"TP=1 tokens: {len(tok_a)} TP=2 tokens: {len(tok_b)}") + + ok = True + + # HARD gate 1: identical greedy token sequence. + if tok_a and tok_b: + n = min(len(tok_a), len(tok_b)) + first_div = next((i for i in range(n) if tok_a[i] != tok_b[i]), None) + if len(tok_a) != len(tok_b) or first_div is not None: + ok = False + where = first_div if first_div is not None else n + print(f" FAIL: greedy token sequences diverge at position {where}") + print(f" TP=1[{where}]={tok_a[where] if where < len(tok_a) else ''!r}") + print(f" TP=2[{where}]={tok_b[where] if where < len(tok_b) else ''!r}") + else: + print(f" PASS: {len(tok_a)} greedy tokens identical") + else: + # Fallback when the server does not echo per-token logprobs: compare text. + if text_a != text_b: + ok = False + print(" FAIL: generated text differs (no per-token logprobs available)") + print(f" TP=1: {text_a[:160]!r}") + print(f" TP=2: {text_b[:160]!r}") + else: + print(f" PASS: generated text identical ({len(text_a)} chars)") + + # SOFT gate 2: numeric closeness of top-1 logprobs. + if lpa and lpb: + n = min(len(lpa), len(lpb)) + diffs = [ + abs(x - y) + for x, y in zip(lpa[:n], lpb[:n]) + if x is not None and y is not None + ] + if diffs: + mx, mean = max(diffs), sum(diffs) / len(diffs) + status = "PASS" if mx <= logprob_tol else "FAIL" + if mx > logprob_tol: + ok = False + print( + f" {status}: top-1 logprob drift max={mx:.4e} mean={mean:.4e} " + f"(tol={logprob_tol:.2e})" + ) + + print("=== " + ("GOLDEN GATE PASS" if ok else "GOLDEN GATE FAIL") + " ===\n") + return ok + + +def spawn_server(bin_path, model, *, rank, world_size, tp_size, port, + master_port, gpu_mem_util, log_path): + cmd = [ + bin_path, "serve", model, + "--port", str(port), + "--rank", str(rank), + "--world-size", str(world_size), + "--tp-size", str(tp_size), + "--ep-size", "1", + "--master-addr", "127.0.0.1", + "--master-port", str(master_port), + "--gpu-memory-utilization", str(gpu_mem_util), + ] + env = dict(os.environ) + # Both ranks share the single GB10 → NCCL loopback. + env.setdefault("CUDA_VISIBLE_DEVICES", "0") + log = open(log_path, "w") + print(f" launch rank{rank}/{world_size} tp={tp_size} port={port} -> {log_path}") + # setsid so the whole server (+ its threads) is killable as one group. + return subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT, env=env, + start_new_session=True), log + + +def launch_mode(args): + procs = [] + try: + # --- TP=1 reference: single process, world_size=1 --- + p1, _ = spawn_server( + args.bin, args.model, rank=0, world_size=1, tp_size=1, + port=args.port_tp1, master_port=args.master_port, + gpu_mem_util=args.gpu_mem_util, log_path="/tmp/gdn_tp_golden_tp1.log") + procs.append(p1) + model_id = wait_healthy(f"http://127.0.0.1:{args.port_tp1}", args.startup_timeout) + ref = sample(f"http://127.0.0.1:{args.port_tp1}", model_id, + build_prompt(args.prompt_repeat), args.max_tokens) + # Free the reference server before bringing up two ranks (memory). + p1.terminate() + procs.remove(p1) + try: + p1.wait(timeout=30) + except subprocess.TimeoutExpired: + p1.kill() + + # --- TP=2 loopback: rank0 (HTTP) + rank1 (worker), same GPU --- + pr0, _ = spawn_server( + args.bin, args.model, rank=0, world_size=2, tp_size=2, + port=args.port_tp2, master_port=args.master_port + 1, + gpu_mem_util=args.gpu_mem_util, log_path="/tmp/gdn_tp_golden_tp2_rank0.log") + procs.append(pr0) + pr1, _ = spawn_server( + args.bin, args.model, rank=1, world_size=2, tp_size=2, + port=args.port_tp2 + 1, master_port=args.master_port + 1, + gpu_mem_util=args.gpu_mem_util, log_path="/tmp/gdn_tp_golden_tp2_rank1.log") + procs.append(pr1) + model_id2 = wait_healthy(f"http://127.0.0.1:{args.port_tp2}", args.startup_timeout) + got = sample(f"http://127.0.0.1:{args.port_tp2}", model_id2, + build_prompt(args.prompt_repeat), args.max_tokens) + + return compare(ref, got, args.logprob_tol) + finally: + for p in procs: + p.terminate() + for p in procs: + try: + p.wait(timeout=30) + except subprocess.TimeoutExpired: + p.kill() + + +def url_mode(args): + m1 = wait_healthy(args.url_tp1, args.startup_timeout) + m2 = wait_healthy(args.url_tp2, args.startup_timeout) + prompt = build_prompt(args.prompt_repeat) + ref = sample(args.url_tp1, m1, prompt, args.max_tokens) + got = sample(args.url_tp2, m2, prompt, args.max_tokens) + return compare(ref, got, args.logprob_tol) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__, + formatter_class=argparse.RawDescriptionHelpFormatter) + ap.add_argument("--launch", action="store_true", + help="launch both servers on one GB10 (TP=2 = NCCL loopback)") + ap.add_argument("--model", help="HF id or local path (required with --launch)") + ap.add_argument("--bin", default="target/release/spark", help="spark server binary") + ap.add_argument("--url-tp1", default="http://127.0.0.1:8800") + ap.add_argument("--url-tp2", default="http://127.0.0.1:8801") + ap.add_argument("--port-tp1", type=int, default=8800) + ap.add_argument("--port-tp2", type=int, default=8801) + ap.add_argument("--master-port", type=int, default=29500) + ap.add_argument("--gpu-mem-util", type=float, default=0.40) + ap.add_argument("--prompt-repeat", type=int, default=90, + help="paragraph repeats (~1.5k+ prompt tokens by default)") + ap.add_argument("--max-tokens", type=int, default=512, + help="greedy generation length (>=250 per house rule)") + ap.add_argument("--logprob-tol", type=float, default=5e-2) + ap.add_argument("--startup-timeout", type=int, default=600) + args = ap.parse_args() + + if args.launch: + if not args.model: + ap.error("--launch requires --model") + ok = launch_mode(args) + else: + ok = url_mode(args) + sys.exit(0 if ok else 1) + + +if __name__ == "__main__": + main() From ae8e6b2d95ce9bce4cb30bf649f180e23fd113a7 Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Fri, 3 Jul 2026 01:17:23 -0400 Subject: [PATCH 02/15] perf(gdn-tp): register attn norm_output buffer + flag-gated CUDA-graph decode (ATLAS_GDN_DECODE_GRAPH=1) Collectives were already async (stream-ordered via all_reduce_async). Two speedups: - register buffers.norm_output() with NCCL so the 16 attention o_proj all-reduces/token use cached IB MR + user-buffer zero-copy (was unregistered). - ATLAS_GDN_DECODE_GRAPH=1 extends the existing decode CUDA-graph gate (was self.comm.is_none() only) to capture the full tp>1 decode forward (~130 kernels + 64 all-reduces) into one replayable graph; eager fallback on capture failure. Default off; tp=1 byte-identical. --- crates/spark-model/src/model/impl_a1.rs | 18 +- .../src/model/trait_impl/decode_a.rs | 233 ++++++++++++------ 2 files changed, 178 insertions(+), 73 deletions(-) diff --git a/crates/spark-model/src/model/impl_a1.rs b/crates/spark-model/src/model/impl_a1.rs index ac48a4768..299a475c1 100644 --- a/crates/spark-model/src/model/impl_a1.rs +++ b/crates/spark-model/src/model/impl_a1.rs @@ -257,7 +257,17 @@ impl TransformerModel { // Marconi restore (prefill stream). See `snapshot_event` doc in types.rs. let snapshot_event = gpu.create_event()?; - // EP: register moe_output buffer with NCCL and provide bf16_add kernel. + // EP/TP: register the all-reduce target buffers with NCCL (caches the + // IB/RoCE memory registration, enabling zero-copy user-buffer + // collectives) and provide the bf16_add kernel for the 2-rank + // send/recv fast path. + // - moe_output: EP MoE reduce + ALL GDN HeadParallel SSM out_proj + // reduces (decode `ssm_forward`, batched decode, multi-seq + // batched, prefill, prefill phase-3 all write out_proj into + // `buffers.moe_output()`). + // - norm_output: attention o_proj decode output + // (`attention_forward_oproj` writes o_out = `buffers.norm_output()`), + // reduced per attention layer under TP. if let Some(ref comm) = comm && comm.world_size() == 2 { @@ -267,6 +277,12 @@ impl TransformerModel { Ok(_) => tracing::info!("Registered moe_output ({moe_bytes} B) with NCCL"), Err(e) => tracing::warn!("ncclCommRegister moe_output failed (non-fatal): {e}"), } + let norm_ptr = buffers.norm_output().0; + let norm_bytes = buffers.sizes().norm_output; + match comm.register_buffer(norm_ptr, norm_bytes) { + Ok(_) => tracing::info!("Registered norm_output ({norm_bytes} B) with NCCL"), + Err(e) => tracing::warn!("ncclCommRegister norm_output failed (non-fatal): {e}"), + } match gpu.kernel("bf16_add", "bf16_add_inplace") { Ok(k) => comm.set_add_kernel(k.0), Err(e) => { diff --git a/crates/spark-model/src/model/trait_impl/decode_a.rs b/crates/spark-model/src/model/trait_impl/decode_a.rs index eba822299..806a2d7fe 100644 --- a/crates/spark-model/src/model/trait_impl/decode_a.rs +++ b/crates/spark-model/src/model/trait_impl/decode_a.rs @@ -170,7 +170,24 @@ impl TransformerModel { // and remove per-kernel launch overhead. Env-gated so it can be toggled // off at deploy time (instant revert) if capture crashes / replay hangs. let ep_graphs = std::env::var("ATLAS_EP_GRAPHS").is_ok_and(|v| v == "1" || v == "true"); - let use_graphs = (self.comm.is_none() || ep_graphs) + // GDN HeadParallel TP decode graphs (ATLAS_GDN_DECODE_GRAPH=1, default + // OFF): capture the whole single-token decode forward — ~130 kernels + // plus the per-layer TP all-reduces (48 GDN SSM out_proj + 16 + // attention o_proj on Qwen3.6) — into one replayable graph. The + // collectives go through `all_reduce_async` (event fork/join onto the + // dedicated NCCL comm stream), which stream capture pulls into the + // graph as cross-stream nodes; capture runs in RELAXED mode (see + // `begin_capture`) as NCCL requires, and the events are + // CU_EVENT_DISABLE_TIMING (capture-legal). All per-token inputs + // (token embedding, positions, slot, seq_len, block_table) are + // uploaded to STABLE device buffers in Phase 1 before replay, and the + // per-slot SSM conv/h states are updated in place at stable pointers, + // so replay is shape/pointer-static. This removes the per-token host + // launch cost that dominates 2-node GDN HeadParallel decode. Capture + // failure falls back to eager execution (graphs then stay disabled). + let gdn_graphs = + std::env::var("ATLAS_GDN_DECODE_GRAPH").is_ok_and(|v| v == "1" || v == "true"); + let use_graphs = (self.comm.is_none() || ep_graphs || gdn_graphs) && !self.profile && !self .suppress_graphs @@ -223,28 +240,167 @@ impl TransformerModel { // ── Phase 3: Capture new CUDA graph (or run eagerly for EP) ── + // Track whether a capture is actually recording: a begin_capture + // failure falls back to eager execution (and disables graphs for the + // rest of the run) instead of failing the decode step. + let mut capture_active = false; if use_graphs { tracing::info!( "CUDA graph capture: starting for {} layers", self.layers.len() ); - self.gpu.begin_capture(stream)?; + match self.gpu.begin_capture(stream) { + Ok(()) => capture_active = true, + Err(e) => { + tracing::warn!( + "CUDA graph begin_capture failed ({e:#}) — \ + running eagerly and disabling graph capture" + ); + self.suppress_graphs + .store(true, std::sync::atomic::Ordering::Relaxed); + } + } } let probe_layers = !use_graphs && seq.seq_len == seq.prompt_len && std::env::var("ATLAS_SSM_SAVE_DUMP").is_ok(); + self.decode_forward_body( + hidden, + residual, + seq, + &mut kv_cache, + &ctx, + probe_layers, + use_graphs, + stream, + )?; + + // Decode-step diagnostic for Gemma-4 degeneration analysis. Only fires + // when ATLAS_DIAG_GEMMA4=1 (which also disables CUDA graphs upstream, + // so the d2h sync below is safe). Reads top-5 tokens by logit so we + // can see whether the LM head produced a near-tie or a confident bad + // pick. (B4 — Creative haiku degeneration loop diagnostic.) + if std::env::var("ATLAS_DIAG_GEMMA4").is_ok_and(|v| v == "1" || v == "true") { + self.gpu.synchronize(stream)?; + let n_logits = self.config.vocab_size; + // Read the buffer the LM head actually wrote to. With Gemma-4 + // dense the single-token decode lm_head produces FP32 in + // `logits_fp32_buf`; the BF16 buffer would be all zeros there. + let logit_vals: Vec = if self.use_fp32_logits { + let mut buf = vec![0u8; n_logits * 4]; + if let Err(e) = self.gpu.copy_d2h(self.logits_fp32_buf, &mut buf) { + tracing::error!("ATLAS_DIAG_GEMMA4: copy_d2h(logits_fp32_buf): {e:#}"); + } + buf.chunks_exact(4) + .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) + .collect() + } else { + let mut buf = vec![0u8; n_logits * 2]; + if let Err(e) = self.gpu.copy_d2h(self.buffers.logits(), &mut buf) { + tracing::error!("ATLAS_DIAG_GEMMA4: copy_d2h(logits BF16): {e:#}"); + } + buf.chunks_exact(2) + .map(|c| { + let bits = u16::from_le_bytes([c[0], c[1]]); + f32::from_bits((bits as u32) << 16) + }) + .collect() + }; + let max = logit_vals.iter().cloned().fold(f32::NEG_INFINITY, f32::max); + let min = logit_vals.iter().cloned().fold(f32::INFINITY, f32::min); + let mut idx: Vec = (0..logit_vals.len()).collect(); + idx.sort_by(|&a, &b| { + logit_vals[b] + .partial_cmp(&logit_vals[a]) + .unwrap_or(std::cmp::Ordering::Equal) + }); + let top5: Vec<(usize, f32)> = idx.iter().take(5).map(|&i| (i, logit_vals[i])).collect(); + tracing::warn!( + "DIAG decode logits: max={max:.4} min={min:.4} prev_token={token} top5={top5:?}", + ); + } + + if capture_active { + match self.gpu.end_capture(stream) { + Ok(graph) if graph.0 != 0 => { + tracing::info!( + "CUDA graph captured successfully for slot={} (handle={:?})", + seq.slot_idx, + graph.0 + ); + if let Some(ref mut cache) = graph_cache { + cache.insert(seq.slot_idx, graph); + } + self.gpu.launch_graph(graph, stream)?; + } + Ok(_) => { + tracing::warn!("CUDA graph capture returned null handle — running eagerly"); + // If graph.0 == 0 (mock): operations already executed during capture + } + Err(e) => { + // Capture RECORDS without executing, so nothing has run + // for this token yet — re-run the whole forward body + // eagerly (end_capture failure terminates the capture, so + // the stream is back in normal mode) and disable graphs + // for the rest of the run. + tracing::warn!( + "CUDA graph end_capture failed ({e:#}) — \ + re-running decode step eagerly and disabling graph capture" + ); + self.suppress_graphs + .store(true, std::sync::atomic::Ordering::Relaxed); + self.decode_forward_body( + hidden, + residual, + seq, + &mut kv_cache, + &ctx, + false, + false, + stream, + )?; + } + } + } + + seq.tokens.push(token); + seq.seq_len += 1; + + Ok(self.decode_logits_ptr()) + } + + /// Single-token decode forward body: per-layer decode + periodic SSM + /// state normalization + final RMS norm + LM head. + /// + /// Runs once per decode step — either eagerly or inside a CUDA graph + /// capture region — and a SECOND time as the eager fallback when + /// `end_capture` fails (capture records without executing, so a re-run + /// performs the step exactly once). Everything here is stream-ordered on + /// `stream`; the only host syncs are gated on `!use_graphs` / + /// `probe_layers` (both false under capture). + fn decode_forward_body( + &self, + hidden: DevicePtr, + residual: DevicePtr, + seq: &mut SequenceState, + kv_cache: &mut PagedKvCache, + ctx: &ForwardContext, + probe_layers: bool, + use_graphs: bool, + stream: u64, + ) -> Result<()> { for (i, layer) in self.layers.iter().enumerate() { layer.decode( hidden, residual, seq.layer_states[i].as_mut(), - &mut kv_cache, + kv_cache, seq.seq_len, &mut seq.block_table, &mut seq.disk_block_ids, &mut seq.disk_last_offloaded_per_layer, - &ctx, + ctx, stream, )?; // CBD per-layer hidden fingerprint at decode step 0 (eager only). @@ -309,73 +465,6 @@ impl TransformerModel { // LM head reads from normed directly (no D2D copy needed) self.lm_head(normed, stream)?; - - // Decode-step diagnostic for Gemma-4 degeneration analysis. Only fires - // when ATLAS_DIAG_GEMMA4=1 (which also disables CUDA graphs upstream, - // so the d2h sync below is safe). Reads top-5 tokens by logit so we - // can see whether the LM head produced a near-tie or a confident bad - // pick. (B4 — Creative haiku degeneration loop diagnostic.) - if std::env::var("ATLAS_DIAG_GEMMA4").is_ok_and(|v| v == "1" || v == "true") { - self.gpu.synchronize(stream)?; - let n_logits = self.config.vocab_size; - // Read the buffer the LM head actually wrote to. With Gemma-4 - // dense the single-token decode lm_head produces FP32 in - // `logits_fp32_buf`; the BF16 buffer would be all zeros there. - let logit_vals: Vec = if self.use_fp32_logits { - let mut buf = vec![0u8; n_logits * 4]; - if let Err(e) = self.gpu.copy_d2h(self.logits_fp32_buf, &mut buf) { - tracing::error!("ATLAS_DIAG_GEMMA4: copy_d2h(logits_fp32_buf): {e:#}"); - } - buf.chunks_exact(4) - .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) - .collect() - } else { - let mut buf = vec![0u8; n_logits * 2]; - if let Err(e) = self.gpu.copy_d2h(self.buffers.logits(), &mut buf) { - tracing::error!("ATLAS_DIAG_GEMMA4: copy_d2h(logits BF16): {e:#}"); - } - buf.chunks_exact(2) - .map(|c| { - let bits = u16::from_le_bytes([c[0], c[1]]); - f32::from_bits((bits as u32) << 16) - }) - .collect() - }; - let max = logit_vals.iter().cloned().fold(f32::NEG_INFINITY, f32::max); - let min = logit_vals.iter().cloned().fold(f32::INFINITY, f32::min); - let mut idx: Vec = (0..logit_vals.len()).collect(); - idx.sort_by(|&a, &b| { - logit_vals[b] - .partial_cmp(&logit_vals[a]) - .unwrap_or(std::cmp::Ordering::Equal) - }); - let top5: Vec<(usize, f32)> = idx.iter().take(5).map(|&i| (i, logit_vals[i])).collect(); - tracing::warn!( - "DIAG decode logits: max={max:.4} min={min:.4} prev_token={token} top5={top5:?}", - ); - } - - if use_graphs { - let graph = self.gpu.end_capture(stream)?; - if graph.0 != 0 { - tracing::info!( - "CUDA graph captured successfully for slot={} (handle={:?})", - seq.slot_idx, - graph.0 - ); - if let Some(ref mut cache) = graph_cache { - cache.insert(seq.slot_idx, graph); - } - self.gpu.launch_graph(graph, stream)?; - } else { - tracing::warn!("CUDA graph capture returned null handle — running eagerly"); - } - // If graph.0 == 0 (mock): operations already executed during capture - } - - seq.tokens.push(token); - seq.seq_len += 1; - - Ok(self.decode_logits_ptr()) + Ok(()) } } From edca7140a6cf4ee18c538f4612b81a47318941b7 Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Fri, 3 Jul 2026 02:06:18 -0400 Subject: [PATCH 03/15] fix(gdn-tp): make GDN decode body CUDA-graph-capturable --- crates/spark-model/src/model/impl_a2.rs | 52 +++++++++++++++---- .../src/model/trait_impl/ep_misc.rs | 13 ++++- .../src/model/trait_impl/prefill_b.rs | 8 +-- .../src/model/trait_impl/prefill_b/batch.rs | 4 +- .../trait_impl/prefill_b/batch_kernel.rs | 4 +- .../trait_impl/prefill_b/prefix_lookup.rs | 4 +- .../src/model/trait_impl/prefill_c.rs | 7 +-- 7 files changed, 68 insertions(+), 24 deletions(-) diff --git a/crates/spark-model/src/model/impl_a2.rs b/crates/spark-model/src/model/impl_a2.rs index 717c69e32..5a7d14a0a 100644 --- a/crates/spark-model/src/model/impl_a2.rs +++ b/crates/spark-model/src/model/impl_a2.rs @@ -164,7 +164,14 @@ impl TransformerModel { return Ok(val); }; let stream = self.gpu.default_stream(); - let world = self.config.ep_world_size; + // Loop over the ranks of the ACTUAL communicator: under pure TP + // (`--tp-size 2 --ep-size 1`) `ep_world_size` is 1 but the comm + // spans `tp_world_size` ranks — looping only `0..1` would leave + // the head min-reducing over its own value alone (asymmetric + // agreement → proc_count mismatch → collective deadlock on a warm + // cache-hit divergence). For EP-only and overlapping TP==EP + // topologies `max()` is identical to the previous value. + let world = self.config.ep_world_size.max(self.config.tp_world_size); let mut min_val = val; for root in 0..world { let v = if comm.rank() == root { @@ -198,16 +205,39 @@ impl TransformerModel { /// var). Disagreement causes the worker to misread the next u32 as a /// command code and is the kind of misconfiguration we want to fail /// loudly in development — there's no graceful fallback. + /// True when the head↔worker command protocol must be live: a + /// multi-rank NCCL world exists — EP **or** pure TP. Under + /// `--tp-size 2 --ep-size 1` (GDN HeadParallel 2-node) rank>0 still + /// runs the command-driven worker loop, so the head MUST emit the + /// same wire protocol as EP mode. + /// + /// Root cause of the 2026-07-03 2-node GDN deadlock: every broadcast + /// helper gated on `ep_world_size > 1` alone, so under pure TP the + /// head silently dropped ALL commands (prefill cmd + args + decode + /// cmds) while `ep_broadcast_tokens` (gated only on `comm`) still + /// fired — the rank-1 worker's 4-byte cmd recv paired with the head's + /// token-bulk broadcast, read `prompt_tokens[0]` as a decode command, + /// decoded (and CUDA-graph-captured) while the head was mid-prefill, + /// and both ranks wedged in shape-mismatched collectives (head stuck + /// in `sample_first_token` D2H, worker in the next cmd recv, both + /// GPUs spinning in NCCL kernels). + pub(crate) fn multi_rank_protocol_active(&self) -> bool { + self.comm.is_some() + && (self.config.ep_world_size > 1 || self.config.tp_world_size > 1) + } + pub(super) fn ep_broadcast_seq_and_cmd(&self, seq_id: u32, cmd: u32, v2: bool) -> Result<()> { - // No-op unless EP is actually active. The per-seq broadcast helpers - // (`ep_broadcast_cmd_for_seq`) are called unconditionally from the - // head's prefill / decode / mtp / lifecycle paths, exactly like the - // original `ep_broadcast_cmd` which no-ops here via - // `ep_broadcast_cmd_dispatch`. Without this guard `ep_broadcast_u32` - // panics ("ep_broadcast_u32 without comm") on every single-GPU - // generation, since `self.comm` is `None`. (Regression from the EP=2 - // slot-mux work, which only exercised the 2-rank path.) - if !(self.comm.is_some() && self.config.ep_world_size > 1) { + // No-op unless a multi-rank worker protocol is active (EP or pure + // TP — see `multi_rank_protocol_active`). The per-seq broadcast + // helpers (`ep_broadcast_cmd_for_seq`) are called unconditionally + // from the head's prefill / decode / mtp / lifecycle paths, + // exactly like the original `ep_broadcast_cmd` which no-ops here + // via `ep_broadcast_cmd_dispatch`. Without this guard + // `ep_broadcast_u32` panics ("ep_broadcast_u32 without comm") on + // every single-GPU generation, since `self.comm` is `None`. + // (Regression from the EP=2 slot-mux work, which only exercised + // the 2-rank path.) + if !self.multi_rank_protocol_active() { return Ok(()); } if v2 { @@ -243,7 +273,7 @@ impl TransformerModel { seq_ids: &[u32], tokens: &[u32], ) -> Result<()> { - if !(self.comm.is_some() && self.config.ep_world_size > 1) { + if !self.multi_rank_protocol_active() { return Ok(()); } debug_assert!( diff --git a/crates/spark-model/src/model/trait_impl/ep_misc.rs b/crates/spark-model/src/model/trait_impl/ep_misc.rs index 6c3415c7c..143ec825b 100644 --- a/crates/spark-model/src/model/trait_impl/ep_misc.rs +++ b/crates/spark-model/src/model/trait_impl/ep_misc.rs @@ -36,7 +36,13 @@ impl TransformerModel { } pub(super) fn is_ep_dispatch(&self) -> bool { - self.comm.is_some() && self.config.ep_world_size > 1 + // "EP" here means "the multi-rank head↔worker command protocol is + // active" — true for EP sharding AND for pure TP (GDN HeadParallel + // `--tp-size 2 --ep-size 1`, where rank>0 also runs the command + // loop). The scheduler consults this to disable single-GPU-only + // fused paths (mixed_forward / co-dispatch) that have no worker + // wire protocol; those must be off for ANY multi-rank world. + self.multi_rank_protocol_active() } pub(super) fn is_mla_dispatch(&self) -> bool { @@ -57,7 +63,10 @@ impl TransformerModel { } pub(super) fn ep_broadcast_cmd_dispatch(&self, cmd: u32) -> Result<()> { - if self.comm.is_some() && self.config.ep_world_size > 1 { + // Gate matches `ep_broadcast_seq_and_cmd`: live for EP AND pure TP + // (see `multi_rank_protocol_active`). Gating on `ep_world_size` + // alone dropped the prefill chunk args under `--tp-size 2`. + if self.multi_rank_protocol_active() { self.ep_broadcast_u32(cmd)?; } Ok(()) diff --git a/crates/spark-model/src/model/trait_impl/prefill_b.rs b/crates/spark-model/src/model/trait_impl/prefill_b.rs index f79377130..1db6eae8c 100644 --- a/crates/spark-model/src/model/trait_impl/prefill_b.rs +++ b/crates/spark-model/src/model/trait_impl/prefill_b.rs @@ -128,9 +128,11 @@ impl TransformerModel { ); } - // Use the caller-provided stream for compute-copy overlap, - // unless EP is active (NCCL requires the default stream). - let stream = if self.comm.is_some() && self.config.ep_world_size > 1 { + // Use the caller-provided stream for compute-copy overlap, unless + // a multi-rank world is active (EP or pure TP — NCCL collectives + // must stay stream-ordered with the cmd broadcasts, which run on + // the default stream). + let stream = if self.multi_rank_protocol_active() { self.gpu.default_stream() } else { stream diff --git a/crates/spark-model/src/model/trait_impl/prefill_b/batch.rs b/crates/spark-model/src/model/trait_impl/prefill_b/batch.rs index e50e1fe8c..ca7e1658d 100644 --- a/crates/spark-model/src/model/trait_impl/prefill_b/batch.rs +++ b/crates/spark-model/src/model/trait_impl/prefill_b/batch.rs @@ -184,8 +184,8 @@ impl TransformerModel { ); } - // EP active → NCCL needs the default stream. - let stream = if self.comm.is_some() && self.config.ep_world_size > 1 { + // Multi-rank world (EP or pure TP) → NCCL needs the default stream. + let stream = if self.multi_rank_protocol_active() { self.gpu.default_stream() } else { stream diff --git a/crates/spark-model/src/model/trait_impl/prefill_b/batch_kernel.rs b/crates/spark-model/src/model/trait_impl/prefill_b/batch_kernel.rs index 32cffcc04..0c9bb5bf1 100644 --- a/crates/spark-model/src/model/trait_impl/prefill_b/batch_kernel.rs +++ b/crates/spark-model/src/model/trait_impl/prefill_b/batch_kernel.rs @@ -87,8 +87,8 @@ impl TransformerModel { .max() .unwrap_or(chunk_len); - // EP active → NCCL needs the default stream. - let stream = if self.comm.is_some() && self.config.ep_world_size > 1 { + // Multi-rank world (EP or pure TP) → NCCL needs the default stream. + let stream = if self.multi_rank_protocol_active() { self.gpu.default_stream() } else { stream diff --git a/crates/spark-model/src/model/trait_impl/prefill_b/prefix_lookup.rs b/crates/spark-model/src/model/trait_impl/prefill_b/prefix_lookup.rs index 8c7037343..2da1d2113 100644 --- a/crates/spark-model/src/model/trait_impl/prefill_b/prefix_lookup.rs +++ b/crates/spark-model/src/model/trait_impl/prefill_b/prefix_lookup.rs @@ -53,7 +53,9 @@ impl TransformerModel { // Calling unconditionally on EP active fixes that — when // either side has matched=0 the agreed value is 0 and we // simply fall through to the no-cache path on both sides. - let ep_active = self.comm.is_some() && self.config.ep_world_size > 1; + // EP *or* pure TP: any multi-rank world must agree on `matched` + // (rank-local prefix caches can diverge in either topology). + let ep_active = self.multi_rank_protocol_active(); if ep_active { let local = prefix_match.matched_tokens as u32; let agreed = self.ep_min_u32(local)? as usize; diff --git a/crates/spark-model/src/model/trait_impl/prefill_c.rs b/crates/spark-model/src/model/trait_impl/prefill_c.rs index 4fa08379e..32d14d10b 100644 --- a/crates/spark-model/src/model/trait_impl/prefill_c.rs +++ b/crates/spark-model/src/model/trait_impl/prefill_c.rs @@ -85,9 +85,10 @@ impl TransformerModel { return self.prefill_chunk(tokens, seq, 0, total_len, true, stream); } - // Use the caller-provided stream for compute-copy overlap, - // unless EP is active (NCCL requires the default stream). - let stream = if self.comm.is_some() && self.config.ep_world_size > 1 { + // Use the caller-provided stream for compute-copy overlap, unless + // a multi-rank world is active (EP or pure TP — NCCL requires the + // default stream). + let stream = if self.multi_rank_protocol_active() { self.gpu.default_stream() } else { stream From cb7804e2822dd0f2d666fa53cd67a1d1b4aef352 Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Fri, 3 Jul 2026 03:06:33 -0400 Subject: [PATCH 04/15] fix(gdn-tp): replicate gated-RMSNorm [vd] weight under HeadParallel + shard debug logs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The GDN gated-RMSNorm gain is over the value HEAD-DIM ([vd]=128), shared across all value heads — NOT per-head [nv*vd]. The HeadParallel loader sharded it on the value-head axis via shard_gdn_value_vector(unit=vd), reading local_nv*vd*2 = 4096 bytes from the [vd] (~256-512B) buffer -> cuMemcpyDtoDAsync INVALID_VALUE at model build (both loader sites: FP8-dense + NVFP4). Fix: replicate norm.weight (every rank keeps the full [vd]); a_log/dt_bias stay sharded (they ARE per-head [nv]). tp=1 unaffected (slicer was a no-op there). Also add permanent tracing::debug! (target spark_model::tp_shard / ::concat) to slice_segments, shard_gdn_value_vector, dense row-parallel, and gpu_concat_rows so TP shard sizing is inspectable via RUST_LOG=spark_model::tp_shard=debug instead of ad-hoc eprintln next time. --- crates/spark-model/src/tp_shard.rs | 19 +++++++++++++++++++ .../qwen35/load_layers/linear_attn_arms.rs | 12 ++++++++++-- crates/spark-model/src/weight_map/fp8_lut.rs | 6 ++++++ 3 files changed, 35 insertions(+), 2 deletions(-) diff --git a/crates/spark-model/src/tp_shard.rs b/crates/spark-model/src/tp_shard.rs index e8537a38e..c8a1f2180 100644 --- a/crates/spark-model/src/tp_shard.rs +++ b/crates/spark-model/src/tp_shard.rs @@ -96,6 +96,12 @@ pub fn shard_dense_bf16( // Per-row strided copy: row r of dst comes from row r of src, // starting at column `tp_rank * local_in`. let col_offset_bytes = tp_rank * local_row_bytes; + tracing::debug!( + target: "spark_model::tp_shard", + out_dim, in_dim, local_in, src_row_bytes, local_row_bytes, + tp_rank, tp_size, src = src.0, + "dense row-parallel shard (per-row strided)" + ); for r in 0..out_dim { let src_off = r * src_row_bytes + col_offset_bytes; let dst_off = r * local_row_bytes; @@ -523,6 +529,12 @@ fn slice_segments( let row_bytes = row_elems * elem_bytes; let (ops, local_rows) = segment_copy_plan(segments, row_bytes, tp_rank, tp_size)?; let dst = gpu.alloc(local_rows * row_bytes)?; + tracing::debug!( + target: "spark_model::tp_shard", + ?segments, full_rows, row_elems, elem_bytes, local_rows, + tp_rank, tp_size, src = src.0, + "gdn segmented row-slice" + ); for op in &ops { gpu.copy_d2d(src.offset(op.src_off), dst.offset(op.dst_off), op.len)?; } @@ -649,6 +661,13 @@ pub fn shard_gdn_value_vector( let local_bytes = local_len * elem_bytes; let dst = gpu.alloc(local_bytes)?; let src_off = dims.tp_rank * local_bytes; + tracing::debug!( + target: "spark_model::tp_shard", + full_nv = dims.full_nv, local_nv = dims.local_nv, unit, elem_bytes, + full_len, local_len, local_bytes, src_off, tp_rank = dims.tp_rank, + src = src.0, + "gdn value-vector shard (per-value-head axis)" + ); gpu.copy_d2d(src.offset(src_off), dst, local_bytes)?; Ok((dst, local_len)) } diff --git a/crates/spark-model/src/weight_loader/qwen35/load_layers/linear_attn_arms.rs b/crates/spark-model/src/weight_loader/qwen35/load_layers/linear_attn_arms.rs index 836a1c875..da9701a35 100644 --- a/crates/spark-model/src/weight_loader/qwen35/load_layers/linear_attn_arms.rs +++ b/crates/spark-model/src/weight_loader/qwen35/load_layers/linear_attn_arms.rs @@ -285,7 +285,11 @@ pub(super) fn build_linear_attention_dense_bf16( let (conv_ptr, _, _) = shard_gdn_conv_rows(ssm35.conv1d.weight, &dims, d_conv, gpu)?; let (a_log_ptr, _) = shard_gdn_value_vector(ssm35.a_log.weight, &dims, 1, 4, gpu)?; let (dt_bias_ptr, _) = shard_gdn_value_vector(ssm35.dt_bias.weight, &dims, 1, 4, gpu)?; - let (norm_ptr, _) = shard_gdn_value_vector(ssm35.norm.weight, &dims, dims.vd, 2, gpu)?; + // norm.weight is the gated-RMSNorm gain over the value HEAD-DIM ([vd]), + // SHARED across all value heads — REPLICATE under HeadParallel. (a_log/dt_bias + // above ARE per-head [nv] scalars, so they slice; slicing norm on the head + // axis read past the [vd] buffer → cuMemcpyDtoDAsync INVALID_VALUE at load.) + let norm_ptr = ssm35.norm.weight; let (out_proj_ptr, _, _) = shard_gdn_out_proj_row_parallel(ssm35.out_proj.weight, &dims, gpu)?; let ssm = SsmWeights { @@ -394,7 +398,11 @@ pub(super) fn build_linear_attention_nvfp4( let (conv_ptr, _, _) = shard_gdn_conv_rows(ssm35.conv1d.weight, &dims, d_conv, gpu)?; let (a_log_ptr, _) = shard_gdn_value_vector(ssm35.a_log.weight, &dims, 1, 4, gpu)?; let (dt_bias_ptr, _) = shard_gdn_value_vector(ssm35.dt_bias.weight, &dims, 1, 4, gpu)?; - let (norm_ptr, _) = shard_gdn_value_vector(ssm35.norm.weight, &dims, dims.vd, 2, gpu)?; + // norm.weight is the gated-RMSNorm gain over the value HEAD-DIM ([vd]), + // SHARED across all value heads — REPLICATE under HeadParallel. (a_log/dt_bias + // above ARE per-head [nv] scalars, so they slice; slicing norm on the head + // axis read past the [vd] buffer → cuMemcpyDtoDAsync INVALID_VALUE at load.) + let norm_ptr = ssm35.norm.weight; let conv1d_local = DenseWeight { weight: conv_ptr }; let a_log_local = DenseWeight { weight: a_log_ptr }; let dt_bias_local = DenseWeight { weight: dt_bias_ptr }; diff --git a/crates/spark-model/src/weight_map/fp8_lut.rs b/crates/spark-model/src/weight_map/fp8_lut.rs index 00bf543ed..ff7e6cfd1 100644 --- a/crates/spark-model/src/weight_map/fp8_lut.rs +++ b/crates/spark-model/src/weight_map/fp8_lut.rs @@ -470,6 +470,12 @@ pub(crate) fn gpu_concat_rows( let b_bytes = b_rows * k * 2; let total = a_bytes + b_bytes; let buf = gpu.alloc(total)?; + tracing::debug!( + target: "spark_model::weight_map::concat", + a_rows, b_rows, k, a_bytes, b_bytes, total, + a = a.weight.0, b = b.weight.0, buf = buf.0, + "gpu concat rows [a;b]" + ); gpu.copy_d2d(a.weight, buf, a_bytes)?; gpu.copy_d2d(b.weight, buf.offset(a_bytes), b_bytes)?; Ok(DenseWeight { weight: buf }) From 8f249588a5f072949f621cfdcb420db73742fba7 Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Mon, 29 Jun 2026 03:12:13 -0400 Subject: [PATCH 05/15] =?UTF-8?q?fix(scheduler):=20sort=20decode=20batch?= =?UTF-8?q?=20by=20SSM=20slot=20=E2=80=94=20enable=20batched-recurrent=20+?= =?UTF-8?q?=20graph=20contiguity?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The active decode list arrives in reverse-arrival order ([7,6,..,0] for 8 seqs), so batch position i mapped to SSM pool slot (n-1-i). Both the batched-recurrent SSM (ssm_batched_recurrent.rs contiguity check: states[i]==base+i*stride) and the multi-seq CUDA-graph capture (slot==i assumption) then FAIL → fall back to the eager per-seq serial loop → no concurrency scaling. The pool assigns consecutive slots {0..n-1}; sorting the active list ascending by SSM slot (SequenceState::ssm_slot_idx, new pub accessor) makes position i ↔ slot i, satisfying the contiguity invariant so the batched paths engage. Reordering the whole ActiveSeq keeps the position->seq logits mapping consistent; safe for all decode paths. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01AoQXbv8a4aVXFPqj8nQqYN (cherry picked from commit 0b5cf608b7125821e13e3585b0eb033af0d48211) --- crates/spark-model/src/traits.rs | 9 +++++++++ crates/spark-server/src/scheduler/decode_step.rs | 14 ++++++++++++++ 2 files changed, 23 insertions(+) diff --git a/crates/spark-model/src/traits.rs b/crates/spark-model/src/traits.rs index 45d0e0016..b9b4464cf 100644 --- a/crates/spark-model/src/traits.rs +++ b/crates/spark-model/src/traits.rs @@ -179,6 +179,15 @@ pub struct SequenceState { } impl SequenceState { + /// SSM-pool slot index for this sequence, if it has GDN/SSM (linear-attn) + /// layers. Used by the scheduler to order the decode batch by slot so the + /// batched-recurrent SSM + CUDA-graph contiguity invariant holds + /// (position i ↔ pool_base + i*stride). `None` for pure-attention models. + #[inline] + pub fn ssm_slot_idx(&self) -> Option { + self.ssm_slot.as_ref().and_then(|g| g.idx()) + } + /// Phase 6.3 sliding-window helper: the absolute logical block index /// of `block_table[0]`. Returns 0 when `--high-speed-swap` is off /// (`disk_block_ids` is empty then; `block_table` is the full history). diff --git a/crates/spark-server/src/scheduler/decode_step.rs b/crates/spark-server/src/scheduler/decode_step.rs index ac68ad6b1..da991862a 100644 --- a/crates/spark-server/src/scheduler/decode_step.rs +++ b/crates/spark-server/src/scheduler/decode_step.rs @@ -17,6 +17,20 @@ pub fn step_decode_only( ) { let t0 = std::time::Instant::now(); let n = active.len(); + // Batched decode (CUDA-graph replay + batched-recurrent SSM) requires the + // active sequences in SSM-pool-slot order, so batch position i maps to a + // contiguous state address (pool_base + i*stride). The pool assigns + // consecutive slots but the active list is in reverse-arrival order + // ([7,6,..,0] for 8 seqs), which fails the contiguity check in + // ssm_batched_recurrent.rs and the graph-capture slot==i assumption, + // forcing the eager per-seq loop (no concurrency scaling). Sort ascending + // by SSM slot (falling back to KV slot for non-SSM models) so the + // contiguous-slot invariant holds and the batched paths engage. The whole + // ActiveSeq is reordered, so the post-decode position->seq mapping stays + // consistent. + if n > 1 { + active.sort_by_key(|a| a.seq.ssm_slot_idx().unwrap_or(a.seq.slot_idx as usize)); + } let tokens: Vec = active.iter().map(|a| a.last_token).collect(); // CONCURRENT-DECODE DIAG: per-step batch state (slot, seq_len, etc). From 813e7985fa054eb6d376c713dbdd05786c701c8f Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Mon, 29 Jun 2026 07:54:56 -0400 Subject: [PATCH 06/15] fix(ssm): conv input-stride bug in batched-recurrent decode (n>=2 corruption) The batched-projection SSM recurrent decode launched conv1d_update_l2norm once with batch=n on the base deinterleaved pointer. The conv kernel strides its input by dim(=conv_dim=8192), but deinterleaved is the QKVZ projection output laid out [Q|K|V|Z] with stride qkvz_size(=12288). So seq b>=1 read from b*conv_dim instead of b*qkvz_size, pulling the previous seq's Z-gate region into the GDN scan: correct at n=1, corrupt at n>=2 (measured 0/3 at n=4). Mirror the proven per-token pattern in trait_decode_batched_conv_gdn.rs: run the cheap per-channel conv per-seq (batch=1) with pre-offset pointers so the qkvz_size input stride and conv_dim output stride are both honored. The expensive GDN scan stays fully batched. Diagnosis via subagent root-cause read. Co-Authored-By: Claude Opus 4.8 (1M context) Claude-Session: https://claude.ai/code/session_01AoQXbv8a4aVXFPqj8nQqYN (cherry picked from commit 632e9861efa0e9306ec78f044fc5a2b6b0ede992) --- .../ssm_batched_recurrent.rs | 43 ++++++++++++------- 1 file changed, 28 insertions(+), 15 deletions(-) diff --git a/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched_recurrent.rs b/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched_recurrent.rs index e1294cec1..48b334d11 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched_recurrent.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched_recurrent.rs @@ -112,21 +112,34 @@ impl Qwen3SsmLayer { detail_step!("recurrent_batched_ba"); let conv_out = ctx.buffers.ssm_conv_out_f32(); - ops::conv1d_update_l2norm( - ctx.gpu, - self.conv1d_l2norm_f32_k, - conv_state_base, - deinterleaved, - &self.ssm.conv1d, - conv_out, - conv_dim, - d_conv, - n as u32, - qk_channels, - kd as u32, - 1e-6, - stream, - )?; + // CONV INPUT-STRIDE FIX: the conv kernel strides its input by `dim` + // (= conv_dim), but `deinterleaved` (the QKVZ-projection output) is + // laid out [Q|K|V|Z] with stride `qkvz_size` (> conv_dim). A single + // batched launch (batch=n) would read seq b>=1 from `b*conv_dim` + // instead of `b*qkvz_size`, pulling in the previous seq's Z-gate + // region → garbage into the GDN scan (correct at n=1, corrupt at + // n>=2). Mirror the proven per-token pattern in + // `trait_decode_batched_conv_gdn.rs`: run the cheap conv per-seq + // (batch=1) with pre-offset pointers so the qkvz_size input stride + // and conv_dim output stride are both honored. The expensive GDN + // scan below stays fully batched. + for i in 0..n { + ops::conv1d_update_l2norm( + ctx.gpu, + self.conv1d_l2norm_f32_k, + conv_state_base.offset(i * self.conv_state_bytes), + deinterleaved.offset(i * qkvz_size * bf16), + &self.ssm.conv1d, + conv_out.offset(i * conv_dim as usize * 4), // FP32 output, conv_dim-strided + conv_dim, + d_conv, + 1, + qk_channels, + kd as u32, + 1e-6, + stream, + )?; + } detail_step!("recurrent_batched_conv"); if self.gdn_f32_strided_norm_k.0 != 0 From cb7361325ad64051a29fb0b6e86e3c1d015fa11c Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Fri, 3 Jul 2026 11:30:29 -0400 Subject: [PATCH 07/15] ci(gdn-tp): file-size split (tp_shard/gdn.rs) + fmt + clippy/typos green MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Split the GDN HeadParallel slicers out of tp_shard.rs into tp_shard/gdn.rs (704 -> 343 LoC; both under the 500 cap) — drops tp_shard.rs from the file-size allow-list; allow-list linear_attn_arms.rs (548) + trait_decode_batched.rs (503). - clippy: repeat().take() -> std::iter::repeat_n() in tp_shard/tests.rs. - typos: alog_src_off -> a_log_src_off. - cargo fmt. --- .github/workflows/file-size-cap.yml | 8 + crates/spark-model/src/model/impl_a2.rs | 3 +- crates/spark-model/src/tp_shard.rs | 361 +---------------- crates/spark-model/src/tp_shard/gdn.rs | 374 ++++++++++++++++++ crates/spark-model/src/tp_shard/tests.rs | 51 ++- .../qwen35/load_layers/linear_attn_arms.rs | 8 +- 6 files changed, 430 insertions(+), 375 deletions(-) create mode 100644 crates/spark-model/src/tp_shard/gdn.rs diff --git a/.github/workflows/file-size-cap.yml b/.github/workflows/file-size-cap.yml index 605d8be1c..9d1544d0f 100644 --- a/.github/workflows/file-size-cap.yml +++ b/.github/workflows/file-size-cap.yml @@ -246,6 +246,14 @@ jobs: # decode dispatchers above; conv/GDN halves were already # extracted to trait_decode_batched_conv_gdn.rs. Clean split # tracked as follow-up. + # 548 LoC — Qwen3.5 linear-attn (SSM) loader arm; grew with the + # HeadParallel config-localize + per-rank segmented weight loads + # (FP8-dense + NVFP4 builders). Two linear load sequences read + # top-to-bottom; clean split tracked as follow-up. + "crates/spark-model/src/weight_loader/qwen35/load_layers/linear_attn_arms.rs" + # 503 LoC — batched-recurrent GDN decode; marginally over (added the + # per-seq SSM out_proj all-reduce for HeadParallel). Back under 500 + # once the batched/graph paths consolidate. Follow-up. "crates/spark-model/src/layers/qwen3_ssm/trait_decode_batched.rs" # -- #300/#301 (mixed-precision NVFP4 weights) carry-over -- # 522 LoC; landed on main 2026-07-12 without an entry, so the diff --git a/crates/spark-model/src/model/impl_a2.rs b/crates/spark-model/src/model/impl_a2.rs index 5a7d14a0a..5f82ef4e0 100644 --- a/crates/spark-model/src/model/impl_a2.rs +++ b/crates/spark-model/src/model/impl_a2.rs @@ -222,8 +222,7 @@ impl TransformerModel { /// in `sample_first_token` D2H, worker in the next cmd recv, both /// GPUs spinning in NCCL kernels). pub(crate) fn multi_rank_protocol_active(&self) -> bool { - self.comm.is_some() - && (self.config.ep_world_size > 1 || self.config.tp_world_size > 1) + self.comm.is_some() && (self.config.ep_world_size > 1 || self.config.tp_world_size > 1) } pub(super) fn ep_broadcast_seq_and_cmd(&self, seq_id: u32, cmd: u32, v2: bool) -> Result<()> { diff --git a/crates/spark-model/src/tp_shard.rs b/crates/spark-model/src/tp_shard.rs index c8a1f2180..018bfaaee 100644 --- a/crates/spark-model/src/tp_shard.rs +++ b/crates/spark-model/src/tp_shard.rs @@ -333,365 +333,8 @@ impl TpMoeDims { } } -// ════════════════════════════════════════════════════════════════════ -// GDN HeadParallel — tensor-parallel sharding of the Gated-DeltaNet -// (linear_attention / SSM) layers for Qwen3.5 / 3.6. -// -// The GDN recurrence is embarrassingly parallel across *value-head groups*: -// each TP rank owns a contiguous range of key/value heads, runs the whole -// scan locally with LOCAL nk/nv/conv_dim, and the ranks reconcile with a -// single all-reduce after `out_proj` (row-parallel, exactly like attention -// after `o_proj`). No cross-rank comm inside the scan. -// -// CRITICAL LAYOUT FACT — the in-projection is stored as *segmented* -// contiguous blocks, NOT one flat matrix: -// -// in_proj_qkv : [Q | K | V] rows = nk·kd + nk·kd + nv·vd (= conv_dim) -// in_proj_z : [Z] rows = nv·vd -// → gpu_concat_rows → QKVZ [Q | K | V | Z] -// -// A naive "first out_dim/tp rows" slice is WRONG: it would give rank 0 the -// whole Q block plus part of K. Each segment (Q, K, V, Z) must be sliced by -// the LOCAL head range *independently*, then the local slices re-concatenated -// in the same [Q|K|V|Z] order. The `segment_copy_plan` below encodes exactly -// that: one contiguous copy per segment, packed back-to-back into the local -// buffer. -// -// The depthwise `conv1d` weight `[conv_dim, d_conv]` is sharded with the SAME -// segment pattern as QKV (its channels ARE the QKV channels — one filter per -// channel), NOT replicated. `a_log`/`dt_bias` (`[nv]` FP32), `norm` -// (`[nv·vd]` BF16) and `out_proj` (`[h, nv·vd]`, row-parallel) shard on the -// value-head axis. The BA gate buffer is per-group interleaved but the rank -// boundary always lands on a group boundary, so it slices contiguously. -// ════════════════════════════════════════════════════════════════════ - -/// Pre-TP-shard GDN (linear-attention / SSM) dimensions reconstructed from -/// `config`. -/// -/// Mirrors [`TpAttentionDims`]: `topology.rs` divides -/// `linear_num_key_heads` / `linear_num_value_heads` by `tp_world_size` at -/// startup, so by the time a loader runs `config` holds **per-rank-local** -/// head counts. The `full_*` fields multiply back up to the pre-shard sizes -/// that the segment slicers expect. Head *dims* (`kd`, `vd`) and the hidden -/// size `h` are never sharded. -#[derive(Debug, Clone, Copy)] -pub struct TpGdnDims { - pub tp_rank: usize, - /// `tp_world_size` clamped to `>= 1`. Loaders treat `tp_size == 1` as the - /// no-shard fast path. - pub tp_size: usize, - /// Hidden size (model embed dim) — never sharded. - pub h: usize, - /// Key head dim (`linear_key_head_dim`) — never sharded. - pub kd: usize, - /// Value head dim (`linear_value_head_dim`) — never sharded. - pub vd: usize, - /// Per-rank key heads (Q and K share this count). - pub local_nk: usize, - /// Full pre-shard key heads = `local_nk * tp_size`. - pub full_nk: usize, - /// Per-rank value heads. - pub local_nv: usize, - /// Full pre-shard value heads = `local_nv * tp_size`. - pub full_nv: usize, -} - -impl TpGdnDims { - pub fn from_config(config: &ModelConfig) -> Self { - let tp_size = config.tp_world_size.max(1); - let local_nk = config.linear_num_key_heads; - let local_nv = config.linear_num_value_heads; - Self { - tp_rank: config.tp_rank, - tp_size, - h: config.hidden_size, - kd: config.linear_key_head_dim, - vd: config.linear_value_head_dim, - local_nk, - full_nk: local_nk * tp_size, - local_nv, - full_nv: local_nv * tp_size, - } - } - - /// Full (pre-shard) key projection width: `full_nk * kd`. - pub fn full_key_dim(&self) -> usize { - self.full_nk * self.kd - } - /// Local key projection width: `local_nk * kd`. - pub fn local_key_dim(&self) -> usize { - self.local_nk * self.kd - } - /// Full (pre-shard) value projection width: `full_nv * vd`. - pub fn full_value_dim(&self) -> usize { - self.full_nv * self.vd - } - /// Local value projection width: `local_nv * vd`. - pub fn local_value_dim(&self) -> usize { - self.local_nv * self.vd - } - /// Full conv / QKV width: `2*full_nk*kd + full_nv*vd`. - pub fn full_conv_dim(&self) -> usize { - 2 * self.full_key_dim() + self.full_value_dim() - } - /// Local conv / QKV width: `2*local_nk*kd + local_nv*vd`. - pub fn local_conv_dim(&self) -> usize { - 2 * self.local_key_dim() + self.local_value_dim() - } - /// Full QKVZ out dim: `2*full_nk*kd + 2*full_nv*vd`. - pub fn full_qkvz_out(&self) -> usize { - self.full_conv_dim() + self.full_value_dim() - } - /// Local QKVZ out dim: `2*local_nk*kd + 2*local_nv*vd`. - pub fn local_qkvz_out(&self) -> usize { - self.local_conv_dim() + self.local_value_dim() - } - - /// Full-row segment list for the `[Q|K|V]` in-projection. - fn qkv_segments(&self) -> [usize; 3] { - [self.full_key_dim(), self.full_key_dim(), self.full_value_dim()] - } - /// Full-row segment list for the concatenated `[Q|K|V|Z]` in-projection. - fn qkvz_segments(&self) -> [usize; 4] { - [ - self.full_key_dim(), - self.full_key_dim(), - self.full_value_dim(), - self.full_value_dim(), - ] - } -} - -/// A single device-to-device copy in a segmented-slice plan. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct CopyOp { - src_off: usize, - dst_off: usize, - len: usize, -} - -/// Build the copy plan for a SEGMENTED row-slice. -/// -/// `segments` lists the full (pre-shard) row count of each contiguous block -/// (Q, K, V[, Z] for QKVZ). Each block is sliced *independently* to the local -/// rank's head range `[tp_rank * seg/tp, (tp_rank+1) * seg/tp)` and the local -/// slices are packed back-to-back into the output buffer, preserving segment -/// order. `row_bytes` is the byte width of one row (`in_dim * elem_bytes`). -/// -/// Returns `(ops, local_total_rows)`. Every segment must be divisible by -/// `tp_size` — the caller has already reconstructed `full_*` as -/// `local_* * tp_size`, so this holds by construction, but it is checked to -/// fail loudly on a mis-wired config rather than silently corrupt heads. -fn segment_copy_plan( - segments: &[usize], - row_bytes: usize, - tp_rank: usize, - tp_size: usize, -) -> Result<(Vec, usize)> { - ensure!(tp_rank < tp_size, "tp_rank {tp_rank} >= tp_size {tp_size}"); - let mut ops = Vec::with_capacity(segments.len()); - let mut src_rows = 0usize; // running offset into the source (in rows) - let mut dst_rows = 0usize; // running offset into the packed dst (in rows) - for (i, &seg) in segments.iter().enumerate() { - ensure!( - seg.is_multiple_of(tp_size), - "segment {i} ({seg} rows) not divisible by tp_size {tp_size}", - ); - let local = seg / tp_size; - ops.push(CopyOp { - src_off: (src_rows + tp_rank * local) * row_bytes, - dst_off: dst_rows * row_bytes, - len: local * row_bytes, - }); - src_rows += seg; - dst_rows += local; - } - Ok((ops, dst_rows)) -} - -/// Execute a segmented row-slice on the GPU. `row_elems` is the number of -/// elements per row (`in_dim`); `elem_bytes` its size (2 = BF16, 4 = FP32). -/// Returns `(local_ptr, local_total_rows)`. For `tp_size <= 1` returns the -/// source untouched (no allocation, caller must not double-free). -fn slice_segments( - src: DevicePtr, - segments: &[usize], - row_elems: usize, - elem_bytes: usize, - tp_rank: usize, - tp_size: usize, - gpu: &dyn GpuBackend, -) -> Result<(DevicePtr, usize)> { - let full_rows: usize = segments.iter().sum(); - if tp_size <= 1 { - return Ok((src, full_rows)); - } - let row_bytes = row_elems * elem_bytes; - let (ops, local_rows) = segment_copy_plan(segments, row_bytes, tp_rank, tp_size)?; - let dst = gpu.alloc(local_rows * row_bytes)?; - tracing::debug!( - target: "spark_model::tp_shard", - ?segments, full_rows, row_elems, elem_bytes, local_rows, - tp_rank, tp_size, src = src.0, - "gdn segmented row-slice" - ); - for op in &ops { - gpu.copy_d2d(src.offset(op.src_off), dst.offset(op.dst_off), op.len)?; - } - Ok((dst, local_rows)) -} - -/// Shard the `[Q|K|V]` (`in_proj_qkv`) BF16 weight `[full_conv_dim, h]` to the -/// local rank's `[local_conv_dim, h]`, slicing Q, K and V independently by the -/// local head range. Returns `(ptr, local_rows, h)`. -pub fn shard_gdn_qkv_rows( - src: DevicePtr, - dims: &TpGdnDims, - gpu: &dyn GpuBackend, -) -> Result<(DevicePtr, usize, usize)> { - let (ptr, rows) = slice_segments( - src, - &dims.qkv_segments(), - dims.h, - BF16_BYTES, - dims.tp_rank, - dims.tp_size, - gpu, - )?; - Ok((ptr, rows, dims.h)) -} - -/// Shard the concatenated `[Q|K|V|Z]` (`in_proj_qkvz`) BF16 weight -/// `[full_qkvz_out, h]` to the local rank's `[local_qkvz_out, h]`, slicing all -/// four segments independently. Returns `(ptr, local_rows, h)`. -pub fn shard_gdn_qkvz_rows( - src: DevicePtr, - dims: &TpGdnDims, - gpu: &dyn GpuBackend, -) -> Result<(DevicePtr, usize, usize)> { - let (ptr, rows) = slice_segments( - src, - &dims.qkvz_segments(), - dims.h, - BF16_BYTES, - dims.tp_rank, - dims.tp_size, - gpu, - )?; - Ok((ptr, rows, dims.h)) -} - -/// Shard the BA gate BF16 weight `[2*full_nv, h]` to `[2*local_nv, h]`. -/// -/// The interleave is per key-head group (`[β₀..β_{vpg-1}, α₀..α_{vpg-1}]` per -/// group, `vpg = nv/nk`), but rank `r` owns key-head groups -/// `[r*local_nk, (r+1)*local_nk)` which map to the contiguous row range -/// `[r*2*local_nv, (r+1)*2*local_nv)` — the rank boundary always lands on a -/// group boundary, so a single contiguous slice preserves the interleave. -pub fn shard_gdn_ba_rows( - src: DevicePtr, - dims: &TpGdnDims, - gpu: &dyn GpuBackend, -) -> Result<(DevicePtr, usize, usize)> { - // Group-boundary alignment guarantee: full_nk divisible by tp_size ⇒ - // each rank gets whole groups. - ensure!( - dims.full_nk.is_multiple_of(dims.tp_size), - "BA: full_nk {} not divisible by tp_size {}", - dims.full_nk, - dims.tp_size, - ); - let (ptr, rows) = slice_segments( - src, - &[2 * dims.full_nv], - dims.h, - BF16_BYTES, - dims.tp_rank, - dims.tp_size, - gpu, - )?; - Ok((ptr, rows, dims.h)) -} - -/// Shard the depthwise `conv1d` BF16 weight `[full_conv_dim, d_conv]` to -/// `[local_conv_dim, d_conv]`. Channels ARE the QKV channels (one filter per -/// channel), so this uses the SAME `[Q|K|V]` segment pattern as the QKV -/// in-projection — the conv is NOT replicated across ranks. -pub fn shard_gdn_conv_rows( - src: DevicePtr, - dims: &TpGdnDims, - d_conv: usize, - gpu: &dyn GpuBackend, -) -> Result<(DevicePtr, usize, usize)> { - let (ptr, rows) = slice_segments( - src, - &dims.qkv_segments(), - d_conv, - BF16_BYTES, - dims.tp_rank, - dims.tp_size, - gpu, - )?; - Ok((ptr, rows, d_conv)) -} - -/// Shard a per-value-head 1D vector on the value-head axis. Handles BF16 -/// (`norm`, `[full_nv*vd]` → `[local_nv*vd]` with `elem_bytes = 2`, -/// `unit = vd`) and FP32 (`a_log` / `dt_bias`, `[full_nv]` → `[local_nv]` with -/// `elem_bytes = 4`, `unit = 1`). `unit` is the number of elements per value -/// head. Returns `(ptr, local_len_elems)`. -pub fn shard_gdn_value_vector( - src: DevicePtr, - dims: &TpGdnDims, - unit: usize, - elem_bytes: usize, - gpu: &dyn GpuBackend, -) -> Result<(DevicePtr, usize)> { - let full_len = dims.full_nv * unit; - if dims.tp_size <= 1 { - return Ok((src, full_len)); - } - ensure!( - dims.tp_rank < dims.tp_size, - "tp_rank {} >= tp_size {}", - dims.tp_rank, - dims.tp_size, - ); - let local_len = dims.local_nv * unit; - let local_bytes = local_len * elem_bytes; - let dst = gpu.alloc(local_bytes)?; - let src_off = dims.tp_rank * local_bytes; - tracing::debug!( - target: "spark_model::tp_shard", - full_nv = dims.full_nv, local_nv = dims.local_nv, unit, elem_bytes, - full_len, local_len, local_bytes, src_off, tp_rank = dims.tp_rank, - src = src.0, - "gdn value-vector shard (per-value-head axis)" - ); - gpu.copy_d2d(src.offset(src_off), dst, local_bytes)?; - Ok((dst, local_len)) -} - -/// Shard the `out_proj` BF16 weight `[h, full_value_dim]` row-parallel on its -/// input dim (value_dim). Rank `r` keeps columns -/// `[r*local_value_dim, (r+1)*local_value_dim)` of every output row; the -/// partial products are summed with an all-reduce after the GEMM (mirrors -/// attention `o_proj`). Returns `(ptr, h, local_value_dim)`. -pub fn shard_gdn_out_proj_row_parallel( - src: DevicePtr, - dims: &TpGdnDims, - gpu: &dyn GpuBackend, -) -> Result<(DevicePtr, usize, usize)> { - shard_dense_bf16( - src, - dims.h, - dims.full_value_dim(), - TpShardKind::RowParallel, - dims.tp_rank, - dims.tp_size, - gpu, - ) -} +mod gdn; +pub use gdn::*; mod quant_shard; pub use quant_shard::{shard_fp8_block_scaled, shard_quantized_nvfp4}; diff --git a/crates/spark-model/src/tp_shard/gdn.rs b/crates/spark-model/src/tp_shard/gdn.rs new file mode 100644 index 000000000..400b8029a --- /dev/null +++ b/crates/spark-model/src/tp_shard/gdn.rs @@ -0,0 +1,374 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! GDN HeadParallel — tensor-parallel sharding of the Gated-DeltaNet (SSM / +//! linear-attention) layers. Split out of `tp_shard.rs` (file-size cap). + +use anyhow::{Result, ensure}; +use atlas_core::config::ModelConfig; +use spark_runtime::gpu::{DevicePtr, GpuBackend}; + +use super::{BF16_BYTES, TpShardKind, shard_dense_bf16}; + +// ════════════════════════════════════════════════════════════════════ +// GDN HeadParallel — tensor-parallel sharding of the Gated-DeltaNet +// (linear_attention / SSM) layers for Qwen3.5 / 3.6. +// +// The GDN recurrence is embarrassingly parallel across *value-head groups*: +// each TP rank owns a contiguous range of key/value heads, runs the whole +// scan locally with LOCAL nk/nv/conv_dim, and the ranks reconcile with a +// single all-reduce after `out_proj` (row-parallel, exactly like attention +// after `o_proj`). No cross-rank comm inside the scan. +// +// CRITICAL LAYOUT FACT — the in-projection is stored as *segmented* +// contiguous blocks, NOT one flat matrix: +// +// in_proj_qkv : [Q | K | V] rows = nk·kd + nk·kd + nv·vd (= conv_dim) +// in_proj_z : [Z] rows = nv·vd +// → gpu_concat_rows → QKVZ [Q | K | V | Z] +// +// A naive "first out_dim/tp rows" slice is WRONG: it would give rank 0 the +// whole Q block plus part of K. Each segment (Q, K, V, Z) must be sliced by +// the LOCAL head range *independently*, then the local slices re-concatenated +// in the same [Q|K|V|Z] order. The `segment_copy_plan` below encodes exactly +// that: one contiguous copy per segment, packed back-to-back into the local +// buffer. +// +// The depthwise `conv1d` weight `[conv_dim, d_conv]` is sharded with the SAME +// segment pattern as QKV (its channels ARE the QKV channels — one filter per +// channel), NOT replicated. `a_log`/`dt_bias` (`[nv]` FP32), `norm` +// (`[nv·vd]` BF16) and `out_proj` (`[h, nv·vd]`, row-parallel) shard on the +// value-head axis. The BA gate buffer is per-group interleaved but the rank +// boundary always lands on a group boundary, so it slices contiguously. +// ════════════════════════════════════════════════════════════════════ + +/// Pre-TP-shard GDN (linear-attention / SSM) dimensions reconstructed from +/// `config`. +/// +/// Mirrors [`TpAttentionDims`]: `topology.rs` divides +/// `linear_num_key_heads` / `linear_num_value_heads` by `tp_world_size` at +/// startup, so by the time a loader runs `config` holds **per-rank-local** +/// head counts. The `full_*` fields multiply back up to the pre-shard sizes +/// that the segment slicers expect. Head *dims* (`kd`, `vd`) and the hidden +/// size `h` are never sharded. +#[derive(Debug, Clone, Copy)] +pub struct TpGdnDims { + pub tp_rank: usize, + /// `tp_world_size` clamped to `>= 1`. Loaders treat `tp_size == 1` as the + /// no-shard fast path. + pub tp_size: usize, + /// Hidden size (model embed dim) — never sharded. + pub h: usize, + /// Key head dim (`linear_key_head_dim`) — never sharded. + pub kd: usize, + /// Value head dim (`linear_value_head_dim`) — never sharded. + pub vd: usize, + /// Per-rank key heads (Q and K share this count). + pub local_nk: usize, + /// Full pre-shard key heads = `local_nk * tp_size`. + pub full_nk: usize, + /// Per-rank value heads. + pub local_nv: usize, + /// Full pre-shard value heads = `local_nv * tp_size`. + pub full_nv: usize, +} + +impl TpGdnDims { + pub fn from_config(config: &ModelConfig) -> Self { + let tp_size = config.tp_world_size.max(1); + let local_nk = config.linear_num_key_heads; + let local_nv = config.linear_num_value_heads; + Self { + tp_rank: config.tp_rank, + tp_size, + h: config.hidden_size, + kd: config.linear_key_head_dim, + vd: config.linear_value_head_dim, + local_nk, + full_nk: local_nk * tp_size, + local_nv, + full_nv: local_nv * tp_size, + } + } + + /// Full (pre-shard) key projection width: `full_nk * kd`. + pub fn full_key_dim(&self) -> usize { + self.full_nk * self.kd + } + /// Local key projection width: `local_nk * kd`. + pub fn local_key_dim(&self) -> usize { + self.local_nk * self.kd + } + /// Full (pre-shard) value projection width: `full_nv * vd`. + pub fn full_value_dim(&self) -> usize { + self.full_nv * self.vd + } + /// Local value projection width: `local_nv * vd`. + pub fn local_value_dim(&self) -> usize { + self.local_nv * self.vd + } + /// Full conv / QKV width: `2*full_nk*kd + full_nv*vd`. + pub fn full_conv_dim(&self) -> usize { + 2 * self.full_key_dim() + self.full_value_dim() + } + /// Local conv / QKV width: `2*local_nk*kd + local_nv*vd`. + pub fn local_conv_dim(&self) -> usize { + 2 * self.local_key_dim() + self.local_value_dim() + } + /// Full QKVZ out dim: `2*full_nk*kd + 2*full_nv*vd`. + pub fn full_qkvz_out(&self) -> usize { + self.full_conv_dim() + self.full_value_dim() + } + /// Local QKVZ out dim: `2*local_nk*kd + 2*local_nv*vd`. + pub fn local_qkvz_out(&self) -> usize { + self.local_conv_dim() + self.local_value_dim() + } + + /// Full-row segment list for the `[Q|K|V]` in-projection. + fn qkv_segments(&self) -> [usize; 3] { + [ + self.full_key_dim(), + self.full_key_dim(), + self.full_value_dim(), + ] + } + /// Full-row segment list for the concatenated `[Q|K|V|Z]` in-projection. + fn qkvz_segments(&self) -> [usize; 4] { + [ + self.full_key_dim(), + self.full_key_dim(), + self.full_value_dim(), + self.full_value_dim(), + ] + } +} + +/// A single device-to-device copy in a segmented-slice plan. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct CopyOp { + src_off: usize, + dst_off: usize, + len: usize, +} + +/// Build the copy plan for a SEGMENTED row-slice. +/// +/// `segments` lists the full (pre-shard) row count of each contiguous block +/// (Q, K, V[, Z] for QKVZ). Each block is sliced *independently* to the local +/// rank's head range `[tp_rank * seg/tp, (tp_rank+1) * seg/tp)` and the local +/// slices are packed back-to-back into the output buffer, preserving segment +/// order. `row_bytes` is the byte width of one row (`in_dim * elem_bytes`). +/// +/// Returns `(ops, local_total_rows)`. Every segment must be divisible by +/// `tp_size` — the caller has already reconstructed `full_*` as +/// `local_* * tp_size`, so this holds by construction, but it is checked to +/// fail loudly on a mis-wired config rather than silently corrupt heads. +fn segment_copy_plan( + segments: &[usize], + row_bytes: usize, + tp_rank: usize, + tp_size: usize, +) -> Result<(Vec, usize)> { + ensure!(tp_rank < tp_size, "tp_rank {tp_rank} >= tp_size {tp_size}"); + let mut ops = Vec::with_capacity(segments.len()); + let mut src_rows = 0usize; // running offset into the source (in rows) + let mut dst_rows = 0usize; // running offset into the packed dst (in rows) + for (i, &seg) in segments.iter().enumerate() { + ensure!( + seg.is_multiple_of(tp_size), + "segment {i} ({seg} rows) not divisible by tp_size {tp_size}", + ); + let local = seg / tp_size; + ops.push(CopyOp { + src_off: (src_rows + tp_rank * local) * row_bytes, + dst_off: dst_rows * row_bytes, + len: local * row_bytes, + }); + src_rows += seg; + dst_rows += local; + } + Ok((ops, dst_rows)) +} + +/// Execute a segmented row-slice on the GPU. `row_elems` is the number of +/// elements per row (`in_dim`); `elem_bytes` its size (2 = BF16, 4 = FP32). +/// Returns `(local_ptr, local_total_rows)`. For `tp_size <= 1` returns the +/// source untouched (no allocation, caller must not double-free). +fn slice_segments( + src: DevicePtr, + segments: &[usize], + row_elems: usize, + elem_bytes: usize, + tp_rank: usize, + tp_size: usize, + gpu: &dyn GpuBackend, +) -> Result<(DevicePtr, usize)> { + let full_rows: usize = segments.iter().sum(); + if tp_size <= 1 { + return Ok((src, full_rows)); + } + let row_bytes = row_elems * elem_bytes; + let (ops, local_rows) = segment_copy_plan(segments, row_bytes, tp_rank, tp_size)?; + let dst = gpu.alloc(local_rows * row_bytes)?; + tracing::debug!( + target: "spark_model::tp_shard", + ?segments, full_rows, row_elems, elem_bytes, local_rows, + tp_rank, tp_size, src = src.0, + "gdn segmented row-slice" + ); + for op in &ops { + gpu.copy_d2d(src.offset(op.src_off), dst.offset(op.dst_off), op.len)?; + } + Ok((dst, local_rows)) +} + +/// Shard the `[Q|K|V]` (`in_proj_qkv`) BF16 weight `[full_conv_dim, h]` to the +/// local rank's `[local_conv_dim, h]`, slicing Q, K and V independently by the +/// local head range. Returns `(ptr, local_rows, h)`. +pub fn shard_gdn_qkv_rows( + src: DevicePtr, + dims: &TpGdnDims, + gpu: &dyn GpuBackend, +) -> Result<(DevicePtr, usize, usize)> { + let (ptr, rows) = slice_segments( + src, + &dims.qkv_segments(), + dims.h, + BF16_BYTES, + dims.tp_rank, + dims.tp_size, + gpu, + )?; + Ok((ptr, rows, dims.h)) +} + +/// Shard the concatenated `[Q|K|V|Z]` (`in_proj_qkvz`) BF16 weight +/// `[full_qkvz_out, h]` to the local rank's `[local_qkvz_out, h]`, slicing all +/// four segments independently. Returns `(ptr, local_rows, h)`. +pub fn shard_gdn_qkvz_rows( + src: DevicePtr, + dims: &TpGdnDims, + gpu: &dyn GpuBackend, +) -> Result<(DevicePtr, usize, usize)> { + let (ptr, rows) = slice_segments( + src, + &dims.qkvz_segments(), + dims.h, + BF16_BYTES, + dims.tp_rank, + dims.tp_size, + gpu, + )?; + Ok((ptr, rows, dims.h)) +} + +/// Shard the BA gate BF16 weight `[2*full_nv, h]` to `[2*local_nv, h]`. +/// +/// The interleave is per key-head group (`[β₀..β_{vpg-1}, α₀..α_{vpg-1}]` per +/// group, `vpg = nv/nk`), but rank `r` owns key-head groups +/// `[r*local_nk, (r+1)*local_nk)` which map to the contiguous row range +/// `[r*2*local_nv, (r+1)*2*local_nv)` — the rank boundary always lands on a +/// group boundary, so a single contiguous slice preserves the interleave. +pub fn shard_gdn_ba_rows( + src: DevicePtr, + dims: &TpGdnDims, + gpu: &dyn GpuBackend, +) -> Result<(DevicePtr, usize, usize)> { + // Group-boundary alignment guarantee: full_nk divisible by tp_size ⇒ + // each rank gets whole groups. + ensure!( + dims.full_nk.is_multiple_of(dims.tp_size), + "BA: full_nk {} not divisible by tp_size {}", + dims.full_nk, + dims.tp_size, + ); + let (ptr, rows) = slice_segments( + src, + &[2 * dims.full_nv], + dims.h, + BF16_BYTES, + dims.tp_rank, + dims.tp_size, + gpu, + )?; + Ok((ptr, rows, dims.h)) +} + +/// Shard the depthwise `conv1d` BF16 weight `[full_conv_dim, d_conv]` to +/// `[local_conv_dim, d_conv]`. Channels ARE the QKV channels (one filter per +/// channel), so this uses the SAME `[Q|K|V]` segment pattern as the QKV +/// in-projection — the conv is NOT replicated across ranks. +pub fn shard_gdn_conv_rows( + src: DevicePtr, + dims: &TpGdnDims, + d_conv: usize, + gpu: &dyn GpuBackend, +) -> Result<(DevicePtr, usize, usize)> { + let (ptr, rows) = slice_segments( + src, + &dims.qkv_segments(), + d_conv, + BF16_BYTES, + dims.tp_rank, + dims.tp_size, + gpu, + )?; + Ok((ptr, rows, d_conv)) +} + +/// Shard a per-value-head 1D vector on the value-head axis. Handles BF16 +/// (`norm`, `[full_nv*vd]` → `[local_nv*vd]` with `elem_bytes = 2`, +/// `unit = vd`) and FP32 (`a_log` / `dt_bias`, `[full_nv]` → `[local_nv]` with +/// `elem_bytes = 4`, `unit = 1`). `unit` is the number of elements per value +/// head. Returns `(ptr, local_len_elems)`. +pub fn shard_gdn_value_vector( + src: DevicePtr, + dims: &TpGdnDims, + unit: usize, + elem_bytes: usize, + gpu: &dyn GpuBackend, +) -> Result<(DevicePtr, usize)> { + let full_len = dims.full_nv * unit; + if dims.tp_size <= 1 { + return Ok((src, full_len)); + } + ensure!( + dims.tp_rank < dims.tp_size, + "tp_rank {} >= tp_size {}", + dims.tp_rank, + dims.tp_size, + ); + let local_len = dims.local_nv * unit; + let local_bytes = local_len * elem_bytes; + let dst = gpu.alloc(local_bytes)?; + let src_off = dims.tp_rank * local_bytes; + tracing::debug!( + target: "spark_model::tp_shard", + full_nv = dims.full_nv, local_nv = dims.local_nv, unit, elem_bytes, + full_len, local_len, local_bytes, src_off, tp_rank = dims.tp_rank, + src = src.0, + "gdn value-vector shard (per-value-head axis)" + ); + gpu.copy_d2d(src.offset(src_off), dst, local_bytes)?; + Ok((dst, local_len)) +} + +/// Shard the `out_proj` BF16 weight `[h, full_value_dim]` row-parallel on its +/// input dim (value_dim). Rank `r` keeps columns +/// `[r*local_value_dim, (r+1)*local_value_dim)` of every output row; the +/// partial products are summed with an all-reduce after the GEMM (mirrors +/// attention `o_proj`). Returns `(ptr, h, local_value_dim)`. +pub fn shard_gdn_out_proj_row_parallel( + src: DevicePtr, + dims: &TpGdnDims, + gpu: &dyn GpuBackend, +) -> Result<(DevicePtr, usize, usize)> { + shard_dense_bf16( + src, + dims.h, + dims.full_value_dim(), + TpShardKind::RowParallel, + dims.tp_rank, + dims.tp_size, + gpu, + ) +} diff --git a/crates/spark-model/src/tp_shard/tests.rs b/crates/spark-model/src/tp_shard/tests.rs index cf748c6ea..e92afd74b 100644 --- a/crates/spark-model/src/tp_shard/tests.rs +++ b/crates/spark-model/src/tp_shard/tests.rs @@ -106,10 +106,26 @@ fn qkvz_segment_plan_is_segmented_not_contiguous() { // Rank-1 local halves: Q[16..32], K[48..64], V[128..192], Z[256..320]. // Packed dst (rows): 0, 16, 32, 96. let want = [ - CopyOp { src_off: 16 * row_bytes, dst_off: 0, len: 16 * row_bytes }, - CopyOp { src_off: 48 * row_bytes, dst_off: 16 * row_bytes, len: 16 * row_bytes }, - CopyOp { src_off: 128 * row_bytes, dst_off: 32 * row_bytes, len: 64 * row_bytes }, - CopyOp { src_off: 256 * row_bytes, dst_off: 96 * row_bytes, len: 64 * row_bytes }, + CopyOp { + src_off: 16 * row_bytes, + dst_off: 0, + len: 16 * row_bytes, + }, + CopyOp { + src_off: 48 * row_bytes, + dst_off: 16 * row_bytes, + len: 16 * row_bytes, + }, + CopyOp { + src_off: 128 * row_bytes, + dst_off: 32 * row_bytes, + len: 64 * row_bytes, + }, + CopyOp { + src_off: 256 * row_bytes, + dst_off: 96 * row_bytes, + len: 64 * row_bytes, + }, ]; assert_eq!(ops, want); @@ -130,14 +146,13 @@ fn qkvz_two_rank_reconcat_tiles_full() { // Synthetic full weight: row r filled with value r (u16), h cols each. let full: Vec = (0..full_rows) - .flat_map(|r| std::iter::repeat(r as u16).take(h)) + .flat_map(|r| std::iter::repeat_n(r as u16, h)) .collect(); // CPU reference slice using the plan (byte offsets → row indices). let cpu_slice = |d: &TpGdnDims| -> Vec { let row_bytes = h * BF16_BYTES; - let (ops, local_rows) = - segment_copy_plan(&segs, row_bytes, d.tp_rank, d.tp_size).unwrap(); + let (ops, local_rows) = segment_copy_plan(&segs, row_bytes, d.tp_rank, d.tp_size).unwrap(); let mut out = vec![0u16; local_rows * h]; for op in &ops { let src_row = op.src_off / row_bytes; @@ -197,9 +212,21 @@ fn conv_segment_plan_uses_qkv_segments() { assert_eq!(local_rows, d.local_conv_dim()); // 96 // Rank 0: Q[0..16], K[32..48], V[64..128]; packed at 0,16,32. let want = [ - CopyOp { src_off: 0, dst_off: 0, len: 16 * row_bytes }, - CopyOp { src_off: 32 * row_bytes, dst_off: 16 * row_bytes, len: 16 * row_bytes }, - CopyOp { src_off: 64 * row_bytes, dst_off: 32 * row_bytes, len: 64 * row_bytes }, + CopyOp { + src_off: 0, + dst_off: 0, + len: 16 * row_bytes, + }, + CopyOp { + src_off: 32 * row_bytes, + dst_off: 16 * row_bytes, + len: 16 * row_bytes, + }, + CopyOp { + src_off: 64 * row_bytes, + dst_off: 32 * row_bytes, + len: 64 * row_bytes, + }, ]; assert_eq!(ops, want); } @@ -238,8 +265,8 @@ fn value_vector_offsets() { assert_eq!(norm_src_off, 128); // a_log: unit = 1, fp32. full = 8, local = 4. let f32_bytes = 4usize; - let alog_src_off = d.tp_rank * d.local_nv * f32_bytes; // rank1 = 4*4 - assert_eq!(alog_src_off, 16); + let a_log_src_off = d.tp_rank * d.local_nv * f32_bytes; // rank1 = 4*4 + assert_eq!(a_log_src_off, 16); } /// A non-divisible segment must be rejected loudly, not silently corrupt. diff --git a/crates/spark-model/src/weight_loader/qwen35/load_layers/linear_attn_arms.rs b/crates/spark-model/src/weight_loader/qwen35/load_layers/linear_attn_arms.rs index da9701a35..287811728 100644 --- a/crates/spark-model/src/weight_loader/qwen35/load_layers/linear_attn_arms.rs +++ b/crates/spark-model/src/weight_loader/qwen35/load_layers/linear_attn_arms.rs @@ -297,7 +297,9 @@ pub(super) fn build_linear_attention_dense_bf16( in_proj_ba: ba_dense, conv1d: DenseWeight { weight: conv_ptr }, a_log: DenseWeight { weight: a_log_ptr }, - dt_bias: DenseWeight { weight: dt_bias_ptr }, + dt_bias: DenseWeight { + weight: dt_bias_ptr, + }, norm: DenseWeight { weight: norm_ptr }, out_proj: QuantizedWeight::null(), }; @@ -405,7 +407,9 @@ pub(super) fn build_linear_attention_nvfp4( let norm_ptr = ssm35.norm.weight; let conv1d_local = DenseWeight { weight: conv_ptr }; let a_log_local = DenseWeight { weight: a_log_ptr }; - let dt_bias_local = DenseWeight { weight: dt_bias_ptr }; + let dt_bias_local = DenseWeight { + weight: dt_bias_ptr, + }; let norm_local = DenseWeight { weight: norm_ptr }; // out_proj is row-parallel: slice its input (value_dim) to local, then From a1c8894651357a9407f1d1f6116293b51b9fd96d Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Fri, 3 Jul 2026 11:36:37 -0400 Subject: [PATCH 08/15] fix(gdn-tp): expose tp_shard/gdn slice-plan internals to unit tests After splitting the GDN slicers into tp_shard/gdn.rs, the unit tests (a sibling module) could no longer reach segment_copy_plan / CopyOp / the qkv[z]_segments methods (were private, previously visible via super::* when co-located). Mark them pub(crate) + re-export CopyOp/segment_copy_plan from tp_shard so tests resolve them via super::*. --- crates/spark-model/src/tp_shard.rs | 3 +++ crates/spark-model/src/tp_shard/gdn.rs | 14 +++++++------- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/crates/spark-model/src/tp_shard.rs b/crates/spark-model/src/tp_shard.rs index 018bfaaee..c38feca68 100644 --- a/crates/spark-model/src/tp_shard.rs +++ b/crates/spark-model/src/tp_shard.rs @@ -335,6 +335,9 @@ impl TpMoeDims { mod gdn; pub use gdn::*; +// Internal slice-plan helpers (private API) — re-exported for the unit tests +// in `tp_shard/tests.rs`, which reach them via `use super::*`. +pub(crate) use gdn::{CopyOp, segment_copy_plan}; mod quant_shard; pub use quant_shard::{shard_fp8_block_scaled, shard_quantized_nvfp4}; diff --git a/crates/spark-model/src/tp_shard/gdn.rs b/crates/spark-model/src/tp_shard/gdn.rs index 400b8029a..dc4865c65 100644 --- a/crates/spark-model/src/tp_shard/gdn.rs +++ b/crates/spark-model/src/tp_shard/gdn.rs @@ -124,7 +124,7 @@ impl TpGdnDims { } /// Full-row segment list for the `[Q|K|V]` in-projection. - fn qkv_segments(&self) -> [usize; 3] { + pub(crate) fn qkv_segments(&self) -> [usize; 3] { [ self.full_key_dim(), self.full_key_dim(), @@ -132,7 +132,7 @@ impl TpGdnDims { ] } /// Full-row segment list for the concatenated `[Q|K|V|Z]` in-projection. - fn qkvz_segments(&self) -> [usize; 4] { + pub(crate) fn qkvz_segments(&self) -> [usize; 4] { [ self.full_key_dim(), self.full_key_dim(), @@ -144,10 +144,10 @@ impl TpGdnDims { /// A single device-to-device copy in a segmented-slice plan. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct CopyOp { - src_off: usize, - dst_off: usize, - len: usize, +pub(crate) struct CopyOp { + pub(crate) src_off: usize, + pub(crate) dst_off: usize, + pub(crate) len: usize, } /// Build the copy plan for a SEGMENTED row-slice. @@ -162,7 +162,7 @@ struct CopyOp { /// `tp_size` — the caller has already reconstructed `full_*` as /// `local_* * tp_size`, so this holds by construction, but it is checked to /// fail loudly on a mis-wired config rather than silently corrupt heads. -fn segment_copy_plan( +pub(crate) fn segment_copy_plan( segments: &[usize], row_bytes: usize, tp_rank: usize, From 2f633c80d1996750715679cdef04ad7ad1dd9e00 Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Fri, 3 Jul 2026 11:39:20 -0400 Subject: [PATCH 09/15] fix(gdn-tp): import gdn test-internals in tests.rs, not via lib re-export The pub(crate) re-export in tp_shard.rs was unused in the lib (non-test) build (deny(unused) -> error). Import CopyOp/segment_copy_plan directly in tp_shard/tests.rs from the gdn submodule instead. --- crates/spark-model/src/tp_shard.rs | 3 --- crates/spark-model/src/tp_shard/tests.rs | 3 +++ 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/crates/spark-model/src/tp_shard.rs b/crates/spark-model/src/tp_shard.rs index c38feca68..018bfaaee 100644 --- a/crates/spark-model/src/tp_shard.rs +++ b/crates/spark-model/src/tp_shard.rs @@ -335,9 +335,6 @@ impl TpMoeDims { mod gdn; pub use gdn::*; -// Internal slice-plan helpers (private API) — re-exported for the unit tests -// in `tp_shard/tests.rs`, which reach them via `use super::*`. -pub(crate) use gdn::{CopyOp, segment_copy_plan}; mod quant_shard; pub use quant_shard::{shard_fp8_block_scaled, shard_quantized_nvfp4}; diff --git a/crates/spark-model/src/tp_shard/tests.rs b/crates/spark-model/src/tp_shard/tests.rs index e92afd74b..aa0ce1f84 100644 --- a/crates/spark-model/src/tp_shard/tests.rs +++ b/crates/spark-model/src/tp_shard/tests.rs @@ -1,6 +1,9 @@ // SPDX-License-Identifier: AGPL-3.0-only use super::*; +// Internal slice-plan helpers live in the `gdn` submodule (pub(crate)); +// import them directly rather than re-exporting from the lib surface. +use super::gdn::{CopyOp, segment_copy_plan}; /// Validate the slice-offset math without exercising the GPU. #[test] From 6827057420c11e9dc0560e82b369894699e68028 Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Fri, 3 Jul 2026 11:40:26 -0400 Subject: [PATCH 10/15] fix(gdn-tp): drop redundant usize cast in decode slot-sort (clippy) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit slot_idx is already usize on this branch — the `as usize` from the original cherry-pick trips clippy::unnecessary_cast under deny(clippy::all). --- crates/spark-server/src/scheduler/decode_step.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/spark-server/src/scheduler/decode_step.rs b/crates/spark-server/src/scheduler/decode_step.rs index da991862a..d0708175d 100644 --- a/crates/spark-server/src/scheduler/decode_step.rs +++ b/crates/spark-server/src/scheduler/decode_step.rs @@ -29,7 +29,7 @@ pub fn step_decode_only( // ActiveSeq is reordered, so the post-decode position->seq mapping stays // consistent. if n > 1 { - active.sort_by_key(|a| a.seq.ssm_slot_idx().unwrap_or(a.seq.slot_idx as usize)); + active.sort_by_key(|a| a.seq.ssm_slot_idx().unwrap_or(a.seq.slot_idx)); } let tokens: Vec = active.iter().map(|a| a.last_token).collect(); From e4826634b0ffc31edb3110e1afdc2963648f46f7 Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Fri, 3 Jul 2026 02:40:38 -0400 Subject: [PATCH 11/15] fix(moe): gate EP all-reduce on ep_world_size, not world_size Under pure-TP (--tp-size 2 --ep-size 1: world_size=2, ep_world_size=1) the MoE experts are REPLICATED (local_expert_range returns all experts when ep<=1), so every rank computes the full identical MoE output. But the all-reduce + shared-expert-exclusion gated on comm.world_size()>1, summing two full outputs -> 2x per MoE layer -> deterministic garbage from the first token. Gate on config.ep_world_size>1 (the true 'experts are sharded' invariant) at all 15 sites across 9 forward variants. Single-node/EP/EP+TP unchanged; only pure-TP flips from wrongly-reducing to correctly-not. --- crates/spark-model/src/layers/moe/forward.rs | 4 ++-- crates/spark-model/src/layers/moe/forward_atomic_c4.rs | 4 ++-- crates/spark-model/src/layers/moe/forward_batched.rs | 2 +- crates/spark-model/src/layers/moe/forward_k2.rs | 4 ++-- crates/spark-model/src/layers/moe/forward_k3.rs | 4 ++-- crates/spark-model/src/layers/moe/forward_prefill.rs | 4 ++-- crates/spark-model/src/layers/moe/forward_prefill_bf16.rs | 2 +- crates/spark-model/src/layers/moe/forward_prefill_fp8.rs | 2 +- crates/spark-model/src/layers/moe/forward_token_major.rs | 4 ++-- 9 files changed, 15 insertions(+), 15 deletions(-) diff --git a/crates/spark-model/src/layers/moe/forward.rs b/crates/spark-model/src/layers/moe/forward.rs index 9e87ac885..58b0d1553 100644 --- a/crates/spark-model/src/layers/moe/forward.rs +++ b/crates/spark-model/src/layers/moe/forward.rs @@ -513,7 +513,7 @@ impl MoeLayer { // times. Solution: pass NULL shared_out for EP, all-reduce the routed sum, // then add shared_out once after all-reduce. let output = ctx.buffers.moe_output(); - let is_ep = ctx.comm.is_some_and(|c| c.world_size() > 1); + let is_ep = ctx.comm.is_some() && ctx.config.ep_world_size > 1; let shared_for_blend = if is_ep && !shared_out.is_null() { // EP: exclude shared expert from blend (will add after all-reduce). // Zero a temp buffer to pass as shared_out (kernel reads it even with NULL gate). @@ -544,7 +544,7 @@ impl MoeLayer { // Each rank only computed its local experts (remote → zero), so // SUM gives the correct global result. if let Some(comm) = ctx.comm - && comm.world_size() > 1 + && ctx.config.ep_world_size > 1 { if ctx.graph_capture { comm.all_reduce(output.0, h as usize * 2)?; diff --git a/crates/spark-model/src/layers/moe/forward_atomic_c4.rs b/crates/spark-model/src/layers/moe/forward_atomic_c4.rs index 31bd49acf..fcc96941e 100644 --- a/crates/spark-model/src/layers/moe/forward_atomic_c4.rs +++ b/crates/spark-model/src/layers/moe/forward_atomic_c4.rs @@ -165,7 +165,7 @@ impl MoeLayer { stream, )?; - let is_ep = ctx.comm.is_some_and(|c| c.world_size() > 1); + let is_ep = ctx.comm.is_some() && ctx.config.ep_world_size > 1; ops::moe_decode_atomic_c4_finalize( ctx.gpu, self.moe_decode_atomic_c4_finalize_k, @@ -181,7 +181,7 @@ impl MoeLayer { )?; if let Some(comm) = ctx.comm - && comm.world_size() > 1 + && ctx.config.ep_world_size > 1 { if ctx.graph_capture { comm.all_reduce(output.0, num_tokens * h as usize * 2)?; diff --git a/crates/spark-model/src/layers/moe/forward_batched.rs b/crates/spark-model/src/layers/moe/forward_batched.rs index d0df9eb06..7dbfe5a26 100644 --- a/crates/spark-model/src/layers/moe/forward_batched.rs +++ b/crates/spark-model/src/layers/moe/forward_batched.rs @@ -426,7 +426,7 @@ impl MoeLayer { // EP all-reduce per-token partial output if let Some(comm) = ctx.comm - && comm.world_size() > 1 + && ctx.config.ep_world_size > 1 { if ctx.graph_capture { comm.all_reduce(output_t.0, h as usize * 2)?; diff --git a/crates/spark-model/src/layers/moe/forward_k2.rs b/crates/spark-model/src/layers/moe/forward_k2.rs index 5836540cb..6e01c8f09 100644 --- a/crates/spark-model/src/layers/moe/forward_k2.rs +++ b/crates/spark-model/src/layers/moe/forward_k2.rs @@ -26,7 +26,7 @@ impl MoeLayer { // per-token BF16 decode kernels). Otherwise fall back to the per-token // BF16 batched path (SSOT: reuses the decode BF16 kernels via // forward_batched), which produces the same moe_output()[2,H]. - let is_ep = ctx.comm.is_some_and(|c| c.world_size() > 1); + let is_ep = ctx.comm.is_some() && ctx.config.ep_world_size > 1; let use_bf16_batch2 = self.bf16_gate_weight_ptrs.is_some() && self.moe_expert_gate_up_shared_bf16_batch2_k.0 != 0 && self.moe_expert_silu_down_shared_bf16_batch2_k.0 != 0 @@ -426,7 +426,7 @@ impl MoeLayer { // EP all-reduce: sum partial outputs for 2 tokens if let Some(comm) = ctx.comm - && comm.world_size() > 1 + && ctx.config.ep_world_size > 1 { if ctx.graph_capture { comm.all_reduce(output.0, 2 * h as usize * 2)?; diff --git a/crates/spark-model/src/layers/moe/forward_k3.rs b/crates/spark-model/src/layers/moe/forward_k3.rs index b0cf5ebc2..9350e1ed5 100644 --- a/crates/spark-model/src/layers/moe/forward_k3.rs +++ b/crates/spark-model/src/layers/moe/forward_k3.rs @@ -103,7 +103,7 @@ impl MoeLayer { let shared_down_out = ctx.buffers.attn_output(); let output = ctx.buffers.moe_output(); - let is_ep = ctx.comm.is_some_and(|c| c.world_size() > 1); + let is_ep = ctx.comm.is_some() && ctx.config.ep_world_size > 1; if let (Some(gp), Some(up), Some(dp), Some(sh)) = ( &self.fp8_gate_weight_ptrs, @@ -299,7 +299,7 @@ impl MoeLayer { // EP all-reduce: sum partial outputs for 3 tokens if let Some(comm) = ctx.comm - && comm.world_size() > 1 + && ctx.config.ep_world_size > 1 { if ctx.graph_capture { comm.all_reduce(output.0, 3 * h as usize * 2)?; diff --git a/crates/spark-model/src/layers/moe/forward_prefill.rs b/crates/spark-model/src/layers/moe/forward_prefill.rs index 8bdd65be8..5e4e073b5 100644 --- a/crates/spark-model/src/layers/moe/forward_prefill.rs +++ b/crates/spark-model/src/layers/moe/forward_prefill.rs @@ -364,7 +364,7 @@ impl MoeLayer { // 8. Blend shared expert: output += sigmoid(dot(input, gate)) * shared // Skip when has_shared == false (no shared expert in this model config). // EP fix: defer shared expert blend until AFTER all-reduce to avoid doubling. - let is_ep_prefill = ctx.comm.is_some_and(|c| c.world_size() > 1); + let is_ep_prefill = ctx.comm.is_some() && ctx.config.ep_world_size > 1; if has_shared && !is_ep_prefill { let shared_down_out = ctx.buffers.attn_output(); if use_overlap { @@ -397,7 +397,7 @@ impl MoeLayer { // EP all-reduce if let Some(comm) = ctx.comm - && comm.world_size() > 1 + && ctx.config.ep_world_size > 1 { let _t0 = if ctx.profile { ctx.gpu.synchronize(stream)?; diff --git a/crates/spark-model/src/layers/moe/forward_prefill_bf16.rs b/crates/spark-model/src/layers/moe/forward_prefill_bf16.rs index 1b47fa263..0baf939c8 100644 --- a/crates/spark-model/src/layers/moe/forward_prefill_bf16.rs +++ b/crates/spark-model/src/layers/moe/forward_prefill_bf16.rs @@ -257,7 +257,7 @@ impl MoeLayer { )?; if let Some(comm) = ctx.comm - && comm.world_size() > 1 + && ctx.config.ep_world_size > 1 { comm.all_reduce_async(output.0, num_tokens * h as usize * 2, stream)?; } diff --git a/crates/spark-model/src/layers/moe/forward_prefill_fp8.rs b/crates/spark-model/src/layers/moe/forward_prefill_fp8.rs index 8ed16baf4..5d87963f6 100644 --- a/crates/spark-model/src/layers/moe/forward_prefill_fp8.rs +++ b/crates/spark-model/src/layers/moe/forward_prefill_fp8.rs @@ -611,7 +611,7 @@ impl MoeLayer { // forward_k3() which already do this in the right order; mirrors // vllm PR #39181. if let Some(comm) = ctx.comm - && comm.world_size() > 1 + && ctx.config.ep_world_size > 1 { comm.all_reduce_async(output.0, num_tokens * h as usize * 2, stream)?; } diff --git a/crates/spark-model/src/layers/moe/forward_token_major.rs b/crates/spark-model/src/layers/moe/forward_token_major.rs index 1f78fa2ad..169ccff7c 100644 --- a/crates/spark-model/src/layers/moe/forward_token_major.rs +++ b/crates/spark-model/src/layers/moe/forward_token_major.rs @@ -153,7 +153,7 @@ impl MoeLayer { stream, )?; - let is_ep = ctx.comm.is_some_and(|c| c.world_size() > 1); + let is_ep = ctx.comm.is_some() && ctx.config.ep_world_size > 1; let shared_for_blend = if is_ep { ctx.gpu .memset_async(expert_gate_out, 0, num_tokens * h as usize * 2, stream)?; @@ -178,7 +178,7 @@ impl MoeLayer { )?; if let Some(comm) = ctx.comm - && comm.world_size() > 1 + && ctx.config.ep_world_size > 1 { if ctx.graph_capture { comm.all_reduce(output.0, num_tokens * h as usize * 2)?; From f1462404097d52dae5f4e488f34bb1360051050c Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Sat, 4 Jul 2026 19:29:43 -0400 Subject: [PATCH 12/15] docs(gdn-tp): fix broken intra-doc link to super::TpAttentionDims --- crates/spark-model/src/tp_shard/gdn.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/spark-model/src/tp_shard/gdn.rs b/crates/spark-model/src/tp_shard/gdn.rs index dc4865c65..496d3e7ff 100644 --- a/crates/spark-model/src/tp_shard/gdn.rs +++ b/crates/spark-model/src/tp_shard/gdn.rs @@ -44,7 +44,7 @@ use super::{BF16_BYTES, TpShardKind, shard_dense_bf16}; /// Pre-TP-shard GDN (linear-attention / SSM) dimensions reconstructed from /// `config`. /// -/// Mirrors [`TpAttentionDims`]: `topology.rs` divides +/// Mirrors [`super::TpAttentionDims`]: `topology.rs` divides /// `linear_num_key_heads` / `linear_num_value_heads` by `tp_world_size` at /// startup, so by the time a loader runs `config` holds **per-rank-local** /// head counts. The `full_*` fields multiply back up to the pre-shard sizes From aff65991aa25f1958a2fc29d45100911e1fcdf5f Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Fri, 3 Jul 2026 11:30:29 -0400 Subject: [PATCH 13/15] ci(gdn-tp): file-size split (tp_shard/gdn.rs) + fmt + clippy/typos green MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Split the GDN HeadParallel slicers out of tp_shard.rs into tp_shard/gdn.rs (704 -> 343 LoC; both under the 500 cap) — drops tp_shard.rs from the file-size allow-list; allow-list linear_attn_arms.rs (548) + trait_decode_batched.rs (503). - clippy: repeat().take() -> std::iter::repeat_n() in tp_shard/tests.rs. - typos: alog_src_off -> a_log_src_off. - cargo fmt. --- crates/spark-model/src/tp_shard/gdn.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/crates/spark-model/src/tp_shard/gdn.rs b/crates/spark-model/src/tp_shard/gdn.rs index 496d3e7ff..400b8029a 100644 --- a/crates/spark-model/src/tp_shard/gdn.rs +++ b/crates/spark-model/src/tp_shard/gdn.rs @@ -44,7 +44,7 @@ use super::{BF16_BYTES, TpShardKind, shard_dense_bf16}; /// Pre-TP-shard GDN (linear-attention / SSM) dimensions reconstructed from /// `config`. /// -/// Mirrors [`super::TpAttentionDims`]: `topology.rs` divides +/// Mirrors [`TpAttentionDims`]: `topology.rs` divides /// `linear_num_key_heads` / `linear_num_value_heads` by `tp_world_size` at /// startup, so by the time a loader runs `config` holds **per-rank-local** /// head counts. The `full_*` fields multiply back up to the pre-shard sizes @@ -124,7 +124,7 @@ impl TpGdnDims { } /// Full-row segment list for the `[Q|K|V]` in-projection. - pub(crate) fn qkv_segments(&self) -> [usize; 3] { + fn qkv_segments(&self) -> [usize; 3] { [ self.full_key_dim(), self.full_key_dim(), @@ -132,7 +132,7 @@ impl TpGdnDims { ] } /// Full-row segment list for the concatenated `[Q|K|V|Z]` in-projection. - pub(crate) fn qkvz_segments(&self) -> [usize; 4] { + fn qkvz_segments(&self) -> [usize; 4] { [ self.full_key_dim(), self.full_key_dim(), @@ -144,10 +144,10 @@ impl TpGdnDims { /// A single device-to-device copy in a segmented-slice plan. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub(crate) struct CopyOp { - pub(crate) src_off: usize, - pub(crate) dst_off: usize, - pub(crate) len: usize, +struct CopyOp { + src_off: usize, + dst_off: usize, + len: usize, } /// Build the copy plan for a SEGMENTED row-slice. @@ -162,7 +162,7 @@ pub(crate) struct CopyOp { /// `tp_size` — the caller has already reconstructed `full_*` as /// `local_* * tp_size`, so this holds by construction, but it is checked to /// fail loudly on a mis-wired config rather than silently corrupt heads. -pub(crate) fn segment_copy_plan( +fn segment_copy_plan( segments: &[usize], row_bytes: usize, tp_rank: usize, From fb2314f83958153166011eb727a7230da48821b7 Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Sat, 4 Jul 2026 17:19:55 -0400 Subject: [PATCH 14/15] feat(loader): native-BF16 dense loader for no-metadata Holo/Ornith checkpoints Load Hcompany Holo-3.1-0.8B/4B/9B (and Ornith dense) at native BF16 instead of the lossy runtime BF16->NVFP4 requant the no-metadata (Bf16Raw) dense path defaulted to. Adds BF16 arms for the FFN, attention (Q/K/V/O) and the SSM in_proj_qkvz/out_proj in qwen35_dense.rs, and gates the batched-decode qkvz path (ssm_batched.rs) so the BF16 build engages the batched dense_gemm instead of dropping to the per-seq loop. The NVFP4 weights stay installed as the spec-decode/batched fallback. Fixes a use-after-free that crashed first-layer prefill with CUDA-700: the FFN overlay installed gate/up/down via dense_auto(), which returns the WeightStore cached BF16 ptr un-copied, but load_dense_ffn Bf16Raw arm had already gpu.free() that buffer while building the NVFP4 fallback. Snapshot fresh D2D copies before the free. Validated on Holo-3.1-4B: coherent across FI prefill on/off and long multi-chunk prefill, zero CUDA-700. With --lm-head-dtype bf16 the agentic multi-turn suite goes 2/3 -> 3/3 vs the NVFP4 head at ~10% single-stream decode cost (27.8 -> 25.1 tok/s); lossless head is worth it for these small models (runtime-NVFP4 default was 72% agentic). --- .../trait_decode_multi_seq/ssm_batched.rs | 11 +- .../src/weight_loader/qwen35_dense.rs | 159 +++++++++++++++++- 2 files changed, 166 insertions(+), 4 deletions(-) diff --git a/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched.rs b/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched.rs index 201985dc8..64f7738b9 100644 --- a/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched.rs +++ b/crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched.rs @@ -34,7 +34,16 @@ impl Qwen3SsmLayer { // FP8 build → batched w8a16 GEMM; NVFP4 build → batched w4a16 GEMV // (batch4/16, M<=16). Either amortizes the QKVZ/out_proj weight read // across the n seqs; otherwise the per-seq loop re-streams it n times. - let qkvz_ok = (self.qkvz_nvfp4.is_none() && self.w8a16_gemm_k.0 != 0) + // qkvz_nvfp4.is_none() covers TWO builds: the FP8 build (qkvz_fp8w + // Some -> batched w8a16 GEMM, needs w8a16_gemm_k) and the pure + // native-BF16 build (both None -> the batched `dense_gemm` fallback + // below, which uses dense_gemm_k, always loaded). Gate each on the + // kernel it actually dispatches so the BF16 build engages the batched + // fast path instead of silently dropping to the per-seq loop. The FP8 + // sub-case (qkvz_fp8w Some) is byte-identical to the old gate. + let qkvz_ok = (self.qkvz_nvfp4.is_none() + && ((self.qkvz_fp8w.is_some() && self.w8a16_gemm_k.0 != 0) + || (self.qkvz_fp8w.is_none() && self.dense_gemm_k.0 != 0))) || (self.qkvz_nvfp4.is_some() && self.w4a16_gemv_batch4_k.0 != 0 && n <= 16); let out_ok = self.out_proj_fp8w.is_some() || self.out_proj_dense.is_some() diff --git a/crates/spark-model/src/weight_loader/qwen35_dense.rs b/crates/spark-model/src/weight_loader/qwen35_dense.rs index f1004e837..413a5a2e7 100644 --- a/crates/spark-model/src/weight_loader/qwen35_dense.rs +++ b/crates/spark-model/src/weight_loader/qwen35_dense.rs @@ -240,6 +240,38 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { // FP8 is enabled we overlay the block-scaled FP8 weights on top; // the hot forward / forward_prefill paths then use FP8, the rare // batched paths fall back to real NVFP4. + // Native-BF16 dense-FFN prefill (Holo/Ornith Bf16Raw): the fast + // tensor-core `dense_gemm_tc` path (dense_ffn::forward_prefill) needs + // LIVE BF16 gate/up/down. `load_dense_ffn`'s Bf16Raw arm runtime- + // quantizes each proj via `quantized_any`, whose Bf16Raw branch + // FREES the store's cached BF16 buffer (`gpu.free(w.ptr)` in + // nvfp4_detect.rs). `dense_auto` returns that SAME (now-freed) store + // ptr for a BF16 tensor, so overlaying it AFTER load_dense_ffn hands + // `set_bf16_weights` freed GPU memory -> dense_gemm_tc CUDA-700 at + // grid=[div_ceil(intermediate_size,64),..] on the first prefill. + // Snapshot fresh D2D copies BEFORE the free (the clone is an + // independent allocation the layer owns for its lifetime). + let ffn_bf16_snapshot = if matches!(variant, Nvfp4Variant::Bf16Raw) { + let inter = if config.intermediate_size > 0 { + config.intermediate_size + } else { + config.moe_intermediate_size + }; + let clone_bf16 = |name: &str, rows: usize, cols: usize| -> Result { + let src = dense_auto(store, &format!("{lp}.mlp.{name}.weight"), gpu)?; + let bytes = rows * cols * 2; // BF16 = 2 bytes/elem + let dst = gpu.alloc(bytes)?; + gpu.copy_d2d(src.weight, dst, bytes)?; + Ok(DenseWeight { weight: dst }) + }; + Some(( + clone_bf16("gate_proj", inter, h)?, + clone_bf16("up_proj", inter, h)?, + clone_bf16("down_proj", h, inter)?, + )) + } else { + None + }; let ffn_weights = load_dense_ffn( store, &lp, gpu, variant, absmax_k, quantize_k, stream, config, )?; @@ -260,6 +292,15 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { // ATLAS_FFN_NVFP4_MMQ: same discipline for the W4A4 FP4-MMQ arm — repack // gate/up to block_nvfp4 + free their `_t` copies (net ~0 footprint). dffn.finalize_nvfp4_mmq_load(gpu, h as u32, config.intermediate_size as u32, stream)?; + // Native-BF16 dense-FFN overlay (Bf16Raw, no-metadata Holo dense): install the + // live BF16 gate/up/down snapshot so forward/forward_prefill's bf16 branch + // (preferred over the NVFP4 fallback) reads valid memory. The NVFP4 weights built + // by load_dense_ffn stay as the spec-decode/batched fallback (never null -> no + // CUDA-700 at concurrency). Snapshot was taken before load_dense_ffn freed the + // store's BF16 buffer (see ffn_bf16_snapshot above). + if let Some((g, u, d)) = ffn_bf16_snapshot { + dffn.set_bf16_weights(g, u, d); + } let ffn = FfnComponent::Dense(dffn); match lt { @@ -267,6 +308,9 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { let p = format!("{lp}.self_attn"); let tp_rank = config.tp_rank; let tp_size = config.tp_world_size.max(1); + // Bf16Raw installs a BF16 dense O-proj after the layer is + // built; other variants leave this None (NVFP4/FP8 dispatch). + let mut o_dense_bf16: Option = None; let (attn, q_nvfp4, k_nvfp4, v_nvfp4) = match variant { Nvfp4Variant::CompressedTensors => { // NVFP4-from-disk path: column-parallel Q/K/V, row-parallel O. @@ -331,9 +375,7 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { }; (attn, Some(q), Some(k), Some(v)) } - Nvfp4Variant::Standard - | Nvfp4Variant::Fp8Dequanted - | Nvfp4Variant::Bf16Raw => { + Nvfp4Variant::Standard | Nvfp4Variant::Fp8Dequanted => { // BF16 → NVFP4 path: shard BF16 then quantize per-rank. let load_bf16_then_nvfp4 = |name: &str, full_n: usize, @@ -427,6 +469,60 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { }; (attn, Some(q_nvfp4), Some(k_nvfp4), Some(v_nvfp4)) } + Nvfp4Variant::Bf16Raw => { + // Native BF16 dense attention: keep Q/K/V/O in BF16 + // and dispatch the dense_gemv/dense_gemm kernels that + // ship in the nvfp4 bundle (common/ dense_*_bf16). No + // runtime BF16 -> NVFP4 quant — that lossily quantized + // these no-metadata Holo dense checkpoints. Mirrors + // qwen35/load_layers.rs:552 (BF16-dequant attention) + // and gemma4/loader_a.rs:300/418. + let load_bf16_dense = + |name: &str, + full_n: usize, + full_k: usize, + kind: TpShardKind| + -> Result { + let src = + dense_auto(store, &format!("{p}.{name}.weight"), gpu)?; + if tp_size == 1 { + return Ok(src); + } + let (sharded_ptr, _local_n, _local_k) = shard_dense_bf16( + src.weight, full_n, full_k, kind, tp_rank, tp_size, gpu, + )?; + if sharded_ptr != src.weight { + gpu.free(src.weight)?; + } + Ok(DenseWeight { + weight: sharded_ptr, + }) + }; + let [q_dense, k_dense, v_dense, o_dense] = + load_qkvo_tp(config, load_bf16_dense)?; + + let (k_scale, v_scale) = load_kv_scales(store, &p, gpu); + + // Keep q/k/v BF16 dense ALIVE (do NOT free): the dense + // forward reads them directly. o_proj stays NULL — the + // set_o_dense_bf16(o_dense) call after layer build + // installs the BF16 O-proj that decode/prefill prefer. + let attn = AttentionWeights { + q_proj: q_dense, + k_proj: k_dense, + v_proj: v_dense, + o_proj: crate::weight_map::QuantizedWeight::null(), + q_norm: dense(store, &format!("{p}.q_norm.weight"))?, + k_norm: dense(store, &format!("{p}.k_norm.weight"))?, + q_norm_full: None, + k_norm_full: None, + k_scale, + v_scale, + }; + o_dense_bf16 = Some(o_dense); + // BF16 native: no NVFP4 q/k/v weights → dense fallback. + (attn, None, None, None) + } }; let mut attn_layer = Qwen3AttentionLayer::new( @@ -459,6 +555,13 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { let ot = op.transpose_for_gemm(gpu, hh, nh * hd)?; attn_layer.set_prefill_weights(Some(qt), Some(kt), Some(vt), Some(ot)); } + // Native-BF16 (Bf16Raw): install the dense O-proj so decode + + // prefill prefer it over the (NULL) NVFP4 o_proj. Mutually + // exclusive with the transposed-NVFP4 block above (q_nvfp4 is + // None on the Bf16Raw path). + if let Some(o_dense) = o_dense_bf16 { + attn_layer.set_o_dense_bf16(o_dense); + } // Overlay native FP8 q/k/v/o on top of the NVFP4 weights when // enabled (single-GPU FP8 checkpoint). Hot decode/prefill paths // dispatch FP8 (w8a16); any path without an FP8 branch falls back @@ -728,6 +831,56 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { gpu.free(qkv_dense.weight)?; gpu.free(z_dense.weight)?; + let ba_dense = interleave_ba(&in_proj_a, &in_proj_b, nv, nk, h, gpu)?; + + // Native-BF16 SSM arm (no-metadata dense Holo checkpoints: + // Nvfp4Variant::Bf16Raw). The standard nvfp4 bundle ships the + // common/ BF16 dense kernels (dense_gemv_bf16 / dense_gemm_bf16 + // / dense_gemm_bf16_pipelined) that every SSM forward arm's + // dense fallback already dispatches, so keep the concatenated + // qkvz_dense [Q|K|V|Z] and out_proj_dense ALIVE and route + // through those instead of the lossy BF16->NVFP4 runtime + // requant. No NVFP4/FP8 copy is built at all: in_proj_qkvz + + // out_proj_dense feed dense_gemv (per-seq decode, + // ssm_forward.rs:107/412), dense_gemm (batched decode, + // trait_decode_batched.rs:113/350 + ssm_batched.rs:180/279) + // and dense_gemm_bf16_pipelined (prefill, + // trait_prefill_proj.rs:298 + trait_prefill_helper.rs:89). + // ssm.out_proj stays null — every out_proj arm prefers + // out_proj_dense when Some; predequant_for_prefill / + // set_fp8_prefill_only_weights are skipped (NVFP4/FP8 only). + if matches!(variant, Nvfp4Variant::Bf16Raw) { + let ssm = SsmWeights { + in_proj_qkvz: qkvz_dense, + in_proj_ba: ba_dense, + conv1d, + a_log, + dt_bias, + norm, + out_proj: crate::weight_map::QuantizedWeight::null(), + }; + let mut layer = Qwen3SsmLayer::new_sequential( + input_norm, + ssm, + post_attn_norm, + ffn, + None, // qkvz_nvfp4 — BF16 dense fallback used instead + None, // qkvz_nvfp4_t + None, // out_proj_nvfp4_t + config, + gpu, + )?; + // pub field (qwen3_ssm/mod.rs:46); selected by every + // out_proj arm (ssm_forward.rs:412, + // trait_decode_batched.rs:350, ssm_batched.rs:279, + // trait_prefill_helper.rs:89). + layer.out_proj_dense = Some(out_proj_dense); + layers.push(Box::new(layer)); + continue; + } + + let qkvz_size = config.ssm_qkvz_size(); + // GDN ≥FP8 precision policy (2026-07-04). The nvidia // Qwen3.6-27B-NVFP4 checkpoint ships GDN in_proj_qkv / // out_proj as native F8_E4M3 (modelopt sensitivity From eb540d4d8adf4ab4cd05a145a790979d0886805a Mon Sep 17 00:00:00 2001 From: Richard Safier Date: Sat, 4 Jul 2026 21:30:08 -0400 Subject: [PATCH 15/15] feat(gdn-tp): TP-shard the dense-model GDN/SSM heads (dense pure-TP) Wire the existing GDN HeadParallel slicers (tp_shard/gdn.rs, already used by the MoE loader) into the dense loader SSM arm so dense Holo models (0.8B/4B/9B) TP-shard under pure tensor-parallel. Before this the dense SSM arm loaded the GDN weights UNSHARDED on every rank while the forward ran HeadParallel all-reduces -> garbage under tp>1 (tp=1 was fine). Loader-only. config holds per-rank-LOCAL linear head counts (topology.rs divides them), so TpGdnDims rebuilds the FULL pre-shard sizes; load/ concat/interleave run at FULL, then shard_gdn_{qkvz_rows,ba_rows, conv_rows,value_vector,out_proj_row_parallel} cut this rank contiguous head range. norm.weight [vd] is REPLICATED (shared across value heads). The post-out_proj all-reduce already fires in the shared Qwen3SsmLayer forward. Byte-identical at tp_size==1 (gated else pass-through). Validated 2-node pure-TP2 (GB10x2): dense 0.8B + 4B coherent (were garbage before), MoE 35B unaffected ("15 plus 27 equals 42"). Scope: BF16/NVFP4 dense; the native-FP8 GDN SSM path (ATLAS_NO_GDN_FP8) is not yet sharded (follow-up). Stacks on #250 + #254. --- .github/workflows/file-size-cap.yml | 1 + crates/spark-model/src/tp_shard/gdn.rs | 14 ++-- .../src/weight_loader/qwen35_dense.rs | 71 ++++++++++++++++--- 3 files changed, 69 insertions(+), 17 deletions(-) diff --git a/.github/workflows/file-size-cap.yml b/.github/workflows/file-size-cap.yml index 9d1544d0f..e0a0da1c1 100644 --- a/.github/workflows/file-size-cap.yml +++ b/.github/workflows/file-size-cap.yml @@ -152,6 +152,7 @@ jobs: "crates/spark-model/src/layers/qwen3_attention/prefill/cache_skip.rs" "crates/spark-model/src/layers/qwen3_attention/prefill/paged.rs" "crates/spark-model/src/model/block_mgmt.rs" + "crates/spark-model/src/model/impl_a1.rs" "crates/spark-model/src/model/impl_a2.rs" "crates/spark-model/src/model/ssm_pool.rs" "crates/spark-model/src/weight_loader/qwen35/load_layers.rs" diff --git a/crates/spark-model/src/tp_shard/gdn.rs b/crates/spark-model/src/tp_shard/gdn.rs index 400b8029a..dc4865c65 100644 --- a/crates/spark-model/src/tp_shard/gdn.rs +++ b/crates/spark-model/src/tp_shard/gdn.rs @@ -124,7 +124,7 @@ impl TpGdnDims { } /// Full-row segment list for the `[Q|K|V]` in-projection. - fn qkv_segments(&self) -> [usize; 3] { + pub(crate) fn qkv_segments(&self) -> [usize; 3] { [ self.full_key_dim(), self.full_key_dim(), @@ -132,7 +132,7 @@ impl TpGdnDims { ] } /// Full-row segment list for the concatenated `[Q|K|V|Z]` in-projection. - fn qkvz_segments(&self) -> [usize; 4] { + pub(crate) fn qkvz_segments(&self) -> [usize; 4] { [ self.full_key_dim(), self.full_key_dim(), @@ -144,10 +144,10 @@ impl TpGdnDims { /// A single device-to-device copy in a segmented-slice plan. #[derive(Debug, Clone, Copy, PartialEq, Eq)] -struct CopyOp { - src_off: usize, - dst_off: usize, - len: usize, +pub(crate) struct CopyOp { + pub(crate) src_off: usize, + pub(crate) dst_off: usize, + pub(crate) len: usize, } /// Build the copy plan for a SEGMENTED row-slice. @@ -162,7 +162,7 @@ struct CopyOp { /// `tp_size` — the caller has already reconstructed `full_*` as /// `local_* * tp_size`, so this holds by construction, but it is checked to /// fail loudly on a mis-wired config rather than silently corrupt heads. -fn segment_copy_plan( +pub(crate) fn segment_copy_plan( segments: &[usize], row_bytes: usize, tp_rank: usize, diff --git a/crates/spark-model/src/weight_loader/qwen35_dense.rs b/crates/spark-model/src/weight_loader/qwen35_dense.rs index 413a5a2e7..f372bbd7f 100644 --- a/crates/spark-model/src/weight_loader/qwen35_dense.rs +++ b/crates/spark-model/src/weight_loader/qwen35_dense.rs @@ -9,7 +9,11 @@ use spark_runtime::weights::{WeightDtype, WeightStore}; use super::{ModelWeightLoader, WeightFormat}; use crate::layer::TransformerLayer; use crate::layers::{DenseFfnLayer, FfnComponent, Qwen3AttentionLayer, Qwen3SsmLayer}; -use crate::tp_shard::{TpShardKind, load_qkvo_tp, shard_dense_bf16, shard_quantized_nvfp4}; +use crate::tp_shard::{ + TpGdnDims, TpShardKind, load_qkvo_tp, shard_dense_bf16, shard_gdn_ba_rows, shard_gdn_conv_rows, + shard_gdn_out_proj_row_parallel, shard_gdn_qkvz_rows, shard_gdn_value_vector, + shard_quantized_nvfp4, +}; use crate::weight_map::{ AttentionWeights, DenseWeight, Fp8Weight, MtpWeights, Nvfp4Variant, SsmWeights, dense, dense_auto, dense_f32_safe, dense_keep_f32, dequant_nvfp4_to_bf16, detect_nvfp4_variant, @@ -595,8 +599,15 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { LayerType::LinearAttention => { let nv = config.linear_num_value_heads; let nk = config.linear_num_key_heads; - let qkv_rows = config.ssm_qkv_size(); - let z_rows = config.ssm_z_size(); + // GDN HeadParallel: config holds per-rank-LOCAL linear head counts + // (topology.rs divided them by tp_size). TpGdnDims rebuilds the FULL + // pre-shard sizes so load/concat/interleave run at FULL, then the + // shard_gdn_* slicers cut this rank's contiguous head range. value_dim + // stays LOCAL (sizes the downstream quantize of the sharded buffers). + let tp_size = config.tp_world_size.max(1); + let dims = TpGdnDims::from_config(config); + let qkv_rows = dims.full_conv_dim(); + let z_rows = dims.full_value_dim(); let value_dim = nv * config.linear_value_head_dim; let la = format!("{lp}.linear_attn"); @@ -704,11 +715,6 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { continue; } - // A, B, conv1d, A_log, dt_bias, norm are independent of the - // qkv/z/out_proj on-disk format below — load them once up - // front so both the native-NVFP4 fast path and the legacy - // dequant/requant path can share them. - // // A_log and dt_bias MUST be FP32 — consumer kernels in // `ssm_preprocess.cu` and `mamba2_ssm_decode.cu` declare // them `const float*`. Loading via `dense()` kept BF16 @@ -822,7 +828,8 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { let qkv_dense = load_ssm_proj(&format!("{la}.in_proj_qkv"), qkv_rows, h)?; let z_dense = load_ssm_proj(&format!("{la}.in_proj_z"), z_rows, h)?; - let out_proj_dense = load_ssm_proj(&format!("{la}.out_proj"), h, value_dim)?; + let out_proj_dense = + load_ssm_proj(&format!("{la}.out_proj"), h, dims.full_value_dim())?; let qkvz_dense = gpu_concat_rows(&qkv_dense, qkv_rows, &z_dense, z_rows, h, gpu)?; @@ -831,7 +838,51 @@ impl ModelWeightLoader for Qwen35DenseWeightLoader { gpu.free(qkv_dense.weight)?; gpu.free(z_dense.weight)?; - let ba_dense = interleave_ba(&in_proj_a, &in_proj_b, nv, nk, h, gpu)?; + let ba_dense = + interleave_ba(&in_proj_a, &in_proj_b, dims.full_nv, dims.full_nk, h, gpu)?; + + // GDN HeadParallel shard: cut the FULL concat/interleave/on-disk buffers + // to this rank's contiguous head range (mirrors the MoE loader, + // linear_attn_arms.rs). Runs BEFORE both consumers below (the Bf16Raw + // branch AND the NVFP4/FP8 path), so both get sharded weights. qkvz + // (segmented [Q|K|V|Z]) + conv (segmented [Q|K|V]) slice per-block; ba + // slices contiguously; a_log/dt_bias are per-value-head FP32 scalars; + // out_proj is row-parallel on value_dim (partials summed by the + // post-out_proj all-reduce already in Qwen3SsmLayer::forward). norm is + // [vd] shared across value heads -> REPLICATE (never sliced). Free only + // the fresh qkvz/ba concat buffers; conv/a_log/dt_bias/out_proj sources + // are store aliases or dequant bufs freed downstream (the NVFP4 path + // frees the sharded qkvz/out_proj later). At tp==1 the else arm is a + // pure pass-through and every slicer no-ops -> byte-identical. + let (qkvz_dense, ba_dense, conv1d, a_log, dt_bias, out_proj_dense) = if tp_size + > 1 + { + let d_conv = config.linear_conv_kernel_dim; + let (qkvz_ptr, _, _) = shard_gdn_qkvz_rows(qkvz_dense.weight, &dims, gpu)?; + gpu.free(qkvz_dense.weight)?; + let (ba_ptr, _, _) = shard_gdn_ba_rows(ba_dense.weight, &dims, gpu)?; + gpu.free(ba_dense.weight)?; + let (conv_ptr, _, _) = + shard_gdn_conv_rows(conv1d.weight, &dims, d_conv, gpu)?; + let (a_log_ptr, _) = + shard_gdn_value_vector(a_log.weight, &dims, 1, 4, gpu)?; + let (dt_bias_ptr, _) = + shard_gdn_value_vector(dt_bias.weight, &dims, 1, 4, gpu)?; + let (out_ptr, _, _) = + shard_gdn_out_proj_row_parallel(out_proj_dense.weight, &dims, gpu)?; + ( + DenseWeight { weight: qkvz_ptr }, + DenseWeight { weight: ba_ptr }, + DenseWeight { weight: conv_ptr }, + DenseWeight { weight: a_log_ptr }, + DenseWeight { + weight: dt_bias_ptr, + }, + DenseWeight { weight: out_ptr }, + ) + } else { + (qkvz_dense, ba_dense, conv1d, a_log, dt_bias, out_proj_dense) + }; // Native-BF16 SSM arm (no-metadata dense Holo checkpoints: // Nvfp4Variant::Bf16Raw). The standard nvfp4 bundle ships the