From 0141ca370fb0ccd625fc94658128f7026ec75f47 Mon Sep 17 00:00:00 2001 From: alisonshao <54658187+alisonshao@users.noreply.github.com> Date: Tue, 2 Dec 2025 09:53:16 -0800 Subject: [PATCH] Revert PR #14044: Restore separate memory pool for piecewise CUDA graph (#14278) --- .../device_communicators/pynccl_allocator.py | 3 +-- .../piecewise_cuda_graph_runner.py | 17 +++++++++++++---- 2 files changed, 14 insertions(+), 6 deletions(-) diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py index e7dbb9cdd..1aca62200 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py +++ b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py @@ -71,8 +71,7 @@ def is_symmetric_memory_enabled(): def set_graph_pool_id(graph_pool_id): global _graph_pool_id - if _graph_pool_id is not None: - _graph_pool_id = graph_pool_id + _graph_pool_id = graph_pool_id def disable_symmetric_memory_context(): diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index 789fd757a..a5c1f4d4e 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -45,10 +45,6 @@ from sglang.srt.layers.dp_attention import ( from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.layers.pooler import EmbeddingPoolerOutput from sglang.srt.layers.torchao_utils import save_gemlite_cache -from sglang.srt.model_executor.cuda_graph_runner import ( - get_global_graph_memory_pool, - set_global_graph_memory_pool, -) from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, @@ -147,6 +143,19 @@ def patch_model(model: torch.nn.Module, compiler: str): _to_torch(model, reverse=True, num_tokens=16) +# Reuse this memory pool across all cuda graph runners. +global_graph_memory_pool = None + + +def get_global_graph_memory_pool(): + return global_graph_memory_pool + + +def set_global_graph_memory_pool(val): + global global_graph_memory_pool + global_graph_memory_pool = val + + def set_torch_compile_config(): import torch._dynamo.config