diff --git a/docs/features/batch_invariance.md b/docs/features/batch_invariance.md index ddfef8efbd..73f7c41877 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 diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index 37fd5cba6a..4a86a5826b 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/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")) diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 66373df6fc..e89360505f 100644 --- a/vllm/engine/arg_utils.py +++ b/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 ( diff --git a/vllm/model_executor/determinism/batch_invariant.py b/vllm/model_executor/determinism/batch_invariant.py index 867b4772a1..e2579f30b3 100644 --- a/vllm/model_executor/determinism/batch_invariant.py +++ b/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 diff --git a/vllm/model_executor/determinism/batch_invariant_configs.py b/vllm/model_executor/determinism/batch_invariant_configs.py index 86dd51147f..39cda984f7 100644 --- a/vllm/model_executor/determinism/batch_invariant_configs.py +++ b/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, + } diff --git a/vllm/model_executor/layers/linear.py b/vllm/model_executor/layers/linear.py index a33f3a1166..fdccabae6f 100644 --- a/vllm/model_executor/layers/linear.py +++ b/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) diff --git a/vllm/model_executor/layers/vocab_parallel_embedding.py b/vllm/model_executor/layers/vocab_parallel_embedding.py index 70c6516e2e..e2d9072c25 100644 --- a/vllm/model_executor/layers/vocab_parallel_embedding.py +++ b/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)