Zunhai Su, Yuxuan Sun, Jianchao Tan, Tao Zhang, Ruihan Hu, Yuchen Xie, Xunliang Cai, Ngai Wong
9 min
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.
Alex: But content affects both matching and what gets retrieved. Why would position, with only one job, sometimes do more damage?
Sam: The authors separate the routes. A content error can change attention weights, the values being combined, and the interaction between those changes. A positional error changes only the positional contribution to matching. But the strength of an error depends on where it lands, not just how many routes it has.
Alex: So a smaller error can matter more if it points in a direction the computation is sensitive to?
Sam: Their model makes that idea explicit. It separates average sensitivity, the allocation of error energy across sensitive coordinates, and directional coupling between coordinates. In controlled tests, rearranging errors while preserving their total energy changed their effect. That supports the explanation beyond simply observing that position was more vulnerable.
Alex: How does that diagnosis become a method? You still need a representation the model can use efficiently.
Sam: Here’s the central move: learn transformations that preserve the original computation before quantization, but make quantization less damaging. For content, the authors allow rotations and channel scaling. They compensate the operations consuming that content, so the uncompressed computation stays equivalent.
Alex: Position already involves rotations. Can they apply the same transformation there?
Sam: Not unrestricted rotations. A transformation applied before rotary encoding must respect its paired-coordinate structure. QuantMLA rotates within those pairs and ties the scaling within each pair. It compensates queries and keys reciprocally, preserving their full-precision matching scores.
Alex: Extra transformations could eat up the savings. Are these operations performed every time a token is generated?
Sam: The learned transformations are folded into existing model parameters offline. There are no learned transformations left to run during inference. That is part of the contribution: the allowed transformation space is shaped by both the architecture and the requirement for offline fusion.
Alex: What tells the learning procedure that a transformation is good? Is it still trying to reproduce the cache values?
Sam: No. For content, it learns to reconstruct the attention output from the frozen reference model. That captures both matching and aggregation errors. For position, it reconstructs the positional query–key scores, using the actual queries rather than checking each key coordinate independently.
Alex: Why use scores for position instead of the final output, as they do for content?
Sam: Those scores are the only route by which positional-cache errors affect attention. The authors prove a bound connecting positional-score error to output distortion, with content and values fixed. But they explicitly describe the score objective as an output-controlling surrogate, not an exact replacement for output reconstruction.
Alex: Let’s move from preserved computations to actual answers. How close does the compressed model stay?
Sam: They evaluate four MLA model families, using general-purpose and extraction tasks for some, and reasoning and code tasks for others. The main setting stores both cache paths at four bits. Their reference is BF16, a sixteen-bit numerical format.
Alex: Are the differences small across those suites, or just on the easier tasks?
Sam: On DeepSeek-V2-Lite, the general-suite average falls by about a quarter of a percentage point. Moonlight stays similarly close to its reference. The reasoning-and-code results have larger losses: LongCat’s average falls by about two percentage points, while GLM also loses some performance. “Close” does not mean unchanged on every task.
Alex: What are the comparisons? A weak baseline could make a learned method look better than it is.
Sam: They compare with simple round-to-nearest quantization and MLA adaptations of SmoothQuant and QuaRot. Those adaptations use scaling and fixed rotations, respectively, with the same quantizer as QuantMLA. They are not reproductions of the original full-model methods. The authors say no directly comparable prior joint-four-bit MLA method had been reported.
Alex: Does the final system rely only on learned transformations, or does it keep some history uncompressed?
Sam: It also protects the earliest tokens and a recent window in BF16. An ablation shows that the learned transformations improve accuracy before this protection is added. Protection then adds another gain. Historical tokens remain visible: the policy changes their storage precision, not whether attention can access them.
Alex: Then what does the memory saving actually buy? Is the reported compression counting those protected buffers?
Sam: Yes, along with metadata and redundant packed copies. At a context of about a hundred and twenty-eight thousand tokens, persistent cache storage is about three-point-six times smaller than BF16. Their native attention kernel unpacks and reconstructs values inside attention, without materializing a full-history high-precision cache.
Alex: The headline also promises much higher throughput. Is that because this attention kernel runs much faster?
Sam: No. Attention latency is near BF16 in the tested configurations. The roughly fivefold whole-job throughput gain comes from a serving workload under memory pressure. The compressed cache avoids preemption and recomputation that occur with the BF16 cache.
Alex: So we should not expect that speed gain whenever we change the cache format?
Sam: Correct. The workload just exceeded BF16’s available capacity. The comparison uses one compressed-cache run against a historical BF16 baseline, without a confidence interval across independent serving runs. I find the compression and accuracy evidence more broadly persuasive than that single throughput result. The authors also warn that deployment needs fresh safety and failure-mode checks.
Alex: Who should read the full document, and where should they start?
Sam: Quantization researchers should start with “Modeling Dual-Path Quantization Error,” then “Function-Aligned Learning Objectives.” Serving engineers should read “Serving Under Cache Pressure” and the native cache-layout details, especially what the timings include. For everyone else, remember this: preserve what compressed numbers do, not just what they look like.
Alex: Thanks for listening.