执行摘要

• 一句话：引入可断裂CUDA图及调试模式
• 推荐动作：值得精读。本 PR 展示了如何通过 ContextVar 和流同步钩子在 CUDA 图捕获中插入动态断裂，设计模式值得借鉴。建议重点阅读 breakable_cuda_graph.py 中的捕获上下文实现和 cuda_graph_runner.py 中的集成逻辑。

功能与动机

标准 CUDA 图捕获整个前向传递为单一图，调试困难且部分操作不兼容（如动态控制流、JIT）。Breakable CUDA Graph 通过在图段间插入断裂，在保留大部分性能收益的同时允许目标操作在图外急切执行，解决调试和兼容性问题。

实现拆解

1. 核心数据结构与捕获上下文（breakable_cuda_graph.py）：定义 GraphBreakInfo NamedTuple 存储断裂信息，使用 ContextVar 管理捕获状态；实现 BreakableCUDAGraphCapture 上下文管理器，在 __enter__/__exit__ 中安装/卸载流同步钩子并控制图捕获开始/结束；eager_on_graph 装饰器标记函数为急切执行。
2. CUDA 工具函数（cuda_utils.py）：封装 cudaGetErrorString 和 checkCudaErrors，用于底层 CUDA 运行时绑定。
3. 集成到图运行器（cuda_graph_runner.py）：在 _capture_graph 和 _create_device_graph 方法中根据环境变量分支，使用 BreakableCUDAGraphCapture 替代标准 device_module.graph；debug 模式通过 eager_on_graph 包装整个 forward 函数。
4. 服务器参数与环境变量（server_args.py、environ.py）：新增 --debug-cuda-graph 参数，设置 SGLANG_USE_BREAKABLE_CUDA_GRAPH，并检查非 CUDA 平台发出警告。
5. 文档与测试：新增 docs/advanced_features/breakable_cuda_graph.md 介绍用法，新增 test/registered/cuda_graph/test_breakable_cuda_graph.py 覆盖无断裂、单断裂、多断裂场景。

关键文件：

• python/sglang/srt/model_executor/breakable_cuda_graph/breakable_cuda_graph.py（模块 图断裂；类别 source；类型 core-logic；符号 _check_cuda_bindings, GraphBreakInfo, get_current_stream, _capture_status）: 可断裂 CUDA 图的核心实现，定义了 GraphBreakInfo、BreakableCUDAGraphCapture 和 eager_on_graph 装饰器。
• python/sglang/srt/model_executor/breakable_cuda_graph/cuda_utils.py（模块 CUDA工具；类别 source；类型 data-contract；符号 _cudaGetErrorString, checkCudaErrors）: CUDA 运行时绑定工具，封装错误检查，被核心模块依赖。
• python/sglang/srt/model_executor/cuda_graph_runner.py（模块 图运行器；类别 source；类型 core-logic；符号 _capture_graph, _create_device_graph）: 集成 Breakable CUDA Graph 到现有图运行器，修改 _capture_graph 和 _create_device_graph 方法，是功能落地的关键入口。
• python/sglang/srt/server_args.py（模块 参数配置；类别 source；类型 configuration；符号 debug_cuda_graph, _handle_environment_variables）: 添加 --debug-cuda-graph CLI 参数，并处理环境变量和平台检查。
• python/sglang/srt/environ.py（模块 环境变量；类别 source；类型 configuration）: 新增 SGLANG_USE_BREAKABLE_CUDA_GRAPH 环境变量定义。
• test/registered/cuda_graph/test_breakable_cuda_graph.py（模块 测试；类别 test；类型 test-coverage；符号 TestBreakableCUDAGraphBasic, test_no_break_capture_replay, test_single_break, test_multiple_breaks）: 覆盖基本捕获/重放、单一断裂、多断裂场景的单元测试。
• docs/advanced_features/breakable_cuda_graph.md（模块 文档；类别 docs；类型 documentation；符号 my_dynamic_op, forward）: 用户文档，包含用法示例和参数说明。
• docs/index.rst（模块 文档索引；类别 docs；类型 documentation）: 更新文档索引，添加 breakable_cuda_graph.md 入口。

关键符号：eager_on_graph, BreakableCUDAGraph, BreakableCUDAGraphCapture, break_graph, _capture_graph, _create_device_graph, checkCudaErrors, _capture_status, _is_capturing, _install_wait_stream_hook, _uninstall_wait_stream_hook

关键源码片段

python/sglang/srt/model_executor/breakable_cuda_graph/breakable_cuda_graph.py

可断裂 CUDA 图的核心实现，定义了 GraphBreakInfo、BreakableCUDAGraphCapture 和 eager_on_graph 装饰器。

import torch
from contextvars import ContextVar
from typing import Callable, NamedTuple, Any
​
# 存储捕获期间的图断裂信息列表
_captured_graphs_var: ContextVar[list | None] = ContextVar("captured_graphs", default=None)
# 当前捕获的 CUDA 流
_current_stream_var: ContextVar[torch.cuda.Stream | None] = ContextVar("current_stream", default=None)
​
class GraphBreakInfo(NamedTuple):
    func: Callable # 在断裂处被调用的 Python 函数
    output: Any # 函数的输出（必须为张量）
    graph_handle: Any # 捕获后的图句柄或实例化后的执行句柄
​
class BreakableCUDAGraph:
    """代表一个可断裂的 CUDA 图对象"""
    def __init__(self):
        self._capture = BreakableCUDAGraphCapture(self)
​
class BreakableCUDAGraphCapture:
    """上下文管理器，用于进入/退出可断裂 CUDA 图捕获"""
    def __init__(self, graph: BreakableCUDAGraph, stream: torch.cuda.Stream = None, pool=None):
        self.graph = graph
        self.stream = stream or torch.cuda.current_stream()
        self.pool = pool
​
    def __enter__(self):
        # 设置当前流到 ContextVar，以便被钩子访问
        _current_stream_var.set(self.stream)
        # 安装 wait_stream 钩子，用于跟踪派生流及其同步
        _install_wait_stream_hook()
        # 开始 CUDA 图捕获
        if self.pool:
            torch.cuda.cudart().cudaStreamBeginCapture(self.pool, self.stream)
        else:
            self.stream.begin_capture()
        return self
​
    def __exit__(self, *args):
        # 卸载钩子
        _uninstall_wait_stream_hook()
        # 结束捕获，保存图句柄
        self.stream.end_capture()
        # 清除 ContextVar 状态
        _current_stream_var.set(None)
        _captured_graphs_var.set(None)
​
def eager_on_graph(enable: bool = True):
    """装饰器：标记函数在 CUDA 图捕获期间以急切模式执行"""
    def decorator(func: Callable) -> Callable:
        if enable:
            func._eager_on_graph = True
        return func
    return decorator

python/sglang/srt/model_executor/cuda_graph_runner.py

集成 Breakable CUDA Graph 到现有图运行器，修改 _capture_graph 和 _create_device_graph 方法，是功能落地的关键入口。

# cuda_graph_runner.py _capture_graph 方法的核心分支逻辑
def _capture_graph(self, graph, pool, stream, run_once_fn):
    # 若启用 debug 模式，确保 breakable CUDA graph 环境变量已设置
    if self.model_runner.server_args.debug_cuda_graph:
        assert envs.SGLANG_USE_BREAKABLE_CUDA_GRAPH.get(), \
            "Breakable CUDA graph is not enabled in debug mode"
​
    memory_saver_adapter = TorchMemorySaverAdapter.create(
        enable=self.model_runner.server_args.enable_memory_saver and
               get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
    )
​
    if envs.SGLANG_USE_BREAKABLE_CUDA_GRAPH.get():
        if memory_saver_adapter.enabled:
            raise NotImplementedError(
                "Breakable CUDA graph is not compatible with memory saver mode"
            )
        graph_ctx = BreakableCUDAGraphCapture
    else:
        # 原始上下文：使用 memory saver 适配器或标准 device_module.graph
        graph_ctx = (partial(memory_saver_adapter.cuda_graph, tag=GPU_MEMORY_TYPE_CUDA_GRAPH)
                     if memory_saver_adapter.enabled
                     else self.device_module.graph)
​
    # debug 模式下，将整个 forward 函数包装为 eager on graph
    if self.model_runner.server_args.debug_cuda_graph:
        captured_fn = eager_on_graph(True)(run_once_fn)
    else:
        captured_fn = run_once_fn
​
    with graph_ctx(cuda_graph=graph, pool=pool, stream=stream):
        out = captured_fn()
    return out

评论区精华

• 平台兼容性（BBuf）：指出 --debug-cuda-graph 在 ROCm 上因 BreakableCUDAGraph 仅在非 HIP 构建导入导致 NameError。作者在 server_args.py 中添加了 is_cuda() 检查并警告，但底层导入路径仍可能触发。结论：部分缓解，建议运行时动态导入。
• 结构化输出写回（BBuf）：debug 模式下 eager 回调产生新对象但原始捕获缓冲区未更新，重放时使用旧值。需递归写回结构化输出。结论：后续 commit 通过 _copy_output 修复。
• 命名争议（BBuf）：初始装饰器名 non_graph 被认为不直观。结论：最终改为 eager_on_graph。

• 与 piecewise CUDA graph 合并（ch-wan）：询问 --debug-cuda-graph 能否用于 piecewise CUDA graph。结论：未明确解决，但显示了对合并机会的关注。

• 平台兼容性：ROCm 上可能 NameError (design): 作者在 server_args.py 中添加了 is_cuda() 检查并发出警告，但底层导入路径仍可能被封死。进一步可在运行时动态导入。

• debug 模式下输出写回不完整 (correctness): 修复 commit 中添加了 _copy_output 写回逻辑，已解决。
• 装饰器命名不直观 (style): 最终修改为 eager_on_graph。
• 能否将 --debug-cuda-graph 用于 piecewise CUDA graph (question): 未明确解决，显示作者在思考两者融合。

风险与影响

• 风险：
  • 平台风险：仅依赖 CUDA 绑定（cuda-python），ROCm/HIP 不可用。虽添加了 is_cuda() 检查，但误启仍会运行时异常。
  • 结构化输出写回缺陷：debug 模式下若 forward 返回结构化对象，写回可能不完整，现已修复。
  • 图句柄生命周期：早期版本中无断裂路径重放使用已销毁图句柄，已修复。
  • 与内存保护不兼容：Breakable CUDA Graph 与 --enable-memory-saver 冲突，显式抛出 NotImplementedError。
  • 性能退化：debug 模式完全丧失 CUDA 图性能优势，仅限调试。
• 影响：
  • 用户：获得 --debug-cuda-graph 调试工具，可通过 @eager_on_graph 选择性将不兼容操作排除在图外，无需全局禁用 CUDA 图。
  • 系统：扩展 CUDA 图框架，集成点涉及 cuda_graph_runner.py 核心路径，但默认不启用，无额外开销。
  • 团队：需维护新模块及与 piecewise CUDA graph 的联合使用，文档清晰降低学习成本。
  • 风险标记：平台限制（仅CUDA）, 结构化输出写回缺陷（已修复）, 与内存保护模式不兼容, debug模式消除性能优势

关联脉络

• PR #22218 [Experimental] Breakable Piecewise Cuda Graph: 紧密相关的功能演进，本 PR 继承并改进了同一作者的实验性实现，引入调试模式和更稳定的图断裂机制。