diff --git a/tests/quantization/test_int8_moe_oracle.py b/tests/quantization/test_int8_moe_oracle.py index 8c838a785b..9580eb9ab3 100644 --- a/tests/quantization/test_int8_moe_oracle.py +++ b/tests/quantization/test_int8_moe_oracle.py @@ -2,9 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for INT8 (W8A8) fused-MoE oracle backend selection. -These exercise ``select_int8_moe_backend`` only (no kernels are launched), so -they run on any platform where the Triton INT8 MoE kernel is available — CUDA -(SM >= 7.5) or ROCm — not just gfx950. +These exercise ``select_int8_moe_backend`` only (no MoE kernels are launched), +so they run on any platform where the Triton INT8 MoE kernel is available — +CUDA (SM >= 7.5), ROCm, or XPU — not just gfx950. """ import pytest @@ -27,17 +27,22 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.platforms import current_platform -# The Triton int8_w8a8 fused-MoE kernel is available on CUDA (Turing+) and on -# ROCm CDNA GPUs. Gate on that rather than on a specific arch. +# The Triton int8_w8a8 fused-MoE kernel is available on CUDA (Turing+), on +# ROCm CDNA GPUs and on XPU. Gate on that rather than on a specific arch. INT8_MOE_SUPPORTED = ( - current_platform.is_cuda() and current_platform.has_device_capability((7, 5)) -) or current_platform.is_rocm() + (current_platform.is_cuda() and current_platform.has_device_capability((7, 5))) + or current_platform.is_rocm() + or current_platform.is_xpu() +) requires_int8_moe = pytest.mark.skipif( not INT8_MOE_SUPPORTED, - reason="Requires a GPU with Triton INT8 MoE support (CUDA SM>=7.5 or ROCm)", + reason="Requires Triton INT8 MoE support (CUDA SM>=7.5, ROCm, or XPU)", ) +# So FusedMoEConfig.device names the device the test actually runs on. +DEVICE = current_platform.device_type + def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig: from vllm.model_executor.layers.fused_moe.activation import MoEActivation @@ -52,7 +57,7 @@ def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig: moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), activation=MoEActivation.SILU, in_dtype=torch.bfloat16, - device="cuda", + device=DEVICE, routing_method=RoutingMethodType.Renormalize, moe_backend=moe_backend, ) @@ -62,15 +67,20 @@ def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig: @pytest.mark.parametrize( "weight_key,activation_key", [ - # per-channel weight + dynamic per-token activation - (kInt8StaticChannelSym, kInt8DynamicTokenSym), - # per-tensor weight + dynamic per-tensor activation - (kInt8StaticTensorSym, kInt8DynamicTensorSym), + pytest.param( + kInt8StaticChannelSym, + kInt8DynamicTokenSym, + id="per_channel_weight-per_token_act", + ), + pytest.param( + kInt8StaticTensorSym, + kInt8DynamicTensorSym, + id="per_tensor_weight-per_tensor_act", + ), ], ) def test_int8_dynamic_schemes_dispatch_to_triton(weight_key, activation_key): - """Both dynamic-activation INT8 MoE schemes (per-channel + per-tensor - weights) select the Triton backend.""" + """Both dynamic-activation INT8 MoE schemes select the Triton backend.""" config = _make_int8_moe_config() backend, experts_cls = select_int8_moe_backend( config, weight_key=weight_key, activation_key=activation_key @@ -79,6 +89,28 @@ def test_int8_dynamic_schemes_dispatch_to_triton(weight_key, activation_key): assert experts_cls is not None +@pytest.mark.skipif(not current_platform.is_xpu(), reason="XPU-only behaviour") +def test_scaled_int8_quant_is_available_on_xpu(): + """The per-tensor scheme's activation quantization works on XPU. + + ``_int8_quantize`` sends per-tensor activations through + ``ops.scaled_int8_quant``, which has an XPU branch. This is what lets the + per-tensor scheme be enabled here alongside CUDA and ROCm; pin it so a + future change cannot silently remove that branch. + """ + from vllm import _custom_ops as ops + + x = torch.randn(4, 16, device=DEVICE, dtype=torch.bfloat16) + scale = torch.full((1,), 0.05, device=DEVICE, dtype=torch.float32) + + q, returned_scale, azp = ops.scaled_int8_quant(x, scale=scale) + + assert q.dtype == torch.int8 + assert q.shape == x.shape + assert azp is None + torch.testing.assert_close(returned_scale, scale) + + @requires_int8_moe def test_int8_explicit_moe_backend_triton(): """An explicit --moe-backend triton selects the Triton INT8 backend.""" diff --git a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py index da9666a66b..700d113921 100644 --- a/vllm/model_executor/layers/fused_moe/experts/triton_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/triton_moe.py @@ -130,10 +130,18 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular): # INT8 requires at least 7.5 (Turing) on CUDA. ROCm CDNA GPUs # (e.g. MI2xx/MI3xx/gfx950) provide native INT8 matrix-core support and # the Triton int8_w8a8 fused MoE kernel handles them. + # XPU reaches the same Triton kernel through the same launcher, and both + # int8 activation-quant paths work there: per-token uses the Triton + # `per_token_quant_int8`, per-tensor resolves to the XPU branch of + # `scaled_int8_quant`, which is plain elementwise arithmetic. device_supports_int8 = ( - current_platform.is_cuda() - and current_platform.has_device_capability((7, 5)) - ) or current_platform.is_rocm() + ( + current_platform.is_cuda() + and current_platform.has_device_capability((7, 5)) + ) + or current_platform.is_rocm() + or current_platform.is_xpu() + ) supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)] if device_supports_int8: