From bcdacfc1ffeef001ee2bbabdf4c8b8ceb896ae66 Mon Sep 17 00:00:00 2001 From: Bob Ok <49168652+blipbyte@users.noreply.github.com> Date: Thu, 24 Sep 2026 04:21:56 +0300 Subject: [PATCH] [Test][Determinism] Cover chunked prefill in the batch-invariance suite (#55612) Signed-off-by: Bob Ok <49168652+blipbyte@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- tests/v1/determinism/test_batch_invariance.py | 87 +++++++++++++++++-- 1 file changed, 80 insertions(+), 7 deletions(-) diff --git a/tests/v1/determinism/test_batch_invariance.py b/tests/v1/determinism/test_batch_invariance.py index e370f87833..9af55b9fbb 100644 --- a/tests/v1/determinism/test_batch_invariance.py +++ b/tests/v1/determinism/test_batch_invariance.py @@ -6,6 +6,7 @@ import random import pytest import torch +from transformers import AutoTokenizer from utils import ( BACKENDS, TEST_MODEL, @@ -203,13 +204,8 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( # Use more realistic prompts for better token generation prompts = [_random_prompt(10, 50) for _ in range(32)] - # TODO: Update prompts to have ragged lengths in order to test chunked prefill - # The above tests are not currently long enough to exercise chunking. - # prompts = ( - # [_random_prompt(10, 50) for _ in range(28)] - # + [_random_prompt(256, 512) for _ in range(50)] - # + [_random_prompt(2048, 4096) for _ in range(50)] - # ) + # Chunked prefill is covered by + # test_logprobs_bitwise_batch_invariance_ragged_chunked_prefill below. sp = SamplingParams( temperature=0.6, @@ -389,6 +385,80 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN( pytest.fail(msg) +@skip_unsupported +@pytest.mark.parametrize( + "backend", + BACKENDS, +) +def test_logprobs_bitwise_batch_invariance_ragged_chunked_prefill(backend): + """Batch invariance must hold when a prefill is split across chunks.""" + random.seed(int(os.getenv("VLLM_TEST_SEED", "12345"))) + + prompts = ( + [_random_prompt(10, 50) for _ in range(8)] + + [_random_prompt(300, 450) for _ in range(4)] + + [_random_prompt(800, 1200) for _ in range(2)] + ) + random.shuffle(prompts) + + # _random_prompt takes a word target, not a token count, and the ratio + # depends on the tokenizer, so derive the budget rather than hardcode it: + # half the longest prompt splits that prefill whatever TEST_MODEL is, while + # the shorter prompts are still co-scheduled whole. Floor is max_num_seqs. + tokenizer = AutoTokenizer.from_pretrained(TEST_MODEL) + prompt_lens = [len(tokenizer(p).input_ids) for p in prompts] + max_num_batched_tokens = max(len(prompts), max(prompt_lens) // 2) + assert max_num_batched_tokens < max(prompt_lens), ( + f"a {max_num_batched_tokens}-token budget does not split the longest " + f"prompt ({max(prompt_lens)} tokens), so nothing would be chunked" + ) + + sp = SamplingParams( + temperature=0.6, top_p=1.0, max_tokens=16, seed=1234, logprobs=5 + ) + llm = LLM_with_max_seqs( + model=TEST_MODEL, + max_num_seqs=len(prompts), + gpu_memory_utilization=float(os.getenv("VLLM_GPU_MEMORY_UTILIZATION", "0.5")), + max_model_len=4096, + attention_config={"backend": backend}, + max_num_batched_tokens=max_num_batched_tokens, + ) + + try: + bs1 = [] + for p in prompts: + out = llm.generate([p], sp, use_tqdm=False)[0] + bs1.append(_extract_step_logprobs(out)) + if any(logprobs is None for logprobs, _ in bs1): + pytest.skip("Logprobs are not available on RequestOutput.") + + outs_batched = llm.generate(prompts, sp, use_tqdm=False) + assert len(outs_batched) == len(prompts) + + mismatches = [] + for i, out in enumerate(outs_batched): + logprobs, tokens = _extract_step_logprobs(out) + bs1_logprobs, bs1_tokens = bs1[i] + n = prompt_lens[i] + if tokens != bs1_tokens: + mismatches.append(f"prompt {i} ({n} tokens): tokens differ") + elif not torch.equal(bs1_logprobs, logprobs): + delta = torch.max(torch.abs(bs1_logprobs - logprobs)).item() + mismatches.append( + f"prompt {i} ({n} tokens): logprobs differ, max|delta|={delta:.3e}" + ) + + assert not mismatches, ( + "Batch invariance violated under chunked prefill " + f"(max_num_batched_tokens={max_num_batched_tokens}): " + + "; ".join(mismatches) + ) + finally: + with contextlib.suppress(Exception): + llm.shutdown() + + @skip_unsupported @pytest.mark.parametrize( "backend", @@ -923,6 +993,7 @@ def LLM_with_max_seqs( max_model_len: int, attention_config: dict | None = None, kernel_config: dict | None = None, + max_num_batched_tokens: int | None = None, ) -> LLM: """Helper to construct an LLM with a specific max_num_seqs (batch-size limit) using the high-level v1 LLM API, while constraining memory usage. @@ -930,6 +1001,8 @@ def LLM_with_max_seqs( extra_kwargs: dict = {} if kernel_config is not None: extra_kwargs["kernel_config"] = kernel_config + if max_num_batched_tokens is not None: + extra_kwargs["max_num_batched_tokens"] = max_num_batched_tokens return LLM( model=model, max_num_seqs=max_num_seqs,