[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

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