DeepSeek-V4 中的算子 Overlap 分析

发表于 更新于 Disqus：

DeepSeek-V4 中的算子 Overlap 分析

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 互补

Kernel 性质 瓶颈资源
wqkv_a GEMM (M=4096, N=2112, K=7168) 中 GEMM MFU ~30-40%，CU 大量闲置
compressor flash_c4 / flash_c128 memory-bound HBM 带宽吃满，MFMA 单元闲置

一个 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 不赚

场景 多 stream 是否有用 原因
两个高 MFU GEMM（都跑 70%+） ❌ SM 都被占满，只是 timesharing
两个 memory-bound op 读不同数据 ❌ HBM 总带宽是上限
两个 op 有数据依赖（A → B） ❌ 必须串行
一个 compute-bound + 一个 memory-bound ✅ 经典 case
一个大 GEMM + 一群 µs 级小 kernel ✅✅ 赚得最多

五、工程实现：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% 的端到端提升：

状态 forward latency (decode) 相对提升
关闭（串行） 基准 —
开启（multi-stream overlap） -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 现状

论文图中的结构 论文章节 承诺的并行性 SGLang 状态
Indexer ∥ wqkv_a §3.2.2 双分支并行 ✅ 算法并行，⚠️ 需 SGLANG_OPT_USE_MULTI_STREAM_OVERLAP=1 开启
Compressor c4 ∥ c128 §3.2.1 双尺度评分并行 ❌ 串行
Shared ∥ Routed Expert §3.3 双 expert 路径并行 ✅ 已实现
Wave dispatch/compute/combine §3.3.3 乒乓 overlap ✅ 已实现
KV store ∥ next layer indexer §3.2 层间 pipeline ❌ 未做

七、总结

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