[Misc] Enable test_silu_mul_fp8_quant_deep_gemm on XPU (#49444)
Signed-off-by: pmanczak <pawel.manczak@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
pmanczak committed
Aug 11, 2026 at 17:19 UTC
5b5eae23944e94f618895c05a69331bb06d1e3bd
2 files changed
+19
-17
tests/kernels/moe/test_silu_mul_fp8_quant_deep_gemm.py
+16
-11
@@ -23,6 +23,12 @@ from vllm.utils.math_utils import cdiv, round_up
23
from vllm.utils.torch_utils import set_random_seed
24
25
fp8_dtype = current_platform.fp8_dtype()
26
+DEVICE = current_platform.device_type
27
+
28
+pytestmark = pytest.mark.skipif(
29
+ not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
30
+ reason="persistent_masked_m_silu_mul_quant requires a CUDA-alike or XPU device.",
31
+)
32
33
CASES = [
34
(1, 1, 128, fp8_dtype),
@@ -144,7 +150,7 @@ def pack_scales(x: torch.Tensor, tokens_per_expert: torch.Tensor) -> torch.Tenso
150
151
# Add i32_padding here so we can view it as a i32 tensor later on.
152
i32_padding = round_up(G, 4) - G
147
- ref_s_i8 = torch.empty((E, T, G + i32_padding), dtype=torch.uint8, device="cuda")
153
+ ref_s_i8 = torch.empty((E, T, G + i32_padding), dtype=torch.uint8, device=DEVICE)
154
for e in range(E):
155
nt = tokens_per_expert[e].item()
156
ref_s_i8[e, :nt, :G] = x[e, :nt].view(torch.int32) >> 23
@@ -177,9 +183,9 @@ def ref_with_scale_fmt(
183
DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0,
184
]
185
180
- ref_q = torch.empty((E, T, H), dtype=fp8_dtype, device="cuda")
186
+ ref_q = torch.empty((E, T, H), dtype=fp8_dtype, device=DEVICE)
187
ref_s_f32 = torch.empty(
182
- (E, T, cdiv(H, group_size)), dtype=torch.float32, device="cuda"
188
+ (E, T, cdiv(H, group_size)), dtype=torch.float32, device=DEVICE
189
)
190
191
for e in range(E):
@@ -202,7 +208,7 @@ def token_random(E, T, H2, tokens_per_expert):
208
Initialize each token in a random range so we test a range of
209
scale values.
210
"""
205
- y = torch.empty((E, T, H2), dtype=torch.bfloat16, device="cuda")
211
+ y = torch.empty((E, T, H2), dtype=torch.bfloat16, device=DEVICE)
212
for e in range(E):
213
for t in range(tokens_per_expert[e].item()):
214
exp = random.choice(range(1, 20))
@@ -221,7 +227,7 @@ def test_silu_mul_fp8_quant_deep_gemm(E: int, T: int, H: int, fp8_type: torch.dt
227
high=T,
228
size=(E,),
229
dtype=torch.int32,
224
- device="cuda",
230
+ device=DEVICE,
231
)
232
233
# Input tensor of shape (E, T, 2*H)
@@ -288,12 +294,11 @@ def test_silu_mul_fp8_quant_deep_gemm(E: int, T: int, H: int, fp8_type: torch.dt
294
for e in range(E):
295
nt = tokens_per_expert[e].item()
296
291
- if current_platform.is_rocm():
292
- # On ROCm the Triton fallback kernel uses f32 math
293
- # intrinsics (tl.exp) that may differ from PyTorch's
294
- # torch.exp by 1 ULP. At FP8 quantization
295
- # boundaries this can flip one representable value.
296
- # Allow 1 FP8 quantum of tolerance.
297
+ if current_platform.is_rocm() or current_platform.is_xpu():
298
+ # ROCm and XPU both run the Triton fallback kernel, whose f32
299
+ # math intrinsics (tl.exp) may differ from PyTorch's torch.exp
300
+ # by 1 ULP. At FP8 quantization boundaries this can flip one
301
+ # representable value. Allow 1 FP8 quantum of tolerance.
302
torch.testing.assert_close(
303
y_q[e, :nt].to(torch.float32),
304
ref_y_q[e, :nt].to(torch.float32),
vllm/model_executor/layers/fused_moe/experts/batched_deep_gemm_moe.py
+3
-6
@@ -211,16 +211,13 @@ def persistent_masked_m_silu_mul_quant(
211
DeepGemmQuantScaleFMT.UE8M0,
212
]
213
214
- device_capability = current_platform.get_device_capability()
215
- assert device_capability is not None
216
- cuda_arch = device_capability.to_int()
217
-
218
- if current_platform.is_cuda() and cuda_arch >= 80:
214
+ # The C++ kernel requires sm_80+; ROCm and XPU take the Triton path below.
215
+ if current_platform.is_cuda() and current_platform.has_device_capability(80):
216
torch.ops._C.persistent_masked_m_silu_mul_quant(
217
y, tokens_per_expert, y_q, y_s, ceil_ue8m0
218
)
219
else:
223
- # Triton fallback for ROCm -- the C++ kernel is guarded by
220
+ # Triton fallback for ROCm and XPU -- the C++ kernel is guarded by
221
# #ifndef USE_ROCM in activation_kernels.cu.
222
# https://github.com/ROCm/aiter/issues/2420
223
stride_cnt_e = tokens_per_expert.stride()[0]