Flash-MSA用稀疏注意力核加速百万Token训练

背景

随着大语言模型上下文窗口不断扩展，处理百万级 token 的训练已成为前沿需求。稀疏注意力（Sparse Attention）机制被多个前沿模型（如 GLM-5.2、DeepSeek-V4 等）用于加速推理，但高效训练稀疏注意力的开源实现一直缺失。现有开源方案仅支持推理阶段，而训练阶段需要同时处理前向和反向传播中的梯度计算，复杂度更高。本文作者独立开发了世界首个高性能开源训练内核，基于 CuTeDSL 在 Hopper 和 Blackwell GPU 上实现了 MiniMax Sparse Attention（MSA）的高效训练，并在 Spheron H100 和 B200 租赁实例上完成全部开发工作，参考了 FA4、MSA 推理代码和 Codex 项目。需要说明的是，这不是 MiniMax 官方实现，作者与 MiniMax 无关联。

核心内容

MSA 机制简介

MSA 类似于 DeepSeek Sparse Attention，但有以下关键变化：

1. 块级稀疏性（Blockwise Sparsity）：不是为每个 token 单独选择 KV，而是通过代理注意力（proxy attention）对每 128 个 token 的块进行最大池化，选择块索引。这种设计带来了良好的缓存特性。
2. 主注意力采用 GQA 而非 MLA：目前西方实验室尚未将 MLA 用于训练，而使用 MLA 的稀疏注意力（如 GLM-5.2、DSv4）无法直接迁移。MSA 采用 GQA，使得主注意力更易实现。
3. 代理头分组专业化（Group-wise Specialization of Proxy Heads）：GQA 为每层引入独立的查询分组，每个代理头可以选择不同的 KV 子集，而不是像 DSA 那样对整个注意力层求和打分。有证据表明注意力头天然关注不同 token，因此这一改变能增强主注意力的表达能力。

内核设计

整体流程：前向时依次执行代理注意力 → 稀疏主注意力 → 将主注意力输出传递到下一层，并保存主 LSE（log-sum-exp）用于反向传播。

前向

• 代理注意力：不同于标准 Flash Attention，不再累积输出，而是跟踪 top-k 注意力分数及其索引（在流式处理 key 时记录）。不累积 LSE，而是在反向传播时通过稀疏激活重新计算代理点积以获取 LSE（实际测试中比前向融合 LSE+topk 更快）。每计算一个 QK^T tile，对每个查询行获取因果局部最大分数，并插入排序到当前 top-k 值中（保存在寄存器中）。为容纳这些 top-k 寄存器，将 key 块减半。此外，MSA 要求每个 token 的局部块必须采用滑动窗口式未掩码，因此将局部 KV 块的注意力分数设为 inf。
• 主注意力：只是一个块稀疏的 Flash Attention 前向。借鉴 MoBA 的技巧，将块稀疏注意力重参数化为变长 Flash（varlen flash）。

反向

为计算代理头的梯度，需要融合代理注意力和主注意力的反向传播，因为代理训练信号需要同时访问代理注意力和主注意力的概率分布。 由于前向已保存块索引，且仅对稀疏 KV 激活训练两个注意力，反向传播可以在线性时间内完成。具体步骤：

1. 取出缓存的 block_indices，将映射 B[batch, proxy head, query, top_k_slot] -> [key block] 逆转为 B*[batch, proxy head, key block] -> [queries that use this block]。利用 B* 调度查询块，以优化共享稀疏 KV 块的复用。
2. 对选中的块运行快速稀疏代理注意力（再次使用 MoBA 变长技巧）获取代理 LSE，然后流式执行融合的代理-主注意力反向任务。加载 QKV、Q_proxy、K_proxy 和 main_lse 的块。由于要加载这么多头到寄存器，需要减少每次使用的 Q 块和 KV 块大小。在每次流中计算主注意力和代理注意力的概率，然后计算 dQ, dK, dV，并基于 KL 训练项计算代理的 dQ, dK。

KL 散度损失的反向传播技巧

原始 DSA 的 KL 损失项为：
L^ι = Σ_t D_KL(p_t, s_t || Softmax(I_t, s_t))
需要同时实例化索引器和主注意力概率分布来累积 KL 散度，会大量读写共享内存并占用额外寄存器，显著降低训练速度。作者发现一个技巧：可以原子化地反向传播，且数学上等价于完整 KL 损失。

推导过程：

• 代理注意力概率为 p_px，主注意力概率为 p，展开 KL 项：
  L^ι = Σ_t p_t * log(p_t / p_px,t) = Σ_t (p_t*log(p_t) - p_t*log(p_px,t))
• 需要计算对预 softmax 分数 z_px,i 的梯度：
  ∂L^ι/∂z_px,i = -∂/∂z_px,i Σ_t p_t*log(p_px,t)（主概率 p_t 在 KL 损失图中视为常数）
• 利用 softmax logprob 偏导：∂log(p_px,t)/∂z_px,i = δ_it - p_px,i
• 代入得：∂L^ι/∂z_px,i = -Σ_t p_t*(δ_it - p_px,i) = -p_i + p_px,i * Σ_t p_t
• 由于 Σ_t p_t = 1，最终得：∂L^ι/∂z_px,i = -p_i + p_px,i

即 KL 损失对代理分数的梯度 = 代理概率 - 主概率。内核中直接使用该公式计算代理梯度，无需完全实例化 KL 损失。

热身内核（Warmup Kernels）

在热身模式下，主注意力前向使用密集注意力（不使用块索引），因此可以完全跳过代理前向，只在反向传播中训练代理。热身前向直接调用 Flash 并保存输出和 LSE，返回占位 KL。反向时先对索引器调用密集 Flash 获取 LSE，然后复用稀疏 MSA 内核的融合代理+主注意力反向。

正确性验证

作者在 eager PyTorch 中实现了 MSA，并在多种配置下对比了两个实现的前向输出和反向梯度（包括目标输出损失和内部 KL 损失）的余弦相似度。扫描在 bf16 精度下进行，通常 bf16 的精度容差为 0.01。

下一步计划

• 增加融合反向并行性：当前反向受限于低并行度，正在优化。
• 更多细节待补充（原文未完整列出后续计划）。

关键要点

• 首个开源高性能训练内核：实现了 MSA 在 Hopper/Blackwell GPU 上的高效训练，填补了稀疏注意力训练开源的空白。
• 块级稀疏性：通过 128 长度的块索引缓存，使整个训练步骤中只有代理前向是上下文长度的二次复杂度，其余部分均利用缓存块，大幅降低计算开销。
• GQA 替代 MLA：使得稀疏注意力方案更易被西方实验室采用，因为 MLA 目前未用于训练。
• 代理头分组专业化：利用 GQA 的独立查询分组，每个代理头可独立选择不同 KV 子集，提升主注意力表达能力