[Hardware][XPU] Register batch-invariant kernels for XPU (#41934)

Signed-off-by: tzielinski-habana <[email protected]>
Signed-off-by: Tomasz Zielinski <[email protected]>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Chendi.Xue <[email protected]>

tests/v1/determinism/test_batch_invariance.py

@@ -11,11 +11,13 @@ from utils import (
     TEST_MODEL,
     _extract_step_logprobs,
     _random_prompt,
+    skip_if_not_cuda,
     skip_unsupported,
 )
 
 import vllm.envs as envs
 from vllm import LLM, SamplingParams
+from vllm.platforms import current_platform
 
 
 @skip_unsupported
@@ -49,6 +51,11 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
       seed.
     - Keep max_tokens and max_model_len bounded for speed and memory use.
     """
+    # Not all batch-invariant kernels are registered on XPU yet
+    # (e.g. attention, custom ops), so e2e determinism is not guaranteed.
+    if current_platform.is_xpu():
+        pytest.xfail("Not all batch-invariant kernels registered on XPU yet")
+
     seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
     random.seed(seed)
 
@@ -157,6 +164,11 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
     block_m,
     block_n,
 ):
+    # Not all batch-invariant kernels are registered on XPU yet
+    # (e.g. attention, custom ops), so e2e determinism is not guaranteed.
+    if current_platform.is_xpu():
+        pytest.xfail("Not all batch-invariant kernels registered on XPU yet")
+
     seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
     random.seed(seed)
     tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
@@ -641,7 +653,7 @@ def test_logprobs_without_batch_invariance_should_fail(
         pytest.fail(fail_msg)
 
 
-@skip_unsupported
+@skip_if_not_cuda
 @pytest.mark.parametrize("backend", ["FLASH_ATTN"])
 def test_decode_logprobs_match_prefill_logprobs(
     backend,

tests/v1/determinism/test_nvfp4_batch_invariant.py

@@ -9,7 +9,7 @@ import torch
 from utils import (
     _extract_step_logprobs,
     _random_prompt,
-    skip_unsupported,
+    skip_if_not_cuda,
 )
 
 from vllm import LLM, SamplingParams
@@ -40,7 +40,7 @@ def _make_llm(max_num_seqs: int, backend: str) -> LLM:
     )
 
 
-@skip_unsupported
+@skip_if_not_cuda
 @pytest.mark.parametrize("backend", ["FLASH_ATTN"])
 def test_dense_nvfp4_generation_is_deterministic_across_batch_sizes_e2e(backend):
     seed = int(os.getenv("VLLM_TEST_SEED", "12345"))

tests/v1/determinism/test_online_batch_invariance.py

@@ -17,7 +17,7 @@ from typing import Any
 
 import openai
 import pytest
-from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_unsupported
+from utils import BACKENDS, TEST_MODEL, _random_prompt, skip_if_not_cuda
 
 from tests.utils import RemoteOpenAIServer
 
@@ -133,7 +133,7 @@ def _compare_bs1_vs_bsn_single_process(
                 )
 
 
-@skip_unsupported
+@skip_if_not_cuda
 @pytest.mark.parametrize("backend", BACKENDS)
 def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
     backend: str,

tests/v1/determinism/test_rms_norm_batch_invariant.py

@@ -9,7 +9,7 @@ with the standard CUDA-based implementation to ensure numerical accuracy.
 
 import pytest
 import torch
-from utils import skip_unsupported
+from utils import skip_if_not_cuda, skip_unsupported
 
 from vllm.model_executor.layers.batch_invariant import (
     rms_norm_batch_invariant,
@@ -20,7 +20,7 @@ from vllm.platforms import current_platform
 DEVICE_TYPE = current_platform.device_type
 
 
-@skip_unsupported
+@skip_if_not_cuda
 @pytest.mark.parametrize("batch_size", [1, 4, 16, 64])
 @pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192])
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
@@ -73,7 +73,7 @@ def test_rms_norm_batch_invariant_vs_standard(
     )
 
 
-@skip_unsupported
+@skip_if_not_cuda
 @pytest.mark.parametrize("hidden_size", [512, 4096])
 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
 @pytest.mark.parametrize("eps", [1e-6])
@@ -166,7 +166,7 @@ def test_fused_add_rms_norm_batch_invariant_residual_path(
     )
 
 
-@skip_unsupported
+@skip_if_not_cuda
 @pytest.mark.parametrize("batch_size", [1, 16, 128])
 @pytest.mark.parametrize("seq_len", [1, 32, 512])
 @pytest.mark.parametrize("hidden_size", [2048, 4096])
@@ -210,7 +210,7 @@ def test_rms_norm_3d_input(
     )
 
 
-@skip_unsupported
+@skip_if_not_cuda
 def test_rms_norm_numerical_stability(default_vllm_config):
     """
     Test RMS norm numerical stability with extreme values.
@@ -303,7 +303,7 @@ def test_rms_norm_formula(default_vllm_config):
     )
 
 
-@skip_unsupported
+@skip_if_not_cuda
 @pytest.mark.parametrize("hidden_size", [128, 1024, 4096, 16384])
 def test_rms_norm_different_hidden_sizes(default_vllm_config, hidden_size: int):
     """
@@ -377,6 +377,35 @@ def test_rms_norm_determinism(default_vllm_config):
         )
 
 
+@skip_unsupported
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+def test_rms_norm_batch_invariance(dtype):
+    """Same row gives identical rms_norm result regardless of batch neighbors.
+
+    This verifies that the output for a given row is independent of what other
+    rows are present in the batch — the core batch-invariance property.
+    """
+    device = torch.device(DEVICE_TYPE)
+    torch.manual_seed(42)
+    hidden_size = 2048
+    eps = 1e-6
+
+    weight = torch.randn(hidden_size, dtype=dtype, device=device)
+    row = torch.randn(1, hidden_size, dtype=dtype, device=device)
+
+    # Compute rms_norm on the single row alone
+    out_single = rms_norm_batch_invariant(row, weight, eps=eps)
+
+    # Embed the same row in a larger batch with random neighbors
+    batch = torch.randn(8, hidden_size, dtype=dtype, device=device)
+    batch[4] = row[0]
+    out_batch = rms_norm_batch_invariant(batch, weight, eps=eps)
+
+    assert torch.equal(out_single[0], out_batch[4]), (
+        "rms_norm output for a row differs when batch context changes"
+    )
+
+
 if __name__ == "__main__":
     # Run a quick smoke test
     print("Running quick smoke test of RMS norm implementations...")

tests/v1/determinism/utils.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import os
 import random
+from typing import NamedTuple
 
 import pytest
 import torch
@@ -11,36 +12,61 @@ from vllm.transformers_utils.config import get_config
 from vllm.transformers_utils.model_arch_config_convertor import (
     ModelArchConfigConvertorBase,
 )
+from vllm.triton_utils import HAS_TRITON
 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla
 
-skip_unsupported = pytest.mark.skipif(
-    not (current_platform.is_cuda() and current_platform.has_device_capability(80)),
-    # Supports testing on Ampere and Ada Lovelace devices.
-    # Note: For devices with SM < 90, batch invariance does not support CUDA Graphs.
-    reason="Requires CUDA and >= Ampere (SM80)",
-)
+
+class DeviceConfig(NamedTuple):
+    available: bool
+    backends: list[str]
+
+
+# Maps each device to its availability and supported backends.
+DEVICE_BACKENDS: dict[str, DeviceConfig] = {
+    "cuda": DeviceConfig(
+        available=current_platform.is_cuda()
+        and current_platform.has_device_capability(80),
+        # FlashInfer backend temporarily disabled due to invariant CTA sizes.
+        # See FlashInfer issue #2424
+        backends=["FLASH_ATTN", "TRITON_ATTN", "FLEX_ATTENTION"],
+    ),
+    "xpu": DeviceConfig(
+        available=current_platform.is_xpu() and HAS_TRITON,
+        backends=["TRITON_ATTN"],
+    ),
+}
 
 DEFAULT_MODEL = "Qwen/Qwen3-1.7B"
 TEST_MODEL = os.getenv("VLLM_TEST_MODEL", DEFAULT_MODEL)
 
-BACKENDS: list[str] = [
-    "FLASH_ATTN",
-    "TRITON_ATTN",
-    "FLEX_ATTENTION",
-]
-
-# FlashInfer temporarily disabled due to invariant CTA sizes.
-# See FlashInfer issue #2424
-# if has_flashinfer():
-#     BACKENDS.append("FLASHINFER")
-
-# only run MLA backends when the requested test model is itself an MLA model.
+# Override backends for MLA models (MLA only supported on CUDA).
 if os.getenv("VLLM_TEST_MODEL"):
     config = get_config(TEST_MODEL, trust_remote_code=False)
     if ModelArchConfigConvertorBase(config, config.get_text_config()).is_deepseek_mla():
-        BACKENDS = ["TRITON_MLA"]
-        if flash_attn_supports_mla():
-            BACKENDS.append("FLASH_ATTN_MLA")
+        DEVICE_BACKENDS["cuda"] = DeviceConfig(
+            available=DEVICE_BACKENDS["cuda"].available,
+            backends=["TRITON_MLA"]
+            + (["FLASH_ATTN_MLA"] if flash_attn_supports_mla() else []),
+        )
+        DEVICE_BACKENDS["xpu"] = DeviceConfig(
+            available=DEVICE_BACKENDS["xpu"].available,
+            backends=[],
+        )
+
+# Only include backends for devices that are actually available.
+BACKENDS: list[str] = sorted(
+    {b for cfg in DEVICE_BACKENDS.values() if cfg.available for b in cfg.backends}
+)
+
+skip_unsupported = pytest.mark.skipif(
+    not any(cfg.available for cfg in DEVICE_BACKENDS.values()),
+    reason="Requires CUDA >= Ampere (SM80) or Intel XPU with Triton",
+)
+
+skip_if_not_cuda = pytest.mark.skipif(
+    not DEVICE_BACKENDS["cuda"].available,
+    reason="Requires CUDA >= Ampere (SM80)",
+)
 
 
 def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:

vllm/model_executor/layers/batch_invariant.py

@@ -904,37 +904,41 @@ def enable_batch_invariant_mode():
     _batch_invariant_MODE = True
     _batch_invariant_LIB = torch.library.Library("aten", "IMPL")
 
-    if current_platform.is_device_capability_family(80):
-        # SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the
-        # triton persistent matmul overrides for mm/addmm/matmul/linear.
-        _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA")
-        _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA")
-        _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, "CUDA")
-        _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, "CUDA")
-    else:
-        # Hopper (SM90) and Blackwell (SM100): the only source of batch
-        # variance is split-k, which we disable via the cuBLAS workspace
-        # config.
-        os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
-        os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1"
+    key = current_platform.dispatch_key
 
-    # Triton bmm/persistent-matmul kernels read this for the FP16 N-tile size;
-    # set unconditionally because bmm is overridden on all CUDA platforms.
     if current_platform.is_cuda():
+        if current_platform.is_device_capability_family(80):
+            # SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the
+            # triton persistent matmul overrides for mm/addmm/matmul/linear.
+            _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key)
+            _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key)
+            _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, key)
+            _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, key)
+        else:
+            # Hopper (SM90) and Blackwell (SM100): the only source of batch
+            # variance is split-k, which we disable via the cuBLAS workspace
+            # config.
+            os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
+            os.environ["CUBLASLT_WORKSPACE_SIZE"] = "1"
+
         _fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128
+    elif current_platform.is_xpu():
+        _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key)
+        _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key)
+        # TODO: register matmul and linear for XPU
+        # once suitable Triton kernels are implemented
 
-    _batch_invariant_LIB.impl(
-        "aten::_log_softmax", _log_softmax_batch_invariant, "CUDA"
-    )
-    _batch_invariant_LIB.impl("aten::softmax", softmax_batch_invariant, "CUDA")
-    _batch_invariant_LIB.impl("aten::_softmax", softmax_batch_invariant, "CUDA")
-    _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA")
+        _fp16_block_size_n = 128
 
+    _batch_invariant_LIB.impl("aten::_log_softmax", _log_softmax_batch_invariant, key)
+    _batch_invariant_LIB.impl("aten::softmax", softmax_batch_invariant, key)
+    _batch_invariant_LIB.impl("aten::_softmax", softmax_batch_invariant, key)
+    _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, key)
     # torch 2.12+ registers a built-in Triton bmm kernel for CUDA
     # (torch._native.ops.bmm_outer_product), so we need allow_override
     # to replace it at the dispatcher level.
     _batch_invariant_LIB.impl(
-        "aten::bmm", bmm_batch_invariant, "CUDA", allow_override=True
+        "aten::bmm", bmm_batch_invariant, key, allow_override=True
     )
     torch.bmm = bmm_batch_invariant
 
@@ -947,7 +951,8 @@ def enable_batch_invariant_mode():
     torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = (
         reduced_precision_val
     )
-    torch.backends.cuda.preferred_blas_library(backend="cublaslt")
+    if current_platform.is_cuda():
+        torch.backends.cuda.preferred_blas_library(backend="cublaslt")
 
 
 def override_envs_for_invariance():