[Bugfix][KV Cache][GLM-5.3-Flash] Disable slot mapping kernel for the kpool tail buffer (#57317)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Simon Danielsson committed
Sep 18, 2026 at 04:42 UTC
70df48dc3d01fe1bb3206f7d8a832c7e8de348e8
3 files changed
+126
tests/v1/attention/test_kpool_tail_slot_mapping.py
+24
@@ -32,6 +32,7 @@ from vllm.v1.attention.backends.mla.indexer import (
32
)
33
from vllm.v1.kv_cache_interface import KpoolTailSpec, compute_layout_strides
34
from vllm.v1.kv_cache_layout import KVCacheLayout
35
+from vllm.v1.worker.block_table import get_block_table_width
36
37
KPOOL = 4
38
@@ -55,6 +56,29 @@ def test_tail_backend_layout_matches_kernel_pointer_arithmetic():
56
assert content_stride == 1
57
58
59
+def test_tail_spec_opts_out_of_generic_slot_mapping():
60
+ """The tail row is one block wide (padded to the block-table alignment), so
61
+ the generic kernel's ``pos // kpool`` column index runs off the end of the
62
+ allocation for long prompts. The spec must opt out of it entirely."""
63
+ spec = KpoolTailSpec(
64
+ block_size=KPOOL,
65
+ num_kv_heads=2,
66
+ head_size=128,
67
+ head_size_v=0,
68
+ dtype=torch.bfloat16,
69
+ sliding_window=KPOOL,
70
+ )
71
+ max_len = 1 << 20
72
+ width = get_block_table_width(
73
+ spec.max_num_blocks_per_req(None, max_len),
74
+ spec.block_size,
75
+ token_alignment=spec.block_table_token_alignment,
76
+ )
77
+
78
+ assert width * KPOOL < max_len
79
+ assert spec.uses_slot_mapping is False
80
+
81
+
82
def make_tail_block_table(own_blocks, width=64):
83
"""Tail-group block table as BlockTables produces it: column 0 holds the
84
request's single KpoolTailManager block, the remaining columns are never
tests/v1/worker/test_gpu_kpool_tail_slot_mapping.py
new
+98
@@ -0,0 +1,98 @@
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
+ ]
vllm/v1/kv_cache_interface.py
+4
@@ -987,6 +987,10 @@ class KpoolTailSpec(SlidingWindowSpec):
987
def prefix_cacheable(self) -> bool:
988
return False
989
990
+ @property
991
+ def uses_slot_mapping(self) -> bool:
992
+ return False
993
+
994
995
@dataclass(frozen=True)
996
class MambaSpec(KVCacheSpec):