[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>

tests/v1/cudagraph/test_cudagraph_manager.py

@@ -0,0 +1,111 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from contextlib import contextmanager
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm.config import (
+    CompilationConfig,
+    CUDAGraphMode,
+    ParallelConfig,
+    SchedulerConfig,
+    VllmConfig,
+)
+from vllm.distributed.device_communicators import pynccl_allocator
+from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils
+from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
+
+pytestmark = pytest.mark.cpu_test
+
+
+@pytest.fixture(autouse=True)
+def _reset_graph_pool_id():
+    pynccl_allocator._graph_pool_id = None
+    yield
+    pynccl_allocator._graph_pool_id = None
+
+
+def _create_vllm_config() -> MagicMock:
+    compilation_config = CompilationConfig(
+        cudagraph_mode="FULL",
+        cudagraph_capture_sizes=[4],
+    )
+    compilation_config.max_cudagraph_capture_size = 4
+    compilation_config.post_init_cudagraph_sizes()
+
+    vllm_config = MagicMock(spec=VllmConfig)
+    vllm_config.compilation_config = compilation_config
+    vllm_config.scheduler_config = SchedulerConfig.default_factory(max_num_seqs=4)
+    vllm_config.parallel_config = ParallelConfig()
+    vllm_config.speculative_config = None
+    vllm_config.num_speculative_tokens = 0
+    return vllm_config
+
+
+def test_full_capture_sets_graph_pool_id_before_cuda_graph(monkeypatch):
+    """FULL capture must set graph_pool_id before entering torch.cuda.graph().
+
+    NCCL symmetric memory checks this global during graph capture; without
+    it, capture fails with:
+    AssertionError: graph_pool_id is not set under graph capture
+    """
+    graph_pool = object()
+    monkeypatch.setattr(
+        gpu_cudagraph_utils,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
+    )
+    monkeypatch.setattr(
+        gpu_cudagraph_utils.current_platform,
+        "get_global_graph_pool",
+        lambda: graph_pool,
+    )
+
+    manager = gpu_cudagraph_utils.CudaGraphManager(
+        vllm_config=_create_vllm_config(),
+        device=torch.device("cpu"),
+        cudagraph_mode=CUDAGraphMode.FULL,
+        decode_query_len=1,
+    )
+
+    desc = BatchExecutionDescriptor(
+        cg_mode=CUDAGraphMode.FULL,
+        num_tokens=4,
+        num_reqs=4,
+        uniform_token_count=1,
+    )
+    manager._capture_descs[CUDAGraphMode.FULL] = [desc]
+
+    def create_forward_fn(desc, warmup):
+        return lambda _mode: None
+
+    @contextmanager
+    def fake_graph_capture(*args, **kwargs):
+        yield SimpleNamespace(stream=MagicMock())
+
+    fake_offloader = MagicMock()
+
+    def cuda_graph_enter(*args, **kwargs):
+        assert pynccl_allocator._graph_pool_id is graph_pool
+
+    mock_cuda_graph_ctx = MagicMock()
+    mock_cuda_graph_ctx.__enter__ = cuda_graph_enter
+    mock_cuda_graph_ctx.__exit__ = MagicMock(return_value=False)
+
+    with (
+        patch.object(gpu_cudagraph_utils, "graph_capture", fake_graph_capture),
+        patch.object(gpu_cudagraph_utils, "get_offloader", lambda: fake_offloader),
+        patch.object(gpu_cudagraph_utils.torch.cuda, "CUDAGraph"),
+        patch.object(
+            gpu_cudagraph_utils.torch.cuda,
+            "graph",
+            return_value=mock_cuda_graph_ctx,
+        ) as mock_cuda_graph,
+    ):
+        manager.capture(create_forward_fn)
+
+    mock_cuda_graph.assert_called_once()

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -17,6 +17,7 @@ from vllm.compilation.breakable_cudagraph import (
 from vllm.compilation.counter import compilation_counter
 from vllm.config import VllmConfig
 from vllm.config.compilation import CUDAGraphMode
+from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id
 from vllm.distributed.parallel_state import (
     get_pp_group,
     graph_capture,
@@ -346,6 +347,10 @@ class CudaGraphManager:
                         # Sync offloader's copy stream before capture.
                         # Ensure any pre-capture prefetches from offloader are complete.
                         get_offloader().sync_prev_onload()
+                        if self.pool is not None:
+                            set_graph_pool_id(self.pool)
+                        else:
+                            set_graph_pool_id(current_platform.graph_pool_handle())
                         with torch.cuda.graph(graph, self.pool):
                             forward_fn(CUDAGraphMode.NONE)
                             # Join offloader's copy stream after forward to avoid