Converting pretrained attention modules such as grouped-query attention (GQA) into multi-head latent attention (MLA) can improve expressivity without increasing KV-cache cost, making it attractive for efficient inference. However, many practical conversion baselines rely on weight-only low-rank approximations (e.g., SVD-style initializations) and uniform rank allocation. They focus on minimizing the difference between weight matrices rather than on how those weights affect input activations, ignore the covariance structure of activations, and enforce uniform rank across layers, causing activation drift and degraded attention fidelity. To address these issues, we propose CARE, a Covariance-Aware, Rank-Enhanced MLA conversion pipeline under a fixed KV width. CARE introduces three key steps: (i) activation-preserving factorization, which aligns the approximation with the actual input activations rather than just the weights; (ii) adjusted-rank allocation, which spreads a fixed KV budget across layers by giving more capacity to layers that need it most; and (iii) KV-parity mapping, which reparameterizes the converted K and V to fit the MLA format while keeping the KV-cache size unchanged. Our method outperforms a uniform-rank SVD baseline on Qwen3-4B/30B-A3B-Instruct-2507 and Llama-3.1-8B/70B-Instruct, reducing one-shot perplexity by up to 215x and improving mean accuracy by up to 1.70x at matched KV budgets. With a brief post-SVD healing fine-tune, we fully recover the original model's accuracy.
Figures & tables
Figure 1: (a) Naive MLA transfer : jointly factorize WK(g) and WV(g) by SVD and truncate to a uniform per-layer rank, optimizing ∥W−W^∥F while ignoring layerwise heterogeneity. (b) CARE : estimate activation covariance C , factorize CW , unwhiten via C−1 to initialize MLA factors, and use the singular spectrum of CW for global dynamicrank scheduling under KV parity. This preserves activation geometry and yields a stronger one-shot initialization with less healing.
Figure 2: (a) Accuracy under 50% rank reduction applied one layer at a time in DeepSeek-V2-Lite, measured on ARC Challenge and MMLU. Sensitivity is strongly layer-dependent. (b) WikiText perplexity under grouped truncation of GQA attention in Llama-3-8B (layers 30–32). Singular-value groups are ordered by magnitude; the resulting non-monotone degradation shows that singular values alone are an imperfect proxy for MLA conversion quality.
Rank
KV Save
Methods
PPL ( ↓ )
ACC ( ↑ )
ARC ( ↓ )
ARE
HellaSwag
PIQA
MMLU
OBQA
RA
WG
AVG
Llama-3.1-8B-Instruct
64
93.75
GQA (Original)
7.21 ( ↓ )
50.34
80.18
60.15
79.65
48.05
34.80
40.10
72.69
58.24
Palu(SVD)
2260.60 ( ↓ )
25.77
27.53
26.50
52.18
24.26
27.80
20.86
50.04
31.87
MHA2MLA
284863.91 ( ↓ )
25.94
23.82
26.34
50.92
25.62
27.80
22.11
50.59
31.64
CARE-U (OURS)
983.55 ( ↓ )
23.89
30.51
26.82
54.73
23.18
26.00
20.96
49.01
31.89
Table 1: One-shot comparison on Llama3.1-8B and Qwen3-4B-Instruct-2507 against original and baselines on multiple tasks. Calibration samples: 256. Sequence length: 2048. Calibration dataset: alpaca. Higher is better for Accuracy (%) (ACC.) ( ↑ ) and lower is better for Perplexity (PPL.) ( ↓ ).
Figure 3: Covariance-aware rank profiles across calibration corpora (Alpaca, WikiText2, PTB, C4) for Llama-3.1-8B-Instruct at target ranks 64, 128, 256, and 512. Across all target budgets, both WK and WV show a depth-dependent increase—small in early layers, rising through mid layers—with stronger late-layer growth for WV . The consistency across corpora suggests a model-intrinsic trend.
Figure 4: One-shot accuracy versus calibration samples and sequence length across eight benchmarks. Sequence length is fixed at 256 unless otherwise noted; the red curve denotes fixed 256 samples with varying sequence length. OOM occurs beyond length =512 during covariance computation.
Appendix figures & tables25 assets
Supplementary material from the paper’s appendix.
Appendix
Rank
Methods
PPL ( ↓ )
ACC ( ↑ )
ARC-E
ARC-C
HellaSwag
PIQA
WG
MMLU
OBQA
RA
AVG
64
Palu(SVD)
2260.60 ( ↓ )
27.53
25.77
26.50
52.18
50.04
24.26
27.80
20.86
31.87
MHA2MLA
284863.91 ( ↓ )
23.82
25.94
26.34
50.92
50.59
25.62
27.80
22.11
31.64
CARE-C-based-U
893.84 ( ↓ )
31.40
23.98
26.91
55.66
50.04
22.91
27.00
21.15
32.38
CARE-C-based-E
837.25 ( ↓ )
32.41
23.46
27.50
56.75
49.96
22.90
27.40
21.34
32.71
CARE-U
983.55 ( ↓ )
30.51
23.89
26.82
54.73
49.01
23.18
26.00
20.96
31.89
Appendix
Table I.1: Detailed zero-shot comparison for Llama-3.1-8B-Instruct with Alpaca calibration. Higher is better for Accuracy (%) (ACC.) ( ↑ ) and lower is better for Perplexity (PPL.) ( ↓ ).
Rank
Methods
PPL ( ↓ )
ACC ( ↑ )
ARC-E
ARC-C
HellaSwag
PIQA
WG
MMLU
OBQA
RA
AVG
64
Palu(SVD)
2260.60 ( ↓ )
27.53
25.77
26.50
52.18
50.04
24.26
27.80
20.86
31.87
MHA2MLA
284863.91 ( ↓ )
23.82
25.94
26.34
50.92
50.59
25.62
27.80
22.11
31.64
CARE-C-based-U
688.95 ( ↓ )
28.32
24.74
27.28
55.93
51.85
22.95
26.40
21.05
32.31
CARE-C-based-E
582.93 ( ↓ )
29.55
24.15
28.15
56.75
52.25
22.95
26.20
20.96
32.62
CARE-U
786.17 ( ↓ )
27.61
24.40
26.98
54.68
49.88
22.95
26.40
20.19
31.64
Appendix
Table I.2: Detailed zero-shot comparison for Llama-3.1-8B-Instruct with C4 calibration. Higher is better for Accuracy (%) (ACC.) ( ↑ ) and lower is better for Perplexity (PPL.) ( ↓ ).
Rank
Methods
PPL ( ↓ )
ACC ( ↑ )
ARC-E
ARC-C
HellaSwag
PIQA
WG
MMLU
OBQA
RA
AVG
64
Palu(SVD)
2260.60 ( ↓ )
27.53
25.77
26.50
52.18
50.04
24.26
27.80
20.86
31.87
MHA2MLA
284863.91 ( ↓ )
23.82
25.94
26.34
50.92
50.59
25.62
27.80
22.11
31.64
CARE-C-based-U
756.56 ( ↓ )
26.73
25.68
26.28
51.58
50.20
22.95
25.80
21.34
31.32
CARE-C-based-E
612.72 ( ↓ )
27.57
25.43
26.51
51.74
50.75
22.95
26.00
21.05
31.50
CARE-U
815.80 ( ↓ )
26.26
25.00
26.50
51.36
49.25
22.95
27.00
20.67
31.12
Appendix
Table I.3: Detailed zero-shot comparison for Llama-3.1-8B-Instruct with PTB calibration. Higher is better for Accuracy (%) (ACC.) ( ↑ ) and lower is better for Perplexity (PPL.) ( ↓ ).
Rank
Methods
PPL ( ↓ )
ACC ( ↑ )
ARC-E
ARC-C
HellaSwag
PIQA
WG
MMLU
OBQA
RA
AVG
64
Palu(SVD)
2260.60 ( ↓ )
27.53
25.77
26.50
52.18
50.04
24.26
27.80
20.86
31.87
MHA2MLA
284863.91 ( ↓ )
23.82
25.94
26.34
50.92
50.59
25.62
27.80
22.11
31.64
CARE-C-based-U
189.39 ( ↓ )
28.11
25.09
26.82
51.47
50.12
22.95
27.80
20.57
31.62
CARE-C-based-E
138.00 ( ↓ )
28.91
25.43
27.16
52.39
50.36
22.95
27.60
20.38
31.90
CARE-U
237.45 ( ↓ )
27.31
25.34
26.46
51.63
48.22
22.94
27.60
20.77
31.28
Appendix
Table I.4: Detailed zero-shot comparison for Llama-3.1-8B-Instruct with WikiText2 calibration. Higher is better for Accuracy (%) (ACC.) ( ↑ ) and lower is better for Perplexity (PPL.) ( ↓ ).
Rank
Methods
PPL ( ↓ )
ACC ( ↑ )
ARC-E
ARC-C
HellaSwag
PIQA
WG
MMLU
OBQA
RA
AVG
64
Palu(SVD)
114579.38 ( ↓ )
25.17
25.26
25.55
50.27
50.20
24.06
27.40
20.67
31.07
MHA2MLA
62077.36 ( ↓ )
26.56
25.34
26.57
51.63
50.20
24.60
29.00
22.68
32.07
CARE-C-based-U
909.76 ( ↓ )
29.00
22.35
26.14
52.50
48.15
23.11
26.60
22.30
31.27
CARE-C-based-E
730.60 ( ↓ )
29.00
22.87
26.83
54.19
51.30
22.85
27.60
24.11
32.34
CARE-U
1098.80 ( ↓ )
26.77
23.38
26.95
52.29
51.07
22.95
25.20
23.16
31.47
Appendix
Table I.5: Detailed zero-shot comparison for Qwen2.5-1.5B-Instruct with Alpaca calibration. Higher is better for Accuracy (%) (ACC.) ( ↑ ) and lower is better for Perplexity (PPL.) ( ↓ ).
Rank
Methods
PPL ( ↓ )
ACC ( ↑ )
ARC-E
ARC-C
HellaSwag
PIQA
WG
MMLU
OBQA
RA
AVG
64
Palu(SVD)
114579.38 ( ↓ )
25.17
25.26
25.55
50.27
50.20
24.06
27.40
20.67
31.07
MHA2MLA
62077.36 ( ↓ )
26.56
25.34
26.57
51.63
50.20
24.60
29.00
22.68
32.07
CARE-C-based-U
679.30 ( ↓ )
27.74
23.12
26.69
52.77
51.22
22.95
26.80
23.16
31.81
CARE-C-based-E
365.98 ( ↓ )
28.54
22.70
27.42
55.01
52.09
22.95
26.40
23.64
32.34
CARE-U
914.36 ( ↓ )
26.43
24.57
26.17
51.90
49.41
22.95
25.80
23.83
31.38
Appendix
Table I.6: Detailed zero-shot comparison for Qwen2.5-1.5B-Instruct with C4 calibration. Higher is better for Accuracy (%) (ACC.) ( ↑ ) and lower is better for Perplexity (PPL.) ( ↓ ).
Rank
Methods
PPL ( ↓ )
ACC ( ↑ )
ARC-E
ARC-C
HellaSwag
PIQA
WG
MMLU
OBQA
RA
AVG
64
Palu(SVD)
114579.38 ( ↓ )
25.17
25.26
25.55
50.27
50.20
24.06
27.40
20.67
31.07
MHA2MLA
62077.36 ( ↓ )
26.56
25.34
26.57
51.63
50.20
24.60
29.00
22.68
32.07
CARE-C-based-U
1608.42 ( ↓ )
27.48
23.46
26.79
51.09
50.59
23.01
26.80
21.63
31.36
CARE-C-based-E
701.84 ( ↓ )
27.95
23.55
27.12
51.52
50.04
22.95
24.80
23.73
31.46
CARE-U
2131.72 ( ↓ )
27.53
23.38
27.00
51.74
49.01
23.14
27.20
23.06
31.51
Appendix
Table I.7: Detailed zero-shot comparison for Qwen2.5-1.5B-Instruct with PTB calibration. Higher is better for Accuracy (%) (ACC.) ( ↑ ) and lower is better for Perplexity (PPL.) ( ↓ ).
Rank
Methods
PPL ( ↓ )
ACC ( ↑ )
ARC-E
ARC-C
HellaSwag
PIQA
WG
MMLU
OBQA
RA
AVG
64
Palu(SVD)
114579.38 ( ↓ )
25.17
25.26
25.55
50.27
50.20
24.06
27.40
20.67
31.07
MHA2MLA
62077.36 ( ↓ )
26.56
25.34
26.57
51.63
50.20
24.60
29.00
22.68
32.07
CARE-C-based-U
209.05 ( ↓ )
27.86
22.18
26.34
52.56
50.67
22.98
28.20
23.64
31.80
CARE-C-based-E
122.65 ( ↓ )
27.86
23.29
27.12
53.48
50.75
22.94
26.20
22.68
31.79
CARE-U
300.99 ( ↓ )
26.77
23.12
26.74
52.18
50.20
22.95
26.60
24.50
31.63
Appendix
Table I.8: Detailed zero-shot comparison for Qwen2.5-1.5B-Instruct with WikiText2 calibration. Higher is better for Accuracy (%) (ACC.) ( ↑ ) and lower is better for Perplexity (PPL.) ( ↓ ).
Rank
Methods
PPL ( ↓ )
ACC ( ↑ )
ARC-E
ARC-C
HellaSwag
PIQA
WG
MMLU
OBQA
RA
AVG
64
Palu(SVD)
56922.11 ( ↓ )
26.77
25.34
25.73
50.44
50.91
24.31
28.40
22.11
31.75
MHA2MLA
21850.16 ( ↓ )
25.46
25.68
26.03
51.96
50.20
22.92
27.00
22.01
31.41
CARE-C-based-U
1052.43 ( ↓ )
29.34
22.70
27.51
54.35
50.67
23.21
26.20
23.64
32.20
CARE-C-based-E
941.16 ( ↓ )
28.75
23.12
27.33
55.33
50.28
23.08
25.60
22.30
31.97
CARE-U
905.74 ( ↓ )
29.00
23.46
27.61
55.22
50.59
23.03
26.60
23.06
32.32
Appendix
Table I.9: Detailed zero-shot comparison for Qwen3-4B-Instruct-2507 with Alpaca calibration. Higher is better for Accuracy (%) (ACC.) ( ↑ ) and lower is better for Perplexity (PPL.) ( ↓ ).
Rank
Methods
PPL ( ↓ )
ACC ( ↑ )
ARC-E
ARC-C
HellaSwag
PIQA
WG
MMLU
OBQA
RA
AVG
64
Palu(SVD)
56922.11 ( ↓ )
26.77
25.34
25.73
50.44
50.91
24.31
28.40
22.11
31.75
MHA2MLA
21850.16 ( ↓ )
25.46
25.68
26.03
51.96
50.20
22.92
27.00
22.01
31.41
CARE-C-based-U
494.84 ( ↓ )
28.91
23.55
27.41
55.11
50.91
22.93
23.60
24.21
32.08
CARE-C-based-E
447.15 ( ↓ )
28.28
21.50
27.90
55.28
50.36
22.95
24.40
23.35
31.75
CARE-U
380.04 ( ↓ )
29.59
22.53
28.13
53.21
51.14
22.92
26.20
23.64
32.17
Appendix
Table I.10: Detailed zero-shot comparison for Qwen3-4B-Instruct-2507 with C4 calibration. Higher is better for Accuracy (%) (ACC.) ( ↑ ) and lower is better for Perplexity (PPL.) ( ↓ ).
Rank
Methods
PPL ( ↓ )
ACC ( ↑ )
ARC-E
ARC-C
HellaSwag
PIQA
WG
MMLU
OBQA
RA
AVG
64
Palu(SVD)
56922.11 ( ↓ )
26.77
25.34
25.73
50.44
50.91
24.31
28.40
22.11
31.75
MHA2MLA
21850.16 ( ↓ )
25.46
25.68
26.03
51.96
50.20
22.92
27.00
22.01
31.41
CARE-C-based-U
1234.23 ( ↓ )
28.49
22.44
26.91
51.74
48.46
22.95
24.20
22.20
30.92
CARE-C-based-E
1036.30 ( ↓ )
28.32
22.87
26.93
52.67
50.83
22.95
24.40
21.82
31.35
CARE-U
961.94 ( ↓ )
28.79
23.29
27.49
51.31
49.09
22.98
24.60
22.58
31.27
Appendix
Table I.11: Detailed zero-shot comparison for Qwen3-4B-Instruct-2507 with PTB calibration. Higher is better for Accuracy (%) (ACC.) ( ↑ ) and lower is better for Perplexity (PPL.) ( ↓ ).
Rank
Methods
PPL ( ↓ )
ACC ( ↑ )
ARC-E
ARC-C
HellaSwag
PIQA
WG
MMLU
OBQA
RA
AVG
64
Palu(SVD)
56922.11 ( ↓ )
26.77
25.34
25.73
50.44
50.91
24.31
28.40
22.11
31.75
MHA2MLA
21850.16 ( ↓ )
25.46
25.68
26.03
51.96
50.20
22.92
27.00
22.01
31.41
CARE-C-based-U
169.94 ( ↓ )
27.99
22.35
27.16
52.18
49.25
22.95
25.60
23.35
31.35
CARE-C-based-E
142.06 ( ↓ )
28.37
23.81
27.50
52.72
49.64
22.95
27.20
22.20
31.80
CARE-U
156.42 ( ↓ )
28.66
23.29
28.50
52.45
49.96
22.95
26.00
22.11
31.74
Appendix
Table I.12: Detailed zero-shot comparison for Qwen3-4B-Instruct-2507 with WikiText2 calibration. Higher is better for Accuracy (%) (ACC.) ( ↑ ) and lower is better for Perplexity (PPL.) ( ↓ ).
Rank
Methods
PPL ( ↓ )
ACC ( ↑ )
ARC-E
ARC-C
HellaSwag
PIQA
WG
MMLU
OBQA
RA
AVG
128
Palu(SVD)
17930.54 ( ↓ )
25.88
26.28
26.22
50.98
48.22
24.29
25.60
22.01
31.19
MHA2MLA
1929.49 ( ↓ )
29.46
24.74
26.57
51.58
48.62
22.85
27.20
21.34
31.55
CARE-C-based-U
29.00 ( ↓ )
58.75
38.74
58.02
71.65
55.33
48.40
36.00
33.78
50.08
CARE-C-based-E
36.65 ( ↓ )
50.00
33.36
56.32
69.86
54.46
43.01
33.00
31.48
46.44
CARE-U
39.87 ( ↓ )
53.75
35.75
54.22
67.36
53.91
42.84
30.80
30.43
46.13
Appendix
Table I.13: Detailed zero-shot comparison for Qwen3-30B-A3B-Instruct-2507 with Alpaca calibration. Higher is better for Accuracy (%) (ACC.) ( ↑ ) and lower is better for Perplexity (PPL.) ( ↓ ).
Figure I.1: Needle-in-a-Haystack retrieval heatmaps for Llama-3.1-8B-Instruct under matched KV budgets. The figure assembles five panel-wise heatmaps comparing CARE variants with different calibration settings against Palu(SVD) across context lengths and needle depths. Greener cells indicate stronger retrieval accuracy.
Figure I.2: Generated Text Example
Figure I.3: Generated Text Example
Figure I.4: Covariance-aware rank profiles for Qwen3-30B-A3B-Instruct-2507 (MoE, 30B total / 3B active) under Alpaca calibration at target ranks 128, 256. Despite the mixture-of-experts architecture, the same depth-dependent pattern persists, confirming the trend generalizes beyond dense models.
Figure I.5: Covariance-aware rank profiles for Llama-3.1-70B-Instruct under Alpaca calibration at target ranks 128, 256. The same depth-dependent pattern observed in the 8B variant holds at 70B scale, with WV exhibiting stronger late-layer growth than WK .
Figure I.6: Covariance-aware rank profiles for Qwen2.5-1.5B-Instruct across Alpaca, WikiText2, PTB, and C4 calibration corpora at target ranks 64 and 128. Despite the smaller rank budget, the same depth-dependent increase remains visible, again with stronger late-layer growth for WV .
Figure I.7: Covariance-aware rank profiles for Qwen3-4B-Instruct-2507 across Alpaca, WikiText2, PTB, and C4 calibration corpora at target ranks 64, 128, 256, and 512. Both WK and WV show a depth-dependent increase, with stronger late-layer growth for WV .
Rank
Methods
ARC ( ↑ )
ARE ( ↑ )
HellaSwag ( ↑ )
PIQA ( ↑ )
MMLU ( ↑ )
OBQA ( ↑ )
RA ( ↑ )
WG ( ↑ )
AVG ( ↑ )
256
CARE- Alpaca
34.81
65.40
40.76
72.47
49.32
21.20
35.98
62.98
47.87
CARE- ARE
38.82
72.90
41.25
73.01
53.13
25.80
31.87
63.69
50.06
CARE- ARC
39.59
71.84
41.33
72.91
52.84
26.60
32.54
62.67
50.04
CARE- WikiText2
28.33
57.11
38.07
68.28
37.72
19.60
33.49
62.59
43.15
CARE- PTB
27.99
53.91
37.70
67.68
42.15
17.00
31.96
62.59
42.62
CARE- C4
31.48
59.51
43.00
73.12
41.69
20.00
35.89
63.30
46.00
Appendix
Table I.14: One-shot Llama3.1-8B comparison on different covariance. Higher is better for Accuracy (ACC.).
Model
Rank
α
ARC ( ↑ )
ARE ( ↑ )
HellaSwag ( ↑ )
PIQA ( ↑ )
MMLU ( ↑ )
OBQA ( ↑ )
RA ( ↑ )
WG ( ↑ )
AVG ( ↑ )
Llama-3.1-8B-Instruct
256
10−3
34.90
62.42
42.44
72.31
57.15
21.00
34.74
64.48
48.68
10−2
34.90
62.46
42.52
72.14
57.20
21.60
35.60
65.35
48.97
10−1
32.94
59.34
41.55
70.46
56.88
20.20
35.31
65.75
47.80
512
10−3
46.16
76.30
53.64
77.69
65.54
29.40
41.44
72.85
57.88
10−2
51.79
80.72
51.70
74.54
70.13
30.20
39.43
68.27
58.35
10−1
47.27
77.10
54.40
77.80
66.15
28.40
42.58
74.35
58.51
Appendix
Table I.15: Shrinkage-coefficient ablation for CARE under the Alpaca-256-32 calibration setup. We report one-shot accuracy across Llama-3.1-8B-Instruct and Qwen3-4B-Instruct at target ranks 256 and 512. Higher is better for Accuracy (ACC.).
Rank
Methods
ARC ( ↑ )
ARE ( ↑ )
HellaSwag ( ↑ )
PIQA ( ↑ )
MMLU ( ↑ )
OBQA ( ↑ )
RA ( ↑ )
WG ( ↑ )
AVG ( ↑ )
256
CARE- Alpaca
34.81
65.40
40.76
72.47
49.32
21.20
35.98
62.98
47.87
CARE- 0.5A+0.5ARC
36.09
67.72
43.22
73.34
58.70
24.40
34.26
63.38
50.14
CARE- 0.2A+0.8ARC
34.73
65.03
42.99
72.25
58.02
23.20
35.41
64.96
49.57
CARE- ARC
39.59
71.84
41.33
72.91
52.84
26.60
32.54
62.67
50.04
Appendix
Table I.16: Task-weighted covariance mixing for multi-objective rank allocation (Llama3.1-8B-Instruct, CARE, Rank 256). A denotes Alpaca and ARC denotes ARC-Challenge ; mixes indicate sampling proportions used to estimate Cmix . Higher is better for Accuracy (ACC.).
Model
Rank
CARE-U ( C )
CARE-C-based-U ( C )
CARE-E ( C )
CARE-C-based-E ( C )
Llama-3.1-8B-Instruct
64
31.89
32.38
32.37
32.71
128
36.72
37.11
38.59
38.56
256
50.00
48.79
52.62
51.41
512
62.33
61.69
57.82
59.00
Qwen2.5-1.5B-Instruct
64
31.47
31.27
32.10
32.34
96
33.83
34.57
42.19
43.03
Appendix
Table I.17: Comparison of C versus C covariance weighting under Alpaca calibration. Values are zero-shot AVG accuracy (%); within each CARE variant the better result is shown in bold.
Multi-Head Latent Attention (MLA) compresses the KV cache into a latent representation, reducing inference memory and bandwidth. Existing MHA2MLA/TransMLA methods do not support QKNorm and require costly recovery on modern LLMs, especially MoE. We present MLAfication, a QKNorm-compatible, parameter-efficient extension of MHA2MLA. It combines calibrated KNorm, parameter-efficient distillation, contribution-aware Partial-RoPE, and activation-aware joint KV initialization. Across MHA and GQA models, including QKNorm- and MoE-equipped variants, MLAfication preserves performance under substantial KV-cache compression, surpasses the current SOTA on commonsense and long-context benchmarks, and extends MLA conversion to modern architectures. On Qwen3-30B-A3B, MLAfication delivers up to 5.35× the serving throughput of BF16 GQA and, combined with 4-bit KV quantization, reduces the KV-cache footprint by 94.73%. Our code is available at https://github.com/JT-Ushio/MLAfication.
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.
Zunhai Su, Yuxuan Sun, Jianchao Tan +5
The University of Hong Kong · Meituan LongCat Team · South China University of Technology +1
Long-context inference in large language models is bottlenecked by Key--Value (KV) cache loading during the decoding stage, where the sequential nature of generation requires repeatedly transferring the KV cache from off-chip High-Bandwidth Memory (HBM) to on-chip Static Random-Access Memory (SRAM) at each step. While Multi-Head Latent Attention (MLA) significantly reduces the total KV cache size, it suffers from a sharding bottleneck during distributed decoding via Tensor Parallelism (TP). Since its single latent head cannot be partitioned, each device is forced to redundantly load the complete KV cache for every token, consuming excessive memory traffic and diminishing TP benefits like weight sharding. In this work, we propose Multi-Head Low-Rank Attention (MLRA), which enables partitionable latent states for efficient 4-way TP decoding. Extensive experiments show that MLRA achieves state-of-the-art perplexity and downstream task performance, while also delivering a 2.8× decoding speedup over MLA. Code is available at https://github.com/SongtaoLiu0823/MLRA. Pretrained weights, along with the training and evaluation data, are available at https://huggingface.co/Soughing/MLRA.
Songtao Liu, Hongwu Peng, Zhiwei Zhang +2
The Pennsylvania State University · University of Connecticut · Carnegie Mellon University +1