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

[Feat][XPU] VLLM_BATCH_INVARIANT support for Dense/MoE models (#55881)

Signed-off-by: Tony Lin <tony.lin@intel.com> Co-authored-by: Flora Feng <4florafeng@gmail.com>

Tony Lin committed Sep 22, 2026 at 23:06 UTC e4340e41c9a32b2247cc2442b35aa91b081c424f
10 files changed +619 -8
.buildkite/intel_jobs/misc_intel.yaml
+29 -1
@@ -2,6 +2,34 @@ group: Miscellaneous Intel
2 depends_on:
3 - image-build-xpu
4 steps:
5 +- label: XPU Batch Invariance
6 + key: xpu-batch-invariance
7 + timeout_in_minutes: 60
8 + num_devices: 2
9 + device: intel_gpu
10 + agent_tags:
11 + label: production
12 + gpu: 2+
13 + mem: 24+
14 + no_plugin: true
15 + working_dir: "."
16 + env:
17 + REGISTRY: "public.ecr.aws/q9t5s3a7"
18 + REPO: "vllm-ci-test-repo"
19 + VLLM_TEST_DEVICE: "xpu"
20 + source_file_dependencies:
21 + - vllm/
22 + - tests/v1/determinism/test_xpu_batch_invariant_ut.py
23 + - .buildkite/scripts/hardware_ci/run-xpu-batch-invariance.sh
24 + - .buildkite/intel_jobs/misc_intel.yaml
25 + commands:
26 + - >-
27 + bash .buildkite/scripts/hardware_ci/run-intel-test.sh
28 + 'export VLLM_BATCH_INVARIANT=1 &&
29 + cd tests &&
30 + pytest -v -s v1/determinism/test_xpu_batch_invariant_ut.py &&
31 + bash ../.buildkite/scripts/hardware_ci/run-xpu-batch-invariance.sh'
32 +
33 - label: V1 Core + KV + Metrics
34 key: v1-core-kv-metrics-intel
35 timeout_in_minutes: 45
@@ -290,4 +318,4 @@ steps:
318 - >-
319 bash .buildkite/scripts/hardware_ci/run-intel-test.sh
320 'cd tests &&
293 - pytest -v -s compile/passes/test_qk_norm_rope_fusion.py'
\ No newline at end of file
321 + pytest -v -s compile/passes/test_qk_norm_rope_fusion.py'
.buildkite/scripts/hardware_ci/run-xpu-batch-invariance.sh new
+243
@@ -0,0 +1,243 @@
1 +#!/usr/bin/env bash
2 +set -euo pipefail
3 +
4 +repo_root="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/../../.." && pwd)"
5 +
6 +model="ibm-research/PowerMoE-3b"
7 +port=8000
8 +num_decode_tokens=128
9 +server_pid=
10 +server_log=
11 +
12 +# XPU MoE kernels may not safely handle the -1 routes emitted for padding.
13 +export VLLM_MOE_SKIP_PADDING=0
14 +
15 +cd "$repo_root"
16 +
17 +cleanup_server() {
18 + local attempt
19 +
20 + if [[ -n "$server_pid" ]] && kill -0 "$server_pid" 2>/dev/null; then
21 + kill "$server_pid" 2>/dev/null || true
22 +
23 + for ((attempt = 0; attempt < 30; attempt++)); do
24 + if ! kill -0 "$server_pid" 2>/dev/null; then
25 + break
26 + fi
27 + sleep 1
28 + done
29 +
30 + if kill -0 "$server_pid" 2>/dev/null; then
31 + echo "vLLM did not stop after 30 seconds; sending SIGKILL" >&2
32 + kill -KILL "$server_pid" 2>/dev/null || true
33 + fi
34 +
35 + wait "$server_pid" 2>/dev/null || true
36 + fi
37 +
38 + server_pid=
39 +}
40 +
41 +on_exit() {
42 + local status=$?
43 +
44 + cleanup_server
45 + if [[ "$status" -ne 0 && -n "$server_log" && -f "$server_log" ]]; then
46 + if command -v buildkite-agent >/dev/null 2>&1; then
47 + buildkite-agent artifact upload "$server_log" ||
48 + echo "Failed to upload $server_log as a Buildkite artifact" >&2
49 + fi
50 + fi
51 + [[ -z "$server_log" ]] || rm -f "$server_log"
52 + exit "$status"
53 +}
54 +trap on_exit EXIT
55 +
56 +wait_for_server() {
57 + for _ in $(seq 1 600); do
58 + if curl -sf "http://127.0.0.1:$port/health" >/dev/null; then
59 + return
60 + fi
61 + if ! kill -0 "$server_pid" 2>/dev/null; then
62 + echo "vLLM exited before becoming healthy; see the server log above." >&2
63 + return 1
64 + fi
65 + sleep 1
66 + done
67 + echo "Timed out waiting for vLLM server health check" >&2
68 + return 1
69 +}
70 +
71 +request_outputs() {
72 + local prompts_json=$1
73 + local expected_choices=$2
74 + local response
75 +
76 + response=$(printf \
77 + '{"model":"%s","prompt":%s,"temperature":0.7,"top_p":0.95,"seed":0,"max_tokens":%s,"logprobs":64,"ignore_eos":true,"return_token_ids":true}' \
78 + "$model" "$prompts_json" "$num_decode_tokens" |
79 + curl --fail --silent --show-error \
80 + --max-time 1200 \
81 + "http://127.0.0.1:$port/v1/completions" \
82 + -H "Content-Type: application/json" \
83 + --data-binary @-)
84 +
85 + local -a token_id_arrays logprob_arrays
86 + mapfile -t token_id_arrays < <(
87 + grep -o '"token_ids":\[[^]]*\]' <<<"$response" |
88 + sed 's/^"token_ids":\[//; s/\]$//'
89 + )
90 + mapfile -t logprob_arrays < <(
91 + grep -o '"token_logprobs":\[[^]]*\]' <<<"$response" |
92 + sed 's/^"token_logprobs":\[//; s/\]$//'
93 + )
94 + if [[ "${#token_id_arrays[@]}" -ne "$expected_choices" ]] ||
95 + [[ "${#logprob_arrays[@]}" -ne "$expected_choices" ]]; then
96 + echo "Expected $expected_choices choices with token IDs and logits" >&2
97 + return 1
98 + fi
99 + local index
100 + local array
101 + local value
102 + local -a values
103 + for index in "${!logprob_arrays[@]}"; do
104 + array="${logprob_arrays[index]}"
105 + IFS=, read -r -a values <<<"$array"
106 + if [[ "${#values[@]}" -ne "$num_decode_tokens" ]]; then
107 + echo "Completion returned an unexpected number of decode steps" >&2
108 + return 1
109 + fi
110 + IFS=, read -r -a values <<<"${token_id_arrays[index]}"
111 + if [[ "${#values[@]}" -ne "$num_decode_tokens" ]]; then
112 + echo "Completion returned an unexpected number of generated tokens" >&2
113 + return 1
114 + fi
115 + for value in "${values[@]}"; do
116 + if [[ ! "$value" =~ ^[0-9]+$ ]]; then
117 + echo "Completion returned an invalid token ID" >&2
118 + return 1
119 + fi
120 + done
121 + IFS=, read -r -a values <<<"$array"
122 + for value in "${values[@]}"; do
123 + if [[ ! "$value" =~ ^-?([0-9]+([.][0-9]*)?|[.][0-9]+)([eE][+-]?[0-9]+)?$ ]]; then
124 + echo "Completion returned non-finite raw logits" >&2
125 + return 1
126 + fi
127 + done
128 + printf '%s|%s\n' "${token_id_arrays[index]}" "${logprob_arrays[index]}"
129 + done
130 +}
131 +
132 +make_prompt() {
133 + local length=$1
134 + local offset=${2:-0}
135 + local i
136 + local prompt="Use the following system-design context to answer the question. Scenario $offset:"
137 + for ((i = 0; i < length; i++)); do
138 + prompt+=" context"
139 + done
140 + prompt+="\\n\\nQuestion: How should batch-invariant inference behave when this same request is decoded alongside unrelated requests?\\nAnswer:"
141 + printf '"%s"' "$prompt"
142 +}
143 +
144 +first_difference() {
145 + local actual=$1
146 + local expected=$2
147 + local position=$3
148 + local actual_tokens=${actual%%|*}
149 + local expected_tokens=${expected%%|*}
150 + local actual_logits=${actual#*|}
151 + local expected_logits=${expected#*|}
152 + local -a actual_token_values
153 + local -a expected_token_values
154 + local -a actual_logit_values
155 + local -a expected_logit_values
156 + local step
157 +
158 + IFS=, read -r -a actual_token_values <<<"$actual_tokens"
159 + IFS=, read -r -a expected_token_values <<<"$expected_tokens"
160 + IFS=, read -r -a actual_logit_values <<<"$actual_logits"
161 + IFS=, read -r -a expected_logit_values <<<"$expected_logits"
162 + for ((step = 0; step < num_decode_tokens; step++)); do
163 + if [[ "${actual_token_values[step]}" != "${expected_token_values[step]}" ]]; then
164 + printf 'batch position %s: step=%s, baseline token=%s, actual token=%s' \
165 + "$position" \
166 + "$step" \
167 + "${expected_token_values[step]}" \
168 + "${actual_token_values[step]}"
169 + return
170 + fi
171 + if [[ "${actual_logit_values[step]}" != "${expected_logit_values[step]}" ]]; then
172 + printf 'batch position %s: step=%s, token=%s, baseline logit=%s, actual logit=%s' \
173 + "$position" \
174 + "$step" \
175 + "${expected_token_values[step]}" \
176 + "${expected_logit_values[step]}" \
177 + "${actual_logit_values[step]}"
178 + return
179 + fi
180 + done
181 +}
182 +
183 +run_e2e() {
184 + local needle
185 + local batch
186 + local baseline_output
187 + local batch_response
188 + local difference
189 + local position
190 + local -a fillers
191 + local -a batch_outputs
192 + local -a differences=()
193 +
194 + echo "Testing PowerMoE-3b: TP=2, VLLM_BATCH_INVARIANT=1, VLLM_MOE_SKIP_PADDING=0"
195 + server_log=$(mktemp "${TMPDIR:-/tmp}/xpu-batch-invariance.XXXXXX.log")
196 + VLLM_BATCH_INVARIANT=1 vllm serve "$model" \
197 + --tensor-parallel-size 2 \
198 + --max-model-len 1024 \
199 + --max-num-seqs 16 \
200 + --max-num-batched-tokens 512 \
201 + --enable-chunked-prefill \
202 + --no-enable-prefix-caching \
203 + --enforce-eager \
204 + --seed 0 \
205 + --max-logprobs 64 \
206 + --logprobs-mode raw_logits \
207 + --port "$port" > >(tee "$server_log") 2>&1 &
208 + server_pid=$!
209 + wait_for_server
210 +
211 + needle=$(make_prompt 257)
212 + local filler_lengths=(7 31 63 95 127 159 191 223 255 287 319 383 447)
213 + for position in "${!filler_lengths[@]}"; do
214 + fillers+=( "$(make_prompt "${filler_lengths[position]}" "$((position + 1))")" )
215 + done
216 + batch="[$needle,${fillers[0]},${fillers[1]},${fillers[2]},${fillers[3]},${fillers[4]},${fillers[5]},$needle,${fillers[6]},${fillers[7]},${fillers[8]},${fillers[9]},${fillers[10]},${fillers[11]},${fillers[12]},$needle]"
217 +
218 + echo "Requesting BS=1 baseline ($num_decode_tokens decode tokens)..."
219 + baseline_output=$(request_outputs "$needle" 1)
220 + echo "Received BS=1 baseline."
221 + echo "Requesting BS=16 batch ($((16 * num_decode_tokens)) decode tokens)..."
222 + batch_response=$(request_outputs "$batch" 16)
223 + mapfile -t batch_outputs <<<"$batch_response"
224 + echo "Received BS=16 batch; comparing generated tokens and logits..."
225 +
226 + for position in 0 7 15; do
227 + difference=$(first_difference "${batch_outputs[position]}" "$baseline_output" "$position")
228 + if [[ -n "$difference" ]]; then
229 + differences+=("$difference")
230 + fi
231 + done
232 +
233 + if [[ "${#differences[@]}" -ne 0 ]]; then
234 + echo "Batch-invariant PowerMoE-3b output changed: ${differences[*]}" >&2
235 + return 1
236 + fi
237 + echo "PowerMoE-3b batch invariance check passed"
238 + cleanup_server
239 + rm -f "$server_log"
240 + server_log=
241 +}
242 +
243 +run_e2e
tests/v1/determinism/test_xpu_batch_invariant_ut.py new
+228
@@ -0,0 +1,228 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""XPU-only batch-invariance tests for norm, collectives, and sampler behavior."""
4 +
5 +import pytest
6 +import torch
7 +
8 +from vllm import envs
9 +from vllm.config import CacheConfig, LoadConfig, ModelConfig, VllmConfig
10 +from vllm.distributed.device_communicators.xpu_communicator import XpuCommunicator
11 +from vllm.platforms import current_platform
12 +
13 +pytestmark = pytest.mark.skipif(
14 + not current_platform.is_xpu(), reason="XPU batch-invariance tests"
15 +)
16 +
17 +
18 +@pytest.fixture
19 +def tiny_qwen3_moe_path(tmp_path):
20 + """Create a local Qwen3-MoE config for dummy-loading tests."""
21 + from transformers import Qwen3MoeConfig
22 +
23 + config = Qwen3MoeConfig(
24 + architectures=["Qwen3MoeForCausalLM"],
25 + vocab_size=256,
26 + hidden_size=512,
27 + intermediate_size=1024,
28 + moe_intermediate_size=512,
29 + num_hidden_layers=2,
30 + num_attention_heads=4,
31 + num_key_value_heads=2,
32 + head_dim=128,
33 + num_experts=4,
34 + num_experts_per_tok=2,
35 + max_position_embeddings=512,
36 + bos_token_id=1,
37 + eos_token_id=2,
38 + tie_word_embeddings=False,
39 + )
40 + config.save_pretrained(tmp_path)
41 + return tmp_path
42 +
43 +
44 +@pytest.mark.parametrize("source", ["argument", "checkpoint"])
45 +def test_rejects_weight_quantization(tiny_qwen3_moe_path, source):
46 + """Reject both explicit and checkpoint-inferred quantization before loading."""
47 + from transformers import Qwen3MoeConfig
48 +
49 + if source == "checkpoint":
50 + config = Qwen3MoeConfig.from_pretrained(tiny_qwen3_moe_path)
51 + config.quantization_config = {
52 + "quant_method": "fp8",
53 + "activation_scheme": "dynamic",
54 + }
55 + config.save_pretrained(tiny_qwen3_moe_path)
56 +
57 + model_config = ModelConfig(
58 + model=str(tiny_qwen3_moe_path),
59 + skip_tokenizer_init=True,
60 + dtype="bfloat16",
61 + quantization="fp8" if source == "argument" else None,
62 + )
63 + assert model_config.quantization == "fp8"
64 + with pytest.raises(ValueError, match="supports only unquantized models"):
65 + VllmConfig(
66 + model_config=model_config,
67 + load_config=LoadConfig(load_format="dummy"),
68 + )
69 +
70 +
71 +@pytest.mark.parametrize("cache_dtype", ["fp8", "int8_per_token_head", "nvfp4"])
72 +def test_rejects_quantized_kv_cache(cache_dtype):
73 + with pytest.raises(ValueError, match="does not support quantized KV caches"):
74 + VllmConfig(cache_config=CacheConfig(cache_dtype=cache_dtype))
75 +
76 +
77 +@pytest.mark.parametrize("cache_dtype", ["auto", "float16", "bfloat16"])
78 +def test_accepts_unquantized_kv_cache(cache_dtype):
79 + config = VllmConfig(cache_config=CacheConfig(cache_dtype=cache_dtype))
80 + assert config.cache_config.cache_dtype == cache_dtype
81 +
82 +
83 +def test_quantized_kv_cache_allowed_without_batch_invariance(monkeypatch):
84 + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", False)
85 + config = VllmConfig(cache_config=CacheConfig(cache_dtype="fp8"))
86 + assert config.cache_config.cache_dtype == "fp8"
87 +
88 +
89 +# These dimensions and seeds locally reproduce non-invariant norm mismatches.
90 +# They make the invariant equality assertion exercise a discriminating case.
91 +@pytest.mark.parametrize(
92 + (
93 + "norm_name",
94 + "batch_size",
95 + "hidden_size",
96 + "position",
97 + "weight_seed",
98 + "input_seed",
99 + "input_pattern",
100 + ),
101 + [
102 + ("rms", 1024, 16384, 512, 7, 42, "normal"),
103 + ("gemma", 257, 8192, 255, 1, 2, "normal"),
104 + ("rms", 5, 8192, 3, 7, 46, "scaled"),
105 + ("gemma", 5, 8192, 3, 1, 46, "scaled"),
106 + ],
107 +)
108 +def test_residual_norm_preserves_batch_invariance(
109 + norm_name,
110 + batch_size,
111 + hidden_size,
112 + position,
113 + weight_seed,
114 + input_seed,
115 + input_pattern,
116 + monkeypatch,
117 + default_vllm_config,
118 +):
119 + """Residual norm output must not change when its row changes batch position."""
120 + from vllm.model_executor.determinism.batch_invariant import init_batch_invariance
121 + from vllm.model_executor.layers.layernorm import GemmaRMSNorm, RMSNorm
122 +
123 + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True)
124 + init_batch_invariance()
125 + norm_cls = RMSNorm if norm_name == "rms" else GemmaRMSNorm
126 + norm = norm_cls(hidden_size, eps=1e-6).to(device="xpu", dtype=torch.bfloat16)
127 + with torch.no_grad():
128 + torch.manual_seed(weight_seed)
129 + norm.weight.uniform_(-0.5, 0.5)
130 +
131 + if input_pattern == "scaled":
132 + # Heterogeneous columns expose native mean's batch-dependent rounding.
133 + generator = torch.Generator(device="cpu").manual_seed(input_seed)
134 + scale = torch.randn(hidden_size, generator=generator).exp()
135 + x = (torch.randn(batch_size, hidden_size, generator=generator) * scale).to(
136 + device="xpu", dtype=torch.bfloat16
137 + )
138 + residual = (
139 + torch.randn(batch_size, hidden_size, generator=generator) * scale
140 + ).to(device="xpu", dtype=torch.bfloat16)
141 + else:
142 + torch.manual_seed(input_seed)
143 + x = torch.randn(batch_size, hidden_size, device="xpu", dtype=torch.bfloat16)
144 + residual = torch.randn_like(x)
145 +
146 + single_output, single_residual = norm(
147 + x[position : position + 1].clone(),
148 + residual[position : position + 1].clone(),
149 + )
150 + batch_output, batch_residual = norm(x.clone(), residual.clone())
151 +
152 + assert torch.equal(single_output, batch_output[position : position + 1])
153 + assert torch.equal(single_residual, batch_residual[position : position + 1])
154 +
155 +
156 +@pytest.mark.parametrize(
157 + "operation", ["all_reduce", "reduce_scatter", "reduce_scatterv"]
158 +)
159 +@pytest.mark.parametrize("rank", range(4))
160 +def test_collectives_preserve_rank_order_and_partition(monkeypatch, operation, rank):
161 + """Cancellation exposes reordered sums; mock ranks require no extra XPUs."""
162 + from vllm.distributed.device_communicators import xpu_communicator
163 +
164 + inputs = [
165 + torch.full((8, 3), value, dtype=torch.float32, device="xpu")
166 + for value in (1e20, -1e20, 3.0, 0.0)
167 + ]
168 + inputs[2].add_(torch.arange(24, device="xpu").reshape(8, 3))
169 +
170 + def all_gather_single(output, input_, group):
171 + torch.testing.assert_close(input_, inputs[rank].reshape(-1), rtol=0, atol=0)
172 + output.copy_(torch.stack(inputs).reshape(-1))
173 +
174 + monkeypatch.setattr(xpu_communicator.dist, "all_gather_single", all_gather_single)
175 + communicator = XpuCommunicator.__new__(XpuCommunicator)
176 + communicator.world_size = 4
177 + communicator.rank_in_group = rank
178 + communicator.device_group = object()
179 + expected = inputs[0].clone()
180 + for value in inputs[1:]:
181 + expected.add_(value)
182 + assert torch.equal(expected, inputs[2])
183 +
184 + if operation == "all_reduce":
185 + actual = communicator.all_reduce(inputs[rank])
186 + elif operation == "reduce_scatter":
187 + actual = communicator.reduce_scatter(inputs[rank], dim=0)
188 + expected = expected.chunk(4)[rank]
189 + else:
190 + sizes = [1, 3, 2, 2]
191 + actual = communicator.reduce_scatterv(inputs[rank], dim=0, sizes=sizes)
192 + expected = expected.split(sizes)[rank]
193 + torch.testing.assert_close(actual, expected, rtol=0, atol=0)
194 +
195 +
196 +# This batch size, vocabulary size, and positions locally reproduce an
197 +# unseeded non-invariant XPU sampler mismatch.
198 +@pytest.mark.parametrize("position", [0, 128, 255])
199 +def test_seeded_sampler_preserves_full_distribution_and_rng(position, monkeypatch):
200 + """Moving a request must preserve its distribution and per-request RNG."""
201 + from vllm.triton_utils import HAS_TRITON
202 + from vllm.v1.sample.ops.topk_topp_sampler import TopKTopPSampler
203 +
204 + assert HAS_TRITON, "XPU sampling must exercise the Triton top-k/top-p path"
205 + monkeypatch.setattr(envs, "VLLM_XPU_USE_SAMPLER_KERNEL", True)
206 + sampler = TopKTopPSampler(logprobs_mode="processed_logprobs")
207 + torch.manual_seed(42)
208 + logits = torch.randn(256, 8192, device="xpu", dtype=torch.float32)
209 + single_rng = torch.Generator(device="xpu").manual_seed(123)
210 + batch_rng = torch.Generator(device="xpu").manual_seed(123)
211 + k = torch.full((256,), 32, dtype=torch.int32, device="xpu")
212 + p = torch.full((256,), 0.9, device="xpu")
213 + for _ in range(3):
214 + single_token, single_logprobs = sampler(
215 + logits[position : position + 1].clone(),
216 + {0: single_rng},
217 + k[position : position + 1],
218 + p[position : position + 1],
219 + )
220 + batch_tokens, batch_logprobs = sampler(
221 + logits.clone(), {position: batch_rng}, k, p
222 + )
223 + assert single_logprobs is not None and batch_logprobs is not None
224 + torch.testing.assert_close(
225 + single_logprobs[0], batch_logprobs[position], rtol=0, atol=0
226 + )
227 + assert single_token.item() == batch_tokens[position].item()
228 + assert torch.equal(single_rng.get_state(), batch_rng.get_state())
vllm/distributed/device_communicators/xpu_communicator.py
+35 -3
@@ -6,6 +6,7 @@ import torch
6 import torch.distributed as dist
7 from torch.distributed import ProcessGroup
8
9 +import vllm.envs as envs
10 from vllm.logger import init_logger
11
12 from .base_device_communicator import DeviceCommunicatorBase
@@ -45,7 +46,25 @@ class XpuCommunicator(DeviceCommunicatorBase):
46 self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
47 logger.info("Using AgRs manager on XPU device.")
48
49 + def _fixed_rank_sum(self, input_: torch.Tensor) -> torch.Tensor:
50 + flat_input = input_.reshape(-1)
51 + gathered = torch.empty(
52 + (self.world_size, flat_input.numel()),
53 + dtype=input_.dtype,
54 + device=input_.device,
55 + )
56 + dist.all_gather_single(gathered.view(-1), flat_input, group=self.device_group)
57 + output = gathered[0].clone()
58 + for rank in range(1, self.world_size):
59 + output.add_(gathered[rank])
60 + return output.view(input_.shape)
61 +
62 def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
63 + if envs.VLLM_BATCH_INVARIANT:
64 + # XCCL changes its reduction order with message size. Gather first
65 + # and accumulate in rank order, independent of the token batch.
66 + return self._fixed_rank_sum(input_)
67 +
68 output = input_.clone()
69 dist.all_reduce(output, group=self.device_group)
70 return output
@@ -69,7 +88,12 @@ class XpuCommunicator(DeviceCommunicatorBase):
88 output_shape, dtype=input_tensor.dtype, device=input_tensor.device
89 )
90
72 - dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group)
91 + if envs.VLLM_BATCH_INVARIANT:
92 + reduced = self._fixed_rank_sum(input_tensor)
93 + start = self.rank_in_group * chunk_size
94 + output.copy_(reduced.narrow(0, start, chunk_size))
95 + else:
96 + dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group)
97
98 # Reshape before returning
99 return output.movedim(0, dim).contiguous()
@@ -99,7 +123,15 @@ class XpuCommunicator(DeviceCommunicatorBase):
123 output = torch.empty(
124 output_shape, dtype=input_tensor.dtype, device=input_tensor.device
125 )
102 - if sizes is None or sizes.count(sizes[0]) == len(sizes):
126 + if envs.VLLM_BATCH_INVARIANT:
127 + reduced = self._fixed_rank_sum(input_tensor)
128 + start = (
129 + self.rank_in_group * chunk_size
130 + if sizes is None
131 + else sum(sizes[: self.rank_in_group])
132 + )
133 + output.copy_(reduced.narrow(0, start, chunk_size))
134 + elif sizes is None or sizes.count(sizes[0]) == len(sizes):
135 dist.reduce_scatter_tensor(output, input_tensor, group=self.device_group)
136 elif torch.xpu.device_count() < self.world_size:
137 # uneven reduce_scatter is not supported when the
@@ -191,7 +223,7 @@ class XpuCommunicator(DeviceCommunicatorBase):
223 (self.world_size,) + input_size, dtype=input_.dtype, device=input_.device
224 )
225 # All-gather.
194 - dist.all_gather_into_tensor(output_tensor, input_, group=self.device_group)
226 + dist.all_gather_single(output_tensor, input_, group=self.device_group)
227 if self.rank_in_group == dst:
228 # Reshape
229 output_tensor = output_tensor.movedim(0, dim)
vllm/model_executor/determinism/batch_invariant.py
+4 -2
@@ -1104,7 +1104,10 @@ def enable_batch_invariant_mode():
1104
1105 _fp16_block_size_n = 128
1106
1107 - # Softmax, log_softmax, mean, and bmm are already batch-invariant on XPU
1107 + # Native mean can change reduction order with the row count on XPU too.
1108 + _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, key)
1109 +
1110 + # Softmax, log_softmax, 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():
@@ -1113,7 +1116,6 @@ def enable_batch_invariant_mode():
1116 )
1117 _batch_invariant_LIB.impl("aten::softmax", softmax_batch_invariant, key)
1118 _batch_invariant_LIB.impl("aten::_softmax", softmax_batch_invariant, key)
1116 - _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, key)
1119 # torch 2.12+ registers a built-in Triton bmm kernel for CUDA
1120 # (torch._native.ops.bmm_outer_product), so we need allow_override
1121 # to replace it at the dispatcher level.
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+7 -1
@@ -94,7 +94,13 @@ def _get_priority_backends(moe_config: FusedMoEConfig) -> list[UnquantizedMoeBac
94 _move_to_back(_AVAILABLE_BACKENDS, UnquantizedMoeBackend.FLASHINFER_CUTLASS)
95
96 elif current_platform.is_xpu():
97 - _AVAILABLE_BACKENDS = [UnquantizedMoeBackend.XPU]
97 + # XPUExperts does not support batch invariance; TRITON is kept as a
98 + # fallback so VLLM_BATCH_INVARIANT=1 still resolves to a valid
99 + # (batch-invariant) backend on XPU.
100 + _AVAILABLE_BACKENDS = [
101 + UnquantizedMoeBackend.XPU,
102 + UnquantizedMoeBackend.TRITON,
103 + ]
104 elif current_platform.is_cpu():
105 _AVAILABLE_BACKENDS = [UnquantizedMoeBackend.CPU]
106 return _AVAILABLE_BACKENDS
vllm/model_executor/layers/layernorm.py
+17
@@ -119,6 +119,14 @@ class RMSNorm(CustomOp):
119 x: torch.Tensor,
120 residual: torch.Tensor | None = None,
121 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
122 + if envs.VLLM_BATCH_INVARIANT and residual is not None:
123 + assert self.variance_size_override is None, (
124 + "Batch invariance is not supported for variance_size_override"
125 + )
126 + weight = self.weight.data if self.pass_weight_add else None
127 + return ir.ops.fused_add_rms_norm.impls["native"].impl_fn(
128 + x, residual, weight, self.variance_epsilon
129 + )
130 return self.forward_cuda(x, residual)
131
132 def extra_repr(self) -> str:
@@ -171,6 +179,15 @@ class GemmaRMSNorm(CustomOp):
179 x: torch.Tensor,
180 residual: torch.Tensor | None = None,
181 ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
182 + if envs.VLLM_BATCH_INVARIANT:
183 + if residual is None:
184 + return self.forward_cuda(x, residual)
185 + return ir.ops.fused_add_rms_norm.impls["native"].impl_fn(
186 + x,
187 + residual,
188 + self.weight.float() + 1.0,
189 + self.variance_epsilon,
190 + )
191 import vllm._xpu_ops # noqa: F401 registers torch.ops.vllm.xpu_gemma_rms_norm
192
193 # Fall back to the native path if the fused gemma kernels are not
vllm/platforms/xpu.py
+40
@@ -162,6 +162,25 @@ class XPUPlatform(Platform):
162 if selected_backend == AttentionBackendEnum.TRITON_ATTN:
163 logger.info_once("Using Triton backend.")
164 return AttentionBackendEnum.TRITON_ATTN.get_path()
165 + elif attn_selector_config.use_batch_invariant:
166 + # Flash Attention on XPU has not been validated for batch
167 + # invariance. Honor an explicit Flash Attention request;
168 + # otherwise fall back to Triton Attention, which implements
169 + # batch-invariant kernels.
170 + if selected_backend == AttentionBackendEnum.FLASH_ATTN:
171 + logger.warning_once(
172 + "Using Flash Attention on XPU with batch invariance "
173 + "enabled because it was explicitly requested. This "
174 + "backend has not been validated for batch invariance "
175 + "on XPU and may produce non-deterministic results "
176 + "across batch sizes."
177 + )
178 + return AttentionBackendEnum.FLASH_ATTN.get_path()
179 + logger.info_once(
180 + "VLLM_BATCH_INVARIANT is enabled. Using Triton Attention "
181 + "backend on XPU, which implements batch-invariant kernels."
182 + )
183 + return AttentionBackendEnum.TRITON_ATTN.get_path()
184 elif attn_selector_config.use_mm_prefix:
185 # Flash Attention on XPU has no FA4 kernel, so it cannot apply the
186 # multimodal prefix-LM bidirectional mask. Honor an explicit Flash
@@ -283,6 +302,27 @@ class XPUPlatform(Platform):
302 # lazy import to avoid circular import
303 from vllm.config import CUDAGraphMode
304
305 + if envs.VLLM_BATCH_INVARIANT:
306 + model_config = vllm_config.model_config
307 + if model_config is not None and (
308 + model_config.quantization is not None
309 + or vllm_config.quant_config is not None
310 + ):
311 + raise ValueError(
312 + "XPU batch invariance currently supports only unquantized "
313 + f"models; got quantization={model_config.quantization!r}. "
314 + "Use an unquantized model or disable VLLM_BATCH_INVARIANT."
315 + )
316 +
317 + cache_dtype = vllm_config.cache_config.cache_dtype
318 + if cache_dtype not in ("auto", "float16", "bfloat16"):
319 + raise ValueError(
320 + "XPU batch invariance currently does not support quantized "
321 + f"KV caches; got kv_cache_dtype={cache_dtype!r}. "
322 + "Use an unquantized KV cache dtype or disable "
323 + "VLLM_BATCH_INVARIANT."
324 + )
325 +
326 compilation_config = vllm_config.compilation_config
327 if compilation_config.compile_sizes is None:
328 compilation_config.compile_sizes = []
vllm/sampling_params.py
+10
@@ -526,6 +526,16 @@ class SamplingParams(
526
527 self._verify_args()
528
529 + if (
530 + envs.VLLM_BATCH_INVARIANT
531 + and self.temperature >= _SAMPLING_EPS
532 + and self.seed is None
533 + ):
534 + logger.warning_once(
535 + "Random sampling without an explicit seed may not be batch "
536 + "invariant. Set seed in SamplingParams or use temperature=0."
537 + )
538 +
539 if self.temperature < _SAMPLING_EPS:
540 # Zero temperature means greedy sampling.
541 self.top_p = 1.0
vllm/v1/sample/ops/topk_topp_sampler.py
+6 -1
@@ -224,9 +224,14 @@ class TopKTopPSampler(nn.Module):
224 else:
225 self.forward = self.forward_cpu
226 elif current_platform.is_xpu():
227 - if xpu_sampler_supported():
227 + if xpu_sampler_supported() and not envs.VLLM_BATCH_INVARIANT:
228 self.forward = self.forward_xpu
229 else:
230 + if envs.VLLM_BATCH_INVARIANT:
231 + logger.info_once(
232 + "VLLM_BATCH_INVARIANT is enabled. Using the "
233 + "PyTorch-native sampler on XPU."
234 + )
235 self.forward = self.forward_native
236 elif (
237 logprobs_mode not in PROCESSED_LOGPROBS_MODES