@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 170 lines 6.69 KB
Rendered Raw
1 # Batch Invariance
2
3 !!! note
4 Batch invariance is currently in beta. Some features are still under active development.
5 Track progress and planned improvements at <https://github.com/vllm-project/vllm/issues/27433>
6
7 This document shows how to enable batch invariance in vLLM. Batch invariance ensures that the output of a model is deterministic and independent of the batch size or the order of requests in a batch.
8
9 ## Motivation
10
11 Batch invariance is crucial for several use cases:
12
13 - **Framework debugging**: Deterministic outputs make it easier to debug issues in the inference framework, as the same input will always produce the same output regardless of batching.
14 - **Model debugging**: Helps identify issues in model implementations by ensuring consistent behavior across different batch configurations.
15 - **Reinforcement Learning (RL)**: RL training often requires deterministic rollouts for reproducibility and stable training.
16 - **Large-scale inference systems**: Systems that use vLLM as a component benefit from deterministic behavior for testing, validation, and consistency guarantees.
17
18 ## Hardware Requirements
19
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
46 Batch invariance can be enabled by setting the `VLLM_BATCH_INVARIANT` environment variable to `1`:
47
48 ```bash
49 export VLLM_BATCH_INVARIANT=1
50 ```
51
52 ### Online Inference (Server Mode)
53
54 To start a vLLM server with batch invariance enabled:
55
56 ```bash
57 VLLM_BATCH_INVARIANT=1 vllm serve meta-llama/Llama-3.1-8B-Instruct
58 ```
59
60 Then use the OpenAI-compatible client:
61
62 ```python
63 from openai import OpenAI
64
65 client = OpenAI(
66 api_key="EMPTY",
67 base_url="http://localhost:8000/v1",
68 )
69
70 # These requests will produce deterministic outputs
71 # regardless of batch size or order
72 response = client.completions.create(
73 model="meta-llama/Llama-3.1-8B-Instruct",
74 prompt="The future of AI is",
75 max_tokens=100,
76 temperature=0.7,
77 seed=42,
78 )
79
80 print(response.choices[0].text)
81 ```
82
83 ### Offline Inference
84
85 For offline batch inference with batch invariance:
86
87 ```python
88 import os
89 os.environ["VLLM_BATCH_INVARIANT"] = "1"
90
91 from vllm import LLM, SamplingParams
92
93 prompts = [
94 "The future of AI is",
95 "Machine learning enables",
96 "Deep learning models can",
97 ]
98
99 sampling_params = SamplingParams(
100 temperature=0.7,
101 top_p=0.95,
102 max_tokens=100,
103 seed=42,
104 )
105
106 llm = LLM(
107 model="meta-llama/Llama-3.1-8B-Instruct",
108 tensor_parallel_size=1,
109 )
110
111 # Outputs will be deterministic regardless of batch size
112 outputs = llm.generate(prompts, sampling_params)
113
114 for output in outputs:
115 prompt = output.prompt
116 generated_text = output.outputs[0].text
117 print(f"Prompt: {prompt!r}")
118 print(f"Generated: {generated_text!r}\n")
119 ```
120
121 ## Tested Models
122
123 Batch invariance has been tested and verified on the following models:
124
125 - **DeepSeek series**: `deepseek-ai/DeepSeek-V3`, `deepseek-ai/DeepSeek-V3-0324`, `deepseek-ai/DeepSeek-R1`, `deepseek-ai/DeepSeek-V3.1`
126 - **Qwen3 (Dense)**: `Qwen/Qwen3-1.7B`, `Qwen/Qwen3-8B`, `Qwen/Qwen3-4B-AWQ`, `Qwen/Qwen3-8B-AWQ`
127 - **Qwen3-VL (Vision-Language)**: `Qwen/Qwen3-VL-2B-Instruct`, `Qwen/Qwen3-VL-4B-Instruct` (single image and video inputs)
128 - **Qwen3 (MoE)**: `Qwen/Qwen3-30B-A3B`, `Qwen/Qwen3-Next-80B-A3B-Instruct`, `Qwen/Qwen3-30B-A3B-Thinking-2507-FP8`
129 - **Qwen2.5**: `Qwen/Qwen2.5-0.5B-Instruct`, `Qwen/Qwen2.5-1.5B-Instruct`, `Qwen/Qwen2.5-3B-Instruct`, `Qwen/Qwen2.5-7B-Instruct`, `Qwen/Qwen2.5-14B-Instruct`, `Qwen/Qwen2.5-32B-Instruct`
130 - **Llama 3**: Llama3.1 and 3.2 series, `meta-llama/Llama-3.2-1B-Instruct`, `meta-llama/Llama-3.2-3B-Instruct` for example
131 - **GPT-OSS**: `openai/gpt-oss-20b`, `openai/gpt-oss-120b`
132 - **Mistral**: `mistralai/Mistral-7B-v0.3`
133 - **Phi series**: `microsoft/Phi-3.5-mini-instruct`
134 - **Granite 3.1 (MoE)**: `ibm-granite/granite-3.1-1b-a400m-instruct`, `ibm-granite/granite-3.1-3b-a800m-instruct`
135 - **Granite 3.1 (Dense)**: `ibm-granite/granite-3.1-2b-instruct`, `ibm-granite/granite-3.1-8b-instruct`
136 - **EXAONE 4.0 series**: `LGAI-EXAONE/EXAONE-4.0-1.2B`, `LGAI-EXAONE/EXAONE-4.0.1-32B`, `LGAI-EXAONE/EXAONE-4.0-32B`
137 - **OLMo 2**: `allenai/OLMo-2-0425-1B-Instruct`
138 - **SmolLM2**: `HuggingFaceTB/SmolLM2-1.7B-Instruct`
139 - **PLaMo3**: `pfnet/plamo-3-nict-2b-base`
140
141 Other models may also work, but these have been explicitly validated. If you encounter issues with a specific model, please report them on the [GitHub issue tracker](https://github.com/vllm-project/vllm/issues/new/choose).
142
143 ## Implementation Details
144
145 When batch invariance is enabled, vLLM:
146
147 1. Uses deterministic kernel implementations for attention and other operations
148 2. Ensures consistent numerical behavior across different batch sizes
149 3. Disables certain optimizations that may introduce non-determinism (such as sequence parallelism / async TP, whose reduce-scatter path is not batch-invariant)
150 4. Under tensor parallelism, keeps custom all-reduce on with a fixed reduction order (the 1-stage kernel is pinned, and large inputs are reduced in fixed-size chunks), and disables FlashInfer, AITER and QuickReduce all-reduce
151 5. On CUDA devices with tuned matmul table entries for the model's bf16 unquantized forward linear layers (Ada, Hopper, Blackwell),
152 runs without `torch.compile` using breakable CUDA graphs so tile configs follow the runtime batch size; set `VLLM_USE_BREAKABLE_CUDAGRAPH=0` to opt out
153 (not applied when sequence parallelism / async TP are enabled, since those are `torch.compile` passes).
154
155 !!! warning
156 Batch invariance under tensor parallelism is not yet supported for all-reduces whose size isn't a multiple of 16 bytes, for example a hidden size that isn't a multiple of 8 in fp16/bf16. Such tensors can switch between NCCL and custom all-reduce depending on batch size. All validated models meet this requirement.
157
158 !!! note
159 Enabling batch invariance may impact performance compared to the default non-deterministic mode. This trade-off is intentional to guarantee reproducibility.
160
161 ## Future Improvements
162
163 The batch invariance feature is under active development. Planned improvements include:
164
165 - Support for additional GPU architectures
166 - Expanded model coverage
167 - Performance optimizations
168 - Additional testing and validation
169
170 For the latest status and to contribute ideas, see the [tracking issue](https://github.com/vllm-project/vllm/issues/27433).