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

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

[Test] Batch Invariant: Unit test using parameterized backend (#27478)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed Oct 28, 2025 at 16:51 UTC 6afc28a9ba236b35bbe77f96d8e7c499d1b0323a

2 files changed +208 -204

tests/v1/generation/test_batch_invariance.py

+207 -203

@@ -17,16 +17,10 @@ skip_unsupported = pytest.mark.skipif(

19 @pytest.fixture(autouse=True)

20 -def enable_batch_invariant_mode():

20 +def enable_batch_invariant_mode(monkeypatch: pytest.MonkeyPatch):

21 """Automatically enable batch invariant kernel overrides for all tests."""

22 - old_value = os.environ.get("VLLM_BATCH_INVARIANT")

23 - os.environ["VLLM_BATCH_INVARIANT"] = "1"

22 + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1")

23 yield

25 - # Restore original value after test

26 - if old_value is None:

27 - os.environ.pop("VLLM_BATCH_INVARIANT", None)

28 - else:

29 - os.environ["VLLM_BATCH_INVARIANT"] = old_value

26 def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:

@@ -76,7 +70,13 @@ def _random_prompt(min_words: int = 1024, max_words: int = 1024 * 2) -> str:

71 @skip_unsupported

72 @pytest.mark.timeout(1000)

79 -def test_v1_generation_is_deterministic_across_batch_sizes_with_needle():

73 +@pytest.mark.parametrize(

74 + "backend",

75 + ["FLASH_ATTN", "FLASHINFER", "FLASH_ATTN_MLA", "FLASHINFER_MLA", "TRITON_MLA"],

76 +)

77 +def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(

78 + backend, monkeypatch: pytest.MonkeyPatch

79 +):

80 """

81 Ensures that the same request (the 'needle' prompt) yields identical output

82 whether run alone (bs=1) or mixed into a larger batch (e.g., bs=64),

@@ -101,6 +101,7 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle():

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

102 random.seed(seed)

103

104 + monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)

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

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

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

@@ -220,11 +221,15 @@ def _extract_step_logprobs(request_output):

221

222

223 @skip_unsupported

223 -@pytest.mark.parametrize("backend", ["FLASH_ATTN", "FLASHINFER"])

224 +@pytest.mark.parametrize(

225 + "backend",

226 + ["FLASH_ATTN", "FLASHINFER", "FLASH_ATTN_MLA", "FLASHINFER_MLA", "TRITON_MLA"],

227 +)

228 @pytest.mark.forked

225 -def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(backend):

226 - backend = os.getenv("VLLM_ATTENTION_BACKEND", backend)

227 - os.environ["VLLM_ATTENTION_BACKEND"] = backend

229 +def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(

230 + backend, monkeypatch: pytest.MonkeyPatch

231 +):

232 + monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)

233

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

235 random.seed(seed)

@@ -435,11 +440,16 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(backend):

440

441

442 @skip_unsupported

438 -def test_simple_generation():

443 +@pytest.mark.parametrize(

444 + "backend",

445 + ["FLASH_ATTN", "FLASHINFER", "FLASH_ATTN_MLA", "FLASHINFER_MLA", "TRITON_MLA"],

446 +)

447 +def test_simple_generation(backend, monkeypatch: pytest.MonkeyPatch):

448 """

449 Simple test that runs the model with a basic prompt and prints the output.

450 Useful for quick smoke testing and debugging.

451 """

452 + monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)

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

454

455 llm = LLM(

@@ -481,9 +491,14 @@ def test_simple_generation():

491

492

493 @skip_unsupported

484 -@pytest.mark.parametrize("backend", ["FLASH_ATTN", "FLASHINFER"])

494 +@pytest.mark.parametrize(

495 + "backend",

496 + ["FLASH_ATTN", "FLASHINFER", "FLASH_ATTN_MLA", "FLASHINFER_MLA", "TRITON_MLA"],

497 +)

498 @pytest.mark.forked

486 -def test_logprobs_WITHOUT_batch_invariance_should_FAIL(backend):

499 +def test_logprobs_without_batch_invariance_should_fail(

500 + backend, monkeypatch: pytest.MonkeyPatch

501 +):

502 """

503 This test is the inverse of test_logprobs_bitwise_batch_invariance_bs1_vs_bsN.

504 It DISABLES batch invariance mode and expects to see non-deterministic behavior

@@ -493,224 +508,214 @@ def test_logprobs_WITHOUT_batch_invariance_should_FAIL(backend):

508 The test will PASS if we detect differences (proving batch invariance matters).

509 The test will FAIL if everything matches (suggesting batch invariance isn't needed).

510 """

496 - backend = os.getenv("VLLM_ATTENTION_BACKEND", backend)

497 - os.environ["VLLM_ATTENTION_BACKEND"] = backend

511 + monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)

512

513 # CRITICAL: Disable batch invariance for this test

500 - old_value = os.environ.get("VLLM_BATCH_INVARIANT")

501 - os.environ["VLLM_BATCH_INVARIANT"] = "0"

514 + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0")

515

503 - try:

504 - seed = int(os.getenv("VLLM_TEST_SEED", "12345"))

505 - random.seed(seed)

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

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

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

509 - print(f"\n{'=' * 80}")

510 - print("BATCH INVARIANCE DISABLED: Expecting non-deterministic behavior")

511 - print(f"{'=' * 80}\n")

521 + print(f"\n{'=' * 80}")

522 + print("BATCH INVARIANCE DISABLED: Expecting non-deterministic behavior")

523 + print(f"{'=' * 80}\n")

524

513 - llm = LLM(

514 - model=model_name,

515 - tensor_parallel_size=tp_size,

516 - enable_prefix_caching=False,

517 - max_num_seqs=32,

518 - max_model_len=8192,

519 - dtype="bfloat16",

520 - )

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

522 - # build ragged prompts to change shapes significantly across BS=1 vs BS=N

523 - long_min = int(os.getenv("VLLM_MIN_PROMPT", "768"))

524 - long_max = int(os.getenv("VLLM_MAX_PROMPT", "2048"))

525 - prompts: list[str] = []

526 - options = [

527 - (max(long_min, 1536), max(long_max, 3072)), # very long

528 - (max(1024, long_min), max(2048, long_max)), # long

529 - (256, 512), # mid

530 - (10, 20), # short

531 - ]

532 -

533 - for _ in range(32):

534 - lo, hi = random.choice(options)

535 - prompts.append(_random_prompt(lo, hi))

536 -

537 - sp = SamplingParams(

538 - temperature=0.6,

539 - top_p=1.0,

540 - max_tokens=8,

541 - seed=1234,

542 - logprobs=5,

543 - )

534 + # build ragged prompts to change shapes significantly across BS=1 vs BS=N

535 + long_min = int(os.getenv("VLLM_MIN_PROMPT", "768"))

536 + long_max = int(os.getenv("VLLM_MAX_PROMPT", "2048"))

537 + prompts: list[str] = []

538 + options = [

539 + (max(long_min, 1536), max(long_max, 3072)), # very long

540 + (max(1024, long_min), max(2048, long_max)), # long

541 + (256, 512), # mid

542 + (10, 20), # short

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")

545 + for _ in range(32):

546 + lo, hi = random.choice(options)

547 + prompts.append(_random_prompt(lo, hi))

548

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]}..."

549 + sp = SamplingParams(

550 + temperature=0.6,

551 + top_p=1.0,

552 + max_tokens=8,

553 + seed=1234,

554 + logprobs=5,

555 + )

556 +

557 + # BS=1: run prompts individually and collect logprobs per step.

558 + print("\n" + "=" * 80)

559 + print("STARTING BS=1 RUNS (each prompt individually)")

560 + print("=" * 80 + "\n")

561 +

562 + bs1_logprobs_per_prompt = []

563 + bs1_tokens_per_prompt = []

564 + for idx, p in enumerate(prompts):

565 + print(f"\n[BS=1] Running prompt {idx}/{len(prompts)} - Preview: {p[:80]}...")

566 + outs = llm.generate([p], sp, use_tqdm=False)

567 + assert len(outs) == 1

568 + step_logprobs, token_ids = _extract_step_logprobs(outs[0])

569 + if step_logprobs is None:

570 + pytest.skip(

571 + "Logits are not available on RequestOutput; "

572 + "enable logprobs return to run this test."

573 )

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,

574 + bs1_logprobs_per_prompt.append(step_logprobs)

575 + bs1_tokens_per_prompt.append(token_ids)

576 + print(f"[BS=1] Prompt {idx} generated tokens: {token_ids}")

577 +

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

579 + print("\n" + "=" * 80)

580 + print(f"STARTING BS={len(prompts)} RUN (all prompts batched)")

581 + print("=" * 80 + "\n")

582 +

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

584 + assert len(outs_batched) == len(prompts)

585 + bsN_logprobs_per_prompt = []

586 + bsN_tokens_per_prompt = []

587 +

588 + print(f"\n[BS={len(prompts)}] Processing batched outputs...")

589 + for idx, o in enumerate(outs_batched):

590 + tokens = o.outputs[0].token_ids if o.outputs else "N/A"

591 + print(f"[BS={len(prompts)}] Prompt {idx} generated tokens: {tokens}")

592 + step_logprobs, token_ids = _extract_step_logprobs(o)

593 + if step_logprobs is None:

594 + pytest.skip(

595 + "Logits are not available on RequestOutput; "

596 + "enable logprobs return to run this test."

597 )

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 - )

598 + bsN_logprobs_per_prompt.append(step_logprobs)

599 + bsN_tokens_per_prompt.append(token_ids)

600 +

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

602 + differences_found = []

603 + for i, (logprobs_bs1, logprobs_bsN, tokens_bs1, tokens_bsN) in enumerate(

604 + zip(

605 + bs1_logprobs_per_prompt,

606 + bsN_logprobs_per_prompt,

607 + bs1_tokens_per_prompt,

608 + bsN_tokens_per_prompt,

609 + )

610 + ):

611 + if len(logprobs_bs1) != len(logprobs_bsN):

612 + reason = (

613 + f"Different number of steps: {len(logprobs_bs1)} (BS=1) "

614 + f"vs {len(logprobs_bsN)} (BS=N)"

615 + )

616 + differences_found.append(

617 + {

618 + "prompt_idx": i,

619 + "step": "all",

620 + "reason": reason,

621 + "prompt_preview": prompts[i][:100],

622 + "bs1_tokens": tokens_bs1,

623 + "bsN_tokens": tokens_bsN,

624 + }

625 + )

626 + continue

627 +

628 + # Check if tokens match first

629 + if tokens_bs1 != tokens_bsN:

630 + differences_found.append(

631 + {

632 + "prompt_idx": i,

633 + "step": "sampling",

634 + "reason": "Different tokens sampled",

635 + "prompt_preview": prompts[i][:100],

636 + "bs1_tokens": tokens_bs1,

637 + "bsN_tokens": tokens_bsN,

638 + }

639 + )

640 + continue

641 +

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

643 + if a.shape != b.shape:

644 differences_found.append(

645 {

646 "prompt_idx": i,

609 - "step": "all",

610 - "reason": reason,

647 + "step": t,

648 + "reason": f"Shape mismatch: {a.shape} vs {b.shape}",

649 "prompt_preview": prompts[i][:100],

650 "bs1_tokens": tokens_bs1,

651 "bsN_tokens": tokens_bsN,

652 }

653 )

616 - continue

654 + break

655

618 - # Check if tokens match first

619 - if tokens_bs1 != tokens_bsN:

656 + if not torch.equal(a, b):

657 + max_diff = torch.abs(a - b).max().item()

658 + print(

659 + f"\n[EXPECTED DIVERGENCE FOUND] Prompt {i}, "

660 + f"Token {t}: max_diff={max_diff:.6e}"

661 + )

662 + bs1_tok = tokens_bs1[t] if t < len(tokens_bs1) else "N/A"

663 + bsN_tok = tokens_bsN[t] if t < len(tokens_bsN) else "N/A"

664 + print(f" Token IDs: bs1={bs1_tok}, bsN={bsN_tok}")

665 + print(f" BS=1 logprob: {a.tolist()}")

666 + print(f" BS=N logprob: {b.tolist()}")

667 differences_found.append(

668 {

669 "prompt_idx": i,

623 - "step": "sampling",

624 - "reason": "Different tokens sampled",

670 + "step": t,

671 + "reason": f"Bitwise mismatch (max_diff={max_diff:.6e})",

672 "prompt_preview": prompts[i][:100],

673 "bs1_tokens": tokens_bs1,

674 "bsN_tokens": tokens_bsN,

675 }

676 )

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)

677 + break

678

702 - finally:

703 - # Restore original value

704 - if old_value is None:

705 - os.environ.pop("VLLM_BATCH_INVARIANT", None)

706 - else:

707 - os.environ["VLLM_BATCH_INVARIANT"] = old_value

679 + # Print summary

680 + print(f"\n{'=' * 80}")

681 + if differences_found:

682 + success_msg = (

683 + f"✓ SUCCESS: Batch invariance is doing something! "

684 + f"Found {len(differences_found)}/{len(prompts)} prompts "

685 + f"with differences when batch invariance was DISABLED."

686 + )

687 + print(success_msg)

688 + print(f"{'=' * 80}")

689 + for diff in differences_found:

690 + print(f"\nPrompt {diff['prompt_idx']} (step {diff['step']}):")

691 + print(f" Reason: {diff['reason']}")

692 + print(f" Preview: {diff['prompt_preview']}...")

693 + if "bs1_tokens" in diff:

694 + print(f" BS=1 tokens: {diff['bs1_tokens']}")

695 + if "bsN_tokens" in diff:

696 + print(f" BS=N tokens: {diff['bsN_tokens']}")

697 + print(f"{'=' * 80}\n")

698 + # Test PASSES because we found differences (batch invariance matters!)

699 + return

700 + else:

701 + # Test FAILS because everything matched even without batch invariance

702 + fail_msg = (

703 + f"✗ UNEXPECTED: All {len(prompts)} prompts matched "

704 + f"between BS=1 and BS=N even with batch invariance DISABLED. "

705 + f"This suggests batch invariance might not be necessary, "

706 + f"or the test needs more sensitive prompts."

707 + )

708 + print(fail_msg)

709 + print(f"{'=' * 80}\n")

710 + pytest.fail(fail_msg)

711

712

713 @skip_unsupported

714 @pytest.mark.parametrize("backend", ["FLASH_ATTN"])

715 @pytest.mark.forked

713 -def test_decode_logprobs_match_prefill_logprobs(backend):

716 +def test_decode_logprobs_match_prefill_logprobs(

717 + backend, monkeypatch: pytest.MonkeyPatch

718 +):

719 """

720 Test that verifies decode logprobs match prefill logprobs.

721

@@ -724,8 +729,7 @@ def test_decode_logprobs_match_prefill_logprobs(backend):

729 This ensures that the logprobs from decode are consistent with what

730 we would get if we ran prefill on each prefix.

731 """

727 - backend = os.getenv("VLLM_ATTENTION_BACKEND", backend)

728 - os.environ["VLLM_ATTENTION_BACKEND"] = backend

732 + monkeypatch.setenv("VLLM_ATTENTION_BACKEND", backend)

733

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

735 random.seed(seed)

vllm/model_executor/layers/batch_invariant.py

+1 -1

@@ -753,13 +753,13 @@ def override_envs_for_invariance():

753 curr_attn_backend = envs.VLLM_ATTENTION_BACKEND

754 supported_backends = [

755 "FLASH_ATTN", # best supported backend

756 - "FLEX_ATTENTION",

756 "FLASHINFER",

757 "FLASH_ATTN_MLA",

758 "FLASHINFER_MLA",

759 "TRITON_MLA",

760 # Not yet supported MLA backends

761 # "FLASHMLA",

762 + # "FLEX_ATTENTION", # IMA issue even if we disable batch invariance

763 ]

764 if curr_attn_backend not in supported_backends:

765 warning = (