[BugFix][MRV2] Don't create dummy requests longer than `max_model_len` (#49751)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jul 27, 2026 at 03:04 UTC
50aa83048219b70a3a68adf2fc8cd860ccc3e238
4 files changed
+72
-7
tests/v1/worker/test_gpu_input_batch_v2.py
new
+52
@@ -0,0 +1,52 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the V2 model runner's InputBatch (vllm.v1.worker.gpu.input_batch)."""
4
+
5
+import pytest
6
+import torch
7
+
8
+from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
9
+
10
+DEVICE = "cuda"
11
+
12
+
13
+@pytest.mark.parametrize(
14
+ "num_reqs,num_tokens",
15
+ [
16
+ (256, 496), # remainder 240: previously gave the last request 241 tokens
17
+ (128, 512), # no remainder
18
+ (3, 8),
19
+ (1, 7),
20
+ ],
21
+)
22
+def test_make_dummy_distributes_remainder(num_reqs: int, num_tokens: int):
23
+ """No dummy request may exceed ceil(num_tokens / num_reqs) tokens.
24
+
25
+ Dumping the remainder on a single request can produce a dummy request with
26
+ seq_len > max_model_len, which the block tables cannot back; attention
27
+ kernels running on the dummy batch during cudagraph capture then read
28
+ block-table entries out of bounds (https://github.com/vllm-project/vllm/pull/49364
29
+ CI failure).
30
+ """
31
+ buffers = InputBuffers(
32
+ max_num_reqs=num_reqs, max_num_tokens=num_tokens, device=torch.device(DEVICE)
33
+ )
34
+ batch = InputBatch.make_dummy(num_reqs, num_tokens, buffers)
35
+
36
+ max_per_req = -(-num_tokens // num_reqs)
37
+ assert batch.num_scheduled_tokens.sum() == num_tokens
38
+ assert batch.num_scheduled_tokens.max() == max_per_req
39
+ assert batch.num_scheduled_tokens.min() >= num_tokens // num_reqs
40
+ # Requests with an extra token are placed at the end of the batch.
41
+ assert (batch.num_scheduled_tokens[:-1] <= batch.num_scheduled_tokens[1:]).all()
42
+
43
+ # seq_len == query_len for the dummy prefill-shaped batch, on GPU and CPU.
44
+ query_lens = batch.query_start_loc_np[1:] - batch.query_start_loc_np[:-1]
45
+ assert (query_lens == batch.num_scheduled_tokens).all()
46
+ assert torch.equal(
47
+ batch.seq_lens, torch.from_numpy(batch.num_scheduled_tokens).to(DEVICE)
48
+ )
49
+ assert batch.query_start_loc_np[-1] == num_tokens
50
+ assert torch.equal(
51
+ batch.query_start_loc.cpu(), torch.from_numpy(batch.query_start_loc_np)
52
+ )
vllm/v1/worker/gpu/input_batch.py
+9
-4
@@ -115,13 +115,18 @@ class InputBatch:
115
expanded_idx_mapping = idx_mapping
116
expanded_local_pos = torch.zeros(num_reqs, dtype=torch.int32, device=device)
117
118
- num_scheduled_tokens = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)
119
- num_scheduled_tokens[-1] += num_tokens % num_reqs
118
+ # Distribute the remainder evenly so that no dummy request exceeds
119
+ # ceil(num_tokens / num_reqs) <= max_model_len tokens.
120
+ base_tokens = num_tokens // num_reqs
121
+ num_extra = num_tokens % num_reqs
122
+ num_scheduled_tokens = np.full(num_reqs, base_tokens, dtype=np.int32)
123
+ if num_extra > 0:
124
+ num_scheduled_tokens[-num_extra:] += 1
125
assert int(num_scheduled_tokens.sum()) == num_tokens
126
127
# seq_len equals to query_len
123
- input_buffers.seq_lens[:num_reqs] = num_tokens // num_reqs
124
- input_buffers.seq_lens[num_reqs - 1] += num_tokens % num_reqs
128
+ input_buffers.seq_lens[: num_reqs - num_extra] = base_tokens
129
+ input_buffers.seq_lens[num_reqs - num_extra : num_reqs] = base_tokens + 1
130
# Pad for full CUDA graph mode.
131
input_buffers.seq_lens[num_reqs:] = 0
132
seq_lens = input_buffers.seq_lens[:num_reqs]
vllm/v1/worker/gpu/lora_utils.py
+5
-1
@@ -59,8 +59,12 @@ def create_lora_capture_hook(
59
return None
60
61
def hook(num_active_loras: int, num_reqs: int, num_tokens: int) -> None:
62
+ # Match InputBatch.make_dummy: distribute the remainder evenly so no
63
+ # dummy request exceeds ceil(num_tokens / num_reqs) tokens.
64
num_scheduled = np.full(num_reqs, num_tokens // num_reqs, dtype=np.int32)
63
- num_scheduled[-1] += num_tokens % num_reqs
65
+ num_extra = num_tokens % num_reqs
66
+ if num_extra > 0:
67
+ num_scheduled[-num_extra:] += 1
68
with runner.maybe_select_dummy_loras(
69
lora_config, num_scheduled, num_active_loras=num_active_loras
70
):
vllm/v1/worker/gpu/model_runner.py
+6
-2
@@ -564,8 +564,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
564
num_tokens = max(num_tokens, self.decode_query_len)
565
num_reqs = num_tokens // self.decode_query_len
566
assert num_tokens % self.decode_query_len == 0
567
- num_tokens_per_request = [num_tokens // num_reqs] * num_reqs
568
- num_tokens_per_request[-1] += num_tokens % num_reqs
567
+ # Distribute the remainder evenly so no dummy request exceeds
568
+ # ceil(num_tokens / num_reqs) <= max_model_len tokens.
569
+ num_tokens_per_request = [
570
+ num_tokens // num_reqs + (i >= num_reqs - num_tokens % num_reqs)
571
+ for i in range(num_reqs)
572
+ ]
573
574
assert sum(num_tokens_per_request) == num_tokens
575
num_scheduled_tokens = {