MiniMax-01: Scaling Foundation Models with Lightning Attention

1 · 动机:context-length 撞上的那堵墙

1.1 历史脉络

2024 年的 frontier model 大多停在 32K–256K context window。Gemini-1.5 Pro 把它顶到 1M 是个例外,但代价是极高的延时和 TPU 集群成本;开源生态(Llama-3.1, Qwen2.5, DeepSeek-V3, Mistral)集体没有突破 256K。

问题不在算法上没人想长上下文 — 是在算力上算不动。Transformer 的 softmax attention 是 O(n²d)。把 context 从 32K 翻到 1M(32×),attention compute 涨 1024×。即使 H800 / B200 翻番,也跟不上 context 翻 8×。

于是 prefilling 1M 个 token 就要花几百秒——本论文 Figure 2 给的硬数据:GPT-4o API 在 128K 上花约 70 秒,而 MiniMax-Text-01 在 1M 上只需 ~50 秒,在 256K 上 ~10× 于其他大模型。

1.2 别的方案为什么"够呛"

从 2020 年起,业界提了一打"线性化 / 稀疏化 attention"的方案。每个都有理论吸引力,但没一个真扩到了百亿千亿规模。论文用一段话点评了对手:

1.3 为什么"扩到 456B / 1M"非平凡?

就算 lightning attention 在 7B / 8K 上能跑,要扩 65× model + 125× context,仍然有三道工程坎:

• Pure linear 在 retrieval 上垮。论文实测 (Figure 7):lightning attention 在 NIAH 上比 softmax 低很多。如果不掺 softmax,7B 模型在 in-context learning 任务上没法用。
• 开源框架不支持。Megatron / DeepSpeed / FlashAttention 全是为 softmax 写的。要在 1500–2500 张 H800 上跑 hybrid,得自己写 LASP / varlen ring / 全套 CUDA kernel。
• MoE 的 all-to-all 通信遇上 1M context 的 sequence parallel 通信,会互相抢 NVLink 带宽。需要把 EP / ETP / EDP / CP 四种并行同时调度。
• Inference 的 prefix cache对 lightning attention 是新问题。Linear attention 的"缓存"是一个 d×d 的 KV state(不再是 token-level KV),怎么和 vLLM 的 paged KV 兼容?

论文 4 / 5 的篇幅都在回答这堆"系统级"问题,所以这不仅是一个 architecture paper,更是一份完整的 production-grade tech report。

2 · 背景速查

2.1 关键术语

2.2 Linear attention 的 right-product 复习

Softmax attention 必须 (QK⊤)V:先算 N×N 的 attention matrix,再乘 V。Linear attention 把 softmax 拿掉后,矩阵乘法满足结合律,可以变成 Q(K⊤V):

左边的 K⊤V ∈ ℝ^d×d,跟 N 无关。所以总 FLOPs 是 O(Nd²) 而不是 O(N²d)。当 N >> d(比如 N=1M, d=128),后者比前者大 8000×。

但有一个致命限制:在 causal 语言建模里,第 t 步只能看到前 t 个 token,即 K⊤V 是逐步累加的(kv_t = kv_{t-1} + k_t v_t⊤)。这个 cumsum 是串行的,GPU 上没法和 N 并行,所以"O(Nd²) 比 O(N²d) 快"在小 N 上反而被并行度吃掉。这就是 9 年来 linear attention 没在大模型里落地的根原因。

3 · Lightning Attention 是什么

3.1 一句话核心

Lightning attention 用 tiling(分块) 把 causal cumsum 切成 "block 内 left-product" + "block 间 right-product" 两条路径,前者并行(因为 block 小),后者递推但只走一次。这是 FlashAttention 思路在 linear attention 上的对应物。

具体一点: 1M context 上的 lightning attention forward

设 N = 1,048,576, d = 128(head dim), B = 256(block size), h = 64 heads。一次 forward:

1. 分块:把 Q, K, V ∈ ℝ^N×d 切成 T = 4096 块,每块 256×128
2. 初始化 KV ∈ ℝ^d×d = ℝ^128×128 = 0,在 SRAM 中常驻
3. 主循环 for t = 1..4096:
   • 从 HBM 加载 Q_t, K_t, V_t(每个 256×128 = 32KB,FP16)到 SRAM
   • O_intra = (Q_t K_t⊤ ⊙ M) V_t ⟶ B²d FLOPs ≈ 8M FLOPs
   • O_inter = Q_t · KV ⟶ Bd² FLOPs ≈ 4M FLOPs
   • O_t = O_intra + O_inter,写回 HBM
   • 更新 KV += K_t⊤ V_t (d²B FLOPs)
4. 总 FLOPs ≈ N(Bd + d²) = 1M·(256·128 + 128²) ≈ 50G FLOPs/head

对比 softmax attention 在 1M 上 ≈ N²d = 1M·1M·128 ≈ 130T FLOPs/head — 差 2600 倍。这就是 Figure 2 里 1M prefilling 拉开的根。

3.2 反向论证:不做 tiling 会怎样?

如果直接用朴素 linear attention 公式 kv_t = kv_{t-1} + k_t v_t⊤; o_t = q_t kv_t,你会得到一个长度 N 的串行 dependency chain。GPU 即便有 1024 个 SM,每一步都要等上一步,实测 训练速度比 softmax(FlashAttn-2)还慢(论文 Figure 8 的 HGRN2 / Mamba2 数据印证)。Tiling 把 dependency chain 缩到 T = N/B 步,每步又能并行 B² 个内积,这才有 33.4K TGS 的吞吐。

4 · Hybrid 架构:7 + 1 的设计

4.1 整体结构

MiniMax-Text-01 总共 80 层,其中:

• 每 8 层一组,前 7 层是 lightning attention block,第 8 层是 softmax attention block(GQA, group=8, RoPE 半维)
• 每个 attention block 后接一个 MoE block (32 experts, top-2)
• 整体: (LA · MoE) × 7 → (Softmax · MoE) × 1 重复 10 次 → 80 层
• Hidden dim 6144, attention head 64, head dim 128, FFN inner dim 9216

4.2 RoPE 只对一半 dim 用

论文一个细节: softmax 层只对 head dim 的前 64 维(共 128)加 RoPE,后 64 维不加。这样长度外推时,前半携带"相对位置信息"会受 RoPE 抑制,后半保留"位置无关的内容信息",得到更平滑的外推曲线。RoPE base = 10000,后续 long-context 训练阶段拉到 5M / 10M(详见 §8)。

4.3 反向思考:为什么不 1:1 / 3:1 / 15:1?

论文做了 ablation(Table 4):

• 15:1 / 全 lightning — NIAH 大幅下降,7B 上仅 ~80(softmax baseline 98)
• 3:1 — 速度优势变小,长上下文 attention 占比上升
• 1:1(类 Jamba) — softmax 占比过半,1M 时 prefill 时间退化为 quadratic
• 7:1 是甜点 — 既保住 retrieval(NIAH 95.7,接近 softmax 的 98),速度又比 hybrid-window 1024 快 ~1.05× 而 NIAH 高 80%

5 · MoE 整合: 32 experts × top-2 + Global Router

5.1 路由公式

每 token 选 top-2 个 expert(K=2),取 softmax 加权。32 个 expert,每 expert FFN 内部 dim = 9216,每个约 14B 参数。所以总参数 32 × 14B + 共享部分 ≈ 456B,激活 2 × 14B + 共享 ≈ 45.9B。

5.2 Global Router (路由崩塌的修复)

标准 GShard auxiliary loss 在大集群 + 大 batch下会失灵:每个 EP group 看到的 token 分布都不一样,group A 的 expert 被打爆,group B 的 expert 闲着。论文加一步 allgather:在 dispatch 之前先全局同步每个 expert 已收到的 token 数,再做全局调度。这一步代价是 O(world_size · E) 个数,通信量极小但能把 token drop rate 砍到接近 0。

5.3 Worked example: 一个 token 的路由

设 input token x_t = "积分" 的 hidden h_t ∈ ℝ⁶¹⁴⁴。Router 步骤:

1. logits = h_t · W_g, W_g ∈ ℝ^6144×32 → 32 个 score
2. 例如 score = [..., expert_3: 4.2, expert_17: 3.8, expert_29: 3.5, ...]
3. top-2 选出 expert 3 和 17,softmax([4.2, 3.8]) = [0.60, 0.40]
4. FFN_3(h_t) ∈ ℝ⁶¹⁴⁴, FFN_17(h_t) ∈ ℝ⁶¹⁴⁴
5. output = 0.60 · FFN_3(h_t) + 0.40 · FFN_17(h_t)

整个 batch 的 ~30k token 同时做这个,每个 expert 期望承接 30k × 2 / 32 ≈ 1900 token;有了 capacity limit (e.g. 1.2× 期望) 后,超出的部分被丢。Global router 通过预先同步避免某个 expert 被打爆。

6 · Capacity 论证: 为什么 hybrid 反而比纯 softmax 还好

这是论文最反直觉的发现之一。常识里 hybrid 应该是权衡(略差但更快),但 Table 4 / Table 5 显示 hybrid-lightning 在 NIAH / SCROLLS 上甚至比纯 softmax 还高。论文给了一个干净的 capacity 论证:

6.1 把 softmax 看成线性 RNN

令 q_t, k_j 为 query/key 行向量。Softmax attention 可以重写成递推:

注意每个时间步 t 都要从 j=1 开始重算一遍 → 论文称这叫 "Going Through a Book"。这是为什么 softmax 在 retrieval 上强:它对每个 query整本书重读一遍,信息不会丢。

但仔细看:t 时刻保存的状态 (s_t, o_t) 都是 d 维向量 → recurrent state 容量 = O(d)。

6.2 Lightning attention 的状态

而 lightning attention 的状态是 KV ∈ ℝ^d×d/h(每个 head),容量 O(d²/h)。在 d=128, h=64 时,d²/h = 256,而 d=128 → lightning 的状态容量是 softmax 的 2×。

所以 hybrid 不是 trade-off,而是 "大容量 streaming + 周期性精确重读" 的组合。这解释了为什么 NIAH 上 hybrid 反而能超 pure softmax:lightning 那 7 层先把信息压缩进高容量状态,softmax 那 1 层再做精确 lookup,等于两次 representation 的协同。

7 · 系统优化: LASP+ / Varlen Ring / EP-ETP

这一节是论文最"工程"的部分。能在 1500–2500 张 H800 上把 hybrid + MoE 跑到 75% MFU,几乎全靠这三件事:

7.1 LASP+: 把序列并行串行依赖打掉

原版 LASP (Sun 2024) 在 lightning attention 的 sequence parallel 模式下,需要 CP rank 之间用 send-recv 串行传 KV state——rank 0 算完 → 发给 rank 1 → rank 1 算完 → 发给 rank 2……一条链。这是因为 KV cumsum 必须按序累加。

LASP+ 的发现:每个 rank本地可以先算自己片段的 KV 局部前缀和(无依赖),再用一次 AllGather 把所有局部前缀和广播,各 rank 再独立组装出全局前缀和。把"链式 send-recv N 步"压成"一次 AllGather + 局部计算"。

论文实测:LASP+ 把 N 节点的速度从 1× 提到 1/N_pcn × 原 LASP 速度(基本恢复线性加速比)。

7.2 Varlen Ring Attention: 让 data-packing 不浪费

训练时把多条样本沿 sequence 维拼起来(data-packing),避免 padding 浪费。但传统 ring attention 假设单条等长样本,套用 data-packing 会让 attention 跨样本混淆,或要重度 padding。

Varlen ring attention 在每个 ring step 中用 variable-length causal mask 标记每条样本的边界 offset,让 attention 只在样本内部计算,跨样本部分被 mask 掉。配合 ring 通信,每个 CP rank 看到的 KV 跨 rank 滚动,整段 1M 长 packed sequence 直接训练,无任何 padding。

7.3 EP-ETP Overlap: 把 a2a 通信藏起来

MoE 的 all-to-all (a2a) 通信和 ETP 的 allgather/reduce-scatter 在不同 ProcessGroup 上,可以同时跑。论文把 token chunk 成 N 组,组间错开调度,让一组的 a2a 和另一组的 expert compute 并发。配合 EP-ETP 双层并行,把 MoE 通信开销从 baseline 减半。

7.4 一句话: 为什么 MFU 能到 75%

三件事叠加:

• LASP+ 让 sequence parallel 几乎无通信瓶颈
• Varlen ring 让 data-packing 不浪费
• EP-ETP overlap 让 MoE 的 a2a 隐藏

剩下的就是 lightning attention kernel 本身的高 SM 利用率(B=256 块对应 WGMMA 256×256 指令的 sweet spot)。

8 · 长上下文训练 recipe

论文最被工业界关注的部分: 怎么把 8K 训出来的模型扩到 1M / 4M。三阶段渐进训练:

关键 trick:

• RoPE base 调整:从 10K 拉到 5M 再 10M,确保长距离 token 的位置编码不撞退化区
• Linear 插值切换数据混合,而不是阶跃,避免训练发散
• NIAH 不够测:NIAH 在 128K step 就饱和到 1.0 了。后续阶段用 MR-NIAH(多轮 needles)和 RULER 13 任务持续监控
• 4M 外推靠的是 RoPE 半维 + lightning attention 的位置无关特性,无需任何额外训练

9 · 实验关键结果

9.1 学术 benchmark (短文本)

MMLU 88.5 / MMLU-Pro 75.7 / MATH 77.4 / HumanEval 86.9 / GPQA 54.4 — 与 GPT-4o (11-20)、Claude-3.5-Sonnet 同档,部分超 DeepSeek-V3。属于"不输"级别,而不是"碾压"。

9.2 长上下文 RULER (核心战场)

看 1M 列: Gemini-1.5-Pro 0.85,MiniMax 0.91。这是论文最有说服力的数字——之前只有 Gemini 能跑 1M,现在开源拿到了 +6 个点。

9.3 LongBench-V2 (推理 + 长文本)

w/ CoT 设置: MiniMax 56.5 vs GPT-4o 51.4,Long 子集 47.2 vs 43.5。在长 + 复杂推理交叉场景下显著领先。

9.4 MTOB (Long In-Context Learning)

这个 benchmark 让模型只能从 context 里的语法书 + 375 个例子学一门小语种 (Kalamang)。MiniMax 在 eng→kalam ChrF 半本书 +44.7 / 全本书 ~50,delta 全部超 GPT-4o / Claude / Gemini。说明 long-context 不只是 retrieval,而是真能做 in-context learning。

9.5 Prefill latency (Figure 2)

1M context: MiniMax-Text-01 H800 ~50s,Llama-3-70B H800 OOM(都跑不了)。256K: GPT-4o API ~70s,MiniMax ~5s。

10 · 与同类工作对比

11 · 局限 / 个人 take / 待验证问题

论文承认 / 隐含的局限

• 4M 外推靠的是"vanilla NIAH 还能找到 needle",并不等于 4M 真正的推理能力;论文没在 4M 上跑 RULER,实际 4M 长上下文质量未知
• Lightning attention 在 retrieval 上的劣势依赖 softmax 层补救,如果未来场景需要 12 层一组(更长 chunk)的 lookup,7:1 不一定够
• RULER 13 任务包含 MK-NIAH 等合成任务,真实长文档(代码库、法律文书)的复杂依赖未必体现
• 456B / 45.9B active 仍然需要 8×80G H800 单机做 INT8 推理,推理成本不便宜
• 训练框架几乎全是自研(LASP+ / varlen ring / EP-ETP / 全套 CUDA kernel),社区复制门槛高

个人 take

• 这是linear attention 的 ImageNet moment——之前都是 toy demo,这次是 frontier-scale + 公开权重。生态会跟上,2025 年内应能看到更多 hybrid 模型(实际后来 Llama-4 / Qwen3-VL 都跟进了 linear/hybrid)
• "7:1 ratio"是当前最优实践,但具体取决于 lightning state 容量。换 backbone(比如 Mamba-2)就要重新 sweep
• 论文在系统工程上的贡献被低估了。LASP+ + Varlen ring + EP-ETP 这一套是 1M context 训练的必要条件,而非锦上添花
• Capacity 论证(state O(d²/h) > O(d))是简洁优美的,可能是 hybrid > pure softmax 反直觉现象的根源
• Long-context 已经从"能不能塞进去"变成"能不能在长文本里推理"的问题。RULER / LongBench-V2 / MTOB 三件套是新的标杆

我的疑问 (落地后想验证)

1. Lightning attention 的 KV state 是 d×d/h ≈ 256 维 — 在长 reasoning(o1 / R1 风格)上,这个容量是否足够缓存中间推理步骤?
2. RoPE 半维的设计在 4M 外推时有没有出现频率混叠?为什么没在论文里看到 4M RULER 数字?
3. Lightning attention 的 prefix cache 怎么和 vLLM 的 paged KV 对齐?KV state 是 dense d×d 的,无法按 page 切分
4. 32 experts × top-2 这个 K 值是否最优?DeepSeek-V3 的 256 experts × top-8 路径走向是否更好?
5. Hybrid 7:1 在 reasoning RL post-training 里会怎样?如果 reward signal 强烈奖励"精确召回",训练会不会把 lightning 层"逼成" softmax?
6. ETP 增加了 expert 维的 tensor parallel,但这要求 expert weight 切分,与 ZeRO-3 怎么协调没讲清楚

记忆点

精读笔记 v1 · 配套论文 PDF 在 /data/szhang967/papers/paper-notes/MiniMax01_2501.08313.pdf