feat(gdn-tp): dense-model GDN/SSM TP-shard (dense pure-TP) [stacks on #254+#250]

.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"
@@ -246,6 +247,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

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)?;

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)?;

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)?;

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)?;

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)?;

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)?;

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)?;
         }

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)?;
         }

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)?;

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)
     }
 }

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();

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()
@@ -307,6 +316,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,

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

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,

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,

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(

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) => {