spark_model/layers/glm5next_kda_ref/

mod.rs

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

//! GLM-5.3-Flash **KDA (Kimi Delta Attention) CPU reference** — Slice 2B.
//!
//! Design artifact, **not** a production path. Nothing here runs on GPU, nothing here is wired
//! into a forward pass, and no checkpoint tensor is bound. Its only job is to prove the GLM KDA
//! equations in Atlas-shaped code against golden vectors produced by HuggingFace itself, before a
//! single CUDA kernel is written.
//!
//! # Why this exists
//!
//! GLM KDA is **not** a thin parameterization of Atlas's Qwen GDN (`layers::qwen3_ssm`). Four
//! structural items are genuinely new:
//!
//! 1. **Decay is per (head, key-channel)**, not scalar-per-head. Atlas's `compute_gdn_gates`
//!    writes `gate_out[num_tokens, num_v_heads]`; KDA needs `[T, H, head_dim]`.
//! 2. **The decay source is a low-rank projection** `f_a: hidden -> head_dim`,
//!    `f_b: head_dim -> H*head_dim`. Atlas derives its scalar from the fused `BA` projection.
//! 3. **The gate law is bounded**: `lower_bound * sigmoid(exp(A_log) * (g + dt_bias))`, versus
//!    Atlas's unbounded `exp(-exp(A_log) * softplus(a + dt_bias))`.
//! 4. **The output gate is low-rank** `g_a`/`g_b`. Atlas takes a full-rank `Z` out of its fused
//!    `QKVZ` projection; the KDA checkpoint has no `Z` tensor at all.
//!
//! See `docs/glm5next/KDA-VS-QWEN-GDN.md` for the full REUSE/ADAPT/NEW table.
//!
//! # Provenance of the goldens
//!
//! `kda_golden.json` is generated by `gen_kda_golden.py` (kept beside it), which calls the real
//! `transformers` **5.16.1** `glm5_next` module — `Glm5NextTextForgetGate`,
//! `Glm5NextTextRMSNormGated`, `l2norm`, `recurrent_kimi_delta_attention`,
//! `chunk_kimi_delta_attention`. No equation is re-derived on the Python side either.
//!
//! `transformers` **5.16.0 contains no `glm5_next` at all** even though the checkpoint declares
//! that version; 5.16.1 is the first source-bearing release.
//!
//! # Config values are read, never defaulted
//!
//! `gate_lower_bound`, `rms_norm_eps` and `hidden_act` must all come from the checkpoint config.
//! vLLM happens to agree with HF on this checkpoint only by coincidence: it looks up the legacy
//! key `lower_bound` while the checkpoint stores `gate_lower_bound`, and falls back to a default
//! of `-5.0` that happens to match. Atlas must not inherit that.
//!
//! # Layout conventions
//!
//! * `q`/`k`/`v`/`gate`: `[T, H, D]`, row-major, index `((t * H) + h) * D + d`.
//! * `beta`: `[T, H]`, index `t * H + h`.
//! * `dt_bias`: `[H * D]` — per **channel**.
//! * `a_log`: `[H]` — per **head**.
//! * `state`: `[H, D_k, D_v]`, index `(h * D_k + kd) * D_v + vd`.
//! * Weights follow the torch `Linear` convention `[out, in]`, row-major.

/// KDA geometry. Production is `hidden 4096 / heads 64 / head_dim 128`; the low-rank width for
/// both `f_a`/`f_b` and `g_a`/`g_b` equals `head_dim`.
#[derive(Clone, Copy, Debug)]
pub struct KdaDims {
    pub hidden: usize,
    pub heads: usize,
    pub head_dim: usize,
    pub tokens: usize,
}

#[inline]
fn sigmoid(x: f32) -> f32 {
    1.0 / (1.0 + (-x).exp())
}

/// `x / sqrt(sum(x^2) + eps)` over the trailing dimension.
///
/// Deliberately **not** `x / max(norm, eps)`. HF's `l2norm` adds `eps` *inside* the square root to
/// match the original FLA triton kernel, and the vLLM kernels do the same
/// (`b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)`). The two forms differ for small-norm rows.
pub fn l2norm_rows(x: &[f32], d: usize, eps: f32) -> Vec<f32> {
    let mut out = vec![0.0f32; x.len()];
    for (row_in, row_out) in x.chunks_exact(d).zip(out.chunks_exact_mut(d)) {
        let inv = 1.0 / (row_in.iter().map(|v| v * v).sum::<f32>() + eps).sqrt();
        for (o, i) in row_out.iter_mut().zip(row_in) {
            *o = i * inv;
        }
    }
    out
}

/// `y = x @ w^T` for `x: [m, k]`, `w: [n, k]` (torch `Linear` weight layout), no bias.
pub fn linear(x: &[f32], m: usize, k: usize, w: &[f32], n: usize) -> Vec<f32> {
    let mut out = vec![0.0f32; m * n];
    for row in 0..m {
        for col in 0..n {
            let mut acc = 0.0f32;
            for i in 0..k {
                acc += x[row * k + i] * w[col * k + i];
            }
            out[row * n + col] = acc;
        }
    }
    out
}

/// The bounded GLM forget gate, item 3 of the four NEW structural items.
///
/// `gate[t, h, d] = lower_bound * sigmoid(exp(a_log[h]) * (g_lowrank[t, h*D + d] + dt_bias[h*D + d]))`
///
/// `a_log` is per **head**; `dt_bias` is per **channel**. Getting that asymmetry wrong is the
/// single highest-risk line in the port, and it is why the two are separate arguments here.
pub fn bounded_gate(
    g_lowrank: &[f32],
    dt_bias: &[f32],
    a_log: &[f32],
    dims: KdaDims,
    lower_bound: f32,
) -> Vec<f32> {
    let (h_n, d) = (dims.heads, dims.head_dim);
    let mut out = vec![0.0f32; dims.tokens * h_n * d];
    for t in 0..dims.tokens {
        for h in 0..h_n {
            let decay = a_log[h].exp();
            for dd in 0..d {
                let ch = h * d + dd;
                let g = g_lowrank[t * h_n * d + ch] + dt_bias[ch];
                out[(t * h_n + h) * d + dd] = lower_bound * sigmoid(decay * g);
            }
        }
    }
    out
}

/// Unbounded Qwen-GDN gate, kept only so the microtest can show the two laws diverge.
///
/// Atlas's `compute_gdn_gates` stores `exp(g)`; this returns `g` itself, matching the KDA
/// convention where the exponential is taken inside the recurrence.
pub fn unbounded_gdn_gate(g_raw: f32, dt_bias: f32, a_log: f32) -> f32 {
    let x = g_raw + dt_bias;
    // HF: softplus(x) with the `x > 20` linear shortcut, beta = 1.
    let softplus = if x > 20.0 { x } else { (1.0 + x.exp()).ln() };
    -a_log.exp() * softplus
}

/// Decode formulation: one token at a time, carrying `state`.
///
/// ```text
/// S <- S * diag(exp(g_t))          decay along the KEY axis, per channel
/// delta <- (v_t - S^T k_t) * beta_t
/// S <- S + k_t (x) delta
/// o_t <- S^T q_t
/// ```
///
/// `q`/`k` are l2-normalised and `q` scaled by `1/sqrt(head_dim)` **inside** this function,
/// matching `use_qk_l2norm_in_kernel=True` on the HF side. Pass raw post-conv q/k/v.
pub fn kda_recurrent(
    q: &[f32],
    k: &[f32],
    v: &[f32],
    gate: &[f32],
    beta: &[f32],
    dims: KdaDims,
    state: &mut [f32],
) -> Vec<f32> {
    let d = dims.head_dim;
    let qn = l2norm_rows(q, d, 1e-6);
    let kn = l2norm_rows(k, d, 1e-6);
    kda_recurrent_prenorm(&qn, &kn, v, gate, beta, dims, state)
}

/// Same recurrence, but `q`/`k` are **already L2-normalised**.
///
/// This is Atlas's contract, where `causal_conv1d_update_l2norm` fuses conv + SiLU + L2
/// upstream, and it is what the `kda_recurrent` GPU kernel consumes. Keep it separate
/// rather than passing pre-normalised vectors into [`kda_recurrent`]: re-normalising an
/// already-unit vector is nearly a no-op in fp32 (it scales by `1/sqrt(1+1e-6)`), but on a
/// bf16-rounded vector — whose norm is off by ~0.4% — it silently RESTORES the norm the
/// rounding destroyed, which makes the reference disagree with the kernel by ~1e-5 and
/// looks exactly like a kernel bug.
pub fn kda_recurrent_prenorm(
    qn: &[f32],
    kn: &[f32],
    v: &[f32],
    gate: &[f32],
    beta: &[f32],
    dims: KdaDims,
    state: &mut [f32],
) -> Vec<f32> {
    let (h_n, d, t_n) = (dims.heads, dims.head_dim, dims.tokens);
    let scale = 1.0 / (d as f32).sqrt();

    let mut out = vec![0.0f32; t_n * h_n * d];
    let mut delta = vec![0.0f32; d];
    for t in 0..t_n {
        for h in 0..h_n {
            let base = (t * h_n + h) * d;
            let s = &mut state[h * d * d..(h + 1) * d * d];

            for kd in 0..d {
                let decay = gate[base + kd].exp();
                for vd in 0..d {
                    s[kd * d + vd] *= decay;
                }
            }
            let b = beta[t * h_n + h];
            for vd in 0..d {
                let mut kv = 0.0f32;
                for kd in 0..d {
                    kv += s[kd * d + vd] * kn[base + kd];
                }
                delta[vd] = (v[base + vd] - kv) * b;
            }
            for kd in 0..d {
                let kk = kn[base + kd];
                for vd in 0..d {
                    s[kd * d + vd] += kk * delta[vd];
                }
            }
            for vd in 0..d {
                let mut acc = 0.0f32;
                for kd in 0..d {
                    acc += s[kd * d + vd] * qn[base + kd] * scale;
                }
                out[base + vd] = acc;
            }
        }
    }
    out
}

/// Prefill formulation: chunked (WY-style) delta rule, mirroring HF's
/// `chunk_kimi_delta_attention`. Included at reference level because the per-channel decay mask is
/// exactly the part that a scalar-decay GDN kernel cannot express — proving it here is cheaper
/// than proving it in CUDA.
///
/// Handles `tokens % chunk != 0` by zero-padding, the same way HF does.
pub fn kda_chunked(
    q: &[f32],
    k: &[f32],
    v: &[f32],
    gate: &[f32],
    beta: &[f32],
    dims: KdaDims,
    chunk: usize,
    state: &mut [f32],
) -> Vec<f32> {
    let d = dims.head_dim;
    // l2norm runs on the REAL tokens, then the result is padded — HF pads after normalising.
    let qn = l2norm_rows(q, d, 1e-6);
    let kn = l2norm_rows(k, d, 1e-6);
    kda_chunked_prenorm(&qn, &kn, v, gate, beta, dims, chunk, state)
}

/// Same chunked formulation, but `q`/`k` are **already L2-normalised** — Atlas's contract, where
/// the conv path (fused on decode, `l2_norm_bf16` on prefill) has already normalised them.
///
/// Split out for the same reason as [`kda_recurrent_prenorm`]: re-normalising an already-unit
/// vector is nearly a no-op in fp32 but on a bf16-rounded vector it RESTORES the norm the
/// rounding destroyed, which makes the reference disagree with the kernel and looks exactly
/// like a kernel bug. Atlas's prefill L2 writes **bf16**, so this path is always the bf16 case.
#[allow(clippy::too_many_arguments)]
pub fn kda_chunked_prenorm(
    qn: &[f32],
    kn: &[f32],
    v: &[f32],
    gate: &[f32],
    beta: &[f32],
    dims: KdaDims,
    chunk: usize,
    state: &mut [f32],
) -> Vec<f32> {
    let (h_n, d, t_n) = (dims.heads, dims.head_dim, dims.tokens);
    let scale = 1.0 / (d as f32).sqrt();
    let pad = (chunk - t_n % chunk) % chunk;
    let tt = t_n + pad;
    let n_chunks = tt / chunk;

    let mut out = vec![0.0f32; t_n * h_n * d];
    for h in 0..h_n {
        // Per-head padded views. Pad rows are zero, which makes their k_beta/v_beta contributions
        // vanish and their cumulative decay a no-op.
        let mut qh = vec![0.0f32; tt * d];
        let mut kh = vec![0.0f32; tt * d];
        let mut vh = vec![0.0f32; tt * d];
        let mut gh = vec![0.0f32; tt * d];
        let mut bh = vec![0.0f32; tt];
        for t in 0..t_n {
            let src = (t * h_n + h) * d;
            for dd in 0..d {
                qh[t * d + dd] = qn[src + dd] * scale;
                kh[t * d + dd] = kn[src + dd];
                vh[t * d + dd] = v[src + dd];
                gh[t * d + dd] = gate[src + dd];
            }
            bh[t] = beta[t * h_n + h];
        }

        let s = &mut state[h * d * d..(h + 1) * d * d];
        for c in 0..n_chunks {
            let off = c * chunk;
            // Cumulative decay within the chunk, per channel.
            let mut gc = vec![0.0f32; chunk * d];
            for i in 0..chunk {
                for dd in 0..d {
                    let prev = if i == 0 { 0.0 } else { gc[(i - 1) * d + dd] };
                    gc[i * d + dd] = prev + gh[(off + i) * d + dd];
                }
            }
            let dmask = |i: usize, j: usize, dd: usize| (gc[i * d + dd] - gc[j * d + dd]).exp();

            // attn[i][j] = -sum_dd k_beta[i][dd] * k[j][dd] * exp(gc[i][dd] - gc[j][dd]), j < i.
            let mut attn = vec![0.0f32; chunk * chunk];
            for i in 0..chunk {
                for j in 0..i {
                    let mut acc = 0.0f32;
                    for dd in 0..d {
                        acc += kh[(off + i) * d + dd]
                            * bh[off + i]
                            * kh[(off + j) * d + dd]
                            * dmask(i, j, dd);
                    }
                    attn[i * chunk + j] = -acc;
                }
            }
            // Forward substitution: attn[i, :i] += attn[i, :i] @ attn[:i, :i].
            for i in 1..chunk {
                let row: Vec<f32> = (0..i).map(|j| attn[i * chunk + j]).collect();
                for j in 0..i {
                    let mut acc = 0.0f32;
                    for (m, r) in row.iter().enumerate() {
                        acc += r * attn[m * chunk + j];
                    }
                    attn[i * chunk + j] = row[j] + acc;
                }
            }
            for i in 0..chunk {
                attn[i * chunk + i] = 1.0;
            }

            // value = attn @ v_beta ;  k_cumdecay = attn @ (k_beta * exp(gc))
            let mut value = vec![0.0f32; chunk * d];
            let mut k_cumdecay = vec![0.0f32; chunk * d];
            for i in 0..chunk {
                for dd in 0..d {
                    let (mut av, mut ak) = (0.0f32, 0.0f32);
                    for j in 0..chunk {
                        let a = attn[i * chunk + j];
                        av += a * vh[(off + j) * d + dd] * bh[off + j];
                        ak += a * kh[(off + j) * d + dd] * bh[off + j] * gc[j * d + dd].exp();
                    }
                    value[i * d + dd] = av;
                    k_cumdecay[i * d + dd] = ak;
                }
            }

            // v_new = value - k_cumdecay @ S
            let mut v_new = vec![0.0f32; chunk * d];
            for i in 0..chunk {
                for vd in 0..d {
                    let mut vp = 0.0f32;
                    for kd in 0..d {
                        vp += k_cumdecay[i * d + kd] * s[kd * d + vd];
                    }
                    v_new[i * d + vd] = value[i * d + vd] - vp;
                }
            }

            // out = (q * exp(gc)) @ S  +  attn_intra @ v_new
            for i in 0..chunk {
                if off + i >= t_n {
                    continue;
                }
                let dst = ((off + i) * h_n + h) * d;
                for vd in 0..d {
                    let mut acc = 0.0f32;
                    for kd in 0..d {
                        acc += qh[(off + i) * d + kd] * gc[i * d + kd].exp() * s[kd * d + vd];
                    }
                    out[dst + vd] = acc;
                }
                for j in 0..=i {
                    let mut intra = 0.0f32;
                    for dd in 0..d {
                        intra += qh[(off + i) * d + dd] * kh[(off + j) * d + dd] * dmask(i, j, dd);
                    }
                    for vd in 0..d {
                        out[dst + vd] += intra * v_new[j * d + vd];
                    }
                }
            }

            // S <- S * exp(gc_last) + sum_i k[i] * exp(gc_last - gc[i]) (x) v_new[i]
            let last = (chunk - 1) * d;
            for kd in 0..d {
                let gl = gc[last + kd];
                for vd in 0..d {
                    s[kd * d + vd] *= gl.exp();
                }
                for i in 0..chunk {
                    let w = kh[(off + i) * d + kd] * (gl - gc[i * d + kd]).exp();
                    for vd in 0..d {
                        s[kd * d + vd] += w * v_new[i * d + vd];
                    }
                }
            }
        }
    }
    out
}

/// Gated RMSNorm over the trailing `d`, strict FP32, sigmoid gate.
///
/// `eps` is `rms_norm_eps` from the checkpoint config (1e-5 for GLM-5.3-Flash). vLLM never passes
/// it and relies on `FusedRMSNormGated`'s default, which coincides here — do not rely on that.
pub fn rms_norm_gated(x: &[f32], weight: &[f32], gate: &[f32], d: usize, eps: f32) -> Vec<f32> {
    let mut out = vec![0.0f32; x.len()];
    for i in (0..x.len()).step_by(d) {
        let var = x[i..i + d].iter().map(|v| v * v).sum::<f32>() / d as f32;
        let inv = 1.0 / (var + eps).sqrt();
        for dd in 0..d {
            out[i + dd] = x[i + dd] * inv * weight[dd] * sigmoid(gate[i + dd]);
        }
    }
    out
}

/// Weights for one KDA layer, reference-only. Torch `Linear` layout `[out, in]` throughout.
///
/// Note what is **absent** relative to Atlas's Qwen GDN: there is no `Z` tensor. The output gate is
/// `g_a`/`g_b`, and the decay source is `f_a`/`f_b`. Neither has a slot in Atlas's fused
/// `QKVZ` + `BA` layout, so weight binding is not a rename map.
pub struct KdaWeights<'a> {
    pub w_f_a: &'a [f32],
    pub w_f_b: &'a [f32],
    pub dt_bias: &'a [f32],
    pub a_log: &'a [f32],
    pub w_b: &'a [f32],
    pub w_g_a: &'a [f32],
    pub w_g_b: &'a [f32],
    pub o_norm_w: &'a [f32],
    pub w_o: &'a [f32],
}

/// End-to-end reference for one KDA layer, from post-conv q/k/v to `o_proj` output.
///
/// The short conv is intentionally excluded: it is REUSE against Atlas's existing fused
/// conv+SiLU+L2 kernel, and folding it in would couple two independent checks.
#[allow(clippy::too_many_arguments)]
pub fn kda_reference_layer(
    hidden: &[f32],
    q: &[f32],
    k: &[f32],
    v: &[f32],
    w: &KdaWeights<'_>,
    dims: KdaDims,
    lower_bound: f32,
    rms_eps: f32,
    state: &mut [f32],
) -> Vec<f32> {
    let (t_n, h_n, d, hid) = (dims.tokens, dims.heads, dims.head_dim, dims.hidden);

    let f_a = linear(hidden, t_n, hid, w.w_f_a, d);
    let g_lowrank = linear(&f_a, t_n, d, w.w_f_b, h_n * d);
    let gate = bounded_gate(&g_lowrank, w.dt_bias, w.a_log, dims, lower_bound);

    let beta: Vec<f32> = linear(hidden, t_n, hid, w.w_b, h_n)
        .iter()
        .map(|x| sigmoid(*x))
        .collect();

    let core = kda_recurrent(q, k, v, &gate, &beta, dims, state);

    let g_a = linear(hidden, t_n, hid, w.w_g_a, d);
    let out_gate = linear(&g_a, t_n, d, w.w_g_b, h_n * d);
    let normed = rms_norm_gated(&core, w.o_norm_w, &out_gate, d, rms_eps);

    linear(&normed, t_n, h_n * d, w.w_o, hid)
}

#[cfg(test)]
mod tests;