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

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

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

Signed-off-by: Bram Wasti <bwasti@meta.com> Signed-off-by: Bram Wasti <bwasti@fb.com>

Bram Wasti committed Oct 13, 2025 at 07:24 UTC 3263799056f48a2653e8d75fc961b5eb973d7a91

4 files changed +81 -35

csrc/moe/topk_softmax_kernels.cu

+1 -3

@@ -21,7 +21,6 @@

21 #include <c10/cuda/CUDAGuard.h>

22 #include "../cuda_compat.h"

23 #include "../cub_helpers.h"

24 -#include "../core/batch_invariant.hpp"

25 #define MAX(a, b) ((a) > (b) ? (a) : (b))

26 #define MIN(a, b) ((a) < (b) ? (a) : (b))

@@ -406,8 +405,7 @@ void topkGatingSoftmaxLauncherHelper(const float* input, const bool* finished, f

405 using Constants = detail::TopkConstants<EXPERTS, BYTES_PER_LDG, WARP_SIZE_PARAM>;

406 static constexpr int VPT = Constants::VPT;

407 static constexpr int ROWS_PER_WARP = Constants::ROWS_PER_WARP;

409 - const bool batch_invariant_launch = vllm::vllm_kernel_override_batch_invariant();

410 - const int num_warps = batch_invariant_launch ? 32 : (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;

408 + const int num_warps = (num_rows + ROWS_PER_WARP - 1) / ROWS_PER_WARP;

409 const int num_blocks = (num_warps + WARPS_PER_TB - 1) / WARPS_PER_TB;

410

411 dim3 block_dim(WARP_SIZE_PARAM, WARPS_PER_TB);

tests/v1/generation/test_batch_invariance.py

+37 -28

@@ -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=max_batch_size,

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=max_batch_size,

133 gpu_memory_utilization=gpu_mem_util,

134 max_model_len=max_model_len,

135 swap_space=swap_space_gb,

@@ -135,15 +138,16 @@ 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(_random_prompt(min_random_prompt, max_random_prompt))

151

152 # Generate with the larger-batch engine

153 outputs = llm_bsN.generate(prompts, sampling)

@@ -154,19 +158,20 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle():

158 text = needle_output.outputs[0].text

159

160 if text != baseline_text:

161 + print(f"{text}\n\n== Not the same as ==\n\n{baseline_text}\n\n")

162 mismatches += 1

163

164 passes = num_trials - mismatches

165 # Dump how many passed vs failed

166 print(

167 f"[determinism] total={num_trials}, passed={passes}, "

163 - f"failed={mismatches}, batch_size={batch_size}"

168 + f"failed={mismatches}, max_batch_size={max_batch_size}"

169 )

170

171 if mismatches > 0:

172 pytest.fail(

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

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

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

175 )

176

177 finally:

@@ -199,8 +204,13 @@ def _extract_step_logprobs(request_output):

204 not torch.cuda.is_available(),

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

206 )

202 -def test_logprobs_bitwise_batch_invariance_bs1_vs_bs2():

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

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

208 +def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(backend):

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

@@ -208,16 +218,14 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bs2():

218 llm = LLM(

219 model=model_name,

220 tensor_parallel_size=tp_size,

211 - enforce_eager=True, # helps reduce nondeterminism from some backends

221 + enforce_eager=True,

222 + enable_prefix_caching=False,

223 )

224

214 - prompts = [

215 - "The capital of France is",

216 - "The capital of Germany is",

217 - ]

225 + prompts = [_random_prompt(10, 1024) for i in range(100)]

226

227 sp = SamplingParams(

220 - temperature=0.0,

228 + temperature=0.6,

229 top_p=1.0,

230 max_tokens=8,

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

@@ -238,11 +246,11 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bs2():

246 )

247 bs1_logprobs_per_prompt.append(step_logprobs)

248

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

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

250 # prompt.

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

252 assert len(outs_batched) == len(prompts)

245 - bs2_logprobs_per_prompt = []

253 + bsN_logprobs_per_prompt = []

254 for o in outs_batched:

255 step_logprobs = _extract_step_logprobs(o)

256 if step_logprobs is None:

@@ -250,17 +258,17 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bs2():

258 "Logits are not available on RequestOutput; "

259 "enable logprobs return to run this test."

260 )

253 - bs2_logprobs_per_prompt.append(step_logprobs)

261 + bsN_logprobs_per_prompt.append(step_logprobs)

262

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

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

257 - zip(bs1_logprobs_per_prompt, bs2_logprobs_per_prompt)

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

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

265 + zip(bs1_logprobs_per_prompt, bsN_logprobs_per_prompt)

266 ):

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

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

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

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

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

270 )

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

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

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

273 f"Logits shape mismatch at prompt {i}, step {t}: {a.shape} vs {b.shape}"

274 )

@@ -297,6 +305,7 @@ def LLM_with_max_seqs(

305 tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),

306 trust_remote_code=os.getenv("VLLM_TRUST_REMOTE_CODE", "0") == "1",

307 enable_prefix_caching=False,

308 + enforce_eager=True,

309 # Enable for MOE models

310 # enable_expert_parallel=True,

311 )

vllm/model_executor/layers/batch_invariant.py

+14 -1

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

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(

19 grid: Callable[..., Any], kernel: Any, args: dict[str, Any]

@@ -562,5 +566,14 @@ def vllm_kernel_override_batch_invariant():

566 def init_batch_invariance():

567 # this will hit all the csrc overrides as well

568 if vllm_kernel_override_batch_invariant():

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

569 + curr_attn_backend = envs.VLLM_ATTENTION_BACKEND

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

571 + if curr_attn_backend not in supported_backends:

572 + warning = (

573 + "Forcibly updating attention backend to"

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

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

576 + )

577 + logger.warning_once(warning)

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

579 enable_batch_invariant_mode()

vllm/v1/attention/backends/flashinfer.py

+29 -3

@@ -25,6 +25,9 @@ from vllm.attention.backends.abstract import (

25 )

26 from vllm.config import CUDAGraphMode, VllmConfig

27 from vllm.logger import init_logger

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

29 + vllm_kernel_override_batch_invariant,

30 +)

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

32 QuantKey,

33 kFp8StaticTensorSym,

@@ -50,6 +53,7 @@ from vllm.v1.attention.backends.utils import (

53 from vllm.v1.kv_cache_interface import AttentionSpec

55 FLASHINFER_WORKSPACE_BUFFER_SIZE = 256 * 1024 * 1024

56 +FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT = 2048 * 1024 * 1024

58 FP8_DTYPE = current_platform.fp8_dtype()

59 FP4_DTYPE = torch.uint8

@@ -288,6 +292,15 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

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

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

294

295 + if vllm_kernel_override_batch_invariant():

296 + self.decode_fixed_split_size = 2048

297 + self.prefill_fixed_split_size = 4096

298 + self.disable_split_kv = True

299 + else:

300 + self.decode_fixed_split_size = -1

301 + self.prefill_fixed_split_size = -1

302 + self.disable_split_kv = False

303 +

304 self.compilation_config = vllm_config.compilation_config

305 max_num_pages_per_req = cdiv(

306 self.model_config.max_model_len, self.kv_cache_spec.block_size

@@ -391,8 +404,11 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

404

405 def _get_workspace_buffer(self):

406 if self._workspace_buffer is None:

407 + buffer_size = FLASHINFER_WORKSPACE_BUFFER_SIZE

408 + if vllm_kernel_override_batch_invariant():

409 + buffer_size = FLASHINFER_WORKSPACE_BUFFER_SIZE_BATCH_INVARIANT

410 self._workspace_buffer = torch.zeros(

395 - FLASHINFER_WORKSPACE_BUFFER_SIZE, dtype=torch.uint8, device=self.device

411 + buffer_size, dtype=torch.uint8, device=self.device

412 )

413 return self._workspace_buffer

414

@@ -669,6 +685,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

685 logits_soft_cap=self.logits_soft_cap,

686 q_data_type=self.q_data_type,

687 kv_data_type=self.kv_cache_dtype,

688 + fixed_split_size=self.prefill_fixed_split_size,

689 + disable_split_kv=self.disable_split_kv,

690 )

691 else:

692 attn_metadata.qo_indptr_gpu = qo_indptr_cpu.to(

@@ -730,6 +748,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

748 logits_soft_cap=self.logits_soft_cap,

749 q_data_type=self.q_data_type,

750 kv_data_type=self.kv_cache_dtype,

751 + fixed_split_size=self.decode_fixed_split_size,

752 + disable_split_kv=self.disable_split_kv,

753 )

754 return attn_metadata

755

@@ -1121,6 +1141,8 @@ def fast_plan_decode(

1141 rope_scale: float | None = None,

1142 rope_theta: float | None = None,

1143 non_blocking: bool = True,

1144 + fixed_split_size: int = -1,

1145 + disable_split_kv: bool = False,

1146 ) -> None:

1147 """

1148 A faster version of BatchDecodeWithPagedKVCacheWrapper::plan used for

@@ -1157,6 +1179,10 @@ def fast_plan_decode(

1179 rope_scale,

1180 rope_theta,

1181 non_blocking,

1182 + None, # block_tables

1183 + None, # seq_lens

1184 + fixed_split_size,

1185 + disable_split_kv,

1186 )

1187 self.vllm_first_call = False

1188 return

@@ -1222,8 +1248,8 @@ def fast_plan_decode(

1248 head_dim,

1249 False, # causal

1250 window_left,

1225 - -1, # fixed_split_size

1226 - False, # disable_split_kv

1251 + fixed_split_size,

1252 + disable_split_kv,

1253 )

1254 except Exception as e:

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