[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>
Artur Fierka committed
Sep 18, 2026 at 16:38 UTC
1cdf1689e55c483a754c6c24baac3e47d39d4ece
2 files changed
+58
-18
tests/quantization/test_int8_moe_oracle.py
+47
-15
@@ -2,9 +2,9 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Tests for INT8 (W8A8) fused-MoE oracle backend selection.
4
5
-These exercise ``select_int8_moe_backend`` only (no kernels are launched), so
6
-they run on any platform where the Triton INT8 MoE kernel is available — CUDA
7
-(SM >= 7.5) or ROCm — not just gfx950.
5
+These exercise ``select_int8_moe_backend`` only (no MoE kernels are launched),
6
+so they run on any platform where the Triton INT8 MoE kernel is available —
7
+CUDA (SM >= 7.5), ROCm, or XPU — not just gfx950.
8
"""
9
10
import pytest
@@ -27,17 +27,22 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
27
)
28
from vllm.platforms import current_platform
29
30
-# The Triton int8_w8a8 fused-MoE kernel is available on CUDA (Turing+) and on
31
-# ROCm CDNA GPUs. Gate on that rather than on a specific arch.
30
+# The Triton int8_w8a8 fused-MoE kernel is available on CUDA (Turing+), on
31
+# ROCm CDNA GPUs and on XPU. Gate on that rather than on a specific arch.
32
INT8_MOE_SUPPORTED = (
33
- current_platform.is_cuda() and current_platform.has_device_capability((7, 5))
34
-) or current_platform.is_rocm()
33
+ (current_platform.is_cuda() and current_platform.has_device_capability((7, 5)))
34
+ or current_platform.is_rocm()
35
+ or current_platform.is_xpu()
36
+)
37
38
requires_int8_moe = pytest.mark.skipif(
39
not INT8_MOE_SUPPORTED,
38
- reason="Requires a GPU with Triton INT8 MoE support (CUDA SM>=7.5 or ROCm)",
40
+ reason="Requires Triton INT8 MoE support (CUDA SM>=7.5, ROCm, or XPU)",
41
)
42
43
+# So FusedMoEConfig.device names the device the test actually runs on.
44
+DEVICE = current_platform.device_type
45
+
46
47
def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:
48
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
@@ -52,7 +57,7 @@ def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:
57
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
58
activation=MoEActivation.SILU,
59
in_dtype=torch.bfloat16,
55
- device="cuda",
60
+ device=DEVICE,
61
routing_method=RoutingMethodType.Renormalize,
62
moe_backend=moe_backend,
63
)
@@ -62,15 +67,20 @@ def _make_int8_moe_config(moe_backend: str = "auto") -> FusedMoEConfig:
67
@pytest.mark.parametrize(
68
"weight_key,activation_key",
69
[
65
- # per-channel weight + dynamic per-token activation
66
- (kInt8StaticChannelSym, kInt8DynamicTokenSym),
67
- # per-tensor weight + dynamic per-tensor activation
68
- (kInt8StaticTensorSym, kInt8DynamicTensorSym),
70
+ pytest.param(
71
+ kInt8StaticChannelSym,
72
+ kInt8DynamicTokenSym,
73
+ id="per_channel_weight-per_token_act",
74
+ ),
75
+ pytest.param(
76
+ kInt8StaticTensorSym,
77
+ kInt8DynamicTensorSym,
78
+ id="per_tensor_weight-per_tensor_act",
79
+ ),
80
],
81
)
82
def test_int8_dynamic_schemes_dispatch_to_triton(weight_key, activation_key):
72
- """Both dynamic-activation INT8 MoE schemes (per-channel + per-tensor
73
- weights) select the Triton backend."""
83
+ """Both dynamic-activation INT8 MoE schemes select the Triton backend."""
84
config = _make_int8_moe_config()
85
backend, experts_cls = select_int8_moe_backend(
86
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):
89
assert experts_cls is not None
90
91
92
+@pytest.mark.skipif(not current_platform.is_xpu(), reason="XPU-only behaviour")
93
+def test_scaled_int8_quant_is_available_on_xpu():
94
+ """The per-tensor scheme's activation quantization works on XPU.
95
+
96
+ ``_int8_quantize`` sends per-tensor activations through
97
+ ``ops.scaled_int8_quant``, which has an XPU branch. This is what lets the
98
+ per-tensor scheme be enabled here alongside CUDA and ROCm; pin it so a
99
+ future change cannot silently remove that branch.
100
+ """
101
+ from vllm import _custom_ops as ops
102
+
103
+ x = torch.randn(4, 16, device=DEVICE, dtype=torch.bfloat16)
104
+ scale = torch.full((1,), 0.05, device=DEVICE, dtype=torch.float32)
105
+
106
+ q, returned_scale, azp = ops.scaled_int8_quant(x, scale=scale)
107
+
108
+ assert q.dtype == torch.int8
109
+ assert q.shape == x.shape
110
+ assert azp is None
111
+ torch.testing.assert_close(returned_scale, scale)
112
+
113
+
114
@requires_int8_moe
115
def test_int8_explicit_moe_backend_triton():
116
"""An explicit --moe-backend triton selects the Triton INT8 backend."""
vllm/model_executor/layers/fused_moe/experts/triton_moe.py
+11
-3
@@ -130,10 +130,18 @@ class TritonExperts(LoRAExpertsMixin, mk.FusedMoEExpertsModular):
130
# INT8 requires at least 7.5 (Turing) on CUDA. ROCm CDNA GPUs
131
# (e.g. MI2xx/MI3xx/gfx950) provide native INT8 matrix-core support and
132
# the Triton int8_w8a8 fused MoE kernel handles them.
133
+ # XPU reaches the same Triton kernel through the same launcher, and both
134
+ # int8 activation-quant paths work there: per-token uses the Triton
135
+ # `per_token_quant_int8`, per-tensor resolves to the XPU branch of
136
+ # `scaled_int8_quant`, which is plain elementwise arithmetic.
137
device_supports_int8 = (
134
- current_platform.is_cuda()
135
- and current_platform.has_device_capability((7, 5))
136
- ) or current_platform.is_rocm()
138
+ (
139
+ current_platform.is_cuda()
140
+ and current_platform.has_device_capability((7, 5))
141
+ )
142
+ or current_platform.is_rocm()
143
+ or current_platform.is_xpu()
144
+ )
145
146
supported: list[tuple[QuantKey | None, QuantKey | None]] = [(None, None)]
147
if device_supports_int8: