FlashMLA

简介

FlashMLA 是 DeepSeek 的优化注意力核函数库，为 DeepSeek-V3 和 DeepSeek-V3.2-Exp 模型提供支持。本仓库包含以下实现：

稀疏注意力核函数

这些核函数驱动了 DeepSeek 稀疏注意力（DSA），如 这篇论文 中所述。

• 预填充阶段的 token 级别稀疏注意力
• 解码阶段的 token 级别稀疏注意力，采用 FP8 KV 缓存

密集注意力核函数

• 预填充阶段的密集注意力
• 解码阶段的密集注意力

最新消息

• 2025年9月29日 发布稀疏注意力核函数：随着 DeepSeek-V3.2 的推出，我们发布了相应的 token 级别稀疏注意力核函数。这些核函数支持模型的 DeepSeek 稀疏注意力（DSA），在预填充时最高可达 640 TFlops，在解码时可达 410 TFlops。我们还发布了一篇关于我们全新 FP8 稀疏解码核函数的深度解析博客，请点击 这里 查看。
• 2025年8月1日 SM100 上的 MHA 核函数：感谢 NVIDIA 的 PR，为我们带来了 SM100 上的 MHA 前向和后向核函数！
• 2025年4月22日 深度解析博客：我们非常乐意分享新 FlashMLA 核函数背后的技术细节！请查看我们的深度解析文章 这里。
• 2025年4月22日 性能更新：我们很高兴宣布 Flash MLA 的新版本发布，该版本在计算密集型工作负载上实现了 5% 至 15% 的性能提升，在 NVIDIA H800 SXM5 GPU 上最高可达 660 TFlops。新版本的接口与旧版本完全兼容，只需升级到新版本即可立即获得性能提升！🚀🚀🚀

性能

测试与基准测试 MLA 解码（稀疏与密集）：

python tests/test_flash_mla_dense_decoding.py
python tests/test_flash_mla_sparse_decoding.py

在内存受限的配置下，密集 MLA 解码核函数可达到最高 3000 GB/s 的吞吐量；而在计算受限的配置下，则可在 H800 SXM5 上以 CUDA 12.8 实现 660 TFLOPS 的计算性能。使用 bfloat16 进行矩阵乘法并采用 FP8 KV 缓存的 token 级别稀疏 MLA 解码核函数，在 H800 SXM5 上以 CUDA 12.8 可实现 410 TFLOPS 的计算性能，而在尚未完全优化的 B200 上则可达到最高 350 TFLOPS。

测试与基准测试 MHA 预填充（密集）：

python tests/test_fmha_sm100.py

据 NVIDIA 报道，它在 B200 上的前向计算最高可达 1460 TFLOPS，后向计算则可达 1000 TFLOPS。

测试与基准测试 MLA 预填充（稀疏）：

python tests/test_flash_mla_sparse_prefill.py

在 H800 SXM5 上以 CUDA 12.8 运行时，其前向计算最高可达 640 TFLOPS；而在 B200 上以 CUDA 12.9 运行时，则最高可达 1450 TFLOPS。

系统要求

• SM90 / SM100（参见下方的支持矩阵）
• CUDA 12.8 及以上版本（SM100 核函数需 CUDA 12.9+）
• PyTorch 2.0 及以上版本

支持矩阵：

核函数	GPU 架构	MLA 模式 [2]	KV 缓存格式
密集解码	SM90	MQA	BF16
稀疏解码	SM90 & SM100	MQA	FP8 [1]
密集预填充	SM100	MHA
稀疏预填充	SM90 & SM100	MQA

[1]: 关于如何使用 FP8 KV 缓存的更多详情，请参阅下方文档。

[2]: 此处“MLA 模式”指用于 MLA 计算的模式。MQA 代表多查询注意力模式（即 head_dim_k = 576，head_dim_v = 512），而 MHA 则代表多头注意力模式（即 head_dim_k = 192 / 128，head_dim_v = 128）。有关这些模式的详细解释，请参阅 DeepSeek V3.2 论文 的附录。

安装

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

使用方法

MLA 解码

要使用 MLA 解码核函数，需在解码循环之前调用一次 get_mla_metadata 来获取分块调度器元数据。然后在每一步解码中调用 flash_mla_with_kvcache。例如：

from flash_mla import get_mla_metadata, flash_mla_with_kvcache

tile_scheduler_metadata, num_splits = get_mla_metadata(
    cache_seqlens,
    s_q * h_q // h_kv,
    h_kv,
    h_q,
    is_fp8,
    topk,
)

for i in range(num_layers):
    ...
    o_i, lse_i = flash_mla_with_kvcache(
        q_i, kvcache_i, block_table, cache_seqlens, dv,
        tile_scheduler_metadata, num_splits,
        is_causal, is_fp8_kvcache, indices,
    )
    ...

其中

• s_q 是每个查询序列中的查询 token 数量。如果禁用了 MTP（推测性解码），则应为 1。
• h_kv 是键值头的数量。
• h_q 是查询头的数量。

FP8 KV 缓存： 如果将 is_fp8_kvcache 设置为 True，核函数将以“FP8 with scale”格式读取 KV 缓存（如下所述）。它会将缓存去量化为 bfloat16，并以 bfloat16 进行注意力计算。输出结果同样为 bfloat16。

在“FP8 with scale”格式中，每个 token 的 KV 缓存大小为 656 字节，结构如下：

• 前 512 字节： “量化 NoPE”部分，包含 512 个 float8_e4m3 值。
• 接下来的 16 字节： 缩放因子，包含 4 个 float32 值。第一个 float32 是前 128 个 float8_e4m3 值的缩放因子，第二个是接下来的 128 个的缩放因子，以此类推。
• 最后 128 字节： “RoPE”部分，包含 64 个 bfloat16 值。这部分未被量化，以保证精度。

有关量化和去量化的详细信息，请参阅 tests/quant.py。

稀疏注意力（indices 张量）： indices 张量（如果提供）可通过指示核函数仅对指定 token 进行注意力计算，从而实现 token 级别的稀疏注意力。

• 形状： indices 应为一个三维张量，形状为 (batch_size, seq_len_q, topk)。
• 格式： indices_in_kvcache[i][j][k] = （token t 所在页块的索引）* 页块大小 + （token t 在页块内的偏移量），其中 t 是第 i 个批次中第 j 个查询序列的第 k 个 token。由于页块的索引已编码进 indices_in_kvcache，因此核函数无需 block_table 参数。
• 无效条目： 将无效索引设为 -1。

返回值： 核函数返回 (out, lse)，其中：

• out 是注意力计算结果。
• lse 是每个查询头的注意力分数的 log-sum-exp 值。

完整示例请参阅 tests/test_flash_mla_decoding.py。

稀疏 MLA 预填充

对于稀疏 MLA 预填充核函数，可直接调用 flash_mla_sparse_fwd，参数如下：

• q: 查询张量，形状为 [s_q, h_q, d_qk]
• kv: 键值张量，形状为 [s_kv, h_kv, d_qk]
• indices: 索引张量，形状为 [s_q, h_kv, topk]
• sm_scale: 一个标量值

关于批处理的说明： 此核函数不支持批量维度。若需进行多批次推理，可对输入张量进行重塑，并调整 indices 参数以模拟批处理过程。

无效索引： 将 indices 中的无效条目设置为 -1 或任何大于等于 s_kv 的数值。

返回值及等效 PyTorch 代码： 该核函数返回 (out, max_logits, lse)。这与以下 PyTorch 操作等价：

Q: [s_q, h_q, d_qk], bfloat16
kv: [s_kv, h_kv, d_qk], bfloat16
indices: [s_q, h_kv, topk], int32

kv = kv.squeeze(1)  # [s_kv, d_qk], h_kv 必须为 1
indices = indices.squeeze(1)    # [s_q, topk]
focused_kv = kv[indices]    # 对于第 i 个序列 (s_q)，根据 indices[i, :] 从 KV 缓存中选择对应的 KV 令牌。此操作会生成一个形状为 [s_q, topk, d_qk] 的张量。

P = (Q @ focused_kv.transpose(-1, -2)) * sm_scale * math.log2(math.e)    # [s_q, h_q, topk]
max_logits = P.max(dim=-1) # [s_q, h_q]
lse = log2sumexp2(P, dim=-1, base=2)   # [s_q, h_q]，“log2sumexp2” 表示指数和对数均以 2 为底
S = exp2(P - lse)      # [s_q, h_q, topk]
out = S @ focused_kv  # [s_q, h_q, d_qk]

return (out, max_logits, lse)

完整示例请参阅 tests/test_flash_mla_prefill.py。

密集 MHA 预填充

该核函数实现了标准的密集多头注意力（MHA）前向和反向运算。可通过以下函数调用：

• flash_attn_varlen_func
• flash_attn_varlen_qkvpacked_func
• flash_attn_varlen_kvpacked_func

使用方法与 flash_attn 包类似。完整示例请参阅 tests/test_fmha_sm100.py。

致谢

FlashMLA 受到 FlashAttention 2&3 和 cutlass 项目的启发。

社区支持

MetaX

针对 MetaX GPU，请访问其官方网站：MetaX。

相应的 FlashMLA 版本可在以下地址找到：MetaX-MACA/FlashMLA。

Moore Threads

针对 Moore Threads GPU，请访问其官方网站：Moore Threads。

相应的 FlashMLA 版本已在 GitHub 上发布：MooreThreads/MT-flashMLA。

Hygon DCU

针对 Hygon DCU，请访问其官方网站：Hygon Developer。

相应的 FlashMLA 版本在此处提供：OpenDAS/MLAttention。

Intellifusion

针对 Intellifusion NNP，请访问其官方网站：Intellifusion。

相应的 FlashMLA 版本已在 Gitee 上发布：Intellifusion/tyllm。

Iluvatar Corex

针对 Iluvatar Corex GPU，请访问其官方网站：Iluvatar Corex。

相应的 FlashMLA 版本已在 GitHub 上发布：Deep-Spark/FlashMLA。

AMD Instinct

针对 AMD Instinct GPU，请访问其官方网站：AMD Instinct。

相应的 FlashMLA 版本可在以下地址找到：AITER/MLA。

引用

@misc{flashmla2025,
      title={FlashMLA: 高效多头潜在注意力核函数},
      author={Jiashi Li, Shengyu Liu},
      year={2025},
      publisher = {GitHub},
      howpublished = {\url{https://github.com/deepseek-ai/FlashMLA}},
}

FlashMLA 快速上手指南

FlashMLA 是 DeepSeek 推出的高性能注意力机制内核库，专为 DeepSeek-V3 及 V3.2 系列模型优化。它支持稀疏（Sparse）和稠密（Dense）两种注意力模式，并在 NVIDIA H800/B200 等最新硬件上实现了极致的推理性能。

环境准备

在开始之前，请确保您的开发环境满足以下硬件和软件要求：

• GPU 架构：
  • NVIDIA SM90 (如 H800, H100)
  • NVIDIA SM100 (如 B200, Blackwell 系列)
• CUDA 版本：
  • 基础要求：CUDA 12.8 及以上
  • SM100 专用：必须使用 CUDA 12.9 及以上
• 框架依赖：
  • PyTorch 2.0 及以上
• 操作系统：Linux (推荐 Ubuntu 20.04/22.04)

注意：不同内核模式对显存格式有特定要求。例如，稀疏解码（Sparse Decoding）支持 FP8 KV Cache，而稠密解码通常使用 BF16。

安装步骤

通过源码安装是最推荐的方式，可确保获取最新优化特性。

# 1. 克隆仓库
git clone https://github.com/deepseek-ai/FlashMLA.git flash-mla
cd flash-mla

# 2. 初始化子模块 (重要：包含底层优化代码)
git submodule update --init --recursive

# 3. 执行安装
pip install -v .

国内加速建议：如果克隆速度较慢，可使用 Gitee 镜像（如有）或配置 Git 代理。安装过程中若 pip 下载慢，建议添加国内源： pip install -v . -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

FlashMLA 提供了针对不同场景的 API，以下是两种最核心用法的极简示例。

1. MLA 解码 (Decoding)

适用于推理阶段的逐 token 生成。需先获取调度元数据，然后在循环中调用内核。

from flash_mla import get_mla_metadata, flash_mla_with_kvcache

# 第一步：在解码循环前获取元数据
# s_q: 每个序列的 query token 数 (无推测解码时为 1)
# h_kv: KV 头数，h_q: Query 头数
tile_scheduler_metadata, num_splits = get_mla_metadata(
    cache_seqlens,
    s_q * h_q // h_kv,
    h_kv,
    h_q,
    is_fp8,       # 是否使用 FP8 KV Cache
    topk,         # 稀疏模式的 topk 值
)

# 第二步：在每一层解码循环中调用
for i in range(num_layers):
    # ... 准备 q_i, kvcache_i, block_table 等张量 ...
    
    o_i, lse_i = flash_mla_with_kvcache(
        q_i, kvcache_i, block_table, cache_seqlens, dv,
        tile_scheduler_metadata, num_splits,
        is_causal,      # 是否因果掩码
        is_fp8_kvcache, # KV Cache 是否为 FP8 格式
        indices,        # 稀疏模式下的索引张量，稠密模式可为 None
    )
    # ... 后续处理 ...

关于 FP8 KV Cache： 若 is_fp8_kvcache=True，内核会自动将 FP8 格式的 KV Cache 反量化为 BF16 进行计算。FP8 格式结构为：512 字节量化数据 + 16 字节缩放因子 + 128 字节未量化的 RoPE 数据。

2. 稀疏 MLA 预填充 (Sparse Prefill)

适用于处理长上下文输入时的稀疏注意力计算。此接口不支持 batch 维度，多 Batch 需自行 reshape。

from flash_mla import flash_mla_sparse_fwd

# 准备输入张量
# q: [s_q, h_q, d_qk], bfloat16
# kv: [s_kv, h_kv, d_qk], bfloat16
# indices: [s_q, h_kv, topk], int32 (无效索引设为 -1)
# sm_scale: 缩放标量

out, max_logits, lse = flash_mla_sparse_fwd(
    q,
    kv,
    indices,
    sm_scale
)

参数说明：

• indices：指定每个 query token 需要关注的 key token 索引。形状为 (seq_len_q, h_kv, topk)。
• 返回值：返回 (out, max_logits, lse)，分别对应注意力输出、最大 logits 值和 log-sum-exp 值。

更多详细测试用例请参考仓库中的 tests/test_flash_mla_decoding.py 和 tests/test_flash_mla_prefill.py。