spark_model/layers/glm5next_dsa/tp.rs
1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! GLM-5.3 DSA tensor-parallel shard plan — **MLA heads sharded, indexer replicated**.
4//!
5//! Shapes measured from `LibertAIDAI/GLM-5.3-Flash-NVFP4@9e0d74e3` layer 3.
6//! Same pure-plan approach as [`crate::layers::glm5next_kda::tp`]: the GPU copy is
7//! `tp_shard` reuse, but a wrong head range yields a running model with mixed-up
8//! heads, so the row arithmetic is data and gets proven without a GPU.
9//!
10//! # 🔴 Why the indexer is REPLICATED, not sharded
11//!
12//! `index_n_heads = 32` divides cleanly at TP=2, so sharding looks free. It is not.
13//!
14//! The indexer emits a **token selection**, not a partial sum. `index_scores`
15//! accumulates `weights[h] * relu(scale · dot)` as a plain sum over heads, so a
16//! head-sharded indexer gives each rank only a *partial* score. If each rank then
17//! takes its own top-k, **the two ranks attend to different tokens** — no crash, no
18//! shape error, just a wrong answer. Sharding is therefore only correct with an
19//! all-reduce of the score tensor *before* the top-k.
20//!
21//! That reduce is the problem. Scores are `[q_rows, n_pools]` with
22//! `n_pools = seq / kpool`, so at a 262 144-token context one decode token costs
23//! `65 536 × 4 B = 256 KB` **per layer** — about 2.8 MB/token across the 11 DSA
24//! layers, on a critical path the DS4F performance review measured as
25//! **latency-bound** (86 × 8 KB collectives/token). The weight saved by sharding the
26//! indexer is a small fraction of one 249.8 MB layer, once. The trade loses.
27//!
28//! A second hazard argues the same way: the reference pins a deterministic tiebreak
29//! (higher score, then **smaller** pool index) precisely because `torch.topk`'s tie
30//! order is undefined. An all-reduced score changes the summation order and can flip
31//! a tie that straddles the cutoff.
32//!
33//! So: indexer replicated, MLA heads sharded, `o_proj` row-parallel with the single
34//! all-reduce `qwen3_attention` already performs.
35//!
36//! # 🪤 Traps this module encodes
37//!
38//! * **`kv_a_proj_with_mqa` is REPLICATED.** It produces the shared latent KV that
39//! every head decompresses from — it is the MQA part of MLA and has no head axis.
40//! Sharding it starves each rank of half the latent.
41//! * **`q_a_proj` / `kv_a_layernorm` / `q_a_layernorm` are replicated** — low-rank
42//! down-projections and their norms, no head structure. Same failure mode as
43//! KDA's `f_a`/`g_a`.
44//! * **`q_b_proj` and `kv_b_proj` shard by head**, at `qk_head_dim` and
45//! `nope + v_dim` per head respectively. The two strides differ; using one for the
46//! other silently mixes heads.
47//! * **`o_proj` is row-parallel** on `heads * v_head_dim`. Column-slicing gives a
48//! plausible, wrong output.
49
50use anyhow::{Result, bail};
51
52use super::Glm5NextDsaConfig;
53
54/// How one DSA tensor maps onto TP ranks.
55#[derive(Debug, Clone, Copy, PartialEq, Eq)]
56pub enum DsaShard {
57 /// Every rank holds the whole tensor: the indexer, the latent KV projection,
58 /// and every low-rank down-projection / norm.
59 Replicated,
60 /// Leading dim is `heads * per_head` — slice by this rank's head range.
61 HeadRows,
62 /// Trailing (input) dim is `heads * v_head_dim` — row-parallel GEMM, slice the
63 /// input dim, then **all-reduce**.
64 HeadCols,
65}
66
67#[derive(Debug, Clone, PartialEq, Eq)]
68pub struct DsaTensorPlan {
69 pub name: &'static str,
70 pub kind: DsaShard,
71 pub elem_bytes: usize,
72 pub full_rows: usize,
73 pub full_row_elems: usize,
74 pub local_rows: usize,
75 pub local_row_elems: usize,
76 pub src_row_offset: usize,
77 pub src_col_offset: usize,
78}
79
80impl DsaTensorPlan {
81 pub fn local_bytes(&self) -> usize {
82 self.local_rows * self.local_row_elems * self.elem_bytes
83 }
84 pub fn full_bytes(&self) -> usize {
85 self.full_rows * self.full_row_elems * self.elem_bytes
86 }
87}
88
89const BF16: usize = 2;
90
91/// Per-rank shard plan for one DSA block.
92#[derive(Debug, Clone)]
93pub struct DsaTpPlan {
94 pub tp_rank: usize,
95 pub tp_size: usize,
96 pub full_heads: usize,
97 pub local_heads: usize,
98 pub tensors: Vec<DsaTensorPlan>,
99}
100
101impl DsaTpPlan {
102 /// `cfg.local_heads` is already per-rank (topology divides before loaders run),
103 /// so the full count is reconstructed as `local * tp_size` — the same convention
104 /// `TpGdnDims::from_config` uses.
105 pub fn new(tp_rank: usize, tp_size: usize, cfg: &Glm5NextDsaConfig) -> Result<Self> {
106 if tp_rank >= tp_size {
107 bail!("tp_rank {tp_rank} >= tp_size {tp_size}");
108 }
109 cfg.validate()?;
110 let local_heads = cfg.local_heads;
111 let full_heads = local_heads * tp_size;
112
113 let h = cfg.hidden;
114 let qk = cfg.qk_head_dim();
115 let kvb_per_head = cfg.qk_nope_head_dim + cfg.v_head_dim;
116 let ihd = cfg.index_head_dim;
117
118 let mk = |name, kind, full_rows: usize, full_row_elems: usize| {
119 let (local_rows, local_row_elems, src_row_offset, src_col_offset) = match kind {
120 DsaShard::Replicated => (full_rows, full_row_elems, 0, 0),
121 DsaShard::HeadRows => {
122 let per = full_rows / tp_size;
123 (per, full_row_elems, tp_rank * per, 0)
124 }
125 DsaShard::HeadCols => {
126 let per = full_row_elems / tp_size;
127 (full_rows, per, 0, tp_rank * per)
128 }
129 };
130 DsaTensorPlan {
131 name,
132 kind,
133 elem_bytes: BF16,
134 full_rows,
135 full_row_elems,
136 local_rows,
137 local_row_elems,
138 src_row_offset,
139 src_col_offset,
140 }
141 };
142
143 let tensors = vec![
144 // ── MLA ──
145 mk("q_a_proj", DsaShard::Replicated, cfg.q_lora_rank, h),
146 mk("q_a_layernorm", DsaShard::Replicated, cfg.q_lora_rank, 1),
147 mk(
148 "q_b_proj",
149 DsaShard::HeadRows,
150 full_heads * qk,
151 cfg.q_lora_rank,
152 ),
153 // 🪤 the shared MQA latent — no head axis, must replicate.
154 mk(
155 "kv_a_proj_with_mqa",
156 DsaShard::Replicated,
157 cfg.kv_cache_dim(),
158 h,
159 ),
160 mk("kv_a_layernorm", DsaShard::Replicated, cfg.kv_lora_rank, 1),
161 mk(
162 "kv_b_proj",
163 DsaShard::HeadRows,
164 full_heads * kvb_per_head,
165 cfg.kv_lora_rank,
166 ),
167 mk("o_proj", DsaShard::HeadCols, h, full_heads * cfg.v_head_dim),
168 // ── indexer: replicated, see the module docs ──
169 mk(
170 "indexer.wq_b",
171 DsaShard::Replicated,
172 cfg.index_heads * ihd,
173 cfg.q_lora_rank,
174 ),
175 mk("indexer.wk", DsaShard::Replicated, ihd, h),
176 mk("indexer.k_norm.weight", DsaShard::Replicated, ihd, 1),
177 // 🪤 LayerNorm, not RMSNorm — the bias is real and is the tell.
178 mk("indexer.k_norm.bias", DsaShard::Replicated, ihd, 1),
179 mk(
180 "indexer.weights_proj",
181 DsaShard::Replicated,
182 cfg.index_heads,
183 h,
184 ),
185 mk(
186 "indexer.index_kpool_compress_gate",
187 DsaShard::Replicated,
188 ihd,
189 h,
190 ),
191 mk(
192 "indexer.index_kpool_compress_ape",
193 DsaShard::Replicated,
194 cfg.index_kpool,
195 ihd,
196 ),
197 ];
198
199 Ok(Self {
200 tp_rank,
201 tp_size,
202 full_heads,
203 local_heads,
204 tensors,
205 })
206 }
207
208 pub fn get(&self, name: &str) -> Option<&DsaTensorPlan> {
209 self.tensors.iter().find(|t| t.name == name)
210 }
211 pub fn local_bytes(&self) -> usize {
212 self.tensors.iter().map(|t| t.local_bytes()).sum()
213 }
214 pub fn full_bytes(&self) -> usize {
215 self.tensors.iter().map(|t| t.full_bytes()).sum()
216 }
217 /// `o_proj` is row-parallel, so its partial output needs the reduce.
218 pub fn needs_output_all_reduce(&self) -> bool {
219 self.tp_size > 1
220 }
221 /// Bytes replicated on every rank — the part TP cannot remove.
222 pub fn replicated_bytes(&self) -> usize {
223 self.tensors
224 .iter()
225 .filter(|t| t.kind == DsaShard::Replicated)
226 .map(|t| t.full_bytes())
227 .sum()
228 }
229}
230
231#[cfg(test)]
232mod tests;