arXiv:2508.15881cs.LGcs.AI2025-08被引 2

让大模型推理更快:通过分片压缩缓存,实现高效分布式推理。

TPLA: Tensor Parallel Latent Attention for Efficient Disaggregated Prefill and Decode Inference

  • 将潜在表示和注意力头跨设备分片,独立计算后聚合
  • 在32K上下文下提速1.79倍到1.93倍,性能几乎不变
  • 兼容现有模型,无需重训练,适合大规模推理部署

多头潜在注意力(MLA)将键值状态压缩为低秩潜在向量,仅缓存该向量以降低内存占用。但在张量并行(TP)中,注意力头分布在多个设备上,每个设备仍需加载完整缓存,削弱了MLA相比分组查询注意力(GQA)的优势。我们提出张量并行潜在注意力(TPLA):将潜在表示和每头输入维度跨设备分片,各分片独立执行注意力,再通过all-reduce合并结果。TPLA保持压缩缓存优势的同时实现TP效率。与分组潜在注意力(GLA)不同,TPLA中每头仍使用完整潜在表示,保留更强表达能力。TPLA可直接用于预训练的MLA模型:支持MLA式prefill,并实现无需重训练的高效张量并行解码。在TP切分前施加简单正交变换(如Hadamard变换或PCA),可显著减轻跨分片干扰,带来极小精度损失。对DeepSeek-V3和Kimi-K2,TPLA在32K上下文长度下分别实现1.79倍和1.93倍加速,同时在常识和LongBench基准上保持性能。结合FlashAttention-3,可实现端到端实用加速。

原文摘要 · Abstract (English)

Multi-Head Latent Attention (MLA), introduced in DeepSeek-V2, compresses key-value states into a low-rank latent vector, caching only this vector to reduce memory. In tensor parallelism (TP), however, attention heads are computed across multiple devices, and each device must load the full cache, eroding the advantage of MLA over Grouped Query Attention (GQA). We propose Tensor-Parallel Latent Attention (TPLA): a scheme that partitions both the latent representation and each head's input dimension across devices, performs attention independently per shard, and then combines results with an all-reduce. TPLA preserves the benefits of a compressed KV cache while unlocking TP efficiency. Unlike Grouped Latent Attention (GLA), every head in TPLA still leverages the full latent representation, maintaining stronger representational capacity. TPLA is drop-in compatible with models pre-trained using MLA: it supports MLA-style prefilling and enables efficient tensor-parallel decoding without retraining. Applying simple orthogonal transforms -- e.g., the Hadamard transform or PCA -- before TP slicing further mitigates cross-shard interference, yielding minimal accuracy degradation. By reducing the per-device KV cache for DeepSeek-V3 and Kimi-K2, we achieve 1.79x and 1.93x speedups, respectively, at a 32K-token context length while maintaining performance on commonsense and LongBench benchmarks. TPLA can be implemented with FlashAttention-3, enabling practical end-to-end acceleration.

大模型推理张量并行缓存优化注意力机制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。