[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>
Bob Ok committed
Sep 24, 2026 at 04:21 UTC
bcdacfc1ffeef001ee2bbabdf4c8b8ceb896ae66
1 file changed
+80
-7
tests/v1/determinism/test_batch_invariance.py
+80
-7
@@ -6,6 +6,7 @@ import random
6
7
import pytest
8
import torch
9
+from transformers import AutoTokenizer
10
from utils import (
11
BACKENDS,
12
TEST_MODEL,
@@ -203,13 +204,8 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
204
# Use more realistic prompts for better token generation
205
prompts = [_random_prompt(10, 50) for _ in range(32)]
206
206
- # TODO: Update prompts to have ragged lengths in order to test chunked prefill
207
- # The above tests are not currently long enough to exercise chunking.
208
- # prompts = (
209
- # [_random_prompt(10, 50) for _ in range(28)]
210
- # + [_random_prompt(256, 512) for _ in range(50)]
211
- # + [_random_prompt(2048, 4096) for _ in range(50)]
212
- # )
207
+ # Chunked prefill is covered by
208
+ # test_logprobs_bitwise_batch_invariance_ragged_chunked_prefill below.
209
210
sp = SamplingParams(
211
temperature=0.6,
@@ -389,6 +385,80 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
385
pytest.fail(msg)
386
387
388
+@skip_unsupported
389
+@pytest.mark.parametrize(
390
+ "backend",
391
+ BACKENDS,
392
+)
393
+def test_logprobs_bitwise_batch_invariance_ragged_chunked_prefill(backend):
394
+ """Batch invariance must hold when a prefill is split across chunks."""
395
+ random.seed(int(os.getenv("VLLM_TEST_SEED", "12345")))
396
+
397
+ prompts = (
398
+ [_random_prompt(10, 50) for _ in range(8)]
399
+ + [_random_prompt(300, 450) for _ in range(4)]
400
+ + [_random_prompt(800, 1200) for _ in range(2)]
401
+ )
402
+ random.shuffle(prompts)
403
+
404
+ # _random_prompt takes a word target, not a token count, and the ratio
405
+ # depends on the tokenizer, so derive the budget rather than hardcode it:
406
+ # half the longest prompt splits that prefill whatever TEST_MODEL is, while
407
+ # the shorter prompts are still co-scheduled whole. Floor is max_num_seqs.
408
+ tokenizer = AutoTokenizer.from_pretrained(TEST_MODEL)
409
+ prompt_lens = [len(tokenizer(p).input_ids) for p in prompts]
410
+ max_num_batched_tokens = max(len(prompts), max(prompt_lens) // 2)
411
+ assert max_num_batched_tokens < max(prompt_lens), (
412
+ f"a {max_num_batched_tokens}-token budget does not split the longest "
413
+ f"prompt ({max(prompt_lens)} tokens), so nothing would be chunked"
414
+ )
415
+
416
+ sp = SamplingParams(
417
+ temperature=0.6, top_p=1.0, max_tokens=16, seed=1234, logprobs=5
418
+ )
419
+ llm = LLM_with_max_seqs(
420
+ model=TEST_MODEL,
421
+ max_num_seqs=len(prompts),
422
+ gpu_memory_utilization=float(os.getenv("VLLM_GPU_MEMORY_UTILIZATION", "0.5")),
423
+ max_model_len=4096,
424
+ attention_config={"backend": backend},
425
+ max_num_batched_tokens=max_num_batched_tokens,
426
+ )
427
+
428
+ try:
429
+ bs1 = []
430
+ for p in prompts:
431
+ out = llm.generate([p], sp, use_tqdm=False)[0]
432
+ bs1.append(_extract_step_logprobs(out))
433
+ if any(logprobs is None for logprobs, _ in bs1):
434
+ pytest.skip("Logprobs are not available on RequestOutput.")
435
+
436
+ outs_batched = llm.generate(prompts, sp, use_tqdm=False)
437
+ assert len(outs_batched) == len(prompts)
438
+
439
+ mismatches = []
440
+ for i, out in enumerate(outs_batched):
441
+ logprobs, tokens = _extract_step_logprobs(out)
442
+ bs1_logprobs, bs1_tokens = bs1[i]
443
+ n = prompt_lens[i]
444
+ if tokens != bs1_tokens:
445
+ mismatches.append(f"prompt {i} ({n} tokens): tokens differ")
446
+ elif not torch.equal(bs1_logprobs, logprobs):
447
+ delta = torch.max(torch.abs(bs1_logprobs - logprobs)).item()
448
+ mismatches.append(
449
+ f"prompt {i} ({n} tokens): logprobs differ, max|delta|={delta:.3e}"
450
+ )
451
+
452
+ assert not mismatches, (
453
+ "Batch invariance violated under chunked prefill "
454
+ f"(max_num_batched_tokens={max_num_batched_tokens}): "
455
+ + "; ".join(mismatches)
456
+ )
457
+ finally:
458
+ with contextlib.suppress(Exception):
459
+ llm.shutdown()
460
+
461
+
462
@skip_unsupported
463
@pytest.mark.parametrize(
464
"backend",
@@ -923,6 +993,7 @@ def LLM_with_max_seqs(
993
max_model_len: int,
994
attention_config: dict | None = None,
995
kernel_config: dict | None = None,
996
+ max_num_batched_tokens: int | None = None,
997
) -> LLM:
998
"""Helper to construct an LLM with a specific max_num_seqs (batch-size limit)
999
using the high-level v1 LLM API, while constraining memory usage.
@@ -930,6 +1001,8 @@ def LLM_with_max_seqs(
1001
extra_kwargs: dict = {}
1002
if kernel_config is not None:
1003
extra_kwargs["kernel_config"] = kernel_config
1004
+ if max_num_batched_tokens is not None:
1005
+ extra_kwargs["max_num_batched_tokens"] = max_num_batched_tokens
1006
return LLM(
1007
model=model,
1008
max_num_seqs=max_num_seqs,