QuantMLA: Function-Aligned Dual-Path Quantization for Low-Bit MLA KV Caching

AI summary

Overview

Research area: Machine learning systems — low-bit quantization of the KV cache for Multi-Head Latent Attention (MLA) transformer models, plus a custom attention kernel.

Technical level: Advanced. The paper combines a matrix-level error analysis (Jacobians, Rayleigh quotients, an output-distortion bound) with quantization algorithm design and GPU kernel engineering.

Scope in one sentence: This paper analyzes why MLA's two cached paths — the content latent and the decoupled RoPE key — respond unequally to quantization error, then builds a dual-path INT4/INT2 quantization method and native low-bit attention backend around that analysis.

What This Paper Is About

MLA stores its KV cache in two parts: a shared low-rank content latent and a separate decoupled RoPE key cache. Existing systems (for example FlashMLA with DeepSeek-V3.2) compress only the content cache to FP8 and keep RoPE keys at BF16, so nobody has characterized what RoPE quantization actually does or compressed both paths jointly at low bit-widths. The paper sets out to explain the functional difference between the two paths' errors and then to exploit that difference to quantize both caches to INT4 (and the content cache to INT2) with minimal accuracy loss.

Key Contributions

1. A systematic analysis of MLA dual-path quantization. The authors characterize the distinct error routes of the content path (attention matching, value aggregation, and their interaction) and the RoPE path (positional logits only), and build an operator–error model that decomposes amplification into average operator sensitivity, error-energy allocation, and directional coupling.

2. Function-aligned dual-path quantization (QuantMLA). They derive path-specific transformation spaces that exactly preserve full-precision computation yet can be fused entirely into model weights offline (an orthogonal rotation plus positive channel scaling for the content path, and a RoPE-commuting block-diagonal rotation with per-frequency-pair scaling for the RoPE path), then learn within each space using a different objective: attention-output reconstruction for content, positional QK reconstruction for RoPE.

3. Joint low-bit MLA caching results. To the authors' knowledge, the first reported joint INT4 caching of both the content and RoPE caches, evaluated across four MLA model families; plus a C2R4 setting that compresses content to INT2 while keeping RoPE keys at INT4.

4. A native low-bit MLA execution backend. Built on FlashMLA's decode path, it fuses unpacking and dequantization into attention, keeps BF16 buffers only for four sink tokens and the most recent 128 tokens, and uses a single global softmax over all valid tokens.

Main Findings

• Matched cache error, unequal output damage. When each path's cache is perturbed to the same NMSE of 0.01, the RoPE-to-content response ratio across six models (DeepSeek-V2-Lite, DeepSeek-V3-Base, DeepSeek-R1, Kimi-K2-Instruct, LongCat-Flash-Lite, GLM-4.7-Flash) ranges from 1.941× to 6.864×. Model-level content/RoPE responses include 0.620/3.925 for DeepSeek-V2-Lite, 0.594/4.077 for LongCat-Flash-Lite, and 1.594/3.095 for Kimi-K2-Instruct.

• Cache reconstruction error is a poor proxy for functional damage. The analysis covers 285 attention layers (27, 61, 61, 61, 28 and 47 per model) and shows substantial layer-wise variation in both responses and their ratios.

• Content-path errors have three routes, RoPE-path errors have one. Content errors propagate through attention matching, value aggregation, and the interaction term; RoPE errors act only through the RoPE-induced component of the attention logits — and still dominate.

• Accuracy at C4R4 stays close to BF16 on general tasks. On the ten-task suite, QuantMLA is within 0.26 percentage points of BF16 on DeepSeek-V2-Lite (66.54 vs 66.80 average) and within 0.13 points on Moonlight-16B-A3B (73.44 vs 73.57). On DeepSeek-V2-Lite GSM8K it reaches 35.25% versus 16.98% (RTN), 24.56% (SmoothQuant†) and 30.71% (QuaRot†); Moonlight matches BF16 at 74.53%.

• Aggressive content compression remains usable. At C2R4, suite averages are 64.42% (DeepSeek-V2-Lite) and 72.70% (Moonlight-16B-A3B), beating the strongest same-precision baselines by 7.20 and 10.84 points.

• Reasoning and code benchmarks follow the same pattern. At C4R4, LongCat-Flash-Lite averages 62.01% (BF16: 64.26%) and GLM-4.7-Flash 53.75% (BF16: 55.32%), exceeding the strongest same-precision baselines by 4.82 and 2.70 points. LiveCodeBench on LongCat-Flash-Lite is 0.29 points below BF16; GLM-4.7-Flash GSM8K is 0.15 points below. At C2R4, GLM-4.7-Flash averages 51.43%, 6.22 points above the strongest same-precision baseline.

• Long-context retrieval degrades gracefully. On DeepSeek-V2-Lite RULER S-NIAH at C4R4 (500 examples per variant, 4K–32K tokens), QuantMLA averages 99.95% (S-NIAH-1), 99.60% (S-NIAH-2) and 87.20% (S-NIAH-3), the last exceeding the strongest quantized baseline by 7.40 points and beating all quantized baselines at every tested length.

• Both transformations and the mixed-precision policy contribute. On DeepSeek-V2-Lite MMLU, the content-path transformation adds 1.22 points (51.90 → 53.12) and the RoPE-path transformation adds 3.19 (→ 56.31) before any BF16 protection; the mixed-precision cache policy adds a further 1.64, reaching 57.95 versus BF16 at 57.90. The transformation-only configuration beats SmoothQuant† and QuaRot† by 2.29 and 0.65 points.

• Memory and speed gains. C4R4 yields 3.22× persistent KV-cache compression at 4K and 3.59× at 128K. Attention latency from 1K to 1M tokens at batch sizes 8 and 32 ranges from 7.42% lower to 4.29% higher than BF16, with differences below 0.5% from 64K onward.

• Serving under cache pressure. For five simultaneous requests with 128K input and 1K output tokens on eight GPUs, BF16 incurs 32 preemption events and C4R4 none; whole-job output throughput rises from 2.61 to 13.51 tokens/s (5.168×). The authors note this gain is specific to the evaluated cache-pressure workload.

Methodology in Plain English

The work proceeds in three layers.

Diagnose first. The researchers take one path's cache at a time, add a quantization-style error of fixed size (normalized to NMSE 0.01), and measure how much the attention output moves. Doing this across six models shows the RoPE path hurts output far more than the content path for the same cache error. They then write out the exact algebra of how each path's error reaches the output, and build a local model: the output damage per unit cache error equals a sensitivity term (how strongly the output reacts to cache changes), multiplied by how well the error aligns with sensitive directions, plus a cross-coordinate coupling term.

Then design around the diagnosis. Each path gets its own reversible "reshaping" of the cache, chosen so that the full-precision computation is unchanged and the reshaping can be folded into the model's weight matrices ahead of time — no extra work at inference. For the content latent, this is a rotation plus per-channel scaling; for the RoPE key, the reshaping must commute with RoPE's pairwise rotations, so it is restricted to 2×2 rotations and per-frequency-pair scaling. Inside these reshaped spaces, the two paths are trained with different targets: the content path minimizes distortion of the final attention output (because its errors hit both matching and aggregation), while the RoPE path minimizes error in the positional query–key scores, a choice justified by a proved bound showing output distortion is controlled by positional-score error.

Finally, engineer the execution. They build a low-bit MLA decode kernel where unpacking and dequantization happen inside the attention computation; packed data and metadata are buffered separately from reconstructed BF16 tiles so prefetching overlaps with computation, and a small number of sink and recent tokens are kept in BF16 while everything else is served from INT4. Calibration is light: pretrained weights are frozen, each path's transformation is optimized for 300 steps on 128 WikiText-2 sequences, and the best checkpoint is chosen on 32 held-out sequences. All quantized methods use the same per-token asymmetric affine quantizer with group size 64.

Why This Matters

Impact on research. The paper reframes KV-cache quantization as a question about function, not reconstruction error: two caches with identical reconstruction error can differ by up to 6.864× in output distortion. It also shows that architecture structure (shared latent vs. decoupled RoPE keys) dictates which transformations are legally fusable offline, giving a template for quantization work on other latent-cache architectures.

Real-world applications:

• Long-context serving of MLA-based assistants, where 3.59× cache compression at 128K tokens directly increases how many requests fit on the same GPU memory.
• High-throughput, memory-pressured inference, where eliminating preemption events (32 for BF16 versus none for C4R4) translates into 5.168× whole-job output throughput on the evaluated workload.
• Reasoning and code-generation services, since the accuracy results are reported on GPQA-Diamond, AIME25, MATH500, HumanEval and LiveCodeBench as well as general benchmarks.
• Retrieval over long documents, given the near-BF16 RULER S-NIAH-1/2 scores (99.95% and 99.60%) at INT4.

Industry relevance. The method targets models that are actually deployed — DeepSeek-V3, Kimi K2, LongCat-Flash and GLM — and is implemented on top of FlashMLA's decode path with vLLM scheduling, so it is aimed at production inference stacks rather than a synthetic setting.

Future Directions

• Extending the analysis beyond the tested families. The error study covers six models from four MLA families; whether the RoPE-to-content ratio range (1.941× to 6.864×) generalizes to other MLA variants, depths and modalities is not established.

• Closing the remaining INT2 accuracy gap. At C2R4, GLM-4.7-Flash reaches 51.43% versus 55.32% at BF16 with larger task-to-task variation, which the authors attribute to greater sensitivity to extreme content-cache compression — an open target for better content-path handling.

• Generalizing the cache-pressure throughput result. The 5.168× gain is explicitly reported as specific to the evaluated workload; mapping when it holds across workloads, request mixes and hardware is left open.

• Broadening safety and calibration evaluation. The ethics statement itself argues that low-bit caching can change model behavior even when aggregate benchmark accuracy is preserved, and calls for reassessing safety, calibration and failure modes under each exact model/precision/serving configuration.

Target Audience

Researchers and engineers working on LLM inference efficiency — particularly those implementing KV-cache quantization, MLA/DeepSeek-style attention, or custom attention kernels — will get the most from this paper. It is also relevant to systems teams evaluating whether low-bit caching is safe to deploy for a specific model. Readers need comfort with attention algebra, matrix decomposition and quantization terminology; the empirical tables in Section 4 are readable on their own for practitioners who want results without the derivation.

Authors’ abstract

Multi-Head Latent Attention (MLA) enables expressive multi-head attention with compact caches for its content and decoupled RoPE paths, yet cache memory still scales linearly with context length and batch size. In this work, we establish a systematic model of MLA's dual-path quantization errors, characterizing their distinct effects on attention-output distortion and explaining the pronounced amplification of RoPE-path errors. Guided by this analysis, we introduce QuantMLA, a function-aligned framework for low-bit dual-path quantization. We derive path-specific transformation spaces that preserve full-precision computation while remaining fully fusible into model parameters offline, eliminating online transformation overhead. Within these spaces, QuantMLA learns path-specific transformations with function-aligned objectives: attention-output reconstruction captures the content path's coupled matching and aggregation errors, while positional QK reconstruction preserves the RoPE-induced component of the attention logits and admits a theoretical bound on output distortion. Across four MLA model families, QuantMLA enables, to our knowledge, the first reported joint INT4 caching of the content and RoPE caches with minimal accuracy degradation. Further compressing the content cache to INT2 while retaining the RoPE key cache at INT4 maintains competitive performance on challenging reasoning and code benchmarks. We develop a native low-bit MLA attention kernel that integrates unpacking and dequantization directly into attention computation. The physical cache layout provides 3.59x compression at 128K context, while a cache-pressure serving workload achieves 5.168x higher whole-job output throughput than BF16. The code will be released upon acceptance.