| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | |
| 7 | from vllm.platforms import current_platform |
| 8 | from vllm.v1.attention.backends.utils import PAD_SLOT_ID |
| 9 | from vllm.v1.worker.gpu.block_table import BlockTables |
| 10 | |
| 11 | pytestmark = pytest.mark.skipif( |
| 12 | not current_platform.is_cuda(), |
| 13 | reason="requires CUDA", |
| 14 | ) |
| 15 | |
| 16 | |
| 17 | def test_block_tables_apply_staged_writes_fuses_kv_groups(monkeypatch): |
| 18 | device = torch.device("cuda") |
| 19 | block_tables = BlockTables( |
| 20 | block_sizes=[16, 32, 8], |
| 21 | max_num_reqs=4, |
| 22 | max_num_batched_tokens=64, |
| 23 | max_num_blocks_per_group=[8, 8, 8], |
| 24 | device=device, |
| 25 | kernel_block_sizes=[16, 16, 8], |
| 26 | ) |
| 27 | |
| 28 | def fail_if_apply_write_called(): |
| 29 | pytest.fail("multi-group writes should use the fused apply kernel") |
| 30 | |
| 31 | for block_table in block_tables.block_tables: |
| 32 | monkeypatch.setattr(block_table, "apply_write", fail_if_apply_write_called) |
| 33 | |
| 34 | block_tables.append_block_ids( |
| 35 | req_index=0, |
| 36 | new_block_ids=([1, 2], [10, 11], []), |
| 37 | overwrite=True, |
| 38 | ) |
| 39 | block_tables.append_block_ids( |
| 40 | req_index=1, |
| 41 | new_block_ids=([3], [12], [5, 6]), |
| 42 | overwrite=True, |
| 43 | ) |
| 44 | block_tables.apply_staged_writes() |
| 45 | torch.accelerator.synchronize() |
| 46 | |
| 47 | assert torch.equal( |
| 48 | block_tables.block_tables[0].gpu[0, :2], |
| 49 | torch.tensor([1, 2], dtype=torch.int32, device=device), |
| 50 | ) |
| 51 | # Group 1 has blocks_per_kv_block == 2, so each KV block expands to two |
| 52 | # kernel block IDs. |
| 53 | assert torch.equal( |
| 54 | block_tables.block_tables[1].gpu[0, :4], |
| 55 | torch.tensor([20, 21, 22, 23], dtype=torch.int32, device=device), |
| 56 | ) |
| 57 | assert torch.equal( |
| 58 | block_tables.block_tables[0].gpu[1, :1], |
| 59 | torch.tensor([3], dtype=torch.int32, device=device), |
| 60 | ) |
| 61 | assert torch.equal( |
| 62 | block_tables.block_tables[1].gpu[1, :2], |
| 63 | torch.tensor([24, 25], dtype=torch.int32, device=device), |
| 64 | ) |
| 65 | assert torch.equal( |
| 66 | block_tables.block_tables[2].gpu[1, :2], |
| 67 | torch.tensor([5, 6], dtype=torch.int32, device=device), |
| 68 | ) |
| 69 | assert block_tables.num_blocks.np[0, 0] == 2 |
| 70 | assert block_tables.num_blocks.np[1, 0] == 4 |
| 71 | assert block_tables.num_blocks.np[2, 0] == 0 |
| 72 | assert block_tables.num_blocks.np[0, 1] == 1 |
| 73 | assert block_tables.num_blocks.np[1, 1] == 2 |
| 74 | assert block_tables.num_blocks.np[2, 1] == 2 |
| 75 | assert torch.equal( |
| 76 | block_tables.num_blocks.gpu[:, :2], |
| 77 | torch.tensor([[2, 1], [4, 2], [0, 2]], dtype=torch.int32, device=device), |
| 78 | ) |
| 79 | |
| 80 | for block_table in block_tables.block_tables: |
| 81 | assert not block_table._staged_write_indices |
| 82 | assert not block_table._staged_write_starts |
| 83 | assert not block_table._staged_write_contents |
| 84 | assert not block_table._staged_write_cu_lens |
| 85 | |
| 86 | block_tables.append_block_ids( |
| 87 | req_index=0, |
| 88 | new_block_ids=([7], [13], [8]), |
| 89 | overwrite=False, |
| 90 | ) |
| 91 | block_tables.apply_staged_writes() |
| 92 | torch.accelerator.synchronize() |
| 93 | |
| 94 | assert torch.equal( |
| 95 | block_tables.block_tables[0].gpu[0, :3], |
| 96 | torch.tensor([1, 2, 7], dtype=torch.int32, device=device), |
| 97 | ) |
| 98 | assert torch.equal( |
| 99 | block_tables.block_tables[1].gpu[0, :6], |
| 100 | torch.tensor([20, 21, 22, 23, 26, 27], dtype=torch.int32, device=device), |
| 101 | ) |
| 102 | assert torch.equal( |
| 103 | block_tables.block_tables[2].gpu[0, :1], |
| 104 | torch.tensor([8], dtype=torch.int32, device=device), |
| 105 | ) |
| 106 | assert block_tables.num_blocks.np[0, 0] == 3 |
| 107 | assert block_tables.num_blocks.np[1, 0] == 6 |
| 108 | assert block_tables.num_blocks.np[2, 0] == 1 |
| 109 | |
| 110 | |
| 111 | def test_block_tables_apply_staged_writes_single_group(): |
| 112 | device = torch.device("cuda") |
| 113 | block_tables = BlockTables( |
| 114 | block_sizes=[16], |
| 115 | max_num_reqs=2, |
| 116 | max_num_batched_tokens=16, |
| 117 | max_num_blocks_per_group=[4], |
| 118 | device=device, |
| 119 | kernel_block_sizes=[16], |
| 120 | ) |
| 121 | |
| 122 | block_tables.append_block_ids( |
| 123 | req_index=0, |
| 124 | new_block_ids=([1, 2],), |
| 125 | overwrite=True, |
| 126 | ) |
| 127 | block_tables.apply_staged_writes() |
| 128 | torch.accelerator.synchronize() |
| 129 | |
| 130 | assert torch.equal( |
| 131 | block_tables.block_tables[0].gpu[0, :2], |
| 132 | torch.tensor([1, 2], dtype=torch.int32, device=device), |
| 133 | ) |
| 134 | |
| 135 | |
| 136 | def test_block_tables_skip_custom_slot_mapping_groups(): |
| 137 | device = torch.device("cuda") |
| 138 | block_tables = BlockTables( |
| 139 | block_sizes=[8, 262144], |
| 140 | max_num_reqs=1, |
| 141 | max_num_batched_tokens=4, |
| 142 | max_num_blocks_per_group=[1, 1], |
| 143 | device=device, |
| 144 | kernel_block_sizes=[8, 262144], |
| 145 | slot_mapping_enabled=[False, True], |
| 146 | ) |
| 147 | block_tables.append_block_ids( |
| 148 | req_index=0, |
| 149 | new_block_ids=([7], [12]), |
| 150 | overwrite=True, |
| 151 | ) |
| 152 | block_tables.apply_staged_writes() |
| 153 | |
| 154 | idx_mapping = torch.tensor([0], dtype=torch.int32, device=device) |
| 155 | query_start_loc = torch.tensor([0, 2], dtype=torch.int32, device=device) |
| 156 | positions = torch.tensor([153797, 165757], dtype=torch.int64, device=device) |
| 157 | slot_mappings = block_tables.compute_slot_mappings( |
| 158 | idx_mapping, |
| 159 | query_start_loc, |
| 160 | positions, |
| 161 | num_tokens_padded=2, |
| 162 | ) |
| 163 | torch.accelerator.synchronize() |
| 164 | |
| 165 | assert slot_mappings[0].tolist() == [-1, -1] |
| 166 | assert slot_mappings[1].tolist() == [ |
| 167 | 12 * 262144 + 153797, |
| 168 | 12 * 262144 + 165757, |
| 169 | ] |
| 170 | |
| 171 | |
| 172 | @pytest.mark.parametrize("cp_rank", range(4)) |
| 173 | def test_dcp_slot_mapping_with_smaller_kernel_blocks(cp_rank: int): |
| 174 | """Only sharded groups use DCP interleave in logical-block coordinates.""" |
| 175 | device = torch.device("cuda") |
| 176 | block_tables = BlockTables( |
| 177 | block_sizes=[128, 128], |
| 178 | max_num_reqs=1, |
| 179 | max_num_batched_tokens=1024, |
| 180 | max_num_blocks_per_group=[2, 8], |
| 181 | device=device, |
| 182 | kernel_block_sizes=[64, 64], |
| 183 | dcp_sharded=[True, False], |
| 184 | cp_size=4, |
| 185 | cp_rank=cp_rank, |
| 186 | cp_interleave=128, |
| 187 | ) |
| 188 | block_tables.append_block_ids( |
| 189 | req_index=0, |
| 190 | new_block_ids=([5, 9], list(range(10, 18))), |
| 191 | overwrite=True, |
| 192 | ) |
| 193 | block_tables.apply_staged_writes() |
| 194 | |
| 195 | idx_mapping = torch.zeros(1, dtype=torch.int32, device=device) |
| 196 | query_start_loc = torch.tensor([0, 1024], dtype=torch.int32, device=device) |
| 197 | positions = torch.arange(1024, dtype=torch.int64, device=device) |
| 198 | actual = block_tables.compute_slot_mappings( |
| 199 | idx_mapping, |
| 200 | query_start_loc, |
| 201 | positions, |
| 202 | num_tokens_padded=1024, |
| 203 | ) |
| 204 | |
| 205 | expected = torch.full((1024,), -1, dtype=torch.int64, device=device) |
| 206 | first_start = cp_rank * 128 |
| 207 | second_start = 512 + first_start |
| 208 | expected[first_start : first_start + 128] = torch.arange( |
| 209 | 5 * 128, 6 * 128, dtype=torch.int64, device=device |
| 210 | ) |
| 211 | expected[second_start : second_start + 128] = torch.arange( |
| 212 | 9 * 128, 10 * 128, dtype=torch.int64, device=device |
| 213 | ) |
| 214 | assert torch.equal(actual[0], expected) |
| 215 | assert torch.equal(actual[1], positions + 10 * 128) |
| 216 | |
| 217 | |
| 218 | def test_v1_block_table_move_row_clears_vacated_row(): |
| 219 | """condense() moves the last row into a freed slot; the vacated row must |
| 220 | not keep stale block ids. Padded dummy-run batches dereference stale rows |
| 221 | as mamba state slots (bypassing the NULL_BLOCK_ID fill of real decode |
| 222 | padding) and write state in place there — corrupting the blocks' new |
| 223 | owner once they are reallocated, e.g. to an in-flight NIXL load.""" |
| 224 | from vllm.v1.worker.block_table import BlockTable |
| 225 | |
| 226 | block_table = BlockTable( |
| 227 | block_size=16, |
| 228 | max_num_reqs=4, |
| 229 | max_num_blocks_per_req=8, |
| 230 | max_num_batched_tokens=64, |
| 231 | pin_memory=False, |
| 232 | device=torch.device("cuda"), |
| 233 | kernel_block_size=16, |
| 234 | cp_kv_cache_interleave_size=1, |
| 235 | ) |
| 236 | block_table.add_row([7, 8, 9], row_idx=0) |
| 237 | block_table.add_row([4, 5], row_idx=1) |
| 238 | |
| 239 | block_table.move_row(1, 0) |
| 240 | |
| 241 | assert block_table.block_table.np[0, :2].tolist() == [4, 5] |
| 242 | assert block_table.num_blocks_per_row[0] == 2 |
| 243 | # The vacated source row routes to the reserved null block. |
| 244 | assert block_table.num_blocks_per_row[1] == 0 |
| 245 | assert (block_table.block_table.np[1] == 0).all() |
| 246 | |
| 247 | |
| 248 | def test_get_dummy_block_tables_returns_zeroed_rows(): |
| 249 | """Dummy runs bypass the gather, so the persistent input_block_tables |
| 250 | hold the previous real step's rows. Mamba/GDN metadata routes in-place |
| 251 | state writes through block_table[:, 0] (dummy slot mappings are |
| 252 | PAD-filled, state indices are not), so stale rows would direct dummy |
| 253 | state writes at freed — possibly reallocated — blocks. |
| 254 | get_dummy_block_tables must hand out zeroed (null block) rows while |
| 255 | preserving the persistent storage address for CUDA graphs.""" |
| 256 | device = torch.device("cuda") |
| 257 | block_tables = BlockTables( |
| 258 | block_sizes=[16], |
| 259 | max_num_reqs=4, |
| 260 | max_num_batched_tokens=64, |
| 261 | max_num_blocks_per_group=[8], |
| 262 | device=device, |
| 263 | kernel_block_sizes=[16], |
| 264 | ) |
| 265 | # Simulate a real step: stage a request's blocks and gather them into |
| 266 | # the persistent input block tables. |
| 267 | block_tables.append_block_ids(req_index=0, new_block_ids=([1, 2],), overwrite=True) |
| 268 | block_tables.apply_staged_writes() |
| 269 | idx_mapping = torch.zeros(1, dtype=torch.int32, device=device) |
| 270 | block_tables.gather_block_tables(idx_mapping, num_reqs_padded=1) |
| 271 | torch.accelerator.synchronize() |
| 272 | assert block_tables.input_block_tables[0][0, 0].item() == 1 |
| 273 | |
| 274 | dummy = block_tables.get_dummy_block_tables(num_reqs=1) |
| 275 | torch.accelerator.synchronize() |
| 276 | assert (dummy[0] == 0).all() |
| 277 | # CUDA graph invariant: same persistent tensor, not a fresh allocation. |
| 278 | assert dummy[0].data_ptr() == block_tables.input_block_tables[0].data_ptr() |
| 279 | |
| 280 | |
| 281 | def test_dummy_request_slot_mapping_is_pad(): |
| 282 | """idx_mapping == -1 marks a dummy (or CUDA-graph padding) request. |
| 283 | |
| 284 | A dummy draft decode step must not resolve slots through the persistent |
| 285 | block-table row of a real request slot, which may be stale and point at |
| 286 | blocks that are now in the prefix cache. |
| 287 | """ |
| 288 | device = torch.device("cuda") |
| 289 | block_tables = BlockTables( |
| 290 | block_sizes=[4], |
| 291 | max_num_reqs=2, |
| 292 | max_num_batched_tokens=16, |
| 293 | max_num_blocks_per_group=[4], |
| 294 | device=device, |
| 295 | kernel_block_sizes=[4], |
| 296 | ) |
| 297 | block_tables.append_block_ids(req_index=0, new_block_ids=([7, 8],), overwrite=True) |
| 298 | block_tables.apply_staged_writes() |
| 299 | query_start_loc = torch.tensor([0, 3], dtype=torch.int32, device=device) |
| 300 | positions = torch.tensor([1, 2, 3], dtype=torch.int64, device=device) |
| 301 | |
| 302 | real = block_tables.compute_slot_mappings( |
| 303 | torch.tensor([0], dtype=torch.int32, device=device), |
| 304 | query_start_loc, |
| 305 | positions, |
| 306 | num_tokens_padded=3, |
| 307 | ) |
| 308 | assert real[0].tolist() == [7 * 4 + 1, 7 * 4 + 2, 7 * 4 + 3] |
| 309 | |
| 310 | dummy = block_tables.compute_slot_mappings( |
| 311 | torch.tensor([-1], dtype=torch.int32, device=device), |
| 312 | query_start_loc, |
| 313 | positions, |
| 314 | num_tokens_padded=3, |
| 315 | ) |
| 316 | assert dummy[0].tolist() == [PAD_SLOT_ID] * 3 |