一句话总结

作者介绍了 MiniMax Sparse Attention (MSA)，这是一种建立在 Grouped Query Attention (GQA) 之上的块级稀疏注意力机制，具有用于 Top-k 块选择的轻量级 Index Branch，在 109B 参数模型上将 1M 上下文中的每个 token 注意力计算减少了 28.4 倍，并在与共同设计的 GPU kernels 配对时在 H800 上实现了 14.2 倍预填充和 7.6 倍解码加速。

核心贡献

• MiniMax Sparse Attention (MSA) 被引入作为一种建立在 Grouped Query Attention 之上的块级稀疏注意力机制，利用轻量级 Index Branch 为每个组选择 top-k 键值块。该设计在通过最大池化评分实现组特定稀疏检索并保留最新块以保持稳定性的同时，保持了高效的块级执行。
• MSA 与专用 GPU 执行路径共同设计，该路径采用 exp-free Top-k 选择和 KV-outer 稀疏注意力，将理论稀疏性转化为实际加速。这种方法在块粒度访问下提高了 tensor-core 利用率，并融合辅助计算以处理偏斜的块流行度，而无需原子更新。
• 在 109B 参数混合专家模型上的实验（训练预算为 3T token）表明，MSA 在下游基准测试中与 Grouped Query Attention 匹配。MSA 在 H800 GPU 上实现了 14.2 倍预填充和 7.6 倍解码墙钟加速，同时将 1M 上下文中的每个 token 注意力计算减少了 28.4 倍。

引言

随着大语言模型转向长程 Agent 工作流，超长上下文窗口对于生产代码部署和网页导航等任务变得至关重要。然而，由于二次方缩放，标准 softmax 注意力造成了严重的计算和内存瓶颈，而现有的混合或稀疏化架构通常无法平衡质量与效率。作者提出了 MiniMax Sparse Attention，该机制保留了关键的稀疏 softmax 组件以最大化硬件重用，同时采用块级 token 选择。他们进一步将该算法与专用 GPU kernels 共同设计以实现高效执行，在 1M 上下文长度下提供显著加速，同时匹配 Grouped Query Attention 的性能。

方法

作者介绍了 MiniMax Sparse Attention (MSA)，这是一种旨在将 Grouped Query Attention (GQA) 扩展到超长上下文的块级稀疏注意力机制。核心架构将注意力层分解为两个不同的组件：负责选择相关上下文的轻量级 Index Branch，以及执行实际注意力计算的主分支 (Main Branch)。这种设计允许模型在动态调整感受野的同时保持高效的块级执行。

如框架图所示，过程始于输入隐藏状态通过线性投影和归一化层。架构随后分为 Index Branch 和 Main Branch。Index Branch 为每个 GQA 组独立运行，以从因果上下文中识别最相关的键值块。它首先使用 index query 和 index key 向量之间的点积为每个可见 key token 计算分数。为了减少路由开销并确保规则内存访问，这些 token 级分数使用 Block Max Pooling 聚合到块级。系统随后为每个 query 组选择得分最高的 Top-K 块。值得注意的是，包含当前 query 位置的本地块始终包含在选择中，以保留即时上下文。

一旦确定块索引，Main Branch 执行注意力机制。它不是关注整个序列，而是将 softmax 注意力计算限制在所选块内的 token 上。只要所选块的数量保持较小，这就会将计算复杂度从相对于序列长度的二次方显著降低到线性。输出通过从所选块中收集 Key 和 Value 向量并计算标准缩放点积注意力来生成。

为了有效训练这个两阶段系统，作者采用了一种专门的过程。由于 Top-K 选择操作是不可微的，Index Branch 无法通过标准语言建模损失直接训练。相反，作者利用 KL 散度损失来对齐 Index Branch 的分布与 Main Branch 在所选 token 上的分布。这种辅助损失为 indexer 提供了直接的学习信号，以模仿全注意力骨干的注意力模式。为了确保稳定性并防止辅助损失干扰主模型参数，Gradient Detach 机制应用于 Index Branch 输入。这确保 KL 损失仅更新索引投影。此外，使用 Indexer Warmup 阶段，其中模型最初运行全注意力以初始化索引投影，然后切换到稀疏注意力。这种架构设计和训练策略的结合使 MSA 能够在不损害模型性能的情况下实现显著加速。

实验

实验评估了稀疏注意力机制在多模态模型上相对于全注意力基线的两个 109B 规模训练路线。结果证实，该设计通过优化的 top-k 选择和减少的计算 FLOPs 保持了训练稳定性和模型质量，同时提供了实质性的效率提升。进一步的分析表明，动态 token 选择优于固定模式，并且关键的注意力结构（如 attention sinks）在没有硬编码架构先验的情况下自然产生。

作者调查了改变键值块大小对模型性能的影响，比较了 32、64 和 128 的配置以平衡效率和质量。结果表明，困惑度指标在不同块大小下保持稳定，表明可以使用更大的块来提高 kernel 效率，而不会显著降低通用语言能力。长上下文检索性能因任务长度而异，8K 设置在较大块大小时显示改进，而 32K 设置显示轻微下降。包括 TAU2 和 SWE 在内的困惑度基准测试在所有测试的块大小下显示一致的结果。8K 设置上的长上下文检索随着块大小从 32 增加到 128 而略有改善。32K 设置上的检索性能随着块大小增加而逐渐下降，尽管报告的整体影响有限。

作者比较了带有 index value head 的模型配置与没有该配置的变体，以评估额外计算的必要性。结果表明，移除 value head 不会导致系统性退化，因为性能差异很小且取决于基准测试。没有 value head 的变体在通用推理方面表现更好，而带有 value head 的版本在数学和代码任务中保留了优势。无 value 变体在通用知识和推理基准测试上实现了更高分。有 value 变体在数学和代码评估任务上表现出更强的性能。多模态和长上下文检索基准测试的结果混合，显示任一配置没有一致的优势。

作者在扩展训练阶段后评估了 MSA-CPT 模型在长上下文基准测试上的有效性。结果表明，稀疏模型与 Full-Attention 基线保持竞争力，尽管在高度受限的注意力预算下运行，仍保留了长上下文能力。MSA-CPT 在 HELMET 和 RULER 基准测试上相对于 Full 基线保持了有竞争力的总体得分。在 In-Context Learning 和多跳推理任务等特定类别中观察到性能提升。虽然一些检索和问答子集显示轻微下降，但该模型通常保留了长上下文能力。

作者通过比较具有这些约束的模型与没有这些约束的模型，评估了硬编码特定注意力块（如序列开始和局部窗口）的必要性。结果表明，移除这些强制选择对推理、代码和困惑度的标准基准测试的影响可以忽略不计。这表明稀疏注意力机制在没有显式干预的情况下自然学会优先考虑重要的 token，如 attention sinks。无论是否强制特定块，通用知识和推理基准测试的性能保持一致。困惑度指标显示强制和非强制配置之间具有可比稳定性。模型自主开发 sinks 和本地上下文的注意力模式，消除了对硬编码选择规则的需求。

作者在 H800 GPU 上对专用 top-k 选择 kernel 与标准 torch.topk 和 TileLang 实现进行了基准测试。结果表明，他们的方法在所有测试的序列长度和块配置下始终实现最低延迟。性能优势在部署设置中尤为明显，其中相对于基线的加速最大化。提出的 kernel 在所有测试设置中始终快于 torch 和 TileLang。在部署配置中观察到最大的性能提升。即使序列长度扩展到 512K 且块计数增加，效率优势也得以保留。

作者通过测试键值块大小、value head 配置和注意力机制来评估架构选择和效率优化，并与标准基线进行比较。结果表明，较大的块大小提高了效率，而不会降低通用能力，而移除 index value head 会造成任务相关的权衡，而不是系统性损失。此外，稀疏注意力模型在与全注意力基线相比时保持了有竞争力的长上下文性能，并在没有硬编码约束的情况下自主学习 token 优先级排序，而专用 top-k kernel 在扩展序列长度上实现了优越的延迟。