# PR #30460 完整报告

- 仓库：`sgl-project/sglang`
- 标题：[DeepSeek V2] Reorder dual-stream MoE to main-first to avoid CUDA graph stream explosion
- 合并时间：2026-07-10 07:59
- 原文链接：http://prhub.com.cn/sgl-project/sglang/pull/30460

---

# 执行摘要

- 一句话：交换 dual-stream MoE 执行顺序，消除 CUDA graph 流爆炸
- 推荐动作：值得精读，特别是 GPU 编程中 CUDA graph 捕获与 tensor 生命周期管理的设计模式。展示了别名和 capture flag 两种解决方案的权衡，以及如何通过细粒度标志避免跨流的内存释放问题。

# 功能与动机

Kimi-K2.5-NVFP4 profiling 显示目标验证解码 CUDA graph 扩展到 ~61 个流，而非预期的 2 个（main + alt）。根因在于 `DeepseekV2MoE.forward_normal_dual_stream` 中 shared expert 分支在 alt 流上先于 routed 分支执行。在 CUDA graph 捕获期间，这种 alt-first 逐层排序导致 `cudaGraphInstantiate` 为每一层分配新的 side stream，与 #30025 中 DSA indexer 的“流爆炸”机制相同。原先 alt-first 是由 #29463 要求的，但该顺序在捕获时造成大量流。

# 实现拆解

1. 在 `deepseek_v2.py` 的 `forward_normal_dual_stream` 中，将 shared expert 的调用从 alt 流开头移到 routed 分支之后，确保 main 流操作先于 alt 流。同时预计算 `has_shared_output` 标志替代之前对 `shared_output is not None` 的运行时判断。
2. 在 `capture_mode.py` 的上下文管理器 `model_capture_mode` 中，设置 `CaptureFlags.disable_dispose_tensor = True`，使 CUDA graph 捕获期间 `dispose_tensor` 变为空操作，避免 routed deep_gemm 释放 `hidden_states` 后被 shared expert 读取到空指针。
3. 在 `common.py` 的 `dispose_tensor` 函数中加入检查，若 `disable_dispose_tensor` 为 `True` 则跳过释放。
4. 在 `runtime_context.py` 的 `CaptureFlags` 类中新增 `disable_dispose_tensor: bool = False` 字段。
5. 未新增测试，但 CI 通过了现有的 DeepSeek-V3/V4 FP4 测试和 `test_moe_ep_extra.py`。

关键文件：
- `python/sglang/srt/models/deepseek_v2.py`（模块 模型执行；类别 source；类型 core-logic；符号 forward_normal_dual_stream）: 核心改动：重排 forward_normal_dual_stream 中 routed 和 shared expert 的执行顺序，预计算 has_shared_output 标志。
- `python/sglang/srt/model_executor/runner_utils/capture_mode.py`（模块 图捕获模式；类别 source；类型 core-logic；符号 model_capture_mode）: 在 model_capture_mode 上下文中设置 disable_dispose_tensor 标志，使 CUDA graph 捕获期间 dispose_tensor 成为空操作。
- `python/sglang/srt/utils/common.py`（模块 工具函数；类别 source；类型 core-logic；符号 dispose_tensor）: 在 dispose_tensor 函数中加入检查，若 disable_dispose_tensor 为 True 则跳过释放，避免 deep_gemm 的 dispose 操作影响图捕获。
- `python/sglang/srt/runtime_context.py`（模块 运行时标志；类别 source；类型 data-contract；符号 CaptureFlags）: 在 CaptureFlags 中添加 disable_dispose_tensor 字段，作为捕获期间的内存管理标志。

关键符号：forward_normal_dual_stream, model_capture_mode, dispose_tensor

## 关键源码片段

### `python/sglang/srt/model_executor/runner_utils/capture_mode.py`

在 model_capture_mode 上下文中设置 disable_dispose_tensor 标志，使 CUDA graph 捕获期间 dispose_tensor 成为空操作。

```python
# (capture_mode.py) - 全局捕获模式上下文管理器，新增 disable_dispose_tensor 设置

@contextmanager
def model_capture_mode():
    global is_capture_mode
    from sglang.srt.runtime_context import get_flags

    # 禁用 dispose_tensor：防止 routed deep_gemm 在捕获期间释放 hidden_states，
    # 否则 shared expert 在 alt 流上读取时得到 data_ptr()==0 并记录到 CUDA graph 中。
    is_capture_mode = True
    get_flags().capture.disable_dispose_tensor = True
    try:
        yield
    finally:
        is_capture_mode = False
        get_flags().capture.disable_dispose_tensor = False

```

### `python/sglang/srt/utils/common.py`

在 dispose_tensor 函数中加入检查，若 disable_dispose_tensor 为 True 则跳过释放，避免 deep_gemm 的 dispose 操作影响图捕获。

```python
# (common.py) - dispose_tensor 新增 capture flag 检查

def dispose_tensor(x: torch.Tensor):
    """
    Dispose a tensor by freeing its memory.
    During piecewise CUDA graph capture/replay, we skip disposal to avoid
    interfering with torch.compile's memory tracking and graph recording.
    """
    # 已有检查：piecewise 捕获时跳过
    from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
        is_in_tc_piecewise_cuda_graph,
    )
    if is_in_tc_piecewise_cuda_graph():
        return

    # 新增检查：decode/spec 图捕获期间跳过 disposed，防止 deep_gemm pre-permute
    # 释放 hidden_states，导致 shared expert 读取空指针。
    from sglang.srt.runtime_context import get_flags
    if get_flags().capture.disable_dispose_tensor:
        return

    x.set_(torch.empty((0,), device=x.device, dtype=x.dtype))

```

# 评论区精华

设计上需要同时满足三个约束：主流优先避免流爆炸、PDL 重叠保持（routed 是最后一个主流核并与残差融合）、`dispose_tensor` 风险消除。初始采用 `shared_in = hidden_states[:]` 别名方案，后由 Jiminator 重构为 capture flag 方案，更简洁且不依赖内存映射。PR 作者在 body 中详细分析了流爆炸根因和验证结果，审查者 ch-wan 批准，无公开分歧。

- 暂无高价值评论线程

# 风险与影响

- 风险：核心风险是 `model_capture_mode` 上下文管理器异常退出可能导致 `disable_dispose_tensor` 未重置，但 try/finally 保证清理。另一风险是 shared expert 后置于 alt 流，可能减少与 routed 的重叠，但 routed 是最后一个主流核且与残差融合，重叠被保留。对非 DeepSeek 模型无影响，因为 `CaptureFlags.disable_dispose_tensor` 仅在 DeepSeek 的双流路径中使用。性能风险极低，profiling 显示 TPOT 和 accept length 在噪声范围内。
- 影响：直接影响使用 DeepSeek-V2 dual-stream MoE 的模型（如 Kimi-K2.5）在 CUDA graph 捕获时的资源消耗：流数从每层 61 降至 2-3，减少图实例化时间和流管理开销。对不启用 CUDA graph 捕获的模式无影响。开发者无配置负担，新增字段为内部实现细节。团队维护成本低。
- 风险标记：流爆炸修复 , capture 标志 , dispose_tensor 豁免 , 无性能回归 , 依赖上下文管理器

# 关联脉络

- PR #30025 fix: reorder DSA indexer dual-stream ops to avoid CUDA graph stream explosion: 同类问题修复：DSA indexer 的双流顺序导致流爆炸，与本 PR 共享相同的根因和解决思路。
- PR #29463 [DeepSeek V3] Reland: run routed experts on main stream in dual-stream MoE: 先前将 routed 移到主流时引入了 dispose_tensor 风险，本 PR 通过 capture flag 解决了该风险。