metrale_model_arch/glm5next_kda/tp.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: GLM-5.3-Flash KDA tensor-parallel shard plan: head-parallel, one all-reduce after
4//! `o_proj`.
5//!
6//! Owner: model-arch (GLM-5.3-Flash KDA).
7//! Invariants:
8//! - [`KdaTpPlan::new`] refuses `tp_rank >= tp_size`, a head count not divisible by `tp_size`,
9//! and a local q|k width that is not a multiple of 256.
10//! - Sharded tensors split into equal, contiguous per-rank ranges ordered by rank.
11//!
12//! The plan is pure data, so the per-rank row arithmetic is tested without a GPU; `tp_bind`
13//! applies it. Each rank owns a contiguous head range, `o_proj` is row-parallel, and one
14//! all-reduce follows it, as in the GDN head-parallel helpers (`crate::tp_shard::gdn`). Unlike
15//! GDN, KDA's `q/k/v_proj` and `q/k/v_conv1d` are separate tensors on disk, so each is sliced on
16//! its own and no segmented copy is needed.
17//!
18//! * `A_log` is per head and `dt_bias` per channel, so they shard in different units.
19//! * `o_norm` is `[head_dim]`, shared by every head, so it is replicated.
20//! * `f_a`/`g_a` are down-projections `[rank, hidden]` and are replicated; only the `_b`
21//! up-projections carry head structure.
22//! * `o_proj` is `[hidden, heads * head_dim]` sliced on its input dimension; each rank's output is
23//! partial until the all-reduce.
24
25use anyhow::{Result, bail};
26use metrale_config::ModelConfig;
27
28/// 2026-09-25: How one KDA tensor maps onto TP ranks.
29#[derive(Debug, Clone, Copy, PartialEq, Eq)]
30pub enum KdaShard {
31 /// 2026-09-25: Every rank holds the whole tensor.
32 Replicated,
33 /// 2026-09-25: Leading dim is `heads`: slice by head range.
34 HeadRows,
35 /// 2026-09-25: Leading dim is `heads * head_dim`: slice by channel range.
36 ChannelRows,
37 /// 2026-09-25: Trailing (input) dim is `heads * head_dim`: a row-parallel GEMM, sliced on the
38 /// input dim, whose output is all-reduced.
39 ChannelCols,
40}
41
42/// 2026-09-25: One tensor's placement, in `[rows, row_elems]` terms. For
43/// [`KdaShard::ChannelCols`] the sharded axis is `row_elems`.
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub struct KdaTensorPlan {
46 pub name: &'static str,
47 pub kind: KdaShard,
48 pub elem_bytes: usize,
49 /// 2026-09-25: Rows of the full, on-disk tensor.
50 pub full_rows: usize,
51 /// 2026-09-25: Elements per row of the full tensor.
52 pub full_row_elems: usize,
53 /// 2026-09-25: Rows this rank keeps.
54 pub local_rows: usize,
55 /// 2026-09-25: Elements per row this rank keeps (differs from full only for `ChannelCols`).
56 pub local_row_elems: usize,
57 /// 2026-09-25: Offset, in rows, of this rank's slice. 0 for `Replicated` and `ChannelCols`.
58 pub src_row_offset: usize,
59 /// 2026-09-25: Offset, in elements, of this rank's column slice. 0 except for `ChannelCols`.
60 pub src_col_offset: usize,
61}
62
63impl KdaTensorPlan {
64 /// 2026-09-25: Bytes this rank stores for this tensor.
65 pub fn local_bytes(&self) -> usize {
66 self.local_rows * self.local_row_elems * self.elem_bytes
67 }
68 /// 2026-09-25: Bytes the full tensor occupies on disk.
69 pub fn full_bytes(&self) -> usize {
70 self.full_rows * self.full_row_elems * self.elem_bytes
71 }
72}
73
74const BF16: usize = 2;
75const F32: usize = 4;
76
77/// 2026-09-25: The per-rank shard plan for one KDA block.
78#[derive(Debug, Clone)]
79pub struct KdaTpPlan {
80 pub tp_rank: usize,
81 pub tp_size: usize,
82 pub hidden: usize,
83 pub head_dim: usize,
84 /// 2026-09-25: Pre-shard head count (all ranks combined).
85 pub full_heads: usize,
86 /// 2026-09-25: Heads this rank owns.
87 pub local_heads: usize,
88 pub conv_kernel: usize,
89 /// 2026-09-25: Low-rank width of the `f`/`g` gate bottleneck.
90 pub gate_rank: usize,
91 pub tensors: Vec<KdaTensorPlan>,
92}
93
94impl KdaTpPlan {
95 /// 2026-09-25: Build from a `ModelConfig` whose linear-head counts are already per rank
96 /// (`serve_phases::topology` divides them), so the full count is `local * tp_size`, as in
97 /// `TpGdnDims::from_config`.
98 ///
99 /// `gate_rank` is not a config key; the loader reads it as the `f_a_proj` row count.
100 pub fn from_config(config: &ModelConfig, gate_rank: usize) -> Result<Self> {
101 let tp_size = config.tp_world_size.max(1);
102 Self::new(
103 config.tp_rank,
104 tp_size,
105 config.hidden_size,
106 config.linear_key_head_dim,
107 config.linear_num_key_heads * tp_size,
108 config.linear_conv_kernel_dim,
109 gate_rank,
110 )
111 }
112
113 #[allow(clippy::too_many_arguments)]
114 pub fn new(
115 tp_rank: usize,
116 tp_size: usize,
117 hidden: usize,
118 head_dim: usize,
119 full_heads: usize,
120 conv_kernel: usize,
121 gate_rank: usize,
122 ) -> Result<Self> {
123 if tp_rank >= tp_size {
124 bail!("tp_rank {tp_rank} >= tp_size {tp_size}");
125 }
126 if tp_size == 0 || head_dim == 0 || full_heads == 0 {
127 bail!("degenerate KDA TP geometry: heads={full_heads} head_dim={head_dim}");
128 }
129 if !full_heads.is_multiple_of(tp_size) {
130 bail!("KDA TP requires heads ({full_heads}) divisible by tp_size ({tp_size})");
131 }
132 let local_heads = full_heads / tp_size;
133
134 // 2026-09-25: `Glm5NextKdaConfig::validate`'s 256-channel rule, re-checked on the local
135 // q|k width.
136 let local_qk_channels = 2 * local_heads * head_dim;
137 if !local_qk_channels.is_multiple_of(256) {
138 bail!(
139 "KDA TP: local qk_channels ({local_qk_channels}) must be a multiple of 256 \
140 (heads={local_heads}, head_dim={head_dim}); causal_conv1d_update_l2norm \
141 hardcodes 2 heads per 256-thread block"
142 );
143 }
144
145 let full_ch = full_heads * head_dim;
146 let local_ch = local_heads * head_dim;
147 let ch_off = tp_rank * local_ch;
148 let head_off = tp_rank * local_heads;
149
150 let rows = |name, kind, elem_bytes, full_rows, full_row_elems| {
151 let (local_rows, local_row_elems, src_row_offset, src_col_offset) = match kind {
152 KdaShard::Replicated => (full_rows, full_row_elems, 0, 0),
153 KdaShard::HeadRows => (full_rows / tp_size, full_row_elems, head_off, 0),
154 KdaShard::ChannelRows => (full_rows / tp_size, full_row_elems, ch_off, 0),
155 // 2026-09-25: Row-parallel: slice the input (column) dim, keep every row.
156 KdaShard::ChannelCols => (full_rows, full_row_elems / tp_size, 0, ch_off),
157 };
158 KdaTensorPlan {
159 name,
160 kind,
161 elem_bytes,
162 full_rows,
163 full_row_elems,
164 local_rows,
165 local_row_elems,
166 src_row_offset,
167 src_col_offset,
168 }
169 };
170
171 let tensors = vec![
172 // 2026-09-25: `[heads * head_dim, hidden]`, column-parallel by output channel.
173 rows("q_proj", KdaShard::ChannelRows, BF16, full_ch, hidden),
174 rows("k_proj", KdaShard::ChannelRows, BF16, full_ch, hidden),
175 rows("v_proj", KdaShard::ChannelRows, BF16, full_ch, hidden),
176 // 2026-09-25: `[heads * head_dim, conv_kernel]` each, separate on disk.
177 rows(
178 "q_conv1d",
179 KdaShard::ChannelRows,
180 BF16,
181 full_ch,
182 conv_kernel,
183 ),
184 rows(
185 "k_conv1d",
186 KdaShard::ChannelRows,
187 BF16,
188 full_ch,
189 conv_kernel,
190 ),
191 rows(
192 "v_conv1d",
193 KdaShard::ChannelRows,
194 BF16,
195 full_ch,
196 conv_kernel,
197 ),
198 // 2026-09-25: Low-rank gates: `_a` down-projects (replicated), `_b` up-projects into
199 // channel space (sharded).
200 rows("f_a_proj", KdaShard::Replicated, BF16, gate_rank, hidden),
201 rows("f_b_proj", KdaShard::ChannelRows, BF16, full_ch, gate_rank),
202 rows("g_a_proj", KdaShard::Replicated, BF16, gate_rank, hidden),
203 rows("g_b_proj", KdaShard::ChannelRows, BF16, full_ch, gate_rank),
204 // 2026-09-25: beta, one row per head.
205 rows("b_proj", KdaShard::HeadRows, BF16, full_heads, hidden),
206 rows("A_log", KdaShard::HeadRows, F32, full_heads, 1),
207 rows("dt_bias", KdaShard::ChannelRows, F32, full_ch, 1),
208 // 2026-09-25: `[head_dim]`, shared by every head, so replicated.
209 rows("o_norm", KdaShard::Replicated, BF16, head_dim, 1),
210 // 2026-09-25: `[hidden, heads * head_dim]`, row-parallel; all-reduce after.
211 rows("o_proj", KdaShard::ChannelCols, BF16, hidden, full_ch),
212 ];
213
214 Ok(Self {
215 tp_rank,
216 tp_size,
217 hidden,
218 head_dim,
219 full_heads,
220 local_heads,
221 conv_kernel,
222 gate_rank,
223 tensors,
224 })
225 }
226
227 pub fn get(&self, name: &str) -> Option<&KdaTensorPlan> {
228 self.tensors.iter().find(|t| t.name == name)
229 }
230
231 /// 2026-09-25: Total bytes this rank stores for one KDA block.
232 pub fn local_bytes(&self) -> usize {
233 self.tensors.iter().map(|t| t.local_bytes()).sum()
234 }
235
236 /// 2026-09-25: Total bytes one KDA block occupies on disk.
237 pub fn full_bytes(&self) -> usize {
238 self.tensors.iter().map(|t| t.full_bytes()).sum()
239 }
240
241 /// 2026-09-25: Whether the layer must all-reduce after `o_proj`. False at `tp_size == 1`, where
242 /// the row-parallel slice is the whole tensor.
243 pub fn needs_output_all_reduce(&self) -> bool {
244 self.tp_size > 1
245 }
246}
247
248#[cfg(test)]
249mod tests;