The Pulse
Baidu Researchers Propose Virtual Head Attention to Cut KV Cache
Baidu- and Peking University-affiliated researchers propose Virtual Head Attention, a design that cuts language models’ key-value cache while aiming to preserve benchmark quality. Their paper reports results across four model scales and a s
AI.info Team · Written with AI ·
Reducing the memory used by language-model attention can come at a cost: lower benchmark scores, more training compute, or both. A team of researchers affiliated with Baidu Inc. and Peking University proposes Virtual Head Attention (VHA), which aims to reduce that tradeoff by cutting the key-value cache without giving up the model quality measured in its tests. The authors submitted their paper to arXiv on October 13, 2026.
“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.”
Xibo Ding et al., “Virtual Head Attention”
Two training operations, standard attention at inference
VHA adds two linear operations during training. Q Premix transforms query heads within each key-value group, letting fewer physical query heads represent more “virtual” heads; Linear Postmix mixes information across attention-head outputs. The approach is designed to recover some of the diversity lost when the architecture uses fewer heads.
The authors say both operations can be folded into the model’s query and output projections after training. That leaves a standard grouped-query attention design at inference, rather than requiring a VHA-specific attention kernel. The paper says this lets the model reuse existing GQA optimizations; the claim describes the architecture, not a measured speedup on deployed systems.
Four model scales show smaller caches
The main comparison covers four model scales, from 0.6 billion parameters to an 8B-A1B mixture-of-experts model. Each was pretrained on 50 billion tokens with a 4,096-token context, then evaluated on eight zero-shot multiple-choice benchmarks. VHA’s average score exceeded the matching GQA baseline at every scale, by 0.05 to 0.81 percentage points, according to the paper.
At 1.7 billion parameters, VHA used 1 KB of BF16 key-value cache per layer per token, compared with 4 KB for GQA-8. Its eight-task average was 54.50, against 54.09 for GQA-8. A GQA version compressed to the same 1 KB cache but without VHA’s training operations scored 52.38, illustrating why simply reducing the head count does not reproduce the reported result.
The authors also calculate 7.9% fewer training and prefill FLOPs for VHA than the GQA baseline at that 1.7B scale and 4,096-token context. Across the dense-model tests, the paper reports a fourfold smaller cache than GQA-8; at the MoE scale, it reports a twofold reduction against GQA-4. Those savings depend on the architecture and comparison baseline, rather than describing one universal reduction.
Larger-model and long-context results use separate tests
In a separate 30B-A3B MoE experiment, trained on 500 billion tokens at an 8,192-token context, VHA scored 50.36 on the paper’s macro-average of 12 base-model benchmarks. The reported averages were 49.66 for multi-head latent attention and 47.76 for GQA-4; the listed cache sizes were 1 KB, 1.12 KB and 2 KB per layer per token, respectively.
A different experiment tested long-context performance on a 10B-A1B MoE model using RULER. VHA averaged 84.4 across context lengths from 8K to 128K, compared with 73.0 for GQA, while using two rather than four key-value heads in full-attention layers. Because these are separate studies with different model setups, their figures should not be read as one directly comparable benchmark series.
PaddlePaddle’s implementation has a parallelism limit
The team has published an implementation in the PaddlePaddle VHA repository, along with a small GPT example and tests. The repository says its current implementation requires tensor-parallel degree 1, a practical limitation for users seeking to run the attention design across multiple model-parallel partitions.
The evidence presented so far is primarily benchmark accuracy and calculated compute, rather than a reported production-serving speed test. VHA’s central engineering bet is that training-time mixing can restore quality while leaving the inference model compatible with ordinary GQA kernels; the paper’s tables and the repository’s single-degree constraint define what has been demonstrated publicly.