@tduncan / Nvidia-Optimization-Stack-VLLM / commits / ffe3bb3c72

[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>

Tomasz Zielinski committed Sep 2, 2026 at 15:31 UTC ffe3bb3c723d9e5e5f5a0c8aeab1d857a4975b8b
7 files changed +272 -33
docs/features/batch_invariance.md
+23 -1
@@ -17,7 +17,29 @@ Batch invariance is crucial for several use cases:
17
18 ## Hardware Requirements
19
20 -Batch invariance requires NVIDIA GPUs with compute capability 8.0 or higher.
20 +Batch invariance is supported on the following platforms:
21 +
22 +- NVIDIA GPUs with compute capability 8.0 or higher.
23 +- Intel XPUs with Triton support.
24 +
25 +### Attention Backend Selection for XPU
26 +
27 +On XPU, Triton Attention backend is required for batch invariance.
28 +Select this backend using Qwen/Qwen3-1.7B model as an example:
29 +
30 +```python
31 +llm = LLM(
32 + model="Qwen/Qwen3-1.7B",
33 + attention_config={"backend": "TRITON_ATTN"},
34 +)
35 +```
36 +
37 +Or via the CLI:
38 +
39 +```bash
40 +VLLM_BATCH_INVARIANT=1 vllm serve Qwen/Qwen3-1.7B \
41 + --attention-config.backend TRITON_ATTN
42 +```
43
44 ## Enabling Batch Invariance
45
tests/v1/determinism/test_batch_invariance.py
-11
@@ -17,7 +17,6 @@ from utils import (
17
18 import vllm.envs as envs
19 from vllm import LLM, SamplingParams
20 -from vllm.platforms import current_platform
20
21
22 @skip_unsupported
@@ -53,11 +52,6 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
52 seed.
53 - Keep max_tokens and max_model_len bounded for speed and memory use.
54 """
56 - # Not all batch-invariant kernels are registered on XPU yet
57 - # (e.g. attention, custom ops), so e2e determinism is not guaranteed.
58 - if current_platform.is_xpu():
59 - pytest.xfail("Not all batch-invariant kernels registered on XPU yet")
60 -
55 seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
56 random.seed(seed)
57
@@ -177,11 +171,6 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
171 block_m,
172 block_n,
173 ):
180 - # Not all batch-invariant kernels are registered on XPU yet
181 - # (e.g. attention, custom ops), so e2e determinism is not guaranteed.
182 - if current_platform.is_xpu():
183 - pytest.xfail("Not all batch-invariant kernels registered on XPU yet")
184 -
174 seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
175 random.seed(seed)
176 tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
vllm/engine/arg_utils.py
+13
@@ -2435,6 +2435,19 @@ class EngineArgs:
2435 self.attention_backend
2436 )
2437
2438 + # Batch-invariant mode requires deterministic attention behavior.
2439 + # If no backend is explicitly requested, prefer Triton Attention.
2440 + if (
2441 + envs.VLLM_BATCH_INVARIANT
2442 + and attention_config.backend is None
2443 + and current_platform.is_xpu()
2444 + ):
2445 + attention_config.backend = AttentionBackendEnum.TRITON_ATTN
2446 + logger.info(
2447 + "VLLM_BATCH_INVARIANT is enabled and no attention backend was "
2448 + "specified; defaulting to TRITON_ATTN."
2449 + )
2450 +
2451 # TurboQuant requires FlashAttention 2 — FA3 boundary layers assert
2452 # FlashAttentionImpl which fails with TurboQuantAttentionImpl.
2453 if resolved_cache_dtype.startswith("turboquant_") and (
vllm/model_executor/determinism/batch_invariant.py
+180 -19
@@ -9,6 +9,7 @@ import torch
9
10 import vllm.envs as envs
11 from vllm.model_executor.determinism.batch_invariant_configs import (
12 + _get_descriptor_matmul_config,
13 _get_matmul_config,
14 resolve_tuned_matmul_configs,
15 )
@@ -133,6 +134,149 @@ def matmul_kernel_persistent(
134 tl.store(c_ptrs, c, mask=c_mask)
135
136
137 +@triton.jit(launch_metadata=_matmul_launch_metadata)
138 +def matmul_kernel_descriptor_persistent(
139 + a_ptr,
140 + b_ptr,
141 + c_ptr, #
142 + bias_ptr,
143 + M,
144 + N,
145 + K, #
146 + BLOCK_SIZE_M: tl.constexpr, #
147 + BLOCK_SIZE_N: tl.constexpr, #
148 + BLOCK_SIZE_K: tl.constexpr, #
149 + GROUP_SIZE_M: tl.constexpr, #
150 + NUM_SMS: tl.constexpr, #
151 + HAS_BIAS: tl.constexpr,
152 +):
153 + """Persistent matmul using tensor descriptors for 2D block I/O.
154 +
155 + Expects b_ptr to point to a transposed B matrix of shape [N, K] with
156 + row-major (K-contiguous) layout. The dot product transposes each loaded
157 + B-tile back: dot(A_tile, B_tile.T).
158 +
159 + ~3x faster than the pointer-based persistent kernel on Intel XPU because
160 + tensor descriptors leverage hardware 2D block load/store with automatic
161 + bounds checking (no explicit masks needed).
162 + """
163 + dtype = c_ptr.dtype.element_ty
164 + start_pid = tl.program_id(axis=0)
165 + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
166 + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
167 + k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
168 + num_tiles = num_pid_m * num_pid_n
169 +
170 + a_desc = tl.make_tensor_descriptor(
171 + a_ptr,
172 + shape=[M, K],
173 + strides=[K, 1],
174 + block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K],
175 + )
176 + b_desc = tl.make_tensor_descriptor(
177 + b_ptr,
178 + shape=[N, K],
179 + strides=[K, 1],
180 + block_shape=[BLOCK_SIZE_N, BLOCK_SIZE_K],
181 + )
182 + c_desc = tl.make_tensor_descriptor(
183 + c_ptr,
184 + shape=[M, N],
185 + strides=[N, 1],
186 + block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_N],
187 + )
188 +
189 + tile_id_c = start_pid - NUM_SMS
190 + num_pid_in_group = GROUP_SIZE_M * num_pid_n
191 +
192 + for tile_id in tl.range(start_pid, num_tiles, NUM_SMS, flatten=True):
193 + pid_m, pid_n = _compute_pid(tile_id, num_pid_in_group, num_pid_m, GROUP_SIZE_M)
194 + offs_am = pid_m * BLOCK_SIZE_M
195 + offs_bn = pid_n * BLOCK_SIZE_N
196 +
197 + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
198 + for ki in range(k_tiles):
199 + offs_k = ki * BLOCK_SIZE_K
200 + a = a_desc.load([offs_am, offs_k])
201 + b = b_desc.load([offs_bn, offs_k])
202 + accumulator = tl.dot(a, b.T, accumulator)
203 +
204 + tile_id_c += NUM_SMS
205 + pid_m, pid_n = _compute_pid(
206 + tile_id_c, num_pid_in_group, num_pid_m, GROUP_SIZE_M
207 + )
208 + offs_cm = pid_m * BLOCK_SIZE_M
209 + offs_cn = pid_n * BLOCK_SIZE_N
210 +
211 + if HAS_BIAS:
212 + bias_offs = offs_cn + tl.arange(0, BLOCK_SIZE_N)
213 + bias_vals = tl.load(bias_ptr + bias_offs, mask=bias_offs < N, other=0.0).to(
214 + tl.float32
215 + )
216 + accumulator += bias_vals[None, :]
217 +
218 + c = accumulator.to(dtype)
219 + c_desc.store([offs_cm, offs_cn], c)
220 +
221 +
222 +def matmul_descriptor_persistent(
223 + a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None
224 +):
225 + """Persistent matmul using tensor descriptors (Intel XPU fast path).
226 +
227 + Args:
228 + a: Input matrix [M, K], must be contiguous.
229 + b: Weight matrix [K, N] (standard layout — transposed internally).
230 + bias: Optional 1D bias vector [N].
231 +
232 + Returns:
233 + Output matrix [M, N] with dtype matching the inputs.
234 + """
235 + assert a.shape[1] == b.shape[0], "Incompatible dimensions"
236 + assert a.dtype == b.dtype, "Incompatible dtypes"
237 + assert bias is None or bias.dim() == 1, "bias must be 1D"
238 +
239 + M, K = a.shape
240 + _, N = b.shape
241 + dtype = a.dtype
242 +
243 + # Tensor descriptors require contiguous row-major layout.
244 + a = a.contiguous()
245 + # Descriptor kernel expects B in [N, K] layout (K-contiguous).
246 + b_t = b.t().contiguous()
247 +
248 + c = torch.empty((M, N), device=a.device, dtype=dtype)
249 + # This path is XPU-only. _NUM_SMS is set by enable_batch_invariant_mode()
250 + # before torch.compile traces this function, so Dynamo can lift it
251 + # as a compile-time constant.
252 + NUM_SMS = _NUM_SMS if _NUM_SMS > 0 else num_compute_units(a.device.index)
253 +
254 + def grid(META):
255 + return (
256 + min(
257 + NUM_SMS,
258 + triton.cdiv(M, META["BLOCK_SIZE_M"])
259 + * triton.cdiv(N, META["BLOCK_SIZE_N"]),
260 + ),
261 + )
262 +
263 + cfg = _get_descriptor_matmul_config(M, N, K, dtype)
264 +
265 + matmul_kernel_descriptor_persistent[grid](
266 + a,
267 + b_t,
268 + c,
269 + bias,
270 + M,
271 + N,
272 + K,
273 + NUM_SMS=NUM_SMS,
274 + HAS_BIAS=bias is not None,
275 + **cfg,
276 + )
277 + return c
278 +
279 +
280 def matmul_persistent(
281 a: torch.Tensor, b: torch.Tensor, bias: torch.Tensor | None = None
282 ):
@@ -589,15 +733,17 @@ def mean_dim(
733 return output
734
735
592 -def mm_batch_invariant(a, b):
593 - return matmul_persistent(a, b)
736 +def mm_batch_invariant(a, b, bias=None):
737 + if a.device.type == "xpu":
738 + return matmul_descriptor_persistent(a, b, bias=bias)
739 + return matmul_persistent(a, b, bias=bias)
740
741
742 def matmul_batch_invariant(a, b, *, out=None):
743 # torch.matmul can handle various dimensions
744 # For 2D x 2D, it's the same as mm
745 if a.ndim == 2 and b.ndim == 2:
600 - result = matmul_persistent(a, b)
746 + result = mm_batch_invariant(a, b)
747 if out is not None:
748 out.copy_(result)
749 return out
@@ -609,7 +755,7 @@ def matmul_batch_invariant(a, b, *, out=None):
755 hidden = a.shape[-1]
756 out_dim = b.shape[-1]
757 a_2d = a.reshape(-1, hidden)
612 - result_2d = matmul_persistent(a_2d, b)
758 + result_2d = mm_batch_invariant(a_2d, b)
759 result = result_2d.reshape(batch_dims + (out_dim,))
760 if out is not None:
761 out.copy_(result)
@@ -732,7 +878,7 @@ def bmm_batch_invariant(a, b, *, out=None):
878
879
880 def addmm_batch_invariant(bias, a, b):
735 - return matmul_persistent(a, b, bias=bias)
881 + return mm_batch_invariant(a, b, bias=bias)
882
883
884 def _log_softmax_batch_invariant(input, dim, _half_to_float):
@@ -903,11 +1049,13 @@ _batch_invariant_LIB = None
1049 _fp16_block_size_n = 256
1050 _fp32_block_size_n = 128
1051 _fp32_num_stages = 3
1052 +_NUM_SMS: int = 0
1053
1054
1055 def enable_batch_invariant_mode():
1056 global _batch_invariant_MODE, _batch_invariant_LIB
1057 global _fp16_block_size_n, _fp32_block_size_n, _fp32_num_stages
1058 + global _NUM_SMS
1059
1060 if _batch_invariant_MODE:
1061 return
@@ -916,7 +1064,6 @@ def enable_batch_invariant_mode():
1064 _batch_invariant_LIB = torch.library.Library("aten", "IMPL")
1065
1066 key = current_platform.dispatch_key
919 -
1067 if current_platform.is_cuda():
1068 if current_platform.is_device_capability_family(80):
1069 # SM80 (Ampere) cannot rely on cuBLASLt-only determinism; install the
@@ -945,24 +1092,38 @@ def enable_batch_invariant_mode():
1092 _fp32_block_size_n = 32
1093 _fp32_num_stages = 2
1094 elif current_platform.is_xpu():
1095 + # Tensor descriptors need a global-memory allocator; must be set outside
1096 + # torch.compile regions (triton.set_allocator modifies global state).
1097 + def _triton_alloc_fn(size: int, alignment: int, stream: int | None):
1098 + return torch.empty(size, device="xpu", dtype=torch.int8)
1099 +
1100 + triton.set_allocator(_triton_alloc_fn)
1101 +
1102 + _NUM_SMS = num_compute_units(0)
1103 _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, key)
1104 _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, key)
950 - # TODO: register matmul and linear for XPU
951 - # once suitable Triton kernels are implemented
1105 + _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, key)
1106 + _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, key)
1107
1108 _fp16_block_size_n = 128
1109
955 - _batch_invariant_LIB.impl("aten::_log_softmax", _log_softmax_batch_invariant, key)
956 - _batch_invariant_LIB.impl("aten::softmax", softmax_batch_invariant, key)
957 - _batch_invariant_LIB.impl("aten::_softmax", softmax_batch_invariant, key)
958 - _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, key)
959 - # torch 2.12+ registers a built-in Triton bmm kernel for CUDA
960 - # (torch._native.ops.bmm_outer_product), so we need allow_override
961 - # to replace it at the dispatcher level.
962 - _batch_invariant_LIB.impl(
963 - "aten::bmm", bmm_batch_invariant, key, allow_override=True
964 - )
965 - torch.bmm = bmm_batch_invariant
1110 + # Softmax, log_softmax, mean, and bmm are already batch-invariant on XPU
1111 + # (oneDNN backend produces bitwise-identical results regardless of batch
1112 + # context). Only register these overrides on CUDA where they are needed.
1113 + if not current_platform.is_xpu():
1114 + _batch_invariant_LIB.impl(
1115 + "aten::_log_softmax", _log_softmax_batch_invariant, key
1116 + )
1117 + _batch_invariant_LIB.impl("aten::softmax", softmax_batch_invariant, key)
1118 + _batch_invariant_LIB.impl("aten::_softmax", softmax_batch_invariant, key)
1119 + _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, key)
1120 + # torch 2.12+ registers a built-in Triton bmm kernel for CUDA
1121 + # (torch._native.ops.bmm_outer_product), so we need allow_override
1122 + # to replace it at the dispatcher level.
1123 + _batch_invariant_LIB.impl(
1124 + "aten::bmm", bmm_batch_invariant, key, allow_override=True
1125 + )
1126 + torch.bmm = bmm_batch_invariant
1127
1128 reduced_precision_val = (
1129 (False, False) if is_torch_equal_or_newer("2.10.0") else False
vllm/model_executor/determinism/batch_invariant_configs.py
+50
@@ -327,3 +327,53 @@ def _get_matmul_config(
327 "num_warps": m_config.num_warps,
328 "num_stages": m_config.num_stages,
329 }
330 +
331 +
332 +def _get_descriptor_matmul_config(
333 + M: int, N: int, K: int, dtype: torch.dtype
334 +) -> dict[str, int]:
335 + """Select a config for the XPU descriptor matmul from shape and dtype.
336 +
337 + The default 128×128×64 config is optimized for large square-ish GEMMs but
338 + wastes resources on skinny shapes (small M during decode, or very tall/thin
339 + matrices during prefill). Shape-dependent tuning closes the gap with oneDNN.
340 +
341 + Only BLOCK_SIZE_M/N and the launch parameters vary with M; dtype alone
342 + decides BLOCK_SIZE_K, so the K-reduction order never depends on batch size.
343 + """
344 + # fp32 uses smaller BLOCK_SIZE_K due to register pressure
345 + block_k = 32 if dtype == torch.float32 else 64
346 +
347 + if M <= 16:
348 + # Decode: M=1-16. Tiny M means most of a 128-row tile is wasted.
349 + # Use small M-block, wide N-block to maximize useful work per tile.
350 + return {
351 + "BLOCK_SIZE_M": 16,
352 + "BLOCK_SIZE_N": 256,
353 + "BLOCK_SIZE_K": block_k,
354 + "GROUP_SIZE_M": 1,
355 + "num_stages": 4,
356 + "num_warps": 8,
357 + }
358 + elif M <= 64:
359 + # Small batch decode or very short prefill.
360 + return {
361 + "BLOCK_SIZE_M": 32,
362 + "BLOCK_SIZE_N": 128,
363 + "BLOCK_SIZE_K": block_k,
364 + "GROUP_SIZE_M": 4,
365 + "num_stages": 4,
366 + "num_warps": 8,
367 + }
368 + else:
369 + # Medium and large prefill (M > 64). 64×128 tiles provide the best
370 + # balance of register pressure vs parallelism on Intel XPU.
371 + # M=2048, N=4096 → 32×32 = 1024 tiles, well above ~160 compute units.
372 + return {
373 + "BLOCK_SIZE_M": 64,
374 + "BLOCK_SIZE_N": 128,
375 + "BLOCK_SIZE_K": block_k,
376 + "GROUP_SIZE_M": 8,
377 + "num_stages": 3,
378 + "num_warps": 8,
379 + }
vllm/model_executor/layers/linear.py
+3 -1
@@ -228,7 +228,9 @@ class UnquantizedLinearMethod(LinearMethodBase):
228 x: torch.Tensor,
229 bias: torch.Tensor | None = None,
230 ) -> torch.Tensor:
231 - if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike():
231 + if envs.VLLM_BATCH_INVARIANT and (
232 + current_platform.is_cuda_alike() or current_platform.is_xpu()
233 + ):
234 return linear_batch_invariant(x, layer.weight, bias)
235 return self._gemm_impl(layer, x, layer.weight, bias)
236
vllm/model_executor/layers/vocab_parallel_embedding.py
+3 -1
@@ -71,7 +71,9 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase):
71 x: torch.Tensor,
72 bias: torch.Tensor | None = None,
73 ) -> torch.Tensor:
74 - if envs.VLLM_BATCH_INVARIANT and current_platform.is_cuda_alike():
74 + if envs.VLLM_BATCH_INVARIANT and (
75 + current_platform.is_cuda_alike() or current_platform.is_xpu()
76 + ):
77 return linear_batch_invariant(x, layer.weight, bias)
78 return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)
79