Speculative decoding 的 drafter 架构正在经历一次范式转移。DFlash 用 block diffusion 把 drafting 从串行变并行，实现了 6× 加速；DSpark 在此基础上补了两刀——半自回归解决并行生成的后缀衰减，置信度调度解决高并发下的验证浪费。本文围绕这两篇论文，结合源码逐行分析，澄清训练注意力结构中的常见困惑，并讨论其架构设计、核心 trade-off 和工程落地。

一、背景：从串行到并行的 Drafter

Speculative decoding 的加速比为 $\eta = L_{\text{target}} / L$，其中每个 cycle 的 per-token 延迟为 $L = (T_{\text{draft}} + T_{\text{verify}}) / \tau$。$\tau$ 是每个 cycle 期望接受的 token 数。

Autoregressive drafter 和 diffusion drafter 的最大区别在于 drafting 的计算方式。自回归 drafter 一个 token 一个 token 地算，$T_{\text{draft}} = \gamma \cdot t_{\text{step}}$ 与 block size 线性增长。为了控制延迟，只能用极浅架构（Eagle3 仅 1 层 transformer），$\tau$ 很快饱和，加速比卡在 $\sim 2{-}3\times$。Diffusion drafter 一次并行算出整个 block 的 token，$T_{\text{draft}}$ 对 $\gamma$ 基本不敏感，因此可以用更深的网络获得更高的 $\tau$。

DFlash 就是这样一个并行 diffusion drafter。它的关键设计是 KV injection：从目标模型提取 hidden context features，注入到 draft 模型每一层的 Key-Value cache 中，让 draft 模型利用目标模型的深度表征来做条件预测，而不是从头猜。但纯并行生成引入了新问题：block 内 token 之间没有依赖建模。

二、DFlash：用 Diffusion 做 Drafter

2.1 核心思路

DFlash 的核心 insight 很简单：目标模型知道未来。

大型自回归模型的 hidden states 隐含了多个未来 token 的信息。DFlash 不让小模型从头猜，而是把目标模型的 hidden features 作为条件，让 draft 模型变成一个"扩散适配器"——利用目标模型的深度表征来并行预测未来 block。

2.2 “Diffusion” 到底在哪？

DFlash 名字里有个 D，但翻遍代码你会发现一个事实：没有多步去噪，没有噪声调度，没有连续时间 SDE。所谓的 diffusion 只体现在两件事上：

1. Mask token 构造：待预测位置填充为 mask token，类似于 BERT 的 [MASK]，作为"全噪"起点
2. 双向注意力（is_causal=False）：block 内 token 互相可见，一次 forward 出所有位置

就这两点，没有迭代去噪。标准 block diffusion（Arriola et al., 2025）还有多步迭代，DFlash 把它压成了单步。传承链条是这样的：

1
2
3

连续扩散 (LLaDA)  ->  Block 级离散扩散 (BlockDiff)  ->  单步 mask-predict (DFlash)
    高斯噪声              多步迭代去噪               一步出结果
    连续时间 SDE          离散 mask token            BERT-style

每一步都在往"更像自回归、更不像 diffusion"的方向走。DFlash 到了极致——名字叫 diffusion，实质是 parallel mask prediction。论文用 “diffusion” 这个词主要是学术传承定位，不是方法描述。

那 mask token 的作用是什么？模型需要知道哪些位置是待预测的，哪些是已知信息。如果不用 mask token，直接放随机 embedding 进去，模型会把这些位置当成已知输入去做 attention。Mask token 是一个学习到的"我不知道"信号，跟 BERT 的 [MASK] 一回事。

2.3 KV Injection：不是输入融合，是每层注入

Eagle3 也用目标模型的 hidden features，但只在输入层融合，随着 draft 模型变深，目标信息逐渐稀释。DFlash 采用了完全不同的策略：

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

目标模型 hidden states (5层)
    │
    │  concat + 线性投影 W_c (5*hidden -> hidden)
    ▼
H_ctx = RMSNorm(W_c [H(l1); ...; H(l5)])  ← 压缩后的上下文特征
    │
    │  注入到 draft 模型每一层的 KV cache
    ▼
Draft Layer 1:  K = [W^K · H_ctx; W^K · H_d]  ← 目标特征 + draft 特征拼接
                 V = [W^V · H_ctx; W^V · H_d]
    │
Draft Layer 2:  同上（H_ctx 共享）
    │
    ...
    │
Draft Layer 5:  同上

关键区别：目标特征作为额外的 KV entry 直接注入每一层，而不是经过 draft 模型的 Q projection、output projection 和 FFN。这意味着目标信息在每一层都是"常驻"的，不会因深度而稀释。

从源码看（dflash/model.py），KV 注入的实现极其直接：

1
2
3
4
5
6

# 每个 draft decoder 层的 attention 中
k_ctx = self.k_proj(target_hidden)    # KV 注入：target context
k_noise = self.k_proj(hidden_states)  # draft 自身的 K/V
k = torch.cat([k_ctx, k_noise], dim=1)  # 拼接 K
v = torch.cat([v_ctx, v_noise], dim=1)  # 拼接 V
# is_causal = False ← 双向注意力，block 内 token 互相可见

所有层共享同一份 target_hidden（经 fc + RMSNorm 投影后），且 attention 设为 is_causal=False——block 内 token 双向可见，这是并行扩散生成的必要条件。

一句话总结：KV injection 把目标模型的 hidden features 变成 draft 模型每一层的"持久上下文"，让深层 draft 模型也能充分利用目标模型的信息。

2.4 共享 Embedding 和 LM Head：设计意图

在深入推理和训练流程之前，需要先澄清一个贯穿两篇论文的基础设计：共享 embedding 和 LM head。

DFlash 论文明确阐述了这一设计的动机：

“To improve training efficiency, the draft model shares the token embedding layer and language modeling head with the target model and keeps them frozen during training. Only the draft Transformer layers are updated. This design reduces the number of trainable parameters and encourages the draft model to function as a lightweight diffusion adapter tightly aligned with the target model’s representation space.”

这里的"target model"就是你想加速的那个大模型——最终产出正确 token 的 LLM（如 Qwen3-4B、DeepSeek-V4）。target.lm_head 不是什么特殊构造，它就是目标模型自带的最后一个线性层——把 hidden state 映射到词汇表 logits 的那一层。

以 Qwen3-4B 为例：

• embed_tokens：vocab_size(151936) × hidden_size(2560) ≈ 390M 参数
• lm_head：hidden_size(2560) × vocab_size(151936) ≈ 390M 参数
• 两者合计占 Qwen3-4B 总参数(4B)的约 20%

关键词是 “lightweight diffusion adapter”——共享 + 冻结 embed/lm_head 的本质目的不是省参数，而是强制 draft 模型在目标模型的表征空间内工作。embed_tokens 决定输入空间，lm_head 决定输出空间，两者都锁定后，draft 模型只能学习"如何把目标模型的 hidden states 转换成未来 token 的预测"，而不能自己学一套独立的表征。这正是 KV injection 设计的配套——KV injection 让目标模型的信息每层注入，共享 embed/lm_head 让 draft 的输入输出空间与目标对齐，两者合在一起确保 draft 是一个纯粹的"适配器"而非独立模型。

两份源码的共享方式不同：

• DFlash（推理时直接借用）：DFlashDraftModel 类本身不持有 embed_tokens 和 lm_head 模块。在 dflash_generate() 函数中直接调用 target 对象的属性：

  1 2 3

  # model.py 第111-112行 - 直接调用，不存副本 noise_embedding = target.model.embed_tokens(block_output_ids) draft_logits = target.lm_head(model(...))

• DSpark（训练时复制 + 冻结）：draft 模型有自己的 embed_tokens 和 lm_head 模块（modeling.py 第 227-246 行定义），初始化时把 target 的权重逐字节复制过来然后冻结：

  1 2 3 4 5 6

  def initialize_embeddings_and_head(self, *, embed_tokens, lm_head, freeze=True): with torch.no_grad(): self.embed_tokens.weight.copy_(embed_tokens.weight.detach()) self.lm_head.weight.copy_(lm_head.weight.detach()) if freeze: self.set_embedding_head_trainable(False) # requires_grad=False

  训练时必须用独立模块供 PyTorch autograd 走完整前向传播；推理时则像 DFlash 一样直接调用 target 的 lm_head。

2.5 推理流程：极简实现

DFlash 的仓库极其精简（4 个 Python 文件，核心逻辑 ~370 行）。推理主循环 dflash_generate() 的核心步骤：

1
2
3
4
5
6
7
8
9

while not done:
  ① 构造 [prev_token, mask, mask, ..., mask] block
  ② Draft 前向：单次并行生成整个 block 的 logits
     - noise_embedding = target.model.embed_tokens(block_output_ids)  # 直接用 target 的 embed
     - draft_logits = target.lm_head(draft_model(...))                # 直接用 target 的 lm_head
  ③ 采样 draft tokens
  ④ 目标模型单次 forward 验证整个 block
  ⑤ 计算 accept_length（cumprod 找到第一个 reject 的位置）
  ⑥ 裁剪 draft 和 target 的 KV cache 到接受位置

注意：draft 模型在推理时直接使用目标模型的 embed_tokens 和 lm_head（通过传入的 target 对象直接访问），自己只持有 5 个 decoder 层。block_size=1 时退化为普通自回归解码（用于 baseline 对比）。

重要说明：DFlash 仓库只包含推理代码，训练 recipe 尚未开源。DSpark 在 DFlash 架构基础上增加了独立的训练 pipeline，完整实现在 DeepSpec 仓库中。两者不是共享同一套训练框架——DSpark 的训练代码是独立开发的，包含了 Markov head、confidence head、anchor sampling 等 DSpark 特有组件。

2.6 并行扩散 drafting

DFlash 用 block diffusion 一次生成 $\gamma$ 个 token：

$$T_{\text{draft}}^{\text{DFlash}} = t_{\text{parallel}} \quad (\text{与 } \gamma \text{ 无关})$$

这意味着 draft 模型可以用更深的架构（5 层 vs Eagle3 的 1 层），而不会让 drafting 延迟失控。实验显示，5 层 DFlash 生成 16 个 token 的延迟，低于 1 层 Eagle3 生成 8 个 token 的延迟。

需要注意的是，历史 context 的 KV cache 仍然是 causal 的。目标模型的 forward pass 是标准 causal attention，产出的 hidden states 已经编码了"只能看前面"的因果历史。DFlash 通过 KV injection 把这些 hidden states 注入到 draft 模型，所以 draft block 整体的 attention pattern 是：

1
2
3
4
5

[历史 context（causal，来自目标模型 KV 注入）]  [draft block（bidirectional）]
                                                  ↓
                                         mask tokens 互相可见
                                         但都只能看到历史 context
                                         看不到"未来"

2.7 训练设计

DFlash 的训练 recipe 未开源。以下分析基于 DSpark 论文和 DeepSpec 仓库源码，两者的训练设计在 backbone 层面一致（KV injection、共享 embed/lm_head、anchor sampling 等核心机制相同），DSpark 额外增加了 Markov head 和 confidence head 的训练。

2.7.1 序列布局：拼接式而非交错式

训练时的输入序列布局是 concatenated（拼接式），不是 interleaved（交错式）：

1
2

[ context: p1 p2 p3 r1 r2 r3 r4 r5 ]  [ draft: B0 B1 B2 ... ]
                                      ← 拼接在 context 之后

• Context 部分：完整的训练样本 [prompt | response]，全部是 ground truth token。目标模型对这段序列做一次 forward，提取 5 个中间层的 hidden states 作为 KV injection 来源。
• Draft 部分：512 个 block 拼接而成，每个 block 是 [anchor, mask, mask, ..., mask]，block_size=7。

2.7.2 随机 Anchor 采样

不从 response 均匀分块，而是随机采样 anchor 位置作为每个 block 的起点。源码（common.py 第 164 行）显示采样后会 .sort() 排序：

1

anchors = gathered[:, :max_n].sort(dim=1).values  # 随机采样后排序

排序后 block 按位置从小到大排列，每个 block 的 anchor 是一个真实的 ground truth token（teacher-forced），紧跟 block_size - 1 个 mask token。一次训练前向传播同时覆盖 512 个位置。

2.7.3 Flex Attention：用函数描述稀疏注意力模式

DSpark 的注意力模式高度稀疏——每个 7-token block 只能看到 anchor 之前的 context + 自己 block 的 7 个 token。用稠密矩阵（Q_LEN × KV_LEN 布尔矩阵）太浪费。

PyTorch 2.5+ 的 flex_attention API 提供了解法：用函数描述注意力模式，而不是构造稠密矩阵。

流程：

1. 提供一个 mask_mod(b, h, q_idx, kv_idx) -> bool 函数，告诉它"query 位置 q 能不能看到 key 位置 k"
2. create_block_mask() 把这个函数编译成块稀疏格式——把整个矩阵切成小块（比如 128×128），只保留含 True 的块
3. 实际 attention 计算时跳过全 False 的块，只算有内容的块

源码核心（common.py 第 86-96 行）：

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

def dspark_mask_mod(b, h, q_idx, kv_idx):
    q_block_id = q_idx // block_size
    anchor_pos = anchor_positions[b, q_block_id]
    
    is_context = kv_idx < seq_len
    mask_context = is_context & (kv_idx < anchor_pos)   # 只看 anchor 之前的 context
    
    is_draft = kv_idx >= seq_len
    kv_block_id = (kv_idx - seq_len) // block_size
    mask_draft = is_draft & (q_block_id == kv_block_id)  # 只看同一个 block
    
    is_valid_block = block_keep_mask[b, q_block_id]
    return (mask_context | mask_draft) & is_valid_block

2.7.4 Attention Mask 的两条规则

整个注意力可见性只有两条规则：

1. Block 之间互不可见（q_block_id == kv_block_id）——不同 block 的 draft token 完全隔离，双向的
2. Anchor 之前的前缀可见（kv_idx < anchor_pos）——context 中 anchor 位置之前的 token 可见

这两条规则产生了注意力矩阵中的阶梯（staircase）结构。

2.7.5 Invisible Tokens：到底是什么

论文训练图中的"白色 = invisible tokens"让人困惑。Invisible tokens 分两类：

关键澄清：Context 边界是 token 级别的，不是 block 级别的

这是理解训练图最容易混淆的地方。看注意力矩阵：

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

         Context (causal)                    Draft blocks (bidirectional)
   c0  c1  c2  c3  c4  c5  c6  c7    B0(a m m)  B1(a m m)  B2(a m m)
B0  ✓   ✓   ✗   ✗   ✗   ✗   ✗   ✗      ✓  ✓  ✓    ✗  ✗  ✗    ✗  ✗  ✗
B0  ✓   ✓   ✗   ✗   ✗   ✗   ✗   ✗      ✓  ✓  ✓    ✗  ✗  ✗    ✗  ✗  ✗
B0  ✓   ✓   ✗   ✗   ✗   ✗   ✗   ✗      ✓  ✓  ✓    ✗  ✗  ✗    ✗  ✗  ✗
                                          
B1  ✓   ✓   ✓   ✓   ✗   ✗   ✗   ✗      ✗  ✗  ✗    ✓  ✓  ✓    ✗  ✗  ✗
B1  ✓   ✓   ✓   ✓   ✗   ✗   ✗   ✗      ✗  ✗  ✗    ✓  ✓  ✓    ✗  ✗  ✗
B1  ✓   ✓   ✓   ✓   ✗   ✗   ✗   ✗      ✗  ✗  ✗    ✓  ✓  ✓    ✗  ✗  ✗
                                          
B2  ✓   ✓   ✓   ✓   ✓   ✓   ✗   ✗      ✗  ✗  ✗    ✗  ✗  ✗    ✓  ✓  ✓
B2  ✓   ✓   ✓   ✓   ✓   ✓   ✗   ✗      ✗  ✗  ✗    ✗  ✗  ✗    ✓  ✓  ✓
B2  ✓   ✓   ✓   ✓   ✓   ✓   ✗   ✗      ✗  ✗  ✗    ✗  ✗  ✗    ✓  ✓  ✓
         ↑               ↑               ↑
      anchor=2        anchor=4        anchor=6
      (B0的边界)      (B1的边界)       (B2的边界)

• B0（anchor=2）：只看到 context [c0, c1]——2 个 token
• B1（anchor=4）：看到 context [c0, c1, c2, c3]——4 个 token（包含 B0 看到的 + 更多）
• B2（anchor=6）：看到 context [c0, c1, c2, c3, c4, c5]——6 个 token

蓝色（可见 context）形成一个阶梯。阶梯的每一级台阶在 anchor 位置（2、4、6），是单个 token 的位置。

为什么 context 看起来也按 block 切了？

这是视觉错觉。两个相邻 anchor 之间的 context 段（比如 [c2, c3]）对 B0 不可见、对 B1 和 B2 可见，在图里看起来像一个"块"。但边界是随机 anchor 的 token 位置，不是固定的 block 边界。如果 anchor 随机采到位置 1、4、9，分段就完全不同。

为什么 anchor 之后的不看？

训练时虽然完整序列都在手里，但必须用 mask 模拟推理条件。推理时 draft 模型只能看到 anchor 之前的 token（因为后面的还没生成），所以训练时也必须只让它看 [0, anchor_pos)。这和标准自回归训练的 causal mask 完全同理——你有完整序列，但人为限制可见性防止作弊，只是这里"未来"的定义从"当前位置之后"变成了"anchor 位置之后"。

每个 block 内所有 token 共享同一个 anchor_pos，所以它们看到的 context 前缀完全一样。在注意力矩阵里，这表现为同一 block 的所有行在 context 区域的可见性模式完全一致——画出来就是一个矩形块，视觉上像是 context 也按 block 对齐了。但决定可见/不可见边界的是 anchor_pos 这一个整数，是 token 级别的。

2.7.6 KV Injection 在训练中的结构

训练时的 KV injection 和推理时完全一致——目标模型的 hidden states 经过 fc(5×hidden → hidden) + RMSNorm 投影后，作为额外的 K/V entry 注入到 draft 模型每一层：

1
2
3

每层 attention 的 K/V 拼接：
  K = [k_proj(target_hidden) ; k_proj(draft_hidden)]   ← 两部分拼接
  V = [v_proj(target_hidden) ; v_proj(draft_hidden)]

目标特征绕过 Q projection、output projection、FFN，直接作为 KV entry 进入 attention。所有层共享同一份投影后的 target_hidden。

这里的"KV"不是推理时增量生成的 KV cache，而是指 KV injection 的结构——目标模型的 hidden states 作为"常驻 KV"注入每一层。

2.7.7 指数衰减位置加权

$$w_k = \exp\left(-\frac{k-1}{\gamma}\right)$$

Speculative decoding 中，早期 token 的错误会级联失效整个 block 后缀。loss 加权反映了这种不对称性——前面的 token 更重要。

2.7.8 训练 vs 推理

2.8 结果与局限

DFlash 在 Qwen3-8B 上实现 6.1× 加速，比 Eagle3 快 2.5×。但存在两个结构性局限：

1. 后缀衰减：纯并行生成无法建模 block 内依赖。当上下文有多个合理续写（如 “of course” vs “no problem”）时，各位置独立预测可能产生不一致的组合（“of problem”）
2. 验证浪费：所有 draft token 都送去验证，高并发场景下低置信度的后缀 token 占用 batch 容量

DSpark 正是来解决这两个问题的。

三、DSpark：半自回归 + 置信度调度

3.1 整体架构

DSpark = DFlash backbone + 轻量串行 head + 置信度 head + 硬件感知调度器

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

目标模型生成 bonus token (anchor)
    │
    ▼
┌─────────────────────────────────────┐
│  并行 backbone (DFlash)              │
│  输入: anchor + (γ-1) mask tokens    │
│  输出: hidden h_1..h_γ, base logits  │
└──────────┬──────────────────────────┘
           │
    ┌──────┴──────┐
    │             │
    ▼             ▼
┌────────┐  ┌──────────────┐
│串行 head│  │置信度 head   │
│B_k(·)  │  │c_k = σ(w·h)  │
└───┬────┘  └──────┬───────┘
    │              │
    ▼              ▼
 采样 x_k      prefix survival
 (条件于       概率估计
  x_<k)        
    │              │
    ▼              ▼
  draft tokens  ┌──────────────┐
  E F G H       │硬件感知调度器 │
                │ 截断低置信后缀 │
                └──────┬───────┘
                       │
                       ▼
                 目标模型验证
                 E F G (H 被砍掉)

3.2 半自回归生成：解决后缀衰减

问题本质

并行 drafter 在每个位置独立预测，相当于对前缀所有可能的 token 做 marginal 平均。当上下文存在多个合理续写路径时，不同位置可能选到不同路径的 token，产生不连贯的组合。

以论文中的例子：上下文允许 “of course” 和 “no problem” 两种续写。并行 drafter 在位置 1 独立采样得到 “of”，位置 2 仍然不知道位置 1 选了什么，可能选 “problem” 而非 “course”。这就是多模态碰撞（multi-modal collision）。

解法：并行 backbone + 串行 head

DSpark 把生成拆成两个阶段：

并行阶段：DFlash backbone 一次 forward 产生所有位置的 hidden states $h_1, \ldots, h_\gamma$ 和 base logits $U_1, \ldots, U_\gamma$。

串行阶段：在 base logits 上叠加一个 transition bias $B_k$，逐 token 左到右采样：

$$p_k(v | x_0, x_{<k}) = \frac{\exp(U_k(v) + B_k(x_0, x_{<k}, v))}{\sum_{u} \exp(U_k(u) + B_k(x_0, x_{<k}, u))}$$

关键在于 $B_k$ 条件于前面已采样的 token，解决了独立预测的问题。一旦位置 1 采样了 “of”，串行 head 在位置 2 boost “course” 并 suppress “problem”。

三种串行 Head 实现

源码在 markov_head.py 中实现了三种变体，复杂度递增：

Markov head（默认）：$B_k$ 只依赖前一个 token，低秩分解 $B = W_1 W_2$，$W_1 \in \mathbb{R}^{V \times r}$，$W_2 \in \mathbb{R}^{r \times V}$，$r=256$。

1
2
3
4

给定前一个 token x_{k-1}:
    B(x_{k-1}, ·) = W_1[x_{k-1}] · W_2    ← 查表 + 矩阵乘，O(r) 复杂度
    p_k(·) = softmax(U_k + B(x_{k-1}, ·))
    采样 x_k ~ p_k

RNN head 比 Markov head 略好但实现更复杂，收益有限（论文 Figure 4 显示差距很小），生产默认用 Markov head。

训练 vs 推理：一个关键区别

训练时是 teacher-forced，可以并行计算。用 ground truth token ids 作为前缀输入，所有位置的 bias 一次算完，不需要串行循环：

1
2
3
4
5

训练（teacher-forced，并行）:
  位置1: bias = W_1[anchor] · W_2        ← 用 ground truth 的 anchor
  位置2: bias = W_1[gt_token_1] · W_2    ← 用 ground truth 的 token_1
  位置3: bias = W_1[gt_token_2] · W_2    ← 用 ground truth 的 token_2
  所有位置一次 forward 算完

推理时必须逐 token 串行，因为位置 k 的 bias 依赖位置 k-1 实际采样出来的 token，不是 ground truth：

1
2
3
4
5
6

推理（autoregressive，串行）:
  位置1: bias = W_1[anchor] · W_2
         采样 x_1 ~ softmax(U_1 + bias)   ← 这步必须先完成
  位置2: bias = W_1[x_1] · W_2           ← 用上一步采样的 x_1
         采样 x_2 ~ softmax(U_2 + bias)
  ...

因为每步只是一个 embedding 查表 + 低秩矩阵乘，非常轻量，在 batch size 128 的生产环境下延迟开销只有 ~1%。这也是 DSpark 叫"半自回归"的原因——backbone 是并行的，串行 head 是自回归的。

一句话总结：半自回归 = 并行 backbone 出 base logits + 轻量串行 head 加条件 bias，用 1% 的延迟代价换来 16-18% 的接受长度提升。

3.3 位置分析：为什么并行能赢自回归？

DSpark 论文里最精彩的实验是位置层面的条件接受率分析（Figure 2）。这个分析揭示了一个反直觉的结论。

反直觉：并行 drafter 在位置 1 显著优于自回归 drafter

原因纯粹是架构容量：自回归 drafter 受限于 $T_{\text{draft}} \propto \gamma$，只能用 1 层 transformer；并行 drafter 延迟与 $\gamma$ 无关，可以用 5 层甚至 8 层。更深的网络在位置 1（只依赖目标模型上下文，不需要 block 内依赖）自然更准。

论文 Figure 2 按不同 domain（Math、Code、Chat）分别给出了位置接受率曲线。大致趋势是：

1
2
3
4

位置:    1     2     3     4     5     6     7
DFlash:  高位   ↘     ↘     ↘     ↘     ↘     ↘   ← 缓慢但持续衰减
Eagle3:  低位   ↗     ↗     ↗     ↗     ↗     ↗   ← 反而上升（利用已采样的前缀）
DSpark:  最高   ↘     ↘     ↘     ↘     ↘     ↘   ← 高位稳定衰减

Eagle3 在后续位置反超 DFlash，因为自回归模型条件于已采样 token，随着前缀确定，后续预测变得更容易。DFlash 在后续位置仍然是独立预测，持续衰减。

DSpark 两头兼得：位置 1 继承并行 backbone 的高容量优势，后续位置由串行 head 维持稳定。

注意：以上趋势基于论文 Figure 2 的定性描述。具体数值请参阅论文原图，不同 domain（Math/Code/Chat）的绝对值差异较大。

为什么位置 1 最重要？

Speculative decoding 是严格的 prefix-matching：位置 1 被 reject，整个 block 作废。位置 1 的接受率对整体 $\tau$ 的杠杆最大。

这个分析也解释了为什么 DFlash 能比 Eagle3 快 2.5×——不是因为并行生成更快（虽然确实更快），而是因为更深的网络在位置 1 的优势被 prefix-matching 机制放大了。

3.4 置信度调度验证：从固定长度到自适应

问题：高并发下的验证浪费

DFlash 和 Eagle3 都用固定长度验证：draft 模型生成 $\gamma$ 个 token，全部送去目标模型验证。但在高并发场景下：

• 每个 extra verification token 都占用目标模型的 batch 容量
• 低置信度的后缀 token 大概率被 reject，验证它们是纯浪费
• 被浪费的 batch 容量本可以服务其他请求

解法：Confidence Head + Hardware-Aware Scheduler

Confidence Head 的源码实现极其极简——就是一个单层线性投影：

1
2
3
4
5
6

# eval/dspark/confidence_head.py
class AcceptRatePredictor(nn.Module):
    def __init__(self, input_dim: int):
        self.proj = nn.Linear(input_dim, 1)  # 单层线性投影
    def forward(self, features):
        return self.proj(features).squeeze(-1)

输入特征是 [hidden_states, markov_prev_embeddings] 拼接，输出经过 sigmoid 后得到每个位置的条件生存概率：

$$c_k = \sigma(w^\top [h_k; W_1[x_{k-1}]])$$

监督信号是解析的 per-step 接受率：$c_k^* = 1 - \frac{1}{2}\|p_k^d - p_k^t\|_1$（TV distance 的补）。训练时用 BCE loss。

推理时的置信度裁剪同样简洁：

1
2
3
4

# draft_ops.py - 找到第一个低于阈值的置信度位置，截断
below_threshold = confidence_logits.sigmoid() < threshold
first_below = torch.nonzero(below_threshold[0])[0].item()
return first_below  # 只验证 [0, first_below) 的 token

Sequential Temperature Scaling (STS) 校准：原始 confidence 通常过自信（ECE 3-8%）。STS 逐位置做 1D grid search，最小化累积乘积 $\prod_{i \leq k} c_i$ 的 ECE，校准后 ECE 降到 ~1%。

Hardware-Aware Prefix Scheduler 把验证长度选择形式化为全局吞吐量最大化问题：

$$\Theta = \tau \cdot \text{SPS}(B), \quad \text{其中 } \tau = \sum_{r=1}^{R}\left(1 + \sum_{j=1}^{\ell_r} a_{r,j}\right), \quad B = \sum_{r=1}^{R}(1 + \ell_r)$$

• $\text{SPS}(B)$：引擎的 steps-per-second 容量曲线，初始化时 profiling 一次
• $a_{r,j} = \prod_{i \leq j} c_{r,i}$：request $r$ 在位置 $j$ 的 prefix survival 概率
• 目标：选择每个 request 的验证长度 $\ell_1, \ldots, \ell_R$，最大化 $\Theta$

因为 $a_{r,j}$ 单调递减，可以贪心求解：全局排序所有 $(r, j)$ 按 $a_{r,j}$ 降序，逐个加入验证 batch，直到 $\Theta$ 不再上升。

1
2

负载低 → SPS(B) 几乎不变 → 多验证 token 划算 → 验证长度大
负载高 → SPS(B) 快速下降 → 少验证 token 划算 → 砍掉低置信度后缀

一句话总结：置信度调度把"验证多少"从静态配置变成动态优化问题——根据每个请求的 draft 质量和当前系统负载，全局分配验证算力。

3.5 训练目标

DSpark 的 loss 三项加权和：

$$\mathcal{L} = \alpha_{\text{ce}} \mathcal{L}_{\text{ce}} + \alpha_{\text{tv}} \mathcal{L}_{\text{tv}} + \alpha_{\text{conf}} \mathcal{L}_{\text{conf}}$$

$\mathcal{L}_{\text{tv}}$ 权重最高，因为 TV distance 直接对应接受率：per-step 接受概率 $= 1 - \frac{1}{2}\|p^d - p^t\|_1$，最小化 TV distance 就是最大化期望接受率。

四、工程落地：从论文到 DeepSeek-V4 线上

4.1 生产部署架构

DSpark 部署在 DeepSeek-V4-Flash 和 V4-Pro 上。

注意：生产环境使用 MoE + mHC 架构和 block_size=5；开源 checkpoint 使用标准 dense 层和 block_size=7，便于社区复现。两者核心算法一致，架构配置不同。

4.2 异步调度：解决 ZOS 冲突

算法 1 的同步版本与生产系统的 Zero-Overhead Scheduling (ZOS) 冲突——ZOS 需要在当前 step 完成前知道下一步的 batch size。DSpark 的解法是用两步前的 confidence 预测来确定当前步的截断长度：

1
2
3
4

Step N-2:  生成 draft + confidence
Step N-1:  用 N-2 的 confidence 确定截断 → 验证
Step N:    用 N-1 的 confidence 确定截断 → 验证
            ↑ 同时生成新 draft + confidence（供 N+1 使用）

这引入了轻微的时间偏差，但选择机制是 rank-preserving 的——最自信的 draft token 总是优先验证。更重要的是，异步设计形成了一道"因果屏障"：截断决策只依赖历史信息，不会泄露未来 token，保证了 lossless guarantee。

4.3 生产性能

关键结论不是倍数本身，而是DSpark 扩展了可行的交互性边界。在 120 TPS/user 的严格 SLA 下，MTP-1 baseline 几乎无法运作，DSpark 仍然稳定——这意味着原来达不到的延迟等级现在可以服务了。

五、两篇论文的对照

DSpark 不是对 DFlash 的替代，而是增量改进。DFlash 解决了"能不能用 diffusion 做 drafter"的问题，DSpark 解决了"用得好不好"的问题。两者共享 KV injection 条件化、共享 embedding/LM head、位置加权 loss 等核心设计。

六、可迁移的启示

1. "目标模型知道未来"是一个深刻的观察。 大模型的 hidden states 隐含了远超 next-token 的信息。DFlash 的 KV injection 和 DSpark 的串行 head 都在利用这一点——draft 模型不需要从头推理，只需要"解读"目标模型已经知道的东西。

2. 并行 vs 自回归不是二选一。 DSpark 的半自回归架构证明，用并行 backbone 做"重活"+ 串行 head 做"精修"，可以在 1% 延迟代价下获得 16-18% 的质量提升。这个思路在更广泛的 LLM 加速领域也适用——不要追求纯并行或纯串行，找正确的分割点。

3. Speculative decoding 是系统问题，不只是算法问题。 DSpark 的置信度调度把验证长度从算法参数变成系统调度参数，在真实流量下实现了负载感知的自适应。这提醒我们：脱离部署环境谈 drafter 架构是不完整的。

4. 位置 1 的杠杆最大。 在 prefix-matching 机制下，位置 1 的接受率对整体 $\tau$ 的影响远大于后续位置。这意味着 draft 模型的架构选择应该优先考虑位置 1 的容量，而非后续位置的依赖建模——这正好是并行 drafter 的天然优势。

5. 共享 embed/lm_head 不是省参数，是锁定表征空间。 冻结 embed 和 lm_head 后，draft 模型被强制在目标模型的表征空间内工作，成为纯粹的"适配器"。这是 KV injection 的配套设计——前者保证空间对齐，后者保证信息每层注入。

参考：

• DFlash: Chen et al., “DFlash: Block Diffusion for Flash Speculative Decoding”, ICML 2026. arXiv:2602.06036
• DSpark: Cheng et al., “DSpark: Confidence-Scheduled Speculative Decoding with Semi-Autoregressive Generation”, 2026. arXiv:2607.05147
• Eagle3: Li et al., “Eagle-3: Scaling up Inference Acceleration of LLMs via Training-Time Test”, 2025. arXiv:2503.01840

开源仓库：

• DeepSpec（deepseek-ai/DeepSpec）：69 个 Python 文件，包含 Eagle3、DFlash backbone、DSpark 三种 drafter 的统一训练框架，支持 Qwen3 和 Gemma4 系列目标模型。完整训练 pipeline（数据准备 → 训练 → 评估）。
• DFlash（z-lab/dflash）：4 个 Python 文件，仅含推理代码（Transformers/SGLang/vLLM/MLX 四种后端），训练 recipe 尚未开源。

MiMo-V2.5 推理优化解读：Hybrid SWA 的工程落地

小米 MiMo 团队的《Full-Pipeline Inference Optimization for MiMo-V2.5 Series》（arXiv:2607.13095）核心论点很直白：Hybrid SWA 理论上能把 KVCache 和 attention 计算量压到 Full Attention 的 1/7，但如果 KVCache 系统不跟着改造，实际反而会变慢。

本文不写综述，重点拆解几个关键工程问题：SWA 的 KV cache 怎么存、怎么 offload、共享前缀命中时的重算问题、numa_balancing 为什么 +10%、以及 Prefill 为什么也要开 MTP。

1. 论文速览

MiMo-V2.5-Pro 的架构组合拳：

1
2

Hybrid SWA + 稀疏 MoE + 多模态编码器
70 层 = 10 Full Attention + 60 SWA（窗口 W=128）

理论上 attention FLOPs 和 KVCache 存储都降到 1/7，但论文的核心观点是：这些理论红利不会自动兑现。Hybrid SWA 在 KVCache 管理、前缀匹配、Full 和 SWA 层语义一致性上都引入了新问题。

论文涉及的优化面很广：KVCache 分层（L1 GPU / L2 Host / L3 GCache）、调度（LLM-Router）、Prefill/Decode pipeline、多模态等。其中 LLM-Router 和 GCache 属于常见的分布式缓存和调度实现，本文不展开。重点聚焦在 SWA 相关 KVCache 这条主线上。

2. SWA 的 KV Cache 怎么存

核心设计：物理分离 + 逻辑统一。

物理层：两个独立池子

每层存储（L1 GPU / L2 Host / L3 GCache）都并排维护两个 KV Pool：

1
2
3
4
5
6
7
8
9

┌─────────────────────────────────────┐
│  Full KV Pool                       │
│  ├─ 大小: O(seq_len × 10 层)        │
│  └─ 淘汰: 按整个序列做 LRU           │
├─────────────────────────────────────┤
│  SWA KV Pool                        │
│  ├─ 大小: O(W × 60 层)  ← 严格 O(W) │
│  └─ 淘汰: 独立按 window 做 eviction │
└─────────────────────────────────────┘

关键点：SWA pool 物理上就只有 W 大小，写入新 token 时按窗口 evict 旧的。不是「存完整历史然后只用尾部 128 个」，而是存进来就已经是窗口截断后的样子。

直观对比（128K token 序列）：

Hybrid SWA 存储量几乎等于只算 Full 层的量，SWA 层几乎不占空间。

逻辑层：单序列视图 + 双索引前缀树

上层（前缀树、调度器）只看到一个逻辑序列，底下用 Full→SWA mapping 做透明分层。每个前缀树节点存两套元数据：

1
2
3
4
5

PrefixTreeNode {
  tokens: [t0, t1, ..., tN]
  full_seg_idx:  [Full KV pool 位置]      # Full 层复用
  swa_seg_map:   [SWA KV pool 位置 or 空] # 判断 window 安全性
}

用途：

1. 命中判定：token 相等 + tail W 个 token 的 swa_seg_map 都非空 → 才算真命中
2. 独立淘汰：window 外的 SWA 段可以先扔，Full 段还留着

一句话总结：SWA 存的就是「每 SWA 层永远只存 W 个 token 的 KV」，SWA pool 和 Full pool 物理分离，上层用双索引前缀树抽象成单序列视图。

3. SWA Cache 怎么 Offload 到 CPU 内存

存什么：只存「窗口内有效 slot」

对一条 100K token 的序列：

1
2

Full 层 (10 层):  每层 100K 个 slot  ← 全部下沉到 CPU
SWA 层 (60 层):  每层只有 W=128 个 slot ← 只存窗口内

SWA 下沉到 CPU 的就是窗口截断后的形态，不是完整历史。

怎么存：Host 侧独立的 SWA Pool

L2 结构完全镜像 L1：

1
2
3

Host DRAM (L2):
  Full KV Pool: [10 层, seq_len, head_dim], paged + pinned memory
  SWA KV Pool:  [60 层, W=128, head_dim], 独立 eviction

工程要点：

• Pin memory：Host 池子用固定内存，D2H/H2D 用异步 DMA
• Paged 分配：SWA pool 拆成 block，滑出窗口的 block 归还池子
• 和 Full pool 完全独立：分开的地址空间、eviction 策略

怎么搬：按 SWA mask 传输

论文 §3.1.1 的关键一句：

“Cross-tier transfers are performed based solely on the SWA mask, ensuring only valid window data is moved.”

D2H / H2D 时不搬窗口外的数据：

1
2
3

D2H writeback: 只把 mask=1 的 slot 打包 DMA 到 host
H2D prefetch:  只拉窗口内的 W 个 token 的 KV
              → 数据量小，layerwise prefetch 能 overlap 计算

这里省的是带宽。按 mask 走，60 层每层就 128 个 token，总量小到可以忽略。

D2H 时机

• Eviction 前的备份：GPU SWA pool 满了，先 D2H 到 host 再释放
• 一致性修复（§3.1.4）：前缀树节点合并 / prefill chunk 完成时，检查 device 和 host 的 SWA 占用差 → host 补齐槽位 → 异步 D2H
• 请求结束 / 会话切换：高价值会话主动 D2H 到 host，甚至下沉到 L3

一句话总结：SWA 下沉到 CPU 的就是窗口截断后的 128 个 slot，Host 侧用独立 pool + pinned memory 存，搬运只按 SWA mask 走有效位置。

4. 共享前缀命中时，SWA 要不要重算

这是 Hybrid SWA 落地生产最难受的一个点。

问题本质

典型场景：共享 system prompt（2000 tokens），下面挂了很多用户会话：

1
2
3
4

PrefixTree:
  system_prompt (2000 tokens) ← 非叶子节点
    ├─ user_A_session (往下延展了 50K tokens)
    └─ user_B_session (往下延展了 30K tokens)

关键观察：user_A 会话延展了 50K tokens，当前窗口在 [52000-128, 52000] 附近。system_prompt 段（位置 0-2000）早就滑出窗口了，SWA KV 被 evict 掉了。

这就是「伪命中」

user_D 带着同一个 system prompt 进来：

1
2
3
4
5

命中判定:
  1. Token equality 检查 → 匹配到 2000 tokens ✅
  2. Full KV 检查 → 都在 → Full 命中 2000 tokens ✅
  3. Window-safe 检查 → tail 128 个 SWA slot 在不在？
     └─► ❌ 已被 evict

按「window-safe length」规则，SWA 匹配长度被 clip。Full 层能复用 2000 tokens，但 SWA 层不行——SWA 层要看窗口内 KV，位置 2001 的 attention 需要 [1873, 2001] 的 KV，这段没了 → 只能重算 SWA 层的 prefill。

代价：

• Full 层：省了（Full pool 还留着）
• SWA 层：60 层全部重跑一次 prefill（占 6/7）
• 60 层重算，几乎等于白干

小米的应对：SWA 保留策略（§3.1.4 第 4 条）

论文直接点名了这个场景：

“Medium/short sequence SWA retention strategy. Based on user request patterns, we retain relatively dense SWA KV Cache at fixed length positions for medium/short sequences… particularly beneficial for long agent sessions, multi-user shared system prompts, and repeated tool calls to the same codebase.”

翻译：对高频复用的前缀（比如 system prompt），SWA pool 不严格 O(W)，而是在关键位置留几份 SWA 快照。

推测的实现方式：

1
2
3
4
5
6
7

普通序列的 SWA pool:
  只留 tail W 个 slot     ← 严格 O(W)

高频前缀（如 system prompt）:
  在几个 anchor 位置多留 W 个 slot
  比如: [0, W], [500-W, 500], [1000-W, 1000], [2000-W, 2000]
  ← 相当于"检查点"，每个都能作为 SWA 计算的起点

这样 user_D 命中时：匹配到 system_prompt 末尾 tail W=128 的 SWA slot 确实存在（是保留下来的 anchor），Full + SWA 都能完整复用 → 不需要重算 SWA 层。

Trade-off 账

核心思路：用一点点 SWA 存储换取显著提升的共享前缀命中率。因为 system prompt 被成千上万用户共享，每次省掉 60 层 SWA prefill 的收益，远远大于多存几份 128-slot 快照的代价。

一句话总结：非叶子共享前缀命中时，默认严格 O(W) 实现会因为 tail 被 evict 导致 SWA 层必须重算。小米的解法是对高频共享前缀主动保留多个 SWA 快照，让 Full + SWA 都能完整复用，真正兑现 Hybrid SWA 在多用户共享场景下的效率红利。

5. numa_balancing 为什么影响 +10%

论文里就一句话：「关掉 numa_balancing，端到端 +10%」。背后是操作系统调度和 GPU 推理的碰撞。

numa_balancing 是干嘛的

现代服务器都是 NUMA 架构，CPU 访问本 NUMA 节点内存 << 访问远端节点内存（延迟差 1.5-2×）。kernel.numa_balancing=1 是 Linux 的「自动 NUMA 均衡」特性：

1
2
3
4
5
6
7

内核周期性做的事:
  1. 扫描进程页表 → 把页面标记为 PROT_NONE
  2. 进程下次访问 → 触发缺页异常
  3. 内核在 fault handler 里统计 NUMA locality
  4. 如果 CPU 和 page 不在同一 node:
     - 迁移页面到 CPU 所在 node
     - 或者迁移线程到 page 所在 node

但在 GPU 推理场景下变成灾难

SGLang 有自己的 --numa-node 配置，会主动把 rank 绑到指定 NUMA node（PIN 死）。冲突来了：

• SGLang 说：这个 rank 就吃 Node 0，别乱跑
• 内核 numa_balancing 说：我来「优化」一下

内核会做几件破坏性的事：

1. 周期性把页面标 PROT_NONE：即使你已经手动 pin 好了，下一次访问必触发 page fault
2. Page fault 引起大 stall：GPU 计算 → 触发 page fault → 内核处理 → GPU 继续（延迟数百 μs 到 ms）
3. 随机页面迁移：pinned memory 被迁移 → GPU 之前记录的 DMA 地址失效

为什么这个 bug 特别难缠

论文里这句话说到点子上了：

“In multi-node multi-GPU deployments, these gaps appear at random positions across ranks, and each inter-rank synchronization is bottlenecked by the slowest rank.”

关键组合拳：

• 随机性：numa_balancing 是周期性异步扫描，不同 rank、不同时刻会随机中招
• 全局同步放大伤害：LLM 推理每层 attention 之后都要 all_reduce，最慢那个 rank 拖累所有人

1
2
3

所有 rank 计算 ──► 集合通信同步 ──► 下一层
              ↑
              最慢那个 rank 拖累所有人

如果 8 个 rank 里有 1 个刚好被 page fault 命中，那 8 个 rank 全都得等它——尾延迟灾难。对一次前向 60+ 层，只要偶尔一层有一个 rank 中招，整个 latency 就翻倍。

为啥关掉直接 +10%

1. SGLang 已经手动 pin 好了 NUMA，调度器安排得明明白白
2. numa_balancing 是「帮倒忙」，它假设进程没做优化，主动帮你搬，反而破坏了原有布局
3. 代价是 hard cost：page fault、TLB shootdown、页面迁移都是不可避免的 CPU 侧开销

关掉后：内核不再主动折腾 → 页表稳定 → 没有意外 page fault → GPU kernel 之间没有随机 gap → 集合通信不再被拖尾 → 直接 +10%。

一句话总结：SGLang 已经手动做了 NUMA pin，numa_balancing 在这基础上做二次干扰，通过 page fault + 页面迁移引入随机 stall，被 GPU 集合通信同步放大成尾延迟灾难，关掉就直接 +10%。

6. Prefill 为什么也要开 MTP

默认直觉是「MTP 是加速 decode 的，跟 prefill 有啥关系」，但正是这个直觉害惨了 agentic 场景。

MTP 是什么

MTP (Multi-Token Prediction) = 多 token 预测。MiMo-V2.5 系列原生支持 3 层 MTP：

1
2
3
4

主模型: 预测第 t+1 个 token
MTP Layer 1: 同时预测第 t+2 个 token
MTP Layer 2: 同时预测第 t+3 个 token  
MTP Layer 3: 同时预测第 t+4 个 token

核心思想：一次 forward 出 4 个 token，然后主模型验证，接受的部分直接输出。理想情况下 decode 速度接近 4×。

关键前提：MTP 层本身有参数、有 KV cache，需要**跟着上下文一起「预热」**才能预测准。

原始问题：Prefill 不跑 MTP 的后果

默认实现：Prefill 阶段只跑主模型，跳过 MTP 层。

问题：MTP 层需要看历史 context 才能预测。如果 prefill 跳过 MTP：

1
2
3
4
5
6
7

Prefill 结束时:
  主模型 KV cache: [完整 prompt 的 KV] ✅
  MTP 层 KV cache: [空 / invalid] ❌

Decode 第 1 个 token:
  主模型: ✅ 
  MTP 层: context 是空的，瞎猜 → 接受率极低 ❌

MTP 需要多久「热身」：论文说是 128 个 token。因为 MTP 层需要逐个 decode 的 token 慢慢累积 KV。

这段时间 MTP 基本白搭：主模型每步生成 1 个 token，MTP 提议的 3 个 token 因为 context 不足几乎全被拒，有效加速 ≈ 1×。

为什么 agentic 场景特别惨

论文这句话点破了关键：

“Since agentic scenarios involve mostly short output sequences, this limitation significantly limited MTP’s effective speedup.”

Agentic 场景的输出特征：

1
2
3

一次工具调用响应: ~20-50 tokens
一次 function call: ~30-80 tokens  
一次思考步骤: ~50-150 tokens

几乎所有输出都在 128 token 之内——这正好是 MTP 需要热身的窗口。结果：

1
2

理论加速比: 3× (3层 MTP)
Agentic 场景实际加速比: ~1× (刚热身完就结束了)

等于 MTP 层白装了。

解法：Prefill 阶段也跑 MTP

小米的改动很直接：prefill 时也让 MTP 层参与前向。

1
2
3
4
5
6
7

Prefill 结束时:
  主模型 KV cache: [完整 prompt 的 KV] ✅
  MTP 层 KV cache: [完整 prompt 的 MTP KV] ✅  ← 新增

Decode 第 1 个 token:
  主模型: ✅
  MTP 层: 已经用整个 prompt 预热过，接受率立即很高 ✅

效果（论文数据）：

• 0-128 token 加速: 2.3× ← 从近 1× 提升到 2.3×，agentic 场景直接受益
• 128-256 token 加速: 1.5×

工程适配

论文说：

“By introducing MTP support during prefill with dedicated adaptations and optimizations for HiCache L2/L3…”

关键点：MTP 层有自己的 KV cache，之前 HiCache 那套 offloading 系统只管主模型的 KV。要 prefill 阶段跑 MTP，就得：

1. MTP KV 也要进 HiCache：L1 (GPU) ↔ L2 (Host) ↔ L3 (GCache)
2. Prefix cache 得覆盖 MTP KV：前缀树节点要标 MTP 层的状态
3. 传输和存储成本：MTP 有 3 层，KVCache 总量变多，L2/L3 得扛住

Trade-off 账

核心逻辑：用 prefill 的少量额外开销，换 decode 从第 1 个 token 就享受 MTP 加速。Agentic 场景短输出多，这笔账非常划算。

具体场景：Agent 调用工具，输入 5K token，输出 60 token

1
2
3

没开 prefill MTP: 100ms (prefill) + 1200ms (decode) = 1300ms
开了 prefill MTP: 104ms (prefill) + 520ms (decode) = 624ms
省 52%

一句话总结：Prefill 开 MTP 是为了给 MTP 层的 KV cache 做「上下文预热」，避免 decode 初期因为 MTP 无历史导致接受率极低。Agentic 场景输出短，正好都落在这个「未热身窗口」里，所以收益特别大——0-128 token 加速直接从 ~1× 提升到 2.3×。

总结

这篇论文的核心贡献不是某个单点优化，而是把 Hybrid SWA + MoE + 多模态这套组合架构从理论拉到生产的完整工程实践。本文聚焦在 SWA 相关 KVCache 这条主线，拆解了几个关键问题：

1. SWA 存储：物理分离 Full/SWA pool，SWA 严格 O(W)，前缀树双索引做透明分层
2. SWA offload：Host 侧独立 SWA pool + pinned memory，搬运只按 mask 走有效位置
3. 共享前缀命中：默认严格 O(W) 会伪命中，小米用「SWA 快照保留」策略换真实命中率
4. numa_balancing：操作系统自动优化 vs 手动 NUMA pin 的冲突，关掉 +10%
5. Prefill MTP：给 MTP 层 KV cache 做上下文预热，agentic 短输出场景 0-128 token 从 ~1× 提升到 2.3×

可迁移的启示：

• 架构红利不会自动兑现，工程系统必须跟着架构一起改造
• Hybrid SWA 的价值在多用户共享场景，但需要专门的缓存策略才能兑现
• 推理系统的 pipeline 各阶段不是孤立的，跨阶段依赖的组件需要协同优化
• **系统级坑（numa_balancing、THP、CPU governor）**对已经手动优化过的高性能场景是纯干扰，该关就关

参考：

• 论文：MiMo Team, Xiaomi. Full-Pipeline Inference Optimization for MiMo-V2.5 Series: Pushing Hybrid SWA Efficiency to the Limit. arXiv:2607.13095, 2026.
• 模型：Xiaomi MiMo Team. MiMo-V2.5 / MiMo-V2.5-Pro. HuggingFace, 2026.

本文基于 MiMo-V2.5 论文、个人笔记《MiMo-V2.5 SWA offload 方法》整理。

DeepSeek-V4 的架构图有一个特点：本身就是一份并行性说明书。画成并行分支的模块，基本都是在告诉系统实现者"这里有 overlap 的空间"。本文梳理 V4 中可 overlap 的算子对，对照论文承诺与 SGLang 实现现状。

一、Overlap 的三个前提

1. 算法独立 — 两条路径无数据依赖（论文画成并行分支）
2. 资源不冲突 — 不争用同一 SM 或同一 buffer
3. 工程可兜住 — CUDA stream / multi-kernel launch 的硬件支持

DeepSeek 的特殊性在于：MoE + MHC（Multi-Head Compressor）双重稀疏结构，天然产生多条独立路径。

二、Attention 分支内的 Overlap

2.1 wqkv_a 与 Compressor 的 Overlap

数据依赖

1
2
3
4
5
6
7

wqkv_a:  输入 = hidden (x)
              ↓
          输出 = q_a + kv_a + k_rope

compressor: 输入 = hidden (x) + past KV (来自 KV cache)
              ↓
          输出 = summary_K_4 + summary_K_128

两者都读 hidden，但 compressor 不依赖 wqkv_a 的输出，可以完全并行。

资源 pattern 互补

一个 compute-bound 倾向，一个 memory-bound 倾向，硬件资源不冲突。

L2 Cache 共享

hidden 的大小：[4096, 7168] BF16 = 56 MB，接近 MI355X 的 L2 cache（32 MB）。

• 串行：wqkv_a 读完 hidden，compressor 再读——大概率已被 q_b 挤出 L2，又一次从 HBM 读
• 并行：两个 kernel 同时从 HBM 拉 hidden，第二次的 cache line 是免费的

2.2 Indexer 的 Overlap：只依赖 q_a，不依赖 q_b

这是 V4 论文里 “MLA Co-Design with Indexer” 的核心设计。

Indexer 两条路径的依赖

1
2
3
4
5
6
7

Indexer K-side:  W_K^I · x  →  RoPE  →  summary_K
                 ↑
              只需要 hidden (x)，不需要 wqkv_a 的输出

Indexer Q-side:  W_Q^I · q_lora  →  Hadamard + RoPE
                 ↑
              只需要 q_lora (q_a 的输出)，不需要 q_b 的输出

q_b 是 critical path 上最重的 GEMM（929 µs），indexer 不等 q_b 完成就可以启动。

为什么 Q-side 用 q_lora 而不是 hidden

V4 论文的算法决策：

• 实验结论：召回率无差异
• 系统收益：每层省一次 [T, 7168] → [T, ...] 的大 GEMM
• 代价：Q-side 多一个 q_lora_ready event 依赖（K-side 完全自由）

代码里精确实现了这个依赖：

1
2
3
4

# K-side：不依赖 q_lora，可以立刻启动
stream_indexer.wait_stream(current_stream)
# Q-side：只等 q_lora_ready，不等 q_b
# （在 indexer 内部用 q_lora_ready 精确控制）

2.3 三者在时间轴上的 Overlap

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

时间轴（prefill 4096 tokens，典型层）:

0 µs         77 µs       83 µs              1012 µs
│ wqkv_a │ q_a │          q_b (929 µs)          │
│─────────┴─────┴───────────────────────────────┤
│          │    │
│          └──► indexer K-side + Q-side         │  ← 不等 q_b
│             (q_lora_ready 触发)
│          │
│          └──► compressor (c4 + c128)          │  ← 完全独立
│             (只等 x)
│
└──► kv_write (等 wqkv_a 输出切片)              │

端到端时间 = max(q_b, indexer, compressor, kv_write) = ~1012 µs（由 q_b 主导）

串行时 = wqkv_a + q_a + q_b + indexer + compressor + kv_write = ~2290 µs

三、MoE 分支内的 Overlap

3.1 Shared Expert 与 Routed Expert

两条 expert 路径完全独立，SGLang 用 alt_stream 实现，是已实现的最好案例。

3.2 MoE Wave 模型：计算与通信 Overlap

Wave 分 chunk 乒乓：dispatch → compute → combine 流水，掩盖全量通信延迟。V4 论文 Fig.5 的时序图明确画出了这一点。

3.3 Combine 通信与 Shared Expert 尾部计算

Combine（NVLink all-to-all）不需要 shared expert 的完整输出，可以和 shared expert 的 down_proj 最后一部分 overlap。SGLang 目前串行等待，未利用。

四、为什么 Multi-Stream Overlap 能赚到时间

CUDA stream 是 GPU 上的 FIFO 工作队列；stream 本身只给调度器自由度，性能要靠"互补的资源占用"赚出来。

机制 1：单 Kernel 资源利用率低（最主要）

q_b 把计算压满但 HBM 闲；swa_scatter 把 HBM 压满但计算闲。不同 stream 让它们同时跑，各用各的硬件资源。

机制 2：小 Kernel Grid 填不满 SM

MI355X 有 256 个 CU，但 trace 里很多 kernel 的 grid 很小（rocprim cumsum 只有 1 个 CU，占用率 0.4%）。多 stream 让调度器把多个小 kernel 同时塞进 CU。

机制 3：隐藏 CPU Launch Overhead

每次 hipLaunchKernel 在 host 端要 ~3-5 µs。Trace 里 ~30000 个 kernel × 4 µs ≈ 120 ms 纯 launch 开销。多 stream 让 host 预先 enqueue，GPU 不饥饿。

机制 4：同源输入的 L2 Cache 共享

wqkv_a 和 compressor 都读 hidden，并行时第二次读直接 cache hit，省 ~20% 输入带宽。

反向判断：什么情况下多 Stream 不赚

五、工程实现：SGLANG_OPT_USE_MULTI_STREAM_OVERLAP

这个关键优化靠一个环境变量控制，不在 --help 里：

1
2
3
4
5

# 开启
SGLANG_OPT_USE_MULTI_STREAM_OVERLAP=1 python -m sglang.launch_server ...

# 验证
python -c "import os; print(os.environ.get('SGLANG_OPT_USE_MULTI_STREAM_OVERLAP'))"

读取位置在 sglang/srt/layers/deepseek_v4.py 的 _forward_prepare_multi_stream 方法入口，通过 os.environ.get() 判断，默认关闭。

实测效果

在 decode 阶段（batch size 中等，61 层 DeepSeek-V4），开启后整体 forward 有 ~3.5% 的端到端提升：

提升主要来自：

• q_b GEMM（929 µs）与 indexer + compressor 并行：decode 时 q_b 仍是瓶颈，但 indexer K-side 和 compressor 可以完全隐藏在其执行期间
• 小 kernel（swa_scatter、cumsum 等）与主体 GEMM 并行：这些 µs 级 kernel 在串行时被 q_b 的 launch gap 放大，overlap 后基本被吸收

注意：3.5% 是 decode 阶段的收益。prefill 阶段因为 q_b 的 GEMM 更大（M=4096），overlap 的相对收益会被稀释，但绝对时间节省更显著（每层 ~1.28 ms，61 层 ~78 ms）。

没开时，整个 _forward_prepare_multi_stream 退化成串行，论文里"sparse attention 模块在 prefill 阶段近乎 free"的 claim 直接失效。

六、对照论文图：承诺 vs 现状

七、总结

DeepSeek-V4 论文在算法层面为 overlap 留了很大空间，尤其是 Fig.3（MHC 结构）和 Fig.5（Wave 时序图）。当前 SGLang 实现只吃到了 Shared Expert / Wave 的部分，仍有明显优化空间。

实测数据验证了 overlap 的价值：

• decode 阶段：开启 SGLANG_OPT_USE_MULTI_STREAM_OVERLAP，整体 forward 提升 ~3.5%
• prefill 阶段：每层节省 ~1.28 ms，61 层共 ~78 ms，sparse attention 模块接近"free"的目标

更一般的启示：算法论文画依赖图时多想一步系统实现，系统实现时多对照论文的并行性承诺。

参考：DeepSeek-V4 论文（arXiv:2606.02405）§3.2 Attention Mechanism, §3.3 Expert Mixture, §3.3.3 Dynamic Expert Routing & Wave Scheduling

背景

训练超长序列 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 跨层递归展开后得到：

$$x_L = (\prod B_i) x_l + \sum (\prod B_j) C_i^T F(...)$$

$Π 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 迭代，操作很简单：

1
2
3
4
5
6
7
8
9

# 输入: 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 的中间激活。最优块大小由公式给出：

$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%。

传统 DeepEP 的 Overlap 能力分析

常见误解：“传统 DeepEP dispatch/combine 是无法和 MoE 计算 overlap 的”

更准确的说法：传统 DeepEP 可以通过 low-latency kernel + 多 stream + two-batch overlap 实现部分重叠，但需要框架手工编排，重叠率有限，对 decode 小 batch 几乎无效。

DeepEP 的两种模式

DeepEP 提供了两套 kernel，overlap 能力完全不同：

你提到的"无法 overlapped"更接近 Normal kernel 的情况。

框架层怎么强行 overlap（传统方案）

SGLang / TRT-LLM 的做法是 micro-batch 切分 + 多 stream：

1
2
3
4

时间线 ──────────────────────────────────►
stream A: [dispatch chunk0]              [combine chunk0]
stream B:                [GEMM chunk0]   [dispatch chunk1] [combine chunk1]
stream A:                              [GEMM chunk1]

• chunk0 在算的时候，chunk1 在通信
• 需要 2 条 CUDA stream + event 同步
• 代价：buffer 翻倍，GEMM tile 变小导致 Tensor Core 利用率下降

这就是 TBO (Two-Batch Overlap)，SGLang 里有 --enable-two-batch-overlap 开关。

为什么传统方案"能 overlap 但不彻底"

四个硬限制：

1. Kernel 边界 = 同步点
   每次 launch 有 5–10μs 开销，小 batch 下被 launch 开销主导

2. SM 资源争抢
   Normal dispatch 想跑满 NVLink 就要用很多 SM，GEMM 也想要全部 SM，互相挤压

3. 依赖链串行

   1	dispatch(x) → GEMM1 → SwiGLU → GEMM2 → combine

   单个 token 的依赖必须串行，overlap 只能靠不同 micro-batch 之间的并行

4. Decode 阶段 batch 小
   再切一半更糟，GEMM 效率暴跌

重叠率实测

对性能标定（profiler）的意义

如果在 perf 数据库里标注 MoE 路径，至少分三档：

1. DeepEP normal：通信和计算串行（吞吐高，延迟差）
2. DeepEP low-latency + TBO：部分 overlap（需要 enable_two_batch_overlap=true）
3. MegaMoE：原生 overlap，不依赖切 batch

这三档在 SLA 模型里会给出显著不同的 ITL/TTFT，混在一起标定会让 profiler 拟合出错。

核心实现机制

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 中的注册