Virtual Head Attention

Xibo Ding; Guoxia Wang; Jinle Zeng; Jiabin Yang; Sirui Tao; Chao Yang; Dianhai Yu

cs.CLOct 13, 2026

Virtual Head Attention

Authors: Xibo Ding, Guoxia Wang, Jinle Zeng, Jiabin Yang, Sirui Tao, Chao Yang, Dianhai Yu

Organizations: Baidu Inc. · Peking University

Abstract

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^\top$ 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\times$ vs.\ GQA-8 at dense scales and $2\times$ 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	$\Delta$
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

H_{kv}

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.

\Delta

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 .

Method	Q-K-V	Cache	FLOPs	$\Delta$ %	Q-K-V	Cache	FLOPs	$\Delta$ %	Q-K-V	Cache	FLOPs	$\Delta$ %	Q-K-V	Cache	FLOPs	$\Delta$ %
	0.6B				1.7B				4B				8B-A1B
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 (

H_{q}{=}12

at 0.6B/1.7B,

H_{q}{=}20

at 4B/8B-A1B versus GQA’s

H_{q}{=}16

H_{q}{=}32

), at the cost of a

6

–

10\times

larger KV cache.

Figure 2: Q Premix weight analysis on a trained Qwen-1.7B VHA model. (a) Per-layer Frobenius deviation

\|W_{g}^{\mathrm{premix}}-I\|_{F}

per KV group (group mean overlaid). (b) Per-layer mean pairwise distance

\|W_{g_{1}}^{\mathrm{premix}}-W_{g_{2}}^{\mathrm{premix}}\|_{F}

between KV groups.

Figure 3: Postmix head-mixing analysis on a trained Qwen-4B VHA model. Sankey diagrams of

|AB^{\top}|

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^{\top}|_{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 $\uparrow$	InfBench $\uparrow$	BABI $\uparrow$	Short $\uparrow$	arXiv $\downarrow$	Books $\downarrow$	Code $\downarrow$	Web $\downarrow$
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{\pm}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	$\alpha{=}0.1$	$\alpha{=}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	$\alpha{=}0.1$	$\alpha{=}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 %).

Architecture
	0.6B	1.7B	4B	8B-A1B (MoE)	30B-A3B (MoE)
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{\times}64$	$1024{\times}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 (

\alpha{=}0.93

). VHA (blue, solid) achieves the lowest loss at convergence at the 4B scale.

Long-Context Language Modeling

Long-context language modeling studies how language models can use, retain, and efficiently process information across very long sequences.

See all 180 papers

Long-Context Fine-Tuning with Limited VRAM
Jul 16, 2026Vladimir Fedosov, Aleksandr Sazhin, Artemiy Grinenko +1Fine-Tuning KV-Cache Offloading
BMW Group
Autonomy-of-Heads: Data-Free Sparse Attention from Frozen Query-Key Geometry
Aug 7, 2026Yehan Yang, Junyuan Shang, Yang Li +3Attention Head Analysis LLM Inference Acceleration
Institute of Computing Technology, Chinese Academy of Sciences · University of Chinese Academy of Sciences · Baidu Inc. +1
Multi-Head Low-Rank Attention
Mar 2, 2026Songtao Liu, Hongwu Peng, Zhiwei Zhang +2Long-Context Language Model Inference LLM Inference Acceleration
The Pennsylvania State University · University of Connecticut · Carnegie Mellon University +1

Jul 16, 2026J/K move · Enter open · S save

Long-Context Fine-Tuning with Limited VRAM

Vladimir Fedosov, Aleksandr Sazhin, Artemiy Grinenko, Frank Woernle

BMW Group

Fine-Tuning KV-Cache Offloading Memory-Efficient Fine-Tuning Efficient Language Model Training Attention Mechanisms Parameter-Efficient Fine-Tuning

Open paper AskPDF arXiv

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.

Abstract

Long-Context Language Modeling

Explore similar work

Long-Context Fine-Tuning with Limited VRAM

Autonomy-of-Heads: Data-Free Sparse Attention from Frozen Query-Key Geometry

Multi-Head Low-Rank Attention