From bfa7e5b7afe40ebe48288801a2baf5c1135cd881 Mon Sep 17 00:00:00 2001 From: liuzhenwei Date: Mon, 12 Oct 2026 10:39:40 +0800 Subject: [PATCH] [XPU][CI] Enable XPU graph in graph tests (#60784) Signed-off-by: zhenwei-intel Co-authored-by: Kunshang Ji --- .../v1/cudagraph/test_breakable_cudagraph.py | 46 ++++++++++++------- tests/v1/cudagraph/test_cudagraph_dispatch.py | 15 ++++-- tests/v1/cudagraph/test_cudagraph_mode.py | 7 +++ tests/v1/cudagraph/test_encoder_cudagraph.py | 27 ++++++----- 4 files changed, 63 insertions(+), 32 deletions(-) diff --git a/tests/v1/cudagraph/test_breakable_cudagraph.py b/tests/v1/cudagraph/test_breakable_cudagraph.py index b0b49a9567..1b9a79a6ef 100644 --- a/tests/v1/cudagraph/test_breakable_cudagraph.py +++ b/tests/v1/cudagraph/test_breakable_cudagraph.py @@ -11,6 +11,18 @@ from unittest.mock import MagicMock, patch import pytest import torch +from vllm.platforms import current_platform + +DEVICE_TYPE = current_platform.device_type +HAS_DEVICE_GRAPH = current_platform.is_cuda_alike() or current_platform.is_xpu() + +if current_platform.is_xpu(): + from vllm.v1.worker.xpu_model_runner import _torch_cuda_wrapper + + # Route torch.cuda graph/stream APIs to torch.xpu, as XPUModelRunner does. + with _torch_cuda_wrapper(): + pass + @pytest.fixture(autouse=True) def _enable_breakable_cudagraph(monkeypatch: pytest.MonkeyPatch): @@ -96,8 +108,8 @@ def cuda_capture_stream(): stream, so all capture-using tests need to run under ``torch.cuda.stream(...)`` for a separate stream. """ - if not torch.cuda.is_available(): - pytest.skip("CUDA required") + if not HAS_DEVICE_GRAPH: + pytest.skip("CUDA or XPU graph required") from vllm.utils.torch_utils import _current_stream_tls prev_stream = getattr(_current_stream_tls, "value", None) @@ -200,7 +212,7 @@ def test_active_state_isolated_across_threads(cuda_capture_stream): def test_capture_with_no_eager_break_records_one_graph(cuda_capture_stream): from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture - x = torch.zeros(4, device="cuda") + x = torch.zeros(4, device=DEVICE_TYPE) cap = BreakableCUDAGraphCapture() with cap: x.add_(1.0) @@ -212,7 +224,7 @@ def test_capture_with_no_eager_break_records_one_graph(cuda_capture_stream): def test_add_eager_creates_alternating_graph_eager_graph(cuda_capture_stream): from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture - x = torch.zeros(4, device="cuda") + x = torch.zeros(4, device=DEVICE_TYPE) counter = {"eager_calls": 0} def eager_step(): @@ -253,7 +265,7 @@ def test_capture_replay_matches_eager_simple(cuda_capture_stream): """ from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture - x = torch.zeros(8, device="cuda") + x = torch.zeros(8, device=DEVICE_TYPE) log: list[str] = [] def eager_break_op(): @@ -269,21 +281,21 @@ def test_capture_replay_matches_eager_simple(cuda_capture_stream): # Capture-time: graph kernels were recorded only; eager segment ran # once on x == 0, leaving x == 0. torch.accelerator.synchronize() - assert torch.equal(x, torch.zeros(8, device="cuda")) + assert torch.equal(x, torch.zeros(8, device=DEVICE_TYPE)) assert log == ["eager"] # Replay with a fresh input: 10 -> 11 -> 22 -> 27. x.fill_(10.0) cap.replay() torch.accelerator.synchronize() - assert torch.equal(x, torch.full((8,), 27.0, device="cuda")) + assert torch.equal(x, torch.full((8,), 27.0, device=DEVICE_TYPE)) assert log == ["eager", "eager"] # Replay again with another input: 100 -> 101 -> 202 -> 207. x.fill_(100.0) cap.replay() torch.accelerator.synchronize() - assert torch.equal(x, torch.full((8,), 207.0, device="cuda")) + assert torch.equal(x, torch.full((8,), 207.0, device=DEVICE_TYPE)) assert log == ["eager", "eager", "eager"] @@ -301,13 +313,13 @@ def test_decorator_breaks_when_invoked_inside_capture(cuda_capture_stream): # In-place double; stands in for "real" attention work. t.mul_(2.0) - x = torch.zeros(4, device="cuda") + x = torch.zeros(4, device=DEVICE_TYPE) # Outside capture: decorator should just call through. x.fill_(3.0) attention_like(x) torch.accelerator.synchronize() - assert torch.equal(x, torch.full((4,), 6.0, device="cuda")) + assert torch.equal(x, torch.full((4,), 6.0, device=DEVICE_TYPE)) # Inside capture: decorator should split the graph. Only the eager # segment actually mutates state during capture. @@ -318,7 +330,7 @@ def test_decorator_breaks_when_invoked_inside_capture(cuda_capture_stream): attention_like(x) # eager: x *= 2 (on x == 0, no-op) x.add_(1.0) # recorded torch.accelerator.synchronize() - assert torch.equal(x, torch.zeros(4, device="cuda")) + assert torch.equal(x, torch.zeros(4, device=DEVICE_TYPE)) # 2 graph segments + 1 eager segment, ordered G E G; the arithmetic # equivalence check below verifies the ordering. assert len(cap.segments) == 3 @@ -329,7 +341,7 @@ def test_decorator_breaks_when_invoked_inside_capture(cuda_capture_stream): x.fill_(2.0) cap.replay() torch.accelerator.synchronize() - assert torch.equal(x, torch.full((4,), 15.0, device="cuda")) + assert torch.equal(x, torch.full((4,), 15.0, device=DEVICE_TYPE)) @pytest.mark.parametrize("keyword", [False, True]) @@ -352,7 +364,7 @@ def test_quantized_activation_replay_does_not_retain_tensors( data_out.copy_(x.data) scale_out.copy_(x.scale) - source = torch.ones((128, 128), device="cuda", dtype=torch.bfloat16) + source = torch.ones((128, 128), device=DEVICE_TYPE, dtype=torch.bfloat16) data_out = source.to(torch.float8_e4m3fn) scale_out = source[:, ::32].to(torch.uint8) output = torch.empty_like(source) @@ -391,7 +403,7 @@ def test_eager_attention_inside_multistream_overlap(cuda_capture_stream): ) from vllm.utils.multi_stream_utils import maybe_execute_in_parallel - x = torch.zeros(1024, device="cuda") + x = torch.zeros(1024, device=DEVICE_TYPE) output = torch.empty_like(x) aux_stream = torch.cuda.Stream() main_event = torch.cuda.Event() @@ -440,7 +452,7 @@ def test_replay_invokes_eager_segments_in_order(cuda_capture_stream): from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture log: list[str] = [] - x = torch.zeros(1, device="cuda") + x = torch.zeros(1, device=DEVICE_TYPE) def make_eager(name): def step(): @@ -502,7 +514,7 @@ def test_nested_decorated_op_runs_inline(cuda_capture_stream): eager_break_during_capture, ) - x = torch.zeros(4, device="cuda") + x = torch.zeros(4, device=DEVICE_TYPE) inner_calls = 0 @eager_break_during_capture @@ -533,7 +545,7 @@ def test_nested_decorated_op_runs_inline(cuda_capture_stream): cap.replay() torch.accelerator.synchronize() # 0 -> +2 -> +1 (inner) -> +10 (outer) -> +100 = 113 - assert torch.equal(x, torch.full((4,), 113.0, device="cuda")) + assert torch.equal(x, torch.full((4,), 113.0, device=DEVICE_TYPE)) assert inner_calls == 2 # replay invokes the outer's lambda again diff --git a/tests/v1/cudagraph/test_cudagraph_dispatch.py b/tests/v1/cudagraph/test_cudagraph_dispatch.py index 40874be5bd..b0a4587ddf 100644 --- a/tests/v1/cudagraph/test_cudagraph_dispatch.py +++ b/tests/v1/cudagraph/test_cudagraph_dispatch.py @@ -24,6 +24,14 @@ from vllm.platforms import current_platform from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher DEVICE_TYPE = current_platform.device_type +HAS_DEVICE_GRAPH = current_platform.is_cuda() or current_platform.is_xpu() + +if current_platform.is_xpu(): + from vllm.v1.worker.xpu_model_runner import _torch_cuda_wrapper + + # Route torch.cuda graph/stream APIs to torch.xpu, as XPUModelRunner does. + with _torch_cuda_wrapper(): + pass # Helper MLP for testing @@ -50,6 +58,7 @@ def _create_vllm_config( mock_config.parallel_config = ParallelConfig() mock_config.speculative_config = None # No speculative decoding mock_config.num_speculative_tokens = 0 + mock_config.use_cumem_cudagraph_pool = False if not lora_config: mock_config.lora_config = None else: @@ -268,7 +277,7 @@ class TestCudagraphDispatcher: assert dispatcher.get_capture_descs() == [] -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA or XPU graph") class TestCUDAGraphWrapper: def setup_method(self): self.vllm_config = _create_vllm_config(CompilationConfig()) @@ -415,7 +424,7 @@ def _run_and_monitor_call( @create_new_process_for_each_test("spawn") -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA or XPU graph") def test_capture_replay_bypass_logic(): comp_config = CompilationConfig( mode=CompilationMode.VLLM_COMPILE, @@ -484,7 +493,7 @@ def test_capture_replay_bypass_logic(): @create_new_process_for_each_test("spawn") -@pytest.mark.skipif(not current_platform.is_cuda(), reason="Skip if not cuda") +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA or XPU graph") def test_nested_wrappers(): """Tests a scenario with a PIECEWISE wrapper inside a FULL one.""" comp_config = CompilationConfig( diff --git a/tests/v1/cudagraph/test_cudagraph_mode.py b/tests/v1/cudagraph/test_cudagraph_mode.py index 0147f90df5..f35ef4d829 100644 --- a/tests/v1/cudagraph/test_cudagraph_mode.py +++ b/tests/v1/cudagraph/test_cudagraph_mode.py @@ -19,6 +19,13 @@ if current_platform.is_rocm(): ("TritonAttn", "FULL", True), ("TritonAttn", "FULL_AND_PIECEWISE", True), ] +elif current_platform.is_xpu(): + combo_cases_1 = [ + ("FA2", "FULL", True), + ("FA2", "FULL_AND_PIECEWISE", True), + ("TritonAttn", "FULL", True), + ("TritonAttn", "FULL_AND_PIECEWISE", True), + ] else: combo_cases_1 = [ ("FA3", "FULL", True), diff --git a/tests/v1/cudagraph/test_encoder_cudagraph.py b/tests/v1/cudagraph/test_encoder_cudagraph.py index 3489443878..2799a3ef4d 100644 --- a/tests/v1/cudagraph/test_encoder_cudagraph.py +++ b/tests/v1/cudagraph/test_encoder_cudagraph.py @@ -32,6 +32,15 @@ from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderItemSpec, ) +HAS_DEVICE_GRAPH = current_platform.is_cuda_alike() or current_platform.is_xpu() + +if current_platform.is_xpu(): + from vllm.v1.worker.xpu_model_runner import _torch_cuda_wrapper + + # Route torch.cuda graph/stream APIs to torch.xpu, as XPUModelRunner does. + with _torch_cuda_wrapper(): + pass + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- @@ -662,12 +671,10 @@ def _make_video_mm_kwargs( # --------------------------------------------------------------------------- -@pytest.mark.skipif( - not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" -) +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA, ROCm or XPU graph") class TestEncoderCudaGraphCaptureReplay: def setup_method(self): - self.device = torch.device("cuda:0") + self.device = torch.device(current_platform.device_type, 0) self.dtype = torch.float16 self.model = SimpleMockViTModel().to(self.device).half() self.mgr = _make_manager_for_gpu( @@ -758,9 +765,7 @@ class TestEncoderCudaGraphCaptureReplay: # --------------------------------------------------------------------------- -@pytest.mark.skipif( - not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" -) +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA, ROCm or XPU graph") @pytest.mark.usefixtures("dist_init", "workspace_init") @pytest.mark.parametrize("profile_only", [False, True]) @torch.inference_mode() @@ -776,7 +781,7 @@ def test_eonly_capture_preserves_outputs_across_replay_and_fallback(profile_only from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner from vllm.v1.worker.workspace import lock_workspace - device = torch.device("cuda:0") + device = torch.device(current_platform.device_type, 0) dtype = torch.float16 model = SimpleMockViTModel().to(device).half() manager = _make_manager_for_gpu(model, _BUDGETS, _MAX_BATCH, device, dtype) @@ -1047,12 +1052,10 @@ _VIDEO_MAX_BATCH = 4 _VIDEO_MAX_FRAMES = 8 # 2 frames per item at max_batch_size=4 -@pytest.mark.skipif( - not current_platform.is_cuda_alike(), reason="Skip if not cuda or rocm" -) +@pytest.mark.skipif(not HAS_DEVICE_GRAPH, reason="Requires CUDA, ROCm or XPU graph") class TestEncoderCudaGraphVideoReplay: def setup_method(self): - self.device = torch.device("cuda:0") + self.device = torch.device(current_platform.device_type, 0) self.dtype = torch.float16 self.model = SimpleMockViTVideoModel().to(self.device).half() self.mgr = _make_manager_for_gpu(