| 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 numpy as np |
| 6 | import pytest |
| 7 | import torch |
| 8 | |
| 9 | from vllm.config import VllmConfig |
| 10 | from vllm.forward_context import set_forward_context |
| 11 | from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.v1.attention.backends.utils import get_dcp_local_seq_lens |
| 14 | from vllm.v1.worker.gpu import cp_utils |
| 15 | from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers |
| 16 | |
| 17 | DEVICE = current_platform.device_type |
| 18 | |
| 19 | |
| 20 | @pytest.mark.parametrize( |
| 21 | "num_reqs,num_tokens", |
| 22 | [ |
| 23 | (256, 496), # remainder 240: previously gave the last request 241 tokens |
| 24 | (128, 512), # no remainder |
| 25 | (3, 8), |
| 26 | (1, 7), |
| 27 | ], |
| 28 | ) |
| 29 | def test_make_dummy_distributes_remainder(num_reqs: int, num_tokens: int): |
| 30 | """No dummy request may exceed ceil(num_tokens / num_reqs) tokens. |
| 31 | |
| 32 | Dumping the remainder on a single request can produce a dummy request with |
| 33 | seq_len > max_model_len, which the block tables cannot back; attention |
| 34 | kernels running on the dummy batch during cudagraph capture then read |
| 35 | block-table entries out of bounds (https://github.com/vllm-project/vllm/pull/49364 |
| 36 | CI failure). |
| 37 | """ |
| 38 | buffers = InputBuffers( |
| 39 | max_num_reqs=num_reqs, max_num_tokens=num_tokens, device=torch.device(DEVICE) |
| 40 | ) |
| 41 | batch = InputBatch.make_dummy(num_reqs, num_tokens, buffers) |
| 42 | |
| 43 | max_per_req = -(-num_tokens // num_reqs) |
| 44 | assert batch.num_scheduled_tokens.sum() == num_tokens |
| 45 | assert batch.num_scheduled_tokens.max() == max_per_req |
| 46 | assert batch.num_scheduled_tokens.min() >= num_tokens // num_reqs |
| 47 | # Requests with an extra token are placed at the end of the batch. |
| 48 | assert (batch.num_scheduled_tokens[:-1] <= batch.num_scheduled_tokens[1:]).all() |
| 49 | |
| 50 | # seq_len == query_len for the dummy prefill-shaped batch, on GPU and CPU. |
| 51 | query_lens = batch.query_start_loc_np[1:] - batch.query_start_loc_np[:-1] |
| 52 | assert (query_lens == batch.num_scheduled_tokens).all() |
| 53 | assert torch.equal( |
| 54 | batch.seq_lens, torch.from_numpy(batch.num_scheduled_tokens).to(DEVICE) |
| 55 | ) |
| 56 | assert batch.query_start_loc_np[-1] == num_tokens |
| 57 | assert torch.equal( |
| 58 | batch.query_start_loc.cpu(), torch.from_numpy(batch.query_start_loc_np) |
| 59 | ) |
| 60 | |
| 61 | |
| 62 | @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA top-k.") |
| 63 | @pytest.mark.parametrize("is_padding", [True, False]) |
| 64 | def test_make_dummy_padding_controls_moe_routing(monkeypatch, is_padding: bool): |
| 65 | """Dummy tokens marked as padding are dropped by the MoE router (top-k id |
| 66 | -1), which is wanted for idle DP ranks. Profile runs must not mark them, or |
| 67 | no token reaches the experts and MoE memory is never profiled.""" |
| 68 | monkeypatch.setenv("VLLM_MOE_SKIP_PADDING", "1") |
| 69 | num_tokens = 16 |
| 70 | buffers = InputBuffers( |
| 71 | max_num_reqs=4, max_num_tokens=num_tokens, device=torch.device(DEVICE) |
| 72 | ) |
| 73 | batch = InputBatch.make_dummy(4, num_tokens, buffers, is_padding=is_padding) |
| 74 | hidden_states = torch.randn(num_tokens, 4, device=DEVICE) |
| 75 | router_logits = torch.randn(num_tokens, 8, device=DEVICE) |
| 76 | |
| 77 | with set_forward_context(None, VllmConfig(), is_padding=batch.is_padding): |
| 78 | _, topk_ids, _ = fused_topk(hidden_states, router_logits, 2, False) |
| 79 | |
| 80 | assert bool((topk_ids == -1).all()) is is_padding |
| 81 | assert bool((topk_ids >= 0).all()) is not is_padding |
| 82 | |
| 83 | |
| 84 | def test_prepare_dcp_local_seq_lens_uses_shared_buffer(monkeypatch): |
| 85 | """The batch must view the caller-owned buffer, sliced to padded length. |
| 86 | |
| 87 | Runtime (and capture) paths all funnel through this helper; attention |
| 88 | metadata indexes padded rows, so the view must reach |
| 89 | num_reqs_after_padding, and it must alias the persistent buffer so CUDA |
| 90 | graph replay sees the recomputed values. |
| 91 | """ |
| 92 | buffers = InputBuffers(max_num_reqs=4, max_num_tokens=4, device=torch.device("cpu")) |
| 93 | batch = InputBatch.make_dummy(2, 4, buffers) |
| 94 | batch.num_reqs_after_padding = 4 |
| 95 | |
| 96 | def fake_kernel( |
| 97 | output, |
| 98 | seq_lens, |
| 99 | dcp_size, |
| 100 | dcp_rank, |
| 101 | cp_interleave, |
| 102 | num_reqs, |
| 103 | max_num_reqs, |
| 104 | block_size, |
| 105 | ): |
| 106 | assert output is buffers.dcp_local_seq_lens |
| 107 | assert seq_lens is batch.seq_lens |
| 108 | assert (num_reqs, dcp_size, dcp_rank, cp_interleave) == (2, 4, 1, 16) |
| 109 | assert (max_num_reqs, block_size) == (4, 128) |
| 110 | output[:] = torch.tensor([1, 2, 0, 0], dtype=output.dtype) |
| 111 | |
| 112 | class FakeKernel: |
| 113 | def __getitem__(self, grid): |
| 114 | assert grid == (1,) |
| 115 | return fake_kernel |
| 116 | |
| 117 | monkeypatch.setattr(cp_utils, "_dcp_local_seq_lens_kernel", FakeKernel()) |
| 118 | batch.dcp_local_seq_lens = cp_utils.prepare_dcp_local_seq_lens( |
| 119 | buffers.dcp_local_seq_lens, |
| 120 | batch.seq_lens, |
| 121 | batch.num_reqs, |
| 122 | 4, |
| 123 | 1, |
| 124 | 16, |
| 125 | num_reqs_padded=batch.num_reqs_after_padding, |
| 126 | ) |
| 127 | |
| 128 | assert batch.dcp_local_seq_lens is not None |
| 129 | assert batch.dcp_local_seq_lens.data_ptr() == buffers.dcp_local_seq_lens.data_ptr() |
| 130 | assert batch.dcp_local_seq_lens.tolist() == [1, 2, 0, 0] |
| 131 | |
| 132 | |
| 133 | @pytest.mark.skipif(not torch.cuda.is_available(), reason="triton kernel needs CUDA") |
| 134 | @pytest.mark.parametrize("dcp_size", [2, 4]) |
| 135 | @pytest.mark.parametrize("cp_interleave", [1, 16]) |
| 136 | def test_prepare_dcp_local_seq_lens_matches_reference( |
| 137 | dcp_size: int, cp_interleave: int |
| 138 | ): |
| 139 | """Every rank's local lengths must equal the reference torch formula.""" |
| 140 | device = torch.device("cuda:0") |
| 141 | seq_lens_np = np.array([7, 16, 33, 64, 512, 1023], dtype=np.int32) |
| 142 | buffers = InputBuffers(max_num_reqs=8, max_num_tokens=32, device=device) |
| 143 | batch = InputBatch.make_dummy(6, 12, buffers) |
| 144 | buffers.seq_lens[: len(seq_lens_np)] = torch.from_numpy(seq_lens_np).to(device) |
| 145 | |
| 146 | for dcp_rank in range(dcp_size): |
| 147 | buffers.dcp_local_seq_lens.fill_(-1) |
| 148 | batch.dcp_local_seq_lens = cp_utils.prepare_dcp_local_seq_lens( |
| 149 | buffers.dcp_local_seq_lens, |
| 150 | batch.seq_lens, |
| 151 | batch.num_reqs, |
| 152 | dcp_size, |
| 153 | dcp_rank, |
| 154 | cp_interleave, |
| 155 | num_reqs_padded=batch.num_reqs_after_padding, |
| 156 | ) |
| 157 | expected = get_dcp_local_seq_lens( |
| 158 | torch.from_numpy(seq_lens_np), dcp_size, dcp_rank, cp_interleave |
| 159 | ) |
| 160 | assert batch.dcp_local_seq_lens is not None |
| 161 | assert torch.equal(batch.dcp_local_seq_lens.cpu(), expected.to(torch.int32)) |