From 61141ed265bfef41a0ca19e992567ea980919b96 Mon Sep 17 00:00:00 2001 From: Tomasz Zielinski <85164140+tzielinski-habana@users.noreply.github.com> Date: Wed, 15 Jul 2026 17:19:44 +0200 Subject: [PATCH] [Hardware][XPU] Register batch-invariant kernels for XPU (#41934) Signed-off-by: tzielinski-habana 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 --- tests/v1/determinism/test_batch_invariance.py | 14 +++- .../determinism/test_nvfp4_batch_invariant.py | 4 +- .../test_online_batch_invariance.py | 4 +- .../test_rms_norm_batch_invariant.py | 41 +++++++++-- tests/v1/determinism/utils.py | 68 +++++++++++++------ vllm/model_executor/layers/batch_invariant.py | 51 +++++++------- 6 files changed, 127 insertions(+), 55 deletions(-) diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index fb12ffd1706..b2706ed89b7 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/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, diff --git a/tests/v1/determinism/test_nvfp4_batch_invariant.py b/tests/v1/determinism/test_nvfp4_batch_invariant.py index d7a1c9e8404..dafd1a16444 100644 --- a/tests/v1/determinism/test_nvfp4_batch_invariant.py +++ b/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")) diff --git a/tests/v1/determinism/test_online_batch_invariance.py b/tests/v1/determinism/test_online_batch_invariance.py index 2bebb2dca53..de80e4918c1 100644 --- a/tests/v1/determinism/test_online_batch_invariance.py +++ b/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, diff --git a/tests/v1/determinism/test_rms_norm_batch_invariant.py b/tests/v1/determinism/test_rms_norm_batch_invariant.py index 7fbf8f04610..5b3b7a8758b 100644 --- a/tests/v1/determinism/test_rms_norm_batch_invariant.py +++ b/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...") diff --git a/tests/v1/determinism/utils.py b/tests/v1/determinism/utils.py index bbef6147723..f03ea05b433 100644 --- a/tests/v1/determinism/utils.py +++ b/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: diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index 917c72dee8c..f05998b1beb 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/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():