QK-Normed MLA

The Ling Team, with Ruoyu Sun's group at CUHK-Shenzhen, shows that QK normalization and Multi-head Latent Attention are compatible after all. Post-projection QK RMSNorm appears to need the fully projected key for every cached token, which is exactly what MLA avoids storing. The paper decomposes RMSNorm into a static affine weight, which is absorbed into MLA's query-side projection, and a dynamic RMS statistic, which collapses to one inverse-RMS scalar per token and KV group. The result is exactly equivalent to explicit QK RMSNorm in exact arithmetic and keeps MLA's latent decode path; the RoPE component stays materialized and locally normalized.

In 400M-parameter MLA models trained on 100B tokens, QK-Normed MLA holds lower training loss than QK-Clip (the MuonClip stabilizer from Kimi K2) for the whole run, lifts the 3-shot average across eight tasks from 44.75 to 46.33, and cuts LAMBADA perplexity from 16.28 to 14.18. Under a deliberately high learning rate QK-Clip diverges with NaN gradients at step 884 while the normed model keeps training. On a DeepSeek-V3-width layer on H800s, the scalar cache adds 2MB to a 72MB per-layer KV cache and 0.95–1.87% decode latency from 4k to 256k context. First author Yizhou Han holds joint CUHK-Shenzhen and Ant Group affiliations; Ant Group co-authors include Jun Zhou. The paper names no production model and stops at 400M, but Ling-3.0's hybrid stack is built on gated MLA layers, the setting this stabilizer is designed for.