[XPU] quick fix online quantization UT break (#51365)
Signed-off-by: Yan Ma <yan.ma@intel.com>
Yan Ma committed
Aug 7, 2026 at 18:25 UTC
8d9b52f7c2514490bdadfd5eb0c931e58625df2e
1 file changed
+11
-6
tests/quantization/test_online.py
+11
-6
@@ -40,6 +40,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
40
from vllm.platforms import current_platform
41
from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
42
43
+DEVICE = current_platform.device_type
44
+
45
46
@pytest.mark.skipif(
47
not is_quant_method_supported("fp8"),
@@ -260,7 +262,7 @@ def test_online_linear_tp_weight_quant_matches_unsharded(
262
) -> None:
263
"""TP shards pack the same FP8 values and scales as the unsharded weight."""
264
torch.manual_seed(0)
263
- weight = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16)
265
+ weight = torch.randn(32, 64, device=DEVICE, dtype=torch.bfloat16)
266
weight[-1, -1] = 64.0
267
268
full_weight, full_scale = _quantize_linear(weight, scheme, False)
@@ -306,17 +308,20 @@ def _quantize_moe(weight, scheme, moe_tp_size):
308
def test_online_moe_tp_weight_quant_matches_ep(monkeypatch, scheme: str) -> None:
309
"""TP shards of w2 pack the same values and scales as full experts."""
310
if scheme == "nvfp4":
309
- if not (
310
- current_platform.is_cuda()
311
- and current_platform.is_device_capability_family(100)
311
+ if (
312
+ not (
313
+ current_platform.is_cuda()
314
+ and current_platform.is_device_capability_family(100)
315
+ )
316
+ or current_platform.is_xpu()
317
):
318
pytest.skip("NVFP4 weight quantization needs a Blackwell (SM100) GPU.")
319
elif not is_quant_method_supported("fp8"):
320
pytest.skip("FP8 is not supported on this GPU type.")
321
322
torch.manual_seed(0)
318
- weight = torch.randn(2, 32, 32, device="cuda", dtype=torch.bfloat16)
319
- weight[:, -1, -1] = torch.tensor([32.0, 64.0], device="cuda")
323
+ weight = torch.randn(2, 32, 32, device=DEVICE, dtype=torch.bfloat16)
324
+ weight[:, -1, -1] = torch.tensor([32.0, 64.0], device=DEVICE)
325
326
ep_out = _quantize_moe(weight, scheme, 1)
327