deepseek-ai/FlashMLA

FlashMLA – 用于 DeepSeek 模型的高性能注意力核心

简介 – FlashMLA 是一个 C++/CUDA（及华为 Ascend）函数库，实现了 DeepSeek 大型语言模型（最近期为 DeepSeek-V4.1-Flash 模型）所使用的高度优化注意力算子。它提供：

• 稀疏注意力核心（token 级别的“DeepSeek Sparse Attention”），适用于 prefill 阶段（模型处理初始提示时）与 decoding 阶段（一次生成一个 token）。这些核心支持 FP8 或 FP4 量化的 key-value 缓存，在保持准确度的同时大幅降低内存带宽需求。
• 密集多头注意力 (MHA) 核心，用于标准（非稀疏）的 prefill 与反向传播。
• 融合核心，将 Q-norm、RoPE（旋转位置编码）、核心注意力计算、输出端的共轭 RoPE 以及最终转型为 FP8 的步骤合并在单次启动中。这消除了许多微小核心带来的额外开销。
• 支持 NVIDIA GPU (SM100/SM103, CUDA 13.1+) 与 华为 Ascend 950 NPU (通过 CANN 9.2+ 与 torch_npu)。

重要性 – 大型语言模型大部分的推理时间都花在注意力运算上。通过以低精度 (FP8/FP4) 重写注意力并利用稀疏性，FlashMLA 在现代硬件上可达到 >1,000 TFlops，接近设备的理论峰值。这意味着服务 DeepSeek-V4.1 等模型时能获得更快的响应时间与更低的成本。

如何使用

1. 安装 – Clone 仓库，拉取子模块，并执行 pip install -v . --no-build-isolation。构建脚本会自动检测您是在 CUDA 还是 Ascend 环境，并编译对应的扩展功能。
2. Prefill (稀疏) – 使用您的 query tensor、KV 缓存以及一个 3-D indices tensor（告知核心每个 query token 应关注哪些 KV token）来调用 flash_mla_sparse_fwd(q, kv, indices, …)。
3. Decoding (稀疏) – 首先通过 get_mla_metadata() 获取调度器元数据，然后在每个生成步骤重复调用 flash_mla_with_kvcache(...)。该函数预期使用 README 中描述的 FP8/FP4 KV 缓存布局。
4. 密集 MHA – 使用熟悉的 flash_attn_varlen_func / flash_attn_varlen_qkvpacked_func API（与热门的 flash_attn 软件包兼容），但请注意额外的 head_dim_qk 参数，且返回值包含输出与 log-sum-exp (lse)。
5. 融合核心 – 在将 Q-projection 与 Wv-projection 权重排列为融合核心所需的布局后（辅助工具位于 fused_norm_rope_attn_rope_cast.permute_*），调用 prefill(...) 或 decode(...) 以获得适用于后续 FP8-einsum 的 FP8 输出。

性能亮点（根据仓库报告）

• 稀疏 prefill：在 NVIDIA B200 GPU (CUDA 13.3) 上高达 1350 TFlops。在 Ascend 950 NPU 上为 410 TFlops (约为硬件峰值的 95%)。
• 稀疏 decoding：在 B200 上高达 1024 TFlops；在 Ascend 950 上为 360 TFlops (约为峰值的 83%)。
• 融合 norm-RoPE-attn-RoPE-cast：在 B200 上高达 1460 TFlops (prefill) 与 950 TFlops (decoding)，比未融合版本额外提升 10-15%。
• 密集 MHA：在 B200 上前向传播高达 1460 TFlops，反向传播为 1000 TFlops。

关键技术细节

• KV 缓存格式 – 仅支持 FP8 (float8_e4m3) 或 FP4 (e2m1) 量化缓存。布局存储原始量化值，后接每个 token 的缩放因子，实现即时快速反量化。
• Indices tensor – 一个 3-D 整数 tensor，编码每个 query token 应关注哪些 KV“页面”。无效项目标记为 -1 并由核心忽略。
• 寄存器溢出安全 – 构建 CUDA 扩展后，构建系统会检查寄存器溢出；除非设置 FLASH_MLA_SKIP_REG_SPILL_CHECK=1，否则构建将会失败。
• 兼容性说明 – 2026 年 9 月的版本取消了对 NVIDIA Hopper GPU 与旧版 DeepSeek 模型 (V3/V3.2/V4.0) 的支持，并更改了 KV 缓存格式。需要这些旧配置的用户必须 checkout README 中链接的旧版 commit。

深入了解 – 仓库包含深入的 markdown 报告 (docs/…)，解释了算法与底层优化，以及可用于验证硬件性能的测试脚本 (tests/*.py)。

---

FlashMLA 本质上是 DeepSeek LLM 推理流水线中性能关键的构建模块，同时针对尖端的 NVIDIA GPU 与华为 Ascend NPU 进行了优化。