iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[XPU][CI] Enable XPU graph in graph tests (#60784)

Browse Source 
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
This commit is contained in:
liuzhenwei
2026-10-12 10:39:40 +08:00
committed by GitHub
co-authored by Kunshang Ji
parent f147949be8
commit bfa7e5b7af
4 changed files with 63 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/cudagraph/test_breakable_cudagraph.py
+29 -17
View File
@@ -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
tests/v1/cudagraph/test_cudagraph_dispatch.py
+12 -3
View File
@@ -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(
tests/v1/cudagraph/test_cudagraph_mode.py
+7
View File
@@ -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),
tests/v1/cudagraph/test_encoder_cudagraph.py
+15 -12
View File
@@ -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(