02 · FlashMLA: MLA 解码为什么受计算限制, 以及 seesaw、FP8 KV 与稀疏化怎么落到 kernel

FlashMLA 于 2025-02-24 在 DeepSeek 开源周第一天公开, 作者包括 Jiashi Li、Shengyu Liu、Yuanhang Sun, 采用 MIT 许可证. 2026-09-30 的 main 版本加入昇腾 kernel, 同时移除了 Hopper 与 V3/V3.2/V4.0 支持并修改 KV cache 格式. 早期 Hopper 解码 kernel 对应提交 ba89a34, 当前 SM100 稀疏 kernel 则对应 main 分支; 两代实现需要按提交区分, 不能混用接口和缓存布局.

FlashMLA 对外只做一件事: 注意力核那一截 $\mathrm{softmax}(QK^\top)V$ 的高性能实现, 不含前后的线性投影. 它支撑的模型从 2025 年的 DeepSeek-V3, 到 V3.2 的稀疏注意力, 再到 2026 年的 V4.1. 这条演进线上, kernel 的形态换过三轮: 先是 bfloat16 稠密 MLA 解码, 再是配 FP8 KV 的稀疏解码, 最终是把 norm、RoPE、cast 融进来的 fused kernel. 底层硬件也从 Hopper (SM90) 换到 Blackwell (SM100), 随后增加华为昇腾实现. 每次变化都同时改动数据布局、调度方式和可支持的模型版本.

1. 它解决的瓶颈: 解码阶段 MLA 为什么受计算限制

把 FlashMLA 和普通的 FlashAttention 分开看, 关键在一个反直觉的事实: 自回归解码阶段, 注意力一般受访存带宽限制, 而 MLA 的解码 kernel 却落在计算受限区间. 这个差别决定了整套优化的方向, 也决定了 kernel 的形状. 要讲明白, 得先看矩阵吸收把 MLA 变成了什么.

两类瓶颈需要不同的优化手段. 受访存限制的 kernel 要减少读写并提高带宽利用率, 例如加快 KV 搬运、提高缓存命中; 受计算限制的 kernel 卡在 Tensor Core 的矩阵乘吞吐, 要让标量运算与矩阵乘重叠, 尽量缩短 Tensor Core 的空闲时间. MLA 解码属于后者, 当时主流解码注意力 kernel 则主要围绕省带宽调优, 因此 FlashMLA 采用了独立实现.

1.1. 矩阵吸收后 MLA 退化成 MQA, 以及 576/512 的形状从哪来

MLA 的低秩潜变量与解耦 RoPE 推导见 llm-guide 的 03-MLA-低秩潜变量与解耦RoPE, 矩阵吸收的工程实现见 04-MLA-矩阵吸收与工程实现. FlashMLA 接收的形状由这两项设计直接决定: MLA 把每个 token 的 KV 压成一个低秩潜变量 $\mathbf c^{KV}$tctKV​ (维度 $d_c=512$), 再加一段解耦 RoPE 的 key $\mathbf k^R_t$ (维度 $d^R_h=64$). 推理时, 原本要为每个 head 还原出 $\mathbf k, \mathbf v$ 的上投影 $W${UK}, W_{UV} 可以被「吸收」进 query 侧和输出侧的矩阵: $W_{UK}$ 并进 $W_{UQ}$, $W_{UV}$ 并进 $W_O$. 吸收后, 注意力分数可以由 query 直接与潜变量做内积, 无须先展开成每个 head 的 key.

吸收的结果是形状塌缩. 社区里 杨文博的分析 把这步讲得很直接: q 的形状从 $128\times192$ 变成 $128\times576$, k 从每 head 一份变成一份共享的 $576$ 维, v 变成 $512$ 维, 而且 v 就是 k 的前 $512$ 维, 第二份 KV 不再单独存在. 这里的 $576 = 512 + 64$ 是潜变量 $d_c$ 加 RoPE 的 $d^R_h$, $512$ 是潜变量本身. 一旦所有 query head 共享同一份 $576$ 维的 key、同一份 $512$ 维的 value, 这就退化成了 Multi-Query Attention. FlashMLA 的接口如实反映了这点: flash_mla_with_kvcache 的 docstring 写明 head_dim 必须 512、num_heads_k 必须 1 (只支持 MQA), 见 flash_mla_interface.py.

需要划清一条边界: 矩阵吸收这步发生在模型建模代码里, 不在 FlashMLA 里. 开源推理框架约定的 attention 算子只负责 $\mathrm{softmax}(QK^\top)V$ 那一行, 前后的 $W_{UQ}, W_{UK}, W_{UV}, W_O$ 投影由框架的 MLA 模块 (如 SGLang 的 DeepseekV2AttentionMLA) 实现. 所以 FlashMLA 收到的 q 已经是吸收后的 $576$ 维 latent query, kv 已经是 $576$ 维的 latent. 这条边界解释了一个容易误解的点: 库名叫 FlashMLA, 但它导出的 flash_mla_with_kvcache 本质是个输入形状特殊的 MQA 算子, MLA 的灵魂 (低秩压缩) 在它上游已经做完了.

这条边界也顺带解释了 KV cache 为什么能小到可以放进显存. 吸收前, 每个 token 要为 128 个 head 各存一份 key 和 value, 若按 MHA 的 $d_h=128$ 算, 一个 token 的 KV 是 $128\times128\times2 = 32768$ 个元素; 吸收后只存一份 $576$ 维的 latent, 一个 token 是 $576$ 个元素, 两者差约 57 倍. 正是这个压缩比让 128K 上下文的单请求 KV 从不可接受降到 GiB 量级, 也正是它让后面 FP8 量化的收益 (再省一半字节) 变得有意义: 先靠低秩把量级压下来, 再靠量化抠细节, 两步叠在一起才撑得起长上下文的大 batch 推理.

1.2. 计算访存比约 $2 h_q s_q$ 与 128 这个阈值

把形状代进 roofline 分析, 就能看出为什么解码阶段也会计算受限. 官方 20250422 deep-dive 给了完整推导: 设 q 头数 $h_q$, 每请求 q token 数 $s_q$, kv token 数 $s_k \gg h_q s_q$, K 与 V 的头维 $d_k, d_v$. 浮点运算量约 $2 h_q s_q s_k (d_k + d_v)$, 访存量 (字节) 约 $2 s_k d_k$ (因为 $s_k$ 份 kv 占大头, 每份 $d_k$ 个 bfloat16). 两者相除, 计算访存比约 $h_q s_q \cdot \frac{d_k + d_v}{d_k} \approx 2 h_q s_q$.

这个比值是普通 MHA 解码的几十倍. 普通 MHA 每个 query head 自带一份 KV, 读一份 KV 只喂一个 head 算, 计算访存比接近 1, 所以解码受带宽限制. MLA 退化成 MQA 后, 一份 $576$ 维的 latent 要喂 $h_q = 128$ 个 query head 一起算, 访存被这 128 个头摊薄, 算力这边却实打实地算了 128 遍. 代进 H800 的数字: 峰值带宽 $3.35$ TB/s, 降频后实际峰值算力约 $865$ TFlops, 临界点在 $h_q s_q \ge \frac{1}{2}\cdot\frac{865}{3.35} = 128$. DeepSeek 的线上推理系统在解码实例上不开张量并行, 所以 $h_q$ 就是满的 $128$, 正好压在计算受限这一侧.

这解释了为什么整套 kernel 的优化目标是「让 Tensor Core 一直忙」, 而不是「省带宽」. 一个受计算限制的 kernel, 瓶颈在 Tensor Core 的 GEMM 吞吐; 要提速, 就得让 CUDA Core 上的 softmax、缩放这些标量运算与 Tensor Core 的矩阵乘重叠起来, 别让 Tensor Core 停下来等 softmax. 下面的 seesaw 调度就是冲这个目标设计的. 这里也埋了一个伏笔: 一旦 KV 用 FP8 存, 反量化又要占 CUDA Core, 计算受限的瓶颈会从「softmax 挡路」变成「反量化挡路」, 这是第 4 节 Hopper crossover 要解决的问题.

2. 解码 kernel 的数据流: paged KV、split-KV、tile 调度与 combine

受计算限制的 kernel 要跑满, 前提是把活均匀铺到所有 SM 上, 并且让每个 SM 的输入数据按需、按时到位. FlashMLA 解码路径的数据流由三件事拼成: 分页的 KV 缓存与 indices 寻址、把长 KV 序列切成多段的 split-KV、以及负责分配与合并的 tile scheduler 加 combine kernel. 这三件事从 2025-02 的首发版一直延续到现在, 接口细节随版本改过, 但骨架没变.

这三件事各管一段. 分页 KV 管「数据在显存里怎么放、怎么按下标取」, 解决的是长序列、变长 batch 下的显存碎片问题; split-KV 管「一条长请求怎么拆给多个 SM 并行」, 解决的是单 SM 串行扫 128K token 太慢、又喂不满 SM 的问题; tile scheduler 加 combine 管「谁算哪一段、算完怎么拼回去」, 解决的是负载均衡和部分结果的正确合并. 三者配合才让计算受限的 kernel 既不空闲也不算错. 下面分两块讲: 先是 KV 的组织与寻址, 再是调度与合并.

2.1. paged KV 块与 indices 寻址

首发版 (2025-02) 的 README 写明 KV 缓存采用 BF16 分页 (paged) 布局, 块大小为 64. KV 被切成每块 64 个 token 的固定 block, block_table 记录每条序列对应的物理 block. 这套布局来自 vLLM 的 PagedAttention, 不同长度的序列可以共享物理显存, 无须按最长序列预留连续空间. 历史 kernel 代码中的断言 FLASH_ASSERT(params.page_block_size == Kernel_traits::kBlockN) 要求页大小等于 kernel 的 N 维 tile, 见 flash_fwd_mla_kernel.h (b31bfe7).

到了稀疏版 (2025-09 之后), 寻址方式变了. 稀疏 kernel 不再用 block_table 做间接寻址, 而是让调用方把 page block 下标直接编进 indices. 当前接口的契约写得很细: indices_in_kvcache[i][j][k] = (token t 所在 page block 的下标) * page_block_size + (token t 在该 block 内的偏移), 其中 $t$ 是第 $i$ 个 batch、第 $j$ 条 query 序列的第 $k$ 个 token. 因为物理地址已经算进 indices, kernel 内部不再需要 block_table, 但为兼容旧签名仍要求传 (传 None 即可). 仓库提供了 abs_indices2indices_in_kvcache 做这步转换, 把逻辑下标 (0 到 $s_k-1$) 映射成带 page block 的物理下标, 见 tests/quant.py.

无效项的处理是个容易踩的坑. 解码 kernel 只把恰好等于 $-1$ 的下标当无效, 不做任何上界检查; 任何其他越界正值都会产生越界的 TMA 地址, 直接非法访存. 这条在 flash_mla_with_kvcache 的 docstring 里专门标了出来. prefill kernel 宽松一些, 无效项可以是 $-1$ 或任意 $\ge s_{kv}$ 的数. 这个差异来自两条路径的实现: decode 为了省掉边界判断的分支把检查外包给了调用方, prefill 则在 gather 前做了 mask. 对接模型时, 调用方必须保证 decode 的无效位严格填 $-1$.

页大小选 64 不是随手定的. 块越大, block_table 越短、寻址间接层越薄, 但显存碎片越多 (一条序列的尾部不满一块也要占整块); 块越小碎片越少, 但寻址表变长、每次 gather 的随机访存更碎. 64 是在这两端之间的折中, 也和 kernel 里 N 维 tile 的粒度对齐, 让一块 KV 正好是一次 TMA 搬运的单位. 稀疏化之后这层的意义又变了: 因为 top-k 下标本来就是随机分布的, 稀疏 kernel 每次只按 indices 取少量 token, 页大小更多是影响 indices 编码时 block 下标乘以多少, 而不再决定连续扫描的粒度.

2.2. tile scheduler、split-KV 与 combine

单条请求的 KV 可能很长 (128K token), 一个 SM 串行扫完太慢, 也喂不满所有 SM. split-KV 的做法是把一条请求的 KV 序列沿 token 维切成若干段, 每段分给一个 SM 算出部分的注意力输出和对应的 log-sum-exp, 最终再合并. FlashMLA 用一个 tile scheduler 来决定怎么切、怎么分. 首发接口里 get_mla_metadata(cache_seqlens, num_heads_per_head_k, num_heads_k) 在解码循环前跑一次, 根据各请求的序列长算出 tile_scheduler_metadata (形状 (num_sm_parts, TileSchedulerMetaDataSize)) 和 num_splits (形状 (batch_size + 1)), 目标是让各 SM 的负载尽量均衡.

当前接口把这步做成了惰性的. get_mla_metadata() 不再带参数, 只返回一个空的 FlashMLASchedMeta 占位对象, 真正的调度元数据在第一次调用 flash_mla_with_kvcache 时由 C++ 侧 sparse_decode_fwd 生成并写回. 元数据对象里带了一份 Config (记录 b、s_q、h_q、page_block_size、topk 等), 后续调用会断言形状一致, 复用同一份元数据能省掉重算. 文档特别提醒: topk_length 和 extra_topk_length 的值不在校验范围内, 复用带不同 topk_length 的元数据会静默沿用过时的 split-KV 调度. 还有一个 enable_batch_invariant 开关, 置 True 会关掉 split-KV 路径, 让结果不依赖 batch 怎么切分, 代价是放弃 split 带来的并行度, 见 flash_mla_interface.py.

合并由一个独立的 combine kernel 完成. 历史代码里 flash_fwd_splitkv_mla_combine_kernel 以 dim3(params.b * params.h * params.seqlen_q) 的网格启动, 每个线程块负责一个 (batch, head, q token) 的输出, 把各 split 的部分结果按 log-sum-exp 重新归一后相加, 见 flash_fwd_mla_kernel.h (b31bfe7). split 产出的 kernel 和 combine kernel 之间用 programmatic dependent launch 重叠: combine 不必等 split 全部结束, 可以在 split 快写完时提前启动, 省掉一次 kernel 启动的串行等待. num_splits 的分桶上限也随版本调过, 2026-07 的提交把 decode-combine 的 num_splits 桶扩到了 256, 对应极长序列下更细的切分.

合并的正确性靠 log-sum-exp 的可组合性撑着. 每个 split 独立算自己那段 KV 的在线 softmax, 产出一个部分输出和一个部分 lse; 合并时, 两段的 lse 决定各自的权重, 按 $\exp(lse_i - lse_{\max})$ 重新归一后相加, 等价于把两段当成一段从头算. 这就是为什么 kernel 既返回 out 也返回 lse: lse 不只是给上层看的调试量, 它是 split-KV 合并的必要中间态. enable_batch_invariant 开关关掉 split 路径, 换来的是结果不依赖 batch 怎么切, 用在需要逐位复现的场景 (比如对拍、调试数值问题); 代价是放弃了 split 的并行度, 长序列下会变慢. 这是一个明确的正确性与速度的取舍, 放在接口上让调用方按场景选.

3. seesaw 双 warpgroup 流水

seesaw 是 2025-04-22 性能更新的核心, 把 H800 计算受限场景从 580 TFlops 推到 660 TFlops. 它要解决的问题很具体: 在寄存器放不下两份输出矩阵的约束下, 怎么让 CUDA Core 的标量运算 (softmax、缩放) 和 Tensor Core 的矩阵乘重叠起来. 这一节先讲约束为什么卡死了常规做法, 再把 12 步调度拆开看它怎么绕过去.

重叠这件事为什么难, 得回到注意力的两类运算. QK 内积和 score-V 乘积是矩阵乘, 跑在 Tensor Core 上; 求行最大、减最大、取指数、归一化、按新最大缩放旧累加器, 这些是逐元素的标量运算, 跑在 CUDA Core 上. 一个 KV 块的处理里, 这两类运算有先后依赖: 要先算完 QK 才能做 softmax, 做完 softmax 才能算 score-V. 如果串行走, Tensor Core 在做 softmax 时就闲着, CUDA Core 在做 GEMM 时也闲着. 要让两种单元都忙, 就得同时处理两个 KV 块, 让一个块的 softmax 和另一个块的 GEMM 错开进行, 而这正好撞上寄存器放不下两份累加器的墙.

3.1. 寄存器约束: 为什么不能照搬 FlashAttention-3 的 ping-pong

FlashAttention-3 的提速靠两招: warpgroup 间的 ping-pong 调度, 和 warpgroup 内的 GEMM-softmax 流水. ping-pong 的前提是手上有两份独立的输出累加器, 让两个 warpgroup 一个做 GEMM、一个做 softmax, 轮流占用 Tensor Core 和 CUDA Core. 但 MLA 这里有个硬约束: WGMMA 指令要求输出矩阵必须在寄存器里. 一个 $64\times512$ 的输出矩阵占 32768 个 32 位寄存器, 而每个 SM 只有 65536 个. 也就是说, 一个 SM 的寄存器只够放一份输出矩阵, 放不下两份. 没有第二份累加器, 经典 ping-pong 就没法直接搬过来.

官方 deep-dive 在这里留了一句「你可以先停下来想想有没有比我们更好的方案」. DeepSeek 的解法是在 FlashAttention 的在线 softmax 之上再加一步数学变换: 既然放不下两份完整的输出矩阵, 就把一份 $64\times512$ 的输出竖切成两半 $O_L$ 和 $O_R$ (各 $64\times256$), 分别交给两个 warpgroup 维护; 对应地, 每一步取两个 KV 块 $K_0, K_1$ (在 MLA 里 $K$ 和 $V$ 是同一份 latent, 名字不同), 把 $V_0, V_1$ 也各切成左右两半. 这样两个 warpgroup 各自只扛半份输出, 寄存器压力减半, 又能互相错开占用 Tensor Core 和 CUDA Core.

3.2. 12 步 seesaw 与细粒度 TMA 流水

seesaw 的 12 步 (编号 0 到 11) 可以这样读: warpgroup 0 维护 $\vec o_L$, warpgroup 1 维护 $\vec o_R$, 两者共享一个运行最大值 $m$. 步骤 1、2 两个 warpgroup 并行算各自的注意力分数 $\vec p_0 = \vec q K_0^\top / qk_scale$