FlashMLA：DeepSeek 开源的基于 C++ 的高效多头潜在注意力（MLA）算子库

摘要

1) 一句话总结 FlashMLA 是 DeepSeek 开源的基于 C++ 的高效多头潜在注意力（MLA）算子库，为 DeepSeek-V3 和 V3.2-Exp 模型提供底层的稀疏与稠密注意力计算支持。

2) 核心要点

项目基础：采用 MIT 协议开源，主要语言为 C++，定位为 AI 模型与推理基础设施，GitHub 拥有超 1.2 万 Stars。
核心功能：提供稠密（Dense）和稀疏（Sparse）注意力算子，均完整覆盖 Prefill（预填充）和 Decoding（解码）阶段。
H800 性能表现：在 H800 SXM5 (CUDA 12.8) 上，稠密 MLA 解码算子最高达 3000 GB/s（访存瓶颈）或 660 TFLOPS（计算瓶颈）；稀疏 MLA 解码达 410 TFLOPS；稀疏 MLA 预填充达 640 TFLOPS。
B200 性能表现：在 B200 上，稠密 MHA 预填充前向计算达 1460 TFLOPS，反向达 1000 TFLOPS；稀疏 MLA 预填充达 1450 TFLOPS。
环境与硬件要求：仅支持 SM90 和 SM100 架构 GPU；要求 CUDA 12.8 及以上（SM100 需 CUDA 12.9+），以及 PyTorch 2.0+。
FP8 KV Cache 机制：稀疏解码算子支持 FP8 KV Cache，每 Token 占用 656 字节（包含 512 字节量化 NoPE、16 字节缩放因子和 128 字节未量化 RoPE），计算时反量化为 BF16 执行。
广泛的硬件生态适配：项目已获得 MetaX、摩尔线程、海光 DCU、云天励飞、天数智芯及 AMD Instinct 等多家国产与国际硬件厂商/社区的移植与支持。

3) 风险与不足

B200 优化不足：文档明确指出，稀疏 MLA 解码算子在 B200 上的表现（最高 350 TFLOPS）尚未经过充分优化。
批处理维度限制：稀疏 MLA 预填充算子（Sparse MLA Prefill）原生不支持批处理维度（batch dimension），进行多批次推理时需要手动重塑输入张量并调整索引参数。

FlashMLA: Efficient Multi-head Latent Attention Kernels