Revert PR #14044: Restore separate memory pool for piecewise CUDA graph (#14278)

python/sglang/srt/distributed/device_communicators/pynccl_allocator.py

@@ -71,7 +71,6 @@ 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
 
 

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