metrale_model_arch/
glm5next_mhc.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Host launchers for the GLM-5.3 mHC kernels: expand, pre (the `hc_mix` +
4//! `hc_finish` pair), post, and the head mean.
5//!
6//! Owner: model-arch (GLM-5.3).
7//! Invariants:
8//! - Every kernel resolves from `GLM5NEXT_MHC_MODULE`, none from DeepSeek-V4's
9//!   `hyper_connection` module.
10//! - `glm_hc_pre` refuses more tokens than `mhc_mix_max_tokens()`, the bound the loader sizes
11//!   each site's `mix` scratch from.
12
13use anyhow::Result;
14use metrale_gpu_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
15use metrale_gpu_runtime::kernel_args::KernelLaunch;
16
17/// 2026-09-25: Every kernel the GLM-5.3 hyper-connection launches. `Copy`: the loader resolves
18/// them once and every layer holds a copy.
19#[derive(Clone, Copy)]
20pub struct Glm5NextMhcKernels {
21    /// 2026-09-25: Broadcast the embedding into the `hc_mult` highway streams; first text layer
22    /// only. A GLM target cannot use `hyper_connection::hc_expand`: that kernel is in the
23    /// DeepSeek-V4 model directory, and a target compiles only `common/` and its own model
24    /// directory.
25    pub hc_expand: KernelHandle,
26    /// 2026-09-25: The fused `glm5next_hc_pre`, one block per token. `glm_hc_pre` does not launch
27    /// it; `examples/glm5next_hc_split_gate.rs` uses it as the oracle for `hc_mix` + `hc_finish`.
28    pub hc_pre: KernelHandle,
29    /// 2026-09-25: `glm5next_hc_mix`: one block per (token, mixing row), grid `(T, mix_hc)`,
30    /// reading an FP32 `hc_fn`.
31    pub hc_mix: KernelHandle,
32    /// 2026-09-25: `glm5next_hc_mix` reading a BF16 `hc_fn`. Widening BF16 to FP32 is exact, so it
33    /// computes the FP32 kernel's result on the widened weights. Optional (`try_kernel`, 0 when
34    /// absent); `glm_hc_pre` uses it when the site's `hc_fn_bf16` is set and the handle is not 0.
35    pub hc_mix_bf16: KernelHandle,
36    /// 2026-09-25: `glm5next_hc_finish`: from the mixes, `post`, `comb` (Sinkhorn) and the
37    /// collapsed row `y`.
38    pub hc_finish: KernelHandle,
39    pub hc_post: KernelHandle,
40    pub hc_head: KernelHandle,
41}
42
43/// 2026-09-25: The kernel module every GLM mHC kernel resolves from.
44pub const GLM5NEXT_MHC_MODULE: &str = "glm5next_mhc";
45
46impl Glm5NextMhcKernels {
47    /// 2026-09-25: Resolve the kernels. A missing one is an error, except `hc_mix_bf16`, which is
48    /// 0 when absent.
49    pub fn resolve(gpu: &dyn GpuBackend) -> Result<Self> {
50        Ok(Self {
51            hc_expand: gpu.kernel(GLM5NEXT_MHC_MODULE, "glm5next_hc_expand")?,
52            hc_pre: gpu.kernel(GLM5NEXT_MHC_MODULE, "glm5next_hc_pre")?,
53            hc_mix: gpu.kernel(GLM5NEXT_MHC_MODULE, "glm5next_hc_mix")?,
54            hc_mix_bf16: metrale_model_layers::layers::try_kernel(
55                gpu,
56                GLM5NEXT_MHC_MODULE,
57                "glm5next_hc_mix_bf16",
58            ),
59            hc_finish: gpu.kernel(GLM5NEXT_MHC_MODULE, "glm5next_hc_finish")?,
60            hc_post: gpu.kernel(GLM5NEXT_MHC_MODULE, "glm5next_hc_post")?,
61            hc_head: gpu.kernel(GLM5NEXT_MHC_MODULE, "glm5next_hc_head")?,
62        })
63    }
64}
65
66/// 2026-09-25: Broadcast each of `num_tokens` BF16 hidden rows into its `hc_mult` FP32 highway
67/// streams.
68pub fn glm_hc_expand(
69    gpu: &dyn GpuBackend,
70    kernel: KernelHandle,
71    hidden: DevicePtr,
72    streams: DevicePtr,
73    num_tokens: u32,
74    hidden_size: u32,
75    hc_mult: u32,
76    stream: u64,
77) -> Result<()> {
78    KernelLaunch::new(gpu, kernel)
79        .grid([num_tokens, 1, 1])
80        .block([256, 1, 1])
81        .arg_ptr(hidden)
82        .arg_ptr(streams)
83        .arg_u32(hidden_size)
84        .arg_u32(hc_mult)
85        .launch(stream)
86}
87
88/// 2026-09-25: Collapse the highway to one BF16 row per token: the unweighted mean of the
89/// `hc_mult` streams. It takes no weights, unlike DeepSeek-V4's `ops::hc_head`.
90pub fn hc_head_mean(
91    gpu: &dyn GpuBackend,
92    kernel: KernelHandle,
93    streams: DevicePtr,
94    y_out: DevicePtr,
95    num_tokens: u32,
96    hidden_size: u32,
97    hc_mult: u32,
98    stream: u64,
99) -> Result<()> {
100    KernelLaunch::new(gpu, kernel)
101        .grid([num_tokens, 1, 1])
102        .block([256, 1, 1])
103        .arg_ptr(streams)
104        .arg_ptr(y_out)
105        .arg_u32(hidden_size)
106        .arg_u32(hc_mult)
107        .launch(stream)
108}
109
110/// 2026-09-25: One mHC site's weights; a layer has one set for the attention site and one for
111/// the FFN site.
112#[derive(Debug, Clone, Copy)]
113pub struct Glm5NextMhcSiteWeights {
114    /// 2026-09-25: `[mix_hc, hc_mult * hidden]` (see `mix_hc`): BF16 when `hc_fn_bf16` is set,
115    /// FP32 otherwise. The loader uploads it as BF16.
116    pub hc_fn: DevicePtr,
117    /// 2026-09-25: Whether `hc_fn` is BF16; the loader sets it.
118    pub hc_fn_bf16: bool,
119    /// 2026-09-25: `[3]` FP32: the logit scales of pre, post and comb, in that order.
120    pub hc_scale: DevicePtr,
121    /// 2026-09-25: `[mix_hc]` FP32.
122    pub hc_base: DevicePtr,
123    /// 2026-09-25: `[mhc_mix_max_tokens(), mix_hc]` FP32 scratch that `hc_mix` writes and
124    /// `hc_finish` reads; the loader allocates one per site.
125    pub mix: DevicePtr,
126}
127
128/// 2026-09-25: The lower bound of `mhc_mix_max_tokens()`.
129pub const MHC_MIX_MAX_TOKENS: usize = 256;
130
131/// 2026-09-25: Token bound of each site's `mix` scratch, `max(prefill_rows(), MHC_MIX_MAX_TOKENS)`,
132/// read once. The loader allocates the scratch from it and `glm_hc_pre` refuses more tokens
133/// than it, so the allocation and the check read one value.
134pub fn mhc_mix_max_tokens() -> usize {
135    static T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
136    *T.get_or_init(|| {
137        let t = crate::glm5next_layer::prefill_rows().max(MHC_MIX_MAX_TOKENS);
138        if t != MHC_MIX_MAX_TOKENS {
139            tracing::warn!(
140                "GLM mHC `mix` scratch widened to {t} tokens (floor {MHC_MIX_MAX_TOKENS}) to \
141                 match METRALE_GLM_PREFILL_ROWS"
142            );
143        }
144        t
145    })
146}
147
148/// 2026-09-25: Collapse each token's `hc_mult` FP32 streams to one BF16 row in `y_out`, and write
149/// this site's `post` and `comb`, by launching `hc_mix` (or `hc_mix_bf16`) then `hc_finish`.
150/// `streams` is read, not written, so `glm_hc_post` can take it as its residual. Errors when
151/// `num_tokens` exceeds `mhc_mix_max_tokens()`.
152#[allow(clippy::too_many_arguments)]
153pub fn glm_hc_pre(
154    gpu: &dyn GpuBackend,
155    kernels: &Glm5NextMhcKernels,
156    streams: DevicePtr,
157    w: &Glm5NextMhcSiteWeights,
158    y_out: DevicePtr,
159    post_out: DevicePtr,
160    comb_out: DevicePtr,
161    num_tokens: u32,
162    hidden_size: u32,
163    hc_mult: u32,
164    sinkhorn_iters: u32,
165    norm_eps: f32,
166    hc_eps: f32,
167    stream: u64,
168) -> Result<()> {
169    let mix_hc = (2 + hc_mult) * hc_mult;
170    let mix_cap = mhc_mix_max_tokens();
171    if num_tokens as usize > mix_cap {
172        anyhow::bail!(
173            "glm_hc_pre: {num_tokens} tokens exceeds the {mix_cap}-token `mix` scratch \
174             (floor {MHC_MIX_MAX_TOKENS}, widened to METRALE_GLM_PREFILL_ROWS). Do not launch \
175             past the allocation."
176        );
177    }
178    // 2026-09-25: Both mix kernels take the same arguments; only the element width of `hc_fn`
179    // differs, so `hc_fn_bf16` must describe the pointer.
180    let mix_kernel = if w.hc_fn_bf16 && kernels.hc_mix_bf16.0 != 0 {
181        kernels.hc_mix_bf16
182    } else {
183        kernels.hc_mix
184    };
185    KernelLaunch::new(gpu, mix_kernel)
186        .grid([num_tokens, mix_hc, 1])
187        .block([256, 1, 1])
188        .arg_ptr(streams)
189        .arg_ptr(w.hc_fn)
190        .arg_ptr(w.mix)
191        .arg_u32(hidden_size)
192        .arg_u32(hc_mult)
193        .arg_f32(norm_eps)
194        .launch(stream)?;
195    KernelLaunch::new(gpu, kernels.hc_finish)
196        // 2026-09-25: Block `y == 0` computes post, comb and the Sinkhorn; blocks `1..` split the
197        // collapse.
198        .grid([num_tokens, 1 + collapse_blocks(hidden_size), 1])
199        .block([256, 1, 1])
200        .arg_ptr(streams)
201        .arg_ptr(w.mix)
202        .arg_ptr(w.hc_scale)
203        .arg_ptr(w.hc_base)
204        .arg_ptr(y_out)
205        .arg_ptr(post_out)
206        .arg_ptr(comb_out)
207        .arg_u32(hidden_size)
208        .arg_u32(hc_mult)
209        .arg_u32(sinkhorn_iters)
210        .arg_f32(hc_eps)
211        .launch(stream)
212}
213
214/// 2026-09-25: `out[j] = post[j] * block_out + Σ_i comb[i][j] * residual[i]` for each token.
215/// `out` may alias `residual`: the kernel reads every residual stream of an element before it
216/// writes that element.
217#[allow(clippy::too_many_arguments)]
218pub fn glm_hc_post(
219    gpu: &dyn GpuBackend,
220    kernel: KernelHandle,
221    block_out: DevicePtr,
222    residual: DevicePtr,
223    post: DevicePtr,
224    comb: DevicePtr,
225    out: DevicePtr,
226    num_tokens: u32,
227    hidden_size: u32,
228    hc_mult: u32,
229    stream: u64,
230) -> Result<()> {
231    KernelLaunch::new(gpu, kernel)
232        .grid([num_tokens, collapse_blocks(hidden_size), 1])
233        .block([256, 1, 1])
234        .arg_ptr(block_out)
235        .arg_ptr(residual)
236        .arg_ptr(post)
237        .arg_ptr(comb)
238        .arg_ptr(out)
239        .arg_u32(hidden_size)
240        .arg_u32(hc_mult)
241        .launch(stream)
242}
243
244/// 2026-09-25: Blocks along grid y for `hc_finish`'s collapse and for `hc_post`: `ceil(H / 256)`,
245/// at least 1, where 256 is the block width both launch at. Each output element is computed on
246/// its own, so the block count does not change the result.
247const fn collapse_blocks(hidden_size: u32) -> u32 {
248    if hidden_size < 256 {
249        1
250    } else {
251        hidden_size.div_ceil(256)
252    }
253}
254
255/// 2026-09-25: Row count of `hc_fn` and `hc_base`, `(2 + hc_mult) * hc_mult`: `hc_mult` rows
256/// each for pre and post, then `hc_mult * hc_mult` for comb.
257pub fn mix_hc(hc_mult: usize) -> usize {
258    (2 + hc_mult) * hc_mult
259}
260
261#[cfg(test)]
262mod mhc_shape_tests {
263    use super::*;
264
265    /// 2026-09-25: `mix_hc` is the pre, post and comb row counts added, for `hc_mult` 1 to 8.
266    #[test]
267    fn mix_hc_splits_into_pre_post_and_comb() {
268        for hc in 1usize..=8 {
269            assert_eq!(mix_hc(hc), hc + hc + hc * hc, "pre + post + comb rows");
270        }
271        assert_eq!(mix_hc(2), 8);
272    }
273
274    /// 2026-09-25: `GLM_HC_MAX_MIX` in `glm5next_mhc.cu` is 24, which is `mix_hc(4)`; `hc_mult` 5
275    /// would exceed it.
276    #[test]
277    fn the_kernels_mix_bound_is_hc_mult_four() {
278        assert_eq!(mix_hc(4), 24, "GLM_HC_MAX_MIX in glm5next_mhc.cu");
279        assert!(mix_hc(5) > 24, "hc_mult 5 would exceed the kernel's bound");
280    }
281
282    /// 2026-09-25: Without `METRALE_GLM_PREFILL_ROWS`, `prefill_rows()` is `PREFILL_ROWS` (16), so
283    /// the bound is the 256-token floor.
284    #[test]
285    fn the_mix_bound_never_goes_below_the_shipped_floor() {
286        assert_eq!(mhc_mix_max_tokens(), MHC_MIX_MAX_TOKENS);
287        assert!(mhc_mix_max_tokens() >= MHC_MIX_MAX_TOKENS);
288    }
289
290    /// 2026-09-25: Total `mix` scratch over 90 sites (45 text layers, 2 sites each) at
291    /// `hc_mult` 4, the config fixture's values.
292    #[test]
293    fn the_mix_scratch_costs_megabytes_not_gigabytes() {
294        let total = |tokens: usize| tokens * mix_hc(4) * 4 * 90;
295        assert_eq!(
296            total(256),
297            2_211_840,
298            "2.2 MB for the whole model at the floor"
299        );
300        assert_eq!(total(1024), 8_847_360, "8.8 MB at a 1024-row sub-chunk");
301    }
302}
metrale_model_arch/glm5next_mhc.rs

metrale_model_arch/
glm5next_mhc.rs
