一、前言

在百万 Token 超长上下文场景中，传统标准注意力复杂度为 \(O(n^2))，海量 KV 缓存会带来显存溢出、带宽瓶颈、计算延迟等一系列问题。

DeepSeek-V4 采用 CSA（Compressed Sparse Attention，压缩稀疏注意力）+ HCA（重度压缩注意力） 混合注意力架构破解该难题，而 Lightning Indexer（闪电索引器，下文简称 Indexer） 是 CSA 实现”先粗筛、后精算”稀疏推理的核心组件。

Indexer 本质是一套轻量级 MQA（Multi-Query Attention，多查询注意力） 打分模块，不参与最终注意力加权，仅负责为历史 KV 计算相关性分数，通过 Top-K 筛选出高关联条目，让后续正式注意力仅处理少量有效 KV，将计算量从全量 \(O(n^2)) 降至亚线性复杂度。同时它搭配 FP4/FP8/BF16 分层混合精度 进一步压榨显存与带宽性能，是 DeepSeek-V4 实现百万上下文常态化部署的关键工程设计。

本文结合官方论文、开源推理代码，全面拆解 Indexer 的定位、MQA 架构、完整计算流程、精度规则、KV 缓存设计与落地细节。

背景

训练超长序列 LLM 时，单卡显存放不下完整的 KV cache，需要对序列维度做并行（Sequence Parallelism）。目前主流有两种方案：

• Ulysses（DeepSpeed-Ulysses）：两次 All-to-All，按 head 切分
• Ring Attention：环形传递 KV blocks，分块计算

两者目标相同，但设计假设完全不同。本文从原理、通信量、代码实现到架构选择动机，做完整对比。

一、核心思想对比

Ulysses Attention

1
2
3

Step1: [N/P, h, d] ──All-to-All──→ [N, h/P, d]
Step2: 本地做标准 Attention（每张卡拿到全序列、部分 head）
Step3: 输出 ──All-to-All──→ 回到 [N/P, h, d]

关键：两次 All-to-All 把”序列切”转成”head 切”，每张卡对全序列做部分 head 的 attention。

Ring Attention

1
2
3
4

Round 0: attn(Q_i, K_i, V_i)           ← 本地 attention
Round 1: 收到 K_{i-1}, V_{i-1} → attn(Q_i, K_{i-1}, V_{i-1})
...
Round P-1: 收到所有 KV → online softmax 合并结果

关键：KV 沿环形传递，每张卡只存自己的 KV chunk，计算和通信完全重叠。

二、通信量数学推导

Ulysses：All-to-All 的 $(P-1)/P^2$

Ulysses 输入是 [N/P, h, d]（序列已被切，head 完整），输出是 [N, h/P, d]（序列完整，head 被切）。

每 GPU 发送 P-1 个 chunk，每个 chunk 大小：

$$
\frac{N}{P} \times \frac{h}{P} \times d = \frac{N \cdot h \cdot d}{P^2}
$$

总发送量：

$$
\text{send per GPU} = (P-1) \times \frac{N \cdot h \cdot d}{P^2} = N \cdot h \cdot d \cdot \frac{P-1}{P^2}
$$

两个 1/P 的来源：

• 第一个 1/P：序列维度已被切（N/P）
• 第二个 1/P：head 维度再切一次（h/P）

Ring：P2P 的 $(P-1)/P$

Ring 只传 KV（Q 不动），每轮传 (N/P) × d_kv：

$$
\text{send per GPU} = (P-1) \times \frac{N}{P} \times d_{kv} = N \cdot d_{kv} \cdot \frac{P-1}{P}
$$

只有一个 1/P（序列切分），没有 head 切分。

对比（P=8, h=64, d=128, d_kv=576）

注意：MHA 的 KV 是 h × d × 2 = 16384/token，MLA 的 c_kv 只有 576/token，这是后续分析的关键。

三、MLA 对 Ulysses 的致命问题

MLA 的 KV cache 结构

DeepSeek-V3/V4 使用 MLA（Multi-head Latent Attention），KV cache 不是 multi-head 的：

1
2
3

c_kv [seq, 512]     ← 压缩的共享 latent
k_rope [seq, 64]    ← RoPE 位置编码部分
合计：576 dim/token（vs MHA 的 16384 dim/token）

只有 1 个”头”，无法按 head 切分。

DeepSpeed Ulysses 代码验证

1
2
3
4

# deepspeed/sequence/layer.py 核心逻辑
q = _SeqAllToAll.apply(group, query, scatter_idx=2, gather_idx=0)
k = _SeqAllToAll.apply(group, key,   scatter_idx=2, gather_idx=0)  # K 和 Q 对称处理
v = _SeqAllToAll.apply(group, value, scatter_idx=2, gather_idx=0)

Q、K、V 走完全相同的 All-to-All 路径，没有”KV 走 All-Gather”的分支。当 num_kv_heads=1 时：

1
2

1 KV head / 4 GPUs → [1, 0, 0, 0]
GPU 1-3：没有 KV head → 无法计算 ❌

Ulysses SP 上限 = num_kv_heads

四、MLA 场景：Ring vs Ulysses 精确对比

通信量（P=8）

Ring Attention（MLA）：

• 只传 c_kv：每轮 (N/8) × 576，共 7 轮
• 总计：$N × 504$ / GPU

Ulysses（MLA，Q All-to-All + KV All-Gather）：

• Q All-to-All：$7 × (N/8) × 64 × 512 = N × 28,672$
• KV All-Gather：$7 × (N/8) × 576 = N × 504$
• Output reverse：$N × 28,672$
• 总计：$N × 57,848$ / GPU

Ring 比 Ulysses 省 115×。

根本原因：MLA 的非对称性——Q 巨大（32768 dim/token）、KV 极小（576 dim/token）。Ring 只移动小的 KV，Ulysses 被迫移动巨大的 Q。

缩放性对比

交叉点：$P = 4hd/d_{kv} ≈ 57$，但 Ulysses SP 上限 = 64，所以实践中 MLA Ring 几乎总是更优。

五、为什么 Ulysses 在视觉模型流行，文本 LLM 不用？

四个结构性原因

1. KV head 数量趋势

1
2
3
4

2020: MHA (GPT-3) num_kv_heads = 96  → Ulysses 随便用
2022: MQA (PaLM)  num_kv_heads = 1   → Ulysses 废了
2023: GQA (LLaMA-2) num_kv_heads = 8 → Ulysses 受限
2024: MLA (DeepSeek-V3) num_kv_heads = 1 → Ulysses 废了

文本 LLM 全面转向 GQA/MLA 压缩 KV heads，Ulysses 的前提条件被釜底抽薪。

2. 文本 LLM 的 TP 已经做了同样的事

1
2

Tensor Parallelism: 按 head 切权重 → 本地算 attention → All-Reduce
Ulysses:           按 head 切数据 → 本地算 attention → All-to-All

本质重叠，TP 已经切了 head 之后，Ulysses 没有额外收益。

3. 推理阶段 Decode 占 80% 时间

1
2

Prefill: ~20% 时间（可以序列并行）
Decode: ~80% 时间（Q=1 token，序列并行无用）

Ulysses 对 Decode 完全无用。视觉扩散模型没有 Decode 阶段，全程受益。

4. 视频生成 token 数极高

1

Sora 级别: (64×64) × 120 frames = 491,520 tokens

必须序列并行，且 DiT 用标准 MHA（32 heads），Ulysses 完美适配。

一句话总结

Ulysses 在视觉模型流行 = MHA（head 够多）+ 无 decode + 没被 TP 覆盖 + 序列极长。文本 LLM 四条全占不到。

六、DeepSeek 的实际选择

训练：全程 Ring/CP

DeepSeek-V3 技术报告显示，长上下文扩展（32K→128K）用的是 Ring Attention（Context Parallel），不是 Ulysses：

• MLA 只有 1 个 KV latent → Ulysses 物理上不可用
• KV 只有 576 dim → Ring 通信量本来就小
• Ring 的通信-计算 overlap → 长序列时通信几乎免费

推理：SGLang 的 CP 配置

1
2
3

--enable-nsa-prefill-context-parallel
--attn-cp-size 8
--nsa-prefill-cp-mode round-robin-split

选择 Ring 风格 CP 的原因：

• MLA 的 KV 只有 1 个共享 latent，切不了 head
• KV 极小（576 dim），Ring 通信可接受
• round-robin 切分 tokens 均衡负载

七、LoRA 压缩能让 Ulysses 复活吗？

思路：在压缩态做通信

MLA 的 Q 路径：hidden(7168) → q_lora(1536) → expand → Q[128, 576]

能否在 q_lora 维度（1536）做 All-Gather，而非展开后的 Q（73728）？

通信量更新（LoRA 压缩版）

vs 原始 Ulysses（展开态）：$N × 57,848$ → 压缩 21×

但仍然比 Ring 大 5.4×

1
2

Ring:              N × 504
Ulysses (LoRA):    N × 2,744  ← Q 和 O 的压缩态还是要传

根本原因：Ring 的 Q 和 O 根本不过网络，留在本地计算。Ulysses 无论怎么压缩，都要把 Q/O（或其压缩态）在网络上搬一次，这是架构级差距，压缩只能缩小、不能逆转。

八、全景决策树

1
2
3
4
5
6
7
8
9
10

num_kv_heads >= SP_size?
│
├─ YES (DiT, ViT, 视觉 MHA)
│   └─ Ulysses ✅（All-to-All，通信量 1/P²）
│
└─ NO (GQA=8, MLA=1, 文本 LLM)
    ├─ 训练：Ring/CP ✅
    └─ 推理：
        ├─ Prefill：Ring/CP
        └─ Decode：partial attn + All-Reduce

总结

Ulysses 的核心优势是 1/P² 通信缩放，但前提是 num_kv_heads ≥ SP。MLA 把 KV 压缩到 1 个 latent，直接废掉这个前提。Ring 只传 KV（576 dim），Q 留本地，反而成了 MLA 的最优搭档。

这不是巧合——MLA 和 Ring CP 是刻意协同设计，不是将就。

相关阅读：DeepSeek-V3 技术报告、DeepSpeed-Ulysses 论文（arXiv:2309.14509）、Ring Attention 论文（arXiv:2310.01889）

从残差到 mHC：一条清晰的三代演进路

如果你训练过深层 Transformer，一定对残差连接（Residual Connection）不陌生。它是 ResNet 留下的最重要的遗产之一，也是今天所有大语言模型的标配。

但残差连接有个根本问题：它只有一条信息流。

DeepSeek 在 2025 年连发两篇论文，把这个问题彻底讲透了。故事的三步是：

1. Residual（2015）：x_{l+1} = x_l + F(x_l) — 一条流，稳定但表达弱
2. HC（2024）：把残差流拓宽到 n 条并行流，表达能力暴涨，但训练直接崩
3. mHC（2025）：给 HC 加上数学约束，又强又稳

这篇文章把这三代讲清楚，重点放在 mHC 上。

残差连接的瓶颈：一条流不够用

标准 Transformer 里，每一层的计算是这样的：

1
2

x = x + Attention(Norm(x))  # 残差连接
x = x + MoE(Norm(x))        # 残差连接

Layer 0 到 Layer 42 的所有信息，全都挤在同一个向量里做加法。

浅层的语法信息、中层的语义信息、深层的推理信息，互相覆盖、互相干扰。这是残差连接的天花板。

HC（Hyper-Connections） 提出了一个很自然的想法：

为什么不把 1 条流扩成 n 条并行流，让不同深度的信息走不同的”车道”？

HC 的公式长这样：

[
X_{l+1} = B_l X_l + C_l F(A_l X_l)
]

• A_l：把 n 条流压缩成 1 条，送给 Attention/MoE
• F：正常的层计算
• C_l：把 F 的输出扩回 n 条流
• B_l：n×n 矩阵，控制 n 条流之间怎么混合（残差项）

n=4 的时候，FLOPs 几乎没增加（F 只算 1 次），但信息容量翻了 4 倍。

听起来完美，对吧？

HC 的致命缺陷：训练会崩

问题出在 B_l 上。

HC 里的 B_l 是完全可学习的，没有任何约束。当模型叠到 60 层、参数量到万亿级时，这个无约束的矩阵会出问题：

多层复合后，信号被无限放大或衰减。

数学上，HC 跨层递归展开后得到：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

$Π B_i$ 是多层 $B$ 矩阵的复合。因为 $B$ 无约束，这个复合矩阵的谱范数可以远大于 1，也可以接近 0。

**实测结果（DeepSeek 27B 实验）：**  
HC 的复合映射增益（Amax Gain Magnitude）达到 **~3000**，意味着信号在前向传播中可以被放大 3000 倍。训练到第 12k 步，loss 直接 spike，梯度范数爆炸。

HC 的作者（ByteDance Seed 团队）在 ICLR 2025 发表了这个想法，但没有解决稳定性问题。

---

## mHC：给 HC 加上"安全阀"

DeepSeek 团队在 2025 年提出的 **mHC（Manifold-Constrained Hyper-Connections）**，只做了一件事：

> 把 HC 中的残差混合矩阵 `B_l` 约束在**双随机矩阵流形**上。

### 什么是双随机矩阵？

一个 n×n 矩阵，满足：
- 所有元素非负
- 每行之和等于 1
- 每列之和等于 1

这样的矩阵，谱范数 **永远 ≤ 1**。也就是说，它永远不会放大信号。

而且，双随机矩阵在乘法下是**封闭**的：两个双随机矩阵相乘，结果还是双随机矩阵。这意味着叠 100 层，复合映射仍然稳定。

**一句话：mHC = HC + 双随机约束，恢复了残差连接的恒等映射性质，同时保留了多流的表达能力。**

---

## 怎么把矩阵变成双随机？Sinkhorn-Knopp 算法

mHC 的核心算法是 Sinkhorn-Knopp 迭代，操作很简单：

```python
# 输入: B_raw (4×4, 可能是负数)
M = exp(B_raw)  # 先保证非负

# 交替行列归一化 20 次
for _ in range(20):
    M = M / M.sum(dim=1, keepdim=True)  # 行归一化
    M = M / M.sum(dim=0, keepdim=True)  # 列归一化

# 输出: B (4×4 双随机矩阵)

为什么第一步用 exp()？因为 B_raw 是神经网络输出的 logit，可能有负数。直接归一化会产生负权重，违反”非负”约束。exp() 把任意实数映射到正数，完美解决。

20 次迭代是实验得出的经验值，足够收敛，又不会太慢。

4 条流到底解决了什么问题？

很多人问：为什么是 4 条流？不是 2 条或 8 条？

三个核心作用：

1. 梯度高速公路

普通残差的梯度路径是 43 个乘法项的连乘，容易消失或爆炸。
4 条流提供了 4 条并行的梯度路径，一条堵了还有其他。类似 DenseNet 的思路，但高效得多（F 只算 1 次）。

2. 信息分离存储

1
2
3
4

stream_0: 可能专注浅层信息（位置、语法）
stream_1: 可能专注中层信息（语义、实体关系）
stream_2: 可能专注深层信息（推理链）
stream_3: 可能做"工作记忆"（当前层临时计算）

Layer 5 学到的特征，通过 B 矩阵的合理混合，可以在 Layer 30 仍然清晰可用。而在 1 条流里，这个特征早就被 25 次加法淹没了。

3. 动态容量分配

B 矩阵是输入依赖的（由当前层的 hidden state 动态生成）：

• 简单 token（”the”, “a”）：B ≈ 单位矩阵，4 条流基本不混合，省计算
• 复杂 token（需要深度推理）：B 大幅混合，让更多层的信息参与

这是 1 条流做不到的——普通残差对所有 token 都是 x + F(x)。

为什么选 n=4？
论文实测：n=4 是性价比最优的点。n=2 效果不够，n=8 的 Sinkhorn 开销（8×8 矩阵 × 20 次迭代）显著增加，但收益递减。

工程优化：为什么只多 6.7% 开销？

mHC 看起来很重：每行代码都有矩阵运算、Sinkhorn 迭代、4 倍 hidden state……

但 DeepSeek 做了三件事，把开销压到了 仅 +6.7%：

1. Kernel Fusion（算子融合）

用 TileLang 把整个 mHC 计算——RMSNorm + 线性投影 + Sigmoid + Sinkhorn——融合成一个 CUDA kernel，减少内存读写。

2. Selective Recomputing（选择性重计算）

前向时只保存每块的第一个层输入，反向时重新计算 mHC 的中间激活。最优块大小由公式给出：

1

$L_r^* \approx \sqrt{\frac{nL}{n+2}}$

3. Overlapping with DualPipe

把 mHC 的计算和流水线通信重叠。在 DualPipe 调度中，F_post,res kernel 放到高优先级流，避免阻塞 All-to-All 通信。

实验效果：又强又稳

DeepSeek 用 27B MoE 模型做了对比实验：

下游任务（BBH / DROP / GSM8K 等 8 个 benchmark）上，mHC 全面超过 baseline 和 HC。

最重要的结论： mHC 让 1.6T 参数、60+ 层、MoE 模型能够稳定训练——这是普通 HC 做不到的。

一句话总结

mHC = HC + 双随机矩阵约束，是残差连接的最终进化形态，让万亿参数 MoE 能稳定训练。

如果你正在设计超大模型，mHC 是目前最值得考虑的残差连接方式。它不改 FLOPs，不改层计算，只改残差流的拓扑结构——但就是这一点，让深层训练从”可能崩”变成了”一定稳”。

参考资料：

• Hyper-Connections, Zhu et al., ByteDance Seed, ICLR 2025, arXiv:2409.19606
• mHC: Manifold-Constrained Hyper-Connections, Xie et al., DeepSeek, 2025, arXiv:2512.24880
• DeepSeek-V4 Technical Report, 2026

写于 2026-06-02，基于 DeepSeek-V4-Pro 代码和 mHC 原始论文整理。

一句话概括

MegaMoE = 把 EP（Expert Parallelism）中的 All-to-All 通信和 MoE 计算融合到一个 kernel 里，让 NVLink 传输和 Tensor Core 计算时间上完全重叠。

传统 EP 里通信和计算串行，GPU 利用率只有 50-60%；MegaMoE 通过 Symmetric Memory + 细粒度 Scheduler + 单 Kernel 状态机，把利用率推到接近 100%。

传统 EP MoE vs MegaMoE

传统 EP（如 DeepEP）

1
2
3
4

时间 →

[dispatch all-to-all] → 等完 → [GEMM1] → [SwiGLU] → [GEMM2] → 等完 → [combine all-to-all]
   NVLink 通信         空闲      计算      计算      计算        空闲      NVLink 通信

问题：通信和计算串行，GPU 要么在算要么在传，利用率约 50-60%。

MegaMoE

1
2
3
4
5
6

时间 →

┌───────────────────── 一个 Mega Kernel ─────────────────────┐
│  NVLink dispatch ←→ GEMM1 ←→ SwiGLU ←→ GEMM2 ←→ NVLink combine  │
│    (通信和计算同时进行，流水线式重叠)                          │
└──────────────────────────────────────────────────────────────┘

通信隐藏在计算背后，GPU 利用率接近 100%。

核心实现机制

1. Symmetric Memory（对称内存）

传统 all-to-all：

1

GPU_0 → cudaMemcpyAsync → GPU_1（需要显式同步，退出 kernel）

Symmetric Memory：

1
2
3

所有 GPU 共享一片 symmetric buffer
GPU_0 直接写入 GPU_1 的 symmetric buffer（RDMA over NVLink）
GPU_1 看到数据就开始算，不需要全局 barrier

关键：不需要等所有 token 传完再开始算。传一批就算一批（streaming）。

实现层级：

• 硬件层：NVSwitch 全连接，任何 GPU pair 都有直达链路，memory controller 识别远端地址自动走 NVLink
• 驱动/Runtime 层：CUDA VMM 把远端 GPU 物理内存映射到本地虚拟地址（cuMemMap / nvshmem_ptr）
• 应用层：MegaMoE 直接用普通指针读写远端 buffer，一条 PTX load/store 指令完成

注意：只有 symmetric allocation 的那块 buffer 地址一致，不是整个 GPU 显存对称。

2. 单 Kernel 状态机

传统方式多个 kernel 之间有全局同步点，无法实现 fine-grained overlap：

1

kernel_1 (dispatch) → kernel 边界（全局 barrier）→ kernel_2 (GEMM1) → ...

MegaMoE 的 scheduler/mega_moe.cuh 用一个状态机让所有 SM 在同一个 kernel 内自主领任务：

1
2
3
4
5
6
7
8
9
10

enum class BlockPhase { None, Linear1, Linear2 };

while (true) {
    auto block = scheduler.get_next_block();
    switch (block.phase) {
        case Linear1: /* 做 W1/W3 GEMM */ break;
        case Linear2: /* 做 W2 GEMM */ break;
        case None: return;
    }
}

没有 kernel 边界，没有全局 barrier，通信和计算可以交错到 warp 级别。

3. Wave-Based Expert Processing

256 个 expert 太多，无法同时处理（Token Pool 显存有限），分 wave 分批：

1
2
3
4

Wave 0: expert [0..31]   → 处理完 → 释放 pool
Wave 1: expert [32..63]  → 处理完 → 释放 pool
...
Wave 7: expert [224..255]

每个 wave 内：

• 所有 SM 抢 L1 block（gate + up projection）
• SwiGLU 在 SMEM 中直接做激活
• 所有 SM 抢 L2 block（down projection）

kNumExpertsPerWave 由 SMEM 容量、BLOCK_M、pool_capacity 共同决定，让 wave 内 expert 数刚好喂饱所有 SM 又不撑爆 Token Pool。

4. Arrival Count 轮询（细粒度 Overlap 的关键）

1
2
3
4

// 轮询等待：不是等所有 token 到齐，而是等够一个 BLOCK_M 就开始算
while (volatile_count < kNumSMs * kNumRanks) {
    // spin-wait 直到其他 SM/rank 报告完成
}

• 不需要全局 barrier（不需要所有 token 都收完才开始）
• 收到一个 block 的 token 就立刻开始算
• 这就是 streaming overlap 的本质

5. L1/L2 交错

1
2

Expert_0: [L1 block 0][L1 block 1][L1 block 2] → SwiGLU → [L2 block 0][L2 block 1]
Expert_1: [L1 block 0][L1 block 1]             → SwiGLU → [L2 block 0]

L1 结果留在 Token Pool 中，L2 直接从 pool 读（在 L2 cache 中热），一个 wave 结束后才释放 pool 空间。

对应关系：

W1 和 W3 拼在一起做一次 GEMM，可以复用 TMA 加载的 activation tile，减少 GMEM → SMEM 搬运。

SM/Warp 级别并行

GPU 硬件层次：

1
2
3
4

GPU (整块芯片)
└── SM (Streaming Multiprocessor) × 132 (H20) / 192 (B200)
    └── Warp × 最多 64 resident / SM
        └── Thread × 32 / Warp（SIMT 锁步执行）

在 MegaMoE 中，一个 Mega Kernel launch 占满所有 SM，每个 SM 内部：

1
2

Warp Group A（若干 warp）：负责通信（polling + NVLink store/load）
Warp Group B（若干 warp）：负责计算（Tensor Core MMA）

Warp Group A 和 B 在 SM 内部并行执行——这就是”通信和计算在 warp 级别交错”的含义。

为什么必须 GB200/GB300

MegaMoE 依赖三个 GB200/GB300 独有（或大幅增强）的硬件能力：

为什么计算足够快才值得 overlap

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

假设: 通信时间 = T_comm, 计算时间 = T_compute

情况 A — 计算快 (FP4 Tensor Core):
T_compute = 2ms, T_comm = 5ms
不 overlap: 7ms
overlap: max(2, 5) = 5ms ← 省 28%

情况 B — 计算慢 (BF16 软件模拟):
T_compute = 20ms, T_comm = 5ms
不 overlap: 25ms
overlap: max(20, 5) = 20ms ← 只省 20%

而且 overlap 不是免费的：
- 需要 SM 专门做通信 warp（减少计算并行度）
- Symmetric Memory 的 load/store 比本地慢 10-20x
- Scheduler 状态机有开销

只有当计算非常快、通信成为瓶颈时，overlap 的净收益才大于开销。

和 DeepEP / flashinfer_mxfp4 的区别

MegaMoE vs DeepEP

MegaMoE vs flashinfer_mxfp4

两者互补：flashinfer_mxfp4 管单卡 MoE kernel 效率，MegaMoE 管 EP 多卡通信+计算融合。

Python API 调用流程

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

# 1. 多进程初始化 + 分配对称内存
buffer = deep_gemm.get_symm_buffer_for_mega_moe(
    group,                  # 进程组 (torch.distributed)
    num_experts=256,
    num_max_tokens_per_rank=8192,
    num_topk=6,
    hidden=4096,
    intermediate_hidden=2048
)

# 2. 权重预处理（FP4 packing + 布局变换）
l1_transformed, l2_transformed = deep_gemm.transform_weights_for_mega_moe(
    l1_weights,  # W1, W3 (gate + up)
    l2_weights   # W2 (down)
)

# 3. 加载输入到 symmetric buffer
buffer.x[:N].copy_(x_fp8)           # FP8 激活
buffer.x_sf[:N].copy_(x_sf)         # 激活的 scale factor
buffer.topk_idx[:N].copy_(topk_idx) # 路由结果
buffer.topk_weights[:N].copy_(topk_weights)

# 4. 一次调用，完成 dispatch + GEMM1 + SwiGLU + GEMM2 + combine
y = torch.empty((N, 4096), dtype=torch.bfloat16, device='cuda')
deep_gemm.fp8_fp4_mega_moe(y, l1_transformed, l2_transformed, buffer)
# y 就是最终结果，全程只有这一个 kernel launch

总结：MegaMoE 的关键创新

一句话：MegaMoE = Symmetric Memory（消除通信 barrier）+ 细粒度 Scheduler（一个 block 到了就算）+ 单 Kernel 状态机（避免 kernel launch 开销），把 EP MoE 的 GPU 利用率从 ~60% 推到接近 100%。

局限性与适用边界

• 权重格式锁定：只支持 DeepGEMM 的 FP4 权重布局，需要预转换
• 激活函数限制：依赖 SwiGLU 无状态特性，L1 结果可直接给 L2 用；若 MoE 层有跨 token 依赖（如全局 norm），pipeline 会断
• 硬件绑定：Symmetric Memory + NVSwitch 全连接 + 原生 FP4 三个条件缺一不可，H20 及更早硬件无法使用
• 跨机不支持：Symmetric Memory 只在单机 NVSwitch 域内有效，跨机仍需回退到 DeepEP
• 负载不均影响：fetch_expert_recv_count() 的 spin-wait 在 expert 负载不均时可能导致 SM 空转

参考

• DeepGEMM PR #304: MegaMoE 实现（基于 NVSHMEM SymmBuffer）
• WideEP：大规模专家并行部署实战
• FlashInfer MXFP4 MoE 后端深度分析
• NVIDIA GB200 架构深度解析

背景

DeepSeek-V4 系列模型（V4 / V4-Pro）在 MoE（Mixture of Experts）层大量使用 MXFP4 量化权重，配合 FP8/BF16 激活 实现高效推理。在 SGLang / vLLM 中，flashinfer_mxfp4 是一个专为这种量化格式设计的 MoE Runner Backend。

本文将基于 H20 (SM90 Hopper) 实测数据和源码分析，完整解析 FlashInfer MXFP4 的实现原理，以及与 Marlin MoE 的架构差异。

MXFP4 量化格式

格式定义

MXFP4 是 OCP（Open Compute Project）标准的 4-bit 浮点格式：

1
2

每个权重元素: [1 sign | 2 exponent | 1 mantissa] → 4-bit 浮点数
每 32 个元素: 共享一个 8-bit 指数缩放因子 (E8M0 block scale)

与 INT4 的对比

优势：浮点表示让 MXFP4 在训练和推理中都有更好的数值稳定性。

FlashInfer MXFP4 架构

调用链（H20 / SM90）

1
2
3
4
5
6
7
8
9
10
11

SGLang --moe-runner-backend flashinfer_mxfp4
    ↓
Python: trtllm_fp4_block_scale_moe()
    ↓
C++: get_cutlass_fused_moe_module()  → 加载预编译 .so
    ↓
CUTLASS 3.x: cute::gemm::kernel::DefaultGemmUniversal
    ↓
    ├─ TMA (Tensor Memory Accelerator) → 异步搬权重 tile 到 SMEM
    ├─ WGMMA (Warpgroup MMA) → SM90 用 FP16 TC 模拟 FP4
    └─ Warp Specialization → Producer warpgroup 专搬数据，Consumer warpgroup 专计算

核心特性

1. Expert-First 遍历：按 expert 分组 token，而非按 token 遍历 expert
2. 全 Fuse：W1 + W3 + SwiGLU + W2 四个步骤 fuse 成一个 kernel
3. TMA + Warp Specialization：数据搬运和计算 overlap
4. Grouped GEMM：一次 launch 处理所有 256 个 expert

完整 Pipeline（DeepSeek-V4 为例）

模型配置

1
2
3
4
5

DeepSeek-V4-Flash:
- 256 experts, topk=6
- hidden_dim=4096, moe_intermediate=2048
- 权重: MXFP4 (4-bit), 激活: BF16
- Router: Sigmoid + Group TopK (非 Softmax)

第一阶段：Routing（路由选择）

1
2
3
4
5
6
7
8
9
10
11
12

# 文件: trtllm_fused_moe_routing_deepseek.cu

# 输入: router_logits [num_tokens, 256]
scores = torch.sigmoid(router_logits + bias)  # DeepSeek-V4 用 Sigmoid, 非 Softmax

# Group TopK: 256 experts → 8 组 × 32
group_scores = scores.view(num_tokens, 8, 32).max(dim=-1)  # 每组最高分
top_groups = group_scores.topk(4).indices  # 选 4 组 (128 experts)

# 从 128 个 expert 中选 top-6
topk_indices = ...  # [num_tokens, 6]
topk_weights = ...  # 归一化后的权重

为什么用 Sigmoid？

• 独立评分（非 Softmax 的零和游戏）
• 多 expert 可同时激活
• 训练更稳定

第二阶段：Gather（Token 重排）

1
2
3
4

# 文件: cutlass_fused_moe_kernels.cuh → expandInputRowsKernel

# 输入: hidden_states [num_tokens, 4096] (按 token 顺序)
# 输出: expanded_input [num_tokens × 6, 4096] (按 expert 连续排列)

目的：把路由到同一个 expert 的 token 收集到一起，形成连续内存，供后续 GEMM 使用。

1
2
3
4
5
6
7
8
9
10

原始 (token 顺序):
[t0 | t1 | t2 | t3]

路由结果:
expert_3: [t0, t1, t3]
expert_7: [t0, t2]

Gather 后 (expert 顺序):
[e3: t0, t1, t3 | e7: t0, t2 | ...]
 ↑ 连续          ↑ 连续

TMA 优势：连续输入让 TMA 一次加载整个 tile，达到满带宽。

第三阶段：Compute（Fused GEMM）

1

# 文件: cutlass_fused_moe_kernels.cuh → CUTLASS Grouped GEMM

这是最核心的部分，W1 + W3 + SwiGLU + W2 全部 fuse。

Step 3a: GEMM1 (Gate + Up Projection)

1
2
3
4
5
6

对每个 expert_i:
  M = expert_i 分到的 token 数 (例如 3)
  input_i = expanded_input[offset : offset+M]  # [M, 4096]
  
  # W13 = [W1; W3] 拼接，一次 GEMM 算出
  gate_up = input_i @ W13_expert_i  # [M, 4096] → split → [M, 2048] + [M, 2048]

MXFP4 权重存储：

1
2
3

// W1_expert_i: stored as [2048, 2048] int8
// 实际逻辑形状: [2048, 4096] FP4 (每 byte 存 2 个 FP4 元素)
// Scale: [2048, 64] float8_e8m0 (每 32 个元素一个 scale)

Step 3b: SwiGLU Activation

1
2
3

# 在 GEMM1 输出还在 SMEM/寄存器时立即执行（fuse 关键）
intermediate = SiLU(gate) * up  # [M, 2048]
# DeepSeek-V4 还有 swiglu_limit=10.0 的 clamp

Step 3c: GEMM2 (Down Projection)

1
2

output_i = intermediate @ W2_expert_i  # [M, 4096]
# W2_expert_i: stored as [4096, 1024] int8 → 逻辑 [4096, 2048] FP4

CUTLASS Grouped GEMM 调度

1
2
3
4
5
6
7
8
9
10
11
12

// 一次 kernel launch 处理所有 256 个 expert
cutlass::gemm::kernel::GroupedGemmKernel<...>::run();

// 内部调度:
// SM 空闲 → 从 problem list 取下一个 expert
// expert_i 的 M=0 (没 token) → 跳过
// expert_j 的 M=12 → 分配 SM 算 GEMM

// TMA + Warp Specialization:
// Warpgroup 0 (Producer): TMA_LOAD(weight_tile, desc)
// Warpgroup 1 (Consumer): WGMMA(tile_C, tile_A, tile_B)
// Producer 搬下一个 expert 的权重时，Consumer 正在算当前 expert

Pipeline 示意（H20）：

1
2
3
4

Cycle 0-100:  Producer 搬 expert_3 权重，Consumer 算 expert_1 (上一轮)
Cycle 100-200: Consumer 算 expert_3 (权重已到)，Producer 搬 expert_7
Cycle 200+:    Consumer 算 expert_7，Producer 搬 expert_5
→ 数据搬运和计算完全 overlap！

第四阶段：Scatter（结果写回）

1
2
3
4

# 文件: cutlass_fused_moe_kernels.cuh → finalizeMoeRoutingKernel

# 输入: expert_outputs [num_tokens × 6, 4096] (按 expert 排列)
# 输出: final_output [num_tokens, 4096] (恢复 token 原始顺序)

加权求和：

1
2
3

for token_j:
    output[j] = Σ (weight_k × expert_output[permute_map[j][k]])
                k=1..6

示例：

1
2

token_0 → expert_3 (0.6), expert_7 (0.4)
output[0] = 0.6 × out_0_e3 + 0.4 × out_0_e7

为什么叫 Scatter？

• expert_3 的输出: [out_0_e3, out_1_e3, out_3_e3]
• 需要写回: output[0], output[1], output[3]（不连续！）
• 这就是 scatter（分散写入），与 gather（聚集读取）相对

性能实测（H20, DeepSeek-V4-Flash, TP=4）

Benchmark 结果

阶段耗时分解（估）

GEMM 占 90%，所以 TMA + Grouped GEMM 对总体性能影响最大。

为什么 conc=2 差距最大？

1. Expert-First + Grouped GEMM：一次 launch 处理所有 expert，TMA 预取下一个 tile
2. 中等 batch：每个 expert 分到 2-3 个 token → GEMM tile 够大，WGMMA 饱和
3. Warp Specialization：Producer 搬数据，Consumer 计算，完美 overlap

Marlin 的劣势：

• Token-First：每个 token 的 topk expert 不同 → 权重反复换入换出
• 无 TMA：用 cp.async 加载，线程要参与地址计算 → SM 利用率 < 50%
• 无法全 fuse：中间结果 (gate, up, mid) 要写回 GMEM

为什么 conc=4 差距消失？

1
2
3
4
5

conc=4: prefill=50k tokens
→ Prefill 时间: ~1500ms (占 85%+)
→ Decode MoE 时间: ~265ms (占 15%-)

即使 MoE 快 50%，总体提升也只有: 265ms × 50% / 1765ms ≈ 7.5%

Prefill 主导后，MoE decode 的优化被稀释。

与 Marlin MoE 的架构对比

核心差异

Expert-First vs Token-First

Token-First (Marlin)：

1
2
3
4
5

for token_0:
    expert_3: W1 → W3 → SwiGLU → W2  ← 写回 GMEM
    expert_7: W1 → W3 → SwiGLU → W2  ← 换权重
for token_1:
    ...  # 重复上述，权重反复换入换出

Expert-First (FlashInfer)：

1
2
3
4
5
6
7

for expert_3:
    # 固定权重 W1/W3/W2，常驻 SMEM
    gate_up = input_e3 @ W13_e3  # 连续 token 一起算
    mid = SiLU(gate) * up        # 在寄存器/SMEM 直接算
    output = mid @ W2_e3         # 中间结果不落地
for expert_7:
    ...  # 换一次权重，算所有路由到它的 token

为什么 Expert-First 能全 fuse？

• 固定权重 → W1/W3/W2 常驻 SMEM，不换出
• 变 batch → 同一 expert 的多个 token 连续处理
• 中间结果 (gate, up, mid) 全在寄存器/SMEM，不写 GMEM

TMA + WGMMA 的硬件优势

TMA (Tensor Memory Accelerator)

SM90 (Hopper) 引入的硬件 DMA 引擎：

• 异步搬数据从 HBM 到 SMEM，不占用 CUDA core
• 一条 tcgen05.1d 指令触发，后台独立运行
• 支持多维 tensor 描述符（TMA Descriptor）

vs 传统 LDG：

1
2
3
4
5
6
7
8

LDG (传统):
  线程发 load → 等 HBM (~400 cycles) → 数据到 SMEM → 继续算
  → 线程 stall，SM 有空闲

TMA:
  线程发 TMA_LOAD → 立即返回 → 可以去 setup 下一个 tile
  → TMA 后台搬，线程继续算上一轮结果
  → SM 利用率 80%+

WGMMA (Warpgroup MMA)

SM90 引入的 warpgroup 级矩阵乘指令：

• 一个 warpgroup = 4 个 warp = 128 线程
• 直接读写 TMEM（Tensor Memory，SM90 新增寄存器文件）
• 比传统 mma.sync 高一个抽象层级

SM100 (Blackwell) 的进化：

1
2

SM90: TMA → SMEM → WGMMA → TMEM
SM100: TMA_GATHER4 → TMEM (跳过 SMEM) → UTCMMA (原生 FP4 TC)

Blackwell 的 TMA_GATHER4 还能一次加载 4 个不连续 token（专为 MoE 的 sparse routing 优化）。

总结

FlashInfer MXFP4 的核心优势

1. Expert-First + 全 Fuse：中间结果不落地 GMEM，省 ~2× hidden_dim × batch 带宽
2. TMA + Warp Specialization：搬运和计算 overlap，SM 利用率最大化
3. Grouped GEMM：一次 launch 处理 256 experts，launch overhead 最小化
4. Blackwell 未来兼容：cute_dsl_fused_moe_nvfp4 已支持 SM100 原生 FP4 TC

适用场景

一句话

FlashInfer MXFP4 在 Hopper 上通过 Expert-First + TMA + Warp Specialization + 全 Fuse，把 MoE 推理的瓶颈从内存带宽转移到了计算，实现了中等并发下 50% 的吞吐提升。

参考资料

• FlashInfer GitHub
• CUTLASS 3.x Documentation
• Introducing Machete - Red Hat Developer
• DeepSeek-V4 源码分析（本文基于 DeepSeek-V4-Flash 实测）
• NVIDIA Hopper Architecture White Paper

如果你对 MoE 推理优化感兴趣，可以看看我之前写的 Marlin MoE Kernel 深度分析，对比两种实现的差异。

什么是 WideEP

WideEP（Wide Expert Parallelism）不是 SGLang 里的一个具体模块，而是社区对一种大规模 MoE 部署模式的俗称。官方文档里叫 Large-Scale EP。

核心思路只有一句话：

把 Expert Parallelism 撑得很宽（几十到上百张 GPU），配合 DP Attention 消除通信，用 DeepEP 解决 All-to-All 瓶颈。

为什么需要 WideEP

以 DeepSeek-V3/V4 为例：256 个 routed expert，FP8 权重约 25GB，加上 dense 层、KV Cache，单机 8 卡跑高吞吐 decode 场景显存压力极大。

传统方案有三个瓶颈：

WideEP 的解法：用更多卡分摊专家 → 每卡只存 4 个专家 → batch 变大 → HBM 利用率高；DP Attention → 省掉 KV Cache 冗余。

WideEP 的核心组成

1. Expert Parallelism（MoE 层）

1

256 experts ÷ 64 张卡 = 每卡 4 个专家

token 经过 gate 计算 TopK 后，通过 DeepEP All-to-All 精确 dispatch 到目标专家所在卡，算出结果再 combine 回来。

2. DP Attention（Attention 层）

传统 TP 下 Attention 需要 AllReduce 汇总，跨 64 卡延迟巨大。WideEP 改用 Data Parallel：

• 每卡存一份完整的 Attention 权重（QKV/O projection）
• 每卡独立算自己的 batch，KV Cache 独立不共享
• 零通信，省掉 AllReduce

为什么放得下？DeepSeek-V3/V4 的 dense 层（Attention + Gate + Shared Expert + Norm）总共才 ~4GB，复制 64 份完全没问题。

3. DeepEP All-to-All

DeepEP 是 DeepSeek 开源的 MoE 专用通信库，解决了 NCCL 做 All-to-All 的三大问题：

各层并行方式一览

关键设计取舍：dense 层（4GB）复制 64 份 → 省掉通信；MoE 层（25GB）必须 EP 切分 → 靠 DeepEP 通信。

关于 --tp 64 的误解

很多人看到 --tp 64 以为权重被切了 64 份，其实不是。

在 SGLang 里，tp 只是定义一个包含 N 张卡的通信组，组内的并行方式由其他参数决定：

1
2
3
4
5
6

# parallel_state.py
moe_tp_size = tp_size // moe_ep_size // moe_dp_size
# tp=64, ep=64, dp=1 → moe_tp_size = 64 // 64 // 1 = 1

attn_tp_size = tp_size // attn_dp_size
# tp=64, dp=64 → attn_tp_size = 64 // 64 = 1

两个都是 1，意味着没有任何层做真正的 Tensor Parallel 切分。tp=64 的真实含义是”这 64 张卡组成一个协作集群”，具体怎么分工由 ep_size、dp_size 等参数决定。

SGLang 配置示例

基础 WideEP（单机 8 卡）

1
2
3
4
5
6
7

python -m sglang.launch_server \
  --model-path deepseek-ai/DeepSeek-V3 \
  --tp 8 \
  --ep-size 8 \
  --moe-a2a-backend deepep \
  --deepep-mode auto \
  --moe-runner-backend deep_gemm

完整 WideEP（DP Attention + 低延迟模式）

1
2
3
4
5
6
7
8
9
10
11

python -m sglang.launch_server \
  --model-path deepseek-ai/DeepSeek-V3 \
  --tp 64 \
  --dp-size 64 \
  --enable-dp-attention \
  --ep-size 64 \
  --moe-a2a-backend deepep \
  --deepep-mode low_latency \
  --enable-two-batch-overlap \
  --enable-eplb \
  --mem-fraction-static 0.85

多节点部署（2 × 8 卡）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

# Node 0
python -m sglang.launch_server \
  --model-path deepseek-ai/DeepSeek-V3 \
  --tp 16 --ep-size 16 \
  --nnodes 2 --node-rank 0 \
  --dist-init-addr <MASTER_IP>:29500 \
  --moe-a2a-backend deepep

# Node 1
python -m sglang.launch_server \
  --model-path deepseek-ai/DeepSeek-V3 \
  --tp 16 --ep-size 16 \
  --nnodes 2 --node-rank 1 \
  --dist-init-addr <MASTER_IP>:29500 \
  --moe-a2a-backend deepep

两个重要的优化特性

TBO（Two-Batch Overlap）

把请求拆成 micro-batch，在 attention 和 dispatch/combine 之间穿插执行：

1

--enable-two-batch-overlap

• 吞吐量最高提升 2×
• 零额外显存开销
• 原理：attention 算完 → yield → dispatch 和下一个 batch 的 attention 并行

EPLB（Expert Parallelism Load Balancer）

运行时收集 expert 激活统计，动态调整 expert 放置/复制策略，解决负载不均：

1

--enable-eplb

配合大 batch size（如 --max-running-requests 128）效果更好。

与之前方案的对比

WideEP vs 无 DeepEP 的 EP（moe-a2a-backend=none）

none 模式下，token 不跨卡，每卡算完自己的 expert 后 AllReduce 汇总——浪费计算（路由到的 expert 不在本卡就白算），跨机扩展不了。

通信 Backend 选择

总结

WideEP = DP Attention + EP MoE（DeepEP All-to-All），本质上就是：

1. MoE 层：专家分散到 64+ 卡，DeepEP dispatch/combine
2. Dense 层：每卡完整副本，DP 独立算，零通信
3. tp=64：只是通信组大小，不代表 TP 切分（moe_tp_size=1）

DeepEP 出来之前，大规模 EP 做不了——NCCL All-to-All 延迟太高。DeepEP 解决了这个瓶颈，WideEP 才成为实用的部署模式。

参考文档：

• SGLang Expert Parallelism 官方文档
• Large-Scale EP Blog - LMSYS
• DeepEP GitHub
• EPLB GitHub

引言

DeepSeek-V4 引入 MXFP4 量化后，MoE 层的计算效率成为推理性能的关键瓶颈。SGLang 的 Marlin Runner Backend 专门针对 INT4/MXFP4 量化权重优化 MoE 的 GEMM 计算。本文深入分析其实现原理、数据流以及设计权衡。

MoE 层的计算流程

一个标准的 MoE 层包含四个阶段：

1
2
3
4
5
6
7
8
9
10
11
12
13

MoE Layer
│
├── ① Router → topk_ids, topk_weights
│
├── ② Token Dispatch (All-to-All) ← A2A backend
│
├── ③ Expert Compute ← Runner Backend (本文主角)
│   ├── W1 GEMM (gate + up 融合)
│   ├── SwiGLU 激活
│   ├── W2 GEMM (down projection)
│   └── Weighted sum reduce
│
└── ④ Token Combine ← A2A backend

Runner Backend 只负责第 ③ 步，不同 backend 优化的是同一组计算在不同量化格式下的执行效率：

Marlin MoE 完整数据流

以 DeepSeek-V4 配置为例：hidden_size=7168, intermediate_size=3072, num_experts=384, topk=6, 权重为 MXFP4 格式。

Step 1: Block Size 启发式选择

1
2
3
4
5
6
7
8
9

M, K = hidden_states.shape  # M=tokens, K=7168
E = w1.shape[0]  # num_experts=384
N = w2.shape[1] * 16  # 3072 (Marlin 打包因子 16)
topk = topk_ids.shape[1]  # 6

# 启发式选择 M 方向分块大小
for block_size_m in [8, 16, 32, 48, 64]:
    if M * topk / E / block_size_m < 0.9:
        break

设计意图：

• Context 阶段（M 大）：用大 block（64），让每个 block 尽量填满
• Generation 阶段（M 小）：用小 block（8），避免最后一个 block 填充率过低

潜在问题：假设 token 在专家间均匀分布，但 Router 可能有热点。热点专家需要多个 block，最后一个 block 可能只填了 10%，浪费计算。

Step 2: Token-Expert 对齐

1
2
3

sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size(
    topk_ids, block_size_m, global_num_experts
)

作用：

1. 按 expert_id 对 token 排序
2. 按 block_size_m 对齐（padding），让每个专家的 token 数是 block 的倍数
3. 返回排序后的 token 索引和每个 block 对应的 expert_id

目的：让后续 Marlin GEMM kernel 能用 block-sparse 方式高效执行。

Step 3: W1 GEMM (gate + up 融合)

1
2
3
4
5
6
7
8
9
10
11
12
13
14

intermediate_cache1 = moe_wna16_marlin_gemm(
    hidden_states,              # [M, 7168]
    intermediate_cache1,        # output buffer
    w1,                         # [E, 7168/pack, 2*3072*pack]
    w1_scale,                   # 量化 scale
    sorted_token_ids,           # token → expert 映射
    expert_ids,                 # 每个 block 的 expert_id
    num_tokens_post_padded,     # padding 后的总 token 数
    topk_weights,               # 路由权重
    moe_block_size=block_size_m,
    top_k=topk,
    mul_topk_weights=False,     # ← W1 阶段不乘路由权重
    size_m=M, size_n=2*N, size_k=K,
)

关键点：

• mul_topk_weights=False：W1 输出是纯 GEMM 结果，不乘权重
• W1 权重是 gate 和 up 融合 的：[E, K/pack, 2*N*pack]
• 输出 [M*topk, 2*N]：前半 N 是 gate，后半 N 是 up

为什么 W1 不乘权重？

假设 token 的 top-2 是 Expert 3 (weight=0.6) 和 Expert 7 (weight=0.4)：

1
2
3
4
5
6
7
8

W1 输出（不乘权重）:
  [Expert 3 的 gate/up, Expert 7 的 gate/up]

经过 SwiGLU:
  [Expert 3 的 activated, Expert 7 的 activated]

W2 输出（乘权重）:
  Expert 3 输出 × 0.6 + Expert 7 输出 × 0.4

如果 W1 就乘了权重，SwiGLU 激活函数作用在”已经被压缩的信号”上，精度会下降。

Step 4: SwiGLU + Clamp

1
2
3
4
5

if clamp_limit is not None:
    # DeepSeek-V4: swiglu_limit=10.0
    swiglu_limit_func(intermediate_cache2, intermediate_cache1, clamp_limit)
else:
    silu_and_mul(intermediate_cache1, intermediate_cache2)

输入 [M*topk, 2*N] 拆成 gate 和 up：

1
2
3
4
5

gate = input[:, :N]      # SwiGLU gate branch
up   = input[:, N:]      # SwiGLU up branch

# SiLU(clamp(gate)) * clamp(up)
output = F.silu(torch.clamp(gate, max=10)) * torch.clamp(up, -10, 10)

Clamp 的作用：防止激活值爆炸，避免量化误差被放大。DeepSeek-V4 的 swiglu_limit=10.0 是经验值。

Step 5: W2 GEMM (down projection)

1
2
3
4
5
6
7
8
9
10
11
12
13
14

intermediate_cache3 = moe_wna16_marlin_gemm(
    intermediate_cache2,        # [M*topk, 3072]
    intermediate_cache3,        # output buffer
    w2,                         # [E, 3072/pack, 7168*pack]
    w2_scale,
    sorted_token_ids,
    expert_ids,
    num_tokens_post_padded,
    topk_weights,
    moe_block_size=block_size_m,
    top_k=1,                    # ← 注意这里是 1
    mul_topk_weights=True,      # ← W2 阶段乘路由权重
    size_m=M*topk, size_n=K, size_k=N,
).view(-1, topk, K)            # reshape 为 [M, topk, K]

关键变化：

• top_k=1：W2 的输入已经是展开的 M×topk 个 token，每个只对应 1 个专家
• mul_topk_weights=True：在 GEMM 内部就乘上路由权重，融合减少一次 memory pass

Step 6: 跨专家归约

1
2
3
4
5
6

if is_mxfp4_marlin:
    # MXFP4：用 torch.sum（atomic_add 不支持 BF16 on SM < 90）
    output = torch.sum(intermediate_cache3, dim=1)  # [M, topk, K] → [M, K]
else:
    # 普通 INT4：用 CUDA kernel 做 sum + scale
    moe_sum_reduce(intermediate_cache3, output, routed_scaling_factor)

routed_scaling_factor=2.5（DeepSeek-V4）在这里乘进去。

MXFP4 特殊处理

DeepSeek-V4 的专家权重是 MXFP4 格式（4-bit 分块量化），Marlin kernel 对此有特殊路径：

1
2
3
4
5
6
7

is_mxfp4_marlin = (
    num_bits == 4
    and w1_zeros is None           # 无 zero point
    and w2_zeros is None
    and w1_scale.dtype == torch.float8_e8m0  # scale 是 E8M0 格式
    and w2_scale.dtype == torch.float8_e8m0
)

MXFP4 + E8M0 scale 要求激活必须是 BF16（不支持 FP16）：

1
2

if is_mxfp4_marlin and hidden_states.dtype == torch.float16:
    marlin_hidden_states = hidden_states.to(torch.bfloat16)  # 强制转 BF16

MXFP4 特殊处理：

• 不用 atomic_add（因为 BF16 的 atomic_add 在 SM < 90 上不支持）
• 改用 torch.sum 替代归约
• Scale 是 E8M0 格式（8-bit exponent-only），和普通 INT4 的 per-tensor/per-channel scale 不同

内存布局优化

1
2
3
4
5
6
7

# 两个中间 buffer 共享底层存储
intermediate_cache13 = torch.empty(
    (M * topk * max(2*N, K),),  # 按 W1 和 W2 输出中较大的分配
    device=device, dtype=dtype,
)
intermediate_cache1 = intermediate_cache13[:M*topk*2*N].view(-1, 2*N)
intermediate_cache3 = intermediate_cache13[:M*topk*K].view(-1, K)

节省显存：W1 的输出 [M*topk, 2*N] 在被 SwiGLU 消费后就不需要了，W2 可以写入同一区域。

但有个细节：W2 GEMM 的输入是 intermediate_cache2（SwiGLU 输出，维度 N），不是 intermediate_cache1（维度 2N）。所以实际复用链是：

1
2
3

intermediate_cache1 [M*topk, 2*N] → (SwiGLU) → intermediate_cache2 [M*topk, N]
                                      ↓
intermediate_cache3 [M*topk, K] ← (W2 GEMM 输出)

intermediate_cache2 如果能复用 intermediate_cache1 的后半部分（up 分支在 SwiGLU 后就没用了），可以进一步节省显存。

Marlin 在 SGLang 中的注册

1
2
3

@register_fused_func("none", "marlin")
def fused_experts_none_to_marlin(dispatch_output, quant_info, runner_config):
    ...

“none” 是 A2A backend（无 EP，纯 TP），“marlin” 是 runner backend。

这意味着 Marlin MoE 目前只支持纯 TP 模式（每张卡有所有专家的 1/N 权重），不支持 EP（Expert Parallelism）。

为什么 Marlin 没做 EP？

1. MXFP4 权重大小不是瓶颈

   • 384 专家 × 33MB/8 = 1.6GB/卡（TP=8）
   • 如果 EP=8，每卡 48 专家 × 33MB = 1.6GB/卡
   • 权重大小一样，EP 的优势（减少单卡显存）消失

2. 通信量对比

   • TP 的 AllReduce：2 × B × hidden（W1 + W2 各一次）
   • EP 的 All-to-All：topk × B × hidden = 6 × B × hidden（topk=6）
   • EP 通信量是 TP 的 3 倍

3. Marlin 是量化 kernel，EP 是通信问题

   • 理论上可以接 deepep，但 MXFP4 + 高 topk 场景下收益不大
   • SGLang 团队可能评估后觉得优先级不高

完整数据流图

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24

hidden_states [M, 7168] (BF16)
│
▼ moe_align_block_size(topk_ids)
│ → sorted_token_ids, expert_ids (按专家排序+对齐)
│
▼ moe_wna16_marlin_gemm (W1, gate+up fused)
│ 每个 block 读对应 expert 的 MXFP4 权重 → 反量化 → GEMM
│ [M, 7168] × [E, 7168/pack, 2*3072*pack] → [M*6, 6144]
│ mul_topk_weights=False (不乘权重)
│
▼ SwiGLU + clamp (swiglu_limit=10.0)
│ gate = clamp([:, :3072], max=10)
│ up = clamp([:, 3072:], -10, 10)
│ → SiLU(gate) * up → [M*6, 3072]
│
▼ moe_wna16_marlin_gemm (W2, down)
│ [M*6, 3072] × [E, 3072/pack, 7168*pack] → [M*6, 7168]
│ mul_topk_weights=True (内部乘路由权重)
│ → reshape [M, 6, 7168]
│
▼ torch.sum(dim=1) 或 moe_sum_reduce
│ [M, 6, 7168] → [M, 7168] (× routed_scaling_factor=2.5)
│
▼ output [M, 7168] (BF16)

总结

Marlin MoE 是 MXFP4 量化 MoE 的高效实现，通过 block-sparse GEMM、SwiGLU 融合、内存复用等手段优化计算。在 DeepSeek-V4 的 MXFP4 + 高 topk 场景下，选择纯 TP 而非 EP 是理性的设计决策。

参考资料

• SGLang Marlin Runner 源码
• SGLang MoE 文档
• DeepSeek-V4 技术报告

DeepSeek-V4 KV Cache 深度分析

DeepSeek-V4 series incorporate several key upgrades: (1) hybrid attention architecture that combines Compressed Sparse Attention (CSA) and Heavily Compressed Attention (HCA); (2) Manifold-Constrained Hyper-Connections (mHC); (3) Muon optimizer.

DeepSeek-V4 彻底改变了 KV Cache 的设计，从 V3 的 MLA（Multi-head Latent Attention）转向 CSA + HCA 混合注意力架构，通过序列维度压缩而非仅靠隐维度压缩来大幅降低 KV Cache 开销。

附录：Python 计算器代码

以下代码可直接运行，计算任意序列长度下的 KV Cache 大小：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

# DeepSeek-V4-Pro KV Cache Size Calculator!

# 参考文献：
# - config.json: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/config.json
# - model.py: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/model.py
# - paper: https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/DeepSeek_V4.pdf!

# ── 模型参数（来自 config.json） ──────────────────────────────!
NUM_LAYERS = 61
HEAD_DIM = 512  # c: KV entry 维度（RoPE 包含在 512 内）
ROPE_DIM = 64   # r: 最后 64 维携带 RoPE
NOPE_DIM = HEAD_DIM - ROPE_DIM  # 448
WINDOW = 128    # w: 滑动窗口大小
INDEX_DIM = 128  # c_I: indexer KV 维度
CSA_RATIO = 4     # CSA 压缩比
HCA_RATIO = 128   # HCA 压缩比
N_CSA = 30   # CSA 层数（2,4,6,...,60）
N_HCA = 31   # HCA 层数（0,1,3,5,...,59）

# ── 精度配置 ──────────────────────────────────────!
def bf16_bytes():
    """BF16 基准：所有元素 2 bytes。"""
    s_kv = HEAD_DIM * 2  # 512 × 2 = 1024
    s_idx = INDEX_DIM * 2  # 128 × 2 = 256
    return s_kv, s_idx

def mixed_bytes():
    """混合精度部署：nope FP8 (1B) + rope BF16 (2B) + indexer FP4 (0.5B)。"""
    s_kv = NOPE_DIM * 1 + ROPE_DIM * 2  # 448×1 + 64×2 = 576
    s_idx = INDEX_DIM * 0.5  # 128 × 0.5 = 64 (FP4)
    return s_kv, s_idx

# ── 核心计算 ──────────────────────────────────────!
def calc_kvcache(seq_len: int, precision="bf16"):
    """
    计算 DeepSeek-V4-Pro 的 KV Cache 大小。

    Args:
        seq_len: 序列长度（token 数）
        precision: "bf16" 或 "mixed"（FP8+BF16+FP4）

    Returns:
        包含各组件和总大小的字典（字节）
    """
    s_kv, s_idx = bf16_bytes() if precision == "bf16" else mixed_bytes()

    # CSA 层 (ratio=4): Shared-KV + Indexer + SWA
    csa_compressed = (seq_len // CSA_RATIO + WINDOW) * s_kv
    csa_indexer = (seq_len // CSA_RATIO) * s_idx
    csa_swa = WINDOW * s_kv
    csa_per_layer = csa_compressed + csa_indexer + csa_swa

    # HCA 层 (ratio=128): Shared-KV + SWA（无 Indexer）
    hca_compressed = (seq_len // HCA_RATIO + WINDOW) * s_kv
    hca_swa = WINDOW * s_kv
    hca_per_layer = hca_compressed + hca_swa

    total = N_CSA * csa_per_layer + N_HCA * hca_per_layer

    return {
        "csa_per_layer": csa_per_layer,
        "hca_per_layer": hca_per_layer,
        "csa_total": N_CSA * csa_per_layer,
        "hca_total": N_HCA * hca_per_layer,
        "total_bytes": total,
        "total_gib": total / (1024**3),
        # 组件分解
        "csa_shared_kv": N_CSA * (seq_len // CSA_RATIO) * s_kv,
        "csa_indexer": N_CSA * (seq_len // CSA_RATIO) * s_idx,
        "hca_shared_kv": N_HCA * (seq_len // HCA_RATIO) * s_kv,
        "all_windows": (N_CSA + N_HCA) * WINDOW * s_kv,
    }

def fmt(b):
    """格式化字节数为可读字符串。"""
    if b >= 1024**3:
        return f"{b / (1024**3):.2f} GiB"
    return f"{b / (1024**2):.2f} MiB"


def report(seq_len: int):
    print("=" * 60)
    print(f" DeepSeek-V4-Pro KV Cache | seq_len = {seq_len:,}")
    print("=" * 60)
    print(f" {N_CSA} CSA 层 (ratio={CSA_RATIO}) + {N_HCA} HCA 层 (ratio={HCA_RATIO})")
    print(f" head_dim={HEAD_DIM} (nope={NOPE_DIM} + rope={ROPE_DIM})")
    print(f" index_dim={INDEX_DIM}, window={WINDOW}")
    print()

    for prec in ["bf16", "mixed"]:
        r = calc_kvcache(seq_len, prec)
        s_kv, s_idx = bf16_bytes() if prec == "bf16" else mixed_bytes()
        label = "BF16 基准" if prec == "bf16" else "FP8 kv + BF16 rope + FP4 indexer"
        print(f" ── {label} (S_kv={s_kv}, S_idx={s_idx}) ──")
        print(f"  CSA 单层: {fmt(r['csa_per_layer']):>10s}  HCA 单层: {fmt(r['hca_per_layer']):>10s}")
        print(f"  组件分解:")
        for name, val in [
            ("CSA Shared-KV 压缩", r["csa_shared_kv"]),
            ("CSA Indexer", r["csa_indexer"]),
            ("HCA Shared-KV 压缩", r["hca_shared_kv"]),
            ("滑动窗口 (61 层)", r["all_windows"]),
        ]:
            print(f"    {name:30s} {fmt(val):>10s} ({val/r['total_bytes']*100:.1f}%)")
        print(f" {'─' * 56}")
        print(f" {'Total':30s} {fmt(r['total_bytes']):>10s}")
        print()

if __name__ == "__main__":
    report(1_048_576)  # 1M tokens

运行结果（Python 3.x）：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

============================================================
 DeepSeek-V4-Pro KV Cache | seq_len = 1,048,576
============================================================
 30 CSA 层 (ratio=4) + 31 HCA 层 (ratio=128)
 head_dim=512 (nope=448 + rope=64)
 index_dim=128, window=128

 ── BF16 基准 (S_kv=1024, S_idx=256) ──
  CSA 单层: 320.12 MiB  HCA 单层: 8.12 MiB
  组件分解:
    CSA Shared-KV 压缩             7,500.00 MiB (77.9%)
    CSA Indexer                         1,920.00 MiB (19.5%)
    HCA Shared-KV 压缩               248.00 MiB (2.5%)
    滑动窗口 (61 层)                       7.50 MiB (0.1%)
─────────────────────────────────────────────────────────────
 Total                         9,675.50 MiB (9.62 GiB)

── FP8 kv + BF16 rope + FP4 indexer (S_kv=576, S_idx=64.0) ──
  CSA 单层: 160.07 MiB  HCA 单层: 4.57 MiB
  组件分解:
    CSA Shared-KV 压缩             4,218.75 MiB (87.4%)
    CSA Indexer                           480.00 MiB (9.7%)
    HCA Shared-KV 压缩                   139.50 MiB (2.8%)
    滑动窗口 (61 层)                               4.25 MiB (0.1%)
─────────────────────────────────────────────────────────────
 Total                         4,822.50 MiB (4.83 GiB)

对比 V3.2（61 层，每 token 1152 bytes）：

• V3.2 BF16: 61 × 1,048,576 × 1152 / 1024³ ≈ 83.9 GiB
• V4-Pro BF16: 9.62 GiB (8.7× 压缩)
• V4-Pro 混合精度: 4.83 GiB (17.4× 压缩) ✅!

一、核心架构参数!

数据来源：config.json

compress_ratios 数组 → 层类型

compress_ratios 有 62 个元素（61 层 + 1 层 MTP）：

1
2
3

索引: 0  1  2  3  4  5  ... 59 60 61
值:   128 128 4  128 4  128 ... 128 4  0
类型: HCA HCA CSA HCA CSA HCA ... HCA CSA MTP

二、KV Cache 结构革命!

V3 MLA vs V4 CSA/HCA!

V4 KV entry 构成（512 维）

1
2
3
4
5
6
7

KV entry (512 dims)
┌────────────────────────────┬──────────────┐
│ nope 维度 (448 个元素) │ rope 维度 (64)│
│ FP8: 448 × 1 byte     │ BF16: 64 × 2  │
│ = 448 bytes            │ = 128 bytes │
└────────────────────────────┴──────────────┘
合计 = 576 bytes/entry (混合精度)

论文原文（Section 2.3.4）：

“我们采用混合存储格式：RoPE 维度使用 BF16 精度，其余维度使用 FP8 精度。”

代码证据!

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

# model.py, Attention.__init__
self.head_dim = args.head_dim  # 512
self.rope_head_dim = args.rope_head_dim  # 64
self.nope_head_dim = self.head_dim - self.rope_head_dim  # 448!

# KV cache 分配：512 维，不是 576
self.register_buffer("kv_cache",
    torch.zeros(args.max_batch_size, kv_cache_size, self.head_dim),
    persistent=False)

# 写入时：前 448 维 FP8，后 64 维 BF16
kv = self.wkv(x)  # Linear(7168 → 512)
kv = self.kv_norm(kv)
apply_rotary_emb(kv[..., -64:], freqs_cis)  # RoPE 对后 64 维原地施加
act_quant(kv[..., :-64], 64, scale_fmt, scale_dtype, True)  # FP8 量化（前 448 维）

结论：V4 的 512 维包含 RoPE（后 64 维），所以存储是 576 bytes/entry（混合精度），相比 V3 的 640 bytes/token 还小。

三、压缩机制详解!

1. HCA（Heavily Compressed Attention）— 重度压缩!

• 压缩比：128:1（每 128 个 token → 1 个压缩 KV entry）
• 无 overlap（论文明确：”does not perform overlapped compression”）
• 无稀疏选择：对所有压缩条目做全量密集注意力（提供全局”低分辨率概览”）
• 滑动窗口：同样保留最近 128 个 token!

单层 KV Cache 计算（1M tokens）：

2. CSA（Compressed Sparse Attention）— 轻度压缩!

• 压缩比：4:1（每 4 个 token → 1 个压缩 KV entry）
• 窗口大小：8 个 token（含 50% overlap，步长=4）
• 稀疏选择：Lightning Indexer（FP4 精度）选取 top-1024 最相关的压缩块
• 滑动窗口：保留最近 128 个 token 的原始未压缩 KV!

单层 KV Cache 计算（1M tokens）：

四、索引器（Indexer）机制（CSA 独有）!

索引器是 CSA 的”导航系统”，用来决定哪些压缩块最值得关注。

结构

• 独立 Compressor：自己的 wkv 和 wgate 权重（head_dim=128，不是 512）
• FP4 量化：Hadamard 旋转 + fp4_act_quant（全 128 维）
• 计算流程：

  1 2 3 4 5 6

  Query token t → indexer query (128 维) → 和所有压缩块的索引器 key 算分 → Top-k (1024) → 选出最相关的 1024 个压缩块 ```! ### Indexer Cache 大小（1M tokens）

  (1M/4) × 128 dims × 0.5 byte (FP4) = 16 MB/层

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

**论文 Section 5.2.1**：
> "索引器中的 QK 激活被缓存、加载，并完全以 FP4 精度计算。"

---

## 五、RoPE 的特殊处理!

### 三处应用 RoPE（Section 2.3.3）

1. **Query 向量** (q_t,i) → last 64 dims
2. **KV entry 向量** (C_Comp) → last 64 dims
3. **Core attention 输出** (o_t,i) → last 64 dims（位置 = -i）

### 为什么输出也要加 RoPE？!

因为 KV entry 同时作为 K 和 V，压缩后的 entry 带绝对位置信息。如果直接输出：

o_t,i = Σ attn_weight_j × C_Comp_j(R(j)) ← 携带绝对位置 R(j)

1
2
3

对输出施加 RoPE(position=-i)：

R(-i) × o_t,i = Σ attn_weight_j × C_Comp_j(R(j-i)) ← 转为相对位置

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

**论文原文**：
> "注意力输出的贡献将与 query 和 KV entry 之间的距离相关。"

---

## 六、精度配置!

### 两种精度模式!

| 符号 | 含义 | BF16 基准 | 混合精度部署 |
|------|------|-----------|--------------|
| $S_{kv}$ | Shared-KV 每元素字节 | $512 \times 2 = 1024$ | $448 \times 1 + 64 \times 2 = 576$ |
| $S_{idx}$ | Indexer 每元素字节 | $128 \times 2 = 256$ | $128 \times 0.5 = 64$ (FP4) |

### 论文依据!

**Section 2.3.4**：
> "RoPE 维度使用 BF16 精度，其余维度使用 FP8 精度。"

**Section 5.2.1**：
> "索引器中的 QK 激活完全以 FP4 精度计算。"

---

## 七、总 KV Cache 计算公式!

### 通用公式!

设序列长度 $N$，CSA 层数 $N_{CSA}=30$，HCA 层数 $N_{HCA}=31$：

$$\text{Total KV} = N_{CSA} \times \text{PerCSA} + N_{HCA} \times \text{PerHCA}$$

其中：

$$\text{PerCSA} = \underbrace{(\frac{N}{4} + w) \times S_{kv}}_{\text{Shared-KV}} + \underbrace{\frac{N}{4} \times S_{idx}}_{\text{Indexer}}$$

$$\text{PerHCA} = \underbrace{(\frac{N}{128} + w) \times S_{kv}}_{\text{Shared-KV only}}$$

---

## 八、数值结果（N = 1,048,576 = 1M tokens）!

### BF16 基准（vLLM 博客算法）

**CSA 30 层**：

| 组件 | 计算 | 结果 |
|------|------|------|
| Shared-KV 压缩 | 262,144 × 1024 | 256.00 MiB |
| Shared-KV 窗口 | 128 × 1024 | 0.125 MiB |
| Indexer 压缩 | 262,144 × 256 | 64.00 MiB |
| **CSA 单层** | | **320.13 MiB** |
| **30 层 CSA** | 320.13 × 30 | **9,603.75 MiB** |

**HCA 31 层**：

| 组件 | 计算 | 结果 |
|------|------|------|
| Shared-KV 压缩 | 8,192 × 1024 | 8.00 MiB |
| Shared-KV 窗口 | 128 × 1024 | 0.125 MiB |
| **HCA 单层** | | **8.13 MiB** |
| **31 层 HCA** | 8.13 × 31 | **251.88 MiB** |

**总计（BF16）**：9,603.75 + 251.88 = **9,855.63 MiB ≈ 9.62 GiB** ✅（和 vLLM 博客完全一致）

### 混合精度实际部署（FP8 + BF16 + FP4）

| 组件 | bytes/entry | 30 层 CSA | 31 层 HCA |
|------|---------------|------------|------------|
| CSA 压缩 KV | 576 | 4.23 GiB | — |
| CSA Indexer | 64 (FP4) | 0.47 GiB | — |
| CSA SWA | 576 | 2.11 MiB | — |
| HCA 压缩 KV | 576 | — | 0.14 GiB |
| HCA SWA | 576 | — | 2.18 MiB |
| **总计** | | **~4.84 GiB** | |

**压缩比**：
- V3.2 KV Cache (61 层) = 83.9 GiB
- V4-Pro 混合精度 = 4.84 GiB
- **压缩比**：83.9 / 4.84 ≈ **17.3×** ✅!

---

## 九、各组件占比（BF16 基准）!

| 组件 | MiB | 百分比 |
|------|-----|--------|
| CSA Shared-KV 压缩 | 7,680.00 | 77.9% ← 绝对主体 |
| CSA Indexer | 1,920.00 | 19.5% ← 不可忽略 |
| HCA Shared-KV 压缩 | 248.00 | 2.5% ← 128× 压缩极高效 |
| 滑动窗口 (全部 61 层) | 7.63 | 0.1% ← 可忽略 |
| **总计** | **9,855.63** | **100%** |

CSA 的 Indexer 占了近 20%，是不能漏算的部分。

---

## 十、为什么能这么小？关键洞察!

1. **序列维度压缩**（V3 只压缩隐维度，V4 同时压缩序列长度）
   - CSA: L → L/4
   - HCA: L → L/128

2. **稀疏注意力**（CSA 只访问 top-1024 个压缩块，而非全部 L/4 个）

3. **分层设计**：CSA 负责精准的中程依赖，HCA 负责全局概览，互为补充

4. **混合精度**：
   - RoPE 维度：BF16（2 bytes，保证位置精度）
   - 其余维度：FP8（1 byte，节省空间）
   - Indexer：FP4（0.5 byte，最激进）

5. **磁盘 KV Cache**：压缩后的 KV 块可以落盘存储，共享前缀的请求复用缓存*

---

## 十一、Forward 计算流程图（V4-Pro）!

                Hidden State [batch, seq_len, 7168]
                          │
          ┌───────────────┴───────────────┐
          ▼                               ▼
     Q Down-proj                    KV 通路
[7168 → 1536]                  [7168 → 512] (wkv)
          │                               │
          ▼                               ▼
 Q Up-proj (1536→128×512)    ┌──→ 压缩 (CSA: ratio=4 / HCA: ratio=128)
          │                    │     ├─→ RoPE (后 64 维, BF16)
Q [batch, seq_len, 128, 512] │     ├─→ FP8 量化 (前 448 维)
          │                    │     └─→ 写入 KV cache
┌─→ RoPE (后 64 维)       │
│         │                    │
│    Q [batch, seq_len, 128, 512]  KV cache [batch, N/ratio, 512]
│         │                    │
│         └──→ Indexer (仅 CSA) ──→ Top-k → 1024 个选中
│                                      │
└──→ Attention(Q, K=selected_1024, V=selected_1024)
          │
          └──→ + SWA (最近 128 个 token)
          │
          ▼
    Output [batch, seq_len, 7168]

---

## 十二、和 vLLM 部署对齐!

vLLM 博客明确验证了 V4 的计算：

> "using fp4 indexer cache and fp8 attention cache, which further reduces the KV cache size by roughly **2×** compared to the bf16 estimate!"

| 来源 | KV Cache (1M tokens) | 说明 |
|------|---------------------|------|
| 论文 Figure 1 | 9.62 GiB (10% of V3.2) | BF16 估计 |
| vLLM 博客 | ~4.8 GiB | FP4+FP8 混合精度 |
| **本文计算** | **4.84 GiB** | 精确对齐 ✅ |

---

## 十三、总结!

DeepSeek-V4 通过**序列维度压缩 + 混合精度存储 + 稀疏选择**，将 1M tokens 的 KV Cache 从 V3.2 的 83.9 GiB 压缩到 **4.84 GiB**（17.3× 压缩）。

核心创新：
- **CSA**：4:1 压缩 + overlap + Indexer 稀疏选择
- **HCA**：128:1 重度压缩，提供全局概览
- **混合精度**：RoPE (BF16) + 内容 (FP8) + Indexer (FP4)

这套架构让百万 token 上下文从"不可能"变成"日常可用" 🐾!

---

## 参考!

1. [DeepSeek-V4 论文](https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/DeepSeek_V4.pdf) (2026) Section 2.3
2. [vLLM 博客：DeepSeek V4 in vLLM](https://github.com/vllm-project/vllm-project.github.io/blob/main/_posts/2026-04-24-deepseek-v4.md) (2026-04-24)
3. [config.json](https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/config.json)
4. 代码追踪: model.py, compressor.py, c4.cuh
5. 论文 Section 2.3.3: Partial Rotary Positional Embedding
6. 论文 Section 5.2.1: FP4 Quantization-Aware Training

概述

SGLang 的 Pipeline Parallelism (PP) 模式下，HiCache 负责 KV cache 的异步 prefetch（Host→GPU）和 backup（GPU→Host）。本文从 PP 调度器的外层事件循环出发，追踪 Load 和 Write 的完整时序，揭示 write_ack 比 load_ack 多延迟的根本原因。

一、PP 事件循环

每个 iteration 执行以下 4 步：

1
2
3
4
5

iter=N:
  ① check_hicache_events()      ← 查询 HiCache 异步事件（load_ack / write_ack）
  ② get_next_batch_to_run()     ← 选下一批 batch（prefix match → eviction → load）
  ③ _pp_launch_batch()          ← launch forward
  ④ _pp_process_batch_result()  ← 处理上一批 batch 的结果（insert → write）

核心设计：process_batch_result 处理的是 mbs[next_mb_id]——上一轮 launch 的 batch。同一个 iter 内，调度器同时在处理两个不同 batch 的生命周期阶段。

二、Load 时序（Host→GPU）

触发时机

iter=N 的 get_next_batch_to_run —— prefill 阶段 prefix match 发现 host_hit，发起 load_back。

完整时序

关键：load 是 prefetch——在 forward 之前触发，CUDA copy 和 forward 可以重叠（逐层等待、逐层执行）。

时序图

1
2
3
4
5
6
7
8
9
10
11
12

iter=N:
  ① check_hicache_events()
     └─ 消费更早的 load_ack
  ② get_next_batch_to_run()
     └─ match_prefix → host_hit → load_back() → start_loading()
     └─ CUDA copy 启动，event 入队
  ③ _pp_launch_batch()
     └─ forward 逐层等待 load 完成

iter=N+1:
  ① check_hicache_events()
     └─ loading_check() → event.query()=True → load_ack ✓

延迟：load_ack 比 load 发起晚 1 iter。

三、Write 时序（GPU→Host）

触发时机

iter=N 的 process_batch_result —— 处理上一轮 launch 的 batch 结果，insert 时触发 write_backup。

完整时序

关键：write 是 post-write——必须在 forward 算完、拿到完整 KV cache 后才能触发。

时序图

1
2
3
4
5
6
7
8
9
10
11
12
13

iter=N-1:
  ② get_next_batch_to_run() → 选出 Prefill A
  ③ _pp_launch_batch() → forward(Prefill A)

iter=N:
  ④ _pp_process_batch_result()
     └─ 处理 iter=N-1 的 Prefill A
     └─ insert() → write_backup() → start_writing()
     └─ CUDA copy 启动，event 入队

iter=N+1:
  ① check_hicache_events()
     └─ writing_check() → event.query()=True → write_ack ✓

延迟：write_ack 比 write 发起晚 1 iter，但 write 本身比 forward 晚 1 iter（因为 process_batch_result 处理上一批 batch）。

四、延迟对比

以同一个 Prefill batch 为基准

追踪一个 Prefill batch 从 launch 到 write_ack 的完整生命周期：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

iter=1:
  ② get_next_batch_to_run()
     └─ 选出 Prefill A
     └─ match_prefix → host_hit → load_back() → start_loading()
  ③ _pp_launch_batch() → forward(Prefill A)

iter=2:
  ① check_hicache_events()
     └─ loading_check() → load_ack ✓
        ↑ load_ack: 1 iter delay

iter=3:
  ④ process_batch_result()
     └─ 处理 iter=1 的 Prefill A
     └─ insert() → write_backup() → start_writing()
  ① check_hicache_events()
     └─ (write_ack 还没完成)

iter=4:
  ① check_hicache_events()
     └─ writing_check() → write_ack ✓
        ↑ write_ack: 3 iter delay from Prefill launch

对比表

为什么 write_ack 多延迟？

两个因素叠加：

1. process_batch_result 滞后一轮

它处理 mbs[next_mb_id]（上一轮 launch 的 batch），以 PP2 为例偏移 2 iter。所以 write 比 load 晚 2 iter 才触发。

2. CUDA event query 需要等下一轮 check

无论 load 还是 write，ack 入队后最早等下一轮 check_hicache_events 才能消费。两者各加 1 iter。

综合：

1
2

Load:  iter=1 发起 → iter=1 ack 入队 → iter=2 消费
Write: iter=1 forward → iter=3 process_result 发起 → iter=3 ack 入队 → iter=4 消费

Load 是 prefetch（forward 之前），Write 是 post-write（forward 之后）。这个架构差异决定了 Write 必然比 Load 多延迟。

五、PP 间同步问题

PP0 和 PP1 共享同一个 HiCache 实例（radix tree + host memory）。由于 output relay 延迟，PP0 和 PP1 的 iter 进度存在 1-2 iter 的偏移，需要三层同步机制保证一致性。

5.1 逻辑时钟保证重放一致性

PP1 的 writing_check/loading_check 不再直接消费 ack，而是将事件通过 Gloo P2P 通道 replay 给 PP0。PP0 作为唯一的事件消费者，按逻辑时钟顺序处理所有 ack，确保 PP0 和 PP1 的 radix tree 操作顺序一致。

问题：早期实现中 PP1 的 writing_check() 绕过 check_hicache_events guard，直接消费 write_ack（即 ack theft），导致 PP0 端 pending event 永远无法完成，radix tree 分叉。

修复：将 PP rank 分支移入 writing_check/loading_check 内部，用 PPHiCacheEventsReq 控制请求替代 dict wrapper，强制 PP1 replay 事件给 PP0。

5.2 Count Sync 保证 CP 一致性

PP0 比 PP1 早 1-2 iter 积累 write ack（output relay 延迟），导致 ack_write_queue 积累差异。PR #22878 通过 piggybacking write-ack consumption counts 在 PP ranks 间同步：

1
2
3

iter=N:
  PP0: writing_check() → 消费 3 个 write_ack → count=3
  PP1: 等待 PP0 的 count → count=3 → radix tree 操作对齐

Count sync 确保 PP0 和 PP1 对同一个 radix tree node 的 checkpoint（CP）操作一致，避免一个 stage 认为 node 已 backup、另一个 stage 还在等待的情况。

5.3 PP1 同步消费 PP0 的 ack

PP1 不再独立消费 ack，而是通过同步机制确保 PP0 消费 ack 后，PP1 的 radix tree 状态与 PP0 对齐：

1
2
3
4
5

PP0: writing_check() → event.query()=True → 消费 write_ack
     → radix tree: insert() → finalize() → node 状态更新
     ↓ TP all_reduce sync
PP1: 等待 PP0 完成 → radix tree: 同步执行 finalize()
     → node 状态与 PP0 一致

三层保障：

核心目标：无论 PP0 和 PP1 的 iter 偏移多少，radix tree 的结构和节点状态在两个 stage 上始终保持一致。

六、总结

核心结论：

1. PP 调度器每个 iter 同时做两件事：get_next_batch_to_run 选下一批 batch，process_batch_result 处理上一批 batch 的结果
2. Load 是 prefetch（forward 之前触发），Write 是 post-write（forward 之后触发）
3. process_batch_result 处理 mbs[next_mb_id] 的 iter 偏移是 Write 延迟的根本原因
4. PP 间需要 count sync 补偿 output relay 带来的时序差异

本文基于 SGLang 源码分析，涉及文件：hiradix_cache.py、cache_controller.py。

Pipeline Parallelism（PP）是多卡推理中的核心技术之一。SGLang 的 PP 实现有一套独立的事件循环和调度机制，和普通的 single-batch 路径完全不同。本文将从代码出发，深入分析 SGLang 的 PP 实现。

为什么需要 Pipeline Parallelism

单机推理的瓶颈

单卡推理时，GPU 利用率受限于模型大小和显存。一个 70B 模型单卡放不下，必须拆分到多卡。

模型并行方案对比

SGLang 通常 TP + PP 组合使用：节点内 TP，跨节点 PP。

SGLang PP 架构概览

PP 路径的独立事件循环

SGLang 的调度器有三条事件循环路径：

• event_loop_normal - 普通单 batch 调度
• event_loop_overlap - 计算和通信 overlap 的调度
• event_loop_pp - PP 路径的独立事件循环

PP 路径有自己独立的事件循环，定义在 scheduler_pp_mixin.py 的 SchedulerPPMixin 类中。它和 normal/overlap 路径完全隔离，因为 PP 的通信模式（send/recv proxy tensors 跨 stage）和单机调度逻辑差异太大，硬塞进同一个 event loop 会很脏。

event_loop_pp 核心流程

1
2
3
4
5
6
7
8
9
10
11

# scheduler_pp_mixin.py 第 72-145 行
while True:
    for mb_id in range(pp_loop_size):  # 遍历 microbatch
        1. recv_requests()                                  # 接收第 i 个请求（从上一个 PP stage 或 tokenizer）
        2. _pp_send_pyobj_to_next_stage()                   # 立即把第 i 个请求转发给下一个 PP rank（async，提前发送）
        3. get_next_batch_to_run()                          # 调度当前 microbatch 的 batch
        4. _pp_recv_proxy_tensors()                         # 接收第 i 个 proxy（hidden states，来自上一个 PP stage）
        5. _pp_launch_batch()                               # 执行第 i 个 batch 的 forward
        6. _pp_commit_send_output_work_and_preprocess...()  # 发送+接收 output + 预处理 next_mb_id 槽位的结果
        7. _pp_process_batch_result()                       # 处理 next_mb_id 槽位的 batch result
        8. _pp_send_dict_to_next_stage(msg_type="proxy")    # 发送第 i 个 proxy（hidden states）给下一个 PP stage

⚠️ step 5 和 step 6 的顺序取决于 pp_async_batch_depth：

• pp_async_batch_depth > 0：step 6（output 收发+预处理）在 step 5（forward）之前执行，与 forward 重叠，隐藏延迟
• pp_async_batch_depth == 0：step 6（output 收发+预处理）在 step 5（forward）之后执行，串行不重叠

上图展示的是 pp_async_batch_depth == 0 的顺序（step 6 在 step 5 之后）。

Microbatch 索引计算：

核心概念：commit = 等待上一轮的异步操作完成

所有 _pp_commit_comm_work(work) 做的事情就是：

1
2
3
4

def _pp_commit_comm_work(self: Scheduler, work: List[P2PWork]) -> None:
    for p2p_work in work:
        p2p_work.work.wait()  # 阻塞等待 NCCL/Gloo 异步通信完成
    work.clear()

即：commit 是一个同步屏障，确保上一轮发起的异步 send 已经完成，发送缓冲区可以安全复用。这是一种延迟等待模式–发送时不阻塞，到下一轮再等待完成，从而让发送和 GPU 计算重叠。

Microbatch 索引计算：

1
2

next_mb_id = (mb_id + 1) % self.pp_loop_size              # 当前 rank 下一个要处理的 microbatch 槽位
next_first_rank_mb_id = (mb_id + self.pp_size) % self.pp_loop_size  # first rank 对应的"下一轮"槽位

• next_mb_id：当前 rank 下一个要处理的 microbatch 槽位（偏移 1，因为每次迭代处理一个 mb）
• next_first_rank_mb_id：first rank 对应的”下一轮”microbatch 槽位（偏移 pp_size，因为流水线深度–first rank 产出的 batch 要经过 pp_size 个迭代才到 last rank 完成）

步骤与函数的对应关系：

关键设计：

• send req 提前：收到请求后立即转发给下一个 rank（step 2），不等 forward 完成，隐藏 P2P 通信延迟
• output 发送在 _pp_commit_send_output_work_and_preprocess_output_tensors 中处理：output 的发送和接收都在 step 6 中完成，last rank 把 output 发给 rank 0，中间 rank 转发
• output 处理时机取决于 pp_async_batch_depth：
  • > 0：output 收发+预处理在 forward 之前执行，与当前 forward 重叠，隐藏延迟
  • == 0：output 收发+预处理在 forward 之后执行，串行不重叠
• 延迟等待异步通信：_pp_commit_comm_work 在下一轮迭代中调用 .wait() 确保上一轮的 async send 已完成，让发送和 GPU 计算重叠
• 每个 microbatch 槽位独立管理状态
• async send + recv 实现计算和通信的 overlap
• 主循环轮转处理每个 microbatch，填充 pipeline bubble

完整流程示例（pp_size=2, pp_async_batch_depth=0）

以 pp_size=2, pp_async_batch_depth=0 为例，pp_loop_size=2，只有槽位 0 和 1：

1
2
3
4
5
6

参数：
  pp_size = 2（PP0 = first rank, PP1 = last rank）
  pp_async_batch_depth = 0
  pp_loop_size = 2（槽位：mb_id ∈ {0, 1}）
  next_mb_id = (mb_id + 1) % 2              # 另一个槽位
  next_first_rank_mb_id = (mb_id + 2) % 2 = mb_id  # 注意：next_first_rank_mb_id == mb_id

单次迭代（mb_id = N）的完整步骤：

步骤 1：恢复上下文

1
2
3
4

running_batch = running_mbs[N]
last_batch = last_mbs[N]
next_first_rank_mb_id = N          # 例如 N=0 时，next_first_rank_mb_id = 0
next_mb_id = (N + 1) % 2           # 例如 N=0 时，next_mb_id = 1（另一个槽位）

步骤 2：接收并处理请求

1
2

recv_reqs = recv_requests()
process_input_requests(recv_reqs)

• PP0：从 tokenizer_manager / detokenizer 收请求（zmq 通道）
• PP1：从 PP0 收转发过来的请求（P2P，含 PPHiCacheEventsReq 控制消息）

步骤 3：HiCache 事件同步 + 转发请求

1
2
3
4
5
6

if enable_hierarchical_cache:
    tree_cache.check_hicache_events()

if not is_last_rank:  # 即 PP0
    _pp_commit_comm_work(send_req_work)     # ← 等待上一轮的 send_req 异步操作完成
    send_req_work = async_send(pp_send_payload) → PP1

commit 含义：上一轮（槽位 (N-1)%2）发起的 send_req_work 可能还没完成，这里阻塞等它完成后，才能安全地复用发送缓冲区发送本轮的请求。

步骤 4：调度决策

1
2
3

mbs[N] = get_next_batch_to_run()
running_mbs[N] = running_batch
cur_batch = mbs[N]

步骤 5：接收 Proxy 张量

1
2
3
4

if cur_batch is not None:
    pp_proxy_tensors = _pp_recv_proxy_tensors()
    # PP0: 直接返回 None（第一个 stage 不需要接收）
    # PP1: 阻塞等待从 PP0 接收 hidden_states + residual (msg_type="proxy")

步骤 6：等待上一轮 Proxy 发送完成

1

_pp_commit_comm_work(send_proxy_work)  # ← 等待上一轮的 proxy 异步发送完成

commit 含义：上一轮（槽位 (N-1)%2）PP0 发起的 proxy 异步发送可能还没完成，这里确保它完成，释放发送缓冲区。

步骤 7：GPU Forward 计算

1
2
3
4
5

if cur_batch:
    result, launch_event = _pp_launch_batch(N, pp_proxy_tensors, ...)
    # 在 forward_stream 上执行 run_batch
    # 记录 CUDA event
    # PP1(last rank): 将 output 张量 push 到 last_rank_comm_queue

步骤 8：发送 Output + 接收 Output + 预处理

因为 pp_async_batch_depth == 0，这一步在 forward 之后执行（无法与 GPU 计算 overlap）。

1
2
3

_pp_commit_send_output_work_and_preprocess_output_tensors(
    next_first_rank_mb_id=N, next_mb_id=(N+1)%2,
)

内部逻辑（_pp_send_recv_and_preprocess_output_tensors）：

1
2
3
4
5
6
7
8
9
10
11
12
13

# 1. 等待上一轮的 output 异步发送完成
_pp_commit_comm_work(send_output_work)

# 2. 发送 output:
# PP1(last rank): 如果 mbs[next_first_rank_mb_id=N] 不为空，
#   从 last_rank_comm_queue 取出 output，异步发送给 PP0 (msg_type="output")
# PP0(非 last rank): 如果有上一轮暂存的 pp_outputs，转发给 PP1

# 3. 接收 output + 预处理:
# 如果 mbs[next_mb_id] 不为空:
#   阻塞接收 output 张量 (msg_type="output")
#   在 copy_stream 上执行 _pp_prep_batch_result()（组装 GenerationBatchResult + D2H 拷贝）
#   记录 d2h_event

commit 含义：上一轮（槽位 (N-1)%2）发起的 output 异步发送可能还没完成，这里确保完成后才能发起新的 output 发送。

步骤 9：后处理 next_mb_id 槽位的 Batch 结果

1
2
3
4
5

if mbs[next_mb_id] is not None:
    d2h_event.synchronize()  # 等待 copy_stream 上的 D2H 拷贝完成
    _pp_process_batch_result(mbs[next_mb_id], next_batch_result)
    # 更新请求状态、判断 EOS、发送给 detokenizer
    last_mbs[next_mb_id] = mbs[next_mb_id]

步骤 10：发送 Proxy 张量（仅 PP0）

1
2
3

if not is_last_rank and cur_batch:
    torch.cuda.current_stream().wait_event(launch_event)  # 等 forward 完成
    send_proxy_work = async_send(hidden_states + residual) → PP1  # msg_type="proxy"

步骤 11：保存状态

1

pp_outputs = next_pp_outputs  # 暂存本轮接收到的 output，供下一轮步骤 8 转发

完整时序（稳态，两个槽位交替）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

┌─────────────────────────────────────────────────────────────────┐
│ mb_id=0                                                          │
├─────────────────────────────────────────────────────────────────┤
│ 1. recv_requests + process                                       │
│ 2. commit(send_req_work[上轮mb=1的]) → async_send(reqs)         │
│ 3. get_next_batch_to_run → cur_batch                            │
│ 4. recv_proxy (PP1阻塞等PP0; PP0跳过)                            │
│ 5. GPU forward (forward_stream)                                  │
│ 6. commit(send_output_work[上轮mb=1的])                          │
│    send_output → recv_output → prep_batch_result (copy_stream)  │
│ 7. d2h_event.sync → process_batch_result(mbs[next_mb_id]的结果) │
│ 8. commit(send_proxy_work[上轮mb=1的])                           │
│ 9. PP0: wait(launch_event) → async_send(proxy) → PP1            │
│10. pp_outputs = next_pp_outputs                                  │
├─────────────────────────────────────────────────────────────────┤
│ mb_id=1                                                          │
├─────────────────────────────────────────────────────────────────┤
│ 1. recv_requests + process                                       │
│ 2. commit(send_req_work[上轮mb=0的]) → async_send(reqs)         │
│ 3. get_next_batch_to_run → cur_batch                            │
│ 4. recv_proxy (PP1阻塞等PP0; PP0跳过)                            │
│ 5. GPU forward (forward_stream)                                  │
│ 6. commit(send_output_work[上轮mb=0的])                          │
│    send_output → recv_output → prep_batch_result (copy_stream)  │
│ 7. d2h_event.sync → process_batch_result(mbs[next_mb_id]的结果) │
│ 8. commit(send_proxy_work[上轮mb=0的])                           │
│ 9. PP0: wait(launch_event) → async_send(proxy) → PP1            │
│10. pp_outputs = next_pp_outputs                                  │
└─────────────────────────────────────────────────────────────────┘

PP0 与 PP1 的交互时序（稳态）

1
2
3
4
5
6
7
8
9
10
11
12
13

时间轴 →

PP0 (mb_id=0):
recv_req → commit(send_req) → send_req→PP1 → schedule → [skip proxy]
→ FORWARD → commit(send_output)
→ send_output(转发) → recv_output(从PP1) → prep → process_result
→ commit(send_proxy) → send_proxy→PP1

PP1 (mb_id=0):
recv_req(从PP0) → process → schedule → recv_proxy(阻塞等PP0)
→ FORWARD → commit(send_output)
→ send_output→PP0 → recv_output(从PP0转发) → prep → process_result
→ commit(send_proxy)

关键依赖链：

• PP1 的 recv_proxy 必须等 PP0 的 send_proxy（上一轮的）完成
• PP0 的 recv_output 必须等 PP1 的 send_output 完成
• 因为 pp_async_batch_depth=0，output 的发送/接收在 forward 之后，无法与 GPU 计算重叠

所有 commit 汇总

模式统一：每个 async_send 在下一轮同一类型操作之前被 commit，形成”发起→（做其他事）→等待完成→再发起”的流水线模式。即使 pp_async_batch_depth=0 没有计算-通信 overlap，这种 commit 模式仍然保证了同类型通信操作之间不会冲突。

Microbatch：填充 Pipeline Bubble

问题：Pipeline Bubble

PP2（2 个 stage）如果只用一个 batch，会出现大量空闲：

1
2
3
4

时间 →
Stage 0: [ batch 0 ] [ batch 1 ]
Stage 1:              [ batch 0 ] [ batch 1 ]
              ↑ 空闲     ↑ 空闲

Stage 0 做完 batch 0 后必须等 Stage 1 处理完才能做 batch 1（因为只有一个 batch 在 pipeline 里流转），一半的时间都在空闲。这就是 pipeline bubble。

解法：Microbatch

把多个 batch 同时塞进 pipeline，这些 batch 就叫 microbatch：

1
2
3

时间 →
Stage 0: [ mb 0 ][ mb 1 ][ mb 2 ][ mb 3 ]
Stage 1:         [ mb 0 ][ mb 1 ][ mb 2 ][ mb 3 ]

Stage 0 做完 mb0 不用等，立刻做 mb1；Stage 1 也紧接着处理，bubble 大幅减少。

Microbatch 数量的确定

1
2

# scheduler_pp_mixin.py 第 514 行
self.pp_loop_size = self.pp_size + self.server_args.pp_async_batch_depth

• pp_size = 2（PP2）→ 至少 2 个 microbatch 槽位
• pp_async_batch_depth → 额外的 buffer 深度，进一步隐藏延迟

每个 microbatch 都有独立的状态：

1
2
3

self.mbs = [None] * self.pp_loop_size          # 当前 batch
self.last_mbs = [None] * self.pp_loop_size     # 最近一次处理的 batch
self.running_mbs = [ScheduleBatch(...) for _ in range(self.pp_loop_size)]  # running 状态

Microbatch vs 正常 Batch

Microbatch 和 batch 在数据结构上没有区别（都是 ScheduleBatch），区别只在于 PP 需要多个 batch 同时在不同 stage 中流转，所以给它们编号叫 microbatch。

Microbatch 数量不是越多越好

1

pp_loop_size = pp_size + pp_async_batch_depth

• 太少（= pp_size）：pipeline 容易断流，一个 stage 稍慢就导致下游空闲
• 太多：显存占用线性增长（每个 microbatch 都要有自己的 KV cache + activation），而且可能因为显存压力导致 batch size 被迫缩小，反而吞吐下降

调优经验：

• pp_async_batch_depth = 1~2 通常足够，再深收益递减
• 显存充裕时可以试更大的 depth，但要注意 activation memory 的开销
• 如果 single batch 的 token 数很大（长上下文），microbatch 多了容易 OOM

PP Rank 的请求接收路径

recv_requests() 根据 PP rank 的不同，有两条完全不同的接收路径。接收的请求包括两类：

• 外部 API 请求：tokenize 后的推理请求（TokenizedGenerateReqInput、TokenizedEmbeddingReqInput 等）
• 内部控制请求：PP 同步用的控制消息（如 PPHiCacheEventsReq 等）

PP Rank 0（第一个 stage）

从上游组件直接接收用户请求，走两条 zmq 通道：

1
2
3
4
5

# 1. 从 tokenizer 接收（zmq 通道）
recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)

# 2. 从 rpc 接收（zmq 通道）
recv_rpc = self.recv_from_rpc.recv_pyobj(zmq.NOBLOCK)

接收的是 tokenize 后的请求对象，类型为：

• TokenizedGenerateReqInput - 生成请求
• TokenizedEmbeddingReqInput - embedding 请求
• 其他控制类消息（flush、abort 等）

PP Rank > 0（后续 stage）

不直接从 tokenizer 接收，而是从前一个 PP rank P2P 转发过来：

1
2
3
4
5
6
7

recv_reqs = point_to_point_pyobj(
    [],
    self.pp_rank * self.tp_size + dp_offset,  # 本 rank
    self.world_group.cpu_group,
    (self.pp_rank - 1) * self.tp_size + dp_offset,  # src: 前一个 PP rank
    self.pp_rank * self.tp_size + dp_offset,  # dst: 本 rank
)

对应 event_loop_pp 中的流转

1
2
3
4
5
6
7
8
9

# event_loop_pp 第 80-88 行
recv_reqs = self.recv_requests()  # PP0 从 zmq 收，PP1 从 PP0 的 P2P 收
self.process_input_requests(recv_reqs)  # 本地处理

if not self.pp_group.is_last_rank:
    # 步骤 2: send reqs to next PP stage
    self.send_req_work = self._pp_send_pyobj_to_next_stage(  # 转发给下一个 PP rank
        recv_reqs, async_send=True,
    )