[UT][XPU] fix b12x UT (#52265)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com> Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Qiming Zhang committed
Aug 14, 2026 at 06:49 UTC
1f7427bc0ad31fdad1b0933c77ea5acae15da2e0
2 files changed
+24
-3
tests/model_executor/kernels/test_b12x_mxfp8_linear.py
+23
-3
@@ -46,7 +46,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
46
kFp8Static128BlockSym,
47
kFp8StaticTensorSym,
48
)
49
-from vllm.platforms import PlatformEnum
49
+from vllm.platforms import PlatformEnum, current_platform
50
from vllm.utils.b12x import b12x_warmup_token_counts
51
52
@@ -100,6 +100,7 @@ def test_b12x_backend_does_not_intercept_unquantized_bf16(
100
dist_init,
101
) -> None:
102
default_vllm_config.kernel_config.linear_backend = "b12x"
103
+ device = current_platform.device_type
104
layer = ReplicatedLinear(
105
128,
106
64,
@@ -107,9 +108,9 @@ def test_b12x_backend_does_not_intercept_unquantized_bf16(
108
params_dtype=torch.bfloat16,
109
quant_config=None,
110
prefix="bf16_linear",
110
- ).cuda()
111
+ ).to(device)
112
layer.weight.data.normal_()
112
- x = torch.randn(8, 128, device="cuda", dtype=torch.bfloat16)
113
+ x = torch.randn(8, 128, device=device, dtype=torch.bfloat16)
114
115
output, output_bias = layer(x)
116
expected = torch.nn.functional.linear(x, layer.weight)
@@ -767,6 +768,24 @@ def test_b12x_mxfp8_reload_reuses_packed_tensor_addresses(monkeypatch) -> None:
768
assert layer.weight_scale.numel() == 0
769
770
771
+@pytest.fixture
772
+def _mock_b12x_cuda_fp8_platform(monkeypatch: pytest.MonkeyPatch) -> None:
773
+ import vllm.model_executor.layers.quantization.utils.fp8_utils as fp8_utils
774
+
775
+ monkeypatch.setattr(
776
+ fp8_utils,
777
+ "current_platform",
778
+ types.SimpleNamespace(
779
+ is_fp8_fnuz=lambda: False,
780
+ is_rocm=lambda: False,
781
+ fp8_dtype=lambda: torch.float8_e4m3fn,
782
+ is_xpu=lambda: False,
783
+ is_cuda_alike=lambda: True,
784
+ ),
785
+ )
786
+
787
+
788
+@pytest.mark.usefixtures("_mock_b12x_cuda_fp8_platform")
789
def test_b12x_block_fp8_process_weights_keeps_native_block_layout() -> None:
790
layer = torch.nn.Module()
791
layer.weight = torch.nn.Parameter(
@@ -796,6 +815,7 @@ def test_b12x_block_fp8_process_weights_keeps_native_block_layout() -> None:
815
816
817
@pytest.mark.parametrize("scale_dtype", [torch.float8_e8m0fnu, torch.uint8])
818
+@pytest.mark.usefixtures("_mock_b12x_cuda_fp8_platform")
819
def test_b12x_block_fp8_upcasts_e8m0_weight_scales(scale_dtype) -> None:
820
layer = torch.nn.Module()
821
layer.weight = torch.nn.Parameter(
vllm/utils/flashinfer.py
+1
@@ -126,6 +126,7 @@ flashinfer_cutedsl_grouped_gemm_nt_masked = _lazy_import_wrapper(
126
"flashinfer.cute_dsl.blockscaled_gemm", "grouped_gemm_nt_masked"
127
)
128
flashinfer_fp4_quantize = _lazy_import_wrapper("flashinfer", "fp4_quantize")
129
+flashinfer_mxfp4_quantize = _lazy_import_wrapper("flashinfer", "mxfp4_quantize")
130
nvfp4_batched_quantize = _lazy_import_wrapper("flashinfer", "nvfp4_batched_quantize")
131
silu_and_mul_scaled_nvfp4_experts_quantize = _lazy_import_wrapper(
132
"flashinfer", "silu_and_mul_scaled_nvfp4_experts_quantize"