稀疏 95% 几乎不掉点,Kimi 已上线的块级注意力 MoBA

MoBA: Mixture of Block Attention for Long-Context LLMs

Enzhe Lu, Zhejun Jiang, Jingyuan Liu, Yulun Du, Tao Jiang, Chao Hong, Shaowei Liu, Weiran He, Enming Yuan, Yuzhi Wang, Zhiqi Huang, Huan Yuan, Suting Xu, Xinran Xu, Guokun Lai, Yanru Chen, Huabin Zheng, Junjie Yan, Jianlin Su, Yuxin Wu, Neo Y. Zhang, Zhilin Yang, Xinyu Zhou, Mingxing Zhang, Jiezhong Qiu

cs.LG, cs.AI, cs.CL

2025-02-18

把 MoE 搬进注意力,上下文切块、每个 token 自选 top-k 块,95% 稀疏几乎不掉点,百万 token prefill 提速 6.5 倍。已用于 Kimi。

这篇在解决什么

注意力机制的开销随上下文长度平方增长。当模型要处理百万 token 级别的输入（Kimi、Claude、Gemini 都在往这个方向走），全注意力的算力账单就变得很难看。

业界的两条出路各有各的坑。一类是静态稀疏注意力，比如滑动窗口、attention sink（把注意力固定泼在开头几个 token 上）。它们硬编码了「该看哪里」，省是省了，但模式是固定的，换个任务往往就不灵。另一类干脆把注意力整个换掉，改成线性近似（Mamba、RWKV），复杂度降下来了，但在复杂推理上到底行不行还没有定论。

Moonshot AI 这篇 MoBA 走第三条路：保留标准 Transformer 和全 softmax 注意力，把「该看哪里」交给模型自己学，不预先规定模式。「less structure」（少加结构）是他们的口号。

方法

核心是把混合专家（MoE，mixture of experts）那套搬进注意力。把长度 N 的上下文切成 n 个块，每块 B = N/n 个 token。对每一个查询 token，用一个门控给所有块打分：

si = ⟨q, meanpool(K of block i)⟩

也就是用查询向量和该块键向量的均值做内积。取得分最高的 top-k 个块，只在这些块里跑标准的 softmax 注意力。

几个设计选择值得展开：

按块路由，不按 token 路由。块的粒度让 MoBA 能复用 MoE 和 FlashAttention 的 kernel，工程上跑得起来。
当前块强制保留。每个 token 一定 attend 自己所在的块（带因果掩码），相当于 MoE 里的「共享专家」。这一条防止模型走极端、把注意力全切光，保住了局部上下文。
不增不减参数。 MoBA 和全注意力参数完全一致，所以实验里唯一的变量就是注意力模块本身，对比干净。门控分数复用了已有的 q 和 K，没有新参数。
可以和全注意力混着用。分层混（比如最后几层留全注意力）、分阶段混（训练时 90% 的 token 走 MoBA、10% 走全注意力）。切换时没有明显的 loss 尖刺。这一条在后面对 SFT 很关键。

滑窗和 attention sink 都可以看作 MoBA 在特定门控下的特例，所以论文认为 MoBA 的表达能力严格更强。

结果

先把话说在前面：MoBA 在评测的生成阶段其实切回了全注意力，6.5 倍的提速是 prefill（预填充）阶段的成绩，不是端到端。

设置	指标	MoBA	全注意力
8K 上下文,81.25% 稀疏	validation loss 差距	与全注意力相差 <1e-3	基准
32K 上下文,95.31% 稀疏	末尾 token loss	略高于全注意力,随规模缩小	基准
Llama-8B-1M	MMLU	0.4903	0.4904
Llama-8B-1M	RULER @128K	0.7818	0.7849

scaling law 实验（568M 到 2.1B 参数）里，8K 上下文、丢掉 81% 注意力的情况下，MoBA 和全注意力的验证 loss 差距稳定在 1e-3 以内。把上下文拉到 32K、稀疏度堆到 95.31%，末尾 token 的 loss 略高于全注意力，但差距随模型变大在收窄。

最大的一组实验是把 Llama 3.1 8B 持续预训练并微调到 100 万 token（block 4096、top-k=12、最后 3 层留全注意力）。它在 16 个常用基准上和全注意力双胞胎几乎打平，needle-in-haystack（大海捞针）测试在百万 token 下仍然稳。

效率这边，prefill 百万 token 提速最高 6.5 倍。再往千万 token 推，注意力的计算时间省到 1/16（靠把张量并行扩到 query head 级别、把 K/V 广播到各分布式头上来绕开显存）。

为什么重要

它是少有的「生产线上验证过」的长上下文稀疏注意力，已经在 Kimi 的长文本请求里跑。而且因为是 drop-in 替换，可以直接接入已有的全注意力模型，不必从零重训整个架构。

对做长上下文应用的人来说，近无损、次二次、能扛百万到千万 token 的注意力，正是最卡脖子的那块。分层混合的设计让你不用一次性把宝全押在稀疏上。

需要泼一盆冷水：这是 Moonshot 的技术报告，所有数字都来自他们自己的 setup，缺少独立复现，也没有和其他动态稀疏方法（如 Native Sparse Attention、Quest）做横向对比，对照对象主要是他们自己的全注意力基线。

局限与存疑

SFT 阶段表现欠佳。 prompt token 不计入 loss，导致梯度稀疏、反传受阻，需要靠分层混合（最后几层留全注意力）来补。它不是一个可以无缝替换进任何训练流程的模块。
评测只在 prefill 用 MoBA。生成阶段切回全注意力，所以那 6.5 倍的提速反映的是长输入的一次性处理，不是端到端推理。
块粒度敏感。同样 75% 稀疏度下，最粗（8 块选 2）和更细的切法之间有 1e-2 的差距，需要调参。
末尾 token loss 仍略高。虽然随规模在缩小，但长上下文里末尾 token 恰恰是最需要看全文的部分。

术语

MoE（混合专家）：把一个层拆成多个「专家」子网络，用一个门控决定每次只激活其中几个。MoBA 把这个思想用到注意力上，每个「专家」就是一段上下文块。
attention sink（注意力汇）：模型倾向于把注意力堆在开头几个 token 上的现象，常被用作省算力的固定模式。
sliding window attention（滑窗注意力）：每个 token 只看前后固定窗口内的邻居，是最常见的静态稀疏方案。
prefill / 生成：处理长输入（prefill）和逐个吐字（生成）是两个阶段，算力瓶颈分布不同。
trailing loss：序列最后一段 token 的损失，比整体 loss 更能反映长上下文里「看到全文再收尾」的能力。
张量并行：把一个层的权重或激活切到多张卡上算，MoBA 在这里把它扩到了 query head 维度。

原文与代码

社区讨论

PWT 稀疏注意力新法：无需微调，视频生成端到端提速 2.5 倍青稞AI