FlashMLA：DeepSeek 高性能 MLA 注意力算子库编译与验证指南

一、项目官方信息

FlashMLA 是 DeepSeek 的优化注意力算子库，支撑 DeepSeek-V4.1 在 NVIDIA GPU 与华为昇腾 NPU 上的推理。仓库包含：

• 带 FP8/FP4 KV cache 的 prefill 与 decoding 的 token 级稀疏注意力 kernel；
• 融合 Q-norm、Q-RoPE、attention、O-RoPE（共轭）与 cast 到 FP8 的融合 kernel；
• prefill 与反向的稠密注意力 kernel。

变更提示：2026.09.30 版本移除了对 Hopper 架构与前代模型（V3/V3.2/V4.0）的支持，并更改了 FP8/FP4 KV cache 格式，与旧版本不兼容。如需跑旧模型或旧 KV 格式，请切换到指定 commit。

二、官方图片

上图为 FlashMLA 仓库在 GitHub 上的社交预览图。

三、环境要求

NVIDIA 平台：

• NVIDIA SM100 / SM103 GPU；
• CUDA 13.1+（推荐 13.2+，量化 decoding kernel 需要原生 cvt.rn.bf16x2）；
• PyTorch 2.10+。

华为平台：

• 华为 Ascend 950 NPU；
• CANN 9.2.0+（提供 bisheng 编译器与 Ascend C 头文件）；
• torch_npu、PyTorch 2.10+。

四、编译安装步骤

1. 克隆并拉取子模块

git clone https://github.com/deepseek-ai/FlashMLA.git flash-mla
cd flash-mla
git submodule update --init --recursive

2. pip 安装（必须加 --no-build-isolation）

pip install -v . --no-build-isolation

--no-build-isolation 是必须的：setup.py 在构建时会 import torch（Ascend 平台还 import torch_npu），而仓库未把它们声明为 PEP 518 构建依赖。

• 构建目标平台自动检测（/dev/davinci_manager 表示 Ascend，否则为 CUDA），可用 FLASH_MLA_BUILD_TARGET_PLATFORM=CUDA 或 =ASCEND 覆盖；
• Ascend 平台需把 ASCEND_HOME_PATH 指向 CANN 根目录（默认 /usr/local/Ascend/ascend-toolkit/latest）；
• 构建后会对生成的 .so 做寄存器溢出检查，失败可 export FLASH_MLA_SKIP_REG_SPILL_CHECK=1 跳过。

五、使用方法

1. 验证/测试融合 kernel（Sparse）

python tests/test-fused-norm-rope-attn-rope-cast.py

（需 TileLang、Tile-Kernels、DeepGEMM。该融合 kernel 在 B200 + CUDA 13.3 上 prefill 达 1460 TFlops、decoding 950 TFlops。）

2. 测试 MLA prefill（Sparse）

python tests/test-sparse-prefill.py

B200 + CUDA 13.3 上达 1350 TFlops；Ascend 950 上达 410 TFlops（95% 硬件峰值）。

3. 测试 MLA decoding（Sparse）

python tests/test-sparse-decode.py

B200 上达 1024 TFlops；Ascend 950 上达 360 TFlops（83% 硬件峰值）。

4. 在推理中调用

from flash_mla import get_mla_metadata, flash_mla_with_kvcache

tile_scheduler_metadata, num_splits = get_mla_metadata()
# 在 decoding 循环中每步调用 flash_mla_with_kvcache

六、常见问题 FAQ

1. 需要什么 CUDA 版本？ CUDA 13.1+，推荐 13.2+（量化 decoding kernel 依赖原生 bf16x2 转换）。

2. 为什么必须 --no-build-isolation？ setup.py 构建时要 import torch/torch_npu，仓库没把它们声明为 PEP 518 构建依赖。

3. 支持哪些 GPU？ NVIDIA SM100/SM103，或华为 Ascend 950 NPU。

4. 构建时报 Register spilling detected？ 可 export FLASH_MLA_SKIP_REG_SPILL_CHECK=1 跳过检查再构建。

5. 和旧版/旧模型兼容吗？ 2026.09.30 版移除 Hopper 与 V3/V3.2/V4.0 支持，并改了 KV 格式，不兼容旧版；需用旧版请切到指定 commit。

6. 性能大概多少？ B200 上稀疏 prefill 最高 1350–1460 TFlops，decoding 950–1024 TFlops。

七、跳转链接

• GitHub 仓库：https://github.com/deepseek-ai/FlashMLA
• 论文：https://github.com/deepseek-ai/DeepSeek-V3.2-Exp
• 关联 DeepGEMM：https://github.com/deepseek-ai/DeepGEMM