[Misc] Add and enable Triton kernel unit tests on XPU (#45694)
Signed-off-by: pmanczak <pmanczak@users.noreply.github.com> Signed-off-by: pmanczak <pawel.manczak@intel.com> Co-authored-by: pmanczak <pmanczak@users.noreply.github.com>
pmanczak committed
Aug 7, 2026 at 11:53 UTC
6b5bec7bedffa949fbd393fce720faf35831d746
5 files changed
+58
-13
tests/kernels/core/test_fused_rms_norm_gated.py
+5
-2
@@ -7,11 +7,14 @@ matching the eager triton kernel output."""
7
import pytest
8
import torch
9
10
+from vllm.platforms import current_platform
11
from vllm.third_party.flash_linear_attention.ops.fused_norm_gate import (
12
FusedRMSNormGated,
13
)
14
from vllm.utils.torch_utils import set_random_seed
15
16
+DEVICE = "xpu:0" if current_platform.is_xpu() else "cuda:0"
17
+
18
DTYPES = [torch.bfloat16]
19
HIDDEN_SIZES = [128, 512]
20
NUM_TOKENS = [64, 128]
@@ -39,7 +42,7 @@ def test_compiled_vs_eager(
42
"""forward_native decomposition matches forward_cuda triton kernel."""
43
torch._dynamo.reset()
44
set_random_seed(seed)
42
- device = torch.device("cuda:0")
45
+ device = torch.device(DEVICE)
46
47
module = FusedRMSNormGated(
48
hidden_size,
@@ -88,7 +91,7 @@ def test_compiled_vs_eager_multidim(
91
"""forward_native decomposition handles multi-dimensional inputs."""
92
torch._dynamo.reset()
93
set_random_seed(seed)
91
- device = torch.device("cuda:0")
94
+ device = torch.device(DEVICE)
95
head_dim = shape[-1]
96
97
module = FusedRMSNormGated(
tests/kernels/quantization/test_block_int8.py
+9
-1
@@ -14,7 +14,15 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (
14
)
15
from vllm.platforms import current_platform
16
17
-if current_platform.get_device_capability() < (7, 0):
17
+if not (current_platform.is_cuda_alike() or current_platform.is_xpu()):
18
+ pytest.skip(
19
+ "INT8 Triton kernels require a CUDA-alike or XPU device",
20
+ allow_module_level=True,
21
+ )
22
+
23
+if current_platform.is_cuda_alike() and not current_platform.has_device_capability(
24
+ (7, 0)
25
+):
26
pytest.skip("INT8 Triton requires CUDA 7.0 or higher", allow_module_level=True)
27
28
vllm_config = VllmConfig()
tests/kernels/quantization/test_int8_kernel.py
+32
-1
@@ -7,6 +7,7 @@ import itertools
7
import pytest
8
import torch
9
10
+from tests.kernels.quant_utils import native_per_token_group_quant_int8
11
from vllm.model_executor.layers.activation import SiluAndMul
12
from vllm.model_executor.layers.fused_moe import fused_experts
13
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
@@ -15,7 +16,15 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (
16
)
17
from vllm.platforms import current_platform
18
18
-if current_platform.get_device_capability() < (7, 0):
19
+if not (current_platform.is_cuda_alike() or current_platform.is_xpu()):
20
+ pytest.skip(
21
+ "INT8 Triton kernels require a CUDA-alike or XPU device",
22
+ allow_module_level=True,
23
+ )
24
+
25
+if current_platform.is_cuda_alike() and not current_platform.has_device_capability(
26
+ (7, 0)
27
+):
28
pytest.skip("INT8 Triton requires CUDA 7.0 or higher", allow_module_level=True)
29
30
@@ -155,3 +164,25 @@ def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed):
164
torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))
165
) / torch.mean(torch.abs(ref_out.to(torch.float32)))
166
assert rel_diff < 0.05
167
+
168
+
169
+@pytest.mark.parametrize("num_tokens", [1, 33, 64])
170
+@pytest.mark.parametrize("hidden_dim", [128, 1024])
171
+@pytest.mark.parametrize("dtype", DTYPES)
172
+@torch.inference_mode()
173
+def test_per_token_quant_int8(num_tokens, hidden_dim, dtype):
174
+ """Covers the round_int8 Triton kernel via per_token_quant_int8."""
175
+ torch.manual_seed(0)
176
+ device = current_platform.device_type
177
+ x = torch.randn((num_tokens, hidden_dim), dtype=dtype, device=device) * 8
178
+
179
+ q, s = per_token_quant_int8(x)
180
+ # Per-token int8 quant is per-token-group quant with group_size = hidden_dim.
181
+ ref_q, ref_s = native_per_token_group_quant_int8(x, hidden_dim)
182
+
183
+ assert q.dtype == torch.int8
184
+ assert q.shape == (num_tokens, hidden_dim)
185
+ assert s.shape == (num_tokens, 1)
186
+ assert torch.allclose(s.float(), ref_s.float(), atol=1e-3, rtol=1e-3)
187
+ # Quantized values may differ by at most 1 ULP due to rounding.
188
+ assert torch.all((q.float() - ref_q.float()).abs() <= 1)
tests/kernels/quantization/test_triton_scaled_mm.py
+1
-1
@@ -13,7 +13,7 @@ import torch
13
from vllm.platforms import current_platform
14
from vllm.utils.torch_utils import set_random_seed
15
16
-device = "cuda"
16
+device = current_platform.device_type
17
18
triton_scaled_mm_module = importlib.import_module(
19
"vllm.model_executor.layers.quantization.compressed_tensors.triton_scaled_mm"
tests/kernels/test_fla_layernorm_guard.py
+11
-8
@@ -5,6 +5,7 @@ import pytest
5
import torch
6
import torch.nn.functional as F
7
8
+from vllm.platforms import current_platform
9
from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (
10
layer_norm_fwd,
11
layernorm_fn,
@@ -12,6 +13,8 @@ from vllm.third_party.flash_linear_attention.ops.layernorm_guard import (
13
)
14
from vllm.utils.torch_utils import set_random_seed
15
16
+DEVICE = "xpu" if current_platform.is_xpu() else "cuda"
17
+
18
19
def layer_norm_ref(
20
x,
@@ -115,7 +118,7 @@ def test_layer_norm_fwd_basic(
118
) -> None:
119
"""Test basic layer norm forward pass without z (gate) tensor."""
120
set_random_seed(seed)
118
- device = torch.device("cuda:0")
121
+ device = torch.device(DEVICE)
122
123
# Create inputs
124
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
@@ -157,7 +160,7 @@ def test_layer_norm_fwd_with_gate(
160
) -> None:
161
"""Test layer norm forward pass with z (gate) tensor."""
162
set_random_seed(42)
160
- device = torch.device("cuda:0")
163
+ device = torch.device(DEVICE)
164
165
# Create inputs
166
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
@@ -214,7 +217,7 @@ def test_layer_norm_fwd_with_groups(
217
)
218
219
set_random_seed(42)
217
- device = torch.device("cuda:0")
220
+ device = torch.device(DEVICE)
221
222
# Create inputs
223
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
@@ -254,7 +257,7 @@ def test_layer_norm_rows_per_block(
257
) -> None:
258
"""Test that rows_per_block logic works correctly for various M sizes."""
259
set_random_seed(42)
257
- device = torch.device("cuda:0")
260
+ device = torch.device(DEVICE)
261
hidden_size = 1024
262
263
# Create inputs
@@ -279,7 +282,7 @@ def test_strided_input(dtype: torch.dtype) -> None:
282
"""Test that the kernel handles non-contiguous (strided)
283
inputs correctly."""
284
set_random_seed(42)
282
- device = torch.device("cuda:0")
285
+ device = torch.device(DEVICE)
286
num_tokens = 128
287
hidden_size = 1024
288
@@ -319,7 +322,7 @@ def test_output_buffer_provided(
322
) -> None:
323
"""Test that the kernel works when an output buffer is provided."""
324
set_random_seed(42)
322
- device = torch.device("cuda:0")
325
+ device = torch.device(DEVICE)
326
327
# Create inputs
328
x = torch.randn(num_tokens, hidden_size, dtype=dtype, device=device)
@@ -360,7 +363,7 @@ def test_multidimensional_input(
363
) -> None:
364
"""Test that the autograd function handles multidimensional inputs."""
365
set_random_seed(42)
363
- device = torch.device("cuda:0")
366
+ device = torch.device(DEVICE)
367
hidden_size = shape[-1]
368
369
# Create inputs
@@ -404,7 +407,7 @@ def test_rmsnorm_gated_forward_native_dtype(
407
408
from vllm.model_executor.layers.layernorm import RMSNormGated
409
407
- device = torch.device("cuda:0")
410
+ device = torch.device(DEVICE)
411
set_random_seed(42)
412
413
layer = RMSNormGated(