diff --git a/tests/v1/worker/test_gpu_input_batch_v2.py b/tests/v1/worker/test_gpu_input_batch_v2.py new file mode 100644 index 0000000000..2c43cd7214 --- /dev/null +++ b/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) + ) diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 64c1096dfb..8c3a9cc030 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/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] diff --git a/vllm/v1/worker/gpu/lora_utils.py b/vllm/v1/worker/gpu/lora_utils.py index fa281f6817..6a00b35e4b 100644 --- a/vllm/v1/worker/gpu/lora_utils.py +++ b/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 ): diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 64a4f55351..f0da049ce4 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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 = {