VHA面向长上下文模型降本
AI 总结
r/LocalLLaMA 帖子介绍了 Virtual Head Attention（VHA），该方法通过 Q Premix 和 Linear Postmix 应对 KV 缓存、计算量与模型质量之间的取舍。作者称，VHA 在推理时可折叠为标准 GQA-2 架构，并在四个预训练规模上超过 GQA 基线；在其测试中，稠密规模的 KV 缓存较 GQA-8 减少 4 倍，MoE 规模较 GQA-4 减少 2 倍，1.7B、4K 上下文实验的训练和预填充 FLOPs 降低 7.9%。
深度见解
VHA 提出了一种兼顾长上下文模型缓存与计算成本、并可沿用现有 GQA 推理优化的方案，但质量和效率结果目前主要来自作者报告。
| 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. Please read below items as above text's formatting is mess. submitted by /u/pmttyji [link] [comments] |