metrale_model_arch/glm5next_kda/
config.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: KDA block geometry (`Glm5NextKdaConfig`) and device weights (`Glm5NextKdaWeights`).
4//!
5//! Owner: model-arch (GLM-5.3-Flash KDA).
6//! Invariants: none beyond the types; `Glm5NextKdaConfig::validate` states what the kernels need.
7
8use super::*;
9
10/// 2026-09-25: KDA geometry and the two numeric config values the block uses. The loader fills
11/// every field from `ModelConfig` except `l2_eps` and `chunk`; there is no `Default`.
12#[derive(Clone, Copy, Debug)]
13pub struct Glm5NextKdaConfig {
14    pub hidden: usize,
15    pub heads: usize,
16    pub head_dim: usize,
17    pub conv_kernel: usize,
18    /// 2026-09-25: `linear_attn_config.gate_lower_bound`.
19    pub gate_lower_bound: f32,
20    /// 2026-09-25: `rms_norm_eps`, read by `o_norm`.
21    pub rms_norm_eps: f32,
22    /// 2026-09-25: The L2 epsilon, applied as `1/sqrt(sum + eps)` rather than `max(norm, eps)`.
23    /// Not a config key.
24    pub l2_eps: f32,
25    /// 2026-09-25: Prefill chunk width, bounded by [`SMEM_CEILING`] through `validate`.
26    pub chunk: usize,
27}
28
29impl Glm5NextKdaConfig {
30    pub fn qkv_dim(&self) -> usize {
31        self.heads * self.head_dim
32    }
33    /// 2026-09-25: q | k | v concatenated: the channel count the depthwise conv sees.
34    pub fn conv_dim(&self) -> usize {
35        3 * self.qkv_dim()
36    }
37    /// 2026-09-25: The q | k channels, the only ones L2-normalised.
38    pub fn qk_channels(&self) -> usize {
39        2 * self.qkv_dim()
40    }
41    /// 2026-09-25: Elements of the FP32 `[heads, head_dim, head_dim]` K-major recurrent state.
42    pub fn recurrent_state_elems(&self) -> usize {
43        self.heads * self.head_dim * self.head_dim
44    }
45    /// 2026-09-25: Elements of the FP32 `[conv_dim, conv_kernel]` conv state (see `KdaSeqState`).
46    pub fn conv_state_elems(&self) -> usize {
47        self.conv_dim() * self.conv_kernel
48    }
49    pub fn smem_prepare(&self) -> usize {
50        (self.chunk * self.head_dim + self.chunk * self.chunk + self.chunk) * 4
51    }
52    pub fn smem_scan(&self) -> usize {
53        (2 * self.chunk * self.head_dim + self.chunk * self.chunk) * 4
54    }
55
56    pub fn validate(&self) -> Result<()> {
57        if !self.qk_channels().is_multiple_of(256) {
58            bail!("causal_conv1d_update_l2norm requires qk_channels % 256 == 0");
59        }
60        if self.head_dim != 128 {
61            bail!("the fused conv+L2 kernel hardcodes 2 heads per 256-thread block");
62        }
63        if self.conv_kernel > 4 {
64            bail!("the conv kernels keep the sliding window in 4 registers");
65        }
66        let (p, s) = (self.smem_prepare(), self.smem_scan());
67        if p > SMEM_CEILING || s > SMEM_CEILING {
68            bail!(
69                "chunk={} needs {p}/{s} B shared, ceiling {SMEM_CEILING}",
70                self.chunk
71            );
72        }
73        Ok(())
74    }
75}
76
77/// 2026-09-25: One KDA block's device weights. Torch `Linear` layout `[out, in]`, BF16, except the
78/// two F32 gate parameters. There is no `Z` tensor; the output gate is low-rank `g_a`/`g_b`.
79pub struct Glm5NextKdaWeights {
80    pub q_proj: DenseWeight,
81    pub k_proj: DenseWeight,
82    pub v_proj: DenseWeight,
83    /// 2026-09-25: `[conv_dim, conv_kernel]` BF16: the q, k and v conv tensors concatenated, with the
84    /// singleton middle dimension dropped.
85    pub conv: DenseWeight,
86    pub f_a: DenseWeight,
87    pub f_b: DenseWeight,
88    /// 2026-09-25: `[heads * head_dim]` F32, one value per channel.
89    pub dt_bias: DevicePtr,
90    /// 2026-09-25: `[heads]` F32, one value per head (unlike `dt_bias`).
91    pub a_log: DevicePtr,
92    pub b_proj: DenseWeight,
93    pub g_a: DenseWeight,
94    pub g_b: DenseWeight,
95    /// 2026-09-25: `[head_dim]` BF16.
96    pub o_norm: DenseWeight,
97    pub o_proj: DenseWeight,
98}
metrale_model_arch/glm5next_kda/config.rs

metrale_model_arch/glm5next_kda/
config.rs
