A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
Remove graph_pool as member of VllmBackend and argument to CUDAGraphWrapper (#23385)
Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: ProExpertProg <11367180+ProExpertProg@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Copilot committed
Aug 25, 2025 at 19:34 UTC
6fad29b11b3680c44782cd6e5fe555779d620d6c
3 files changed
+7
-20
vllm/compilation/backends.py
+2
-12
@@ -294,13 +294,12 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):
294
295
def __init__(self, module: torch.fx.GraphModule,
296
compile_submod_names: list[str], vllm_config: VllmConfig,
297
- graph_pool, vllm_backend: "VllmBackend"):
297
+ vllm_backend: "VllmBackend"):
298
super().__init__(module)
299
from torch._guards import detect_fake_mode
300
self.fake_mode = detect_fake_mode()
301
self.compile_submod_names = compile_submod_names
302
self.compilation_config = vllm_config.compilation_config
303
- self.graph_pool = graph_pool
303
self.vllm_config = vllm_config
304
self.vllm_backend = vllm_backend
305
# When True, it annoyingly dumps the torch.fx.Graph on errors.
@@ -359,7 +358,6 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):
358
runnable=piecewise_backend,
359
vllm_config=self.vllm_config,
360
runtime_mode=CUDAGraphMode.PIECEWISE,
362
- graph_pool=self.graph_pool,
361
cudagraph_options=CUDAGraphOptions(
362
debug_log_enable=piecewise_backend.is_first_graph,
363
gc_disable=not piecewise_backend.is_first_graph,
@@ -405,7 +403,6 @@ class VllmBackend:
403
404
vllm_config: VllmConfig
405
compilation_config: CompilationConfig
408
- graph_pool: Any
406
_called: bool = False
407
# the graph we compiled
408
graph: fx.GraphModule
@@ -433,13 +430,6 @@ class VllmBackend:
430
# them, e.g. backbone (default), eagle_head, etc.
431
self.prefix = prefix or model_tag
432
436
- global_graph_pool = current_platform.get_global_graph_pool()
437
-
438
- # TODO: in the future, if we want to use multiple
439
- # streams, it might not be safe to share a global pool.
440
- # only investigate this when we use multiple streams
441
- self.graph_pool = global_graph_pool
442
-
433
# Passes to run on the graph post-grad.
434
self.post_grad_pass_manager = PostGradPassManager()
435
@@ -586,7 +576,7 @@ class VllmBackend:
576
# propagate the split graph to the piecewise backend,
577
# compile submodules with symbolic shapes
578
PiecewiseCompileInterpreter(self.split_gm, submod_names_to_compile,
589
- self.vllm_config, self.graph_pool,
579
+ self.vllm_config,
580
self).run(*example_inputs)
581
582
graph_path = os.path.join(local_cache_dir, "computation_graph.py")
vllm/compilation/base_static_graph.py
+1
-4
@@ -13,7 +13,7 @@ class AbstractStaticGraphWrapper(Protocol):
13
"""
14
15
def __init__(self, runnable: Callable, vllm_config: VllmConfig,
16
- runtime_mode: CUDAGraphMode, graph_pool: Any, **kwargs):
16
+ runtime_mode: CUDAGraphMode, **kwargs):
17
"""
18
Initializes the StaticGraphWrapper class with graph capturing and
19
execution-related configurations.
@@ -25,9 +25,6 @@ class AbstractStaticGraphWrapper(Protocol):
25
graph runtime. See CUDAGraphMode in vllm/config.py.
26
Note that only the subset enum `NONE`, `PIECEWISE` and `FULL`
27
are used as concrete runtime mode for cudagraph dispatching.
28
- graph_pool (Any):
29
- Graph memory pool handle, e.g.,
30
- `torch.cuda.graph_pool_handle()`.
28
Keyword Args:
29
kwargs: Additional keyword arguments for platform-specific
30
configurations.
vllm/compilation/cuda_graph.py
+4
-4
@@ -67,11 +67,9 @@ class CUDAGraphWrapper:
67
runnable: Callable,
68
vllm_config: VllmConfig,
69
runtime_mode: CUDAGraphMode,
70
- graph_pool: Any = None,
70
cudagraph_options: Optional[CUDAGraphOptions] = None):
71
self.runnable = runnable
72
self.vllm_config = vllm_config
74
- self.graph_pool = graph_pool
73
self.runtime_mode = runtime_mode
74
self.compilation_config = vllm_config.compilation_config
75
@@ -81,8 +79,10 @@ class CUDAGraphWrapper:
79
# assert runtime_mode is not NONE(no cudagraph), otherwise, we don't
80
# need to initialize a CUDAGraphWrapper.
81
assert self.runtime_mode != CUDAGraphMode.NONE
84
- if self.graph_pool is None:
85
- self.graph_pool = current_platform.get_global_graph_pool()
82
+ # TODO: in the future, if we want to use multiple
83
+ # streams, it might not be safe to share a global pool.
84
+ # only investigate this when we use multiple streams
85
+ self.graph_pool = current_platform.get_global_graph_pool()
86
87
if cudagraph_options is None:
88
cudagraph_options = CUDAGraphOptions()