RAQ-MLA: Role-Aware Quantization for Multi-Head Latent Attention Caches

Tianqi Chen; Xunchao Li; Linchuan Xie; Guanghua Yu; Jianchen Zhu; Peisong Wang; Jian Cheng

cs.AIOct 24, 2026

RAQ-MLA: Role-Aware Quantization for Multi-Head Latent Attention Caches

Authors: Tianqi Chen, Xunchao Li, Linchuan Xie, Guanghua Yu, Jianchen Zhu, Peisong Wang, Jian Cheng

Organizations: C2DL, Institute of Automation, Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences · Tencent

Abstract

Multi-Head Latent Attention (MLA) has been adopted by recent frontier large language models. Instead of per-head keys and values, MLA caches a compressed KV latent that is up-projected into keys and values at decode time, together with a narrow positional key that carries position information into the attention score. Existing KV-cache quantization methods are designed around the keys and values of standard attention and do not account for what MLA actually caches, which costs accuracy at low bit widths. We analyze the MLA cache both in activation distribution and in how quantization error propagates to the attention output, and show that the two cached tensors play different roles: the latent mainly supplies the values, whereas the positional key dominates the attention score. We therefore propose RAQ-MLA (Role-Aware Quantization for Multi-Head Latent Attention Caches), which fits an invertible transform to each cached tensor against the error that tensor delivers to the attention output: attention-weighted value error for the latent and score perturbation for the positional key. Without training, RAQ-MLA quantizes the full MLA cache to about 2.5 bits per element, a 6.5x increase in theoretical KV-cache capacity, while remaining nearly lossless on reasoning and coding tasks across four MLA models.

Figures & tables

Figure 1: Cache capacity against accuracy on GLM-4.7-Flash.

Figure 2: Fixed-channel outliers appear in the decoupled positional key but not in the compressed KV latent. The positional key behaves like the Key cache of standard attention, and the compressed latent like the Value cache.

Panel of Figure 3 — Figure 3: The positional key dominates the attention score, while the latent mainly plays the value role. (a) Logit attribution to the outlier and normal channels of the latent $\mathbf{c}^{\mathrm{KV}}$ and the positional key $\mathbf{k}^{\mathrm{R}}$ , per layer. (b) Relative error at each stage under INT4, with the latent quantized per token and the positional key per channel.

Figure 4: RAQ-MLA at one MLA layer. $A_{\mathrm{c}}$ is applied to the latent after RMSNorm and $A_{\mathrm{R}}$ to the positional key after RoPE, so both cached tensors are quantized in a learned basis. On the read path $A_{\mathrm{c}}^{-1}$ is fused into $\mathbf{W}^{\mathrm{UK}}$ and $\mathbf{W}^{\mathrm{UV}}$ , while the positional key is compensated on the query side.

Panel of Figure 5 — Figure 5: Evaluation of the latent role metric. (a) At layer $23$ , the role-metric loss and the exact attention-output error it is meant to predict agree across $128$ random transforms (Pearson $r=0.992$ , Spearman $\rho=0.964$ ), and RAQ-MLA minimizes both. (b) Layer-wise attention-output error of each method on held-out validation prompts; shading marks the reduction from the better baseline to RAQ-MLA.

Model	Method	BPE	AIME25	GPQA	MATH-500	GSM8K	HumanE	$\Delta$
	BF16	16.00	93.33	77.27	97.40	95.22	93.29	—
	KVarN INT3	4.67	90.00	73.23	98.20	95.60	92.07	$-1.48$
	QuaRot INT2	3.78	86.67	69.19	98.20	94.09	86.59	$-4.35$
	OSCAR INT2	3.78	87.78	72.06	97.40	95.75	94.31	$-1.84$
GLM-4.7	RAQ-MLA	2.47	92.22	75.25	97.67	96.08	94.92	$\mathbf{-0.08}$
	BF16	16.00	90.00	86.87	97.80	97.50	99.39	—

Table 1: Reasoning and coding tasks. BPE is effective bits per cached element, and

\Delta

is the mean change from BF16 over the five tasks.

Method	BPE	4K	8K	16K	32K	64K	128K	Avg.
BF16	16.00	100.00	99.75	99.31	98.44	96.97	87.22	96.95
OSCAR	3.78	99.56	98.64	98.34	96.01	90.39	78.65	93.60
RAQ-MLA	2.66	99.38	99.06	97.96	95.74	90.69	79.13	93.66

Table 2: RULER-NIAH on GLM-4.7-Flash. Each cell averages eight subtasks of

100

samples.

Method	Kernel ( $\mu$ s/layer)	TPOT (ms)
FlashMLA	17.4	6.56
RAQ-MLA	12.6	5.97
Speedup	1.38 $\times$	1.10 $\times$

Table 4: Kernel latency and TPOT on H20 GPUs.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 7: Analysis of the RAQ-MLA latent fit on GLM-4.7-Flash. (a, b) Calibration objective relative to iteration

0

and Val attention-output error per round, mean

\pm

std over the

47

INT2 layers along each layer’s Val-selected damping trajectory, from Hadamard and OSCAR initializations. (c) Relative error of the latent at each stage under the INT4 quantizer of Figure 3(b) , with the fitted

A_{\mathrm{c}}

(solid) and without a transform (dashed).

Figure 8: The latent and positional key retain different channel structures throughout depth. Token–channel magnitude surfaces for the normalized compressed latent $\mathbf{c}^{\mathrm{KV}}$ (top) and post-RoPE positional key $\mathbf{k}^{\mathrm{R}}$ (bottom) at five approximately uniformly spaced layers of GLM-4.7-Flash. Each panel uses its own $99$ th-percentile vertical scale, so the comparison concerns structure rather than absolute magnitude. The latent remains diffuse across channels, whereas the positional key exhibits persistent fixed-channel bands from early to late layers.

Panel of Figure 9 — Figure 9: Value-readout errors are nearly uncorrelated across tokens, although attention reads nearby tokens together. (a–b) Layer $23$ of GLM-4.7-Flash, first $2048$ tokens of a calibration prompt, deployed INT2 quantizer with the fitted RAQ-MLA transform, visualized at $16$ -token resolution. (a) Attention Gram $S^{\top}S$ , normalized by its diagonal and averaged over heads. (b) Magnitude of the correlation $\rho_{t,t^{\prime}}$ between the value-readout errors at tokens $t$ and $t^{\prime}$ , pooled over heads. (c) Val attention-output error of the latent transform fitted with $M_{t}=I$ or with the role metric under the same fitting pipeline.

Figure 10: Evaluation of the positional-key role metric. (a) Score-role loss

\mathcal{D}_{\mathrm{R}}

against the exact attention-output error at layer

23

across

128

random transforms; Hadamard, OSCAR, and RAQ-MLA are excluded from the correlations. (b) Layer-wise attention-output error; baselines are shown on calibration prompts only.

	Latent $\mathbf{c}^{\mathrm{KV}}$	Positional key $\mathbf{k}^{\mathrm{R}}$
Quantized tensor	$\mathbf{c}^{\mathrm{KV}}_{t}A_{\mathrm{c}}$	$\mathbf{k}^{\mathrm{R}}_{t}A_{\mathrm{R}}$ (after RoPE)
Bits (levels)	$2$ ( $4$ )	$4$ ( $16$ )
Grouping	per token, groups of $128$ channels	per channel, within a cache page
Clipping	per token, $0.94$ quantile of $\|x\|$	none
Metadata	FP16 scale and minimum per group	FP16 scale and minimum per channel and page
Effective bits	$2+32/128=2.25$	$4+32/P$ ( $4.25$ at $P=128$ , $4.50$ at $P=64$ )

Appendix

Table 5: Quantizer configuration of RAQ-MLA.

P

is the cache-page size in tokens.

	GLM-4.7-Flash	GLM-5.2-FP8	Moonlight	DS-V2-Lite
max_model_len	$202{,}752$	$202{,}752$	$8{,}192$	$32{,}768$
max_tokens	$131{,}072$	$65{,}536$	$7{,}168$	$7{,}168$
Cache page (tokens)	$128$	$64$	$128$	$128$
Temperature	$1.0$	$1.0$	$1.0$	$1.0$
Top- $p$	$0.95$	$0.95$	$0.95$	$0.95$

Appendix

Table 6: Serving and sampling configuration for the reasoning and coding tasks. ARC-Challenge, PIQA, and RULER NIAH use greedy decoding instead.

Method	BPE	Single-QA	Multi-QA	Summ.	Few-shot	Synth.	Code	Avg
OSCAR	3.78	49.16	57.24	22.88	60.46	77.20	54.75	52.38
RAQ-MLA	2.66	49.07	56.83	22.80	61.84	76.30	52.44	52.19

Appendix

Table 7: LongBench-v1 macro average (%) over

21

tasks on GLM-4.7-Flash , thinking on.

Output tokens	FlashMLA (ms)	RAQ-MLA (ms)	Speedup
$1024$	$6.61$	$\mathbf{6.04}$	$1.09\times$
$2048$	$6.69$	$\mathbf{6.11}$	$1.09\times$
$4096$	$6.87$	$\mathbf{6.25}$	$1.10\times$
$8192$	$7.02$	$\mathbf{6.54}$	$1.07\times$

Appendix

Table 8: Decode TPOT across output lengths on GLM-4.7-Flash , with a

512

-token prompt and batch size

1

. Speedup is the FlashMLA TPOT divided by the RAQ-MLA TPOT.

Decode batch	Standalone ( $\mu$ s)	Fused ( $\mu$ s)	Share
$1$	$2.66$	$\mathbf{2.36}$	$9.9\%$
$2$	$2.68$	$\mathbf{1.72}$	$6.6\%$
$4$	$2.64$	$\mathbf{1.77}$	$5.5\%$
$8$	$2.70$	$\mathbf{1.81}$	$4.0\%$
$64$	$2.83$	$\mathbf{2.24}$	$1.0\%$

Appendix

Table 9: Per-layer cost of the latent transform on a single GLM-shaped MLA layer with

2

K cached tokens. Standalone applies

A_{\mathrm{c}}

as a separate GEMM before the cache write; fused applies it inside the write kernel. Share is the fused cost relative to the per-layer cache write plus attention.

LLM Quantization

Large language model quantization reduces numerical precision to lower memory use and accelerate training or inference.

See all 388 papers

QuantMLA: Function-Aligned Dual-Path Quantization for Low-Bit MLA KV Caching
Sep 29, 2026Zunhai Su, Yuxuan Sun, Jianchao Tan +5LLM Inference Acceleration Attention Quantization
The University of Hong Kong · Meituan LongCat Team · South China University of Technology +1
RoLaQ: Joint Rotation and Range Calibration for 2-Bit KV Caches in Multi-Head Latent Attention
Oct 14, 2026Shiyu Ji, Peng Ling, Chen Lin +4KV-Cache Quantization Low-Bit Quantization
Research Center for Social Computing and Interactive Robotics, Harbin Institute of Technology, China · Tsinghua University, China · Baidu Inc., China
MLAfication: Efficiently Enabling High-Fidelity MLA for Modern LLMs
Oct 12, 2026Zile Wang, Zhide Xie, Tao JiMulti-Head Latent Attention LLM Inference Acceleration
Fudan University

Sep 29, 2026J/K move · Enter open · S save

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

Zunhai Su, Yuxuan Sun, Jianchao Tan, Tao Zhang+4

The University of Hong Kong · Meituan LongCat Team · South China University of Technology+1

LLM Inference Acceleration Attention Quantization KV-Cache Quantization Multi-Head Latent Attention Post-Training Quantization Low-Bit Quantization+1

Open paper AskPDF arXiv

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.

Abstract

LLM Quantization

Explore similar work

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

RoLaQ: Joint Rotation and Range Calibration for 2-Bit KV Caches in Multi-Head Latent Attention

MLAfication: Efficiently Enabling High-Fidelity MLA for Modern LLMs