metrale_model_arch/glm5next_kda/
binding.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Typed weight binding for the GLM-5.3-Flash KDA attention blocks.
4//!
5//! Owner: model-arch (GLM-5.3-Flash KDA).
6//! Invariants:
7//! - [`bind_kda_weights`] returns weights only when a layer's `self_attn.*` tensors are exactly
8//!   [`KDA_TENSORS`], each with its listed dtype and shape; anything else is an error.
9//! - Bytes are uploaded unchanged: no cast, requantisation or dequantisation.
10//!
11//! The MTP layer is DSA-shaped (it loads through `glm5next_dsa`, see
12//! `weight_loader/glm5_next_mtp.rs`) and is not bindable here; [`classify_attn_block`] tells the
13//! two apart from the tensor names alone.
14//!
15//! The checkpoint stores three rank-3 conv tensors (`[qkv, 1, kernel]`). The binder concatenates
16//! them in q, k, v order and uses the result as `[conv_dim, kernel]`: dropping the singleton
17//! dimension moves no bytes. `A_log` (per head) and `dt_bias` (per channel) are F32; a BF16 copy
18//! of either is refused.
19
20use std::collections::BTreeSet;
21
22use anyhow::{Context, Result, bail};
23use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend};
24
25use super::{Glm5NextKdaConfig, Glm5NextKdaWeights};
26use metrale_model_layers::weight_map::DenseWeight;
27
28/// 2026-09-25: The dtypes the KDA binder accepts.
29#[derive(Clone, Copy, PartialEq, Eq, Debug)]
30pub enum KdaDtype {
31    Bf16,
32    F32,
33}
34
35impl KdaDtype {
36    pub fn parse(s: &str) -> Option<Self> {
37        match s {
38            "BF16" => Some(Self::Bf16),
39            "F32" => Some(Self::F32),
40            _ => None,
41        }
42    }
43    pub fn name(self) -> &'static str {
44        match self {
45            Self::Bf16 => "BF16",
46            Self::F32 => "F32",
47        }
48    }
49}
50
51/// 2026-09-25: One tensor as it sits in the checkpoint: dtype, shape and raw little-endian bytes.
52pub struct RawTensor<'a> {
53    pub dtype: KdaDtype,
54    pub shape: Vec<usize>,
55    pub bytes: &'a [u8],
56}
57
58/// 2026-09-25: A checkpoint slice scoped to one decoder layer. Names are layer-relative
59/// (`self_attn.q_proj.weight`), so the same binder works for any layer index.
60pub trait KdaTensorSource {
61    fn get(&self, name: &str) -> Option<RawTensor<'_>>;
62    /// 2026-09-25: Every layer-relative name present, including non-attention ones.
63    fn names(&self) -> Vec<String>;
64}
65
66/// 2026-09-25: One entry of [`KDA_TENSORS`]: a `self_attn` tensor's name, dtype and shape.
67///
68/// Shapes are expressed against [`Glm5NextKdaConfig`], so a geometry mismatch fails at binding
69/// rather than at launch. `H` = heads, `D` = head_dim, `Q` = H*D, `X` = hidden, `K` = conv
70/// kernel, `One` = 1.
71#[derive(Clone, Copy, Debug)]
72pub struct TensorSpec {
73    pub name: &'static str,
74    pub dtype: KdaDtype,
75    dims: &'static [Dim],
76}
77
78#[derive(Clone, Copy, Debug)]
79enum Dim {
80    Q,
81    X,
82    D,
83    H,
84    K,
85    One,
86}
87
88use Dim::{D as DD, H as DH, K as DK, One as D1, Q as DQ, X as DX};
89
90pub const KDA_TENSORS: &[TensorSpec] = &[
91    TensorSpec {
92        name: "self_attn.q_proj.weight",
93        dtype: KdaDtype::Bf16,
94        dims: &[DQ, DX],
95    },
96    TensorSpec {
97        name: "self_attn.k_proj.weight",
98        dtype: KdaDtype::Bf16,
99        dims: &[DQ, DX],
100    },
101    TensorSpec {
102        name: "self_attn.v_proj.weight",
103        dtype: KdaDtype::Bf16,
104        dims: &[DQ, DX],
105    },
106    TensorSpec {
107        name: "self_attn.q_conv1d.weight",
108        dtype: KdaDtype::Bf16,
109        dims: &[DQ, D1, DK],
110    },
111    TensorSpec {
112        name: "self_attn.k_conv1d.weight",
113        dtype: KdaDtype::Bf16,
114        dims: &[DQ, D1, DK],
115    },
116    TensorSpec {
117        name: "self_attn.v_conv1d.weight",
118        dtype: KdaDtype::Bf16,
119        dims: &[DQ, D1, DK],
120    },
121    TensorSpec {
122        name: "self_attn.f_a_proj.weight",
123        dtype: KdaDtype::Bf16,
124        dims: &[DD, DX],
125    },
126    TensorSpec {
127        name: "self_attn.f_b_proj.weight",
128        dtype: KdaDtype::Bf16,
129        dims: &[DQ, DD],
130    },
131    TensorSpec {
132        name: "self_attn.g_a_proj.weight",
133        dtype: KdaDtype::Bf16,
134        dims: &[DD, DX],
135    },
136    TensorSpec {
137        name: "self_attn.g_b_proj.weight",
138        dtype: KdaDtype::Bf16,
139        dims: &[DQ, DD],
140    },
141    TensorSpec {
142        name: "self_attn.b_proj.weight",
143        dtype: KdaDtype::Bf16,
144        dims: &[DH, DX],
145    },
146    // 2026-09-25: F32 on disk and F32 in `kda_gate_bf16`'s signature; uploaded unconverted.
147    TensorSpec {
148        name: "self_attn.A_log",
149        dtype: KdaDtype::F32,
150        dims: &[DH],
151    },
152    TensorSpec {
153        name: "self_attn.dt_bias",
154        dtype: KdaDtype::F32,
155        dims: &[DQ],
156    },
157    TensorSpec {
158        name: "self_attn.o_norm.weight",
159        dtype: KdaDtype::Bf16,
160        dims: &[DD],
161    },
162    TensorSpec {
163        name: "self_attn.o_proj.weight",
164        dtype: KdaDtype::Bf16,
165        dims: &[DX, DQ],
166    },
167];
168
169/// 2026-09-25: Names that identify a DSA (`deepseek_sparse_attention`) block, the MTP layer
170/// included, so a caller can classify a layer without its index.
171pub const DSA_MARKERS: &[&str] = &[
172    "self_attn.kv_a_proj_with_mqa.weight",
173    "self_attn.indexer.wk.weight",
174];
175
176impl TensorSpec {
177    pub fn expected_shape(&self, c: &Glm5NextKdaConfig) -> Vec<usize> {
178        self.dims
179            .iter()
180            .map(|d| match d {
181                Dim::Q => c.qkv_dim(),
182                Dim::X => c.hidden,
183                Dim::D => c.head_dim,
184                Dim::H => c.heads,
185                Dim::K => c.conv_kernel,
186                Dim::One => 1,
187            })
188            .collect()
189    }
190}
191
192/// 2026-09-25: What kind of attention block a layer's tensor names describe.
193#[derive(Clone, Copy, PartialEq, Eq, Debug)]
194pub enum AttnBlockKind {
195    Kda,
196    /// 2026-09-25: `deepseek_sparse_attention`, and the MTP layer, which is DSA-shaped.
197    Dsa,
198    Unknown,
199}
200
201/// 2026-09-25: Classify from tensor names alone: `Dsa` when every [`DSA_MARKERS`] name is present,
202/// else `Kda` when every [`KDA_TENSORS`] name is present, else `Unknown`.
203pub fn classify_attn_block(names: &[String]) -> AttnBlockKind {
204    let set: BTreeSet<&str> = names.iter().map(String::as_str).collect();
205    if DSA_MARKERS.iter().all(|m| set.contains(m)) {
206        return AttnBlockKind::Dsa;
207    }
208    if KDA_TENSORS.iter().all(|t| set.contains(t.name)) {
209        return AttnBlockKind::Kda;
210    }
211    AttnBlockKind::Unknown
212}
213
214/// 2026-09-25: Per-layer counts from [`bind_kda_weights`]. Tensors outside `self_attn.` are
215/// counted in `non_attn_seen` and never bound.
216#[derive(Clone, Debug, Default)]
217pub struct KdaBindReport {
218    pub layer_idx: usize,
219    pub bound: usize,
220    pub self_attn_seen: usize,
221    pub non_attn_seen: usize,
222    pub unknown_self_attn: Vec<String>,
223    pub bytes: usize,
224}
225
226fn upload(gpu: &dyn GpuBackend, bytes: &[u8]) -> Result<DevicePtr> {
227    let p = gpu.alloc(bytes.len().max(1))?;
228    gpu.copy_h2d(bytes, p)?;
229    Ok(p)
230}
231
232/// 2026-09-25: Bind one KDA block: exact tensor set, dtypes, shapes and byte counts, or an error.
233///
234/// Weights are uploaded unchanged: BF16 stays BF16 and F32 stays F32.
235pub fn bind_kda_weights(
236    gpu: &dyn GpuBackend,
237    cfg: &Glm5NextKdaConfig,
238    layer_idx: usize,
239    src: &dyn KdaTensorSource,
240) -> Result<(Glm5NextKdaWeights, KdaBindReport)> {
241    cfg.validate()?;
242    let names = src.names();
243    let mut rep = KdaBindReport {
244        layer_idx,
245        ..Default::default()
246    };
247
248    let known: BTreeSet<&str> = KDA_TENSORS.iter().map(|t| t.name).collect();
249    for n in &names {
250        if n.starts_with("self_attn.") {
251            rep.self_attn_seen += 1;
252            if !known.contains(n.as_str()) {
253                rep.unknown_self_attn.push(n.clone());
254            }
255        } else {
256            rep.non_attn_seen += 1;
257        }
258    }
259    if !rep.unknown_self_attn.is_empty() {
260        bail!(
261            "layer {layer_idx}: {} unrecognised self_attn tensor(s): {:?} — a KDA block has \
262             exactly {} and this binder refuses to skip anything",
263            rep.unknown_self_attn.len(),
264            rep.unknown_self_attn,
265            KDA_TENSORS.len()
266        );
267    }
268
269    let mut fetch = |spec: &TensorSpec| -> Result<Vec<u8>> {
270        let t = src
271            .get(spec.name)
272            .with_context(|| format!("layer {layer_idx}: missing {}", spec.name))?;
273        if t.dtype != spec.dtype {
274            bail!(
275                "layer {layer_idx}: {} is {} but a KDA block requires {} — casting it would \
276                 change the numerics",
277                spec.name,
278                t.dtype.name(),
279                spec.dtype.name()
280            );
281        }
282        let want = spec.expected_shape(cfg);
283        if t.shape != want {
284            bail!(
285                "layer {layer_idx}: {} has shape {:?}, expected {want:?}",
286                spec.name,
287                t.shape
288            );
289        }
290        let elem = match spec.dtype {
291            KdaDtype::Bf16 => 2,
292            KdaDtype::F32 => 4,
293        };
294        let expect_bytes = want.iter().product::<usize>() * elem;
295        if t.bytes.len() != expect_bytes {
296            bail!(
297                "layer {layer_idx}: {} is {} B, shape {want:?} implies {expect_bytes} B",
298                spec.name,
299                t.bytes.len()
300            );
301        }
302        rep.bound += 1;
303        rep.bytes += t.bytes.len();
304        Ok(t.bytes.to_vec())
305    };
306
307    let by_name = |n: &str| -> &TensorSpec { KDA_TENSORS.iter().find(|t| t.name == n).unwrap() };
308    let mut raw = |n: &str| fetch(by_name(n));
309
310    let q_proj = raw("self_attn.q_proj.weight")?;
311    let k_proj = raw("self_attn.k_proj.weight")?;
312    let v_proj = raw("self_attn.v_proj.weight")?;
313    // 2026-09-25: Concatenate in q, k, v order, the order the projections are packed in. The
314    // shape check above pins rank 3 with a middle dimension of 1, and `[dim, 1, ks]` and
315    // `[dim, ks]` are the same row-major bytes.
316    let mut conv = raw("self_attn.q_conv1d.weight")?;
317    conv.extend_from_slice(&raw("self_attn.k_conv1d.weight")?);
318    conv.extend_from_slice(&raw("self_attn.v_conv1d.weight")?);
319    debug_assert_eq!(conv.len(), cfg.conv_dim() * cfg.conv_kernel * 2);
320    let f_a = raw("self_attn.f_a_proj.weight")?;
321    let f_b = raw("self_attn.f_b_proj.weight")?;
322    let g_a = raw("self_attn.g_a_proj.weight")?;
323    let g_b = raw("self_attn.g_b_proj.weight")?;
324    let b_proj = raw("self_attn.b_proj.weight")?;
325    let a_log = raw("self_attn.A_log")?;
326    let dt_bias = raw("self_attn.dt_bias")?;
327    let o_norm = raw("self_attn.o_norm.weight")?;
328    let o_proj = raw("self_attn.o_proj.weight")?;
329
330    let dw = |b: &[u8]| -> Result<DenseWeight> {
331        Ok(DenseWeight {
332            weight: upload(gpu, b)?,
333        })
334    };
335    let w = Glm5NextKdaWeights {
336        q_proj: dw(&q_proj)?,
337        k_proj: dw(&k_proj)?,
338        v_proj: dw(&v_proj)?,
339        conv: dw(&conv)?,
340        f_a: dw(&f_a)?,
341        f_b: dw(&f_b)?,
342        dt_bias: upload(gpu, &dt_bias)?,
343        a_log: upload(gpu, &a_log)?,
344        b_proj: dw(&b_proj)?,
345        g_a: dw(&g_a)?,
346        g_b: dw(&g_b)?,
347        o_norm: dw(&o_norm)?,
348        o_proj: dw(&o_proj)?,
349    };
350    if rep.bound != KDA_TENSORS.len() {
351        bail!(
352            "layer {layer_idx}: bound {} of {} tensors",
353            rep.bound,
354            KDA_TENSORS.len()
355        );
356    }
357    Ok((w, rep))
358}
359
360#[cfg(test)]
361mod tests {
362    use super::*;
363
364    fn cfg() -> Glm5NextKdaConfig {
365        Glm5NextKdaConfig {
366            hidden: 4096,
367            heads: 64,
368            head_dim: 128,
369            conv_kernel: 4,
370            gate_lower_bound: -5.0,
371            rms_norm_eps: 1e-5,
372            l2_eps: 1e-6,
373            chunk: 32,
374        }
375    }
376
377    /// 2026-09-25: The spec table at the `glm-5.3-flash` geometry (hidden 4096, 64 heads of 128,
378    /// conv kernel 4, as `kernels/gb10/glm-5.3-flash/MODEL.toml` records it).
379    #[test]
380    fn tensor_spec_matches_the_audited_checkpoint_shapes() {
381        let c = cfg();
382        let want: &[(&str, &str, &[usize])] = &[
383            ("self_attn.q_proj.weight", "BF16", &[8192, 4096]),
384            ("self_attn.k_proj.weight", "BF16", &[8192, 4096]),
385            ("self_attn.v_proj.weight", "BF16", &[8192, 4096]),
386            ("self_attn.q_conv1d.weight", "BF16", &[8192, 1, 4]),
387            ("self_attn.k_conv1d.weight", "BF16", &[8192, 1, 4]),
388            ("self_attn.v_conv1d.weight", "BF16", &[8192, 1, 4]),
389            ("self_attn.f_a_proj.weight", "BF16", &[128, 4096]),
390            ("self_attn.f_b_proj.weight", "BF16", &[8192, 128]),
391            ("self_attn.g_a_proj.weight", "BF16", &[128, 4096]),
392            ("self_attn.g_b_proj.weight", "BF16", &[8192, 128]),
393            ("self_attn.b_proj.weight", "BF16", &[64, 4096]),
394            ("self_attn.A_log", "F32", &[64]),
395            ("self_attn.dt_bias", "F32", &[8192]),
396            ("self_attn.o_norm.weight", "BF16", &[128]),
397            ("self_attn.o_proj.weight", "BF16", &[4096, 8192]),
398        ];
399        assert_eq!(
400            KDA_TENSORS.len(),
401            want.len(),
402            "the KDA block has exactly 15 tensors"
403        );
404        for (n, dt, sh) in want {
405            let s = KDA_TENSORS.iter().find(|t| &t.name == n).expect(n);
406            assert_eq!(s.dtype.name(), *dt, "{n} dtype");
407            assert_eq!(s.expected_shape(&c), sh.to_vec(), "{n} shape");
408        }
409    }
410
411    /// 2026-09-25: A DSA name set classifies as `Dsa`, never `Kda`.
412    #[test]
413    fn dsa_and_mtp_blocks_do_not_classify_as_kda() {
414        let dsa: Vec<String> = [
415            "self_attn.kv_a_proj_with_mqa.weight",
416            "self_attn.kv_a_layernorm.weight",
417            "self_attn.kv_b_proj.weight",
418            "self_attn.q_a_proj.weight",
419            "self_attn.q_b_proj.weight",
420            "self_attn.indexer.wk.weight",
421            "self_attn.indexer.wq_b.weight",
422            "self_attn.o_proj.weight",
423        ]
424        .iter()
425        .map(|s| s.to_string())
426        .collect();
427        assert_eq!(classify_attn_block(&dsa), AttnBlockKind::Dsa);
428
429        let kda: Vec<String> = KDA_TENSORS.iter().map(|t| t.name.to_string()).collect();
430        assert_eq!(classify_attn_block(&kda), AttnBlockKind::Kda);
431
432        // 2026-09-25: A KDA block missing one tensor is `Unknown`, not `Kda`.
433        assert_eq!(classify_attn_block(&kda[1..]), AttnBlockKind::Unknown);
434    }
435
436    #[test]
437    fn chunk_width_is_bounded_by_the_shared_memory_ceiling() {
438        let mut c = cfg();
439        assert!(c.validate().is_ok(), "C=32 must fit");
440        assert!(c.smem_scan() <= SMEM_CEILING);
441        c.chunk = 64;
442        assert!(
443            c.validate().is_err(),
444            "C=64 needs 81920 B and must be rejected, not truncated"
445        );
446    }
447
448    use super::super::SMEM_CEILING;
449}
metrale_model_arch/glm5next_kda/binding.rs

metrale_model_arch/glm5next_kda/
binding.rs
