优化大模型推理速度，让671B模型在单H20显卡上跑得更快更稳。
FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs
- 通过转置重排注意力计算，减少冗余操作提升效率。
- 64K序列下速度比FlashMLA快2.78倍，误差降低15.2倍。
- 可无缝接入主流框架，适合资源受限场景下的高效部署。
在单台多GPU服务器上部署DeepSeek-R1 671B模型时，多头潜在注意力（MLA）的高效推理面临挑战。本文提出FlashMLA-ETAP框架，引入高效转置注意力流水线（ETAP），通过转置重构注意力计算，使键值上下文长度与WGMMA操作中的M维对齐，显著减少冗余计算。在64K序列长度（批大小16）下，FlashMLA-ETAP相比FlashMLA实现2.78倍加速，分别优于FlashAttention-3和FlashInfer达5.24倍和4.94倍，同时保持数值稳定性，均方根误差（RMSE）仅为FlashAttention-3的1/15.2（1.25×10⁻⁵）。此外，ETAP设计支持与FlashAttention-3、FlashInfer等框架无缝集成，并提供详尽理论分析。本工作填补了资源受限推理中的关键空白，为中端GPU提供了可扩展的优化方案，推动硬件感知优化的广泛应用。代码已开源：https://github.com/pengcuo/FlashMLA-ETAP。
原文摘要 · Abstract (English)
Efficient inference of Multi-Head Latent Attention (MLA) is challenged by deploying the DeepSeek-R1 671B model on a single Multi-GPU server. This paper introduces FlashMLA-ETAP, a novel framework that enhances MLA inference for the single-instance deployment scenario on NVIDIA H20 GPUs. We propose the Efficient Transpose Attention Pipeline (ETAP), which reconfigures attention computation through transposition to align the KV context length with the \(M\)-dimension in WGMMA operations, significantly reducing redundant computations. FlashMLA-ETAP achieves a 2.78x speedup over FlashMLA at 64K sequence length (batch size 16), with 5.24x and 4.94x improvements over FlashAttention-3 and FlashInfer, respectively, while maintaining numerical stability with a 15.2x lower RMSE (\(1.25 \times 10^{-5}\)) than FlashAttention-3. Furthermore, ETAP's design enables seamless integration into frameworks like FlashAttention-3 and FlashInfer, supported by a detailed theoretical analysis. Our work addresses a critical gap in resource-constrained inference, offering a scalable solution for mid-tier GPUs and paving the way for broader adoption in hardware-aware optimization. Code is available at https://github.com/pengcuo/FlashMLA-ETAP.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。