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 8, 11:15 UTC
[Model Runner V2] Init cuda graph pool when necessary (#33217)
Signed-off-by: Xinyu Chen <xinyu1.chen@intel.com>
Xinyu Chen committed
Feb 12, 2026 at 01:12 UTC
ffb3d553cc9258049bf4d48214c9f4106cc67cfb
2 files changed
+6
-2
vllm/v1/worker/gpu/cudagraph_utils.py
+3
-1
@@ -45,7 +45,9 @@ class CudaGraphManager:
45
)
46
47
self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
48
- self.pool = torch.cuda.graph_pool_handle()
48
+ self.pool = None
49
+ if self.cudagraph_mode != CUDAGraphMode.NONE:
50
+ self.pool = torch.cuda.graph_pool_handle()
51
self.hidden_states: torch.Tensor | None = None
52
53
def needs_capture(self) -> bool:
vllm/v1/worker/gpu/spec_decode/eagle_cudagraph.py
+3
-1
@@ -44,7 +44,9 @@ class EagleCudaGraphManager:
44
)
45
46
self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
47
- self.pool = torch.cuda.graph_pool_handle()
47
+ self.pool = None
48
+ if self.cudagraph_mode != CUDAGraphMode.NONE:
49
+ self.pool = torch.cuda.graph_pool_handle()
50
51
def get_cudagraph_size(self, num_tokens: int) -> int | None:
52
return self.cudagraph_sizes.get(num_tokens)