@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7d8975de84
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

Deepseek-v3 Batch Invariant on 8xH100 (#26609)

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

Bram Wasti committed Oct 15, 2025 at 22:06 UTC 7d8975de84dadb87df23eae4a663a3c6620433dd
21 files changed +1567 -102
tests/v1/generation/test_batch_invariance.py
+769 -73
@@ -3,56 +3,73 @@
3 import contextlib
4 import os
5 import random
6 -import string
6
7 import pytest
8 import torch
9
10 from vllm import LLM, SamplingParams
11 +from vllm.platforms import current_platform
12 +
13 +
14 +@pytest.fixture(autouse=True)
15 +def enable_batch_invariant_mode():
16 + """Automatically enable batch invariant kernel overrides for all tests."""
17 + old_value = os.environ.get("VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT")
18 + os.environ["VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT"] = "1"
19 + yield
20 + # Restore original value after test
21 + if old_value is None:
22 + os.environ.pop("VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT", None)
23 + else:
24 + os.environ["VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT"] = old_value
25
26
27 def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:
15 - # Lightweight random prompt generator to vary prompt lengths and content.
16 - vocab = [
17 - "alpha",
18 - "bravo",
19 - "charlie",
20 - "delta",
21 - "echo",
22 - "foxtrot",
23 - "golf",
24 - "hotel",
25 - "india",
26 - "juliet",
27 - "kilo",
28 - "lima",
29 - "mike",
30 - "november",
31 - "oscar",
32 - "papa",
33 - "quebec",
34 - "romeo",
35 - "sierra",
36 - "tango",
37 - "uniform",
38 - "victor",
39 - "whiskey",
40 - "xray",
41 - "yankee",
42 - "zulu",
28 + # Generate more realistic prompts that will actually produce varied tokens
29 + # Use a mix of common English text patterns
30 +
31 + prompt_templates = [
32 + # Question-answer style
33 + "Question: What is the capital of France?\nAnswer: The capital of France is",
34 + "Q: How does photosynthesis work?\nA: Photosynthesis is the process by which",
35 + "User: Can you explain quantum mechanics?\nAssistant: Quantum mechanics is",
36 + # Story/narrative style
37 + "Once upon a time in a distant galaxy, there lived",
38 + "The old man walked slowly down the street, remembering",
39 + "In the year 2157, humanity finally discovered",
40 + # Technical/code style
41 + "To implement a binary search tree in Python, first we need to",
42 + "The algorithm works by iterating through the array and",
43 + "Here's how to optimize database queries using indexing:",
44 + # Factual/informative style
45 + "The Renaissance was a period in European history that",
46 + "Climate change is caused by several factors including",
47 + "The human brain contains approximately 86 billion neurons which",
48 + # Conversational style
49 + "I've been thinking about getting a new laptop because",
50 + "Yesterday I went to the store and bought",
51 + "My favorite thing about summer is definitely",
52 ]
44 - n = random.randint(min_words, max_words)
45 - words = random.choices(vocab, k=n)
53
47 - # Add some noise and punctuation variability
48 - if random.random() < 0.5:
49 - words[0] = words[0].capitalize()
50 - if random.random() < 0.2:
51 - words.append("".join(random.choices(string.ascii_lowercase, k=5)))
52 - punct = random.choice([".", "?", "!", "...", ""])
53 - return " ".join(words) + punct
54 + # Pick a random template
55 + base_prompt = random.choice(prompt_templates)
56 +
57 + # Add some padding to vary the length if needed
58 + if min_words > 50:
59 + # For longer prompts, repeat context
60 + padding_text = (
61 + " This is an interesting topic that deserves more explanation. "
62 + * (min_words // 50)
63 + )
64 + base_prompt = base_prompt + padding_text
65 +
66 + return base_prompt
67
68
69 +@pytest.mark.skipif(
70 + not current_platform.has_device_capability(90),
71 + reason="Batch invariance tests only supported on Hopper (SM90)",
72 +)
73 @pytest.mark.timeout(1000)
74 def test_v1_generation_is_deterministic_across_batch_sizes_with_needle():
75 """
@@ -91,7 +108,6 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle():
108 # Keep GPU memory usage low to avoid startup allocation failures.
109 gpu_mem_util = float(os.getenv("VLLM_GPU_MEMORY_UTILIZATION", "0.4"))
110 max_model_len = int(os.getenv("VLLM_MAX_MODEL_LEN", "5120"))
94 - swap_space_gb = int(os.getenv("VLLM_SWAP_SPACE_GB", "4"))
111
112 # Sampling parameters: longer outputs with a more random-sounding
113 # continuation,but still deterministic due to fixed seed.
@@ -117,7 +133,6 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle():
133 max_num_seqs=max_batch_size,
134 gpu_memory_utilization=gpu_mem_util,
135 max_model_len=max_model_len,
120 - swap_space=swap_space_gb,
136 )
137
138 # Baseline generation for the needle prompt alone.
@@ -132,7 +147,6 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle():
147 max_num_seqs=max_batch_size,
148 gpu_memory_utilization=gpu_mem_util,
149 max_model_len=max_model_len,
135 - swap_space=swap_space_gb,
150 )
151
152 mismatches = 0
@@ -195,16 +209,21 @@ def _extract_step_logprobs(request_output):
209 ],
210 dtype=torch.float32,
211 )
198 - return t
212 + return t, inner.token_ids
213
200 - return None
214 + return None, None
215
216
217 +@pytest.mark.skipif(
218 + not current_platform.has_device_capability(90),
219 + reason="Batch invariance tests only supported on Hopper (SM90)",
220 +)
221 @pytest.mark.skipif(
222 not torch.cuda.is_available(),
223 reason="Requires CUDA to match production inference path.",
224 )
207 -@pytest.mark.parametrize("backend", ["FLEX_ATTENTION", "FLASHINFER"])
225 +@pytest.mark.parametrize("backend", ["FLASH_ATTN", "FLASHINFER"])
226 +@pytest.mark.forked
227 def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(backend):
228 backend = os.getenv("VLLM_ATTENTION_BACKEND", backend)
229 os.environ["VLLM_ATTENTION_BACKEND"] = backend
@@ -214,69 +233,754 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(backend):
233 model_name = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
234 tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
235
217 - # Force float32 to avoid precision-induced differences.
236 + # For batch invariance, disable custom all-reduce to ensure deterministic
237 + # all-reduce operations (custom all-reduce may not be deterministic)
238 + from vllm.model_executor.layers.batch_invariant import (
239 + vllm_kernel_override_batch_invariant,
240 + )
241 +
242 + disable_custom_ar = vllm_kernel_override_batch_invariant()
243 +
244 + if disable_custom_ar:
245 + print(f"\n{'=' * 80}")
246 + print(f"BATCH INVARIANCE MODE: Disabling custom all-reduce (TP={tp_size})")
247 + print(f"{'=' * 80}\n")
248 +
249 llm = LLM(
250 model=model_name,
251 tensor_parallel_size=tp_size,
221 - enforce_eager=True,
252 enable_prefix_caching=False,
253 + max_num_seqs=32,
254 + max_model_len=8192,
255 + dtype="bfloat16", # not everything is supported
256 )
257
225 - prompts = [_random_prompt(10, 1024) for i in range(100)]
258 + # Use more realistic prompts for better token generation
259 + prompts = [_random_prompt(10, 50) for i in range(32)]
260
261 sp = SamplingParams(
262 temperature=0.6,
263 top_p=1.0,
264 max_tokens=8,
231 - # Seed shouldn't matter at temperature=0, but keeping it stable anyway.
265 seed=1234,
266 logprobs=5,
267 )
268
269 # BS=1: run prompts individually and collect logprobs per step.
270 + print("\n" + "=" * 80)
271 + print("STARTING BS=1 RUNS (each prompt individually)")
272 + print("=" * 80 + "\n")
273 +
274 bs1_logprobs_per_prompt = []
238 - for p in prompts:
275 + bs1_tokens_per_prompt = []
276 + for idx, p in enumerate(prompts):
277 + print(f"\n[BS=1] Running prompt {idx}/{len(prompts)} - Preview: {p[:80]}...")
278 outs = llm.generate([p], sp, use_tqdm=False)
279 assert len(outs) == 1
241 - step_logprobs = _extract_step_logprobs(outs[0])
280 + step_logprobs, token_ids = _extract_step_logprobs(outs[0])
281 if step_logprobs is None:
282 pytest.skip(
283 "Logits are not available on RequestOutput; "
284 "enable logprobs return to run this test."
285 )
286 bs1_logprobs_per_prompt.append(step_logprobs)
287 + bs1_tokens_per_prompt.append(token_ids)
288 + print(f"[BS=1] Prompt {idx} generated tokens: {token_ids}")
289
290 # BS=N: run prompts in a batch and collect logprobs per step for each
291 # prompt.
292 + print("\n" + "=" * 80)
293 + print(f"STARTING BS={len(prompts)} RUN (all prompts batched)")
294 + print("=" * 80 + "\n")
295 +
296 outs_batched = llm.generate(prompts, sp, use_tqdm=False)
297 assert len(outs_batched) == len(prompts)
298 bsN_logprobs_per_prompt = []
254 - for o in outs_batched:
255 - step_logprobs = _extract_step_logprobs(o)
299 + bsN_tokens_per_prompt = []
300 +
301 + print(f"\n[BS={len(prompts)}] Processing batched outputs...")
302 + for idx, o in enumerate(outs_batched):
303 + tokens = o.outputs[0].token_ids if o.outputs else "N/A"
304 + print(f"[BS={len(prompts)}] Prompt {idx} generated tokens: {tokens}")
305 + step_logprobs, token_ids = _extract_step_logprobs(o)
306 if step_logprobs is None:
307 pytest.skip(
308 "Logits are not available on RequestOutput; "
309 "enable logprobs return to run this test."
310 )
311 bsN_logprobs_per_prompt.append(step_logprobs)
312 + bsN_tokens_per_prompt.append(token_ids)
313
314 # 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 - ):
267 - assert len(logprobs_bs1) == len(logprobs_bsN), (
268 - f"Different number of generation steps for prompt index {i}: "
269 - f"{len(logprobs_bs1)} (BS=1) vs {len(logprobs_bsN)} (BS=N)"
315 + failed_prompts = []
316 + for i, (logprobs_bs1, logprobs_bsN, tokens_bs1, tokens_bsN) in enumerate(
317 + zip(
318 + bs1_logprobs_per_prompt,
319 + bsN_logprobs_per_prompt,
320 + bs1_tokens_per_prompt,
321 + bsN_tokens_per_prompt,
322 )
323 + ):
324 + if len(logprobs_bs1) != len(logprobs_bsN):
325 + reason = (
326 + f"Different number of steps: {len(logprobs_bs1)} (BS=1) "
327 + f"vs {len(logprobs_bsN)} (BS=N)"
328 + )
329 + failed_prompts.append(
330 + {
331 + "prompt_idx": i,
332 + "step": "all",
333 + "reason": reason,
334 + "prompt_preview": prompts[i][:100],
335 + "bs1_tokens": tokens_bs1,
336 + "bsN_tokens": tokens_bsN,
337 + }
338 + )
339 + continue
340 +
341 + # Check if tokens match first
342 + if tokens_bs1 != tokens_bsN:
343 + failed_prompts.append(
344 + {
345 + "prompt_idx": i,
346 + "step": "sampling",
347 + "reason": "Different tokens sampled",
348 + "prompt_preview": prompts[i][:100],
349 + "bs1_tokens": tokens_bs1,
350 + "bsN_tokens": tokens_bsN,
351 + "bs1_all_logprobs": [
352 + logprobs_bs1[s].tolist() for s in range(len(logprobs_bs1))
353 + ],
354 + "bsN_all_logprobs": [
355 + logprobs_bsN[s].tolist() for s in range(len(logprobs_bsN))
356 + ],
357 + }
358 + )
359 + continue
360 +
361 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}"
362 + if a.shape != b.shape:
363 + failed_prompts.append(
364 + {
365 + "prompt_idx": i,
366 + "step": t,
367 + "reason": f"Shape mismatch: {a.shape} vs {b.shape}",
368 + "prompt_preview": prompts[i][:100],
369 + "bs1_tokens": tokens_bs1,
370 + "bsN_tokens": tokens_bsN,
371 + }
372 + )
373 + break
374 +
375 + if not torch.equal(a, b):
376 + max_diff = torch.abs(a - b).max().item()
377 + # Print which token failed
378 + print(f"\n[DIVERGENCE] Prompt {i}, Token {t}: max_diff={max_diff:.6e}")
379 + bs1_tok = tokens_bs1[t] if t < len(tokens_bs1) else "N/A"
380 + bsN_tok = tokens_bsN[t] if t < len(tokens_bsN) else "N/A"
381 + print(f" Token IDs: bs1={bs1_tok}, bsN={bsN_tok}")
382 + print(f" BS=1 logprob: {a.tolist()}")
383 + print(f" BS=N logprob: {b.tolist()}")
384 + failed_prompts.append(
385 + {
386 + "prompt_idx": i,
387 + "step": t,
388 + "reason": f"Bitwise mismatch (max_diff={max_diff:.6e})",
389 + "prompt_preview": prompts[i][:100],
390 + "bs1_tokens": tokens_bs1,
391 + "bsN_tokens": tokens_bsN,
392 + "bs1_all_logprobs": [
393 + logprobs_bs1[s].tolist() for s in range(len(logprobs_bs1))
394 + ],
395 + "bsN_all_logprobs": [
396 + logprobs_bsN[s].tolist() for s in range(len(logprobs_bsN))
397 + ],
398 + }
399 + )
400 + break
401 +
402 + # Print summary of all failures
403 + if failed_prompts:
404 + print(f"\n{'=' * 80}")
405 + fail_msg = (
406 + f"BATCH INVARIANCE FAILURES: {len(failed_prompts)}/"
407 + f"{len(prompts)} prompts failed"
408 + )
409 + print(fail_msg)
410 + print(f"{'=' * 80}")
411 + for fail in failed_prompts:
412 + print(f"\nPrompt {fail['prompt_idx']} (step {fail['step']}):")
413 + print(f" Reason: {fail['reason']}")
414 + print(f" Preview: {fail['prompt_preview']}...")
415 +
416 + # Always show the tokens
417 + if "bs1_tokens" in fail:
418 + print(f" BS=1 tokens: {fail['bs1_tokens']}")
419 + if "bsN_tokens" in fail:
420 + print(f" BS=N tokens: {fail['bsN_tokens']}")
421 +
422 + if "bs1_all_logprobs" in fail:
423 + print(f" BS=1 logprobs for all {len(fail['bs1_all_logprobs'])} steps:")
424 + for step_idx, logprobs in enumerate(fail["bs1_all_logprobs"]):
425 + print(f" Step {step_idx}: {logprobs}")
426 + print(f" BS=N logprobs for all {len(fail['bsN_all_logprobs'])} steps:")
427 + for step_idx, logprobs in enumerate(fail["bsN_all_logprobs"]):
428 + print(f" Step {step_idx}: {logprobs}")
429 + print(f"{'=' * 80}\n")
430 +
431 + # Fail the test with summary
432 + msg = (
433 + f"Batch invariance violated in {len(failed_prompts)}/"
434 + f"{len(prompts)} prompts. See output above for details."
435 + )
436 + pytest.fail(msg)
437 +
438 +
439 +@pytest.mark.skipif(
440 + not current_platform.has_device_capability(90),
441 + reason="Batch invariance tests only supported on Hopper (SM90)",
442 +)
443 +def test_simple_generation():
444 + """
445 + Simple test that runs the model with a basic prompt and prints the output.
446 + Useful for quick smoke testing and debugging.
447 + """
448 + model = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
449 +
450 + llm = LLM(
451 + model=model,
452 + max_num_seqs=1,
453 + tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
454 + enforce_eager=True,
455 + gpu_memory_utilization=0.9,
456 + max_model_len=2048,
457 + dtype="bfloat16",
458 + enable_prefix_caching=False,
459 + )
460 +
461 + prompt = "the capital of france is"
462 + sampling_params = SamplingParams(
463 + temperature=0.0,
464 + max_tokens=20,
465 + )
466 +
467 + print(f"\n{'=' * 80}")
468 + print("Running simple generation test")
469 + print(f"Prompt: '{prompt}'")
470 + print(f"{'=' * 80}\n")
471 +
472 + try:
473 + outputs = llm.generate([prompt], sampling_params)
474 +
475 + assert len(outputs) == 1
476 + output_text = outputs[0].outputs[0].text
477 +
478 + print(f"Output: '{output_text}'")
479 + print(f"\n{'=' * 80}")
480 + print(f"Full completion: '{prompt}{output_text}'")
481 + print(f"{'=' * 80}\n")
482 +
483 + finally:
484 + with contextlib.suppress(Exception):
485 + llm.shutdown()
486 +
487 +
488 +@pytest.mark.skipif(
489 + not current_platform.has_device_capability(90),
490 + reason="Batch invariance tests only supported on Hopper (SM90)",
491 +)
492 +@pytest.mark.skipif(
493 + not torch.cuda.is_available(),
494 + reason="Requires CUDA to match production inference path.",
495 +)
496 +@pytest.mark.parametrize("backend", ["FLASH_ATTN", "FLASHINFER"])
497 +@pytest.mark.forked
498 +def test_logprobs_WITHOUT_batch_invariance_should_FAIL(backend):
499 + """
500 + This test is the inverse of test_logprobs_bitwise_batch_invariance_bs1_vs_bsN.
501 + It DISABLES batch invariance mode and expects to see non-deterministic behavior
502 + between BS=1 and BS=N runs. This demonstrates that batch invariance is actually
503 + doing something useful.
504 +
505 + The test will PASS if we detect differences (proving batch invariance matters).
506 + The test will FAIL if everything matches (suggesting batch invariance isn't needed).
507 + """
508 + backend = os.getenv("VLLM_ATTENTION_BACKEND", backend)
509 + os.environ["VLLM_ATTENTION_BACKEND"] = backend
510 +
511 + # CRITICAL: Disable batch invariance for this test
512 + old_value = os.environ.get("VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT")
513 + os.environ["VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT"] = "0"
514 +
515 + try:
516 + seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
517 + random.seed(seed)
518 + model_name = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
519 + tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
520 +
521 + print(f"\n{'=' * 80}")
522 + print("BATCH INVARIANCE DISABLED: Expecting non-deterministic behavior")
523 + print(f"{'=' * 80}\n")
524 +
525 + llm = LLM(
526 + model=model_name,
527 + tensor_parallel_size=tp_size,
528 + enable_prefix_caching=False,
529 + max_num_seqs=32,
530 + max_model_len=8192,
531 + dtype="bfloat16",
532 + )
533 +
534 + # Use more realistic prompts for better token generation
535 + prompts = [_random_prompt(10, 50) for i in range(32)]
536 +
537 + sp = SamplingParams(
538 + temperature=0.6,
539 + top_p=1.0,
540 + max_tokens=8,
541 + seed=1234,
542 + logprobs=5,
543 + )
544 +
545 + # BS=1: run prompts individually and collect logprobs per step.
546 + print("\n" + "=" * 80)
547 + print("STARTING BS=1 RUNS (each prompt individually)")
548 + print("=" * 80 + "\n")
549 +
550 + bs1_logprobs_per_prompt = []
551 + bs1_tokens_per_prompt = []
552 + for idx, p in enumerate(prompts):
553 + print(
554 + f"\n[BS=1] Running prompt {idx}/{len(prompts)} - Preview: {p[:80]}..."
555 + )
556 + outs = llm.generate([p], sp, use_tqdm=False)
557 + assert len(outs) == 1
558 + step_logprobs, token_ids = _extract_step_logprobs(outs[0])
559 + if step_logprobs is None:
560 + pytest.skip(
561 + "Logits are not available on RequestOutput; "
562 + "enable logprobs return to run this test."
563 + )
564 + bs1_logprobs_per_prompt.append(step_logprobs)
565 + bs1_tokens_per_prompt.append(token_ids)
566 + print(f"[BS=1] Prompt {idx} generated tokens: {token_ids}")
567 +
568 + # BS=N: run prompts in a batch and collect logprobs per step for each prompt.
569 + print("\n" + "=" * 80)
570 + print(f"STARTING BS={len(prompts)} RUN (all prompts batched)")
571 + print("=" * 80 + "\n")
572 +
573 + outs_batched = llm.generate(prompts, sp, use_tqdm=False)
574 + assert len(outs_batched) == len(prompts)
575 + bsN_logprobs_per_prompt = []
576 + bsN_tokens_per_prompt = []
577 +
578 + print(f"\n[BS={len(prompts)}] Processing batched outputs...")
579 + for idx, o in enumerate(outs_batched):
580 + tokens = o.outputs[0].token_ids if o.outputs else "N/A"
581 + print(f"[BS={len(prompts)}] Prompt {idx} generated tokens: {tokens}")
582 + step_logprobs, token_ids = _extract_step_logprobs(o)
583 + if step_logprobs is None:
584 + pytest.skip(
585 + "Logits are not available on RequestOutput; "
586 + "enable logprobs return to run this test."
587 + )
588 + bsN_logprobs_per_prompt.append(step_logprobs)
589 + bsN_tokens_per_prompt.append(token_ids)
590 +
591 + # Compare step-by-step logprobs for each prompt between BS=1 and BS=N runs.
592 + differences_found = []
593 + for i, (logprobs_bs1, logprobs_bsN, tokens_bs1, tokens_bsN) in enumerate(
594 + zip(
595 + bs1_logprobs_per_prompt,
596 + bsN_logprobs_per_prompt,
597 + bs1_tokens_per_prompt,
598 + bsN_tokens_per_prompt,
599 )
275 - # Bitwise exact equality.
276 - assert torch.equal(a, b), (
277 - f"Bitwise logprobs mismatch at prompt {i}, step {t} "
278 - f"(dtype={a.dtype}, shape={a.shape})."
600 + ):
601 + if len(logprobs_bs1) != len(logprobs_bsN):
602 + reason = (
603 + f"Different number of steps: {len(logprobs_bs1)} (BS=1) "
604 + f"vs {len(logprobs_bsN)} (BS=N)"
605 + )
606 + differences_found.append(
607 + {
608 + "prompt_idx": i,
609 + "step": "all",
610 + "reason": reason,
611 + "prompt_preview": prompts[i][:100],
612 + "bs1_tokens": tokens_bs1,
613 + "bsN_tokens": tokens_bsN,
614 + }
615 + )
616 + continue
617 +
618 + # Check if tokens match first
619 + if tokens_bs1 != tokens_bsN:
620 + differences_found.append(
621 + {
622 + "prompt_idx": i,
623 + "step": "sampling",
624 + "reason": "Different tokens sampled",
625 + "prompt_preview": prompts[i][:100],
626 + "bs1_tokens": tokens_bs1,
627 + "bsN_tokens": tokens_bsN,
628 + }
629 + )
630 + continue
631 +
632 + for t, (a, b) in enumerate(zip(logprobs_bs1, logprobs_bsN)):
633 + if a.shape != b.shape:
634 + differences_found.append(
635 + {
636 + "prompt_idx": i,
637 + "step": t,
638 + "reason": f"Shape mismatch: {a.shape} vs {b.shape}",
639 + "prompt_preview": prompts[i][:100],
640 + "bs1_tokens": tokens_bs1,
641 + "bsN_tokens": tokens_bsN,
642 + }
643 + )
644 + break
645 +
646 + if not torch.equal(a, b):
647 + max_diff = torch.abs(a - b).max().item()
648 + print(
649 + f"\n[EXPECTED DIVERGENCE FOUND] Prompt {i}, "
650 + f"Token {t}: max_diff={max_diff:.6e}"
651 + )
652 + bs1_tok = tokens_bs1[t] if t < len(tokens_bs1) else "N/A"
653 + bsN_tok = tokens_bsN[t] if t < len(tokens_bsN) else "N/A"
654 + print(f" Token IDs: bs1={bs1_tok}, bsN={bsN_tok}")
655 + print(f" BS=1 logprob: {a.tolist()}")
656 + print(f" BS=N logprob: {b.tolist()}")
657 + differences_found.append(
658 + {
659 + "prompt_idx": i,
660 + "step": t,
661 + "reason": f"Bitwise mismatch (max_diff={max_diff:.6e})",
662 + "prompt_preview": prompts[i][:100],
663 + "bs1_tokens": tokens_bs1,
664 + "bsN_tokens": tokens_bsN,
665 + }
666 + )
667 + break
668 +
669 + # Print summary
670 + print(f"\n{'=' * 80}")
671 + if differences_found:
672 + success_msg = (
673 + f"✓ SUCCESS: Batch invariance is doing something! "
674 + f"Found {len(differences_found)}/{len(prompts)} prompts "
675 + f"with differences when batch invariance was DISABLED."
676 )
677 + print(success_msg)
678 + print(f"{'=' * 80}")
679 + for diff in differences_found:
680 + print(f"\nPrompt {diff['prompt_idx']} (step {diff['step']}):")
681 + print(f" Reason: {diff['reason']}")
682 + print(f" Preview: {diff['prompt_preview']}...")
683 + if "bs1_tokens" in diff:
684 + print(f" BS=1 tokens: {diff['bs1_tokens']}")
685 + if "bsN_tokens" in diff:
686 + print(f" BS=N tokens: {diff['bsN_tokens']}")
687 + print(f"{'=' * 80}\n")
688 + # Test PASSES because we found differences (batch invariance matters!)
689 + return
690 + else:
691 + # Test FAILS because everything matched even without batch invariance
692 + fail_msg = (
693 + f"✗ UNEXPECTED: All {len(prompts)} prompts matched "
694 + f"between BS=1 and BS=N even with batch invariance DISABLED. "
695 + f"This suggests batch invariance might not be necessary, "
696 + f"or the test needs more sensitive prompts."
697 + )
698 + print(fail_msg)
699 + print(f"{'=' * 80}\n")
700 + pytest.fail(fail_msg)
701 +
702 + finally:
703 + # Restore original value
704 + if old_value is None:
705 + os.environ.pop("VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT", None)
706 + else:
707 + os.environ["VLLM_KERNEL_OVERRIDE_BATCH_INVARIANT"] = old_value
708 +
709 +
710 +@pytest.mark.skipif(
711 + not current_platform.has_device_capability(90),
712 + reason="Batch invariance tests only supported on Hopper (SM90)",
713 +)
714 +@pytest.mark.skipif(
715 + not torch.cuda.is_available(),
716 + reason="Requires CUDA to match production inference path.",
717 +)
718 +@pytest.mark.parametrize("backend", ["FLASH_ATTN"])
719 +@pytest.mark.forked
720 +def test_decode_logprobs_match_prefill_logprobs(backend):
721 + """
722 + Test that verifies decode logprobs match prefill logprobs.
723 +
724 + For each decoded token at position i:
725 + 1. Run decode to generate N tokens and collect their logprobs
726 + 2. For each position i in [0, N):
727 + - Take prefix = prompt + tokens[0:i]
728 + - Run prefill(prefix + tokens[i]) to get logprob of tokens[i]
729 + - Verify prefill logprob matches decode logprob bitwise
730 +
731 + This ensures that the logprobs from decode are consistent with what
732 + we would get if we ran prefill on each prefix.
733 + """
734 + backend = os.getenv("VLLM_ATTENTION_BACKEND", backend)
735 + os.environ["VLLM_ATTENTION_BACKEND"] = backend
736 +
737 + seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
738 + random.seed(seed)
739 + model_name = os.getenv("VLLM_TEST_MODEL", "Qwen/Qwen3-1.7B")
740 + tp_size = int(os.getenv("VLLM_TEST_TP_SIZE", "1"))
741 +
742 + from vllm.model_executor.layers.batch_invariant import (
743 + vllm_kernel_override_batch_invariant,
744 + )
745 +
746 + disable_custom_ar = vllm_kernel_override_batch_invariant()
747 +
748 + if disable_custom_ar:
749 + print(f"\n{'=' * 80}")
750 + print(f"BATCH INVARIANCE MODE: Disabling custom all-reduce (TP={tp_size})")
751 + print(f"{'=' * 80}\n")
752 +
753 + llm = LLM(
754 + model=model_name,
755 + tensor_parallel_size=tp_size,
756 + enable_prefix_caching=False,
757 + max_num_seqs=32,
758 + max_model_len=8192,
759 + dtype="bfloat16",
760 + )
761 +
762 + # Use a few test prompts
763 + num_test_prompts = int(os.getenv("VLLM_DECODE_PREFILL_NUM_PROMPTS", "4"))
764 + prompts = [_random_prompt(10, 50) for _ in range(num_test_prompts)]
765 +
766 + # Generate longer sequences to test multiple decode steps
767 + max_tokens = int(os.getenv("VLLM_DECODE_PREFILL_MAX_TOKENS", "16"))
768 +
769 + sp = SamplingParams(
770 + temperature=0.0, # Greedy for determinism
771 + max_tokens=max_tokens,
772 + logprobs=5,
773 + )
774 +
775 + print("\n" + "=" * 80)
776 + print("STEP 1: Running decode to generate tokens and collect logprobs")
777 + print("=" * 80 + "\n")
778 +
779 + # Step 1: Run decode and collect logprobs
780 + decode_outputs = llm.generate(prompts, sp, use_tqdm=False)
781 +
782 + failed_comparisons = []
783 +
784 + for prompt_idx, (prompt, decode_output) in enumerate(zip(prompts, decode_outputs)):
785 + print(f"\n[Prompt {prompt_idx}] Testing: {prompt[:80]}...")
786 +
787 + # Extract decode logprobs and tokens
788 + decode_logprobs, token_ids = _extract_step_logprobs(decode_output)
789 + if decode_logprobs is None:
790 + pytest.skip(
791 + "Logprobs are not available on RequestOutput; "
792 + "enable logprobs return to run this test."
793 + )
794 +
795 + print(f"[Prompt {prompt_idx}] Generated {len(token_ids)} tokens: {token_ids}")
796 + print(f"[Prompt {prompt_idx}] Decode logprobs: {decode_logprobs.tolist()}")
797 +
798 + # Step 2: For each token position, run prefill and compare
799 + print(f"\n[Prompt {prompt_idx}] Verifying each token via prefill...")
800 +
801 + for token_idx in range(len(token_ids)):
802 + # Construct the prefix up to (but not including) this token
803 + current_token = token_ids[token_idx]
804 +
805 + # We need to detokenize to get the text prefix
806 + # For this, we'll use the tokenizer from the LLM
807 + # However, the LLM API doesn't expose tokenizer easily, so we'll
808 + # construct the prefix by decoding from the original prompt
809 +
810 + # Get text up to this point by using the output text
811 + # This is approximate but should work for verification
812 + if token_idx == 0:
813 + prefix_prompt = prompt
814 + else:
815 + # Use the partial output text up to this token
816 + # We'll need to construct this from the full output
817 + prefix_output = decode_output.outputs[0]
818 + # Get the text for tokens 0 to token_idx-1
819 + # Unfortunately, we don't have per-token text, so we'll use
820 + # a different approach: run prefill with prompt + tokens[0:token_idx]
821 +
822 + # Actually, we need to get the actual text. Let's use a workaround:
823 + # Run a generation with max_tokens = token_idx to get that prefix
824 + prefix_sp = SamplingParams(
825 + temperature=0.0,
826 + max_tokens=token_idx,
827 + logprobs=1,
828 + )
829 + prefix_output = llm.generate([prompt], prefix_sp, use_tqdm=False)[0]
830 + prefix_prompt = prompt + prefix_output.outputs[0].text
831 +
832 + # Now run prefill with max_tokens=1 to get the logprob of the next token
833 + prefill_sp = SamplingParams(
834 + temperature=0.0,
835 + max_tokens=1,
836 + logprobs=5,
837 + )
838 +
839 + print(
840 + f" [Token {token_idx}] Running prefill for prefix "
841 + f"(len={len(prefix_prompt)})..."
842 + )
843 + prefill_output = llm.generate([prefix_prompt], prefill_sp, use_tqdm=False)[
844 + 0
845 + ]
846 + prefill_logprobs, prefill_token_ids = _extract_step_logprobs(prefill_output)
847 +
848 + if prefill_logprobs is None:
849 + print(f" [Token {token_idx}] Warning: No prefill logprobs available")
850 + continue
851 +
852 + # The first token from prefill should match the current token
853 + prefill_token = prefill_token_ids[0]
854 + prefill_logprob = prefill_logprobs[0].item()
855 + decode_logprob = decode_logprobs[token_idx].item()
856 +
857 + print(
858 + f" [Token {token_idx}] Decode token: {current_token}, "
859 + f"logprob: {decode_logprob:.8f}"
860 + )
861 + print(
862 + f" [Token {token_idx}] Prefill token: {prefill_token}, "
863 + f"logprob: {prefill_logprob:.8f}"
864 + )
865 +
866 + # Check if tokens match
867 + if current_token != prefill_token:
868 + failed_comparisons.append(
869 + {
870 + "prompt_idx": prompt_idx,
871 + "token_idx": token_idx,
872 + "reason": "Token mismatch",
873 + "decode_token": current_token,
874 + "prefill_token": prefill_token,
875 + "decode_logprob": decode_logprob,
876 + "prefill_logprob": prefill_logprob,
877 + "prompt_text": prompt[:100],
878 + "prefix_text": prefix_prompt[:100],
879 + }
880 + )
881 + print(f" [Token {token_idx}] ✗ TOKEN MISMATCH!")
882 + continue
883 +
884 + # Check if logprobs match bitwise
885 + if decode_logprob != prefill_logprob:
886 + diff = abs(decode_logprob - prefill_logprob)
887 + failed_comparisons.append(
888 + {
889 + "prompt_idx": prompt_idx,
890 + "token_idx": token_idx,
891 + "reason": "Logprob mismatch",
892 + "decode_token": current_token,
893 + "prefill_token": prefill_token,
894 + "decode_logprob": decode_logprob,
895 + "prefill_logprob": prefill_logprob,
896 + "diff": diff,
897 + "prompt_text": prompt[:100],
898 + "prefix_text": prefix_prompt[:100],
899 + "decode_all_tokens": token_ids,
900 + "decode_all_logprobs": decode_logprobs.tolist(),
901 + }
902 + )
903 + print(f" [Token {token_idx}] ✗ LOGPROB MISMATCH! diff={diff:.8e}")
904 + else:
905 + print(f" [Token {token_idx}] ✓ Match (bitwise equal)")
906 +
907 + # Print summary
908 + print(f"\n{'=' * 80}")
909 + if failed_comparisons:
910 + print(f"DECODE-PREFILL MISMATCH: {len(failed_comparisons)} failures detected")
911 + print(f"{'=' * 80}")
912 +
913 + # Group failures by prompt for better readability
914 + failures_by_prompt: dict[int, list[dict]] = {}
915 + for fail in failed_comparisons:
916 + pid = fail["prompt_idx"]
917 + if pid not in failures_by_prompt:
918 + failures_by_prompt[pid] = []
919 + failures_by_prompt[pid].append(fail)
920 +
921 + for prompt_idx, failures in failures_by_prompt.items():
922 + print(f"\n{'=' * 80}")
923 + print(f"PROMPT {prompt_idx}: {failures[0]['prompt_text']}...")
924 + print(f"{'=' * 80}")
925 + print(f"Total failures for this prompt: {len(failures)}")
926 +
927 + # Show where mismatches occur (which token positions)
928 + mismatch_positions = [f["token_idx"] for f in failures]
929 + print(f"Mismatch at token positions: {mismatch_positions}")
930 +
931 + # Show first few failures in detail
932 + for i, fail in enumerate(failures[:5]): # Show first 5 failures per prompt
933 + print(f"\n [Failure {i + 1}] Token position {fail['token_idx']}:")
934 + print(f" Reason: {fail['reason']}")
935 + print(f" Prefix text: '{fail['prefix_text']}...'")
936 + print(
937 + f" Decode: token={fail['decode_token']}, "
938 + f"logprob={fail['decode_logprob']:.10f}"
939 + )
940 + print(
941 + f" Prefill: token={fail['prefill_token']}, "
942 + f"logprob={fail['prefill_logprob']:.10f}"
943 + )
944 + if "diff" in fail:
945 + print(f" Difference: {fail['diff']:.10e}")
946 + # Show in hex to see bitwise difference
947 + import struct
948 +
949 + decode_hex = struct.pack("f", fail["decode_logprob"]).hex()
950 + prefill_hex = struct.pack("f", fail["prefill_logprob"]).hex()
951 + print(f" Decode logprob (hex): 0x{decode_hex}")
952 + print(f" Prefill logprob (hex): 0x{prefill_hex}")
953 +
954 + # If we have all tokens/logprobs, show the context
955 + if "decode_all_tokens" in fail and "decode_all_logprobs" in fail:
956 + token_idx = fail["token_idx"]
957 + all_tokens = fail["decode_all_tokens"]
958 + all_logprobs = fail["decode_all_logprobs"]
959 +
960 + # Show context: 2 tokens before and after
961 + start = max(0, token_idx - 2)
962 + end = min(len(all_tokens), token_idx + 3)
963 +
964 + print(f" Context (tokens {start} to {end - 1}):")
965 + for j in range(start, end):
966 + marker = " <-- MISMATCH" if j == token_idx else ""
967 + print(
968 + f" [{j}] token={all_tokens[j]}, "
969 + f"logprob={all_logprobs[j]:.8f}{marker}"
970 + )
971 +
972 + if len(failures) > 5:
973 + print(f"\n ... and {len(failures) - 5} more failures for this prompt")
974 +
975 + print(f"\n{'=' * 80}\n")
976 +
977 + pytest.fail(
978 + f"Decode logprobs do not match prefill logprobs: "
979 + f"{len(failed_comparisons)} mismatches found."
980 + )
981 + else:
982 + print("✓ SUCCESS: All decode logprobs match prefill logprobs bitwise!")
983 + print(f"{'=' * 80}\n")
984
985
986 def LLM_with_max_seqs(
@@ -284,7 +988,6 @@ def LLM_with_max_seqs(
988 max_num_seqs: int,
989 gpu_memory_utilization: float,
990 max_model_len: int,
287 - swap_space: int,
991 ) -> LLM:
992 """
993 Helper to construct an LLM with a specific max_num_seqs (batch-size limit)
@@ -293,17 +996,10 @@ def LLM_with_max_seqs(
996 return LLM(
997 model=model,
998 max_num_seqs=max_num_seqs,
296 - # Constrain GPU memory pool so test can run even on busy GPUs.
999 gpu_memory_utilization=gpu_memory_utilization,
298 - # Keep KV cache footprint small while allowing longer outputs.
1000 max_model_len=max_model_len,
300 - # Allow some CPU offload if needed.
301 - swap_space=swap_space,
302 - # Keep things lean and CI-friendly.
303 - dtype="auto",
304 - # Single-GPU by default; override externally if desired.
1001 + dtype="bfloat16",
1002 tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
306 - trust_remote_code=os.getenv("VLLM_TRUST_REMOTE_CODE", "0") == "1",
1003 enable_prefix_caching=False,
1004 enforce_eager=True,
1005 # Enable for MOE models
tests/v1/generation/test_rms_norm_batch_invariant.py new
+346
@@ -0,0 +1,346 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Test batch-invariant RMS normalization against standard implementations.
5 +
6 +This test compares the Triton-based batch-invariant RMS norm implementation
7 +with the standard CUDA-based implementation to ensure numerical accuracy.
8 +"""
9 +
10 +import pytest
11 +import torch
12 +
13 +from vllm.model_executor.layers.batch_invariant import rms_norm as triton_rms_norm
14 +from vllm.model_executor.layers.layernorm import RMSNorm
15 +from vllm.platforms import current_platform
16 +
17 +
18 +@pytest.mark.skipif(
19 + not current_platform.has_device_capability(90),
20 + reason="Batch invariance tests only supported on Hopper (SM90)",
21 +)
22 +@pytest.mark.skipif(
23 + not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
24 +)
25 +@pytest.mark.parametrize("batch_size", [1, 4, 16, 64])
26 +@pytest.mark.parametrize("hidden_size", [512, 2048, 4096, 8192])
27 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
28 +@pytest.mark.parametrize("eps", [1e-6, 1e-5])
29 +def test_rms_norm_batch_invariant_vs_standard(
30 + batch_size: int, hidden_size: int, dtype: torch.dtype, eps: float
31 +):
32 + """
33 + Compare batch-invariant Triton RMS norm against standard CUDA implementation.
34 +
35 + Tests that the Triton-based batch-invariant RMS norm produces numerically
36 + equivalent results to the standard CUDA implementation across various
37 + configurations.
38 + """
39 + device = torch.device("cuda")
40 +
41 + # Create test input and weight
42 + torch.manual_seed(42)
43 + input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
44 + weight = torch.randn(hidden_size, dtype=dtype, device=device)
45 +
46 + # Standard implementation (CUDA ops)
47 + rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
48 + rms_norm_layer.weight.data = weight.clone()
49 +
50 + standard_output = rms_norm_layer.forward_cuda(input_tensor)
51 +
52 + # Batch-invariant implementation (Triton)
53 + triton_output = triton_rms_norm(input_tensor, weight, eps=eps)
54 +
55 + # Compare outputs
56 + # Use looser tolerance for bfloat16 due to its lower precision
57 + if dtype == torch.bfloat16:
58 + rtol, atol = 1e-1, 1e-1 # 10% relative tolerance for bfloat16
59 + else:
60 + rtol, atol = 1e-2, 1e-2 # 1% for float16/float32
61 +
62 + torch.testing.assert_close(
63 + triton_output,
64 + standard_output,
65 + rtol=rtol,
66 + atol=atol,
67 + msg=f"RMS norm mismatch for batch_size={batch_size}, "
68 + f"hidden_size={hidden_size}, "
69 + f"dtype={dtype}, eps={eps}",
70 + )
71 +
72 +
73 +@pytest.mark.skipif(
74 + not current_platform.has_device_capability(90),
75 + reason="Batch invariance tests only supported on Hopper (SM90)",
76 +)
77 +@pytest.mark.skipif(
78 + not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
79 +)
80 +@pytest.mark.parametrize("batch_size", [1, 16, 128])
81 +@pytest.mark.parametrize("seq_len", [1, 32, 512])
82 +@pytest.mark.parametrize("hidden_size", [2048, 4096])
83 +def test_rms_norm_3d_input(batch_size: int, seq_len: int, hidden_size: int):
84 + """
85 + Test RMS norm with 3D input tensors (batch, seq_len, hidden_size).
86 +
87 + Ensures that the batch-invariant RMS norm correctly handles multi-dimensional
88 + inputs that are common in transformer models.
89 + """
90 + device = torch.device("cuda")
91 + dtype = torch.bfloat16
92 + eps = 1e-6
93 +
94 + torch.manual_seed(42)
95 + input_tensor = torch.randn(
96 + batch_size, seq_len, hidden_size, dtype=dtype, device=device
97 + )
98 + weight = torch.randn(hidden_size, dtype=dtype, device=device)
99 +
100 + # Standard implementation
101 + rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
102 + rms_norm_layer.weight.data = weight.clone()
103 + standard_output = rms_norm_layer.forward_cuda(input_tensor)
104 +
105 + # Batch-invariant implementation
106 + triton_output = triton_rms_norm(input_tensor, weight, eps=eps)
107 +
108 + # Use looser tolerance for bfloat16
109 + rtol, atol = 1e-1, 1e-1 # 10% tolerance for bfloat16
110 +
111 + torch.testing.assert_close(
112 + triton_output,
113 + standard_output,
114 + rtol=rtol,
115 + atol=atol,
116 + msg=f"RMS norm mismatch for 3D input with batch_size={batch_size}, "
117 + f"seq_len={seq_len}, hidden_size={hidden_size}",
118 + )
119 +
120 +
121 +@pytest.mark.skipif(
122 + not current_platform.has_device_capability(90),
123 + reason="Batch invariance tests only supported on Hopper (SM90)",
124 +)
125 +@pytest.mark.skipif(
126 + not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
127 +)
128 +def test_rms_norm_numerical_stability():
129 + """
130 + Test RMS norm numerical stability with extreme values.
131 +
132 + Ensures that both implementations handle edge cases like very small or large
133 + values without producing NaN or Inf.
134 + """
135 + device = torch.device("cuda")
136 + dtype = torch.float16
137 + eps = 1e-6
138 + hidden_size = 2048
139 +
140 + # Test cases with extreme values
141 + test_cases = [
142 + # Very small values
143 + torch.ones(4, hidden_size, dtype=dtype, device=device) * 1e-5,
144 + # Very large values
145 + torch.ones(4, hidden_size, dtype=dtype, device=device) * 1e4,
146 + # Mixed small and large
147 + torch.randn(4, hidden_size, dtype=dtype, device=device) * 100,
148 + # Values near zero
149 + torch.randn(4, hidden_size, dtype=dtype, device=device) * 1e-6,
150 + ]
151 +
152 + weight = torch.ones(hidden_size, dtype=dtype, device=device)
153 +
154 + for idx, input_tensor in enumerate(test_cases):
155 + # Standard implementation
156 + rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
157 + rms_norm_layer.weight.data = weight.clone()
158 + standard_output = rms_norm_layer.forward_cuda(input_tensor)
159 +
160 + # Batch-invariant implementation
161 + triton_output = triton_rms_norm(input_tensor, weight, eps=eps)
162 +
163 + # Check for NaN or Inf
164 + assert not torch.isnan(standard_output).any(), (
165 + f"Standard RMS norm produced NaN for test case {idx}"
166 + )
167 + assert not torch.isinf(standard_output).any(), (
168 + f"Standard RMS norm produced Inf for test case {idx}"
169 + )
170 + assert not torch.isnan(triton_output).any(), (
171 + f"Triton RMS norm produced NaN for test case {idx}"
172 + )
173 + assert not torch.isinf(triton_output).any(), (
174 + f"Triton RMS norm produced Inf for test case {idx}"
175 + )
176 +
177 + # Compare outputs - very lenient for extreme values with float16
178 + torch.testing.assert_close(
179 + triton_output,
180 + standard_output,
181 + rtol=2e-1, # 20% tolerance for extreme values
182 + atol=2e-1,
183 + msg=f"RMS norm mismatch for extreme value test case {idx}",
184 + )
185 +
186 +
187 +@pytest.mark.skipif(
188 + not current_platform.has_device_capability(90),
189 + reason="Batch invariance tests only supported on Hopper (SM90)",
190 +)
191 +@pytest.mark.skipif(
192 + not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
193 +)
194 +def test_rms_norm_formula():
195 + """
196 + Test that RMS norm follows the correct mathematical formula.
197 +
198 + Verifies: output = input / sqrt(mean(input^2) + eps) * weight
199 + """
200 + device = torch.device("cuda")
201 + dtype = torch.float32 # Use float32 for higher precision in formula check
202 + eps = 1e-6
203 + hidden_size = 1024
204 +
205 + torch.manual_seed(42)
206 + input_tensor = torch.randn(8, hidden_size, dtype=dtype, device=device)
207 + weight = torch.randn(hidden_size, dtype=dtype, device=device)
208 +
209 + # Compute expected output using the formula
210 + variance = (input_tensor.pow(2).mean(dim=-1, keepdim=True)).to(dtype)
211 + expected_output = input_tensor * torch.rsqrt(variance + eps) * weight
212 +
213 + # Batch-invariant implementation
214 + triton_output = triton_rms_norm(input_tensor, weight, eps=eps)
215 +
216 + # Compare against formula
217 + torch.testing.assert_close(
218 + triton_output,
219 + expected_output,
220 + rtol=1e-4,
221 + atol=1e-4,
222 + msg="Triton RMS norm doesn't match expected formula",
223 + )
224 +
225 +
226 +@pytest.mark.skipif(
227 + not current_platform.has_device_capability(90),
228 + reason="Batch invariance tests only supported on Hopper (SM90)",
229 +)
230 +@pytest.mark.skipif(
231 + not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
232 +)
233 +@pytest.mark.parametrize("hidden_size", [128, 1024, 4096, 16384])
234 +def test_rms_norm_different_hidden_sizes(hidden_size: int):
235 + """
236 + Test RMS norm with various hidden sizes to ensure block size handling.
237 +
238 + The Triton kernel uses a fixed BLOCK_SIZE=1024, so this tests that it
239 + correctly handles hidden sizes both smaller and larger than the block size.
240 + """
241 + device = torch.device("cuda")
242 + dtype = torch.bfloat16
243 + eps = 1e-6
244 + batch_size = 16
245 +
246 + torch.manual_seed(42)
247 + input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
248 + weight = torch.randn(hidden_size, dtype=dtype, device=device)
249 +
250 + # Standard implementation
251 + rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
252 + rms_norm_layer.weight.data = weight.clone()
253 + standard_output = rms_norm_layer.forward_cuda(input_tensor)
254 +
255 + # Batch-invariant implementation
256 + triton_output = triton_rms_norm(input_tensor, weight, eps=eps)
257 +
258 + # Use looser tolerance for bfloat16
259 + rtol, atol = 1e-1, 1e-1 # 10% tolerance for bfloat16
260 +
261 + torch.testing.assert_close(
262 + triton_output,
263 + standard_output,
264 + rtol=rtol,
265 + atol=atol,
266 + msg=f"RMS norm mismatch for hidden_size={hidden_size}",
267 + )
268 +
269 +
270 +@pytest.mark.skipif(
271 + not current_platform.has_device_capability(90),
272 + reason="Batch invariance tests only supported on Hopper (SM90)",
273 +)
274 +@pytest.mark.skipif(
275 + not torch.cuda.is_available(), reason="Requires CUDA for RMS norm kernels"
276 +)
277 +def test_rms_norm_determinism():
278 + """
279 + Test that batch-invariant RMS norm produces deterministic results.
280 +
281 + Runs the same input through the kernel multiple times and verifies
282 + identical outputs.
283 + """
284 + device = torch.device("cuda")
285 + dtype = torch.bfloat16
286 + eps = 1e-6
287 + hidden_size = 4096
288 + batch_size = 32
289 +
290 + torch.manual_seed(42)
291 + input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
292 + weight = torch.randn(hidden_size, dtype=dtype, device=device)
293 +
294 + # Run multiple times
295 + outputs = []
296 + for _ in range(5):
297 + output = triton_rms_norm(input_tensor.clone(), weight, eps=eps)
298 + outputs.append(output)
299 +
300 + # All outputs should be identical
301 + reference = outputs[0]
302 + for idx, output in enumerate(outputs[1:], start=1):
303 + torch.testing.assert_close(
304 + output,
305 + reference,
306 + rtol=0.0,
307 + atol=0.0,
308 + msg=f"RMS norm not deterministic: run {idx} differs from reference",
309 + )
310 +
311 +
312 +if __name__ == "__main__":
313 + # Run a quick smoke test
314 + print("Running quick smoke test of RMS norm implementations...")
315 +
316 + device = torch.device("cuda")
317 + batch_size = 8
318 + hidden_size = 4096
319 + dtype = torch.bfloat16
320 + eps = 1e-6
321 +
322 + torch.manual_seed(42)
323 + input_tensor = torch.randn(batch_size, hidden_size, dtype=dtype, device=device)
324 + weight = torch.randn(hidden_size, dtype=dtype, device=device)
325 +
326 + # Standard implementation
327 + rms_norm_layer = RMSNorm(hidden_size, eps=eps, dtype=dtype).to(device)
328 + rms_norm_layer.weight.data = weight.clone()
329 + standard_output = rms_norm_layer.forward_cuda(input_tensor)
330 +
331 + # Batch-invariant implementation
332 + triton_output = triton_rms_norm(input_tensor, weight, eps=eps)
333 +
334 + # Compare
335 + max_diff = (triton_output - standard_output).abs().max().item()
336 + mean_diff = (triton_output - standard_output).abs().mean().item()
337 +
338 + print(f"Max difference: {max_diff:.6e}")
339 + print(f"Mean difference: {mean_diff:.6e}")
340 + print(f"Standard output sample: {standard_output[0, :5].tolist()}")
341 + print(f"Triton output sample: {triton_output[0, :5].tolist()}")
342 +
343 + if max_diff < 1e-3:
344 + print("✓ Smoke test passed!")
345 + else:
346 + print("✗ Smoke test failed - differences too large")
vllm/compilation/caching.py
+3 -1
@@ -3,6 +3,7 @@
3
4 import hashlib
5 import inspect
6 +import os
7 import pickle
8 from unittest.mock import patch
9
@@ -168,7 +169,8 @@ def _compute_code_hash(files: set[str]) -> str:
169 )
170 file_contents = {}
171 for filepath in files:
171 - if filepath == "<string>":
172 + # Skip files that don't exist (e.g., <string>, <frozen modules>, etc.)
173 + if not os.path.isfile(filepath):
174 file_contents[filepath] = ""
175 else:
176 with open(filepath) as f:
vllm/config/model.py
+7
@@ -20,6 +20,9 @@ from vllm.config.pooler import PoolerConfig
20 from vllm.config.scheduler import RunnerType
21 from vllm.config.utils import assert_hashable, config, getattr_iter
22 from vllm.logger import init_logger
23 +from vllm.model_executor.layers.batch_invariant import (
24 + vllm_kernel_override_batch_invariant,
25 +)
26 from vllm.platforms import current_platform
27 from vllm.transformers_utils.config import (
28 ConfigFormat,
@@ -419,6 +422,10 @@ class ModelConfig:
422 skip_mm_profiling: bool | None,
423 video_pruning_rate: float | None,
424 ) -> None:
425 + # Enable batch invariance settings if requested
426 + if vllm_kernel_override_batch_invariant():
427 + self.enforce_eager = True
428 +
429 # Set the default seed to 0 in V1.
430 # NOTE(woosuk): In V0, we set the default seed to None because the
431 # driver worker shares the same process as the user process, and thus
vllm/config/parallel.py
+7 -1
@@ -14,6 +14,9 @@ from typing_extensions import Self
14 import vllm.envs as envs
15 from vllm.config.utils import config
16 from vllm.logger import init_logger
17 +from vllm.model_executor.layers.batch_invariant import (
18 + vllm_kernel_override_batch_invariant,
19 +)
20 from vllm.platforms import current_platform
21 from vllm.utils import cuda_device_count_stateless, get_open_ports_list
22
@@ -560,7 +563,10 @@ class ParallelConfig:
563 def _verify_args(self) -> Self:
564 # Lazy import to avoid circular import
565 from vllm.executor.executor_base import ExecutorBase
563 - from vllm.platforms import current_platform
566 +
567 + # Enable batch invariance settings if requested
568 + if vllm_kernel_override_batch_invariant():
569 + self.disable_custom_all_reduce = True
570
571 if (
572 self.distributed_executor_backend is not None
vllm/distributed/device_communicators/all_reduce_utils.py
+6
@@ -19,6 +19,9 @@ import torch.multiprocessing as mp
19 import vllm.envs as envs
20 from vllm.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
21 from vllm.logger import init_logger
22 +from vllm.model_executor.layers.batch_invariant import (
23 + vllm_kernel_override_batch_invariant,
24 +)
25 from vllm.utils import cuda_device_count_stateless, update_environment_variables
26
27 logger = init_logger(__name__)
@@ -71,6 +74,9 @@ def should_nccl_symm_mem_allreduce(world_size: int, input_tensor: torch.Tensor)
74 is_symmetric_memory_enabled,
75 )
76
77 + if vllm_kernel_override_batch_invariant():
78 + return False
79 +
80 if not is_symmetric_memory_enabled():
81 return False
82 if world_size < NCCL_SYMM_MEM_ALL_REDUCE_CONFIG["min_world_size"]:
vllm/distributed/device_communicators/symm_mem.py
+5
@@ -9,6 +9,9 @@ from vllm.distributed.device_communicators.all_reduce_utils import (
9 SYMM_MEM_ALL_REDUCE_MAX_SIZES,
10 )
11 from vllm.logger import init_logger
12 +from vllm.model_executor.layers.batch_invariant import (
13 + vllm_kernel_override_batch_invariant,
14 +)
15 from vllm.platforms import current_platform
16
17 try:
@@ -100,6 +103,8 @@ class SymmMemCommunicator:
103 return
104 self.force_multimem = force_multimem
105 self.disabled = False
106 + if vllm_kernel_override_batch_invariant():
107 + self.disabled = True
108
109 def should_use_symm_mem(self, inp: torch.Tensor):
110 if self.disabled:
vllm/engine/arg_utils.py
+1 -1
@@ -1694,7 +1694,7 @@ class EngineArgs:
1694 ) -> None:
1695 """Set Default Arguments for V1 Engine."""
1696
1697 - # V1 always uses chunked prefills and prefix caching
1697 + # V1 uses chunked prefills and prefix caching by default
1698 # for non-pooling tasks.
1699 # For pooling tasks the default is False
1700 if model_config.runner_type != "pooling":
vllm/model_executor/layers/batch_invariant.py
+240 -13
@@ -395,7 +395,6 @@ def mean_dim(
395 Tensor with mean values along specified dimension
396 """
397 # Validate inputs
398 - assert input.is_cuda, "Input must be a CUDA tensor"
398 assert -input.ndim <= dim < input.ndim, (
399 f"Invalid dimension {dim} for tensor with {input.ndim} dimensions"
400 )
@@ -470,6 +469,45 @@ def mm_batch_invariant(a, b):
469 return matmul_persistent(a, b)
470
471
472 +def matmul_batch_invariant(a, b, *, out=None):
473 + # torch.matmul can handle various dimensions
474 + # For 2D x 2D, it's the same as mm
475 + if a.ndim == 2 and b.ndim == 2:
476 + result = matmul_persistent(a, b)
477 + if out is not None:
478 + out.copy_(result)
479 + return out
480 + return result
481 + elif a.ndim == 3 and b.ndim == 3:
482 + # Handle batched case like bmm
483 + return bmm_batch_invariant(a, b, out=out)
484 + else:
485 + raise ValueError(
486 + f"matmul_batch_invariant currently only supports 2D x 2D and 3D x 3D, "
487 + f"got shapes {a.shape} and {b.shape}"
488 + )
489 +
490 +
491 +def bmm_batch_invariant(a, b, *, out=None):
492 + # Batched matrix multiply: (B, M, K) x (B, K, N) -> (B, M, N)
493 + # Process each batch separately with our persistent kernel
494 + if a.ndim == 3 and b.ndim == 3:
495 + results = []
496 + for i in range(a.shape[0]):
497 + results.append(matmul_persistent(a[i], b[i]))
498 + result = torch.stack(results, dim=0)
499 +
500 + if out is not None:
501 + out.copy_(result)
502 + return out
503 + return result
504 + else:
505 + raise ValueError(
506 + f"bmm_batch_invariant expects 3D tensors, "
507 + f"got shapes {a.shape} and {b.shape}"
508 + )
509 +
510 +
511 def addmm_batch_invariant(bias, a, b):
512 return matmul_persistent(a, b, bias=bias)
513
@@ -479,11 +517,24 @@ def _log_softmax_batch_invariant(input, dim, _half_to_float):
517 return log_softmax(input, dim=dim)
518
519
520 +def softmax_batch_invariant(input, dim, dtype=None):
521 + # Compute softmax in a deterministic way
522 + # First subtract max for numerical stability (standard practice)
523 + input_max = torch.amax(input, dim=dim, keepdim=True)
524 + input = input - input_max
525 + exp_x = torch.exp(input)
526 + sum_exp_x = torch.sum(exp_x, dim=dim, keepdim=True)
527 + return exp_x / sum_exp_x
528 +
529 +
530 def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None = None):
531 assert dtype is None or dtype == torch.float32, f"unsupported dtype: {dtype}"
532
533 result = input.to(torch.float32)
534
535 + if len(dim) == 0:
536 + dim = [i for i in range(len(input.shape))]
537 +
538 # Sort dimensions to reduce from largest to smallest to handle shifting dims
539 # during iterative reduction.
540 sorted_dims = sorted([d % input.ndim for d in dim], reverse=True)
@@ -500,8 +551,134 @@ def mean_batch_invariant(input, dim, keepdim=False, dtype: torch.dtype | None =
551 return result
552
553
554 +@triton.jit
555 +def _rms_norm_kernel(
556 + input_ptr,
557 + weight_ptr,
558 + output_ptr,
559 + input_row_stride,
560 + output_row_stride,
561 + n_cols,
562 + eps,
563 + BLOCK_SIZE: tl.constexpr,
564 +):
565 + """
566 + Compute RMS normalization along the last dimension of a 2D tensor.
567 + RMS Norm: y = x / sqrt(mean(x^2) + eps) * weight
568 + Each block handles one row of the input tensor.
569 + """
570 + row_idx = tl.program_id(0).to(tl.int64)
571 + row_start_ptr = input_ptr + row_idx * input_row_stride
572 + output_row_start_ptr = output_ptr + row_idx * output_row_stride
573 +
574 + # Step 1: Compute sum of squares in float32 to avoid overflow
575 + sum_sq = tl.zeros([1], dtype=tl.float32)
576 + for col_offset in range(0, n_cols, BLOCK_SIZE):
577 + col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
578 + mask = col_idx < n_cols
579 +
580 + vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0)
581 + # Convert to float32 for accumulation to prevent overflow
582 + vals_f32 = vals.to(tl.float32)
583 + sq_vals = vals_f32 * vals_f32
584 + sum_sq += tl.sum(tl.where(mask, sq_vals, 0.0))
585 +
586 + # Step 2: Compute RMS (root mean square) in float32
587 + mean_sq = sum_sq / n_cols
588 + rms = tl.sqrt(mean_sq + eps)
589 + inv_rms = 1.0 / rms
590 +
591 + # Step 3: Normalize and apply weight
592 + for col_offset in range(0, n_cols, BLOCK_SIZE):
593 + col_idx = col_offset + tl.arange(0, BLOCK_SIZE)
594 + mask = col_idx < n_cols
595 + vals = tl.load(row_start_ptr + col_idx, mask=mask, other=0.0)
596 + weight = tl.load(weight_ptr + col_idx, mask=mask, other=1.0)
597 + # Compute in float32 then convert back to input dtype
598 + vals_f32 = vals.to(tl.float32)
599 + weight_f32 = weight.to(tl.float32)
600 + output_f32 = vals_f32 * inv_rms * weight_f32
601 + output = output_f32.to(vals.dtype)
602 + tl.store(output_row_start_ptr + col_idx, output, mask=mask)
603 +
604 +
605 +def rms_norm(
606 + input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
607 +) -> torch.Tensor:
608 + """
609 + Compute RMS normalization using Triton kernel.
610 +
611 + RMS Norm normalizes the input by the root mean square and scales by weight:
612 + output = input / sqrt(mean(input^2) + eps) * weight
613 +
614 + Args:
615 + input: Input tensor of shape (..., hidden_size)
616 + weight: Weight tensor of shape (hidden_size,)
617 + eps: Small constant for numerical stability
618 +
619 + Returns:
620 + Tensor with RMS normalization applied along the last dimension
621 + """
622 + assert weight.dim() == 1, "Weight must be 1-dimensional"
623 + assert input.shape[-1] == weight.shape[0], (
624 + f"Input last dimension ({input.shape[-1]}) must match "
625 + f"weight dimension ({weight.shape[0]})"
626 + )
627 +
628 + # Flatten all dimensions except the last one
629 + original_shape = input.shape
630 + input_2d = input.reshape(-1, input.shape[-1])
631 + input_2d = input_2d.contiguous()
632 + weight = weight.contiguous()
633 +
634 + n_rows, n_cols = input_2d.shape
635 +
636 + output = torch.empty_like(input_2d)
637 + BLOCK_SIZE = 1024
638 + grid = (n_rows,)
639 + _rms_norm_kernel[grid](
640 + input_2d,
641 + weight,
642 + output,
643 + input_2d.stride(0),
644 + output.stride(0),
645 + n_cols,
646 + eps,
647 + BLOCK_SIZE=BLOCK_SIZE,
648 + )
649 + return output.reshape(original_shape)
650 +
651 +
652 +def rms_norm_batch_invariant(
653 + input: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6
654 +) -> torch.Tensor:
655 + """
656 + Batch-invariant wrapper for RMS normalization.
657 +
658 + This function provides a deterministic, batch-invariant implementation
659 + of RMS normalization for use with the batch_invariant mode.
660 +
661 + Args:
662 + input: Input tensor of shape (..., hidden_size)
663 + weight: Weight tensor of shape (hidden_size,)
664 + eps: Small constant for numerical stability
665 +
666 + Returns:
667 + RMS normalized tensor
668 + """
669 + return rms_norm(input, weight, eps=eps)
670 +
671 +
672 +def linear_batch_invariant(input, weight, bias=None):
673 + output = mm_batch_invariant(input, weight.t())
674 + if bias is not None:
675 + output = output + bias
676 + return output
677 +
678 +
679 _batch_invariant_MODE = False
680 _batch_invariant_LIB = None
681 +_original_torch_bmm = None
682
683
684 def is_batch_invariant_mode_enabled():
@@ -509,7 +686,7 @@ def is_batch_invariant_mode_enabled():
686
687
688 def enable_batch_invariant_mode():
512 - global _batch_invariant_MODE, _batch_invariant_LIB
689 + global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm
690 if _batch_invariant_MODE:
691 return
692
@@ -517,16 +694,28 @@ def enable_batch_invariant_mode():
694 _batch_invariant_LIB = torch.library.Library("aten", "IMPL")
695 _batch_invariant_LIB.impl("aten::mm", mm_batch_invariant, "CUDA")
696 _batch_invariant_LIB.impl("aten::addmm", addmm_batch_invariant, "CUDA")
697 + _batch_invariant_LIB.impl("aten::matmul", matmul_batch_invariant, "CUDA")
698 + _batch_invariant_LIB.impl("aten::bmm", bmm_batch_invariant, "CUDA")
699 + _batch_invariant_LIB.impl("aten::linear", linear_batch_invariant, "CUDA")
700 _batch_invariant_LIB.impl(
701 "aten::_log_softmax", _log_softmax_batch_invariant, "CUDA"
702 )
703 + _batch_invariant_LIB.impl("aten::softmax", softmax_batch_invariant, "CUDA")
704 + _batch_invariant_LIB.impl("aten::_softmax", softmax_batch_invariant, "CUDA")
705 _batch_invariant_LIB.impl("aten::mean.dim", mean_batch_invariant, "CUDA")
706
707 + # Also monkeypatch torch.bmm directly as a fallback
708 + _original_torch_bmm = torch.bmm
709 + torch.bmm = bmm_batch_invariant
710 +
711
712 def disable_batch_invariant_mode():
527 - global _batch_invariant_MODE, _batch_invariant_LIB
713 + global _batch_invariant_MODE, _batch_invariant_LIB, _original_torch_bmm
714 if _batch_invariant_LIB is not None:
715 _batch_invariant_LIB._destroy()
716 + if _original_torch_bmm is not None:
717 + torch.bmm = _original_torch_bmm
718 + _original_torch_bmm = None
719 _batch_invariant_MODE = False
720 _batch_invariant_LIB = None
721
@@ -563,17 +752,55 @@ def vllm_kernel_override_batch_invariant():
752 return is_overridden
753
754
755 +def override_envs_for_invariance():
756 + curr_attn_backend = envs.VLLM_ATTENTION_BACKEND
757 + supported_backends = [
758 + "FLASH_ATTN", # best supported backend
759 + "FLEX_ATTENTION",
760 + "FLASHINFER",
761 + "FLASH_ATTN_MLA",
762 + "TRITON_MLA",
763 + # Not yet supported MLA backends
764 + # "FLASHMLA",
765 + # "FLASHINFER_MLA",
766 + ]
767 + if curr_attn_backend not in supported_backends:
768 + warning = (
769 + "Forcibly updating attention backend to"
770 + f" {supported_backends[0]} for batch_invariant. "
771 + f" Supported backends: {supported_backends}."
772 + )
773 + logger.warning_once(warning)
774 + os.environ["VLLM_ATTENTION_BACKEND"] = supported_backends[0]
775 + if os.environ["VLLM_ATTENTION_BACKEND"] != supported_backends[0]:
776 + warning = (
777 + "You are using a decode-invariant form of batch invariance. "
778 + "This will not be invariant between prefill and decode."
779 + )
780 + logger.warning_once(warning)
781 + os.environ["VLLM_ALLREDUCE_USE_SYMM_MEM"] = "0"
782 +
783 + os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
784 +
785 + # NCCL determinism settings
786 + os.environ["NCCL_LAUNCH_MODE"] = "GROUP"
787 + os.environ["NCCL_COLLNET_ENABLE"] = "0"
788 + os.environ["NCCL_NVLS_ENABLE"] = "0"
789 + os.environ["NCCL_P2P_NET_DISABLE"] = "1"
790 + os.environ["NCCL_MIN_NCHANNELS"] = "1"
791 + os.environ["NCCL_MAX_NCHANNELS"] = "1"
792 + os.environ["NCCL_PROTO"] = "Simple"
793 + os.environ["NCCL_ALGO"] = "allreduce:tree"
794 + os.environ["NCCL_NTHREADS"] = "1"
795 + os.environ["NCCL_SOCKET_NTHREADS"] = "1"
796 +
797 +
798 def init_batch_invariance():
799 # this will hit all the csrc overrides as well
800 if vllm_kernel_override_batch_invariant():
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]
801 + override_envs_for_invariance()
802 enable_batch_invariant_mode()
803 +
804 + # Disable TF32 for batch invariance - it causes non-deterministic rounding
805 + torch.backends.cuda.matmul.allow_tf32 = False
806 + torch.backends.cudnn.allow_tf32 = False
vllm/model_executor/layers/fused_moe/fused_moe.py
+28 -4
@@ -15,6 +15,9 @@ import vllm.envs as envs
15 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
16 from vllm import _custom_ops as ops
17 from vllm.logger import init_logger
18 +from vllm.model_executor.layers.batch_invariant import (
19 + vllm_kernel_override_batch_invariant,
20 +)
21 from vllm.model_executor.layers.fused_moe.config import (
22 FUSED_MOE_UNQUANTIZED_CONFIG,
23 FusedMoEQuantConfig,
@@ -837,6 +840,10 @@ def get_moe_configs(
840 be picked and the associated configuration chosen to invoke the kernel.
841 """
842
843 + # Avoid optimizing for the batch invariant case. Use default config
844 + if vllm_kernel_override_batch_invariant():
845 + return None
846 +
847 # First look up if an optimized configuration is available in the configs
848 # directory
849 block_shape = [block_n, block_k] if block_n and block_k else None
@@ -969,6 +976,15 @@ def get_default_config(
976 dtype: str | None,
977 block_shape: list[int] | None = None,
978 ) -> dict[str, int]:
979 + if vllm_kernel_override_batch_invariant():
980 + config = {
981 + "BLOCK_SIZE_M": 64,
982 + "BLOCK_SIZE_N": 64,
983 + "BLOCK_SIZE_K": 32,
984 + "GROUP_SIZE_M": 8,
985 + }
986 + return config
987 +
988 if dtype == "fp8_w8a8" and block_shape is not None:
989 # Block-wise quant: BLOCK_SIZE_N must be divisible by block_shape[0]
990 # BLOCK_SIZE_K must be divisible by block_shape[1]
@@ -1118,7 +1134,10 @@ def fused_topk_bias(
1134 scores_for_choice = scores.view(
1135 -1, n_routed_experts
1136 ) + e_score_correction_bias.unsqueeze(0)
1121 - topk_indices = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=False)[1]
1137 +
1138 + # For batch invariance, use sorted=True to ensure deterministic expert selection
1139 + use_sorted = vllm_kernel_override_batch_invariant()
1140 + topk_indices = torch.topk(scores_for_choice, k=topk, dim=-1, sorted=use_sorted)[1]
1141 topk_weights = scores.gather(1, topk_indices)
1142 if renormalize:
1143 topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
@@ -1179,7 +1198,10 @@ def grouped_topk(
1198 group_scores = (
1199 scores.view(num_token, num_expert_group, -1).max(dim=-1).values
1200 ) # [n, n_group]
1182 - group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=False)[
1201 +
1202 + # For batch invariance, use sorted=True to ensure deterministic expert selection
1203 + use_sorted = vllm_kernel_override_batch_invariant()
1204 + group_idx = torch.topk(group_scores, k=topk_group, dim=-1, sorted=use_sorted)[
1205 1
1206 ] # [n, top_k_group]
1207 group_mask = torch.zeros_like(group_scores) # [n, n_group]
@@ -1192,11 +1214,13 @@ def grouped_topk(
1214 tmp_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) # [n, e]
1215
1216 if e_score_correction_bias is not None:
1195 - topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)[1]
1217 + topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=use_sorted)[1]
1218 # Use original unbiased scores for the routing weights
1219 topk_weights = original_scores.gather(1, topk_ids)
1220 else:
1199 - topk_weights, topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)
1221 + topk_weights, topk_ids = torch.topk(
1222 + tmp_scores, k=topk, dim=-1, sorted=use_sorted
1223 + )
1224
1225 if renormalize:
1226 topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)
vllm/model_executor/layers/layernorm.py
+10
@@ -8,6 +8,10 @@ import torch.nn.functional as F
8
9 import vllm.envs as envs
10 from vllm.model_executor.custom_op import CustomOp
11 +from vllm.model_executor.layers.batch_invariant import (
12 + rms_norm_batch_invariant,
13 + vllm_kernel_override_batch_invariant,
14 +)
15 from vllm.platforms import current_platform
16 from vllm.utils import direct_register_custom_op
17
@@ -21,6 +25,8 @@ def rms_norm(
25 ) -> torch.Tensor:
26 from vllm import _custom_ops as ops
27
28 + if vllm_kernel_override_batch_invariant():
29 + return rms_norm_batch_invariant(x, weight, variance_epsilon)
30 out = torch.empty_like(x)
31 ops.rms_norm(
32 out,
@@ -39,6 +45,10 @@ def fused_add_rms_norm(
45 ) -> tuple[torch.Tensor, torch.Tensor]:
46 from vllm import _custom_ops as ops
47
48 + if vllm_kernel_override_batch_invariant():
49 + return rms_norm_batch_invariant(
50 + x + residual, weight, variance_epsilon
51 + ), x + residual
52 ops.fused_add_rms_norm(
53 x,
54 residual,
vllm/model_executor/layers/mla.py
+1
@@ -160,6 +160,7 @@ class MultiHeadLatentAttentionWrapper(CustomOp):
160 k_pe,
161 output_shape=(hidden_states.shape[0], self.num_heads * self.v_head_dim),
162 )
163 +
164 return self.o_proj(attn_out)[0]
165
166 def forward_cuda(self, *args, **kwargs):
vllm/model_executor/layers/quantization/fp8.py
+65
@@ -14,6 +14,9 @@ import vllm.model_executor.layers.fused_moe.modular_kernel as mk
14 from vllm import _custom_ops as ops
15 from vllm.distributed import get_tensor_model_parallel_world_size
16 from vllm.logger import init_logger
17 +from vllm.model_executor.layers.batch_invariant import (
18 + vllm_kernel_override_batch_invariant,
19 +)
20 from vllm.model_executor.layers.fused_moe import (
21 FusedMoE,
22 FusedMoEActivationFormat,
@@ -353,6 +356,8 @@ class Fp8LinearMethod(LinearMethodBase):
356 # Disable marlin for rocm
357 if current_platform.is_rocm():
358 self.use_marlin = False
359 + if vllm_kernel_override_batch_invariant():
360 + self.use_marlin = False
361
362 self.use_aiter_and_is_supported = check_aiter_fp8_linear_support()
363
@@ -534,6 +539,66 @@ class Fp8LinearMethod(LinearMethodBase):
539 x: torch.Tensor,
540 bias: torch.Tensor | None = None,
541 ) -> torch.Tensor:
542 + # If batch invariant mode is enabled, dequantize and use BF16 compute
543 + if vllm_kernel_override_batch_invariant():
544 + # Dequantize FP8 weights to BF16
545 + weight_fp8 = layer.weight.to(torch.bfloat16)
546 + weight_scale = layer.weight_scale.to(torch.bfloat16)
547 +
548 + # Handle different quantization granularities
549 + if self.block_quant:
550 + # Block-wise quantization:
551 + # - Weight is NOT transposed, shape is [N, K] (output_size, input_size)
552 + # - Scale has shape [num_blocks_k, num_blocks_n] (TRANSPOSED!)
553 + assert self.weight_block_size is not None
554 + block_n, block_k = self.weight_block_size # Note: order is [N, K]
555 +
556 + N, K = weight_fp8.shape
557 +
558 + # Scale is stored transposed: [num_blocks_k, num_blocks_n]
559 + # We need to transpose it to [num_blocks_n, num_blocks_k] first
560 + weight_scale = weight_scale.t()
561 +
562 + # Expand scale to match weight dimensions
563 + # scale_expanded should have shape [N, K]
564 + scale_expanded = weight_scale.repeat_interleave(
565 + block_n, dim=0
566 + ).repeat_interleave(block_k, dim=1)
567 + # Trim to exact weight size (in case of padding)
568 + scale_expanded = scale_expanded[:N, :K]
569 + weight_bf16 = weight_fp8 * scale_expanded
570 + else:
571 + # Per-tensor quantization: weight IS transposed to [K, N]
572 + # scale should be scalar or [1] or per-output-channel [N]
573 + if weight_scale.numel() == 1:
574 + # Per-tensor: simple scalar multiplication
575 + weight_bf16 = weight_fp8 * weight_scale
576 + else:
577 + # Multiple scales (fused modules like QKV)
578 + # Try to infer correct broadcasting
579 + # weight is [K, N], scale could be [num_logical_weights]
580 + # Need to figure out how to broadcast - for now just try
581 + # direct multiplication
582 + if (
583 + weight_scale.dim() == 1
584 + and weight_scale.shape[0] == weight_fp8.shape[0]
585 + ):
586 + # Per-row scaling
587 + weight_bf16 = weight_fp8 * weight_scale.unsqueeze(1)
588 + else:
589 + # Fallback
590 + weight_bf16 = weight_fp8 * weight_scale
591 +
592 + # For block quant, weight is [N, K], for per-tensor it's [K, N]
593 + # F.linear expects weight to be [N, K], so:
594 + if self.block_quant:
595 + # Already in correct shape [N, K]
596 + output = torch.nn.functional.linear(x, weight_bf16, bias)
597 + else:
598 + # Need to transpose back: [K, N] -> [N, K]
599 + output = torch.nn.functional.linear(x, weight_bf16.t(), bias)
600 + return output
601 +
602 if self.use_marlin:
603 return apply_fp8_marlin_linear(
604 input=x,
vllm/model_executor/models/gpt_oss.py
+1
@@ -216,6 +216,7 @@ class TransformerBlock(torch.nn.Module):
216 else:
217 hidden_states, residual = self.input_layernorm(hidden_states, residual)
218 hidden_states = self.attn(hidden_states, positions)
219 +
220 # Fully Connected
221 hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
222 output = self.mlp(hidden_states)
vllm/v1/attention/backends/flash_attn.py
+12 -1
@@ -31,10 +31,12 @@ if is_flash_attn_varlen_func_available():
31 get_scheduler_metadata,
32 reshape_and_cache_flash,
33 )
34 -
34 from vllm.config import VllmConfig, get_layers_from_vllm_config
35 from vllm.distributed.parallel_state import get_dcp_group
36 from vllm.logger import init_logger
37 +from vllm.model_executor.layers.batch_invariant import (
38 + vllm_kernel_override_batch_invariant,
39 +)
40 from vllm.utils import cdiv
41 from vllm.v1.attention.backends.utils import (
42 AttentionCGSupport,
@@ -306,6 +308,9 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
308 # we only set num_splits when using cuda graphs.
309 max_num_splits = self.max_num_splits
310
311 + if vllm_kernel_override_batch_invariant():
312 + max_num_splits = 1
313 +
314 def schedule(
315 batch_size, cu_query_lens, max_query_len, seqlens, max_seq_len, causal
316 ):
@@ -478,6 +483,9 @@ class FlashAttentionImpl(AttentionImpl):
483
484 self.attn_type = attn_type
485 self.vllm_flash_attn_version = get_flash_attn_version()
486 + # Cache the batch invariant result for use in forward passes
487 + self.batch_invariant_enabled = vllm_kernel_override_batch_invariant()
488 +
489 if is_quantized_kv_cache(self.kv_cache_dtype) and not flash_attn_supports_fp8():
490 raise NotImplementedError(
491 "FlashAttention does not support fp8 kv-cache on this device."
@@ -810,6 +818,7 @@ class FlashAttentionImpl(AttentionImpl):
818 q_descale=layer._q_scale.expand(descale_shape),
819 k_descale=layer._k_scale.expand(descale_shape),
820 v_descale=layer._v_scale.expand(descale_shape),
821 + num_splits=1 if self.batch_invariant_enabled else 0,
822 )
823
824 return output
@@ -954,6 +963,7 @@ def cascade_attention(
963 # s_aux is incorporated into prefix_lse inside the GPU kernel,
964 # enabling its effect during the final attention merge.
965 s_aux=s_aux,
966 + num_splits=1 if vllm_kernel_override_batch_invariant() else 0,
967 )
968
969 descale_shape = (cu_query_lens.shape[0] - 1, key_cache.shape[-2])
@@ -978,6 +988,7 @@ def cascade_attention(
988 q_descale=q_descale.expand(descale_shape) if q_descale is not None else None,
989 k_descale=k_descale.expand(descale_shape) if k_descale is not None else None,
990 v_descale=v_descale.expand(descale_shape) if v_descale is not None else None,
991 + num_splits=1 if vllm_kernel_override_batch_invariant() else 0,
992 )
993
994 # Merge prefix and suffix outputs, and store the result in output.
vllm/v1/attention/backends/mla/common.py
+10 -1
@@ -211,6 +211,9 @@ from vllm.attention.utils.fa_utils import get_flash_attn_version
211 from vllm.config import VllmConfig, get_current_vllm_config
212 from vllm.distributed.parallel_state import get_dcp_group, is_global_first_rank
213 from vllm.logger import init_logger
214 +from vllm.model_executor.layers.batch_invariant import (
215 + vllm_kernel_override_batch_invariant,
216 +)
217 from vllm.model_executor.layers.linear import (
218 ColumnParallelLinear,
219 LinearBase,
@@ -1187,6 +1190,7 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
1190 def _v_up_proj(self, x: torch.Tensor, out: torch.Tensor):
1191 # Convert from (B, N, L) to (N, B, L)
1192 x = x.view(-1, self.num_heads, self.kv_lora_rank).transpose(0, 1)
1193 +
1194 if is_rocm_aiter_fp8bmm_enabled():
1195 # Multiply + Transpose (N, B, L) x (N, L, V)->(N, B, V)->(B, N, V)
1196 x = aiter_triton_fp8_bmm(
@@ -1279,6 +1283,8 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1283 # ROCm leverages the upstream flash_attn, which takes a parameter
1284 # called "return_attn_probs" instead of return_softmax_lse
1285 kwargs["return_attn_probs"] = return_softmax_lse
1286 + if vllm_kernel_override_batch_invariant():
1287 + kwargs["num_splits"] = 1
1288
1289 attn_out = self.flash_attn_varlen_func(
1290 q=q,
@@ -1841,9 +1847,11 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1847
1848 if has_decode:
1849 assert attn_metadata.decode is not None
1850 +
1851 decode_q_nope, decode_q_pe = decode_q.split(
1852 [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
1853 )
1854 +
1855 # Convert from (B, N, P) to (N, B, P)
1856 decode_q_nope = decode_q_nope.transpose(0, 1)
1857
@@ -1868,17 +1876,18 @@ class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):
1876 # Pads the head_dim if necessary (for the underlying kernel)
1877 N, B, P = decode_q_nope.shape
1878 _, _, L = self.W_UK_T.shape
1879 +
1880 if self.q_pad_num_heads is not None:
1881 decode_ql_nope = decode_q_nope.new_empty(
1882 (self.q_pad_num_heads, B, L)
1883 )
1884 decode_ql_nope.resize_((N, B, L))
1876 -
1885 else:
1886 decode_ql_nope = decode_q_nope.new_empty((N, B, L))
1887
1888 # Multiply (N, B, P) x (N, P, L) -> (N, B, L)
1889 torch.bmm(decode_q_nope, self.W_UK_T, out=decode_ql_nope)
1890 +
1891 # Convert from (N, B, L) to (B, N, L)
1892 decode_ql_nope = decode_ql_nope.transpose(0, 1)
1893
vllm/v1/attention/backends/mla/flashattn_mla.py
+11 -1
@@ -18,6 +18,9 @@ from vllm.attention.utils.fa_utils import (
18 )
19 from vllm.config import VllmConfig
20 from vllm.logger import init_logger
21 +from vllm.model_executor.layers.batch_invariant import (
22 + vllm_kernel_override_batch_invariant,
23 +)
24 from vllm.v1.attention.backends.mla.common import (
25 MLACommonBackend,
26 MLACommonDecodeMetadata,
@@ -107,6 +110,9 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
110 # pre-allocated during capture.
111 self.max_num_splits = envs.VLLM_FLASH_ATTN_MAX_NUM_SPLITS_FOR_CUDA_GRAPH
112
113 + if vllm_kernel_override_batch_invariant():
114 + self.max_num_splits = 1
115 +
116 def _schedule_decode(
117 self, num_reqs, cu_query_lens, max_query_len, seqlens, max_seq_len, causal
118 ):
@@ -175,7 +181,10 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
181 # we only set num_splits when using cuda graphs.
182 max_num_splits = self.max_num_splits
183
178 - return FlashAttnMLADecodeMetadata(
184 + if vllm_kernel_override_batch_invariant():
185 + max_num_splits = 1
186 +
187 + metadata = FlashAttnMLADecodeMetadata(
188 block_table=block_table_tensor,
189 seq_lens=seq_lens_device,
190 query_start_loc=query_start_loc_device,
@@ -185,6 +194,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]
194 max_num_splits=max_num_splits,
195 dcp_tot_seq_lens=dcp_tot_seq_lens_device,
196 )
197 + return metadata
198
199
200 class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):
vllm/v1/attention/backends/mla/flashmla.py
+36 -2
@@ -14,6 +14,9 @@ from vllm.attention.ops.flashmla import (
14 )
15 from vllm.config import VllmConfig
16 from vllm.logger import init_logger
17 +from vllm.model_executor.layers.batch_invariant import (
18 + vllm_kernel_override_batch_invariant,
19 +)
20 from vllm.v1.attention.backends.mla.common import (
21 MLACommonBackend,
22 MLACommonDecodeMetadata,
@@ -223,19 +226,50 @@ class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):
226 if type(q) is tuple:
227 q = torch.cat(q, dim=-1)
228
229 + # mypy assertion: q is now always a tensor
230 assert isinstance(q, torch.Tensor)
231
232 num_decodes = attn_metadata.num_decodes
233 q = reshape_query_for_spec_decode(q, num_decodes)
234
235 + tile_scheduler_metadata = attn_metadata.decode.tile_scheduler_metadata
236 + num_splits = attn_metadata.decode.num_splits
237 + if vllm_kernel_override_batch_invariant():
238 + device = q.device
239 + dtype = torch.int32
240 +
241 + B = q.shape[0]
242 + # block_table shape: [batch_size, max_num_blocks_per_seq]
243 + # The number of blocks per sequence is in the second dimension
244 + topk = attn_metadata.decode.block_table.shape[-1]
245 + B_TOPK = 64
246 + assert topk % B_TOPK == 0, f"topk ({topk}) must be divisible by {B_TOPK}"
247 + end_block_idx = topk // B_TOPK
248 +
249 + # Single partition => num_sm_parts = 1
250 + # TileSchedulerMetaDataSize = 8, layout:
251 + # [begin_idx, begin_block_idx, end_idx, end_block_idx,
252 + # begin_n_split_idx, _, _, _]
253 + tile_scheduler_metadata = torch.zeros((1, 8), dtype=dtype, device=device)
254 + tile_scheduler_metadata[0, 0] = 0 # begin_idx
255 + tile_scheduler_metadata[0, 1] = 0 # sched_begin_block_idx
256 + tile_scheduler_metadata[0, 2] = B - 1 # end_idx
257 + tile_scheduler_metadata[0, 3] = end_block_idx
258 + tile_scheduler_metadata[0, 4] = 0 # begin_n_split_idx
259 + # fields [5..7] stay 0
260 +
261 + # Non-split path ignores num_splits, but the API requires it:
262 + # zeros of length B+1
263 + num_splits = torch.zeros((B + 1,), dtype=dtype, device=device)
264 +
265 o, lse = flash_mla_with_kvcache(
266 q=q,
267 k_cache=kv_c_and_k_pe_cache.unsqueeze(-2), # Add head dim of 1
268 block_table=attn_metadata.decode.block_table,
269 cache_seqlens=attn_metadata.decode.seq_lens,
270 head_dim_v=self.kv_lora_rank,
237 - tile_scheduler_metadata=attn_metadata.decode.tile_scheduler_metadata,
238 - num_splits=attn_metadata.decode.num_splits,
271 + tile_scheduler_metadata=tile_scheduler_metadata,
272 + num_splits=num_splits,
273 softmax_scale=self.scale,
274 causal=True,
275 descale_q=layer._q_scale.reshape(1),
vllm/v1/attention/backends/mla/triton_mla.py
+6 -1
@@ -13,6 +13,9 @@ from vllm.attention.backends.abstract import (
13 from vllm.attention.ops.triton_decode_attention import decode_attention_fwd
14 from vllm.attention.ops.triton_flash_attention import triton_attention
15 from vllm.logger import init_logger
16 +from vllm.model_executor.layers.batch_invariant import (
17 + vllm_kernel_override_batch_invariant,
18 +)
19 from vllm.platforms import current_platform
20 from vllm.triton_utils import HAS_TRITON
21 from vllm.v1.attention.backends.mla.common import (
@@ -158,7 +161,9 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
161 B, q_num_heads, self.kv_lora_rank, dtype=q.dtype, device=q.device
162 )
163 lse = torch.zeros(B, q_num_heads, dtype=q.dtype, device=q.device)
161 - num_kv_splits = 4 # TODO: heuristic
164 +
165 + # For batch invariance, use only 1 split to ensure deterministic reduction
166 + num_kv_splits = 1 if vllm_kernel_override_batch_invariant() else 4
167
168 # TODO(lucas) Allocate ahead of time
169 attn_logits = torch.empty(
vllm/v1/worker/gpu_model_runner.py
-3
@@ -231,9 +231,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
231 from vllm.model_executor.models.utils import set_cpu_offload_max_bytes
232
233 set_cpu_offload_max_bytes(int(self.cache_config.cpu_offload_gb * 1024**3))
234 - from vllm.model_executor.layers.batch_invariant import init_batch_invariance
235 -
236 - init_batch_invariance()
234
235 model_config = self.model_config
236 cache_config = self.cache_config
vllm/v1/worker/gpu_worker.py
+3
@@ -764,6 +764,9 @@ def init_worker_distributed_environment(
764 ) -> None:
765 """Initialize the distributed environment."""
766 parallel_config = vllm_config.parallel_config
767 + from vllm.model_executor.layers.batch_invariant import init_batch_invariance
768 +
769 + init_batch_invariance()
770 set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)
771
772 init_distributed_environment(