From 0188c010d2a4dd774b2ce77951208940b11e7a05 Mon Sep 17 00:00:00 2001 From: tzielinski-habana Date: Wed, 15 Jul 2026 19:32:03 +0300 Subject: [PATCH 01/13] Added batch invariant matmul and linear kernels Signed-off-by: tzielinski-habana --- tests/v1/determinism/test_batch_invariance.py | 5 --- vllm/model_executor/layers/batch_invariant.py | 40 +++++++++++++++++-- vllm/model_executor/layers/linear.py | 4 +- .../layers/vocab_parallel_embedding.py | 4 +- 4 files changed, 42 insertions(+), 11 deletions(-) diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index b2706ed89b7b..1d88d6fe6a4e 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -164,11 +164,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")) diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index f05998b1beb6..8156254c0387 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -138,7 +138,11 @@ def matmul_persistent( assert bias is None or bias.dim() == 1, ( "Currently assuming bias is 1D, let Horace know if you run into this" ) - NUM_SMS = num_compute_units(a.device.index) + # _NUM_SMS is set by enable_batch_invariant_mode() before torch.compile + # traces this function. Dynamo lifts it as a compile-time constant. + # The fallback (num_compute_units) only runs in eager-mode tests that + # skip enable_batch_invariant_mode(). + NUM_SMS = _NUM_SMS if _NUM_SMS > 0 else num_compute_units(0) M, K = a.shape K, N = b.shape dtype = a.dtype @@ -881,6 +885,30 @@ def rms_norm_batch_invariant( return output.reshape(original_shape) +def _linear_backward_xpu(input, grad_output, weight, output_mask): + """XPU implementation of aten::linear_backward.""" + grad_input = grad_weight = grad_bias = None + go_2d = grad_output.reshape(-1, grad_output.shape[-1]) + if output_mask[0]: + grad_input = torch.mm(go_2d, weight).reshape(input.shape) + if output_mask[1]: + grad_weight = torch.mm(go_2d.t(), input.reshape(-1, input.shape[-1])) + if output_mask[2]: + grad_bias = go_2d.sum(0) + return grad_input, grad_weight, grad_bias + + +def _matmul_backward_xpu(grad, self, other, mask): + """XPU implementation of aten::matmul_backward.""" + grad_self = ( + matmul_batch_invariant(grad, other.transpose(-1, -2)) if mask[0] else None + ) + grad_other = ( + matmul_batch_invariant(self.transpose(-1, -2), grad) if mask[1] else None + ) + return grad_self, grad_other + + def linear_batch_invariant(input, weight, bias=None): output = matmul_batch_invariant(input, weight.t()) @@ -892,16 +920,18 @@ def linear_batch_invariant(input, weight, bias=None): _batch_invariant_MODE = False _batch_invariant_LIB = None _fp16_block_size_n = 256 +_NUM_SMS: int = 0 def enable_batch_invariant_mode(): global _batch_invariant_MODE, _batch_invariant_LIB - global _fp16_block_size_n + global _fp16_block_size_n, _NUM_SMS if _batch_invariant_MODE: return _batch_invariant_MODE = True + _NUM_SMS = num_compute_units(0) _batch_invariant_LIB = torch.library.Library("aten", "IMPL") key = current_platform.dispatch_key @@ -925,8 +955,10 @@ def enable_batch_invariant_mode(): 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::matmul", matmul_batch_invariant, key) + _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, key) + _batch_invariant_LIB.impl("aten::linear_backward", _linear_backward_xpu, key) + _batch_invariant_LIB.impl("aten::matmul_backward", _matmul_backward_xpu, key) _fp16_block_size_n = 128 diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index d92b9fc7d008..f1087750c626 100644 --- a/vllm/model_executor/layers/linear.py +++ b/vllm/model_executor/layers/linear.py @@ -223,7 +223,9 @@ def apply( 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) diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 8d9a7ccbacae..d691adb43e3a 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/vllm/model_executor/layers/vocab_parallel_embedding.py @@ -70,7 +70,9 @@ def apply( 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) From dd2943ff8a9138194bacffa68428ecb3602dcaad Mon Sep 17 00:00:00 2001 From: tzielinski-habana Date: Thu, 14 May 2026 17:53:08 +0300 Subject: [PATCH 02/13] Added faster matmul kernel Signed-off-by: tzielinski-habana --- vllm/model_executor/layers/batch_invariant.py | 186 +++++++++++++++++- 1 file changed, 182 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index 8156254c0387..65cb5c40c715 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -129,6 +129,177 @@ 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) + NUM_SMS = num_compute_units(a.device.index) + + # Tensor descriptors need a global-memory allocator. + def alloc_fn(size: int, alignment: int, stream: int | None): + return torch.empty(size, device=a.device, dtype=torch.int8) + + triton.set_allocator(alloc_fn) + + def grid(META): + return ( + min( + NUM_SMS, + triton.cdiv(M, META["BLOCK_SIZE_M"]) + * triton.cdiv(N, META["BLOCK_SIZE_N"]), + ), + ) + + configs = { + torch.bfloat16: { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_stages": 3, + "num_warps": 8, + }, + torch.float16: { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 64, + "GROUP_SIZE_M": 8, + "num_stages": 3, + "num_warps": 8, + }, + torch.float32: { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + "num_stages": 3, + "num_warps": 8, + }, + } + + matmul_kernel_descriptor_persistent[grid]( + a, + b_t, + c, + bias, + M, + N, + K, + NUM_SMS=NUM_SMS, + HAS_BIAS=bias is not None, + **configs[dtype], + ) + return c + + def matmul_persistent( a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None ): @@ -589,15 +760,22 @@ def mean_dim( return output +def _matmul_dispatch(a, b, bias=None): + """Dispatch to the best matmul kernel based on device type.""" + if a.device.type == "xpu": + return matmul_descriptor_persistent(a, b, bias=bias) + return matmul_persistent(a, b, bias=bias) + + def mm_batch_invariant(a, b): - return matmul_persistent(a, b) + return _matmul_dispatch(a, b) 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 = _matmul_dispatch(a, b) if out is not None: out.copy_(result) return out @@ -609,7 +787,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 = _matmul_dispatch(a_2d, b) result = result_2d.reshape(batch_dims + (out_dim,)) if out is not None: out.copy_(result) @@ -732,7 +910,7 @@ def bmm_batch_invariant(a, b, *, out=None): def addmm_batch_invariant(bias, a, b): - return matmul_persistent(a, b, bias=bias) + return _matmul_dispatch(a, b, bias=bias) def _log_softmax_batch_invariant(input, dim, _half_to_float): From 16e12ac79a2f10c646f2687794e0bd73f3b9f111 Mon Sep 17 00:00:00 2001 From: tzielinski-habana Date: Wed, 10 Jun 2026 18:17:23 +0300 Subject: [PATCH 03/13] Fixed matmul_descriptor_persistent kernel Signed-off-by: tzielinski-habana --- vllm/model_executor/layers/batch_invariant.py | 20 +++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index 65cb5c40c715..35a2b58159e8 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -241,13 +241,11 @@ def matmul_descriptor_persistent( b_t = b.t().contiguous() c = torch.empty((M, N), device=a.device, dtype=dtype) - NUM_SMS = num_compute_units(a.device.index) - - # Tensor descriptors need a global-memory allocator. - def alloc_fn(size: int, alignment: int, stream: int | None): - return torch.empty(size, device=a.device, dtype=torch.int8) - - triton.set_allocator(alloc_fn) + # _NUM_SMS is set by enable_batch_invariant_mode() before torch.compile + # traces this function. Dynamo lifts it as a compile-time constant. + # The fallback (num_compute_units) only runs in eager-mode tests that + # skip enable_batch_invariant_mode(). + NUM_SMS = _NUM_SMS if _NUM_SMS > 0 else num_compute_units(0) def grid(META): return ( @@ -1112,8 +1110,14 @@ def enable_batch_invariant_mode(): _NUM_SMS = num_compute_units(0) _batch_invariant_LIB = torch.library.Library("aten", "IMPL") - key = current_platform.dispatch_key + # 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) + + 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 From d97066dcbe0dcc99b45db936ce05106fcc0f5f54 Mon Sep 17 00:00:00 2001 From: tzielinski-habana Date: Fri, 17 Jul 2026 18:19:08 +0300 Subject: [PATCH 04/13] Removed xfail from test Signed-off-by: tzielinski-habana --- tests/v1/determinism/test_batch_invariance.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index 1d88d6fe6a4e..77af5a57cc23 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -17,7 +17,6 @@ import vllm.envs as envs from vllm import LLM, SamplingParams -from vllm.platforms import current_platform @skip_unsupported @@ -51,11 +50,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) From 1f1574447892b98b24d78488b80fc3b802ee75ca Mon Sep 17 00:00:00 2001 From: tzielinski-habana Date: Mon, 20 Jul 2026 18:10:51 +0300 Subject: [PATCH 05/13] Added documentation Signed-off-by: tzielinski-habana --- docs/features/batch_invariance.md | 24 +++++++++++++++++++++++- 1 file changed, 23 insertions(+), 1 deletion(-) diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index 5a1fe71a1eab..5da4b55e3f9c 100644 --- a/docs/features/batch_invariance.md +++ b/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 From 864fdb47b5e4bd3b89dd6f43cc4a0f87e8b04b95 Mon Sep 17 00:00:00 2001 From: tzielinski-habana Date: Mon, 27 Jul 2026 16:31:50 +0300 Subject: [PATCH 06/13] Changed handling of _NUM_SMS Signed-off-by: tzielinski-habana --- vllm/model_executor/layers/batch_invariant.py | 17 ++++++----------- 1 file changed, 6 insertions(+), 11 deletions(-) diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index 6a2e3b02484f..e9d69bd8774b 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -241,11 +241,10 @@ def matmul_descriptor_persistent( b_t = b.t().contiguous() c = torch.empty((M, N), device=a.device, dtype=dtype) - # _NUM_SMS is set by enable_batch_invariant_mode() before torch.compile - # traces this function. Dynamo lifts it as a compile-time constant. - # The fallback (num_compute_units) only runs in eager-mode tests that - # skip enable_batch_invariant_mode(). - NUM_SMS = _NUM_SMS if _NUM_SMS > 0 else num_compute_units(0) + # This path is XPU-only (via _matmul_dispatch). _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 ( @@ -307,11 +306,7 @@ def matmul_persistent( assert bias is None or bias.dim() == 1, ( "Currently assuming bias is 1D, let Horace know if you run into this" ) - # _NUM_SMS is set by enable_batch_invariant_mode() before torch.compile - # traces this function. Dynamo lifts it as a compile-time constant. - # The fallback (num_compute_units) only runs in eager-mode tests that - # skip enable_batch_invariant_mode(). - NUM_SMS = _NUM_SMS if _NUM_SMS > 0 else num_compute_units(0) + NUM_SMS = num_compute_units(a.device.index) M, K = a.shape K, N = b.shape dtype = a.dtype @@ -1112,7 +1107,6 @@ def enable_batch_invariant_mode(): return _batch_invariant_MODE = True - _NUM_SMS = num_compute_units(0) _batch_invariant_LIB = torch.library.Library("aten", "IMPL") # Tensor descriptors need a global-memory allocator; must be set outside @@ -1140,6 +1134,7 @@ def _triton_alloc_fn(size: int, alignment: int, stream: int | None): _fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128 elif current_platform.is_xpu(): + _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) _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, key) From 9236cc579774e14e7796997af52c1c0055d2ebc1 Mon Sep 17 00:00:00 2001 From: tzielinski-habana Date: Mon, 27 Jul 2026 16:33:02 +0300 Subject: [PATCH 07/13] Set TRITON_ATTN as default attn backend for XPU with batch invariance Signed-off-by: tzielinski-habana --- vllm/engine/arg_utils.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index c947e2035cf1..8be673a8c71f 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -2247,6 +2247,19 @@ def create_engine_config( 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 ( From 5caa52dc3f89cd1d726e0156ba165b7b80281857 Mon Sep 17 00:00:00 2001 From: tzielinski-habana Date: Mon, 27 Jul 2026 17:26:53 +0300 Subject: [PATCH 08/13] Moved triton allocator under XPU branch Signed-off-by: tzielinski-habana --- vllm/model_executor/layers/batch_invariant.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index e9d69bd8774b..be260f9a09dd 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -1109,13 +1109,6 @@ def enable_batch_invariant_mode(): _batch_invariant_MODE = True _batch_invariant_LIB = torch.library.Library("aten", "IMPL") - # 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) - key = current_platform.dispatch_key if current_platform.is_cuda(): if current_platform.is_device_capability_family(80): @@ -1134,6 +1127,13 @@ def _triton_alloc_fn(size: int, alignment: int, stream: int | None): _fp16_block_size_n = 256 if get_max_shared_memory_bytes() > 106496 else 128 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) From daea0b688587794f8fd5e1388adabf5746db1daa Mon Sep 17 00:00:00 2001 From: tzielinski-habana Date: Mon, 3 Aug 2026 19:12:33 +0300 Subject: [PATCH 09/13] Better config for matmul_kernel_descriptor_persistent Signed-off-by: tzielinski-habana --- vllm/model_executor/layers/batch_invariant.py | 104 +++++++++++------- 1 file changed, 66 insertions(+), 38 deletions(-) diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index be260f9a09dd..2a2aa08887df 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -214,6 +214,53 @@ def matmul_kernel_descriptor_persistent( c_desc.store([offs_cm, offs_cn], c) +def _get_descriptor_matmul_config( + M: int, N: int, K: int, dtype: torch.dtype +) -> dict[str, int]: + """Select Triton kernel config based on problem 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. + """ + # 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, + } + + def matmul_descriptor_persistent( a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None ): @@ -255,32 +302,7 @@ def grid(META): ), ) - configs = { - torch.bfloat16: { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 64, - "GROUP_SIZE_M": 8, - "num_stages": 3, - "num_warps": 8, - }, - torch.float16: { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 64, - "GROUP_SIZE_M": 8, - "num_stages": 3, - "num_warps": 8, - }, - torch.float32: { - "BLOCK_SIZE_M": 128, - "BLOCK_SIZE_N": 128, - "BLOCK_SIZE_K": 32, - "GROUP_SIZE_M": 8, - "num_stages": 3, - "num_warps": 8, - }, - } + cfg = _get_descriptor_matmul_config(M, N, K, dtype) matmul_kernel_descriptor_persistent[grid]( a, @@ -292,7 +314,7 @@ def grid(META): K, NUM_SMS=NUM_SMS, HAS_BIAS=bias is not None, - **configs[dtype], + **cfg, ) return c @@ -1144,17 +1166,23 @@ def _triton_alloc_fn(size: int, alignment: int, stream: int | None): _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 From c43b879a819e92415b5f322647c0935dfbfd41d7 Mon Sep 17 00:00:00 2001 From: tzielinski-habana Date: Wed, 5 Aug 2026 16:25:47 +0300 Subject: [PATCH 10/13] Removed backward kernels Signed-off-by: tzielinski-habana --- vllm/model_executor/layers/batch_invariant.py | 26 ------------------- 1 file changed, 26 deletions(-) diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index 2a2aa08887df..d7f2042cd687 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -1083,30 +1083,6 @@ def rms_norm_batch_invariant( return output.reshape(original_shape) -def _linear_backward_xpu(input, grad_output, weight, output_mask): - """XPU implementation of aten::linear_backward.""" - grad_input = grad_weight = grad_bias = None - go_2d = grad_output.reshape(-1, grad_output.shape[-1]) - if output_mask[0]: - grad_input = torch.mm(go_2d, weight).reshape(input.shape) - if output_mask[1]: - grad_weight = torch.mm(go_2d.t(), input.reshape(-1, input.shape[-1])) - if output_mask[2]: - grad_bias = go_2d.sum(0) - return grad_input, grad_weight, grad_bias - - -def _matmul_backward_xpu(grad, self, other, mask): - """XPU implementation of aten::matmul_backward.""" - grad_self = ( - matmul_batch_invariant(grad, other.transpose(-1, -2)) if mask[0] else None - ) - grad_other = ( - matmul_batch_invariant(self.transpose(-1, -2), grad) if mask[1] else None - ) - return grad_self, grad_other - - def linear_batch_invariant(input, weight, bias=None): output = matmul_batch_invariant(input, weight.t()) @@ -1161,8 +1137,6 @@ def _triton_alloc_fn(size: int, alignment: int, stream: int | None): _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) - _batch_invariant_LIB.impl("aten::linear_backward", _linear_backward_xpu, key) - _batch_invariant_LIB.impl("aten::matmul_backward", _matmul_backward_xpu, key) _fp16_block_size_n = 128 From f8148e0d779b724cd58ce78dfa9b095fc3a86a91 Mon Sep 17 00:00:00 2001 From: tzielinski-habana Date: Thu, 13 Aug 2026 17:31:15 +0300 Subject: [PATCH 11/13] Applied review suggestions Signed-off-by: tzielinski-habana --- vllm/model_executor/layers/batch_invariant.py | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/vllm/model_executor/layers/batch_invariant.py b/vllm/model_executor/layers/batch_invariant.py index d7f2042cd687..5f295e855163 100644 --- a/vllm/model_executor/layers/batch_invariant.py +++ b/vllm/model_executor/layers/batch_invariant.py @@ -288,9 +288,9 @@ def matmul_descriptor_persistent( b_t = b.t().contiguous() c = torch.empty((M, N), device=a.device, dtype=dtype) - # This path is XPU-only (via _matmul_dispatch). _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. + # 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): @@ -775,22 +775,17 @@ def mean_dim( return output -def _matmul_dispatch(a, b, bias=None): - """Dispatch to the best matmul kernel based on device type.""" +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 mm_batch_invariant(a, b): - return _matmul_dispatch(a, b) - - 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_dispatch(a, b) + result = mm_batch_invariant(a, b) if out is not None: out.copy_(result) return out @@ -802,7 +797,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_dispatch(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) @@ -925,7 +920,7 @@ def bmm_batch_invariant(a, b, *, out=None): def addmm_batch_invariant(bias, a, b): - return _matmul_dispatch(a, b, bias=bias) + return mm_batch_invariant(a, b, bias=bias) def _log_softmax_batch_invariant(input, dim, _half_to_float): From 53f2a7796e9a7c3097d057bc370debe7a42fd89a Mon Sep 17 00:00:00 2001 From: tzielinski-habana Date: Wed, 26 Aug 2026 15:01:08 +0300 Subject: [PATCH 12/13] refactor: move XPU descriptor matmul config to configs module _get_descriptor_matmul_config now sits next to _get_matmul_config in batch_invariant_configs.py, which became the home for persistent-matmul config selection in #53247. Pure move: the returned configs are identical for every M and dtype. Also record why BLOCK_SIZE_K is derived from dtype alone: an M-dependent BLOCK_K would change the K-reduction order and break batch invariance. Co-Authored-By: Claude Opus 5 Signed-off-by: tzielinski-habana --- .../determinism/batch_invariant.py | 48 +----------------- .../determinism/batch_invariant_configs.py | 50 +++++++++++++++++++ 2 files changed, 51 insertions(+), 47 deletions(-) diff --git a/vllm/model_executor/determinism/batch_invariant.py b/vllm/model_executor/determinism/batch_invariant.py index 64bf39f6a1e7..e2579f30b3cb 100644 --- a/vllm/model_executor/determinism/batch_invariant.py +++ b/vllm/model_executor/determinism/batch_invariant.py @@ -9,6 +9,7 @@ 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, ) @@ -218,53 +219,6 @@ def matmul_kernel_descriptor_persistent( c_desc.store([offs_cm, offs_cn], c) -def _get_descriptor_matmul_config( - M: int, N: int, K: int, dtype: torch.dtype -) -> dict[str, int]: - """Select Triton kernel config based on problem 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. - """ - # 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, - } - - def matmul_descriptor_persistent( a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None ): diff --git a/vllm/model_executor/determinism/batch_invariant_configs.py b/vllm/model_executor/determinism/batch_invariant_configs.py index 4570a491cec4..8f222988ead3 100644 --- a/vllm/model_executor/determinism/batch_invariant_configs.py +++ b/vllm/model_executor/determinism/batch_invariant_configs.py @@ -248,3 +248,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, + } From b2137aaf1c99cc1f9b211e802a06c9310ec6dda4 Mon Sep 17 00:00:00 2001 From: Tomasz Zielinski <85164140+tzielinski-habana@users.noreply.github.com> Date: Tue, 1 Sep 2026 14:45:02 +0200 Subject: [PATCH 13/13] Update docs/features/batch_invariance.md Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Signed-off-by: Tomasz Zielinski <85164140+tzielinski-habana@users.noreply.github.com> --- docs/features/batch_invariance.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index 9b5348b1aa31..73f7c41877e2 100644 --- a/docs/features/batch_invariance.md +++ b/docs/features/batch_invariance.md @@ -38,7 +38,7 @@ Or via the CLI: ```bash VLLM_BATCH_INVARIANT=1 vllm serve Qwen/Qwen3-1.7B \ - --attention-config '{"backend": "TRITON_ATTN"}' + --attention-config.backend TRITON_ATTN ``` ## Enabling Batch Invariance