From f651a8a9a444e5ef7eb88003d45ab908b8d4dc25 Mon Sep 17 00:00:00 2001 From: Yejing Lai Date: Wed, 1 Jul 2026 15:38:03 +0800 Subject: [PATCH] [XPU][UT]Enable ut qk_norm_rope_fusion (#42486) Signed-off-by: Lai, Yejing Co-authored-by: Kunshang Ji --- tests/compile/passes/test_qk_norm_rope_fusion.py | 4 ++-- vllm/compilation/passes/vllm_inductor_pass.py | 13 ++++++++----- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/tests/compile/passes/test_qk_norm_rope_fusion.py b/tests/compile/passes/test_qk_norm_rope_fusion.py index 25b8ea56fe..def025ad39 100644 --- a/tests/compile/passes/test_qk_norm_rope_fusion.py +++ b/tests/compile/passes/test_qk_norm_rope_fusion.py @@ -122,7 +122,7 @@ class QKNormRoPETestModel(torch.nn.Module): @pytest.mark.parametrize("enable_rope_custom_op", [True]) @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) @pytest.mark.skipif( - not current_platform.is_cuda_alike(), + not (current_platform.is_cuda_alike() or current_platform.is_xpu()), reason="Only test on cuda and rocm platform", ) def test_qk_norm_rope_fusion( @@ -136,7 +136,7 @@ def test_qk_norm_rope_fusion( if not hasattr(torch.ops._C, "fused_qk_norm_rope"): pytest.skip("fused_qk_norm_rope custom op not available") - torch.set_default_device("cuda") + torch.set_default_device(current_platform.device_type) torch.set_default_dtype(dtype) torch.manual_seed(0) diff --git a/vllm/compilation/passes/vllm_inductor_pass.py b/vllm/compilation/passes/vllm_inductor_pass.py index 4f90b2a27e..95a9e0b37f 100644 --- a/vllm/compilation/passes/vllm_inductor_pass.py +++ b/vllm/compilation/passes/vllm_inductor_pass.py @@ -18,12 +18,15 @@ from torch._inductor.pattern_matcher import PatternMatcherPass, PatternPrettyPri from vllm.config import VllmConfig from vllm.logger import init_logger +from vllm.platforms import current_platform from .fx_utils import is_func from .inductor_pass import InductorPass, enable_fake_mode logger = init_logger(__name__) +DEVICE_TYPE = current_platform.device_type + @dataclass class InductorCompilationConfig: @@ -227,23 +230,23 @@ class VllmPatternReplacement(ABC, Generic[P, R]): # Helpers for get_inputs: uninitialized tensors of common dtypes. @staticmethod def empty(*args, **kwargs) -> torch.Tensor: - return torch.empty(*args, device="cuda", **kwargs) + return torch.empty(*args, device=DEVICE_TYPE, **kwargs) @staticmethod def empty_bf16(*args, **kwargs) -> torch.Tensor: - return torch.empty(*args, dtype=torch.bfloat16, device="cuda", **kwargs) + return torch.empty(*args, dtype=torch.bfloat16, device=DEVICE_TYPE, **kwargs) @staticmethod def empty_fp16(*args, **kwargs) -> torch.Tensor: - return torch.empty(*args, dtype=torch.float16, device="cuda", **kwargs) + return torch.empty(*args, dtype=torch.float16, device=DEVICE_TYPE, **kwargs) @staticmethod def empty_fp32(*args, **kwargs) -> torch.Tensor: - return torch.empty(*args, dtype=torch.float32, device="cuda", **kwargs) + return torch.empty(*args, dtype=torch.float32, device=DEVICE_TYPE, **kwargs) @staticmethod def empty_i32(*args, **kwargs) -> torch.Tensor: - return torch.empty(*args, dtype=torch.int32, device="cuda", **kwargs) + return torch.empty(*args, dtype=torch.int32, device=DEVICE_TYPE, **kwargs) def _fx_view_to_reshape(gm: fx.GraphModule) -> None: