执行摘要

• 一句话：将 CUDA graph 捕获逻辑提取至独立模块
• 推荐动作：值得精读，了解如何采用“先原地重构再机械移动”策略安全拆分大型类。注意无测试配套，在严格项目需补充回归测试。设计决策上，使用 msgspec 不可变结构体传递捕获结果，有助于防止意外修改。

功能与动机

根据 PR body，目标是“Extract cuda-graph setup into a module”，并通过两个提交实现：先在原地 reshape 捕获链（提取静态方法并引入结果结构体），再机械移动到新模块。此做法分离关注点，使 ModelRunner 更专注于高层编排。

实现拆解

1. 准备阶段（原地重构）：在 model_runner.py 内将 init_cuda_graphs 拆分为 capture_cuda_graphs、capture_prefill_graph、capture_decode_graph 三个 @staticmethod，并使用新引入的 CudaGraphsCapture 和 DecodeGraphCapture msgspec 结构体作为返回值，统一结果传递方式。
2. 提取模块：创建 python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py，将上述三个函数及结构体定义、相关导入（如 EagerRunner、NPUGraphRunner、GraphSharedOutput 等）整体移动至新文件，保留完全相同逻辑。
3. 清理原类：在 model_runner.py 中删除被搬走的函数和不再需要的导入（如 prealloc_symmetric_memory_pool、NPUGraphRunner 等），并加入对新模块的导入语句。
4. 调用点适配：原调用点（如 init_decode_cuda_graph）改为直接调用 capture_decode_graph(model_runner=self) 等，其余初始化顺序不受影响。

关键文件：

• python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py（模块 CUDA图模块；类别 source；类型 data-contract；符号 DecodeGraphCapture, CudaGraphsCapture, capture_cuda_graphs, capture_prefill_graph）: 新增模块，包含 CUDA graph 捕获的核心函数及返回结构体，是本次重构的主要产出。
• python/sglang/srt/model_executor/model_runner.py（模块 模型执行器；类别 source；类型 data-contract）: 原文件大幅删除 CUDA graph 相关代码并改为导入新模块，是重构的目标文件。

关键符号：capture_cuda_graphs, capture_prefill_graph, capture_decode_graph, DecodeGraphCapture, CudaGraphsCapture

关键源码片段

python/sglang/srt/model_executor/model_runner_components/cuda_graph_setup.py

新增模块，包含 CUDA graph 捕获的核心函数及返回结构体，是本次重构的主要产出。

from __future__ import annotations
import msgspec
from typing import TYPE_CHECKING, Optional
# ... (imports omitted for brevity)
​
class DecodeGraphCapture(msgspec.Struct, frozen=True, kw_only=True):
    """Decode 阶段图捕获的结果，包含 runner 实例与显存用量。"""
    runner: Optional[BaseRunner]
    graph_mem_usage: float
​
​
class CudaGraphsCapture(msgspec.Struct, frozen=True, kw_only=True):
    """整个 CUDA graph 捕获的结果，组合 eager、prefill 与 decode runner。"""
    eager_runner: EagerRunner
    prefill_runner: Optional[BaseRunner]
    decode: DecodeGraphCapture
​
​
def capture_cuda_graphs(
    *, model_runner: ModelRunner, capture_decode_cuda_graph: bool = True
) -> CudaGraphsCapture:
    """捕获 CUDA graph，需先运行 init_attention_backends()。
​
    Spec draft runner 传入 capture_decode_cuda_graph=False，
    因为它们会单独捕获自己的 decode 图。
    """
    # 创建共享输出缓冲区（必须在 runner 创建之前）
    model_runner.graph_shared_output = GraphSharedOutput.create_for_model_runner(
        model_runner
    )
    # 先创建 EagerRunner，用于 warmup kernel 并分配静态缓冲区
    eager_runner = EagerRunner(model_runner)
    # 捕获 prefill 图（可能直接返回 eager runner）
    prefill_runner = capture_prefill_graph(
        model_runner=model_runner, eager_runner=eager_runner
    )
    # 捕获 decode 图（或直接使用 eager runner）
    decode = DecodeGraphCapture(runner=None, graph_mem_usage=0)
    if capture_decode_cuda_graph:
        if model_runner.device in ("cuda", "musa", "cpu", "npu", "xpu"):
            decode = capture_decode_graph(model_runner=model_runner)
        elif (
            current_platform.is_out_of_tree() and current_platform.support_cuda_graph()
        ):
            decode = capture_decode_graph(model_runner=model_runner)
    else:
        decode = DecodeGraphCapture(runner=eager_runner, graph_mem_usage=0)
    # 注册前向 hook（必须在捕获之后，避免 hook 被追踪进图）
    if model_runner.server_args.forward_hooks:
        register_forward_hooks(
            model_runner.model, model_runner.server_args.forward_hooks
        )
    # 预分配对称内存池
    prealloc_symmetric_memory_pool(
        is_draft_worker=model_runner.is_draft_worker,
        enable_symm_mem=model_runner.server_args.enable_symm_mem,
        device=model_runner.device,
        forward_stream=model_runner.forward_stream,
    )
    # canary 管理初始化完成标记
    if model_runner.canary_manager is not None and not model_runner.is_draft_worker:
        model_runner.canary_manager.mark_init_finished()
    # 返回统一结果
    return CudaGraphsCapture(
        eager_runner=eager_runner, prefill_runner=prefill_runner, decode=decode
    )

评论区精华

Gemini Code Assist 机器人提出两点冗余赋值问题：init_decode_cuda_graph 中 self.decode_cuda_graph_runner = None 和 self.graph_mem_usage = 0 立即被 capture_decode_graph 结果覆盖，init_prefill_cuda_graph 中类似。建议删除这些赋值。PR 作者未回应也未修改，但无实质影响。

• init_decode_cuda_graph 中冗余赋值 (style): 建议未采纳，但无功能影响，仅为代码整洁度。
• init_prefill_cuda_graph 中冗余赋值 (style): 建议未采纳，但无功能影响。

风险与影响

• 风险：作为纯机械移动，风险较低。但 CUDA graph 捕获是推理性能关键路径，若新模块的导入路径或函数签名出错，可能导致启动失败或 graph 捕获异常。需确保捕获顺序（eager runner 先于 graph runner 创建）不受干扰。由于无新增测试，回归风险依赖现有 CI。
• 影响：无用户可见行为变化。开发人员将获得更清晰的模块划分，减少 model_runner.py 的行数和职责，便于独立维护和测试 CUDA graph 相关逻辑。未来修改 CUDA graph 策略时可专注于此模块。
• 风险标记：核心路径变更, 机械移动, 缺少测试覆盖, 依赖初始化顺序

关联脉络

• PR #31169 Split initialize() into orchestration helpers: 同系列 ModelRunner 重构，将 initialize 拆分为多个辅助方法，与本 PR 一起降低 model_runner.py 的复杂度。