QK-Normed MLA: QK normalization without full key caching

2026-06-15

AI Summary

This work addresses the incompatibility between query-key (QK) normalization and multi-head latent attention (MLA), where standard QK normalization typically requires caching full key vectors, undermining MLA’s efficient decoding advantage. The authors propose decomposing RMSNorm into a static affine weight and a dynamic scalar statistic: the former is absorbed into the query projection, while the latter reduces to a single scalar per token and KV group. This decomposition enables mathematically equivalent QK normalization without caching full keys. To the best of our knowledge, this is the first method to achieve full compatibility between QK normalization and the MLA architecture. Evaluated on a 400M-parameter model trained on 100B tokens, it significantly lowers training loss and improves downstream task accuracy compared to QK clipping. On an H800 GPU with 256k-context decoding, the approach incurs less than 2% latency overhead.

Abstract

Query-key (QK) normalization stabilizes attention by controlling the scale of queries and keys before the dot product, but is not immediately compatible with Multi-head Latent Attention (MLA). MLA achieves efficient decoding by caching low-dimensional latent states instead of full keys, whereas post-projection QK RMSNorm appears to require the fully projected key for every cached token. We show this apparent incompatibility is an implementation artifact, not an architectural constraint. RMSNorm decomposes into a static affine weight and a dynamic scalar RMS statistic. The static key-side weight can be absorbed into the MLA query-side projection; the dynamic key statistic reduces to one inverse-RMS scalar per token and KV group. The resulting formulation is exactly equivalent to explicit post-projection QK RMSNorm in exact arithmetic and preserves MLA's latent decode path. In our 400M runs trained for up to 100B tokens, QK-Normed MLA achieves lower training loss and better downstream accuracy than QK clipping, while H800 decode benchmarks show less than 2% latency overhead up to 256k context. These results make QK normalization a practical stabilization option for MLA models without requiring full-key caching.

Problem

• QK normalization
• Multi-head Latent Attention
• key caching
• efficient decoding
• attention stabilization

Innovation

• QK normalization
• Multi-head Latent Attention
• RMSNorm
• efficient decoding
• latent state caching