[Github] [MiniMax 论文] [训练器] Flash-MSA 与 Flash-Attention 的隔离训练步骤对比。1 一些前沿模型 [1, 2, 3, 4, 5] 使用稀疏注意显著加速推理，但没有人发布过高效训练的代码。今天我介绍了世界上第一个针对Hopper和Blackwell GPU的Minimax稀疏注意的高效开源训练内核，所有开发工作都是在Spheron H100和B200租赁机上完成的，并借鉴了FA4、MSA推理和Codex。免责声明：这不是官方实现，我与MiniMax没有关联。关于MSA MSA类似于Deepseek稀疏注意，具有一些核心变化。图1来源于MSA论文。1. 块状稀疏性：主注意力的选择不再是单独的KV，而是使用代理分数的最大池化以128为一组进行选择。这为内核引入了一些良好的缓存属性。2. 主注意力使用GQA而不是MLA：这一点尤其重要，因为据我所知，没有西方实验室将MLA纳入其训练中，这使得像GLM-5.2、DSv4这样的前沿模型所采用的稀疏注意公式不可用于此处的模型。3. 代理头的分组专业化：用GQA替代MLA引入了每层内查询的独立组，使我们可以选择每个代理头中不同的KV子集，而不是像DSA那样对整个注意力层进行求和和评分。有证据表明注意头自然会关注不同的标记，因此这一变化应该会增加主注意力的表达能力。内核设计：内核序列的高级概述：为了高效运行MSA，我们需要尽量减少重复工作，并不过度占用寄存器/共享内存。除了常规的快速寄存器（Q块、KV块、O累加器、LSE累加器），我们还需要考虑在前向过程中流式处理的top-k累加器。在向后过程中，我们必须为双重注意的组合传递腾出空间，以便我们可以计算主注意力的梯度和代理注意力的梯度，因为代理梯度需要访问代理和主注意力的概率。块稀疏性的一个好处是，由于我们只需缓存块的索引而不是像DSA那样的单个标记，因此在整个反向传播过程中，缓存块索引是可行的——这意味着在整个训练步骤中，只有代理前向是关于上下文长度的二次的，其他一切都使用代理前向的缓存稀疏块。前向过程：操作顺序是代理注意力 -> 稀疏主注意力 -> 将主注意力输出发送到下一层，保存主LSE以便用于后向。代理注意力：代理点积与常规Flash Attn略有不同，因为我们不再需要累积输出，但确实需要追踪前k个注意力分数及其各自的索引。在处理关键时，我也不累积向后的LSE，而是运行一个非常便宜的稀疏激活上的代理点积重计算以获取向后的LSE。在实践中，这比在前向过程中融合LSE+topk更快。随着每个QK^T块的计算，我为每个块抓取因果本地最大分数，并对每个查询行的当前top-k值进行插入排序，存储在寄存器中。为了存储这些top-k寄存器，我不得不将关键块切成两半。此外，MSA要求每个标记的本地块必须是以滑动窗口样式遮蔽，所以我为每个查询的本地KV块设置了无穷大的注意力分数。主注意力：主注意力仅是块稀疏的快速注意力前向。这以前在MoBA中已经实现，所以我复制了他们聪明的技巧，将块稀疏注意重新参数化为变长快速注意。反向过程：为了计算代理头的梯度，我们需要融合代理和主注意力的反向传递，因为代理训练信号需要同时访问代理注意力概率和主注意力概率。由于我们从前向过程保存了块索引，并且只在稀疏KV激活上训练这两种注意力，因此反向过程可以在线性时间内完成。首先，我们提取缓存的块索引，并将$B$[批次，代理头，查询，top_k_slot]->[关键块]的映射反转为$B^{ ext{*}}$[批次，代理头，关键块]->[使用此块的查询]。我们使用$B^{ ext{*}}$调度查询块，以优化共享稀疏KV块的重用。然后，我们在所选块上运行一个快速的稀疏代理注意力传递（再一次使用MoBA变长技巧）以获取代理LSE，然后我们流式处理融合的代理-主反向任务，加载QKV、Q_proxy、K_proxy和main_lse的块。为了将如此多的头加载到寄存器中，我们必须减少每次使用的Q块和KV块的大小。在每个流中，我们计算主和代理注意力概率，计算dQ、dK、dV，然后计算...