[XPU][UT]Enable ut qk_norm_rope_fusion (#42486)

Signed-off-by: Lai, Yejing <yejing.lai@intel.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>

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)
 

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: