[Quantization][XPU] Enable int8_w8a8 MoE on the Triton backend for XPU (#53162)

Signed-off-by: Artur Fierka <artur.fierka@intel.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>

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."""

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: