diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index c10835821f58..717f1536d033 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/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" diff --git a/vllm/compilation/breakable_cudagraph.py b/vllm/compilation/breakable_cudagraph.py index 6da3ec717861..f69788fdb43e 100644 --- a/vllm/compilation/breakable_cudagraph.py +++ b/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 diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py index b63d86199720..36b5870ed838 100644 --- a/vllm/compilation/cuda_graph.py +++ b/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 diff --git a/vllm/v1/worker/encoder_cudagraph.py b/vllm/v1/worker/encoder_cudagraph.py index d3c54812894f..ea130abe902a 100644 --- a/vllm/v1/worker/encoder_cudagraph.py +++ b/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, diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index dad1777b47e2..55f08a0aa769 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/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 diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index 657fc8267345..aa8a02d45483 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/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 diff --git a/vllm/v1/worker/xpu_model_runner.py b/vllm/v1/worker/xpu_model_runner.py index f93d04395571..9f0fe34ca41a 100644 --- a/vllm/v1/worker/xpu_model_runner.py +++ b/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