执行摘要

• 一句话：重构CUDA Graph Runner/Backend分层架构
• 推荐动作：值得深度精读。重点关注BaseCudaGraphBackend接口设计（如何平衡通用性与灵活性）、CudaGraphConfig配置数据类（Phase/Backend枚举设计、diff-based导出）以及BaseCudaGraphRunner与后端之间的capture_session/replay_session上下文管理。这些设计模式可直接应用于其他硬件后端或自定义捕获策略。

功能与动机

现有三个CUDA图实现（full cuda graph, breakable cuda graph, torch-compile-based piecewise cuda graph）之间存在大量重复代码，且每个策略需要独立的Runner实现。RFC Issue #23004提出通过Runner-Backend分离消除重复，并支持灵活的per-phase后端选择（如decode使用full，prefill使用breakable/tc_piecewise）。

实现拆解

1. 配置层迁移：创建model_executor/cuda_graph_config.py，定义Phase和Backend枚举类，PhaseConfig和CudaGraphConfig数据类。重构ServerArgs，将原先分散的disable_cuda_graph、enable_breakable_cuda_graph等字段归一为cuda_graph_config，并支持JSON和便捷CLI两种输入方式。

2. 后端接口定义：在runner_backend/下新增BaseCudaGraphBackend抽象基类，声明capture_session、capture_one、can_run、replay_session、replay、cleanup等契约方法。三种具体后端（Full、Breakable、TC_Piecewise）分别实现该接口。

3. 阶段运行器重构：创建BaseCudaGraphRunner基类，提供can_run、capture、capture_one_shape、replay_prepare、replay等框架。PrefillCudaGraphRunner和DecodeCudaGraphRunner继承自该基类，并在初始化时通过resolve_prefill_backend/resolve_decode_backend工厂函数绑定对应后端。

4. 旧文件删除与重命名：删除piecewise_cuda_graph_runner.py（860行）和breakable_cuda_graph_runner.py（541行），将cuda_graph_runner.py迁移并重命名为decode_cuda_graph_runner.py（新增294行/删除463行）。

5. NPU硬件支持：新增NPUCudaGraphBackend，集成NPU的torch.npu.NPUGraph，并在capture_one中支持post_warmup_hook回调以满足后端接口对齐。

6. 测试与配置工具：迁移测试文件到新目录，更新mock_server_args.py以支持新配置结构，新增CudaGraphBufferRegistry等buffer管理类。

关键文件：

• python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py（模块 预填充运行器；类别 source；类型 core-logic；符号 PrefillCudaGraphRunner, init, _is_mamba_track_enabled, _cache_loc_dtype）: 新增核心文件：PrefillCudaGraphRunner，负责prefill阶段的CUDA图捕获与重放，通过resolve_prefill_backend绑定后端（默认TcPiecewiseCudaGraphBackend）。
• python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py（模块 解码运行器；类别 source；类型 rename-or-move；符号 _make_graph_key, get_is_capture_mode, compile_in_capture_mode, model_capture_mode）: 原cuda_graph_runner.py重命名并重构为DecodeCudaGraphRunner，负责decode/TARGET_VERIFY/DLLM_EXTEND阶段的CUDA图捕获。
• python/sglang/srt/model_executor/cuda_graph_config.py（模块 配置模型；类别 source；类型 data-contract；符号 Phase, Backend, PhaseConfig, CudaGraphConfig）: 新增配置数据模型：Phase/Backend枚举、PhaseConfig/CudaGraphConfig dataclass，以及default、parse、check_cuda_graph_backend辅助函数。
• python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py（模块 运行器基类；类别 source；类型 data-contract；符号 freeze_gc, get_batch_sizes_to_capture, BaseCudaGraphRunner, init）: 新增抽象基类BaseCudaGraphRunner，定义了can_run、capture、capture_one_shape、replay_prepare、replay等方法框架，并包含freeze_gc、get_batch_sizes_to_capture工具函数。
• python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py（模块 旧运行器；类别 source；类型 deletion；符号 freeze_gc, _to_torch, patch_model, get_global_graph_memory_pool）: 被删除的核心文件：原PiecewiseCudaGraphRunner（860行）被新PrefillCudaGraphRunner + TcPiecewiseCudaGraphBackend替代。
• python/sglang/srt/model_executor/breakable_cuda_graph_runner.py（模块 旧运行器；类别 source；类型 deletion；符号 BreakableCudaGraphRunner, init, _has_inactive_dp_rank, _init_buffers）: 被删除的核心文件：原BreakableCudaGraphRunner（541行）被PrefillCudaGraphRunner + BreakableCudaGraphBackend替代。
• python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py（模块 后端实现；类别 source；类型 core-logic；符号 _toggle_multi_platform_ops, TcPiecewiseCudaGraphBackend, init, build_compilation_config）: 新增核心后端：TcPiecewiseCudaGraphBackend，基于torch.compile的piecewise CUDA图捕获，支持eager/inductor编译器。
• python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py（模块 NPU后端；类别 source；类型 core-logic；符号 NPUCudaGraphBackend, init, capture_session, capture_one）: 新增NPU后端适配：NPUCudaGraphBackend继承BaseCudaGraphBackend，使用torch.npu.NPUGraph捕获，处理了post_warmup_hook对齐。

关键符号：BaseCudaGraphRunner.can_run, BaseCudaGraphRunner.capture, BaseCudaGraphRunner.capture_one_shape, BaseCudaGraphRunner.replay_prepare, BaseCudaGraphRunner.replay, BaseCudaGraphBackend.capture_session, BaseCudaGraphBackend.capture_one, BaseCudaGraphBackend.replay_session, BaseCudaGraphBackend.replay, FullCudaGraphBackend.capture_one, BreakableCudaGraphBackend.capture_one, TcPiecewiseCudaGraphBackend.capture_one, CudaGraphConfig.from_dict, CudaGraphConfig.to_dict, PrefillCudaGraphRunner._run_forward, DecodeCudaGraphRunner._make_graph_key

关键源码片段

python/sglang/srt/model_executor/cuda_graph_config.py

新增配置数据模型：Phase/Backend枚举、PhaseConfig/CudaGraphConfig dataclass，以及default、parse、check_cuda_graph_backend辅助函数。

# python/sglang/srt/model_executor/cuda_graph_config.py
# 依赖纯 stdlib，便于 ServerArgs 导入而不拉入 torch/srt 后端类。
​
class Phase:
    """模型forward的两个阶段。"""
    DECODE = "decode"
    PREFILL = "prefill"
    ALL = (DECODE, PREFILL)
​
class Backend:
    """每个阶段可使用的CUDA图捕获后端。"""
    FULL = "full"
    BREAKABLE = "breakable"
    TC_PIECEWISE = "tc_piecewise"
    DISABLED = "disabled"
    ALL = (FULL, BREAKABLE, TC_PIECEWISE, DISABLED)
​
# 每个阶段允许的后端不一样：prefill 不允许 full（形状可变）。
ALLOWED_BACKENDS_PER_PHASE = {
    Phase.DECODE: (Backend.FULL, Backend.BREAKABLE, Backend.TC_PIECEWISE, Backend.DISABLED),
    Phase.PREFILL: (Backend.BREAKABLE, Backend.TC_PIECEWISE, Backend.DISABLED),
}
​
@dataclass
class PhaseConfig:
    """每个阶段的CUDA图设置：后端、最大batch size、捕获batch size列表、torch.compile编译器。"""
    backend: str = Backend.DISABLED
    max_bs: Optional[int] = None
    bs: Optional[List[int]] = None
    # 仅当 backend == tc_piecewise 时有效。
    tc_compiler: str = "eager"
​
@dataclass
class CudaGraphConfig:
    """顶层CUDA图配置：decode 和 prefill 各一个 PhaseConfig。"""
    decode: PhaseConfig = field(default_factory=lambda: PhaseConfig(backend=Backend.FULL))
    prefill: PhaseConfig = field(default_factory=lambda: PhaseConfig(backend=Backend.TC_PIECEWISE))
​
    def __getitem__(self, phase: str) -> PhaseConfig:
        if phase not in Phase.ALL:
            raise KeyError(phase)
        return getattr(self, phase)

python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py

新增抽象基类BaseCudaGraphRunner，定义了can_run、capture、capture_one_shape、replay_prepare、replay等方法框架，并包含freeze_gc、get_batch_sizes_to_capture工具函数。

# python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py
from abc import ABC, abstractmethod
​
@contextmanager
def freeze_gc(enable_cudagraph_gc: bool):
    """优化CUDA图捕获期间的垃圾回收。先collect然后冻结剩余对象。"""
    gc.collect()
    should_freeze = not enable_cudagraph_gc
    if should_freeze:
        gc.freeze()
    try:
        yield
    finally:
        if should_freeze:
            gc.unfreeze()
            gc.collect()
​
class BaseCudaGraphRunner(ABC):
    """CUDA图运行器抽象基类。子类（Decode/Prefill）拥有一个Backend处理捕获/回放机制。"""
    # 子类在 capture() 前必须设置：
    buffers: ForwardInputBuffers
    backend: BaseCudaGraphBackend
​
    def __init__(self, model_runner: ModelRunner) -> None:
        self.model_runner = model_runner
        self.device = model_runner.device
        self.device_module = torch.get_device_module(self.device)
        self.tp_size = model_runner.server_args.tp_size
        self.dp_size = model_runner.server_args.dp_size
        self.pp_size = model_runner.server_args.pp_size
        self.attn_tp_size = get_attention_tp_size()
        self.attn_tp_rank = get_attention_tp_rank()
        self.tbo_plugin = TboCudaGraphRunnerPlugin()
​
    @abstractmethod
    def can_run(self, forward_batch: ForwardBatch) -> bool:
        """判断当前batch是否可使用已捕获的CUDA图。"""
        ...
​
    @abstractmethod
    def capture(self) -> None:
        """一次性捕获所有需要形状的CUDA图。"""
        ...
​
    @abstractmethod
    def capture_one_shape(self, size: int, ...) -> None:
        """捕获单个形状：构造dummy batch并调用backend.capture_one。"""
        ...
​
    @abstractmethod
    def replay_prepare(self, forward_batch: ForwardBatch, **kwargs) -> ForwardBatch:
        """填充/重排输入以匹配已捕获的形状，返回静态批次。"""
        ...
​
    @abstractmethod
    def replay(self, forward_batch: ForwardBatch, ...) -> ModelOutput:
        """通过捕获的CUDA图执行batch的forward。"""
        ...

评论区精华

• merrymercy：提出默认prefill应为breakable而非tc_piecewise，并建议将通用定义从cuda_graph重命名为device graph以支持其它硬件。
• ch-wan：批评配置解析逻辑过于复杂，要求增加单元测试验证各种输入组合；建议将JSON配置解析为dataclass而非dict以利用类型提示。
• VDV1985：质疑NPUsNPUGraphRunner为何继承DecodeCudaGraphRunner而非BaseCudaGraphRunner，询问是否意味着NPU不支持prefill图。Oasis-Git回应称NPU目前仅支持decode。
• BBuf：指出NPUCudaGraphBackend.capture_one未接受post_warmup_hook参数，与BaseCudaGraphBackend契约不兼容，导致NPU decode图捕获失败；Oasis-Git后续修复。
• merrymercy：多次强调“不要删除有用的注释”，指出重构过程中AI不慎删除了描述LoRA阶段、DLLM模式等关键注释，要求恢复。

• ch-wan：建议将后端上下文管理器runtime_session重命名为replay_session，避免与运行阶段混淆；Oasis-Git采纳。

  • 默认prefill后端应为breakable而非tc_piecewise (design): 默认为tc_piecewise（原piecewise行为），用户可通过--cuda-graph-backend-prefill=breakable切换。
  • 配置解析应使用dataclass而非dict (design): 已改为dataclass，支持decode.max_bs等类型安全访问。
  • NPU后端capture_one缺少post_warmup_hook参数 (correctness): 已添加post_warmup_hook可选参数并在NPU后端调用。
  • 重构中误删了有用注释 (style): Oasis-Git后续恢复了部分注释，但在最终版本中仍有一些简化。merrymercy要求保留。
  • 配置解析逻辑过于复杂，需增加单元测试 (testing): ch-wan同意未来PR增加测试，在当前PR添加TODO标记。

风险与影响

• 风险：

  1. 核心路径变更：CUDA图是推理性能的核心，重构可能引入回归，尤其在多模型（DeepSeek、MLA、LoRA）与多种注意力后端的组合下。
  2. 配置迁移兼容性：大量旧CLI参数被DeprecatedAction替代，但Python API直接传入ServerArgs构造时可能因字段不存在而崩溃（需通过cuda_graph_config透传）。
  3. NPU后端未完全对齐：NPUCudaGraphBackend的capture_one缺少post_warmup_hook形参，虽然在后续修复，但NPU特有replay_with_input_update路径依赖特殊输入更新逻辑。
  4. 依赖方同步风险：EagleWorker、FrozenKvMtpWorker等依赖旧CudaGraphRunner构造方式，重构后必须转为使用DecodeCudaGraphRunner。
  5. 测试覆盖不足：配置解析逻辑复杂度高但没有单元测试，多线程/多流捕获场景未经充分验证。
     - 影响：用户：启动命令需要调整，旧参数如--disable-piecewise-cuda-graph变为--cuda-graph-config或--cuda-graph-backend-prefill=disabled；但DeprecatedAction会警告并继续工作，短期兼容。
     系统：架构清晰度大幅提升，新增CUDA图策略只需实现Backend子类并注册工厂，无需改动Runner代码。
     团队：维护成本降低，但需要培训确保新扩展点被正确使用。
     影响范围：全仓库160个文件变更，涵盖ServerArgs、ModelRunner、注意力后端、EAGLE推测解码worker、NPU硬件后端、CI测试等。

• 风险标记：核心路径变更, 配置迁移兼容性, NPU后端对齐, 依赖方同步风险, 测试覆盖不足

关联脉络

• PR #23004 [RFC] Cuda Graph Runner Backend Refactor: 本PR的动机和设计方案均源自该RFC Issue，定义了Runner-Backend分离的目标和步骤。
• PR #28081 [refactor] Fold FrozenKVMTPCudaGraphRunner onto the shared DecodeCudaGraphRunner base: 后续依赖本重构的PR，将FrozenKVMTPCudaGraphRunner合并到共享的DecodeCudaGraphRunner基类，体现本重构提供的扩展性。
• PR #28093 [Spec] Move draft-extend prep to EagleDraftWorkerBase; unify prepare_for_* names: 后续依赖本重构的speculative-decoding PR，利用统一的Runner基类简化EagleWorker代码。