[Hardware][XPU] Register matmul and linear batch-invariant kernels for XPU (#49209)

Signed-off-by: tzielinski-habana <tomasz.zielinski@intel.com>
Signed-off-by: Tomasz Zielinski <85164140+tzielinski-habana@users.noreply.github.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

docs/features/batch_invariance.md

@@ -17,7 +17,29 @@ Batch invariance is crucial for several use cases:
 
 ## Hardware Requirements
 
-Batch invariance requires NVIDIA GPUs with compute capability 8.0 or higher.
+Batch invariance is supported on the following platforms:
+
+- NVIDIA GPUs with compute capability 8.0 or higher.
+- Intel XPUs with Triton support.
+
+### Attention Backend Selection for XPU
+
+On XPU, Triton Attention backend is required for batch invariance.
+Select this backend using Qwen/Qwen3-1.7B model as an example:
+
+```python
+llm = LLM(
+    model="Qwen/Qwen3-1.7B",
+    attention_config={"backend": "TRITON_ATTN"},
+)
+```
+
+Or via the CLI:
+
+```bash
+VLLM_BATCH_INVARIANT=1 vllm serve Qwen/Qwen3-1.7B \
+    --attention-config.backend TRITON_ATTN
+```
 
 ## Enabling Batch Invariance
 

tests/v1/determinism/test_batch_invariance.py

@@ -17,7 +17,6 @@ from utils import (
 
 import vllm.envs as envs
 from vllm import LLM, SamplingParams
-from vllm.platforms import current_platform
 
 
 @skip_unsupported
@@ -53,11 +52,6 @@ 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)
 
@@ -177,11 +171,6 @@ 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"))

vllm/engine/arg_utils.py

@@ -2435,6 +2435,19 @@ class EngineArgs:
                 self.attention_backend
             )
 
+        # Batch-invariant mode requires deterministic attention behavior.
+        # If no backend is explicitly requested, prefer Triton Attention.
+        if (
+            envs.VLLM_BATCH_INVARIANT
+            and attention_config.backend is None
+            and current_platform.is_xpu()
+        ):
+            attention_config.backend = AttentionBackendEnum.TRITON_ATTN
+            logger.info(
+                "VLLM_BATCH_INVARIANT is enabled and no attention backend was "
+                "specified; defaulting to TRITON_ATTN."
+            )
+
         # TurboQuant requires FlashAttention 2 — FA3 boundary layers assert
         # FlashAttentionImpl which fails with TurboQuantAttentionImpl.
         if resolved_cache_dtype.startswith("turboquant_") and (

vllm/model_executor/determinism/batch_invariant.py

@@ -9,6 +9,7 @@ import torch
 
 import vllm.envs as envs
 from vllm.model_executor.determinism.batch_invariant_configs import (
+    _get_descriptor_matmul_config,
     _get_matmul_config,
     resolve_tuned_matmul_configs,
 )
@@ -133,6 +134,149 @@ def matmul_kernel_persistent(
         tl.store(c_ptrs, c, mask=c_mask)
 
 
+@triton.jit(launch_metadata=_matmul_launch_metadata)
+def matmul_kernel_descriptor_persistent(
+    a_ptr,
+    b_ptr,
+    c_ptr,  #
+    bias_ptr,
+    M,
+    N,
+    K,  #
+    BLOCK_SIZE_M: tl.constexpr,  #
+    BLOCK_SIZE_N: tl.constexpr,  #
+    BLOCK_SIZE_K: tl.constexpr,  #
+    GROUP_SIZE_M: tl.constexpr,  #
+    NUM_SMS: tl.constexpr,  #
+    HAS_BIAS: tl.constexpr,
+):
+    """Persistent matmul using tensor descriptors for 2D block I/O.
+
+    Expects b_ptr to point to a transposed B matrix of shape [N, K] with
+    row-major (K-contiguous) layout.  The dot product transposes each loaded
+    B-tile back: dot(A_tile, B_tile.T).
+
+    ~3x faster than the pointer-based persistent kernel on Intel XPU because
+    tensor descriptors leverage hardware 2D block load/store with automatic
+    bounds checking (no explicit masks needed).
+    """
+    dtype = c_ptr.dtype.element_ty
+    start_pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
+    num_tiles = num_pid_m * num_pid_n
+
+    a_desc = tl.make_tensor_descriptor(
+        a_ptr,
+        shape=[M, K],
+        strides=[K, 1],
+        block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K],
+    )
+    b_desc = tl.make_tensor_descriptor(
+        b_ptr,
+        shape=[N, K],
+        strides=[K, 1],
+        block_shape=[BLOCK_SIZE_N, BLOCK_SIZE_K],
+    )
+    c_desc = tl.make_tensor_descriptor(
+        c_ptr,
+        shape=[M, N],
+        strides=[N, 1],
+        block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_N],
+    )
+
+    tile_id_c = start_pid - NUM_SMS
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+
+    for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True):
+        pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M)
+        offs_am = pid_m * BLOCK_SIZE_M
+        offs_bn = pid_n * BLOCK_SIZE_N
+
+        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+        for ki in range(k_tiles):
+            offs_k = ki * BLOCK_SIZE_K
+            a = a_desc.load([offs_am, offs_k])
+            b = b_desc.load([offs_bn, offs_k])
+            accumulator = tl.dot(a, b.T, accumulator)
+
+        tile_id_c += NUM_SMS
+        pid_m, pid_n = _compute_pid(
+            tile_id_c, num_pid_in_group, num_pid_m, GROUP_SIZE_M
+        )
+        offs_cm = pid_m * BLOCK_SIZE_M
+        offs_cn = pid_n * BLOCK_SIZE_N
+
+        if HAS_BIAS:
+            bias_offs = offs_cn + tl.arange(0, BLOCK_SIZE_N)
+            bias_vals = tl.load(bias_ptr + bias_offs, mask=bias_offs < N, other=0.0).to(
+                tl.float32
+            )
+            accumulator += bias_vals[None, :]
+
+        c = accumulator.to(dtype)
+        c_desc.store([offs_cm, offs_cn], c)
+
+
+def matmul_descriptor_persistent(
+    a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None
+):
+    """Persistent matmul using tensor descriptors (Intel XPU fast path).
+
+    Args:
+        a: Input matrix [M, K], must be contiguous.
+        b: Weight matrix [K, N] (standard layout — transposed internally).
+        bias: Optional 1D bias vector [N].
+
+    Returns:
+        Output matrix [M, N] with dtype matching the inputs.
+    """
+    assert a.shape[1] == b.shape[0], "Incompatible dimensions"
+    assert a.dtype == b.dtype, "Incompatible dtypes"
+    assert bias is None or bias.dim() == 1, "bias must be 1D"
+
+    M, K = a.shape
+    _, N = b.shape
+    dtype = a.dtype
+
+    # Tensor descriptors require contiguous row-major layout.
+    a = a.contiguous()
+    # Descriptor kernel expects B in [N, K] layout (K-contiguous).
+    b_t = b.t().contiguous()
+
+    c = torch.empty((M, N), device=a.device, dtype=dtype)
+    # This path is XPU-only. _NUM_SMS is set by enable_batch_invariant_mode()
+    # before torch.compile traces this function, so Dynamo can lift it
+    # as a compile-time constant.
+    NUM_SMS = _NUM_SMS if _NUM_SMS > 0 else num_compute_units(a.device.index)
+
+    def grid(META):
+        return (
+            min(
+                NUM_SMS,
+                triton.cdiv(M, META["BLOCK_SIZE_M"])
+                * triton.cdiv(N, META["BLOCK_SIZE_N"]),
+            ),
+        )
+
+    cfg = _get_descriptor_matmul_config(M, N, K, dtype)
+
+    matmul_kernel_descriptor_persistent[grid](
+        a,
+        b_t,
+        c,
+        bias,
+        M,
+        N,
+        K,
+        NUM_SMS=NUM_SMS,
+        HAS_BIAS=bias is not None,
+        **cfg,
+    )
+    return c
+
+
 def matmul_persistent(
     a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None
 ):
@@ -589,15 +733,17 @@ def mean_dim(
     return output
 
 
-def mm_batch_invariant(a, b):
-    return matmul_persistent(a, b)
+def mm_batch_invariant(a, b, bias=None):
+    if a.device.type == "xpu":
+        return matmul_descriptor_persistent(a, b, bias=bias)
+    return matmul_persistent(a, b, bias=bias)
 
 
 def matmul_batch_invariant(a, b, *, out=None):
     # torch.matmul can handle various dimensions
     # For 2D x 2D, it's the same as mm
     if a.ndim == 2 and b.ndim == 2:
-        result = matmul_persistent(a, b)
+        result = mm_batch_invariant(a, b)
         if out is not None:
             out.copy_(result)
             return out
@@ -609,7 +755,7 @@ def matmul_batch_invariant(a, b, *, out=None):
         hidden = a.shape[-1]
         out_dim = b.shape[-1]
         a_2d = a.reshape(-1, hidden)
-        result_2d = matmul_persistent(a_2d, b)
+        result_2d = mm_batch_invariant(a_2d, b)
         result = result_2d.reshape(batch_dims + (out_dim,))
         if out is not None:
             out.copy_(result)
@@ -732,7 +878,7 @@ def bmm_batch_invariant(a, b, *, out=None):
 
 
 def addmm_batch_invariant(bias, a, b):
-    return matmul_persistent(a, b, bias=bias)
+    return mm_batch_invariant(a, b, bias=bias)
 
 
 def _log_softmax_batch_invariant(input, dim, _half_to_float):
@@ -903,11 +1049,13 @@ _batch_invariant_LIB = None
 _fp16_block_size_n = 256
 _fp32_block_size_n = 128
 _fp32_num_stages = 3
+_NUM_SMS: int = 0
 
 
 def enable_batch_invariant_mode():
     global _batch_invariant_MODE, _batch_invariant_LIB
     global _fp16_block_size_n, _fp32_block_size_n, _fp32_num_stages
+    global _NUM_SMS
 
     if _batch_invariant_MODE:
         return
@@ -916,7 +1064,6 @@ def enable_batch_invariant_mode():
     _batch_invariant_LIB = torch.library.Library("aten", "IMPL")
 
     key = current_platform.dispatch_key
-
     if current_platform.is_cuda():
         if current_platform.is_device_capability_family(80):
             # SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the
@@ -945,24 +1092,38 @@ def enable_batch_invariant_mode():
             _fp32_block_size_n = 32
             _fp32_num_stages = 2
     elif current_platform.is_xpu():
+        # Tensor descriptors need a global-memory allocator; must be set outside
+        # torch.compile regions (triton.set_allocator modifies global state).
+        def _triton_alloc_fn(size: int, alignment: int, stream: int | None):
+            return torch.empty(size, device="xpu", dtype=torch.int8)
+
+        triton.set_allocator(_triton_alloc_fn)
+
+        _NUM_SMS = num_compute_units(0)
         _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::matmul", matmul_batch_invariant, key)
+        _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, key)
 
         _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, key, allow_override=True
-    )
-    torch.bmm = bmm_batch_invariant
+    # Softmax, log_softmax, mean, and bmm are already batch-invariant on XPU
+    # (oneDNN backend produces bitwise-identical results regardless of batch
+    # context). Only register these overrides on CUDA where they are needed.
+    if not current_platform.is_xpu():
+        _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, key, allow_override=True
+        )
+        torch.bmm = bmm_batch_invariant
 
     reduced_precision_val = (
         (False, False) if is_torch_equal_or_newer("2.10.0") else False

vllm/model_executor/determinism/batch_invariant_configs.py

@@ -327,3 +327,53 @@ def _get_matmul_config(
         "num_warps": m_config.num_warps,
         "num_stages": m_config.num_stages,
     }
+
+
+def _get_descriptor_matmul_config(
+    M: int, N: int, K: int, dtype: torch.dtype
+) -> dict[str, int]:
+    """Select a config for the XPU descriptor matmul from shape and dtype.
+
+    The default 128×128×64 config is optimized for large square-ish GEMMs but
+    wastes resources on skinny shapes (small M during decode, or very tall/thin
+    matrices during prefill).  Shape-dependent tuning closes the gap with oneDNN.
+
+    Only BLOCK_SIZE_M/N and the launch parameters vary with M; dtype alone
+    decides BLOCK_SIZE_K, so the K-reduction order never depends on batch size.
+    """
+    # fp32 uses smaller BLOCK_SIZE_K due to register pressure
+    block_k = 32 if dtype == torch.float32 else 64
+
+    if M <= 16:
+        # Decode: M=1-16. Tiny M means most of a 128-row tile is wasted.
+        # Use small M-block, wide N-block to maximize useful work per tile.
+        return {
+            "BLOCK_SIZE_M": 16,
+            "BLOCK_SIZE_N": 256,
+            "BLOCK_SIZE_K": block_k,
+            "GROUP_SIZE_M": 1,
+            "num_stages": 4,
+            "num_warps": 8,
+        }
+    elif M <= 64:
+        # Small batch decode or very short prefill.
+        return {
+            "BLOCK_SIZE_M": 32,
+            "BLOCK_SIZE_N": 128,
+            "BLOCK_SIZE_K": block_k,
+            "GROUP_SIZE_M": 4,
+            "num_stages": 4,
+            "num_warps": 8,
+        }
+    else:
+        # Medium and large prefill (M > 64). 64×128 tiles provide the best
+        # balance of register pressure vs parallelism on Intel XPU.
+        # M=2048, N=4096 → 32×32 = 1024 tiles, well above ~160 compute units.
+        return {
+            "BLOCK_SIZE_M": 64,
+            "BLOCK_SIZE_N": 128,
+            "BLOCK_SIZE_K": block_k,
+            "GROUP_SIZE_M": 8,
+            "num_stages": 3,
+            "num_warps": 8,
+        }

vllm/model_executor/layers/linear.py

@@ -228,7 +228,9 @@ class UnquantizedLinearMethod(LinearMethodBase):
         x: torch.Tensor,
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike():
+        if envs.VLLM_BATCH_INVARIANT and (
+            current_platform.is_cuda_alike() or current_platform.is_xpu()
+        ):
             return linear_batch_invariant(x, layer.weight, bias)
         return self._gemm_impl(layer, x, layer.weight, bias)
 

vllm/model_executor/layers/vocab_parallel_embedding.py

@@ -71,7 +71,9 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase):
         x: torch.Tensor,
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
-        if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike():
+        if envs.VLLM_BATCH_INVARIANT and (
+            current_platform.is_cuda_alike() or current_platform.is_xpu()
+        ):
             return linear_batch_invariant(x, layer.weight, bias)
         return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)