The design of attention mechanisms in long-context large language models faces an ``impossible triangle'' among KV cache size, training and prefill FLOPs, and model quality. GQA achieves strong quality but incurs high KV cache overhead; alternative approaches such as MLA and MFA can compress the cache, but MLA suffers from quality loss while MFA substantially increases training and prefill compute. We propose Virtual Head Attention (VHA), which simultaneously improves all three vertices through two lightweight linear operations: Q Premix applies a near-identity transformation to queries within each KV group to recover virtual head diversity from halved physical query heads, and Linear Postmix fuses inter-head features via a low-rank I+AB⊤ residual structure. Since VHA introduces no nonlinear operations, Premix can be folded into the query projection and Postmix into the output projection at inference time, reducing the model to a standard GQA-2 architecture that directly reuses all existing GQA inference optimizations. With only two KV heads, the group size in large models readily reaches 64, fully exploiting Tensor Core parallelism during decoding. Across four scales with 4K-context 50B-token pretraining, VHA consistently surpasses the standard GQA baseline at every scale (+0.05% to +0.81%), while reducing KV cache by 4× vs.\ GQA-8 at dense scales and 2× vs.\ GQA-4 at the MoE scale, and lowering training and prefill FLOPs by 7.9% (1.7B, S=4096). On an industrial 30B-A3B MoE model trained on 500B tokens at 8K context, VHA outperforms both MLA and GQA-4 by +0.70,/,+2.60,% on a 12-benchmark base-model evaluation suite; in a separate matched 10B-A1B MoE comparison after long-context mid-training, VHA attains an average RULER score of 84.4 over 8K--128K versus 73.0 for GQA, while using two rather than four KV heads in the full-attention layers.
Figures & tables
Figure 1: Overview of Virtual Head Attention. (a) VHA improves the trade-off among KV cache size, training/prefill FLOPs, and model quality by combining Q Premix and Linear Postmix during training, while absorbing both linear modules into standard GQA projections at inference. (b) Schematic of the KV cache–FLOPs–quality impossible triangle faced by existing attention designs. (c) Method comparison at 1.7B and 4B scales, plotting KV cache size against training FLOPs, with bubble area proportional to 8-task average accuracy.
Scale
Attn
Q-K-V
Cache
ARC-C
BoolQ
HSwag
OBQA
PIQA
Wino
SciQ
SIQA
Avg
0.6B
MHA
12-12-12
6.00
29.10
60.40
44.88
33.20
68.66
51.85
74.50
39.15
50.22
GQA
16-8-8
4.00
29.86
58.01
44.77
35.80
68.17
52.33
75.70
38.69
50.42
MLA
14-14-14
1.12
29.27
55.96
43.20
33.80
68.23
50.67
76.40
38.59
49.52
MFA
17-1-1
1.00
30.80
60.18
45.47
32.80
69.04
54.70
73.20
39.05
50.66
VHA
8-2-2
1.00
31.91
61.19
45.26
35.60
68.44
50.99
75.60
39.36
51.04
1.7B
MHA
12-12-12
6.00
35.24
58.53
51.86
35.20
71.33
54.06
80.30
39.36
53.24
Table 1: Pretraining accuracy across four model scales. We report zero-shot accuracy (%) on eight standard multiple-choice benchmarks. Cache is KV cache per layer per token in BF16 (KB). Bold marks the best per scale per column; ties are bolded jointly. VHA achieves the highest average accuracy at every scale while using the smallest KV cache. Per-method full configurations and benchmark citations are in Appendix D ; per-method training-FLOPs comparison is in Table 5 .
Method
Cache
ARC-C
BoolQ
HSwag
OBQA
PIQA
Wino
SciQ
SIQA
Avg
Δ
GQA 16-8-8
4.00
35.15
62.11
51.98
35.60
71.71
55.25
81.10
39.82
54.09
—
GQA 16-2-2
1.00
33.19
58.96
50.92
36.00
71.27
54.22
79.60
40.53
53.09
− 1.00
GQA 8-2-2
1.00
33.02
59.11
49.98
34.00
71.44
53.28
77.80
40.43
52.38
− 1.71
VHA 8-2-2
1.00
35.92
60.64
51.71
39.00
71.38
54.70
81.70
40.94
54.50
+ 0.41
Table 2: Recovering from a doubly-compressed shape (1.7B). GQA 16-8-8 is the standard GQA-8 baseline; GQA 16-2-2 reduces only Hkv to obtain a 1 KB KV cache; GQA 8-2-2 additionally halves the physical query heads to match VHA’s shape but lacks Premix and Postmix. Δ is the gap to the GQA 16-8-8 baseline.
Method
Q-K-V
Cache
General
Math & STEM
Coding
Multilingual
Avg
GQA
48-4-4
2.00
52.86
43.09
51.85
43.26
47.76
MLA
28-28-28
1.12
55.57
43.49
54.37
45.20
49.66
VHA
24-2-2
1.00
56.25
46.40
53.23
45.56
50.36
Table 3: Pretraining accuracy at the 30B-A3B MoE scale. We report accuracy (%) averaged within four task categories (General, Math & STEM, Coding, Multilingual) over 12 base-model benchmarks; Avg is the macro-average over the four task categories. Cache is KV cache per layer per token in BF16 (KB). Bold marks the best per column. Per-benchmark scores and benchmark citations are in Appendix C .
0.6B
1.7B
4B
8B-A1B
Method
Q-K-V
Cache
FLOPs
Δ %
Q-K-V
Cache
FLOPs
Δ %
Q-K-V
Cache
FLOPs
Δ %
Q-K-V
Cache
FLOPs
Δ %
MHA
12-12-12
6.00
6.87
− 9.3
12-12-12
6.00
9.91
− 6.6
20-20-20
10.00
21.66
− 11.1
20-20-20
10.00
8.76
− 14.7
GQA
16-8-8
4.00
7.57
0
16-8-8
4.00
10.61
0
32-8-8
4.00
24.38
0
32-4-4
2.00
10.27
0
MLA
14-14-14
1.12
8.25
+ 8.9
12-12-12
1.12
10.42
− 1.8
28-28-28
1.12
25.09
+ 2.9
28-28-28
1.12
11.39
+ 10.9
MFA
17-1-1
1.00
11.54
+ 52.3
19-1-1
1.00
15.90
+ 49.8
33-1-1
1.00
36.04
+ 47.9
33-1-1
1.00
17.82
+ 73.5
VHA
8-2-2
1.00
6.96
− 8.1
8-2-2
1.00
9.78
− 7.9
16-2-2
1.00
22.52
− 7.6
16-2-2
1.00
9.60
− 6.5
Table 5: Computational efficiency across four model scales ( S=4096 ). For each scale we list per-method query/key/value head counts, KV cache per layer per token (BF16, KB), training FLOPs per token aggregated over all transformer layers (G), and the relative change vs. the per-scale GQA baseline (GQA-8 dense, GQA-4 MoE). Bold highlights VHA, the method that achieves the best joint KV cache + FLOPs Pareto optimum at every scale. MHA’s lower FLOPs at three scales reflect its smaller query head count under our parameter-matched configurations ( Hq=12 at 0.6B/1.7B, Hq=20 at 4B/8B-A1B versus GQA’s Hq=16 or Hq=32 ), at the cost of a 6 – 10× larger KV cache.
Figure 2: Q Premix weight analysis on a trained Qwen-1.7B VHA model. (a) Per-layer Frobenius deviation ∥Wgpremix−I∥F per KV group (group mean overlaid). (b) Per-layer mean pairwise distance ∥Wg1premix−Wg2premix∥F between KV groups.
Figure 3: Postmix head-mixing analysis on a trained Qwen-4B VHA model. Sankey diagrams of ∣AB⊤∣ across five representative layers, with ribbons flowing from source heads (left, colored by KV group) to recombined output heads (right); width and color saturation proportional to ∣AB⊤∣j,i .
Appendix figures & tables8 assets
Supplementary material from the paper’s appendix.
Appendix
RULER
HELMET
Checkpoint
Length
Model
Overall
Retr.
Aggr.
QA
VT
RAG
LongQA
Summ.
ICL
Rerank
32K
8K
GQA
80.9
92.4
73.1
56.5
53.6
24.5
9.0
7.2
60.7
0.3
VHA
85.7
95.9
78.0
60.1
69.8
29.6
11.0
6.6
57.1
0.4
16K
GQA
73.7
90.0
51.8
53.7
27.6
22.6
11.0
8.0
70.4
0.0
VHA
83.2
93.3
71.5
55.4
80.8
32.3
10.5
7.2
68.6
0.0
32K
GQA
69.7
81.9
51.9
49.0
49.4
22.3
17.0
8.1
78.3
0.0
Appendix
Table 6: Long-context results on RULER and HELMET (%). RULER Overall averages all 13 tasks. Bold marks the better result within each matched checkpoint and evaluation length; exact ties are unbolded.
Checkpoint
Model
LB-v2 ↑
InfBench ↑
BABI ↑
Short ↑
arXiv ↓
Books ↓
Code ↓
Web ↓
32K
GQA
28.0
29.9
36.5
62.7
0.568
0.812
0.294
0.763
VHA
22.3
27.3
37.4
63.7
0.568
0.812
0.295
0.765
128K
GQA
18.9
31.8
31.1
62.6
0.582
0.819
0.311
0.768
VHA
24.2
31.5
38.2
64.1
0.564
0.804
0.296
0.760
Appendix
Table 7: Complementary evaluation results. LB-v2, InfBench, and BABI denote LongBench-v2, InfiniteBench, and BABILong; Short is the short-context eight-task average. The remaining columns report bits per byte (BPB) on ProLong domains. Bold marks the better matched result.
Model
Run
ARC-C
BoolQ
HSwag
OBQA
PIQA
Wino
SciQ
SIQA
Avg
VHA
01
30.12
60.18
44.46
33.80
68.77
52.25
74.20
39.56
50.42
VHA
02
29.35
59.27
44.36
32.80
68.55
55.96
74.00
39.56
50.48
VHA
03
29.27
59.36
44.79
33.80
67.14
52.25
76.00
40.23
50.35
VHA
04
30.12
56.09
44.68
34.80
68.55
52.64
76.10
39.76
50.34
VHA mean
29.72
58.73
44.57
33.80
68.25
53.28
75.08
39.78
50.40±0.07
GQA
01
29.44
58.59
44.49
32.40
68.28
51.14
75.50
38.79
49.83
Appendix
Table 8: Repeated-run results at 0.6B (%). Every run is reported. Bold in the summary rows marks the higher across-run mean.
Step
α=0.1
α=1.0
Random
10
12.2976 / 14.74
12.3072 / 14.80
12.3132 / 14.98
20
11.1501 / 15.97
11.1581 / 16.26
11.1576 / 16.25
30
9.9346 / 4.98
9.9389 / 5.01
9.9400 / 5.07
50
9.1168 / 2.67
9.1243 / 3.76
9.1255 / 2.03
100
7.9670 / 2.96
7.9788 / 2.93
7.9669 / 4.02
200
6.4777 / 2.99
6.4878 / 2.21
6.4898 / 1.65
Appendix
Table 9: Early training trajectory under different Premix initializations. Each entry is loss / pre-clipping gradient norm.
Task
α=0.1
α=1.0
Random
ARC-C
34.81
34.13
34.39
BoolQ
60.06
58.04
57.03
HellaSwag
51.80
51.72
51.76
OpenBookQA
37.60
37.00
36.00
PIQA
71.06
71.33
72.14
Winogrande
55.72
54.85
53.28
Appendix
Table 10: Downstream accuracy under different Premix initializations (%). Bold marks the best result per row.
Category
Benchmark
Shot
GQA
MLA
VHA
General
MMLU
5-shot
55.97
59.34
60.24
MMLU-Redux
5-shot
63.05
65.25
66.10
MMLU-Pro
5-shot
36.20
41.12
40.03
BBH
3-shot CoT
56.23
56.57
58.61
Math & STEM
GPQA
10-shot CoT
34.62
31.62
34.11
GSM8K
4-shot CoT
56.48
55.42
56.25
Appendix
Table 11: Per-benchmark scores on the 30B-A3B MoE evaluation suite (accuracy %).
0.6B
1.7B
4B
8B-A1B (MoE)
30B-A3B (MoE)
Architecture
Precision
BF16
BF16
BF16
BF16
FP8
Num layers
28
28
36
24
46
Hidden size
1024
2048
2560
2048
1536
FFN dim
3072
6144
9728
768×64
1024×128
Num experts
–
–
–
64
128 (+1 shared)
Appendix
Table 12: Model architecture and training hyperparameters across all five scales. “–” denotes not applicable.
Figure 4: Training loss curves across four model scales. Each panel compares MHA, GQA, MLA, MFA, and VHA on FineWeb-Edu (50B tokens, 4K context). All curves are plotted from step 4K onward after the learning-rate warm-up; smoothing is applied with an exponential moving average ( α=0.93 ). VHA (blue, solid) achieves the lowest loss at convergence at the 4B scale.
Parameter-efficient fine-tuning reduces model and optimizer memory, but dense attention still makes long training sequences expensive. We combine Hierarchical Global Attention (HGA) with segment-wise backpropagation and tiered KV storage. Only the active segment remains differentiable in VRAM; older KV is detached into RAM or NVMe, and HGA loads a bounded set of exact historical tokens for each query block. On Qwen3-8B with 4-bit QLoRA and PG19, dense training on a 16 GB Quadro RTX 5000 fits 2,048 tokens but fails at 4,096, whereas HGA reaches 16,384 tokens with 15.28 GB peak VRAM. Under evaluation the same adapter runs through 131,072 tokens on this card; VRAM is not constant but grows gently with the resident chunk summaries, so RAM and NVMe capacity set the practical limit beyond these lengths. At the shared 2K training length, HGA-trained and dense-trained adapters obtain 2.7405 and 2.7383 nat under the same dense-attention readout, while the stock model obtains 2.9541. At this boundary HGA training is already marginally faster (217.75 vs. 207.02 tokens/s), and the HGA-to-dense throughput ratio improves from 1K to 2K; because HGA keeps the attended historical set per token approximately constant while dense work per token grows, we expect this lead to widen as context grows. Dense attention is used for the main quality and retrieval comparisons so that they measure the learned weights and remain compatible with standard generation frameworks. HGA can also be used for retrieval and generation; an optimized production-grade serving implementation is under development.
Vladimir Fedosov, Aleksandr Sazhin, Artemiy Grinenko +1
Long-context LLM inference is bottlenecked by quadratic attention computation and growing KV-cache costs. Existing sparse attention and KV-compression methods typically decide which tokens or heads to preserve from runtime attention scores, observation windows, calibration prompts, or learned gates, making head diagnosis input-dependent and costly to deploy. We propose Autonomy-of-Heads (AoH), a data-free method that identifies retrieval and streaming heads from the spectral geometry of query-key projections. AoH defines the kernel attention operator Mh=WKh⊤WQh and uses its effective-rank as a weight-space measure of head function: concentrated spectra indicate a small number of dominant query-key matching directions and are associated with retrieval heads, whereas diffuse spectra indicate the absence of a dominant global matching direction and are associated with streaming heads. We further derive an efficient dhead-dimensional computation that avoids constructing the full dmodel×dmodel matrix. We conducted extensive experiments across models demonstrating that at 50% sparsity, AoH retains 96.5% of Full Attention performance on average while reducing prefill and decode latency by up to 41.4% and 66.0%, respectively, and KV-cache memory by 50.0% at 256K tokens.
Yehan Yang, Junyuan Shang, Yang Li +3
Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences · Baidu Inc. +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