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

[rsafier]

Re-opened replacement for #259. #259 was merged into its base feature branch by mistake on 2026-07-09 and reverted; the base branch has since been restored and restacked onto main. GitHub cannot reopen a merged PR, so this is a fresh PR carrying the identical commits (verified patch-identical via git range-diff). Marked ready for review — the original's "draft" wording was stale.

Stacks on #254 + #250 (both unmerged). Base is #254 (feat/gdn-headparallel-tp-on-main); this branch also carries #250 (8ce0a600, native-BF16 dense loader) as its first commit because the dense-TP shard sits on #250s restructured dense SSM arm. The net-new commit is the tip (feat(gdn-tp): TP-shard the dense-model GDN/SSM heads). Once #250 + #254 land on main, rebase collapses this to just that commit.

What

Makes dense Holo models (0.8B/4B/9B) TP-shard under pure tensor-parallel. #254 gave the SSM-hybrid MoE models GDN HeadParallel; the dense loader was never wired for it, so dense TP2 produced garbage (TP1 fine). This wires the same shard_gdn_* slicers (tp_shard/gdn.rs) into the dense loaders SSM arm.

Loader-only. config holds per-rank-LOCAL linear head counts (topology.rs divides them), so TpGdnDims::from_config rebuilds the FULL pre-shard sizes; load/concat/interleave run at FULL, then shard_gdn_qkvz_rows / _ba_rows / _conv_rows / _value_vector (a_log, dt_bias) / _out_proj_row_parallel cut this ranks 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. Gated tp_size>1; byte-identical at tp==1 (else-arm pass-through, every slicer no-ops).

Validation — 2-node pure-TP2 (GB10 × 2)

Model	Before	After
dense 0.8B	garbage (multilingual)	coherent — "red, blue, and yellow"
dense 4B	garbage	coherent — "Paris"
MoE 35B (regression)	—	unaffected — "15 plus 27 equals 42"

Rebuilt the branch binary (not the wip prototype) and re-confirmed 0.8B TP2 coherent on this exact base. Compiles + fmt clean.

Scope

BF16/NVFP4 dense only. The native-FP8 GDN SSM path (ATLAS_NO_GDN_FP8, gated on proj_is_fp8_any_scale) is not yet sharded — dense-FP8-TP is a follow-up (the BF16 dense models dont take that path). Mirrors #254 being MoE-TP-scoped.

[SeedSource]

[BETA] Automated PR review — SeedSource AI pipeline

Advisory only, not a merge gate. AI-generated (internal risk-triage, Atlas serve gate, GLM-5.2 review), attributed per Atlas convention.

Triage

• Risk: med - serve-test-required (triage advisory): True - path-heuristic (authority): True
• Rationale: Fixes BF16 snapshot and dispatch logic for FP8/NVFP4 builds to prevent memory issues and fallbacks.
• Serving-path files: crates/spark-model/src/layers/qwen3_ssm/trait_decode_multi_seq/ssm_batched.rs, crates/spark-model/src/weight_loader/qwen35_dense.rs

Serve gate

• Build: Nones - Qwen3.6-35B-A3B-FP8 on GB10 node1
• Top-k KL vs main: None - decode delta None%
• Coherence probe: None (note: our test_coherence.py is flaky here; KL is the reliable serve signal)

GLM-5.2 code review

Correctness concerns:

1. ssm_batched.rs gate logic: The new gate (qkvz_fp8w.is_some() && w8a16_gemm_k.0 != 0) || (qkvz_fp8w.is_none() && dense_gemm_k.0 != 0) — if both qkvz_fp8w and qkvz_nvfp4 are None but dense_gemm_k.0 == 0, qkvz_ok is false and falls to per-seq loop. Verify dense_gemm_k is always loaded in the BF16 build; otherwise this is a silent perf regression.

2. ffn_bf16_snapshot D2D before load_dense_ffn: The comment claims dense_auto returns the store's cached BF16 ptr, and load_dense_ffn's quantized_any frees it. If dense_auto returns a copy (not the cached ptr), the snapshot is unnecessary but harmless; if it returns the same ptr, the snapshot must happen before any free. Ordering looks correct, but confirm clone_bf16 reads happen before load_dense_ffn — they do (snapshot is built before the call). Risk: if dense_auto doesn't cache or the free path differs, the snapshot is dead code masking a real UAF.

3. sharded_ptr != src.weight free: In load_bf16_dense, after shard_dense_bf16, if tp_size == 1 returns src directly (no free). For tp_size > 1, if shard_dense_bf16 returns the same ptr (no-op shard), the guard prevents double-free. OK.

4. o_proj: QuantizedWeight::null(): Bf16Raw sets o_proj null, then set_o_dense_bf16 installs BF16. Verify decode/prefill paths check o_dense_bf16.is_some() before falling back to NVFP4 o_proj — a null NVFP4 o_proj with no BF16 override would crash.

5. TpGdnDims::from_config: Rebuilds full sizes from per-rank-local config. If tp_size doesn't evenly divide head counts, full_conv_dim() = local * tp_size is wrong. Confirm topology.rs guarantees even division.

Quality: Comment density is high but mostly justified for the subtle memory-lifetime issues. The clone_bf16 closure captures lp, inter, h — verify h is the correct column dim for gate/up and row dim for down (looks right: gate/up are [inter, h], down is [h, inter]).

Risk: Medium. The FFN snapshot logic is the riskiest part — a wrong assumption about dense_auto/quantized_any interaction either wastes memory or reintroduces CUDA-700. Needs a test confirming Bf16Raw + prefill doesn't crash and that snapshot ptrs differ from store ptrs.

_{Pipeline: build+serve on GB10, GLM-5.2 via NVIDIA NIM, internal risk-triage classifier. External diff treated as untrusted input. Public repo.}