From 70df48dc3d01fe1bb3206f7d8a832c7e8de348e8 Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Fri, 18 Sep 2026 04:42:54 +0200 Subject: [PATCH] [Bugfix][KV Cache][GLM-5.3-Flash] Disable slot mapping kernel for the kpool tail buffer (#57317) Signed-off-by: simondanielsson --- .../attention/test_kpool_tail_slot_mapping.py | 24 +++++ .../test_gpu_kpool_tail_slot_mapping.py | 98 +++++++++++++++++++ vllm/v1/kv_cache_interface.py | 4 + 3 files changed, 126 insertions(+) create mode 100644 tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py diff --git a/tests/v1/attention/test_kpool_tail_slot_mapping.py b/tests/v1/attention/test_kpool_tail_slot_mapping.py index 0966f9b20a..cb39e54037 100644 --- a/tests/v1/attention/test_kpool_tail_slot_mapping.py +++ b/tests/v1/attention/test_kpool_tail_slot_mapping.py @@ -32,6 +32,7 @@ from vllm.v1.attention.backends.mla.indexer import ( ) from vllm.v1.kv_cache_interface import KpoolTailSpec, compute_layout_strides from vllm.v1.kv_cache_layout import KVCacheLayout +from vllm.v1.worker.block_table import get_block_table_width KPOOL = 4 @@ -55,6 +56,29 @@ def test_tail_backend_layout_matches_kernel_pointer_arithmetic(): assert content_stride == 1 +def test_tail_spec_opts_out_of_generic_slot_mapping(): + """The tail row is one block wide (padded to the block-table alignment), so + the generic kernel's ``pos // kpool`` column index runs off the end of the + allocation for long prompts. The spec must opt out of it entirely.""" + spec = KpoolTailSpec( + block_size=KPOOL, + num_kv_heads=2, + head_size=128, + head_size_v=0, + dtype=torch.bfloat16, + sliding_window=KPOOL, + ) + max_len = 1 << 20 + width = get_block_table_width( + spec.max_num_blocks_per_req(None, max_len), + spec.block_size, + token_alignment=spec.block_table_token_alignment, + ) + + assert width * KPOOL < max_len + assert spec.uses_slot_mapping is False + + def make_tail_block_table(own_blocks, width=64): """Tail-group block table as BlockTables produces it: column 0 holds the request's single KpoolTailManager block, the remaining columns are never diff --git a/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py b/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py new file mode 100644 index 0000000000..8a7445f7d4 --- /dev/null +++ b/tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py @@ -0,0 +1,98 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GPU regression test for the kpool tail group's slot mapping. + +The kpool tail cache holds one circular block per request, so its block table +row is orders of magnitude shorter than the sequence it serves. +`_compute_slot_mappings_kernel` indexes block tables by +`position // kernel_block_size` with an effectively unmasked load, so +this leads to a read of neighboring requests' rows (and for the later requests, outside +the buffer entirely). +""" + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.attention.backends.utils import PAD_SLOT_ID +from vllm.v1.kv_cache_interface import KpoolTailSpec +from vllm.v1.worker.block_table import get_block_table_width +from vllm.v1.worker.gpu.block_table import BlockTables + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="requires CUDA or ROCm", +) + +KPOOL = 4 +MAX_MODEL_LEN = 1 << 20 +MLA_BLOCK_SIZE = 640 +MAX_NUM_REQS = 4 +POISON_BLOCK_ID = 1 << 20 + + +def make_tail_spec(): + return KpoolTailSpec( + block_size=KPOOL, + num_kv_heads=2, + head_size=128, + head_size_v=0, + dtype=torch.bfloat16, + sliding_window=KPOOL, + ) + + +def test_kpool_tail_group_never_position_indexes_its_block_table(): + device = torch.device("cuda") + spec = make_tail_spec() + tail_width = get_block_table_width( + spec.max_num_blocks_per_req(None, MAX_MODEL_LEN), + spec.block_size, + token_alignment=spec.block_table_token_alignment, + ) + mla_width = get_block_table_width(MAX_MODEL_LEN // MLA_BLOCK_SIZE, MLA_BLOCK_SIZE) + + # Positions chosen against the tail row: `in_row` still lands inside + # request 0's own row, `next_row` lands on request 1's poisoned row, and + # `past_table` is past the whole [MAX_NUM_REQS, tail_width] allocation. + in_row = KPOOL - 1 + next_row = tail_width * KPOOL + past_table = MAX_NUM_REQS * tail_width * KPOOL + KPOOL + assert past_table < MAX_MODEL_LEN + positions = [0, in_row, next_row, past_table] + + block_tables = BlockTables( + block_sizes=[MLA_BLOCK_SIZE, spec.block_size], + max_num_reqs=MAX_NUM_REQS, + max_num_batched_tokens=len(positions), + max_num_blocks_per_group=[mla_width, tail_width], + device=device, + kernel_block_sizes=[MLA_BLOCK_SIZE, spec.block_size], + slot_mapping_enabled=[True, spec.uses_slot_mapping], + ) + mla_blocks = list(range(1, mla_width + 1)) + block_tables.append_block_ids(0, (mla_blocks, [7]), overwrite=True) + block_tables.append_block_ids(1, ([2], [POISON_BLOCK_ID]), overwrite=True) + block_tables.apply_staged_writes() + + idx_mapping = torch.zeros(1, dtype=torch.int32, device=device) + query_start_loc = torch.tensor( + [0, len(positions)], dtype=torch.int32, device=device + ) + slot_mappings = block_tables.compute_slot_mappings( + idx_mapping, + query_start_loc, + torch.tensor(positions, dtype=torch.int64, device=device), + num_tokens_padded=len(positions), + ) + torch.accelerator.synchronize() + + tail_slots = slot_mappings[1].tolist() + assert tail_slots == [PAD_SLOT_ID] * len(positions) + assert POISON_BLOCK_ID * KPOOL not in tail_slots + + mla_slots = slot_mappings[0].tolist() + assert mla_slots == [ + mla_blocks[p // MLA_BLOCK_SIZE] * MLA_BLOCK_SIZE + p % MLA_BLOCK_SIZE + for p in positions + ] diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 17771d9ae6..6595be14d2 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -987,6 +987,10 @@ class KpoolTailSpec(SlidingWindowSpec): def prefix_cacheable(self) -> bool: return False + @property + def uses_slot_mapping(self) -> bool: + return False + @dataclass(frozen=True) class MambaSpec(KVCacheSpec):