[Hardware][XPU] Register batch-invariant kernels for XPU (#41934)
Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com> Signed-off-by: Tomasz Zielinski <85164140+tzielinski-habana@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Chendi.Xue <chendi.xue@intel.com>
Tomasz Zielinski committed
Jul 15, 2026 at 17:19 UTC
61141ed265bfef41a0ca19e992567ea980919b96
6 files changed
+130
-58
tests/v1/determinism/test_batch_invariance.py
+13
-1
@@ -11,11 +11,13 @@ from utils import (
11
TEST_MODEL,
12
_extract_step_logprobs,
13
_random_prompt,
14
+ skip_if_not_cuda,
15
skip_unsupported,
16
)
17
18
import vllm.envs as envs
19
from vllm import LLM, SamplingParams
20
+from vllm.platforms import current_platform
21
22
23
@skip_unsupported
@@ -49,6 +51,11 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
51
seed.
52
- Keep max_tokens and max_model_len bounded for speed and memory use.
53
"""
54
+ # Not all batch-invariant kernels are registered on XPU yet
55
+ # (e.g. attention, custom ops), so e2e determinism is not guaranteed.
56
+ if current_platform.is_xpu():
57
+ pytest.xfail("Not all batch-invariant kernels registered on XPU yet")
58
+
59
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
60
random.seed(seed)
61
@@ -157,6 +164,11 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
164
block_m,
165
block_n,
166
):
167
+ # Not all batch-invariant kernels are registered on XPU yet
168
+ # (e.g. attention, custom ops), so e2e determinism is not guaranteed.
169
+ if current_platform.is_xpu():
170
+ pytest.xfail("Not all batch-invariant kernels registered on XPU yet")
171
+
172
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
173
random.seed(seed)
174
tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
@@ -641,7 +653,7 @@ def test_logprobs_without_batch_invariance_should_fail(
653
pytest.fail(fail_msg)
654
655
644
-@skip_unsupported
656
+@skip_if_not_cuda
657
@pytest.mark.parametrize("backend", ["FLASH_ATTN"])
658
def test_decode_logprobs_match_prefill_logprobs(
659
backend,
tests/v1/determinism/test_nvfp4_batch_invariant.py
+2
-2
@@ -9,7 +9,7 @@ import torch
9
from utils import (
10
_extract_step_logprobs,
11
_random_prompt,
12
- skip_unsupported,
12
+ skip_if_not_cuda,
13
)
14
15
from vllm import LLM, SamplingParams
@@ -40,7 +40,7 @@ def _make_llm(max_num_seqs: int, backend: str) -> LLM:
40
)
41
42
43
-@skip_unsupported
43
+@skip_if_not_cuda
44
@pytest.mark.parametrize("backend", ["FLASH_ATTN"])
45
def test_dense_nvfp4_generation_is_deterministic_across_batch_sizes_e2e(backend):
46
seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
tests/v1/determinism/test_online_batch_invariance.py
+2
-2
@@ -17,7 +17,7 @@ from typing import Any
17
18
import openai
19
import pytest
20
-from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_unsupported
20
+from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_if_not_cuda
21
22
from tests.utils import RemoteOpenAIServer
23
@@ -133,7 +133,7 @@ def _compare_bs1_vs_bsn_single_process(
133
)
134
135
136
-@skip_unsupported
136
+@skip_if_not_cuda
137
@pytest.mark.parametrize("backend", BACKENDS)
138
def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
139
backend: str,
tests/v1/determinism/test_rms_norm_batch_invariant.py
+35
-6
@@ -9,7 +9,7 @@ with the standard CUDA-based implementation to ensure numerical accuracy.
9
10
import pytest
11
import torch
12
-from utils import skip_unsupported
12
+from utils import skip_if_not_cuda, skip_unsupported
13
14
from vllm.model_executor.layers.batch_invariant import (
15
rms_norm_batch_invariant,
@@ -20,7 +20,7 @@ from vllm.platforms import current_platform
20
DEVICE_TYPE = current_platform.device_type
21
22
23
-@skip_unsupported
23
+@skip_if_not_cuda
24
@pytest.mark.parametrize("batch_size", [1, 4, 16, 64])
25
@pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192])
26
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@@ -73,7 +73,7 @@ def test_rms_norm_batch_invariant_vs_standard(
73
)
74
75
76
-@skip_unsupported
76
+@skip_if_not_cuda
77
@pytest.mark.parametrize("hidden_size", [512, 4096])
78
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
79
@pytest.mark.parametrize("eps", [1e-6])
@@ -166,7 +166,7 @@ def test_fused_add_rms_norm_batch_invariant_residual_path(
166
)
167
168
169
-@skip_unsupported
169
+@skip_if_not_cuda
170
@pytest.mark.parametrize("batch_size", [1, 16, 128])
171
@pytest.mark.parametrize("seq_len", [1, 32, 512])
172
@pytest.mark.parametrize("hidden_size", [2048, 4096])
@@ -210,7 +210,7 @@ def test_rms_norm_3d_input(
210
)
211
212
213
-@skip_unsupported
213
+@skip_if_not_cuda
214
def test_rms_norm_numerical_stability(default_vllm_config):
215
"""
216
Test RMS norm numerical stability with extreme values.
@@ -303,7 +303,7 @@ def test_rms_norm_formula(default_vllm_config):
303
)
304
305
306
-@skip_unsupported
306
+@skip_if_not_cuda
307
@pytest.mark.parametrize("hidden_size", [128, 1024, 4096, 16384])
308
def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int):
309
"""
@@ -377,6 +377,35 @@ def test_rms_norm_determinism(default_vllm_config):
377
)
378
379
380
+@skip_unsupported
381
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
382
+def test_rms_norm_batch_invariance(dtype):
383
+ """Same row gives identical rms_norm result regardless of batch neighbors.
384
+
385
+ This verifies that the output for a given row is independent of what other
386
+ rows are present in the batch — the core batch-invariance property.
387
+ """
388
+ device = torch.device(DEVICE_TYPE)
389
+ torch.manual_seed(42)
390
+ hidden_size = 2048
391
+ eps = 1e-6
392
+
393
+ weight = torch.randn(hidden_size, dtype=dtype, device=device)
394
+ row = torch.randn(1, hidden_size, dtype=dtype, device=device)
395
+
396
+ # Compute rms_norm on the single row alone
397
+ out_single = rms_norm_batch_invariant(row, weight, eps=eps)
398
+
399
+ # Embed the same row in a larger batch with random neighbors
400
+ batch = torch.randn(8, hidden_size, dtype=dtype, device=device)
401
+ batch[4] = row[0]
402
+ out_batch = rms_norm_batch_invariant(batch, weight, eps=eps)
403
+
404
+ assert torch.equal(out_single[0], out_batch[4]), (
405
+ "rms_norm output for a row differs when batch context changes"
406
+ )
407
+
408
+
409
if __name__ == "__main__":
410
# Run a quick smoke test
411
print("Running quick smoke test of RMS norm implementations...")
tests/v1/determinism/utils.py
+47
-21
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import os
4
import random
5
+from typing import NamedTuple
6
7
import pytest
8
import torch
@@ -11,36 +12,61 @@ from vllm.transformers_utils.config import get_config
12
from vllm.transformers_utils.model_arch_config_convertor import (
13
ModelArchConfigConvertorBase,
14
)
15
+from vllm.triton_utils import HAS_TRITON
16
from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
17
16
-skip_unsupported = pytest.mark.skipif(
17
- not (current_platform.is_cuda() and current_platform.has_device_capability(80)),
18
- # Supports testing on Ampere and Ada Lovelace devices.
19
- # Note: For devices with SM < 90, batch invariance does not support CUDA Graphs.
20
- reason="Requires CUDA and >= Ampere (SM80)",
21
-)
18
23
-DEFAULT_MODEL = "Qwen/Qwen3-1.7B"
24
-TEST_MODEL = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL)
19
+class DeviceConfig(NamedTuple):
20
+ available: bool
21
+ backends: list[str]
22
26
-BACKENDS: list[str] = [
27
- "FLASH_ATTN",
28
- "TRITON_ATTN",
29
- "FLEX_ATTENTION",
30
-]
23
32
-# FlashInfer temporarily disabled due to invariant CTA sizes.
33
-# See FlashInfer issue #2424
34
-# if has_flashinfer():
35
-# BACKENDS.append("FLASHINFER")
24
+# Maps each device to its availability and supported backends.
25
+DEVICE_BACKENDS: dict[str, DeviceConfig] = {
26
+ "cuda": DeviceConfig(
27
+ available=current_platform.is_cuda()
28
+ and current_platform.has_device_capability(80),
29
+ # FlashInfer backend temporarily disabled due to invariant CTA sizes.
30
+ # See FlashInfer issue #2424
31
+ backends=["FLASH_ATTN", "TRITON_ATTN", "FLEX_ATTENTION"],
32
+ ),
33
+ "xpu": DeviceConfig(
34
+ available=current_platform.is_xpu() and HAS_TRITON,
35
+ backends=["TRITON_ATTN"],
36
+ ),
37
+}
38
37
-# only run MLA backends when the requested test model is itself an MLA model.
39
+DEFAULT_MODEL = "Qwen/Qwen3-1.7B"
40
+TEST_MODEL = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL)
41
+
42
+# Override backends for MLA models (MLA only supported on CUDA).
43
if os.getenv("VLLM_TEST_MODEL"):
44
config = get_config(TEST_MODEL, trust_remote_code=False)
45
if ModelArchConfigConvertorBase(config, config.get_text_config()).is_deepseek_mla():
41
- BACKENDS = ["TRITON_MLA"]
42
- if flash_attn_supports_mla():
43
- BACKENDS.append("FLASH_ATTN_MLA")
46
+ DEVICE_BACKENDS["cuda"] = DeviceConfig(
47
+ available=DEVICE_BACKENDS["cuda"].available,
48
+ backends=["TRITON_MLA"]
49
+ + (["FLASH_ATTN_MLA"] if flash_attn_supports_mla() else []),
50
+ )
51
+ DEVICE_BACKENDS["xpu"] = DeviceConfig(
52
+ available=DEVICE_BACKENDS["xpu"].available,
53
+ backends=[],
54
+ )
55
+
56
+# Only include backends for devices that are actually available.
57
+BACKENDS: list[str] = sorted(
58
+ {b for cfg in DEVICE_BACKENDS.values() if cfg.available for b in cfg.backends}
59
+)
60
+
61
+skip_unsupported = pytest.mark.skipif(
62
+ not any(cfg.available for cfg in DEVICE_BACKENDS.values()),
63
+ reason="Requires CUDA >= Ampere (SM80) or Intel XPU with Triton",
64
+)
65
+
66
+skip_if_not_cuda = pytest.mark.skipif(
67
+ not DEVICE_BACKENDS["cuda"].available,
68
+ reason="Requires CUDA >= Ampere (SM80)",
69
+)
70
71
72
def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:
vllm/model_executor/layers/batch_invariant.py
+31
-26
@@ -904,37 +904,41 @@ def enable_batch_invariant_mode():
904
_batch_invariant_MODE = True
905
_batch_invariant_LIB = torch.library.Library("aten", "IMPL")
906
907
- if current_platform.is_device_capability_family(80):
908
- # SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the
909
- # triton persistent matmul overrides for mm/addmm/matmul/linear.
910
- _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA")
911
- _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA")
912
- _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, "CUDA")
913
- _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, "CUDA")
914
- else:
915
- # Hopper (SM90) and Blackwell (SM100): the only source of batch
916
- # variance is split-k, which we disable via the cuBLAS workspace
917
- # config.
918
- os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
919
- os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1"
920
-
921
- # Triton bmm/persistent-matmul kernels read this for the FP16 N-tile size;
922
- # set unconditionally because bmm is overridden on all CUDA platforms.
923
- if current_platform.is_cuda():
924
- _fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128
907
+ key = current_platform.dispatch_key
908
926
- _batch_invariant_LIB.impl(
927
- "aten::_log_softmax", _log_softmax_batch_invariant, "CUDA"
928
- )
929
- _batch_invariant_LIB.impl("aten::softmax", softmax_batch_invariant, "CUDA")
930
- _batch_invariant_LIB.impl("aten::_softmax", softmax_batch_invariant, "CUDA")
931
- _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA")
909
+ if current_platform.is_cuda():
910
+ if current_platform.is_device_capability_family(80):
911
+ # SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the
912
+ # triton persistent matmul overrides for mm/addmm/matmul/linear.
913
+ _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key)
914
+ _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key)
915
+ _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, key)
916
+ _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, key)
917
+ else:
918
+ # Hopper (SM90) and Blackwell (SM100): the only source of batch
919
+ # variance is split-k, which we disable via the cuBLAS workspace
920
+ # config.
921
+ os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
922
+ os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1"
923
924
+ _fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128
925
+ elif current_platform.is_xpu():
926
+ _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key)
927
+ _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key)
928
+ # TODO: register matmul and linear for XPU
929
+ # once suitable Triton kernels are implemented
930
+
931
+ _fp16_block_size_n = 128
932
+
933
+ _batch_invariant_LIB.impl("aten::_log_softmax", _log_softmax_batch_invariant, key)
934
+ _batch_invariant_LIB.impl("aten::softmax", softmax_batch_invariant, key)
935
+ _batch_invariant_LIB.impl("aten::_softmax", softmax_batch_invariant, key)
936
+ _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, key)
937
# torch 2.12+ registers a built-in Triton bmm kernel for CUDA
938
# (torch._native.ops.bmm_outer_product), so we need allow_override
939
# to replace it at the dispatcher level.
940
_batch_invariant_LIB.impl(
937
- "aten::bmm", bmm_batch_invariant, "CUDA", allow_override=True
941
+ "aten::bmm", bmm_batch_invariant, key, allow_override=True
942
)
943
torch.bmm = bmm_batch_invariant
944
@@ -947,7 +951,8 @@ def enable_batch_invariant_mode():
951
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = (
952
reduced_precision_val
953
)
950
- torch.backends.cuda.preferred_blas_library(backend="cublaslt")
954
+ if current_platform.is_cuda():
955
+ torch.backends.cuda.preferred_blas_library(backend="cublaslt")
956
957
958
def override_envs_for_invariance():