Module mla 

MLA (multi-head latent attention) CPU f32 reference — GLM-5.2 bring-up lane increment 1. Naive vs absorbed decode forms + NORM/NEOX rope permutation, unit-tested; the permanent oracle for the MLA kernel family (research/mla-bringup-20260801/DESIGN.md). No CUDA deps. MLA (multi-head latent attention, DeepSeek lineage / GLM-5 “MLA-256”) — CPU f32 reference.

Increment 1 of the GLM-5.2 bring-up lane (research/mla-bringup-20260801/DESIGN.md). This module pins the decode-path math BEFORE any kernel work: both the naive form (decompress the latent cache to per-head K/V, then attend — vLLM “forward_mha”) and the absorbed form (fold W_UK into the query, attend in latent space as MQA, decompress the output through W_UV — vLLM “forward_mqa”, llama.cpp glm-dsa.cpp). The unit tests prove the two forms agree to f32 tolerance on random inputs across shapes (t=1 decode and small causal prefill), including full GLM-5.2 dims (64 heads, nope 192, rope 64, v 256, rank 512).

Also pinned here: the interleaved (“NORM”, rope_interleave: true) vs NEOX rope pairing and the load-time permutation that maps one onto the other (DESIGN.md §1.4) — memra only ships a NEOX kernel, GLM-5.2 needs NORM, and the permutation trick lets the existing kernel serve.

Everything is plain CPU f32, no CUDA, no engine deps: this is the permanent oracle for the MLA kernel family’s maxdiff gates.

Structs

MlaDims

MLA head geometry. GLM-5.2: n_head=64, d_nope=192, d_rope=64, d_v=256, kv_rank=512.

MlaInputs

Inputs shared by both forms. Rope is already applied to q_pe/k_pe (it happens upstream of the attention core and is identical in both forms). c_kv is already RMS-normed.

Functions

mla_attend_absorbed

Absorbed form (decode form): q̃_h = w_uk[h]ᵀ·q_nope_h (rank-space, kv_rank wide), scores are MQA dots against the raw latent rows [c_kv | k_pe] (kv_rank + d_rope wide), the attention output is accumulated in latent space (kv_rank wide) and decompressed once through w_uv. Identical result to mla_attend_naive by associativity + linearity (DESIGN.md §1.3).

mla_attend_naive

Naive form: decompress k_nope/v per head from the latent cache, attend at qk dim d_nope+d_rope, output [t_q][n_head][d_v]. Quadratic decompression cost — prefill-only shape in production; here it is the independent oracle.

norm_to_neox_perm

The load-time permutation: source (interleaved-layout) index -> NEOX-layout index. pi(2j) = j, pi(2j+1) = j + n_dims/2. Applied to the rope rows of wq_b / wkv_a_mqa at load, it makes the existing NEOX kernel compute exactly the interleaved rotation (dot-product consumers only — which is all of them).

rope_interleaved

Interleaved (“NORM”) rope over the first n_dims of x: pair (x[2j], x[2j+1]) rotated by theta_j = pos * base^(-2j/n_dims). Matches ggml GGML_ROPE_TYPE_NORM / HF interleaved.

rope_neox

NEOX rope over the first n_dims of x: pair (x[j], x[j+half]) rotated by the same theta_j sequence. Matches memra’s rope_neox_f32 (kernels.cu) angle recurrence.