[XPU][CI] fix jit_warmup_triton_launcher (#56670)
Signed-off-by: zhenwei-intel <zhenwei.liu@intel.com>
liuzhenwei committed
Sep 13, 2026 at 21:44 UTC
de500290ffffeb7bd834872c0dea5331590dd40d
1 file changed
+5
-3
tests/model_executor/test_jit_warmup_triton_launcher.py
+5
-3
@@ -21,6 +21,7 @@ from vllm.model_executor.warmup.jit_warmup_triton_helper import (
21
triton_kernel_dispatcher_with_warmup,
22
triton_warmup_inputs,
23
)
24
+from vllm.platforms import current_platform
25
from vllm.triton_utils import tl, triton
26
27
@@ -257,10 +258,11 @@ def test_triton_kernel_decorator_propagates_dispatch_assertions(
258
dispatch.get_warmup_keys()
259
260
260
-def test_triton_kernel_dispatch_uses_cuda_fake_tensors(
261
+def test_triton_kernel_dispatch_uses_device_fake_tensors(
262
monkeypatch: pytest.MonkeyPatch,
263
) -> None:
264
kernel = _FakeTritonKernel()
265
+ device_type = current_platform.device_type
266
267
def warmup_inputs() -> dict[str, Any]:
268
return dict(
@@ -271,7 +273,7 @@ def test_triton_kernel_dispatch_uses_cuda_fake_tensors(
273
@triton_kernel_dispatcher_with_warmup(kernel=kernel, warmup_inputs=warmup_inputs)
274
def dispatch(first: torch.Tensor, second: int) -> LaunchSpec:
275
assert isinstance(first, torch.Tensor)
274
- assert first.is_cuda
276
+ assert first.device.type == device_type
277
assert first[0].is_contiguous()
278
return (first.shape[0],), dict(CONST=first[0].numel())
279
@@ -282,7 +284,7 @@ def test_triton_kernel_dispatch_uses_cuda_fake_tensors(
284
285
keys = dispatch.get_warmup_keys()
286
assert len(keys) == 1
285
- assert dict(keys[0].inputs)["first"].device.type == "cuda"
287
+ assert dict(keys[0].inputs)["first"].device.type == device_type
288
289
290
def test_triton_kernel_decorates_native_launchers(