From 0f8e53947da53dc900f51a6e888a120523887a5b Mon Sep 17 00:00:00 2001 From: Binyao Jiang Date: Sun, 30 Nov 2025 23:04:10 -0800 Subject: [PATCH] =?UTF-8?q?[Piecewise]=20Use=20same=20global=20graph=20mem?= =?UTF-8?q?ory=20pool=20as=20the=20main=20cuda=20graph=20=E2=80=A6=20(#140?= =?UTF-8?q?44)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Stefan He Co-authored-by: BBuf <1182563586@qq.com> --- .../device_communicators/pynccl_allocator.py | 3 ++- .../piecewise_cuda_graph_runner.py | 17 ++++------------- 2 files changed, 6 insertions(+), 14 deletions(-) diff --git a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py index 1aca62200..e7dbb9cdd 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py +++ b/python/sglang/srt/distributed/device_communicators/pynccl_allocator.py @@ -71,7 +71,8 @@ def is_symmetric_memory_enabled(): def set_graph_pool_id(graph_pool_id): global _graph_pool_id - _graph_pool_id = graph_pool_id + if _graph_pool_id is not None: + _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 a5c1f4d4e..789fd757a 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -45,6 +45,10 @@ 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, @@ -143,19 +147,6 @@ 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