FlashMLA

项目简介

FlashMLA 是 DeepSeek 的优化注意力算子库，支撑 DeepSeek-V4.1 在 NVIDIA GPU 和华为昇腾 NPU 上的推理。它包含用于 prefill 和 decoding 的 token 级稀疏注意力算子（使用 FP8 或 FP4 KV 缓存）、将 Q-norm（仅 V4 使用，V4.1 不使用）、Q-RoPE、核心注意力、O-RoPE（共轭）和 FP8 转换融合在一起的算子，以及用于 prefill 和反向传播的稠密注意力算子。稀疏算子支撑 DeepSeek Sparse Attention（DSA）。

核心功能

• 面向 prefill 和 decoding 的 token 级稀疏注意力算子，支持 FP8 或 FP4 KV 缓存
• 融合 Q-norm、Q-RoPE、核心注意力、O-RoPE（共轭）和 FP8 转换的融合算子
• 用于 prefill 和反向传播的稠密注意力算子
• 稀疏算子支撑 DeepSeek Sparse Attention（DSA）
• 可在 NVIDIA GPU 和华为昇腾 NPU 上运行

适用场景

• 在 NVIDIA GPU 和华为昇腾 NPU 上推理 DeepSeek-V4.1
• 面向 DSA 的 token 级稀疏注意力 decoding 与 prefill
• 稠密多头注意力的前向与反向计算
• 使用预排列的 Q_b 和 Wv 权重的 DeepSeek-V4.1 融合注意力流水线

快速开始

• 克隆仓库、初始化子模块，并使用 pip 以 --no-build-isolation 安装。
• DSA decoding 时，在 decoding 循环前调用一次 get_mla_metadata，然后在每个 decoding 步骤调用 flash_mla_with_kvcache。
• DSA prefill 时，直接调用 flash_mla_sparse_fwd，并传入 query、key-value、indices、scale 及可选参数。
• 稠密 MHA prefill 可使用 flash_attn_varlen_func、flash_attn_varlen_qkvpacked_func 或 flash_attn_varlen_kvpacked_func。

部署与要求

• NVIDIA 平台需要 SM100 或 SM103 GPU、CUDA 13.1 及以上版本，以及 PyTorch 2.10 及以上版本。
• 华为平台需要昇腾 950 NPU、CANN 9.2.0 及以上版本、torch_npu，以及 PyTorch 2.10 及以上版本。
• 安装时需要使用 --no-build-isolation，因为 setup.py 在构建期间导入 torch，而仓库未将其声明为 PEP 518 构建要求。

采用前须知

• 许可证：MIT。在使用、修改或分发项目前，请确认其具体条款。