[BugFix] Set graph_pool_id before FULL CUDA graph capture in ModelRunner V2 (#48843)
Signed-off-by: Markov Ilya <markovilya19@gmail.com> Co-authored-by: Markov Ilya <markovilya19@gmail.com>
Ilya Markov committed
Jul 21, 2026 at 14:59 UTC
f890e1dbe2fe5deddd5417250b4e1e2db71e9314
2 files changed
+116
tests/v1/cudagraph/test_cudagraph_manager.py
new
+111
@@ -0,0 +1,111 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from contextlib import contextmanager
5
+from types import SimpleNamespace
6
+from unittest.mock import MagicMock, patch
7
+
8
+import pytest
9
+import torch
10
+
11
+from vllm.config import (
12
+ CompilationConfig,
13
+ CUDAGraphMode,
14
+ ParallelConfig,
15
+ SchedulerConfig,
16
+ VllmConfig,
17
+)
18
+from vllm.distributed.device_communicators import pynccl_allocator
19
+from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils
20
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
21
+
22
+pytestmark = pytest.mark.cpu_test
23
+
24
+
25
+@pytest.fixture(autouse=True)
26
+def _reset_graph_pool_id():
27
+ pynccl_allocator._graph_pool_id = None
28
+ yield
29
+ pynccl_allocator._graph_pool_id = None
30
+
31
+
32
+def _create_vllm_config() -> MagicMock:
33
+ compilation_config = CompilationConfig(
34
+ cudagraph_mode="FULL",
35
+ cudagraph_capture_sizes=[4],
36
+ )
37
+ compilation_config.max_cudagraph_capture_size = 4
38
+ compilation_config.post_init_cudagraph_sizes()
39
+
40
+ vllm_config = MagicMock(spec=VllmConfig)
41
+ vllm_config.compilation_config = compilation_config
42
+ vllm_config.scheduler_config = SchedulerConfig.default_factory(max_num_seqs=4)
43
+ vllm_config.parallel_config = ParallelConfig()
44
+ vllm_config.speculative_config = None
45
+ vllm_config.num_speculative_tokens = 0
46
+ return vllm_config
47
+
48
+
49
+def test_full_capture_sets_graph_pool_id_before_cuda_graph(monkeypatch):
50
+ """FULL capture must set graph_pool_id before entering torch.cuda.graph().
51
+
52
+ NCCL symmetric memory checks this global during graph capture; without
53
+ it, capture fails with:
54
+ AssertionError: graph_pool_id is not set under graph capture
55
+ """
56
+ graph_pool = object()
57
+ monkeypatch.setattr(
58
+ gpu_cudagraph_utils,
59
+ "get_pp_group",
60
+ lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
61
+ )
62
+ monkeypatch.setattr(
63
+ gpu_cudagraph_utils.current_platform,
64
+ "get_global_graph_pool",
65
+ lambda: graph_pool,
66
+ )
67
+
68
+ manager = gpu_cudagraph_utils.CudaGraphManager(
69
+ vllm_config=_create_vllm_config(),
70
+ device=torch.device("cpu"),
71
+ cudagraph_mode=CUDAGraphMode.FULL,
72
+ decode_query_len=1,
73
+ )
74
+
75
+ desc = BatchExecutionDescriptor(
76
+ cg_mode=CUDAGraphMode.FULL,
77
+ num_tokens=4,
78
+ num_reqs=4,
79
+ uniform_token_count=1,
80
+ )
81
+ manager._capture_descs[CUDAGraphMode.FULL] = [desc]
82
+
83
+ def create_forward_fn(desc, warmup):
84
+ return lambda _mode: None
85
+
86
+ @contextmanager
87
+ def fake_graph_capture(*args, **kwargs):
88
+ yield SimpleNamespace(stream=MagicMock())
89
+
90
+ fake_offloader = MagicMock()
91
+
92
+ def cuda_graph_enter(*args, **kwargs):
93
+ assert pynccl_allocator._graph_pool_id is graph_pool
94
+
95
+ mock_cuda_graph_ctx = MagicMock()
96
+ mock_cuda_graph_ctx.__enter__ = cuda_graph_enter
97
+ mock_cuda_graph_ctx.__exit__ = MagicMock(return_value=False)
98
+
99
+ with (
100
+ patch.object(gpu_cudagraph_utils, "graph_capture", fake_graph_capture),
101
+ patch.object(gpu_cudagraph_utils, "get_offloader", lambda: fake_offloader),
102
+ patch.object(gpu_cudagraph_utils.torch.cuda, "CUDAGraph"),
103
+ patch.object(
104
+ gpu_cudagraph_utils.torch.cuda,
105
+ "graph",
106
+ return_value=mock_cuda_graph_ctx,
107
+ ) as mock_cuda_graph,
108
+ ):
109
+ manager.capture(create_forward_fn)
110
+
111
+ mock_cuda_graph.assert_called_once()
vllm/v1/worker/gpu/cudagraph_utils.py
+5
@@ -17,6 +17,7 @@ from vllm.compilation.breakable_cudagraph import (
17
from vllm.compilation.counter import compilation_counter
18
from vllm.config import VllmConfig
19
from vllm.config.compilation import CUDAGraphMode
20
+from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
21
from vllm.distributed.parallel_state import (
22
get_pp_group,
23
graph_capture,
@@ -346,6 +347,10 @@ class CudaGraphManager:
347
# Sync offloader's copy stream before capture.
348
# Ensure any pre-capture prefetches from offloader are complete.
349
get_offloader().sync_prev_onload()
350
+ if self.pool is not None:
351
+ set_graph_pool_id(self.pool)
352
+ else:
353
+ set_graph_pool_id(current_platform.graph_pool_handle())
354
with torch.cuda.graph(graph, self.pool):
355
forward_fn(CUDAGraphMode.NONE)
356
# Join offloader's copy stream after forward to avoid