spark_model/layers/glm5next_kda/
binding.rs

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

spark_model/layers/glm5next_kda/
binding.rs
