提出混合注意力核，让大模型推理更快更省显存。
TyphoonMLA: A Mixed Naive-Absorb MLA Kernel For Shared Prefix
- 结合两种计算方式，按需分配任务：计算密集用高效模式，非共享部分用省带宽模式。
- 在NPU和GPU上使注意力吞吐提升最高达3.24倍，显存仅增加3%。
- 适合追求推理速度与显存效率的LLM部署场景。
多头潜在注意力（MLA）是DeepSeek-v3、Kimi K2等先进大模型采用的新式注意力机制。其支持两种功能等效但计算特性不同的实现方式：朴素（naive）与吸收（absorb）。朴素核（如FlashAttention）因计算效率高常用于训练和预填充；而现有解码核（如FlashMLA）则依赖吸收法以减少显存带宽消耗。然而，吸收法的计算密集型特性限制了注意力计算中共享前缀等数据重用机会的收益。本文提出TyphoonMLA，一种融合朴素与吸收范式的混合方法：对计算密集部分使用朴素法，对非共享部分使用吸收法，从而兼顾计算效率与带宽优化。实验表明，该方法在NPU和GPU上将MLA架构的注意力吞吐分别提升最高达3倍和3.24倍，且仅带来3%的显存开销。
原文摘要 · Abstract (English)
Multi-Head Latent Attention (MLA) is a recent attention mechanism adopted in state-of-the-art LLMs such as DeepSeek-v3 and Kimi K2. Thanks to its novel formulation, MLA allows two functionally equivalent but computationally distinct kernel implementations: naive and absorb. While the naive kernels (e.g., FlashAttention) are typically preferred in training and prefill for their computational efficiency, existing decoding kernels (e.g., FlashMLA) rely on the absorb method to minimize HBM bandwidth usage. However, the compute-bound nature of the absorb implementations prohibits performance benefits from data reuse opportunities in attention calculations, such as shared prefixes. In this work, we introduce TyphoonMLA, a hybrid approach that combines naive and absorb formulations to harness the strengths of both. TyphoonMLA effectively leverages the shared prefix by applying the naive formulation to the compute-bound parts of attention calculations, while reducing the bandwidth requirements for non-shared parts by using the absorb formulation. As a result, TyphoonMLA improves the throughput of attention calculations in MLA architectures by up to 3x and 3.24x on NPU and GPUs, with only a 3% overhead in HBM size.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。