@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2f7dbc9b42

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

Add batch invariant kernel override for FlashInfer backend [2/n] (#25769)

Signed-off-by: Bram Wasti <bwasti@meta.com> Signed-off-by: Bram Wasti <bwasti@fb.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>

Bram Wasti committed Oct 3, 2025 at 21:49 UTC 2f7dbc9b42c51ba192e3dded515e4e07cdfdabea

3 files changed +84 -29

tests/v1/generation/test_batch_invariance.py

+40 -23

@@ -76,18 +76,21 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle():

76 seed.

77 - Keep max_tokens and max_model_len bounded for speed and memory use.

78 """

79 - random.seed(12345)

79 + seed = int(os.getenv("VLLM_TEST_SEED", "12345"))

80 + random.seed(seed)

82 # Allow overrides from environment (useful for CI tuning)

83 # "facebook/opt-125m" is too small, doesn't reliably test determinism

84 model = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")

85 num_trials = int(os.getenv("VLLM_NEEDLE_TRIALS", "5"))

85 - batch_size = int(os.getenv("VLLM_NEEDLE_BATCH_SIZE", "64"))

86 - assert batch_size >= 2, "Batch size should be >= 2 to mix needle."

86 + max_batch_size = int(os.getenv("VLLM_NEEDLE_BATCH_SIZE", "128"))

87 + min_random_prompt = int(os.getenv("VLLM_MIN_PROMPT", "1024"))

88 + max_random_prompt = int(os.getenv("VLLM_MAX_PROMPT", "2048"))

89 + assert max_batch_size >= 2, "Batch size should be >= 2 to mix needle."

91 # Keep GPU memory usage low to avoid startup allocation failures.

89 - gpu_mem_util = float(os.getenv("VLLM_GPU_MEMORY_UTILIZATION", "0.3"))

90 - max_model_len = int(os.getenv("VLLM_MAX_MODEL_LEN", "4096"))

92 + gpu_mem_util = float(os.getenv("VLLM_GPU_MEMORY_UTILIZATION", "0.4"))

93 + max_model_len = int(os.getenv("VLLM_MAX_MODEL_LEN", "5120"))

94 swap_space_gb = int(os.getenv("VLLM_SWAP_SPACE_GB", "4"))

96 # Sampling parameters: longer outputs with a more random-sounding

@@ -111,7 +114,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle():

114 # Engine with bs=1 behavior

115 llm_bs1 = LLM_with_max_seqs(

116 model=model,

114 - max_num_seqs=1,

117 + max_num_seqs=128,

118 gpu_memory_utilization=gpu_mem_util,

119 max_model_len=max_model_len,

120 swap_space=swap_space_gb,

@@ -126,7 +129,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle():

129 # Engine with larger batch limit (e.g., 64)

130 llm_bsN = LLM_with_max_seqs(

131 model=model,

129 - max_num_seqs=batch_size,

132 + max_num_seqs=128,

133 gpu_memory_utilization=gpu_mem_util,

134 max_model_len=max_model_len,

135 swap_space=swap_space_gb,

@@ -135,15 +138,17 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle():

138 mismatches = 0

139

140 for trial in range(num_trials):

138 - # Create a batch of size `batch_size` and insert the needle at

141 + # Create a batch of size `max_batch_size` and insert the needle at

142 # a random index

143 prompts: list[str] = []

144 + batch_size = random.randint(max_batch_size // 2, max_batch_size)

145 needle_pos = random.randint(0, batch_size - 1)

146 for i in range(batch_size):

147 if i == needle_pos:

148 prompts.append(needle_prompt)

149 else:

146 - prompts.append(_random_prompt())

150 + prompts.append(

151 + _random_prompt(min_random_prompt, max_random_prompt))

152

153 # Generate with the larger-batch engine

154 outputs = llm_bsN.generate(prompts, sampling)

@@ -154,17 +159,19 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle():

159 text = needle_output.outputs[0].text

160

161 if text != baseline_text:

162 + print(

163 + f"{text}\n\n== Not the same as ==\n\n{baseline_text}\n\n")

164 mismatches += 1

165

166 passes = num_trials - mismatches

167 # Dump how many passed vs failed

168 print(f"[determinism] total={num_trials}, passed={passes}, "

162 - f"failed={mismatches}, batch_size={batch_size}")

169 + f"failed={mismatches}, max_batch_size={max_batch_size}")

170

171 if mismatches > 0:

172 pytest.fail(

173 f"Nondeterministic outputs detected: {mismatches} failed out "

167 - f"of {num_trials} trials (batch_size={batch_size}).")

174 + f"of {num_trials} trials (max_batch_size={max_batch_size}).")

175

176 finally:

177 # Ensure engines are shutdown to free GPU/VRAM across test sessions

@@ -196,9 +203,14 @@ def _extract_step_logprobs(request_output):

203 not torch.cuda.is_available(),

204 reason="Requires CUDA to match production inference path.",

205 )

199 -def test_logprobs_bitwise_batch_invariance_bs1_vs_bs2():

206 +@pytest.mark.parametrize("backend", ["FLEX_ATTENTION", "FLASHINFER"])

207 +def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(backend):

208

201 - #model_name = os.getenv("VLLM_TEST_MODEL", "facebook/opt-125m")

209 + backend = os.getenv("VLLM_ATTENTION_BACKEND", backend)

210 + os.environ["VLLM_ATTENTION_BACKEND"] = backend

211 +

212 + seed = int(os.getenv("VLLM_TEST_SEED", "12345"))

213 + random.seed(seed)

214 model_name = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")

215 tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))

216

@@ -212,10 +224,15 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bs2():

224 prompts = [

225 "The capital of France is",

226 "The capital of Germany is",

227 + _random_prompt(10, 1024),

228 + _random_prompt(10, 1024),

229 + _random_prompt(10, 1024),

230 + _random_prompt(10, 1024),

231 + _random_prompt(10, 1024),

232 ]

233

234 sp = SamplingParams(

218 - temperature=0.0,

235 + temperature=0.6,

236 top_p=1.0,

237 max_tokens=8,

238 # Seed shouldn't matter at temperature=0, but keeping it stable anyway.

@@ -234,25 +251,25 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bs2():

251 "enable logprobs return to run this test.")

252 bs1_logprobs_per_prompt.append(step_logprobs)

253

237 - # BS=2: run prompts in a batch and collect logprobs per step for each

254 + # BS=N: run prompts in a batch and collect logprobs per step for each

255 # prompt.

256 outs_batched = llm.generate(prompts, sp, use_tqdm=False)

257 assert len(outs_batched) == len(prompts)

241 - bs2_logprobs_per_prompt = []

258 + bsN_logprobs_per_prompt = []

259 for o in outs_batched:

260 step_logprobs = _extract_step_logprobs(o)

261 if step_logprobs is None:

262 pytest.skip("Logits are not available on RequestOutput; "

263 "enable logprobs return to run this test.")

247 - bs2_logprobs_per_prompt.append(step_logprobs)

264 + bsN_logprobs_per_prompt.append(step_logprobs)

265

249 - # Compare step-by-step logprobs for each prompt between BS=1 and BS=2 runs.

250 - for i, (logprobs_bs1, logprobs_bs2) in enumerate(

251 - zip(bs1_logprobs_per_prompt, bs2_logprobs_per_prompt)):

252 - assert len(logprobs_bs1) == len(logprobs_bs2), (

266 + # Compare step-by-step logprobs for each prompt between BS=1 and BS=N runs.

267 + for i, (logprobs_bs1, logprobs_bsN) in enumerate(

268 + zip(bs1_logprobs_per_prompt, bsN_logprobs_per_prompt)):

269 + assert len(logprobs_bs1) == len(logprobs_bsN), (

270 f"Different number of generation steps for prompt index {i}: "

254 - f"{len(logprobs_bs1)} (BS=1) vs {len(logprobs_bs2)} (BS=2)")

255 - for t, (a, b) in enumerate(zip(logprobs_bs1, logprobs_bs2)):

271 + f"{len(logprobs_bs1)} (BS=1) vs {len(logprobs_bsN)} (BS=N)")

272 + for t, (a, b) in enumerate(zip(logprobs_bs1, logprobs_bsN)):

273 assert a.shape == b.shape, (

274 f"Logits shape mismatch at prompt {i}, step {t}: "

275 f"{a.shape} vs {b.shape}")

vllm/model_executor/layers/batch_invariant.py

+12 -1

@@ -8,8 +8,12 @@ from typing import Any, Union

9 import torch

11 +import vllm.envs as envs

12 +from vllm.logger import init_logger

13 from vllm.triton_utils import tl, triton

15 +logger = init_logger(__name__)

16 +

18 def _matmul_launch_metadata(grid: Callable[..., Any], kernel: Any,

19 args: dict[str, Any]) -> dict[str, Any]:

@@ -557,5 +561,12 @@ def vllm_kernel_override_batch_invariant():

561 def init_batch_invariance():

562 # this will hit all the csrc overrides as well

563 if vllm_kernel_override_batch_invariant():

560 - os.environ["VLLM_ATTENTION_BACKEND"] = "FLEX_ATTENTION"

564 + curr_attn_backend = envs.VLLM_ATTENTION_BACKEND

565 + supported_backends = ["FLEX_ATTENTION", "FLASHINFER"]

566 + if curr_attn_backend not in supported_backends:

567 + warning = "Forcibly updating attention backend to" \

568 + f" {supported_backends[0]} for batch_invariant. " \

569 + f" Supported backends: {supported_backends}."

570 + logger.warning_once(warning)

571 + os.environ["VLLM_ATTENTION_BACKEND"] = supported_backends[0]

572 enable_batch_invariant_mode()

vllm/v1/attention/backends/flashinfer.py

+32 -5

@@ -20,6 +20,8 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,

20 AttentionType)

21 from vllm.config import CUDAGraphMode, VllmConfig

22 from vllm.logger import init_logger

23 +from vllm.model_executor.layers.batch_invariant import (

24 + vllm_kernel_override_batch_invariant)

25 from vllm.model_executor.layers.quantization.utils.quant_utils import (

26 QuantKey, kFp8StaticTensorSym, kNvfp4Quant)

27 from vllm.platforms import current_platform

@@ -42,6 +44,7 @@ from vllm.v1.attention.backends.utils import (AttentionCGSupport,

44 from vllm.v1.kv_cache_interface import AttentionSpec

46 FLASHINFER_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024

47 +FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT = 2048 * 1024 * 1024

49 FP8_DTYPE = current_platform.fp8_dtype()

50 FP4_DTYPE = torch.uint8

@@ -263,6 +266,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

266 self._prefill_wrapper = None # Wrapper for prefill/append

267 self._decode_wrapper = None # Wrapper for decode (general shape)

268

269 + if vllm_kernel_override_batch_invariant():

270 + self.decode_fixed_split_size = 2048

271 + self.prefill_fixed_split_size = 4096

272 + self.disable_split_kv = True

273 + else:

274 + self.decode_fixed_split_size = -1

275 + self.prefill_fixed_split_size = -1

276 + self.disable_split_kv = False

277 +

278 self.compilation_config = vllm_config.compilation_config

279 max_num_pages_per_req = cdiv(self.model_config.max_model_len,

280 self.kv_cache_spec.block_size)

@@ -356,10 +368,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

368

369 def _get_workspace_buffer(self):

370 if self._workspace_buffer is None:

359 - self._workspace_buffer = torch.zeros(

360 - FLASHINFER_WORKSPACE_BUFFER_SIZE,

361 - dtype=torch.uint8,

362 - device=self.device)

371 + buffer_size = FLASHINFER_WORKSPACE_BUFFER_SIZE

372 + if vllm_kernel_override_batch_invariant():

373 + buffer_size = FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT

374 + self._workspace_buffer = torch.zeros(buffer_size,

375 + dtype=torch.uint8,

376 + device=self.device)

377 return self._workspace_buffer

378

379 def _get_prefill_wrapper(self):

@@ -615,6 +629,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

629 logits_soft_cap=self.logits_soft_cap,

630 q_data_type=self.q_data_type,

631 kv_data_type=self.kv_cache_dtype,

632 + fixed_split_size=self.prefill_fixed_split_size,

633 + disable_split_kv=self.disable_split_kv,

634 )

635 else:

636 attn_metadata.qo_indptr_gpu = qo_indptr_cpu.to(

@@ -668,6 +684,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

684 logits_soft_cap=self.logits_soft_cap,

685 q_data_type=self.q_data_type,

686 kv_data_type=self.kv_cache_dtype,

687 + fixed_split_size=self.decode_fixed_split_size,

688 + disable_split_kv=self.disable_split_kv,

689 )

690 return attn_metadata

691

@@ -1048,6 +1066,8 @@ def fast_plan_decode(

1066 rope_scale: Optional[float] = None,

1067 rope_theta: Optional[float] = None,

1068 non_blocking: bool = True,

1069 + fixed_split_size: int = -1,

1070 + disable_split_kv: bool = False,

1071 ) -> None:

1072 """

1073 A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for

@@ -1085,6 +1105,10 @@ def fast_plan_decode(

1105 rope_scale,

1106 rope_theta,

1107 non_blocking,

1108 + None, # block_tables

1109 + None, # seq_lens

1110 + fixed_split_size,

1111 + disable_split_kv,

1112 )

1113 self.vllm_first_call = False

1114 return

@@ -1130,7 +1154,7 @@ def fast_plan_decode(

1154 qo_indptr_host = _get_range_buf(batch_size + 1, "cpu")

1155

1156 try:

1133 - # Make sure we pass exactly 15 arguments for tensor core version

1157 + # Make sure we pass exactly 18 arguments for tensor core version

1158 self._plan_info = self._cached_module.plan(

1159 self._float_workspace_buffer,

1160 self._int_workspace_buffer,

@@ -1147,6 +1171,9 @@ def fast_plan_decode(

1171 head_dim,

1172 head_dim,

1173 False, # causal

1174 + window_left,

1175 + fixed_split_size,

1176 + disable_split_kv,

1177 )

1178 except Exception as e:

1179 raise RuntimeError(f"Error in tensor core plan: {e}") from e