ferrox_models/

mla.rs

//! DeepSeek-style Multi-head Latent Attention (MLA): low-rank Q/KV
//! compression, with an optional sigmoid output gate (Kimi K3's real
//! addition) and an optional RoPE rotation of the decoupled `q_rot`/
//! `k_rot` slices (`MlaConfig::rope`; GLM-5.2's real addition -- see
//! below). Transcribed directly from real reference code, not guessed
//! or derived by analogy:
//!
//! 1. **Kimi K3** (`moonshotai/Kimi-K3`'s `modeling_kimi_linear.py`,
//!    `KimiMLAAttention.forward`, fetched live from the model repo): no
//!    rotary embedding is actually applied. `q_rot`/`k_rot` are named
//!    for the historical DeepSeek "rope part" split, but the real
//!    module asserts `self.use_nope` and never calls a rotary
//!    embedding function in `forward()` — the "rot" slice is just
//!    extra head-dim content, never position-rotated. Represented here
//!    as `MlaConfig::rope: None`.
//! 2. **GLM-5.2** (`zai-org/GLM-5.2`'s real `config.json`, confirmed
//!    against llama.cpp PR #25407's `src/models/glm-dsa.cpp`) DOES
//!    rotate its decoupled `q_rot`/`k_rot` slices, with the interleaved
//!    convention (`rope_interleave: true`,
//!    `ferrox_core::attention::apply_rope_interleaved`) — the opposite
//!    of the natural-but-wrong assumption the Kimi K3 module doc above
//!    warns against, for a *different* real architecture. Represented
//!    here as `MlaConfig::rope: Some(MlaRopeConfig { theta })`. `k_rot`
//!    is MQA-style (one shared vector per position, broadcast to every
//!    head — see point 4) so it's rotated once, before broadcasting;
//!    rotating a shared vector once then copying it into every head
//!    is exactly equivalent to rotating each head's copy separately,
//!    since RoPE's rotation angle depends only on position, not on the
//!    vector's per-head value.
//! 3. `kv_b_proj` expands to `num_heads * (...)`, not
//!    `num_key_value_heads * (...)`, despite `num_key_value_heads`/
//!    `num_key_value_groups` being computed in Kimi K3's real
//!    `__init__` — they go unused in `forward()`. Every query head gets
//!    its own decompressed K/V; there is no GQA-style grouping layered
//!    on top of the latent compression, which is why this module uses
//!    `ferrox_core::attention::causal_mla_attention` rather than
//!    `causal_gqa_attention`.
//!
//! When `rope` is `None` (Kimi K3's real path), a further simplification
//! applies, not present in the reference code's literal structure but
//! mathematically identical to it: the real `forward()` splits
//! `q_b_proj`'s output into `q_pass`/`q_rot` and immediately
//! re-concatenates them in the same order to form `query_states`. Since
//! nothing is inserted between the split and the concat (no rotation),
//! that round-trip is a no-op — `concat(x[..a], x[a..]) == x` — so this
//! implementation uses `q_b_proj`'s raw output directly as the query in
//! that case. When `rope` is `Some` (GLM-5.2's real path), the split is
//! no longer a no-op (rotation happens in between), so the `q_rot`
//! slice is rotated in place before attention runs.
//!
//! Not yet wired into `Decoder`'s forward pass (`AttentionKind` doesn't
//! dispatch to this yet) or into `ferrox_core::cache::KvCache` (which
//! assumes K and V share one `head_dim`, whereas MLA's K head dim
//! `qk_nope_head_dim + qk_rope_head_dim` and V head dim `v_head_dim`
//! generally differ) — both are handled by `kimi_decoder` (Kimi K3;
//! `rope: None`) and `glm_dsa`/`glm52_decoder` (GLM-5.2; `rope: Some`),
//! the dedicated decoders that consume this module. Tested here against
//! synthetic weights, cross-validated against independent Python
//! transcriptions of the same real reference algorithms for both
//! rope-disabled and rope-enabled paths.

use ferrox_core::attention::{apply_rope_interleaved, causal_mla_attention};
use ferrox_core::matmul::rms_norm;
use ferrox_core::weight_matrix::WeightMatrix;

use crate::config::MlaConfig;

pub struct MlaAttnWeights {
    pub q_a_proj: WeightMatrix,           // [q_lora_rank, hidden_dim]
    pub q_a_layernorm: Vec<f32>,          // [q_lora_rank]
    pub q_b_proj: WeightMatrix,           // [n_heads*q_head_dim, q_lora_rank]
    pub kv_a_proj_with_mqa: WeightMatrix, // [kv_lora_rank+qk_rope_head_dim, hidden_dim]
    pub kv_a_layernorm: Vec<f32>,         // [kv_lora_rank]
    pub kv_b_proj: WeightMatrix,          // [n_heads*(qk_nope_head_dim+v_head_dim), kv_lora_rank]
    pub o_proj: WeightMatrix,             // [hidden_dim, n_heads*v_head_dim]
    /// Present iff `MlaConfig::use_output_gate`.
    pub g_proj: Option<WeightMatrix>, // [n_heads*v_head_dim, hidden_dim]
}

/// One decode step. `k_cache`/`v_cache` are growable, caller-owned
/// buffers in `[seq_len_so_far, n_heads, head_dim]` layout (head_dim =
/// `qk_nope_head_dim + qk_rope_head_dim` for `k`, `v_head_dim` for `v`)
/// — plain `Vec<f32>`, not yet `ferrox_core::cache::KvCache` (see module
/// doc comment). This function appends the current position's K/V to
/// both before running attention over every position pushed so far.
#[allow(clippy::too_many_arguments)]
pub fn mla_forward_token(
    weights: &MlaAttnWeights,
    cfg: &MlaConfig,
    hidden: &[f32],
    rms_norm_eps: f32,
    k_cache: &mut Vec<f32>,
    v_cache: &mut Vec<f32>,
) -> Vec<f32> {
    let q_head_dim = cfg.qk_nope_head_dim + cfg.qk_rope_head_dim;
    // Position of the token being processed this call = how many
    // positions this layer's cache already holds, before this call
    // appends one more -- the same implicit convention `seq_len` below
    // already relies on (cache length as a proxy for "how many tokens
    // this layer has processed so far").
    let pos = k_cache.len() / (cfg.num_heads * q_head_dim);

    let q_a = weights.q_a_proj.apply(hidden);
    let q_a_normed = rms_norm(&q_a, &weights.q_a_layernorm, rms_norm_eps);
    // Without rope: `query_states` == raw `q_b_proj` output; see module
    // doc comment for why the reference's split+re-concat round-trip is
    // skipped in that case. With rope (GLM-5.2): the split is no longer
    // a no-op, so `q_rot` is rotated in place per head before use.
    let mut query = weights.q_b_proj.apply(&q_a_normed); // [n_heads*q_head_dim]
    if let Some(rope) = &cfg.rope {
        for h in 0..cfg.num_heads {
            let q_rot_h = &mut query[h * q_head_dim + cfg.qk_nope_head_dim..(h + 1) * q_head_dim];
            apply_rope_interleaved(q_rot_h, pos, rope.theta);
        }
    }

    let compressed_kv = weights.kv_a_proj_with_mqa.apply(hidden);
    let (k_pass_c, k_rot_raw) = compressed_kv.split_at(cfg.kv_lora_rank);
    // `k_rot` is MQA-style: one shared vector broadcast to every head,
    // not a per-head projection (real `kv_a_proj_with_mqa` name says so
    // directly, and the reference `.expand(...)`s it across heads) --
    // so with rope enabled, it's rotated once here before broadcasting
    // (see module doc comment point 2 for why that's equivalent to
    // rotating each head's copy separately).
    let mut k_rot = k_rot_raw.to_vec();
    if let Some(rope) = &cfg.rope {
        apply_rope_interleaved(&mut k_rot, pos, rope.theta);
    }
    let k_pass_c_normed = rms_norm(k_pass_c, &weights.kv_a_layernorm, rms_norm_eps);
    let k_pass_full = weights.kv_b_proj.apply(&k_pass_c_normed); // [n_heads*(qk_nope_head_dim+v_head_dim)]

    let mut key_step = vec![0f32; cfg.num_heads * q_head_dim];
    let mut value_step = vec![0f32; cfg.num_heads * cfg.v_head_dim];
    let kpf_stride = cfg.qk_nope_head_dim + cfg.v_head_dim;
    for h in 0..cfg.num_heads {
        let k_pass = &k_pass_full[h * kpf_stride..h * kpf_stride + cfg.qk_nope_head_dim];
        let v_h = &k_pass_full[h * kpf_stride + cfg.qk_nope_head_dim..(h + 1) * kpf_stride];

        let key_h = &mut key_step[h * q_head_dim..(h + 1) * q_head_dim];
        key_h[..cfg.qk_nope_head_dim].copy_from_slice(k_pass);
        key_h[cfg.qk_nope_head_dim..].copy_from_slice(&k_rot);

        value_step[h * cfg.v_head_dim..(h + 1) * cfg.v_head_dim].copy_from_slice(v_h);
    }

    k_cache.extend_from_slice(&key_step);
    v_cache.extend_from_slice(&value_step);
    let seq_len = k_cache.len() / (cfg.num_heads * q_head_dim);

    let attn_out = causal_mla_attention(
        &query,
        k_cache,
        v_cache,
        cfg.num_heads,
        q_head_dim,
        cfg.v_head_dim,
        seq_len,
    );

    let gated = match &weights.g_proj {
        Some(g_proj) => {
            let g = g_proj.apply(hidden);
            attn_out
                .iter()
                .zip(g.iter())
                .map(|(a, g)| a * (1.0 / (1.0 + (-g).exp())))
                .collect::<Vec<f32>>()
        }
        None => attn_out,
    };

    weights.o_proj.apply(&gated)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::MlaRopeConfig;
    use ferrox_core::tensor::Tensor;

    const HIDDEN_SIZE: usize = 8;
    const NUM_HEADS: usize = 2;
    const QK_NOPE_HEAD_DIM: usize = 3;
    const QK_ROPE_HEAD_DIM: usize = 2;
    const KV_LORA_RANK: usize = 4;
    const Q_LORA_RANK: usize = 6;
    const V_HEAD_DIM: usize = 3;
    const EPS: f32 = 1e-5;

    fn wm(data: &[f32], rows: usize, cols: usize) -> WeightMatrix {
        assert_eq!(data.len(), rows * cols);
        WeightMatrix::F32(Tensor::new(data.to_vec(), vec![rows, cols]))
    }

    fn cfg() -> MlaConfig {
        MlaConfig {
            num_heads: NUM_HEADS,
            q_lora_rank: Q_LORA_RANK,
            kv_lora_rank: KV_LORA_RANK,
            qk_nope_head_dim: QK_NOPE_HEAD_DIM,
            qk_rope_head_dim: QK_ROPE_HEAD_DIM,
            v_head_dim: V_HEAD_DIM,
            use_output_gate: true,
            rope: None,
        }
    }

    // Generated by an independent Python reference -- do not hand-edit.
    const MLA_Q_A_PROJ: [f32; 48] = [
        -0.237937, 0.0721714, -0.568898, 0.418732, 0.191488, -0.0876142, -0.0935848, 0.0911506,
        -0.0802981, -0.0677727, 0.21602, 0.154412, -0.0192384, -0.025643, 0.0482749, -0.184206,
        -0.121125, 0.164478, -0.0391448, -0.412328, -0.143184, 0.196986, -0.0696848, -0.0446198,
        0.192551, 0.547383, -0.213957, 0.404462, -0.369004, 0.0524933, -0.350859, 0.405437,
        0.250177, 0.341315, -0.26566, 0.205367, -0.155704, -0.137216, 0.151961, 0.263015,
        0.0613261, -0.188396, -0.247745, 0.433295, 0.178184, 0.215918, 0.655046, -0.244759,
    ];
    const MLA_Q_A_LAYERNORM_W: [f32; 6] = [1.25595, 1.31509, 1.16184, 1.08271, 0.933618, 1.09945];
    const MLA_Q_B_PROJ: [f32; 60] = [
        -0.132808,
        -0.00649524,
        -0.08713,
        0.085149,
        0.386423,
        -0.166675,
        -0.295621,
        -0.300887,
        -0.290483,
        -0.429331,
        -0.27388,
        0.38798,
        -0.177994,
        0.0771323,
        -0.365068,
        0.0508952,
        -0.522234,
        -0.209627,
        0.676362,
        -0.174891,
        0.335993,
        0.136516,
        -0.0458957,
        -0.195632,
        0.386073,
        -0.053215,
        0.458227,
        -0.215724,
        0.0172017,
        0.13965,
        0.111948,
        -0.37014,
        -0.199219,
        -0.0587941,
        -0.25611,
        0.203198,
        0.176406,
        -0.587125,
        -0.541576,
        -0.384469,
        0.0351779,
        0.609952,
        -0.114707,
        0.0751952,
        -0.318934,
        -0.314051,
        -0.587168,
        -0.00850327,
        0.284165,
        -0.107014,
        0.418936,
        0.0593565,
        -0.0109237,
        0.155176,
        0.146213,
        0.344342,
        -0.240586,
        -0.686415,
        0.034424,
        -0.183503,
    ];
    const MLA_KV_A_PROJ: [f32; 48] = [
        -0.00815316,
        0.49929,
        -0.330853,
        0.229439,
        0.28376,
        0.138221,
        0.335552,
        -0.137509,
        -0.204453,
        0.311695,
        0.21609,
        0.417113,
        0.0636869,
        0.487069,
        -0.0846241,
        -0.318099,
        -0.611127,
        -0.33042,
        0.245551,
        -0.439479,
        -0.135709,
        0.631975,
        0.254039,
        0.537295,
        -0.297469,
        -0.737993,
        0.454914,
        -0.425083,
        0.0272771,
        0.065128,
        -0.284669,
        0.46579,
        0.464801,
        0.165433,
        -0.0088472,
        0.057431,
        -0.328772,
        -0.0717191,
        -0.0314799,
        -0.25567,
        0.254764,
        -0.427944,
        -0.136026,
        -0.675863,
        0.149046,
        0.222782,
        0.186388,
        0.875426,
    ];
    const MLA_KV_A_LAYERNORM_W: [f32; 4] = [1.09398, 0.905801, 1.26635, 0.907968];
    const MLA_KV_B_PROJ: [f32; 48] = [
        0.236622, 0.222938, -0.263917, 0.418605, 0.106378, 0.171286, -0.20272, 0.39079, -0.258281,
        -0.533748, 0.427538, -0.128019, 0.013011, 0.576827, 0.340947, 0.189813, 0.320974,
        -0.218253, 0.225927, 0.332981, 0.226627, -0.0612404, -1.12422, 0.167191, 0.0344629,
        -0.110866, 0.23343, 0.0797109, 0.0611867, -0.445077, -0.437797, 0.111161, 0.0228171,
        0.0416223, 0.0498789, 0.123437, 0.0108197, 0.0473313, 0.0768343, 0.278823, -0.154407,
        0.235919, 0.190915, 0.0316206, -0.142534, 0.110742, 0.259192, -0.322493,
    ];
    const MLA_O_PROJ: [f32; 48] = [
        0.0186437, 0.0194232, -0.198361, -0.185786, -0.465344, -0.330938, -0.461243, -0.184972,
        -0.167996, -0.162386, 0.111553, -0.05627, -0.218674, -0.216565, -0.339453, -0.0588504,
        -0.0642164, 0.486887, 0.460246, 0.443371, 0.623773, -0.36512, 0.201895, 0.228316, 0.142865,
        0.638001, 0.761363, 0.29381, 0.195944, -0.336719, -0.336108, 0.210067, -0.0226898,
        -0.399168, -0.0372073, 0.108002, 0.280778, -0.0323038, -0.0807652, 0.00186235, 0.051895,
        -0.34894, 0.249993, 0.657098, 0.24426, -0.434459, 0.0283526, -0.529688,
    ];
    const MLA_G_PROJ: [f32; 48] = [
        -0.548764, 0.197775, 0.0232947, -0.206133, -0.626234, -0.38068, 0.0486588, 0.00956812,
        0.20691, 0.457018, -0.101777, 0.14069, -0.0413395, -0.148962, -0.0734372, 0.288226,
        -0.308526, 0.342312, 0.443186, 0.0107199, 0.0764755, -0.42354, 0.70816, 0.293666,
        0.0516828, 0.0172313, 0.135292, -0.195371, 0.0849785, -0.277073, 0.149196, -0.203464,
        0.268656, -0.0361107, 0.0806238, 0.888267, 0.24127, 0.0803401, 0.0133166, -0.311749,
        0.325266, 0.480702, 0.0193065, 0.503025, -0.0336833, -0.531916, -0.195972, -0.317098,
    ];

    const MLA_HIDDEN_0: [f32; 8] = [
        -0.13727, 0.31349, -0.464487, -0.0401226, 0.0996761, -0.366514, -0.45466, 0.0435914,
    ];
    const MLA_HIDDEN_1: [f32; 8] = [
        0.436185, -0.301614, -0.375696, -0.0706267, 0.144974, -0.271508, -0.167074, -0.115634,
    ];
    const MLA_HIDDEN_2: [f32; 8] = [
        -0.542896, -0.0887686, -0.280024, 0.791695, -0.258661, 0.228556, 0.189759, -0.559016,
    ];

    const MLA_GOLDEN_OUT_0: [f32; 8] = [
        0.01809, 0.174022, -0.240674, 0.212529, 0.493314, 0.191567, -0.184672, 0.118346,
    ];
    const MLA_GOLDEN_OUT_1: [f32; 8] = [
        0.0683158, 0.13817, -0.241981, 0.162395, 0.452095, 0.159199, -0.139237, 0.145646,
    ];
    const MLA_GOLDEN_OUT_2: [f32; 8] = [
        0.074717, 0.0909497, -0.208236, 0.166908, 0.390386, 0.129038, -0.124391, 0.12221,
    ];

    fn make_weights() -> MlaAttnWeights {
        MlaAttnWeights {
            q_a_proj: wm(&MLA_Q_A_PROJ, Q_LORA_RANK, HIDDEN_SIZE),
            q_a_layernorm: MLA_Q_A_LAYERNORM_W.to_vec(),
            q_b_proj: wm(
                &MLA_Q_B_PROJ,
                NUM_HEADS * (QK_NOPE_HEAD_DIM + QK_ROPE_HEAD_DIM),
                Q_LORA_RANK,
            ),
            kv_a_proj_with_mqa: wm(&MLA_KV_A_PROJ, KV_LORA_RANK + QK_ROPE_HEAD_DIM, HIDDEN_SIZE),
            kv_a_layernorm: MLA_KV_A_LAYERNORM_W.to_vec(),
            kv_b_proj: wm(
                &MLA_KV_B_PROJ,
                NUM_HEADS * (QK_NOPE_HEAD_DIM + V_HEAD_DIM),
                KV_LORA_RANK,
            ),
            o_proj: wm(&MLA_O_PROJ, HIDDEN_SIZE, NUM_HEADS * V_HEAD_DIM),
            g_proj: Some(wm(&MLA_G_PROJ, NUM_HEADS * V_HEAD_DIM, HIDDEN_SIZE)),
        }
    }

    #[test]
    fn matches_independent_python_reference_across_three_decode_steps() {
        // With `cfg().rope == None` (Kimi K3's real, rope-less path),
        // this also serves as the regression guard required before ever
        // touching `mla_forward_token` to add optional RoPE support --
        // these golden values and this cfg() are completely unchanged
        // from before that change, so a passing test here proves the
        // `rope: None` path is still byte-for-byte what it always was.
        let weights = make_weights();
        let cfg = cfg();
        assert!(
            cfg.rope.is_none(),
            "this test's whole point is pinning the rope-less path"
        );
        let mut k_cache = Vec::new();
        let mut v_cache = Vec::new();

        let hiddens = [&MLA_HIDDEN_0[..], &MLA_HIDDEN_1[..], &MLA_HIDDEN_2[..]];
        let goldens = [
            &MLA_GOLDEN_OUT_0[..],
            &MLA_GOLDEN_OUT_1[..],
            &MLA_GOLDEN_OUT_2[..],
        ];

        for (pos, (hidden, golden)) in hiddens.iter().zip(goldens.iter()).enumerate() {
            let out = mla_forward_token(&weights, &cfg, hidden, EPS, &mut k_cache, &mut v_cache);
            assert_eq!(out.len(), golden.len());
            for (i, (a, b)) in out.iter().zip(golden.iter()).enumerate() {
                assert!(
                    (a - b).abs() < 1e-3,
                    "position {pos} element {i}: rust={a} python={b}"
                );
            }
        }
    }

    #[test]
    fn without_output_gate_skips_the_sigmoid_multiply() {
        let mut weights = make_weights();
        weights.g_proj = None;
        let mut cfg = cfg();
        cfg.use_output_gate = false;
        let mut k_cache = Vec::new();
        let mut v_cache = Vec::new();

        let out = mla_forward_token(
            &weights,
            &cfg,
            &MLA_HIDDEN_0,
            EPS,
            &mut k_cache,
            &mut v_cache,
        );
        assert_eq!(out.len(), HIDDEN_SIZE);
        assert!(out.iter().all(|v| v.is_finite()));
        // Without gating this must differ from the golden (gated) output.
        assert!((out[0] - MLA_GOLDEN_OUT_0[0]).abs() > 1e-6);
    }

    // --- RoPE-enabled path (GLM-5.2's real convention), cross-validated
    // with an independent Python transcription applying interleaved RoPE
    // to q_rot/k_rot. No
    // output gate here (GLM-5.2's real tensor list has none), unlike
    // the Kimi K3 fixtures above.
    const ROPE_HIDDEN_SIZE: usize = 8;
    const ROPE_NUM_HEADS: usize = 2;
    const ROPE_QK_NOPE_HEAD_DIM: usize = 4;
    const ROPE_QK_ROPE_HEAD_DIM: usize = 4;
    const ROPE_KV_LORA_RANK: usize = 4;
    const ROPE_Q_LORA_RANK: usize = 6;
    const ROPE_V_HEAD_DIM: usize = 3;
    const ROPE_THETA: f32 = 10000.0;

    // Generated by an independent Python reference -- do not hand-edit.
    const MLA_ROPE_Q_A_PROJ: [f32; 48] = [
        -0.0282433, -0.200786, 0.229784, 0.162609, 0.161474, 0.190923, -0.28271, -0.259493,
        -0.299452, -0.445281, 0.438341, 0.223402, -0.114472, 0.108293, 0.138669, -0.0983588,
        -0.139929, -0.287077, 0.233259, -0.0924404, -0.48316, 0.236404, 0.487659, -0.22182,
        -0.425881, 0.474637, 0.224366, 0.227792, 0.155548, 0.139447, -0.0826791, -0.435336,
        0.438776, -0.328727, 0.415205, -0.280294, 0.25502, -0.264693, -0.374199, 0.626933,
        0.175746, 0.198436, -0.0882081, -0.131098, 0.416386, 0.0938253, -0.086288, -0.173465,
    ];
    const MLA_ROPE_Q_A_LAYERNORM_W: [f32; 6] =
        [0.983826, 0.879749, 1.11838, 1.05149, 1.1818, 0.834347];
    const MLA_ROPE_Q_B_PROJ: [f32; 96] = [
        0.291994,
        -0.728247,
        -0.316577,
        0.0813452,
        -0.266261,
        0.389674,
        0.0833719,
        0.351552,
        0.356938,
        0.350334,
        -0.0473987,
        -0.0266404,
        -0.169264,
        0.0701104,
        0.0207743,
        0.44759,
        0.372409,
        0.283663,
        0.161893,
        0.0691206,
        0.164776,
        -0.159844,
        0.244357,
        0.254148,
        0.781266,
        -0.0410461,
        0.00448047,
        0.167438,
        0.134256,
        -0.117364,
        0.613949,
        -0.207111,
        0.42746,
        0.45351,
        0.237126,
        -0.50974,
        0.328859,
        -0.250491,
        -0.356138,
        0.122879,
        0.254109,
        0.120117,
        -0.244618,
        0.090442,
        0.572282,
        -0.175117,
        0.150304,
        0.127176,
        -0.230927,
        -0.181049,
        0.0503238,
        -0.252932,
        -0.00813607,
        -0.169141,
        0.178562,
        -0.172518,
        -0.163208,
        -0.286795,
        0.358209,
        0.355661,
        -0.0321808,
        0.025399,
        -0.227651,
        -0.0153813,
        -0.0254572,
        -0.364581,
        -0.450488,
        0.155816,
        0.0033226,
        0.481021,
        -0.000260049,
        -0.230117,
        -0.0422523,
        0.269254,
        -0.225551,
        -0.265757,
        -0.192519,
        -0.300859,
        -0.152023,
        0.31445,
        -0.229592,
        -0.417754,
        -0.219984,
        0.0230321,
        0.162062,
        -0.162489,
        -0.504785,
        0.117479,
        -0.152083,
        0.203557,
        -0.232979,
        -0.537171,
        -0.131909,
        -0.0782392,
        0.187798,
        -0.364894,
    ];
    const MLA_ROPE_KV_A_PROJ: [f32; 64] = [
        -0.0108888,
        0.122853,
        -0.388147,
        -0.320502,
        0.288834,
        0.0587081,
        0.0027565,
        -0.0303023,
        0.252204,
        0.103756,
        -0.563955,
        0.539224,
        -0.515732,
        0.475067,
        -0.179422,
        0.512039,
        0.564391,
        -0.309453,
        0.178157,
        0.4829,
        0.304922,
        -0.327155,
        0.235531,
        -0.223351,
        0.113775,
        -0.326219,
        0.129363,
        0.343847,
        -0.555633,
        -0.00371874,
        -0.480022,
        0.0622793,
        0.121396,
        0.902273,
        -0.271857,
        -0.0787809,
        -0.148056,
        -0.246381,
        -0.388923,
        -0.326308,
        0.754771,
        -0.188557,
        0.157124,
        -0.242718,
        -0.196856,
        0.168396,
        0.116464,
        0.406121,
        -0.0524445,
        -0.226537,
        -0.220791,
        -0.42747,
        -0.109609,
        0.327875,
        0.238249,
        0.262922,
        0.0603609,
        0.259383,
        -0.125942,
        0.0253563,
        -0.672037,
        -0.0822506,
        -0.313883,
        -0.079927,
    ];
    const MLA_ROPE_KV_A_LAYERNORM_W: [f32; 4] = [1.11964, 1.00794, 0.927425, 0.974059];
    const MLA_ROPE_KV_B_PROJ: [f32; 56] = [
        0.152348, 0.223828, 0.104805, 0.146991, 0.167048, 0.0596911, 0.13428, -0.165432, 0.185469,
        -0.438659, 0.0505363, -0.378193, -0.367254, 0.291605, 0.25605, -0.278039, 0.32183,
        0.077576, 0.657456, -0.210592, -0.241762, -0.113545, 0.305935, -0.142537, -0.131723,
        -0.0698308, -0.231153, 0.0832406, -0.184562, -0.395013, -0.434206, 0.643208, -0.0451786,
        0.016892, -0.503596, 0.38556, 0.0672434, -0.345241, 0.261026, 0.113356, 0.195666, 0.124068,
        0.169155, 0.0241996, 0.0460725, -0.199187, 0.525018, 0.704311, 0.214609, 0.155865,
        -0.13945, -0.361349, 0.200727, 0.669041, -0.35694, 0.405635,
    ];
    const MLA_ROPE_O_PROJ: [f32; 48] = [
        0.156636,
        0.435755,
        0.254836,
        -0.28038,
        -0.00686566,
        0.254093,
        0.13879,
        0.298608,
        -0.654407,
        0.544604,
        -0.40823,
        0.557235,
        -0.401607,
        0.0393622,
        -0.0108063,
        -0.425778,
        -0.0790213,
        0.183181,
        0.770074,
        0.431033,
        -0.191665,
        -0.321149,
        -0.243943,
        -0.0704616,
        0.180775,
        -0.216385,
        0.0824125,
        -0.320591,
        -0.182163,
        -0.0257085,
        -0.0184709,
        0.292862,
        -0.215734,
        0.652291,
        -0.0461593,
        0.249014,
        -0.205017,
        0.0634068,
        0.087137,
        0.529326,
        0.477227,
        0.171185,
        0.0539693,
        0.0189488,
        -0.138254,
        -0.173556,
        0.65771,
        -0.0616593,
    ];

    const MLA_ROPE_HIDDEN_0: [f32; 8] = [
        -0.145978, -0.0867699, 0.281822, -0.765789, -0.590168, -0.252274, -0.397543, 1.30624,
    ];
    const MLA_ROPE_HIDDEN_1: [f32; 8] = [
        0.0788124, -0.477519, -0.142939, 0.0694206, 0.639385, 0.515165, -0.118041, -0.755483,
    ];
    const MLA_ROPE_HIDDEN_2: [f32; 8] = [
        -0.0406029, 0.440107, 0.194938, 0.0549023, 0.270816, 0.624518, 0.0925645, 0.0192617,
    ];

    const MLA_ROPE_GOLDEN_OUT_0: [f32; 8] = [
        0.0871067, -0.172861, 0.884245, -1.2416, 0.0539017, -0.224481, 0.230219, -0.137034,
    ];
    const MLA_ROPE_GOLDEN_OUT_1: [f32; 8] = [
        -0.24112, -0.298445, -0.0592754, -0.296603, -0.0395015, -0.0695785, 0.0639684, -0.0166364,
    ];
    const MLA_ROPE_GOLDEN_OUT_2: [f32; 8] = [
        0.0310026, -0.604048, 0.119056, 0.10382, 0.137893, -0.456617, -0.176747, 0.347884,
    ];

    fn rope_cfg() -> MlaConfig {
        MlaConfig {
            num_heads: ROPE_NUM_HEADS,
            q_lora_rank: ROPE_Q_LORA_RANK,
            kv_lora_rank: ROPE_KV_LORA_RANK,
            qk_nope_head_dim: ROPE_QK_NOPE_HEAD_DIM,
            qk_rope_head_dim: ROPE_QK_ROPE_HEAD_DIM,
            v_head_dim: ROPE_V_HEAD_DIM,
            use_output_gate: false,
            rope: Some(MlaRopeConfig { theta: ROPE_THETA }),
        }
    }

    fn make_rope_weights() -> MlaAttnWeights {
        MlaAttnWeights {
            q_a_proj: wm(&MLA_ROPE_Q_A_PROJ, ROPE_Q_LORA_RANK, ROPE_HIDDEN_SIZE),
            q_a_layernorm: MLA_ROPE_Q_A_LAYERNORM_W.to_vec(),
            q_b_proj: wm(
                &MLA_ROPE_Q_B_PROJ,
                ROPE_NUM_HEADS * (ROPE_QK_NOPE_HEAD_DIM + ROPE_QK_ROPE_HEAD_DIM),
                ROPE_Q_LORA_RANK,
            ),
            kv_a_proj_with_mqa: wm(
                &MLA_ROPE_KV_A_PROJ,
                ROPE_KV_LORA_RANK + ROPE_QK_ROPE_HEAD_DIM,
                ROPE_HIDDEN_SIZE,
            ),
            kv_a_layernorm: MLA_ROPE_KV_A_LAYERNORM_W.to_vec(),
            kv_b_proj: wm(
                &MLA_ROPE_KV_B_PROJ,
                ROPE_NUM_HEADS * (ROPE_QK_NOPE_HEAD_DIM + ROPE_V_HEAD_DIM),
                ROPE_KV_LORA_RANK,
            ),
            o_proj: wm(
                &MLA_ROPE_O_PROJ,
                ROPE_HIDDEN_SIZE,
                ROPE_NUM_HEADS * ROPE_V_HEAD_DIM,
            ),
            g_proj: None,
        }
    }

    #[test]
    fn rope_enabled_matches_independent_python_reference_across_three_decode_steps() {
        let weights = make_rope_weights();
        let cfg = rope_cfg();
        let mut k_cache = Vec::new();
        let mut v_cache = Vec::new();

        let hiddens = [
            &MLA_ROPE_HIDDEN_0[..],
            &MLA_ROPE_HIDDEN_1[..],
            &MLA_ROPE_HIDDEN_2[..],
        ];
        let goldens = [
            &MLA_ROPE_GOLDEN_OUT_0[..],
            &MLA_ROPE_GOLDEN_OUT_1[..],
            &MLA_ROPE_GOLDEN_OUT_2[..],
        ];

        for (pos, (hidden, golden)) in hiddens.iter().zip(goldens.iter()).enumerate() {
            let out = mla_forward_token(&weights, &cfg, hidden, EPS, &mut k_cache, &mut v_cache);
            assert_eq!(out.len(), golden.len());
            for (i, (a, b)) in out.iter().zip(golden.iter()).enumerate() {
                assert!(
                    (a - b).abs() < 1e-3,
                    "position {pos} element {i}: rust={a} python={b}"
                );
            }
        }
    }

    #[test]
    fn rope_enabled_output_changes_with_position() {
        // A direct, position-dependence check independent of the golden
        // values above: feed the *same* hidden state at two different
        // positions (by pre-filling the cache with a dummy earlier
        // position first) and confirm the outputs differ -- RoPE is the
        // only thing in this function that makes output depend on
        // absolute position rather than just on content, so if this
        // ever failed it would mean RoPE silently stopped being applied.
        let weights = make_rope_weights();
        let cfg = rope_cfg();

        let mut k_cache_pos0 = Vec::new();
        let mut v_cache_pos0 = Vec::new();
        let out_pos0 = mla_forward_token(
            &weights,
            &cfg,
            &MLA_ROPE_HIDDEN_0,
            EPS,
            &mut k_cache_pos0,
            &mut v_cache_pos0,
        );

        // Prime the cache with one earlier (dummy) position so the next
        // call happens at position 1 instead of 0, then feed the exact
        // same hidden state as above.
        let mut k_cache_pos1 = Vec::new();
        let mut v_cache_pos1 = Vec::new();
        mla_forward_token(
            &weights,
            &cfg,
            &MLA_ROPE_HIDDEN_1,
            EPS,
            &mut k_cache_pos1,
            &mut v_cache_pos1,
        );
        let out_pos1 = mla_forward_token(
            &weights,
            &cfg,
            &MLA_ROPE_HIDDEN_0,
            EPS,
            &mut k_cache_pos1,
            &mut v_cache_pos1,
        );

        assert_eq!(out_pos0.len(), out_pos1.len());
        let differs = out_pos0
            .iter()
            .zip(out_pos1.iter())
            .any(|(a, b)| (a - b).abs() > 1e-4);
        assert!(
            differs,
            "identical hidden state at two different positions must produce \
             different output when RoPE is enabled"
        );
    }

    #[test]
    fn rope_disabled_config_is_unaffected_by_position_change() {
        // The mirror-image check: with `rope: None` (Kimi K3's real
        // path), the *only* thing that should make output vary across
        // calls is the growing KV cache/causal history -- feeding the
        // identical hidden state as the very first token in two
        // otherwise-empty caches must give byte-identical output
        // regardless of "which call this was," since there is no
        // position-dependent rotation at all.
        let weights = make_weights();
        let cfg = cfg();
        assert!(cfg.rope.is_none());

        let mut k_cache_a = Vec::new();
        let mut v_cache_a = Vec::new();
        let out_a = mla_forward_token(
            &weights,
            &cfg,
            &MLA_HIDDEN_0,
            EPS,
            &mut k_cache_a,
            &mut v_cache_a,
        );

        let mut k_cache_b = Vec::new();
        let mut v_cache_b = Vec::new();
        let out_b = mla_forward_token(
            &weights,
            &cfg,
            &MLA_HIDDEN_0,
            EPS,
            &mut k_cache_b,
            &mut v_cache_b,
        );

        for (a, b) in out_a.iter().zip(out_b.iter()) {
            assert_eq!(a.to_bits(), b.to_bits());
        }
    }
}