执行摘要

• 一句话：共享 prefill 和 decode 的 CUDA graph 内存池
• 推荐动作：值得精读。该 PR 展示了如何通过共享全局资源来减少内存占用，且改动极小（+26/-8）。适合作为类似共享池模式（如 EAGLE 的 draft pool）的参考。

功能与动机

自 PR#23906 重构后，prefill 和 decode 的 graph 后端各自独立分配 graph_pool_handle()，导致每个被捕获的 graph 独立预留内存。而全局辅助函数 _global_graph_memory_pool 闲置未用。PR body 指出："The two phases never replay concurrently — so sharing reserves only the larger phase's footprint instead of the sum."

实现拆解

1. 在 python/sglang/srt/model_executor/runner_utils/pool.py 中新增 get_or_create_global_graph_memory_pool(device_module) 函数：首次调用时通过 device_module.graph_pool_handle() 创建 pool，后续复用已创建的句柄。
2. 修改 FullCudaGraphBackend（full_cuda_graph_backend.py）：在其 capture_session() 方法中将原来的 self._pool = self._device_module.graph_pool_handle() 替换为调用 get_or_create_global_graph_memory_pool(self._device_module)。
3. 修改 BreakableCudaGraphBackend（breakable_cuda_graph_backend.py）：同样在 capture_session() 中替换为共享 pool。
4. 修改 TcPiecewiseCudaGraphBackend（tc_piecewise_cuda_graph_backend.py）：在 _run_compile_pass() 方法中将 self._pool = self._device_module.graph_pool_handle() 替换为共享 pool。
   所有后端都新增导入 get_or_create_global_graph_memory_pool。

关键文件：

• python/sglang/srt/model_executor/runner_utils/pool.py（模块 内存池；类别 source；类型 data-contract；符号 get_or_create_global_graph_memory_pool）: 新增 get_or_create_global_graph_memory_pool 函数，是共享 pool 的核心入口。
• python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py（模块 后端执行器；类别 source；类型 data-contract）: FullCudaGraphBackend 的 capture_session 改为使用共享 pool，是三个后端之一。
• python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py（模块 后端执行器；类别 source；类型 data-contract）: BreakableCudaGraphBackend 的 capture_session 改为使用共享 pool，是三个后端之二。
• python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py（模块 后端执行器；类别 source；类型 data-contract）: TcPiecewiseCudaGraphBackend 的 _run_compile_pass 改为使用共享 pool，是三个后端之三。

关键符号：get_or_create_global_graph_memory_pool, FullCudaGraphBackend.capture_session, BreakableCudaGraphBackend.capture_session, TcPiecewiseCudaGraphBackend._run_compile_pass

关键源码片段

python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py

FullCudaGraphBackend 的 capture_session 改为使用共享 pool，是三个后端之一。

# python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py
# ... ( 导入部分 )
from sglang.srt.model_executor.runner_utils.pool import (
    get_or_create_global_graph_memory_pool,
)
# ...
class FullCudaGraphBackend(BaseCudaGraphBackend):
    # ...
    @contextmanager
    def capture_session(self, stream: torch.cuda.Stream):
        if self._pool is None:
            # Before: self._pool = self._device_module.graph_pool_handle()
            self._pool = get_or_create_global_graph_memory_pool(self._device_module)
        set_graph_pool_id(self._pool)
        self._capture_stream = stream
        try:
            yield
        finally:
            self._capture_stream = None

评论区精华

PR 无 review 评论。

• 暂无高价值评论线程

风险与影响

• 风险：风险较低。变更仅涉及 graph capture 阶段的内存池分配方式，不改变 graph 捕获或回放的逻辑。共享 pool 仅在 prefill 和 decode 不同时回放时安全——当前设计确实如此。潜在的回归风险在于：若未来引入并发回放（如同步执行 prefill 和 decode），则共享 pool 可能导致竞争；但当前架构无此场景。
• 影响：直接影响 CUDA graph 捕获时的 GPU 内存占用，prefill 和 decode 的总预留内存将从两者之和减少为最大值。对用户无功能影响，但可降低 OOM 风险。对团队而言，这是一个低风险、高收益的内存优化。
• 风险标记：核心路径变更, 缺少测试覆盖

关联脉络

• PR #23906 [Refactor] Refactor CUDA graph runner/backend: 该 PR 引入了 prefill 和 decode 各自独立分配 graph_pool_handle，本 PR 修复了其内存浪费问题。