执行摘要

• 一句话：交换 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 成为空操作。

# (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 操作影响图捕获。

# (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 解决了该风险。