[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>
Yejing Lai committed
Jul 1, 2026 at 15:38 UTC
f651a8a9a444e5ef7eb88003d45ab908b8d4dc25
2 files changed
+10
-7
tests/compile/passes/test_qk_norm_rope_fusion.py
+2
-2
@@ -122,7 +122,7 @@ class QKNormRoPETestModel(torch.nn.Module):
122
@pytest.mark.parametrize("enable_rope_custom_op", [True])
123
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
124
@pytest.mark.skipif(
125
- not current_platform.is_cuda_alike(),
125
+ not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
126
reason="Only test on cuda and rocm platform",
127
)
128
def test_qk_norm_rope_fusion(
@@ -136,7 +136,7 @@ def test_qk_norm_rope_fusion(
136
if not hasattr(torch.ops._C, "fused_qk_norm_rope"):
137
pytest.skip("fused_qk_norm_rope custom op not available")
138
139
- torch.set_default_device("cuda")
139
+ torch.set_default_device(current_platform.device_type)
140
torch.set_default_dtype(dtype)
141
torch.manual_seed(0)
142
vllm/compilation/passes/vllm_inductor_pass.py
+8
-5
@@ -18,12 +18,15 @@ from torch._inductor.pattern_matcher import PatternMatcherPass, PatternPrettyPri
18
19
from vllm.config import VllmConfig
20
from vllm.logger import init_logger
21
+from vllm.platforms import current_platform
22
23
from .fx_utils import is_func
24
from .inductor_pass import InductorPass, enable_fake_mode
25
26
logger = init_logger(__name__)
27
28
+DEVICE_TYPE = current_platform.device_type
29
+
30
31
@dataclass
32
class InductorCompilationConfig:
@@ -227,23 +230,23 @@ class VllmPatternReplacement(ABC, Generic[P, R]):
230
# Helpers for get_inputs: uninitialized tensors of common dtypes.
231
@staticmethod
232
def empty(*args, **kwargs) -> torch.Tensor:
230
- return torch.empty(*args, device="cuda", **kwargs)
233
+ return torch.empty(*args, device=DEVICE_TYPE, **kwargs)
234
235
@staticmethod
236
def empty_bf16(*args, **kwargs) -> torch.Tensor:
234
- return torch.empty(*args, dtype=torch.bfloat16, device="cuda", **kwargs)
237
+ return torch.empty(*args, dtype=torch.bfloat16, device=DEVICE_TYPE, **kwargs)
238
239
@staticmethod
240
def empty_fp16(*args, **kwargs) -> torch.Tensor:
238
- return torch.empty(*args, dtype=torch.float16, device="cuda", **kwargs)
241
+ return torch.empty(*args, dtype=torch.float16, device=DEVICE_TYPE, **kwargs)
242
243
@staticmethod
244
def empty_fp32(*args, **kwargs) -> torch.Tensor:
242
- return torch.empty(*args, dtype=torch.float32, device="cuda", **kwargs)
245
+ return torch.empty(*args, dtype=torch.float32, device=DEVICE_TYPE, **kwargs)
246
247
@staticmethod
248
def empty_i32(*args, **kwargs) -> torch.Tensor:
246
- return torch.empty(*args, dtype=torch.int32, device="cuda", **kwargs)
249
+ return torch.empty(*args, dtype=torch.int32, device=DEVICE_TYPE, **kwargs)
250
251
252
def _fx_view_to_reshape(gm: fx.GraphModule) -> None: