metrale_model_arch/glm5next_dsa/
binding.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The tensor contract of one GLM-5.3 DSA block: the 14 `self_attn` tensors
4//! [`dsa_tensor_specs`] lists, their dtype and unsharded shape, and a verifier.
5//!
6//! Owner: model-arch (GLM-5.3 DSA).
7//! Invariants:
8//! - [`verify_dsa_block`] returns `Ok` only when every spec tensor is present with the
9//!   spec's dtype, shape and BF16 byte length, and the source holds no other `self_attn.*`
10//!   tensor.
11//!
12//! It uses the tensor types of [`crate::glm5next_kda::binding`]. Shapes are computed from
13//! [`Glm5NextDsaConfig`].
14//!
15//! # What the table catches
16//!
17//! * `indexer.k_norm` is a LayerNorm with a bias, and the bias is a spec entry, so a
18//!   checkpoint without it fails the bind.
19//! * `index_kpool_compress_ape` is `[index_kpool, index_head_dim]`, indexed by pool slot.
20//! * `q_b_proj` has `qk_head_dim` rows per head and `kv_b_proj` has
21//!   `qk_nope_head_dim + v_head_dim` (256 and 512 on GLM-5.3). Both are
22//!   `[heads * w, lora]`, so only the shape check tells a swapped width apart.
23//! * NoPE: `kv_a_proj_with_mqa` has `kv_cache_dim()` rows, which is `kv_lora_rank` (512),
24//!   not 512 + a rope section.
25
26use std::collections::BTreeSet;
27
28use anyhow::{Result, bail};
29
30use super::Glm5NextDsaConfig;
31use crate::glm5next_kda::binding::{KdaDtype as Dtype, KdaTensorSource as TensorSource};
32
33/// 2026-09-25: One expected tensor: layer-relative name, dtype, and full (unsharded) shape.
34#[derive(Debug, Clone, PartialEq, Eq)]
35pub struct DsaSpec {
36    pub name: String,
37    pub dtype: Dtype,
38    pub shape: Vec<usize>,
39}
40
41/// 2026-09-25: Every `self_attn` tensor a DSA block has, and the only ones it may have.
42/// `full_heads` is the unsharded head count: shapes are checked against the full tensors,
43/// which `build::build_dsa_weights` shards afterwards.
44pub fn dsa_tensor_specs(cfg: &Glm5NextDsaConfig, full_heads: usize) -> Vec<DsaSpec> {
45    let x = cfg.hidden;
46    let ql = cfg.q_lora_rank;
47    let kvl = cfg.kv_lora_rank;
48    let qk = cfg.qk_head_dim();
49    let kvb = cfg.qk_nope_head_dim + cfg.v_head_dim;
50    let ihd = cfg.index_head_dim;
51    let ih = cfg.index_heads;
52
53    let s = |name: &str, shape: Vec<usize>| DsaSpec {
54        name: name.to_string(),
55        dtype: Dtype::Bf16,
56        shape,
57    };
58
59    vec![
60        s("self_attn.q_a_proj.weight", vec![ql, x]),
61        s("self_attn.q_a_layernorm.weight", vec![ql]),
62        s("self_attn.q_b_proj.weight", vec![full_heads * qk, ql]),
63        s(
64            "self_attn.kv_a_proj_with_mqa.weight",
65            vec![cfg.kv_cache_dim(), x],
66        ),
67        s("self_attn.kv_a_layernorm.weight", vec![kvl]),
68        s("self_attn.kv_b_proj.weight", vec![full_heads * kvb, kvl]),
69        s(
70            "self_attn.o_proj.weight",
71            vec![x, full_heads * cfg.v_head_dim],
72        ),
73        s("self_attn.indexer.wq_b.weight", vec![ih * ihd, ql]),
74        s("self_attn.indexer.wk.weight", vec![ihd, x]),
75        s("self_attn.indexer.k_norm.weight", vec![ihd]),
76        s("self_attn.indexer.k_norm.bias", vec![ihd]),
77        s("self_attn.indexer.weights_proj.weight", vec![ih, x]),
78        s("self_attn.indexer.index_kpool_compress_gate", vec![ihd, x]),
79        s(
80            "self_attn.indexer.index_kpool_compress_ape",
81            vec![cfg.index_kpool, ihd],
82        ),
83    ]
84}
85
86/// 2026-09-25: What a successful bind saw.
87#[derive(Debug, Clone, PartialEq, Eq)]
88pub struct DsaBindReport {
89    pub bound: usize,
90    pub total_bytes: usize,
91    /// 2026-09-25: Always empty: [`verify_dsa_block`] fails when the source has an
92    /// unclaimed `self_attn.*` tensor.
93    pub unclaimed: Vec<String>,
94}
95
96/// 2026-09-25: Verify one DSA block against the spec table.
97///
98/// Validates `cfg`, then fails unless every spec is present with its dtype, shape and
99/// BF16 byte length and no `self_attn.*` tensor is left unclaimed.
100pub fn verify_dsa_block(
101    cfg: &Glm5NextDsaConfig,
102    full_heads: usize,
103    source: &dyn TensorSource,
104) -> Result<DsaBindReport> {
105    cfg.validate()?;
106    let specs = dsa_tensor_specs(cfg, full_heads);
107    let mut total_bytes = 0usize;
108
109    for spec in &specs {
110        let Some(raw) = source.get(&spec.name) else {
111            bail!("DSA bind: missing required tensor `{}`", spec.name);
112        };
113        if raw.dtype != spec.dtype {
114            bail!(
115                "DSA bind: `{}` is {:?}, expected {:?}",
116                spec.name,
117                raw.dtype,
118                spec.dtype
119            );
120        }
121        if raw.shape != spec.shape {
122            bail!(
123                "DSA bind: `{}` has shape {:?}, expected {:?}",
124                spec.name,
125                raw.shape,
126                spec.shape
127            );
128        }
129        let elems: usize = spec.shape.iter().product();
130        if raw.bytes.len() != elems * 2 {
131            bail!(
132                "DSA bind: `{}` carries {} bytes, expected {} for {:?} BF16",
133                spec.name,
134                raw.bytes.len(),
135                elems * 2,
136                spec.shape
137            );
138        }
139        total_bytes += raw.bytes.len();
140    }
141
142    let claimed: BTreeSet<&str> = specs.iter().map(|s| s.name.as_str()).collect();
143    let unclaimed: Vec<String> = source
144        .names()
145        .into_iter()
146        .filter(|n| n.starts_with("self_attn.") && !claimed.contains(n.as_str()))
147        .collect();
148    if !unclaimed.is_empty() {
149        bail!(
150            "DSA bind: {} unclaimed self_attn tensor(s), first: {:?}. An unexpected \
151             attention tensor means the architecture moved — do not skip it.",
152            unclaimed.len(),
153            &unclaimed[..unclaimed.len().min(5)]
154        );
155    }
156
157    Ok(DsaBindReport {
158        bound: specs.len(),
159        total_bytes,
160        unclaimed,
161    })
162}
163
164#[cfg(test)]
165mod tests;
metrale_model_arch/glm5next_dsa/binding.rs

metrale_model_arch/glm5next_dsa/
binding.rs
