Module mla 

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.

Structs

MlaAttnWeights

Functions

mla_forward_token

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.