metrale_model_arch/glm5next_kda/
kernels.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The kernel handles a KDA block launches (`Glm5NextKdaKernels`).
4//!
5//! Owner: model-arch (GLM-5.3-Flash KDA).
6//! Invariants:
7//! - `resolve` fails if any entry point other than `dense_gemv_bf16_batchm` and
8//!   `kda_recurrent_decode_bf16_smem` is missing; those two resolve to handle 0 when absent.
9
10use super::*;
11
12#[derive(Clone, Copy)]
13pub struct Glm5NextKdaKernels {
14    pub gemm: KernelHandle,
15    /// 2026-09-25: The M = 1 kernel: `dense_mm_bf16` sends every decode projection here. Measured
16    /// 2026-08-28: `dense_gemm_bf16` at M = 1 moved 58 GB/s and took 32% of the GLM decode step.
17    pub gemv: KernelHandle,
18    /// 2026-09-25: The batched GEMV: `2..=DENSE_GEMV_BATCHM_MAX_M` rows in one pass over the
19    /// weight, so a K-row `decode_k` reads each projection weight once. `0` on a target without
20    /// it; [`ops::dense_mm_bf16`] then sends those rows to the tile GEMM.
21    pub gemv_batchm: KernelHandle,
22    pub conv_decode: KernelHandle,
23    pub conv_prefill: KernelHandle,
24    pub l2: KernelHandle,
25    pub gate: KernelHandle,
26    pub chunk_prepare: KernelHandle,
27    pub chunk_scan: KernelHandle,
28    pub recurrent: KernelHandle,
29    /// 2026-09-25: 1R+1W sibling of `recurrent`: the decayed state column stays in shared memory
30    /// between the two passes. Resolved with `try_kernel`; `0` selects the 2R+2W kernel.
31    pub recurrent_smem: KernelHandle,
32    pub o_norm: KernelHandle,
33    pub split_widen: KernelHandle,
34    pub sigmoid: KernelHandle,
35    pub fill: KernelHandle,
36    pub pack: KernelHandle,
37}
38
39impl Glm5NextKdaKernels {
40    pub const ENTRY_POINTS: usize = 14;
41
42    pub fn resolve(gpu: &dyn GpuBackend) -> Result<Self> {
43        Ok(Self {
44            gemm: gpu.kernel("gemm", "dense_gemm_bf16")?,
45            gemv: gpu.kernel("gemv", "dense_gemv_bf16")?,
46            gemv_batchm: metrale_model_layers::layers::try_kernel(
47                gpu,
48                "dense_gemv_bf16_batchm",
49                "dense_gemv_bf16_batchm",
50            ),
51            conv_decode: gpu.kernel("causal_conv1d", "causal_conv1d_update_l2norm")?,
52            conv_prefill: gpu.kernel("causal_conv1d", "causal_conv1d_update_prefill")?,
53            l2: gpu.kernel("norm", "l2_norm_bf16")?,
54            gate: gpu.kernel("kda_gate", "kda_gate_bf16")?,
55            chunk_prepare: gpu.kernel("kda_chunk", "kda_chunk_prepare")?,
56            chunk_scan: gpu.kernel("kda_chunk", "kda_chunk_scan")?,
57            recurrent: gpu.kernel("kda_recurrent", "kda_recurrent_decode_bf16")?,
58            recurrent_smem: metrale_model_layers::layers::try_kernel(
59                gpu,
60                "kda_recurrent",
61                "kda_recurrent_decode_bf16_smem",
62            ),
63            o_norm: gpu.kernel("kda_layer_ops", "kda_o_norm_gated_bf16")?,
64            split_widen: gpu.kernel("kda_layer_ops", "kda_split_widen")?,
65            sigmoid: gpu.kernel("kda_layer_ops", "kda_sigmoid_bf16_f32")?,
66            fill: gpu.kernel("kda_layer_ops", "kda_fill_f32")?,
67            pack: gpu.kernel("kda_layer_ops", "kda_pack_qkv_bf16")?,
68        })
69    }
70}
metrale_model_arch/glm5next_kda/kernels.rs

metrale_model_arch/glm5next_kda/
kernels.rs
