Replace device-specific CUDA/XPU graph APIs with unified torch.accelerator.Graph

tests/v1/cudagraph/test_cudagraph_dispatch.py

@@ -299,12 +299,12 @@ def test_capture_and_replay(self):
                 cudagraph_runtime_mode=CUDAGraphMode.FULL,
                 batch_descriptor=batch_descriptor,
             ),
-            patch("torch.cuda.graph", wraps=torch.cuda.graph) as mock_cuda_graph,
+            patch("torch.accelerator.Graph", wraps=torch.accelerator.Graph) as mock_graph,
         ):
             output1 = wrapper(self.input_tensor)
             # capturing phase should generate a zero output
             assert torch.allclose(output1, torch.zeros_like(output1))
-            mock_cuda_graph.assert_called_once()
+            mock_graph.assert_called_once()
 
         assert batch_descriptor in wrapper.concrete_cudagraph_entries
         entry = wrapper.concrete_cudagraph_entries[batch_descriptor]
@@ -342,13 +342,13 @@ def test_bypass_on_mode_mismatch(self):
                 cudagraph_runtime_mode=CUDAGraphMode.PIECEWISE,
                 batch_descriptor=batch_descriptor,
             ),
-            patch("torch.cuda.graph", wraps=torch.cuda.graph) as mock_cuda_graph,
+            patch("torch.accelerator.Graph", wraps=torch.accelerator.Graph) as mock_graph,
             patch.object(
                 self.model, "forward", wraps=self.model.forward
             ) as mock_forward,
         ):
             wrapper(self.input_tensor)
-            mock_cuda_graph.assert_not_called()
+            mock_graph.assert_not_called()
             mock_forward.assert_called_once()
         assert not wrapper.concrete_cudagraph_entries
 
@@ -365,10 +365,10 @@ def test_bypass_on_mode_none(self):
                 cudagraph_runtime_mode=CUDAGraphMode.NONE,
                 batch_descriptor=batch_descriptor,
             ),
-            patch("torch.cuda.graph", wraps=torch.cuda.graph) as mock_cuda_graph,
+            patch("torch.accelerator.Graph", wraps=torch.accelerator.Graph) as mock_graph,
         ):
             wrapper(self.input_tensor)
-            mock_cuda_graph.assert_not_called()
+            mock_graph.assert_not_called()
         assert not wrapper.concrete_cudagraph_entries
 
 
@@ -378,7 +378,7 @@ def _run_and_monitor_call(
     """Helper to run a single call and monitor the action."""
 
     with (
-        patch("torch.cuda.graph", wraps=torch.cuda.graph) as mock_graph_context,
+        patch("torch.accelerator.Graph", wraps=torch.accelerator.Graph) as mock_graph,
         patch.object(wrapper, "runnable", wraps=wrapper.runnable) as mock_runnable,
     ):
         entry = wrapper.concrete_cudagraph_entries.get(batch_descriptor, None)
@@ -402,7 +402,7 @@ def _run_and_monitor_call(
             with context:
                 wrapper(input_tensor)
 
-        if mock_graph_context.called:
+        if mock_graph.called:
             # note that this is globally mocked, so it will be detected
             # even whether called by the inner or outer wrapper
             return "capture_global"

vllm/compilation/breakable_cudagraph.py

@@ -152,7 +152,7 @@ def __init__(self, pool: Any | None = None) -> None:
         self.segments: list[Callable[[], Any]] = []
         self._num_graphs: int = 0
         self._num_eager_breaks: int = 0
-        self._current_graph: torch.cuda.CUDAGraph | None = None
+        self._current_graph: Any | None = None
         self._capturing: bool = False
 
     # --- context manager protocol ----------------------------------------
@@ -174,11 +174,8 @@ def __exit__(self, exc_type, exc, tb) -> None:
 
     def _begin_segment(self) -> None:
         assert not self._capturing
-        g = torch.cuda.CUDAGraph()
-        if self.pool is not None:
-            g.capture_begin(pool=self.pool)
-        else:
-            g.capture_begin()
+        g = torch.accelerator.Graph(pool=self.pool, capture_error_mode="global")
+        g.capture_begin()
         self._current_graph = g
         self._capturing = True
 

vllm/compilation/cuda_graph.py

@@ -127,7 +127,7 @@ def log(self, log_fn: Callable[..., Any] = logger.info) -> None:
 @dataclasses.dataclass
 class CUDAGraphEntry:
     batch_descriptor: BatchDescriptor
-    cudagraph: torch.cuda.CUDAGraph | None = None
+    cudagraph: Any | None = None
     output: Any | None = None
 
     # for cudagraph debugging, track the input addresses
@@ -280,7 +280,9 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any | None:
                 x.data_ptr() for x in args if isinstance(x, torch.Tensor)
             ]
             entry.input_addresses = input_addresses
-            cudagraph = torch.cuda.CUDAGraph()
+            cudagraph = torch.accelerator.Graph(
+                pool=self.graph_pool, capture_error_mode="global"
+            )
 
             with ExitStack() as stack:
                 if self.cudagraph_options.gc_disable:
@@ -310,11 +312,8 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any | None:
                 get_offloader().sync_prev_onload()
 
                 # mind-exploding: carefully manage the reference and memory.
-                with torch.cuda.graph(
-                    cudagraph,
-                    pool=self.graph_pool,
-                    stream=current_stream(),
-                ):
+                with torch.cuda.stream(current_stream()):
+                    cudagraph.capture_begin()
                     # `output` is managed by pytorch's cudagraph pool
                     output = self.runnable(*args, **kwargs)
                     # Join offloader's copy stream after forward to avoid
@@ -330,6 +329,7 @@ def __call__(self, *args: Any, **kwargs: Any) -> Any | None:
                         # the output of the last graph will not be used by
                         # any other cuda graph.
                         output = weak_ref_tensors(output)
+                    cudagraph.capture_end()
 
             # here we always use weak ref for the output
             # to save memory

vllm/v1/worker/encoder_cudagraph.py

@@ -40,7 +40,7 @@ class BudgetGraphMetadata:
     token_budget: int
     max_batch_size: int  # Max number of images/videos per batch
     max_frames_per_batch: int  # Max total frames per batch (for video)
-    graph: torch.cuda.CUDAGraph
+    graph: Any
     # Buffers recorded into the CUDA graph (e.g. embeddings, sequence metadata).
     # Before replay the manager updates these in-place. By default buffers are
     # zeroed before slice-copying the actual values; model-specific padding
@@ -276,10 +276,12 @@ def _capture_budget_graph(self, token_budget: int, path: str = "default"):
             output = self.model.encoder_cudagraph_forward({**values}, path=path)
             output_buffer = torch.empty_like(output)
 
-        graph = torch.cuda.CUDAGraph()
-        with torch.inference_mode(), torch.cuda.graph(graph, pool=self.graph_pool):
+        graph = torch.accelerator.Graph(pool=self.graph_pool, capture_error_mode="global")
+        with torch.inference_mode():
+            graph.capture_begin()
             output = self.model.encoder_cudagraph_forward({**values}, path=path)
             output_buffer.copy_(output)
+            graph.capture_end()
 
         graph_set[token_budget] = BudgetGraphMetadata(
             token_budget=token_budget,

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -134,7 +134,7 @@ def __init__(
         # dispatch() is a plain dict lookup instead of a per-call bisect.
         self._lora_dispatch_map, self._max_lora_case = self._build_lora_dispatch_map()
 
-        self.graphs: dict[BatchExecutionDescriptor, torch.cuda.CUDAGraph] = {}
+        self.graphs: dict[BatchExecutionDescriptor, Any] = {}
         self.pool = current_platform.get_global_graph_pool() if cudagraph_mode else None
 
         self._graphs_captured = False
@@ -313,17 +313,20 @@ def capture(
                         assert desc not in self.graphs, (
                             f"Graph already captured for {desc}"
                         )
-                        graph = torch.cuda.CUDAGraph()
+                        graph = torch.accelerator.Graph(
+                            pool=self.pool, capture_error_mode="global"
+                        )
                         # Sync offloader's copy stream before capture.
                         # Ensure any pre-capture prefetches from offloader are complete.
                         get_offloader().sync_prev_onload()
-                        with torch.cuda.graph(graph, self.pool):
-                            forward_fn(CUDAGraphMode.NONE)
-                            # Join offloader's copy stream after forward to avoid
-                            # unjoined stream error. The last layer's start_prefetch
-                            # forks copy_stream, but wait_prefetch only happens in
-                            # the next forward pass.
-                            get_offloader().join_after_forward()
+                        graph.capture_begin()
+                        forward_fn(CUDAGraphMode.NONE)
+                        # Join offloader's copy stream after forward to avoid
+                        # unjoined stream error. The last layer's start_prefetch
+                        # forks copy_stream, but wait_prefetch only happens in
+                        # the next forward pass.
+                        get_offloader().join_after_forward()
+                        graph.capture_end()
                         self.graphs[desc] = graph
                         compilation_counter.num_cudagraph_captured += 1
         self._graphs_captured = True

vllm/v1/worker/gpu_ubatch_wrapper.py

@@ -60,7 +60,7 @@ class UbatchMetadata:
 
 @dataclass
 class CUDAGraphMetaData:
-    cudagraph: torch.cuda.CUDAGraph
+    cudagraph: Any
     ubatch_metadata: UbatchMetadata
     outputs: Any | None = None
 
@@ -262,7 +262,9 @@ def _capture_ubatch_thread(results, ubatch_metadata):
 
             # Capture the cudagraph
             cudagraph_metadata = CUDAGraphMetaData(
-                cudagraph=torch.cuda.CUDAGraph(),
+                cudagraph=torch.accelerator.Graph(
+                    pool=self.graph_pool, capture_error_mode="global"
+                ),
                 ubatch_metadata=ubatch_metadata,
             )
             if self.graph_pool is not None:
@@ -274,11 +276,8 @@ def _capture_ubatch_thread(results, ubatch_metadata):
             # Ensure any pre-capture prefetches from offloader are complete.
             get_offloader().sync_prev_onload()
 
-            with torch.cuda.graph(
-                cudagraph_metadata.cudagraph,
-                stream=compute_stream,
-                pool=self.graph_pool,
-            ):
+            with torch.cuda.stream(compute_stream):
+                cudagraph_metadata.cudagraph.capture_begin()
                 ubatch_metadata[0].context.cpu_wait_event.set()
                 for thread in ubatch_threads:
                     thread.join()
@@ -289,6 +288,7 @@ def _capture_ubatch_thread(results, ubatch_metadata):
                 # stream error. The last layer's start_prefetch forks copy_stream,
                 # but wait_prefetch only happens in the next forward pass.
                 get_offloader().join_after_forward()
+                cudagraph_metadata.cudagraph.capture_end()
             self.cudagraphs[num_tokens] = cudagraph_metadata
         return cudagraph_metadata.outputs
 

vllm/v1/worker/xpu_model_runner.py

@@ -5,7 +5,6 @@
 import torch
 
 from vllm.config import VllmConfig
-from vllm.utils.torch_utils import supports_xpu_graph
 from vllm.v1.worker.gpu.model_runner import (
     GPUModelRunner as GPUModelRunnerV2,
 )
@@ -48,8 +47,4 @@ def _torch_cuda_wrapper():
     torch.cuda.mem_get_info = torch.xpu.mem_get_info
     torch.cuda.Event = torch.Event
     torch.cuda.set_stream = torch.xpu.set_stream
-    if supports_xpu_graph():
-        torch.cuda.graph = torch.xpu.graph
-        torch.cuda.CUDAGraph = torch.xpu.XPUGraph
-        torch.cuda.graph_pool_handle = torch.xpu.graph_pool_handle
     yield