| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """GPU regression test for the kpool tail group's slot mapping. |
| 4 | |
| 5 | The kpool tail cache holds one circular block per request, so its block table |
| 6 | row is orders of magnitude shorter than the sequence it serves. |
| 7 | `_compute_slot_mappings_kernel` indexes block tables by |
| 8 | `position // kernel_block_size` with an effectively unmasked load, so |
| 9 | this leads to a read of neighboring requests' rows (and for the later requests, outside |
| 10 | the buffer entirely). |
| 11 | """ |
| 12 | |
| 13 | import pytest |
| 14 | import torch |
| 15 | |
| 16 | from vllm.platforms import current_platform |
| 17 | from vllm.v1.attention.backends.utils import PAD_SLOT_ID |
| 18 | from vllm.v1.kv_cache_interface import KpoolTailSpec |
| 19 | from vllm.v1.worker.block_table import get_block_table_width |
| 20 | from vllm.v1.worker.gpu.block_table import BlockTables |
| 21 | |
| 22 | pytestmark = pytest.mark.skipif( |
| 23 | not current_platform.is_cuda_alike(), |
| 24 | reason="requires CUDA or ROCm", |
| 25 | ) |
| 26 | |
| 27 | KPOOL = 4 |
| 28 | MAX_MODEL_LEN = 1 << 20 |
| 29 | MLA_BLOCK_SIZE = 640 |
| 30 | MAX_NUM_REQS = 4 |
| 31 | POISON_BLOCK_ID = 1 << 20 |
| 32 | |
| 33 | |
| 34 | def make_tail_spec(): |
| 35 | return KpoolTailSpec( |
| 36 | block_size=KPOOL, |
| 37 | num_kv_heads=2, |
| 38 | head_size=128, |
| 39 | head_size_v=0, |
| 40 | dtype=torch.bfloat16, |
| 41 | sliding_window=KPOOL, |
| 42 | ) |
| 43 | |
| 44 | |
| 45 | def test_kpool_tail_group_never_position_indexes_its_block_table(): |
| 46 | device = torch.device("cuda") |
| 47 | spec = make_tail_spec() |
| 48 | tail_width = get_block_table_width( |
| 49 | spec.max_num_blocks_per_req(None, MAX_MODEL_LEN), |
| 50 | spec.block_size, |
| 51 | token_alignment=spec.block_table_token_alignment, |
| 52 | ) |
| 53 | mla_width = get_block_table_width(MAX_MODEL_LEN // MLA_BLOCK_SIZE, MLA_BLOCK_SIZE) |
| 54 | |
| 55 | # Positions chosen against the tail row: `in_row` still lands inside |
| 56 | # request 0's own row, `next_row` lands on request 1's poisoned row, and |
| 57 | # `past_table` is past the whole [MAX_NUM_REQS, tail_width] allocation. |
| 58 | in_row = KPOOL - 1 |
| 59 | next_row = tail_width * KPOOL |
| 60 | past_table = MAX_NUM_REQS * tail_width * KPOOL + KPOOL |
| 61 | assert past_table < MAX_MODEL_LEN |
| 62 | positions = [0, in_row, next_row, past_table] |
| 63 | |
| 64 | block_tables = BlockTables( |
| 65 | block_sizes=[MLA_BLOCK_SIZE, spec.block_size], |
| 66 | max_num_reqs=MAX_NUM_REQS, |
| 67 | max_num_batched_tokens=len(positions), |
| 68 | max_num_blocks_per_group=[mla_width, tail_width], |
| 69 | device=device, |
| 70 | kernel_block_sizes=[MLA_BLOCK_SIZE, spec.block_size], |
| 71 | slot_mapping_enabled=[True, spec.uses_slot_mapping], |
| 72 | ) |
| 73 | mla_blocks = list(range(1, mla_width + 1)) |
| 74 | block_tables.append_block_ids(0, (mla_blocks, [7]), overwrite=True) |
| 75 | block_tables.append_block_ids(1, ([2], [POISON_BLOCK_ID]), overwrite=True) |
| 76 | block_tables.apply_staged_writes() |
| 77 | |
| 78 | idx_mapping = torch.zeros(1, dtype=torch.int32, device=device) |
| 79 | query_start_loc = torch.tensor( |
| 80 | [0, len(positions)], dtype=torch.int32, device=device |
| 81 | ) |
| 82 | slot_mappings = block_tables.compute_slot_mappings( |
| 83 | idx_mapping, |
| 84 | query_start_loc, |
| 85 | torch.tensor(positions, dtype=torch.int64, device=device), |
| 86 | num_tokens_padded=len(positions), |
| 87 | ) |
| 88 | torch.accelerator.synchronize() |
| 89 | |
| 90 | tail_slots = slot_mappings[1].tolist() |
| 91 | assert tail_slots == [PAD_SLOT_ID] * len(positions) |
| 92 | assert POISON_BLOCK_ID * KPOOL not in tail_slots |
| 93 | |
| 94 | mla_slots = slot_mappings[0].tolist() |
| 95 | assert mla_slots == [ |
| 96 | mla_blocks[p // MLA_BLOCK_SIZE] * MLA_BLOCK_SIZE + p % MLA_BLOCK_SIZE |
| 97 | for p in positions |
| 98 | ] |