[BugFix][MRV2] Don't create dummy requests longer than max_model_len (#49751)

Signed-off-by: Nick Hill <nickhill123@gmail.com>

tests/v1/worker/test_gpu_input_batch_v2.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Tests for the V2 model runner's InputBatch (vllm.v1.worker.gpu.input_batch)."""
+
+import pytest
+import torch
+
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
+
+DEVICE = "cuda"
+
+
+@pytest.mark.parametrize(
+    "num_reqs,num_tokens",
+    [
+        (256, 496),  # remainder 240: previously gave the last request 241 tokens
+        (128, 512),  # no remainder
+        (3, 8),
+        (1, 7),
+    ],
+)
+def test_make_dummy_distributes_remainder(num_reqs: int, num_tokens: int):
+    """No dummy request may exceed ceil(num_tokens / num_reqs) tokens.
+
+    Dumping the remainder on a single request can produce a dummy request with
+    seq_len > max_model_len, which the block tables cannot back; attention
+    kernels running on the dummy batch during cudagraph capture then read
+    block-table entries out of bounds (https://github.com/vllm-project/vllm/pull/49364
+    CI failure).
+    """
+    buffers = InputBuffers(
+        max_num_reqs=num_reqs, max_num_tokens=num_tokens, device=torch.device(DEVICE)
+    )
+    batch = InputBatch.make_dummy(num_reqs, num_tokens, buffers)
+
+    max_per_req = -(-num_tokens // num_reqs)
+    assert batch.num_scheduled_tokens.sum() == num_tokens
+    assert batch.num_scheduled_tokens.max() == max_per_req
+    assert batch.num_scheduled_tokens.min() >= num_tokens // num_reqs
+    # Requests with an extra token are placed at the end of the batch.
+    assert (batch.num_scheduled_tokens[:-1] <= batch.num_scheduled_tokens[1:]).all()
+
+    # seq_len == query_len for the dummy prefill-shaped batch, on GPU and CPU.
+    query_lens = batch.query_start_loc_np[1:] - batch.query_start_loc_np[:-1]
+    assert (query_lens == batch.num_scheduled_tokens).all()
+    assert torch.equal(
+        batch.seq_lens, torch.from_numpy(batch.num_scheduled_tokens).to(DEVICE)
+    )
+    assert batch.query_start_loc_np[-1] == num_tokens
+    assert torch.equal(
+        batch.query_start_loc.cpu(), torch.from_numpy(batch.query_start_loc_np)
+    )

vllm/v1/worker/gpu/input_batch.py

@@ -115,13 +115,18 @@ class InputBatch:
         expanded_idx_mapping = idx_mapping
         expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device)
 
-        num_scheduled_tokens = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)
-        num_scheduled_tokens[-1] += num_tokens % num_reqs
+        # Distribute the remainder evenly so that no dummy request exceeds
+        # ceil(num_tokens / num_reqs) <= max_model_len tokens.
+        base_tokens = num_tokens // num_reqs
+        num_extra = num_tokens % num_reqs
+        num_scheduled_tokens = np.full(num_reqs, base_tokens, dtype=np.int32)
+        if num_extra > 0:
+            num_scheduled_tokens[-num_extra:] += 1
         assert int(num_scheduled_tokens.sum()) == num_tokens
 
         # seq_len equals to query_len
-        input_buffers.seq_lens[:num_reqs] = num_tokens // num_reqs
-        input_buffers.seq_lens[num_reqs - 1] += num_tokens % num_reqs
+        input_buffers.seq_lens[: num_reqs - num_extra] = base_tokens
+        input_buffers.seq_lens[num_reqs - num_extra : num_reqs] = base_tokens + 1
         # Pad for full CUDA graph mode.
         input_buffers.seq_lens[num_reqs:] = 0
         seq_lens = input_buffers.seq_lens[:num_reqs]

vllm/v1/worker/gpu/lora_utils.py

@@ -59,8 +59,12 @@ def create_lora_capture_hook(
         return None
 
     def hook(num_active_loras: int, num_reqs: int, num_tokens: int) -> None:
+        # Match InputBatch.make_dummy: distribute the remainder evenly so no
+        # dummy request exceeds ceil(num_tokens / num_reqs) tokens.
         num_scheduled = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)
-        num_scheduled[-1] += num_tokens % num_reqs
+        num_extra = num_tokens % num_reqs
+        if num_extra > 0:
+            num_scheduled[-num_extra:] += 1
         with runner.maybe_select_dummy_loras(
             lora_config, num_scheduled, num_active_loras=num_active_loras
         ):

vllm/v1/worker/gpu/model_runner.py

@@ -564,8 +564,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             num_tokens = max(num_tokens, self.decode_query_len)
             num_reqs = num_tokens // self.decode_query_len
             assert num_tokens % self.decode_query_len == 0
-        num_tokens_per_request = [num_tokens // num_reqs] * num_reqs
-        num_tokens_per_request[-1] += num_tokens % num_reqs
+        # Distribute the remainder evenly so no dummy request exceeds
+        # ceil(num_tokens / num_reqs) <= max_model_len tokens.
+        num_tokens_per_request = [
+            num_tokens // num_reqs + (i >= num_reqs - num_tokens % num_reqs)
+            for i in range(num_reqs)
+        ]
 
         assert sum(num_tokens_per_request) == num_tokens
         num_scheduled_tokens = {