Kimi推出MoBA：突破性实现无限上下文！

专家混合（Mixture of Experts）和稀疏注意力（Sparse attention）使得几乎无限的上下文成为可能。它让 RAG AI Agent 能够在没有上下文限制的情况下吞噬整个代码库和文档。

📌 长上下文注意力的挑战

当序列变得非常大时，Transformers 仍然面临沉重的计算负担。默认的注意力模式会将每个 Token 与其他所有 token 进行比较，导致计算成本呈二次方增长。当阅读整个代码库、多章节文档或大量法律文本时，这种开销就成了一个问题。

📌 MoBA

MoBA（Mixture of Block Attention，块注意力混合）将专家混合（Mixture of Experts）理念应用于注意力机制。该模型将输入序列划分为多个块，然后一个可训练的门控函数计算每个查询 token 与每个块之间的相关性得分。只有得分最高的块才会被用于注意力计算，从而避免了对完整序列中每个 token 的注意。

块是通过将序列分割成相等的跨度来定义的。每个查询 token 会查看每个块中键的汇聚表示（例如，采用均值池化），然后对其重要性进行排序，选择几个块进行详细的注意力计算。包含查询的块始终会被选中。因果掩码确保 token 不会看到未来的信息，保持了从左到右的生成顺序。

📌 在稀疏注意力与完整注意力之间无缝切换

MoBA 替代了标准的注意力机制，但并未改变参数的数量。它与标准的 Transformator 接口兼容，因此可以在不同的层或训练阶段之间切换稀疏和完整注意力。一些层可能会为特定任务（如有监督微调）保留完整的注意力，而大多数层则使用 MoBA 来降低计算成本。

📌 这适用于更大的 Transformer 堆栈，通过替换标准的注意力调用。门控机制确保每个查询仅关注一小部分块。因果性通过过滤掉未来的块以及在查询当前块内应用局部掩码来处理。

📌 下图展示了查询仅被路由到少数几个“专家”块的键/值，而不是整个序列。 门控机制将每个查询分配到最相关的块，从而将注意力计算的复杂度从二次方降低到子二次方。

📌 门控机制计算每个查询与每个块的凝聚表示之间的相关性得分。 然后，它会为每个查询选择得分最高的前 k 个块，无论这些块在序列中的位置有多远。

由于每个查询仅处理少数几个块，因此计算仍然是子二次方的，但如果门控得分显示出较高的相关性，模型仍然可以跳转到远离当前块的 token。

PyTorch 实现  

这段伪代码将键（keys）和值（values）划分为多个块，计算每个块的均值池化表示，并通过将查询（Q）与池化表示相乘来计算门控得分（S）。

📌 然后，它应用因果掩码，以确保查询不能关注未来的块，使用 top-k 操作符为每个查询选择最相关的块，并将数据组织起来以进行高效的注意力计算。

📌 FlashAttention 被分别应用于自注意力块（当前的位置）和 MoBA 选定的块，最后使用在线 softmax 将输出合并。

📌 最终结果是一个稀疏注意力机制，它保留了因果结构并捕获了长距离依赖关系，同时避免了标准注意力的完整二次方计算成本。

这段代码将专家混合（mixture-of-experts）逻辑与稀疏注意力结合，使得每个查询仅关注少数几个块。

门控机制为每个块与查询进行评分，并选择前 k 个“专家”，从而减少键/值比较的数量。

这使得注意力的计算开销保持在子二次方的水平，能够在不增加计算或内存负担的情况下处理极长的输入。

同时，门控机制确保查询在必要时仍能关注到远距离的 token，从而保留了 Transformer 对全局上下文的处理能力。

这种基于块和门控的策略正是 MoBA 在 LLM 中实现近无限上下文的方式。

实验观察  

使用 MoBA 的模型在语言建模损失和下游任务性能上几乎与完整注意力相当。即使在数十万或数百万个 token 的上下文长度下，结果仍保持一致。通过“尾部 token”评估的实验确认，当查询识别相关块时，重要的远距离依赖关系仍能被捕获。

可扩展性测试表明其成本曲线是子二次方的。研究人员报告，在一百万个 token 的情况下，速度提高了最多六倍，且在该范围之外的增益更大。

MoBA 通过避免使用完整的注意力矩阵并利用标准的 GPU 内核进行基于块的计算，保持了内存友好性。

最终观察  

MoBA 通过一个简单的思路减少了注意力开销：让查询学习哪些块是重要的，并忽略其他所有块。

它保留了基于标准 softmax 的注意力接口，并避免强制使用僵化的局部模式。许多大型语言模型可以以插拔式的方式集成这一机制。

这使得 MoBA 对于需要处理极长上下文的工作负载非常有吸引力，比如扫描整个代码库或总结庞大的文档，而无需对预训练权重进行重大修改或消耗大量的重训练开销。