spark_model/layers/ops/

glm5next_mhc.rs

// SPDX-License-Identifier: AGPL-3.0-only

//! GLM-5.3-Flash mHC kernel dispatch — the parts that are NOT DeepSeek-V4's.
//!
//! Separate file so `ops/hyper_connection.rs` (V4's proven dispatch) stays byte-untouched.
//!
//! `hc_pre` and `hc_post` need no wrapper here: the GLM kernels
//! `glm5next_mhc::{glm5next_hc_pre, glm5next_hc_post}` have signatures IDENTICAL to their
//! `hyper_connection` counterparts, so the GLM path calls `ops::hc_pre` / `ops::hc_post` with a
//! GLM `KernelHandle`. Only `hc_head` needs its own entry point, because GLM's takes no weights
//! at all.

use anyhow::Result;
use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle};
use spark_runtime::kernel_args::KernelLaunch;

/// Every kernel GLM-5.3's hyper-connection needs, all from the single module `glm5next_mhc`.
///
/// The point of this struct is target independence: a GLM kernel target must not have to carry
/// the DeepSeek-V4 `hyper_connection` module to resolve half of its own mHC.
/// `Copy` so all 45 layers can share one resolution — these are opaque handles, not state.
#[derive(Clone, Copy)]
pub struct Glm5NextMhcKernels {
    /// Broadcast the embedding into the `hc_mult` streams. First text layer only.
    ///
    /// 🪤 This is an architecture-neutral broadcast that Atlas already had — and GLM still
    /// needs its own, because `hyper_connection::hc_expand` lives in the **DeepSeek-V4
    /// target directory**. A kernel target merges `common/` plus its OWN model dir and
    /// cannot reach into another target's, so for the GLM target that module does not
    /// exist. Resolving it would fail at first construction, not fall back.
    pub hc_expand: KernelHandle,
    /// The FUSED single-block `hc_pre`. Still resolved, and still the oracle
    /// `examples/glm5next_hc_split_gate.rs` gates the split pair against — but the serve path
    /// goes through `hc_mix` + `hc_finish`, which are bit-identical and ~mix_hc times wider.
    pub hc_pre: KernelHandle,
    /// One block per mixing row: grid `(T, mix_hc)`.
    pub hc_mix: KernelHandle,
    /// Same kernel reading `hc_fn` at the width the checkpoint stores it (BF16).
    /// **Bit-identical** — widening BF16 to F32 is lossless, so it multiplies the same floats.
    /// `try_kernel`; selected per site by `Glm5NextMhcSiteWeights::hc_fn_bf16`.
    pub hc_mix_bf16: KernelHandle,
    /// Split + Sinkhorn + collapse, reading the mixes from global.
    pub hc_finish: KernelHandle,
    pub hc_post: KernelHandle,
    pub hc_head: KernelHandle,
}

/// The one module name GLM's mHC resolves from.
pub const GLM5NEXT_MHC_MODULE: &str = "glm5next_mhc";

impl Glm5NextMhcKernels {
    /// Resolve all six. `kernel()` (not `try_kernel`) — a missing mHC kernel is a hard error,
    /// never a silent fallback onto the DeepSeek variant.
    pub fn resolve(gpu: &dyn GpuBackend) -> Result<Self> {
        Ok(Self {
            hc_expand: gpu.kernel(GLM5NEXT_MHC_MODULE, "glm5next_hc_expand")?,
            hc_pre: gpu.kernel(GLM5NEXT_MHC_MODULE, "glm5next_hc_pre")?,
            hc_mix: gpu.kernel(GLM5NEXT_MHC_MODULE, "glm5next_hc_mix")?,
            hc_mix_bf16: crate::layers::try_kernel(
                gpu,
                GLM5NEXT_MHC_MODULE,
                "glm5next_hc_mix_bf16",
            ),
            hc_finish: gpu.kernel(GLM5NEXT_MHC_MODULE, "glm5next_hc_finish")?,
            hc_post: gpu.kernel(GLM5NEXT_MHC_MODULE, "glm5next_hc_post")?,
            hc_head: gpu.kernel(GLM5NEXT_MHC_MODULE, "glm5next_hc_head")?,
        })
    }
}

/// Expand one BF16 hidden state into the `hc_mult` FP32 highway streams. First layer only.
pub fn glm_hc_expand(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    hidden: DevicePtr,
    streams: DevicePtr,
    num_tokens: u32,
    hidden_size: u32,
    hc_mult: u32,
    stream: u64,
) -> Result<()> {
    KernelLaunch::new(gpu, kernel)
        .grid([num_tokens, 1, 1])
        .block([256, 1, 1])
        .arg_ptr(hidden)
        .arg_ptr(streams)
        .arg_u32(hidden_size)
        .arg_u32(hc_mult)
        .launch(stream)
}

/// Final collapse before the LM head: an **unweighted mean** over the `hc_mult` streams.
///
/// 🔴 Deliberately takes NO weight pointers. GLM's `Glm5NextTextHyperHead` has no parameters and
/// the checkpoint carries zero `hc_head` tensors; DeepSeek-V4's `ops::hc_head` reads
/// `hc_head.{fn,base,scale}`. The absent arguments are the guard against reaching for weights
/// that do not exist.
pub fn hc_head_mean(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    streams: DevicePtr,
    y_out: DevicePtr,
    num_tokens: u32,
    hidden_size: u32,
    hc_mult: u32,
    stream: u64,
) -> Result<()> {
    KernelLaunch::new(gpu, kernel)
        .grid([num_tokens, 1, 1])
        .block([256, 1, 1])
        .arg_ptr(streams)
        .arg_ptr(y_out)
        .arg_u32(hidden_size)
        .arg_u32(hc_mult)
        .launch(stream)
}

/// Per-site mHC weights. One set for the attention site, one for the FFN site.
///
/// 🪤 `hc_fn` is **FP32 to the kernel** and BF16 on disk — the loader upcasts. Reading it at
/// the on-disk width is the #341/#347 dtype-mismatch class, and the shapes do not say so.
#[derive(Debug, Clone, Copy)]
pub struct Glm5NextMhcSiteWeights {
    /// `[mix_hc, hc_mult * hidden]`, where `mix_hc = (2 + hc_mult) * hc_mult`. FP32 unless
    /// `hc_fn_bf16`, in which case BF16 — the width the checkpoint actually stores.
    pub hc_fn: DevicePtr,
    /// Is `hc_fn` BF16? Production sets it; the microtests and the split gate build their own
    /// F32 weights and leave it false, so the oracle they compare against is unchanged.
    pub hc_fn_bf16: bool,
    /// `[3]` FP32 — the three logit scales (pre, post, comb), in that order.
    pub hc_scale: DevicePtr,
    /// `[mix_hc]` FP32.
    pub hc_base: DevicePtr,
    /// `[mhc_mix_max_tokens(), mix_hc]` FP32 scratch: `hc_mix` writes it, `hc_finish` reads it.
    /// Per-site, so the two sites of a layer cannot alias; both run on one stream in order.
    /// 🪤 Sized by the FUNCTION, not by [`MHC_MIX_MAX_TOKENS`] — the const is only its floor.
    pub mix: DevicePtr,
}

/// FLOOR of the `mix` scratch's token bound — the value the stack shipped with, and the value
/// it still uses at every prefill sub-chunk of 256 rows or fewer.
pub const MHC_MIX_MAX_TOKENS: usize = 256;

/// Token bound the `mix` scratch is actually allocated at AND checked against.
///
/// 🔴 MEASURED 2026-09-22, WS1/P2: this was a bare `256` const, and it is a **hard ceiling on
/// `AVAROK_GLM_PREFILL_ROWS`** that has nothing to do with memory. At 512 rows the serve boots
/// fine and then every prefill dies at layer 0 with
/// *"glm_hc_pre: 512 tokens exceeds the 256-token `mix` scratch"* — HTTP 200, empty completion,
/// `runs/ws1-p2-prefill-n1n2/out/shared-r512/`. P1's "rows=1024 is not memory-feasible" never
/// saw this because the LOAD died first; memory was the second blocker, not the first.
///
/// The scratch is `[tokens, mix_hc(hc_mult)] FP32` **per site**: at GLM-5.3's `hc_mult = 4`
/// that is `tokens * 96 B` per site, 90 sites — **2.2 MB for the whole model at 256, 8.8 MB at
/// 1024**. The ceiling was never buying anything; it was slack that hardened into a limit.
///
/// 🪤 The allocation in `glm5_next_load` and the refusal in [`glm_hc_pre`] MUST read the same
/// number, which is why this is one `OnceLock` function and not two `max()` expressions. Two
/// expressions that agree today are a buffer overrun the day one of them is edited.
pub fn mhc_mix_max_tokens() -> usize {
    static T: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
    *T.get_or_init(|| {
        let t = crate::layers::glm5next_layer::prefill_rows().max(MHC_MIX_MAX_TOKENS);
        if t != MHC_MIX_MAX_TOKENS {
            tracing::warn!(
                "GLM mHC `mix` scratch widened to {t} tokens (floor {MHC_MIX_MAX_TOKENS}) to \
                 match AVAROK_GLM_PREFILL_ROWS"
            );
        }
        t
    })
}

/// `hc_pre`: collapse the `hc_mult` FP32 streams to one BF16 sequence and emit this site's
/// `post` / `comb` mixing coefficients.
///
/// 🪤 Named `glm_hc_pre`, not `hc_pre`: `ops` already exports DeepSeek-V4's `hc_pre`/`hc_post`
/// through a glob, and the two are NOT interchangeable — V4's reads `hc_head.{fn,base,scale}`
/// and uses a different mixing law. The rename is the compiler-enforced version of this
/// module's "never silently fall back onto the DeepSeek variant" rule; a glob collision here
/// resolved the wrong way would be a silent architecture swap.
///
/// The residual that `glm_hc_post` mixes is the stream tensor **as it entered here** — snapshot it
/// before calling, which is the skeleton's `ResidualStep::SaveResidual`. Overwriting the streams
/// in place before `hc_post` runs silently changes what the residual means.
#[allow(clippy::too_many_arguments)]
pub fn glm_hc_pre(
    gpu: &dyn GpuBackend,
    kernels: &Glm5NextMhcKernels,
    streams: DevicePtr,
    w: &Glm5NextMhcSiteWeights,
    y_out: DevicePtr,
    post_out: DevicePtr,
    comb_out: DevicePtr,
    num_tokens: u32,
    hidden_size: u32,
    hc_mult: u32,
    sinkhorn_iters: u32,
    norm_eps: f32,
    hc_eps: f32,
    stream: u64,
) -> Result<()> {
    let mix_hc = (2 + hc_mult) * hc_mult;
    let mix_cap = mhc_mix_max_tokens();
    if num_tokens as usize > mix_cap {
        anyhow::bail!(
            "glm_hc_pre: {num_tokens} tokens exceeds the {mix_cap}-token `mix` scratch \
             (floor {MHC_MIX_MAX_TOKENS}, widened to AVAROK_GLM_PREFILL_ROWS). Do not launch \
             past the allocation."
        );
    }
    // 🪤 The two kernels take the SAME arguments; only `hc_fn`'s element width differs, and it
    // is the pointer's own dtype, not something the signature can catch. Pairing the wrong
    // flag with the pointer reads BF16 as F32 (or the reverse) and produces plausible garbage.
    let mix_kernel = if w.hc_fn_bf16 && kernels.hc_mix_bf16.0 != 0 {
        kernels.hc_mix_bf16
    } else {
        kernels.hc_mix
    };
    KernelLaunch::new(gpu, mix_kernel)
        .grid([num_tokens, mix_hc, 1])
        .block([256, 1, 1])
        .arg_ptr(streams)
        .arg_ptr(w.hc_fn)
        .arg_ptr(w.mix)
        .arg_u32(hidden_size)
        .arg_u32(hc_mult)
        .arg_f32(norm_eps)
        .launch(stream)?;
    KernelLaunch::new(gpu, kernels.hc_finish)
        // `1 +` — `blockIdx.y == 0` runs the Sinkhorn and does NOT take a share of the collapse.
        .grid([num_tokens, 1 + collapse_blocks(hidden_size), 1])
        .block([256, 1, 1])
        .arg_ptr(streams)
        .arg_ptr(w.mix)
        .arg_ptr(w.hc_scale)
        .arg_ptr(w.hc_base)
        .arg_ptr(y_out)
        .arg_ptr(post_out)
        .arg_ptr(comb_out)
        .arg_u32(hidden_size)
        .arg_u32(hc_mult)
        .arg_u32(sinkhorn_iters)
        .arg_f32(hc_eps)
        .launch(stream)
}

/// `glm_hc_post`: `out[j] = post[j] * block_out + Σ_i comb[i][j] * residual[i]`.
///
/// `residual` is the pre-`hc_pre` stream snapshot, `block_out` this site's sublayer output.
/// `out` may alias `streams` — each output stream is a fresh combination, so writing back over
/// the highway is the intended flow.
#[allow(clippy::too_many_arguments)]
pub fn glm_hc_post(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    block_out: DevicePtr,
    residual: DevicePtr,
    post: DevicePtr,
    comb: DevicePtr,
    out: DevicePtr,
    num_tokens: u32,
    hidden_size: u32,
    hc_mult: u32,
    stream: u64,
) -> Result<()> {
    KernelLaunch::new(gpu, kernel)
        .grid([num_tokens, collapse_blocks(hidden_size), 1])
        .block([256, 1, 1])
        .arg_ptr(block_out)
        .arg_ptr(residual)
        .arg_ptr(post)
        .arg_ptr(comb)
        .arg_ptr(out)
        .arg_u32(hidden_size)
        .arg_u32(hc_mult)
        .launch(stream)
}

/// Blocks to spread an `H`-wide, per-element-independent pass over — `hc_finish`'s collapse and
/// all of `hc_post`.
///
/// 🔴 Both used to run on grid `(T, 1, 1)`: one block, i.e. ONE of the GB10's 48 SMs, moving
/// `hc_mult * H` floats, 90 times per token each. Nothing in either is a reduction — every `d`
/// is an independent output element — so the block count is free parallelism and the result is
/// bit-identical at any value of it. 256 is the block width both kernels launch at.
const fn collapse_blocks(hidden_size: u32) -> u32 {
    // `max(1)` by hand: `Ord::max` is not const yet.
    if hidden_size < 256 {
        1
    } else {
        hidden_size.div_ceil(256)
    }
}

/// `mix_hc` — the row count of `hc_fn` and `hc_base`: `(2 + hc_mult) * hc_mult`.
///
/// `pre` and `post` contribute one row per stream each, `comb` contributes `hc_mult` rows per
/// stream. Sizing either tensor with a different formula still yields a well-formed 2-D weight.
pub fn mix_hc(hc_mult: usize) -> usize {
    (2 + hc_mult) * hc_mult
}

#[cfg(test)]
mod mhc_shape_tests {
    use super::*;

    /// The `[pre | post | comb]` split of `hc_fn`'s rows. Getting `mix_hc` wrong shifts every
    /// coefficient the kernel reads, with no shape error.
    #[test]
    fn mix_hc_splits_into_pre_post_and_comb() {
        for hc in 1usize..=8 {
            assert_eq!(mix_hc(hc), hc + hc + hc * hc, "pre + post + comb rows");
        }
        // GLM-5.3 carries hc_mult = 2.
        assert_eq!(mix_hc(2), 8);
    }

    /// The kernel caps `hc_mult` at 4 via `GLM_HC_MAX_MIX = 24 = (2 + 4) * 4`. A larger
    /// multiplicity would overrun its fixed-size register arrays.
    #[test]
    fn the_kernels_mix_bound_is_hc_mult_four() {
        assert_eq!(mix_hc(4), 24, "GLM_HC_MAX_MIX in glm5next_mhc.cu");
        assert!(mix_hc(5) > 24, "hc_mult 5 would exceed the kernel's bound");
    }

    /// The `mix` bound must never drop BELOW the shipped floor, whatever the sub-chunk is:
    /// `glm_hc_pre` checks against this number and `glm5_next_load` allocates from it, and a
    /// bound under the floor would turn a refusal into an overrun. With no
    /// `AVAROK_GLM_PREFILL_ROWS` in the environment — the state of a test process — the
    /// default sub-chunk is 16, so the answer must still be the 256-token floor.
    #[test]
    fn the_mix_bound_never_goes_below_the_shipped_floor() {
        assert_eq!(mhc_mix_max_tokens(), MHC_MIX_MAX_TOKENS);
        assert!(mhc_mix_max_tokens() >= MHC_MIX_MAX_TOKENS);
    }

    /// The scratch's real cost, so "raise the ceiling" is never argued from a guess. Per SITE:
    /// `tokens * mix_hc(hc_mult) * 4 B`. GLM-5.3 is `hc_mult = 4`, two sites per layer, 45
    /// layers = 90 sites.
    #[test]
    fn the_mix_scratch_costs_megabytes_not_gigabytes() {
        let total = |tokens: usize| tokens * mix_hc(4) * 4 * 90;
        assert_eq!(
            total(256),
            2_211_840,
            "2.2 MB for the whole model at the floor"
        );
        assert_eq!(total(1024), 8_847_360, "8.8 MB at a 1024-row sub-chunk");
    }
}