在大型语言模型推理系统中，注意力机制的计算效率直接决定了整体吞吐量与延迟表现。DeepSeek 开源的 FlashMLA 库通过一系列针对 NVIDIA Hopper 架构的深度优化，将 MLA（Multi-head Latent Attention）解码内核的性能推至 660 TFlops，较上一版本提升约 14%。这一性能跃升的背后，是内核调度策略的根本性重构 —— 从传统的双输出矩阵乒乓调度演进为单输出矩阵的「跷跷板」（Seesaw）调度方案。本文将从计算特性分析出发，逐步拆解这一调度创新的工程实现细节。
计算边界特性与优化方向抉择
理解 MLA 内核的瓶颈性质是制定优化策略的前提。GPU 内核可大致划分为两类：计算密集型（compute-bound）与内存带宽密集型（memory-bound）。对于解码阶段的注意力操作，常见的直觉认知是内存 - bound，因为需要频繁访问 KV 缓存。然而，DeepSeek 的理论分析揭示了一个关键阈值：当查询头数 $h_q$ 与每请求查询标记数 $s_q$ 的乘积满足 $h_q s_q \ge 128$ 时，内核将转变为计算密集型。
具体而言，MLA 算法的计算量约为 $2 h_q s_q s_k (d_k + d_v)$ 次浮点运算，而内存访问量约为 $2 s_k d_k$ 字节（以 bfloat16 精度计）。两者的比值简化为 $2 h_q s_q$。在 NVIDIA H800 SXM5 GPU 上，峰值内存带宽为 3.35 TB/s，受限频后的峰值算力约为 865 TFlops。当 $h_q s_q \ge 128$ 时，算力成为系统瓶颈。这意味着传统的内存优化手段（如增加数据复用）收益有限，必须从算力利用率的角度重新设计内核。
DeepSeek 的在线推理系统配置决定了这一优化方向的选择。由于解码实例不使用张量并行（Tensor Parallel），$h_q$ 固定为 128，内核必然处于计算密集状态。因此，内核设计的核心目标转变为：最大化 Tensor Core 的利用率，将 CUDA Core 操作与内存访问尽可能重叠，使算力单元始终保持忙碌。
传统方案的寄存器资源困境
FlashAttention-3 论文提出的乒乓调度（ping-pong scheduling）和组内 GEMM-Softmax 流水线技术，通过交替使用两组输出矩阵实现了计算与数据移动的重叠。然而，这一方案在 Hopper 架构上面临严峻的寄存器资源约束。
WGMMA（Warpgroup Matrix Multiply Accumulate）指令要求输出矩阵常驻寄存器。每个 $64 \times 512$ 的输出矩阵需要占用 32,768 个 32 位寄存器。Hopper 架构的流式多处理器（SM）仅提供 65,536 个 32 位寄存器，这意味着单个 SM 上最多只能容纳一个完整的输出矩阵。传统的乒乓调度需要同时维护两个输出矩阵以实现交替计算，在寄存器资源上不可行。
这一约束迫使 DeepSeek 团队寻找一种创新的调度方式：仅使用一个输出矩阵，同时仍能实现 CUDA Core 操作与 Tensor Core 操作的充分重叠。
跷跷板调度的数学重构
Seesaw 调度的核心思想是对输出矩阵进行垂直切分，并引入额外的数学变换，从而在保持在线 softmax 算法等价性的前提下，实现两组 warpgroup 的交替协作。
具体而言，输出矩阵 $O$ 被垂直拆分为左侧部分 $O_L$（$64 \times 256$）和右侧部分 $O_R$（$64 \times 256$）。同理，KV 块 $K_0$、$K_1$ 和 $V_0$、$V_1$ 也被相应拆分为 $K_{0L}$、$K_{0R}$ 等子块。每个子块所需的寄存器数量减半，使得两个 warpgroup 可以分别持有 $O_L$ 和 $O_R$。
调度过程可概述为以下步骤（简化至单查询头情形）：
首先，初始化运行时最大值 $m$（初始为 $-\infty$）和输出矩阵 $\vec o_L$、$\vec o_R$（初始为零）。第一步由 warpgroup 0 执行 $\vec p_0 = \vec q K_0^T /qk_scale$，第二步由 warpgroup 1 执行 $\vec p_1 = \vec q K_1^T /qk_scale$。随后，两组分别进行 softmax 归一化与输出更新。
关键在于第 8 步与第 10 步的协同设计。当 warpgroup 1 更新 $O_R$ 时引入了缩放因子 $scale_0 \cdot scale_1$，而 warpgroup 0 在第 9 步将 $\vec p_0$ 乘以 $scale_1$，使得 warpgroup 1 在第 10 步能够使用更新后的 $\vec p_0$ 直接加上 $V_{0R}$。这一链条确保了两个 warpgroup 的操作在数学上严格等价于标准在线 softmax 累加，同时实现了计算资源的交错利用。
从时序视角观察，两个 warpgroup 形成了类似跷跷板的动作模式：一个 warpgroup 进行 GEMM 操作时，另一个 warpgroup 正在执行 softmax 归一化或输出更新。TMA（Tensor Memory Accelerator）数据传输指令的触发时机也被精心安排：在数据不再被当前 warpgroup 使用后立即启动下一批数据的预取，从而隐藏内存访问延迟。
细粒度流水线与缓存优化
尽管内核处于计算密集状态，内存延迟仍不可忽视。若数据未能在需要时准备就绪，Tensor Core 将被迫空闲等待。FlashMLA 采用了两项关键技术应对这一挑战。
细粒度 TMA 复制与 GEMM 流水线技术将较大的 KV 块分割为多个小片段。对于 $64 \times 576$ 的 K 块，系统发起 9 次 TMA 复制（每次传输 $64 \times 64$ 的子块），而 GEMM 操作在各子块传输完成后立即启动。这种流水化设计使得首个 TMA 复制完成时即可开始首批 GEMM 计算，有效提升了内存带宽的利用效率。
缓存提示策略通过 cute::TMA::CacheHintSm90::EVICT_FIRST 指定 TMA 复制的缓存行为，确保新数据在进入 L2 缓存时优先驱逐旧数据。实验表明，这一设置显著提升了 L2 缓存命中率，减少了对主存的访问压力。
这两项优化的综合效果是：在 H800 SXM5 GPU 上实现 80% 的 Tensor Core 利用率（相对于受限后的理论峰值）和 3 TB/s 的有效内存带宽。虽然该方案在纯内存密集场景下较双输出矩阵方案慢约 2%，但考虑到计算密集场景才是主要应用场景，这一权衡是可接受的。
辅助优化：可编程依赖启动与瓦片调度
除核心调度外，FlashMLA 还引入了若干辅助优化以提升整体性能。可编程依赖启动（Programmatic Dependent Launch）技术将 splitkv_mla 内核与 combine 内核的执行重叠，使后者在前者完成特定工作后立即启动，无需显式的 CPU 同步点。
瓦片调度器（Tile Scheduler）负责将请求和计算块分配至各 SM。良好的负载均衡对于充分利用 GPU 资源至关重要，尤其是在处理变长序列时。调度器根据各 SM 的当前负载动态调整分配策略，避免出现部分 SM 过载而其他 SM 空闲的局面。
FlashMLA 的算法与调度策略受 FlashAttention、Flash-Decoding 和 CUTLASS 等项目启发，在其基础上针对 MLA 的特殊结构进行了深度定制。这些底层优化使得 DeepSeek-V3 与 DeepSeek-V3.2-Exp 模型能够在保持高质量输出的同时，实现高效的推理服务。
参考资料
- DeepSeek FlashMLA 官方仓库：https://github.com/deepseek-ai/FlashMLA
- FlashMLA 内核深度解析文档：https://github.com/deepseek-ai/FlashMLA/blob/main/docs/20250422-new-kernel-deep-dive.md