metrale_model_arch/glm5next_kda/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The GLM-5.3-Flash KDA attention block: one bound `self_attn` block per KDA layer.
4//!
5//! Owner: model-arch (GLM-5.3-Flash KDA).
6//! Invariants:
7//! - [`Glm5NextKdaLayer::new`] and [`Glm5NextKdaWorkspace::new`] refuse a config that fails
8//!   [`Glm5NextKdaConfig::validate`].
9//! - The forward calls (`decode`, `decode_k`, `prefill`) allocate nothing: they write into the
10//!   caller's [`Glm5NextKdaWorkspace`] and [`KdaSeqState`].
11//!
12//! ```text
13//! q|k|v_proj -> pack -> conv1d + SiLU -> L2(q,k only) -> kda_gate / sigmoid(b_proj)
14//!            -> kda_chunk (prefill) | kda_recurrent (decode)
15//!            -> sigmoid-gated RMSNorm(o_norm, g_b(g_a(h))) -> o_proj
16//! ```
17//!
18//! * The binder accepts only BF16 and F32 tensors ([`binding::KDA_TENSORS`]), and every
19//!   projection runs a BF16 kernel, so there is no dequantisation on this path.
20//! * The checkpoint stores `q_conv1d`/`k_conv1d`/`v_conv1d` separately; the binder concatenates
21//!   them in q, k, v order, the order `front_end` packs the projections in.
22//! * `o_norm` gates with a sigmoid, not SiLU, so it runs `kda_o_norm_gated_bf16`
23//!   (`kernels/gb10/common/kda_layer_ops.cu`) rather than `gated_rms_norm`.
24//! * `dense_mm_bf16` and `cublas_bf16_proj_dense` take no output row stride, so the three q/k/v
25//!   projections go to separate `[T, qkv]` buffers and `kda_pack_qkv_bf16` interleaves them.
26//! * The conv state keeps `conv_kernel` slots and the conv kernels shift left before convolving;
27//!   see [`KdaSeqState`].
28//! * Decode fuses conv, SiLU and L2 in one kernel; prefill runs a separate `l2_norm_bf16` over the
29//!   q|k channels. V is never normalised.
30
31pub mod binding;
32pub mod tp;
33pub mod tp_bind;
34
35mod config;
36mod decode;
37mod kernels;
38mod prefill;
39pub use config::{Glm5NextKdaConfig, Glm5NextKdaWeights};
40pub use kernels::Glm5NextKdaKernels;
41
42use anyhow::{Result, bail};
43use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
44use metrale_gpu_runtime::kernel_args::{KernelLaunch, div_ceil};
45
46use metrale_model_layers::layers::ops;
47use metrale_model_layers::weight_map::DenseWeight;
48
49/// 2026-09-25: Largest shared memory a prefill chunk may need. [`Glm5NextKdaConfig::validate`]
50/// refuses a `chunk` whose `smem_prepare` or `smem_scan` exceeds it.
51pub const SMEM_CEILING: usize = 49_152;
52
53const BLOCK: u32 = 128;
54
55/// 2026-09-25: V columns one block of the 1R+1W recurrent kernel owns: one warp. At
56/// `head_dim = 128` its request is `(3 * 128 + 32 * 129) * 4` = 18,048 B.
57const KDA_V_PER_BLOCK: usize = 32;
58/// 2026-09-25: Largest shared memory `stateful_row` requests for the 1R+1W kernel; a larger need
59/// launches the 2R+2W kernel instead.
60const KDA_SMEM_BUDGET: usize = 48 * 1024;
61
62/// 2026-09-25: `METRALE_GLM_KDA_NO_SMEM=1` selects the 2R+2W recurrent kernel. Read once per
63/// process; it is checked on every decode row.
64fn kda_no_smem() -> bool {
65    static F: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
66    *F.get_or_init(|| std::env::var("METRALE_GLM_KDA_NO_SMEM").as_deref() == Ok("1"))
67}
68
69/// 2026-09-25: The per-sequence state a KDA layer carries. The kernels update both buffers in place.
70///
71/// The conv buffer holds `conv_kernel` slots per channel. The conv kernels shift the window left
72/// and write the new input into the last slot before convolving, so slot 0 is shifted out
73/// unread. A state kept as `conv_kernel - 1` slots (the golden generators' layout) maps to
74/// slots `1..conv_kernel` here, before the shift.
75#[derive(Clone, Copy, Debug)]
76pub struct KdaSeqState {
77    /// 2026-09-25: `[conv_dim, conv_kernel]` FP32.
78    pub conv: DevicePtr,
79    /// 2026-09-25: `[heads, head_dim, head_dim]` FP32, K-major (`S[k * head_dim + v]`).
80    pub recurrent: DevicePtr,
81}
82
83/// 2026-09-25: Scratch for the forward path, sized once for `max_tokens` (prefill buffers padded
84/// to a whole number of chunks). The loader shares one workspace across the KDA layers.
85///
86/// The intermediate buffers are `pub` so the numeric tests can compare each stage, not only the
87/// layer output.
88pub struct Glm5NextKdaWorkspace {
89    /// 2026-09-25: `[3, T, qkv]` BF16: the q, k and v projection outputs, before the pack.
90    pub qkv_parts: DevicePtr,
91    /// 2026-09-25: `[T, conv_dim]` BF16: q|k|v per token, before the conv.
92    pub qkv_proj: DevicePtr,
93    /// 2026-09-25: `[T, conv_dim]` BF16: after conv + SiLU, with q|k L2-normalised (inside the
94    /// conv kernel on decode, by `l2_norm_bf16` in place on prefill).
95    pub conv_out: DevicePtr,
96    /// 2026-09-25: `[T_pad, qkv]` FP32: q, k and v widened for the chunked prefill. Prefill only.
97    pub q_f32: DevicePtr,
98    pub k_f32: DevicePtr,
99    pub v_f32: DevicePtr,
100    /// 2026-09-25: `[T_pad, heads, head_dim]` FP32: the bounded log-decay `kda_gate` writes.
101    pub gate: DevicePtr,
102    /// 2026-09-25: `[T_pad, heads]` FP32: `sigmoid(b_proj(h))`.
103    pub beta: DevicePtr,
104    /// 2026-09-25: `[T_pad, heads, head_dim]` FP32: the KDA core output, before `o_norm`.
105    pub core: DevicePtr,
106    /// 2026-09-25: `[T, qkv]` BF16: `f_b(f_a(h))`, the forget-gate input `kda_gate` reads. It has
107    /// the same shape as `out_gate` and must not alias it.
108    pub g_raw: DevicePtr,
109    /// 2026-09-25: `[T, qkv]` BF16: `g_b(g_a(h))`, the low-rank output gate.
110    pub out_gate: DevicePtr,
111    /// 2026-09-25: `[T, qkv]` BF16: after the sigmoid-gated RMSNorm.
112    pub o_norm_out: DevicePtr,
113    /// 2026-09-25: `[T, hidden]` BF16: the block output.
114    pub final_out: DevicePtr,
115    lowrank: DevicePtr,
116    beta_bf16: DevicePtr,
117    chunk_gc: DevicePtr,
118    chunk_u: DevicePtr,
119    chunk_w: DevicePtr,
120    max_tokens: usize,
121    t_pad: usize,
122}
123
124impl Glm5NextKdaWorkspace {
125    pub fn new(gpu: &dyn GpuBackend, cfg: &Glm5NextKdaConfig, max_tokens: usize) -> Result<Self> {
126        cfg.validate()?;
127        if max_tokens == 0 {
128            bail!("workspace needs max_tokens >= 1");
129        }
130        let (qkv, cd, hd, h) = (cfg.qkv_dim(), cfg.conv_dim(), cfg.head_dim, cfg.heads);
131        let t = max_tokens;
132        let t_pad = t.div_ceil(cfg.chunk) * cfg.chunk;
133        let n = t_pad * qkv;
134        Ok(Self {
135            qkv_parts: gpu.alloc(3 * t * qkv * 2)?,
136            qkv_proj: gpu.alloc(t * cd * 2)?,
137            conv_out: gpu.alloc(t * cd * 2)?,
138            q_f32: gpu.alloc(n * 4)?,
139            k_f32: gpu.alloc(n * 4)?,
140            v_f32: gpu.alloc(n * 4)?,
141            gate: gpu.alloc(n * 4)?,
142            beta: gpu.alloc(t_pad * h * 4)?,
143            core: gpu.alloc(n * 4)?,
144            g_raw: gpu.alloc(t * qkv * 2)?,
145            out_gate: gpu.alloc(t * qkv * 2)?,
146            o_norm_out: gpu.alloc(t * qkv * 2)?,
147            final_out: gpu.alloc(t * cfg.hidden * 2)?,
148            lowrank: gpu.alloc(t * hd * 2)?,
149            beta_bf16: gpu.alloc(t * h * 2)?,
150            chunk_gc: gpu.alloc(n * 4)?,
151            chunk_u: gpu.alloc(n * 4)?,
152            chunk_w: gpu.alloc(n * 4)?,
153            max_tokens: t,
154            t_pad,
155        })
156    }
157
158    pub fn max_tokens(&self) -> usize {
159        self.max_tokens
160    }
161    pub fn t_pad(&self) -> usize {
162        self.t_pad
163    }
164}
165
166/// 2026-09-25: One bound KDA attention block.
167pub struct Glm5NextKdaLayer {
168    /// 2026-09-25: The checkpoint layer index this block was bound from.
169    pub layer_idx: usize,
170    pub cfg: Glm5NextKdaConfig,
171    pub weights: Glm5NextKdaWeights,
172    pub kernels: Glm5NextKdaKernels,
173}
174
175impl Glm5NextKdaLayer {
176    pub fn new(
177        layer_idx: usize,
178        cfg: Glm5NextKdaConfig,
179        weights: Glm5NextKdaWeights,
180        kernels: Glm5NextKdaKernels,
181    ) -> Result<Self> {
182        cfg.validate()?;
183        Ok(Self {
184            layer_idx,
185            cfg,
186            weights,
187            kernels,
188        })
189    }
190
191    #[allow(clippy::too_many_arguments)]
192    fn gemm(
193        &self,
194        gpu: &dyn GpuBackend,
195        input: DevicePtr,
196        weight: &DenseWeight,
197        out: DevicePtr,
198        m: usize,
199        n: usize,
200        k: usize,
201        stream: u64,
202    ) -> Result<()> {
203        // 2026-09-25: Only M above `DENSE_GEMV_BATCHM_MAX_M` goes to cuBLASLt. Below it
204        // `dense_mm_bf16` runs the M = 1 GEMV or the batched GEMV, whose rows carry the same
205        // bits, so the cuBLASLt switch never changes those widths.
206        if m > ops::DENSE_GEMV_BATCHM_MAX_M as usize && crate::glm5next_layer::cublas_wide_proj() {
207            return ops::cublas_bf16_proj_dense(
208                input,
209                weight.weight,
210                out,
211                m as u32,
212                n as u32,
213                k as u32,
214                stream,
215            );
216        }
217        ops::dense_mm_bf16(
218            gpu,
219            &ops::DenseMmKernels {
220                gemm: self.kernels.gemm,
221                gemv: self.kernels.gemv,
222                batchm: self.kernels.gemv_batchm,
223            },
224            input,
225            weight.weight,
226            out,
227            m,
228            n,
229            k,
230            stream,
231        )
232    }
233
234    /// 2026-09-25: Projections, forget gate, beta and output gate, shared by decode and prefill.
235    /// All of them read the block's input hidden state, never the post-conv activations.
236    fn front_end(
237        &self,
238        gpu: &dyn GpuBackend,
239        hidden: DevicePtr,
240        t: usize,
241        ws: &Glm5NextKdaWorkspace,
242        stream: u64,
243    ) -> Result<()> {
244        let c = &self.cfg;
245        let (hid, qkv, hd) = (c.hidden, c.qkv_dim(), c.head_dim);
246
247        // 2026-09-25: Three separate `[T, qkv]` projections, then one pack (see the module doc).
248        for (i, w) in [
249            &self.weights.q_proj,
250            &self.weights.k_proj,
251            &self.weights.v_proj,
252        ]
253        .into_iter()
254        .enumerate()
255        {
256            self.gemm(
257                gpu,
258                hidden,
259                w,
260                ws.qkv_parts.offset(i * t * qkv * 2),
261                t,
262                qkv,
263                hid,
264                stream,
265            )?;
266        }
267        KernelLaunch::new(gpu, self.kernels.pack)
268            .grid([div_ceil(qkv as u32, 256), t as u32, 1])
269            .block([256, 1, 1])
270            .arg_ptr(ws.qkv_parts)
271            .arg_ptr(ws.qkv_parts.offset(t * qkv * 2))
272            .arg_ptr(ws.qkv_parts.offset(2 * t * qkv * 2))
273            .arg_ptr(ws.qkv_proj)
274            .arg_u32(t as u32)
275            .arg_u32(qkv as u32)
276            .launch(stream)?;
277
278        // 2026-09-25: Low-rank forget gate, hidden -> head_dim -> heads * head_dim, then
279        // `lower_bound * sigmoid(exp(A_log[h]) * (g[c] + dt_bias[c]))` in `kda_gate_bf16`.
280        // `A_log` is per head and `dt_bias` per channel.
281        self.gemm(
282            gpu,
283            hidden,
284            &self.weights.f_a,
285            ws.lowrank,
286            t,
287            hd,
288            hid,
289            stream,
290        )?;
291        self.gemm(
292            gpu,
293            ws.lowrank,
294            &self.weights.f_b,
295            ws.g_raw,
296            t,
297            qkv,
298            hd,
299            stream,
300        )?;
301        KernelLaunch::new(gpu, self.kernels.gate)
302            .grid([(t * c.heads) as u32, 1, 1])
303            .block([BLOCK, 1, 1])
304            .arg_ptr(ws.g_raw)
305            .arg_ptr(self.weights.dt_bias)
306            .arg_ptr(self.weights.a_log)
307            .arg_ptr(ws.gate)
308            .arg_u32(t as u32)
309            .arg_u32(c.heads as u32)
310            .arg_u32(hd as u32)
311            .arg_f32(c.gate_lower_bound)
312            .launch(stream)?;
313
314        // 2026-09-25: beta = sigmoid(b_proj(hidden)); the KDA kernels read it after the sigmoid.
315        self.gemm(
316            gpu,
317            hidden,
318            &self.weights.b_proj,
319            ws.beta_bf16,
320            t,
321            c.heads,
322            hid,
323            stream,
324        )?;
325        let n = t * c.heads;
326        KernelLaunch::new(gpu, self.kernels.sigmoid)
327            .grid([div_ceil(n as u32, 256), 1, 1])
328            .block([256, 1, 1])
329            .arg_ptr(ws.beta_bf16)
330            .arg_ptr(ws.beta)
331            .arg_u32(n as u32)
332            .launch(stream)?;
333
334        // 2026-09-25: Low-rank output gate; a KDA block has no `Z` tensor (`KDA_TENSORS`).
335        self.gemm(
336            gpu,
337            hidden,
338            &self.weights.g_a,
339            ws.lowrank,
340            t,
341            hd,
342            hid,
343            stream,
344        )?;
345        self.gemm(
346            gpu,
347            ws.lowrank,
348            &self.weights.g_b,
349            ws.out_gate,
350            t,
351            qkv,
352            hd,
353            stream,
354        )
355    }
356
357    /// 2026-09-25: Sigmoid-gated RMSNorm, then `o_proj`.
358    fn back_end(
359        &self,
360        gpu: &dyn GpuBackend,
361        t: usize,
362        ws: &Glm5NextKdaWorkspace,
363        stream: u64,
364    ) -> Result<()> {
365        let c = &self.cfg;
366        KernelLaunch::new(gpu, self.kernels.o_norm)
367            .grid([(t * c.heads) as u32, 1, 1])
368            .block([c.head_dim as u32, 1, 1])
369            .arg_ptr(ws.core)
370            .arg_ptr(ws.out_gate)
371            .arg_ptr(self.weights.o_norm.weight)
372            .arg_ptr(ws.o_norm_out)
373            .arg_u32(c.head_dim as u32)
374            .arg_f32(c.rms_norm_eps)
375            .launch(stream)?;
376        self.gemm(
377            gpu,
378            ws.o_norm_out,
379            &self.weights.o_proj,
380            ws.final_out,
381            t,
382            c.hidden,
383            c.qkv_dim(),
384            stream,
385        )
386    }
387}
metrale_model_arch/glm5next_kda/mod.rs

metrale_model_arch/glm5next_kda/
mod.rs
