[Bugfix] Generalize KV block zeroing to `AttentionSpec` (#51749)
Signed-off-by: mgoin <mgoin64@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Michael Goin committed
Aug 11, 2026 at 12:13 UTC
e3fe212eaf0dc06da03d4df615f7d3d1e8bd5f20
4 files changed
+134
-10
tests/v1/core/test_single_type_kv_cache_manager.py
+68
-2
@@ -26,7 +26,12 @@ from vllm.v1.kv_cache_interface import (
26
pytestmark = pytest.mark.cpu_test
27
28
29
-def get_sliding_window_manager(sliding_window_spec, block_pool, enable_caching=True):
29
+def get_sliding_window_manager(
30
+ sliding_window_spec,
31
+ block_pool,
32
+ enable_caching=True,
33
+ needs_kv_cache_zeroing=False,
34
+):
35
# Tests don't exercise admission gating; pass a large cap that is a no-op.
36
return SlidingWindowManager(
37
sliding_window_spec,
@@ -34,12 +39,16 @@ def get_sliding_window_manager(sliding_window_spec, block_pool, enable_caching=T
39
enable_caching=enable_caching,
40
kv_cache_group_id=0,
41
scheduler_block_size=sliding_window_spec.block_size,
42
+ needs_kv_cache_zeroing=needs_kv_cache_zeroing,
43
max_admission_blocks_per_request=10**9,
44
)
45
46
47
def get_chunked_local_attention_manager(
42
- chunked_local_attention_spec, block_pool, enable_caching=True
48
+ chunked_local_attention_spec,
49
+ block_pool,
50
+ enable_caching=True,
51
+ needs_kv_cache_zeroing=False,
52
):
53
return ChunkedLocalAttentionManager(
54
chunked_local_attention_spec,
@@ -47,10 +56,67 @@ def get_chunked_local_attention_manager(
56
enable_caching=enable_caching,
57
kv_cache_group_id=0,
58
scheduler_block_size=chunked_local_attention_spec.block_size,
59
+ needs_kv_cache_zeroing=needs_kv_cache_zeroing,
60
max_admission_blocks_per_request=10**9,
61
)
62
63
64
+def test_sliding_window_records_new_blocks_for_zeroing():
65
+ block_size = 2
66
+ spec = SlidingWindowSpec(
67
+ block_size=block_size,
68
+ num_kv_heads=1,
69
+ head_size=1,
70
+ dtype=torch.float32,
71
+ sliding_window=4,
72
+ )
73
+ block_pool = BlockPool(
74
+ num_gpu_blocks=10, enable_caching=False, hash_block_size=block_size
75
+ )
76
+ manager = get_sliding_window_manager(
77
+ spec,
78
+ block_pool,
79
+ enable_caching=False,
80
+ needs_kv_cache_zeroing=True,
81
+ )
82
+
83
+ blocks = manager.allocate_new_blocks(
84
+ "request", num_tokens=4, num_tokens_main_model=4
85
+ )
86
+
87
+ assert manager.records_new_block_ids
88
+ assert manager.take_new_block_ids() == [block.block_id for block in blocks]
89
+ assert manager.take_new_block_ids() == []
90
+
91
+
92
+def test_chunked_local_attention_records_new_blocks_for_zeroing():
93
+ block_size = 2
94
+ spec = ChunkedLocalAttentionSpec(
95
+ block_size=block_size,
96
+ num_kv_heads=1,
97
+ head_size=1,
98
+ dtype=torch.float32,
99
+ attention_chunk_size=4,
100
+ )
101
+ block_pool = BlockPool(
102
+ num_gpu_blocks=10, enable_caching=False, hash_block_size=block_size
103
+ )
104
+ manager = get_chunked_local_attention_manager(
105
+ spec,
106
+ block_pool,
107
+ enable_caching=False,
108
+ needs_kv_cache_zeroing=True,
109
+ )
110
+
111
+ blocks = manager.allocate_new_blocks(
112
+ "request", num_tokens=4, num_tokens_main_model=4
113
+ )
114
+
115
+ assert manager.records_new_block_ids
116
+ assert manager.take_new_block_ids() == [block.block_id for block in blocks]
117
+ assert manager.take_new_block_ids() == []
118
+
119
+
120
def test_chunked_local_attention_possible_cached_prefix():
121
block_size = 2
122
chunked_local_attention_spec = ChunkedLocalAttentionSpec(
tests/v1/worker/test_kv_block_zeroer.py
+62
-1
@@ -1,10 +1,71 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from types import SimpleNamespace
5
+
6
import pytest
7
import torch
8
7
-from vllm.v1.worker.utils import KVBlockZeroer, _zero_kv_blocks_kernel
9
+from vllm.v1.kv_cache_interface import (
10
+ ChunkedLocalAttentionSpec,
11
+ SlidingWindowSpec,
12
+)
13
+from vllm.v1.worker.utils import (
14
+ AttentionGroup,
15
+ KVBlockZeroer,
16
+ _zero_kv_blocks_kernel,
17
+)
18
+
19
+
20
+class _BlockFirstBackend:
21
+ @staticmethod
22
+ def get_kv_cache_block_dim(*args, **kwargs):
23
+ return 0
24
+
25
+
26
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
27
+@pytest.mark.parametrize(
28
+ "spec",
29
+ [
30
+ SlidingWindowSpec(
31
+ block_size=2,
32
+ num_kv_heads=1,
33
+ head_size=1,
34
+ dtype=torch.uint8,
35
+ sliding_window=4,
36
+ ),
37
+ ChunkedLocalAttentionSpec(
38
+ block_size=2,
39
+ num_kv_heads=1,
40
+ head_size=1,
41
+ dtype=torch.uint8,
42
+ attention_chunk_size=4,
43
+ ),
44
+ ],
45
+ ids=["sliding-window", "chunked-local"],
46
+)
47
+def test_attention_blocks_are_zeroed(spec):
48
+ device = torch.device("cuda")
49
+ storage = torch.ones((4, 1, 2, 2), dtype=torch.uint8, device=device)
50
+ layer_name = "draft.self_attn"
51
+ zeroer = KVBlockZeroer(
52
+ device,
53
+ attn_groups_iter=[
54
+ AttentionGroup(_BlockFirstBackend, [layer_name], spec, 0) # type: ignore[arg-type]
55
+ ],
56
+ kernel_block_sizes=[2],
57
+ cache_dtype="fp8",
58
+ static_forward_context={
59
+ layer_name: SimpleNamespace(kv_cache=storage),
60
+ },
61
+ )
62
+
63
+ zeroer.zero_block_ids([1])
64
+ torch.accelerator.synchronize()
65
+
66
+ expected = torch.ones_like(storage)
67
+ expected[1] = 0
68
+ assert torch.equal(storage, expected)
69
70
71
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
vllm/v1/core/single_type_kv_cache_manager.py
+3
-5
@@ -16,6 +16,7 @@ from vllm.v1.core.kv_cache_utils import (
16
resolve_block_hashes,
17
)
18
from vllm.v1.kv_cache_interface import (
19
+ AttentionSpec,
20
ChunkedLocalAttentionSpec,
21
CrossAttentionSpec,
22
FullAttentionSpec,
@@ -83,11 +84,8 @@ class SingleTypeKVCacheManager(ABC):
84
self._max_admission_blocks_per_request = max_admission_blocks_per_request
85
# Record newly allocated block ids only when worker-side zeroing will
86
# consume them and this manager holds a spec type that gets zeroed.
86
- self._record_new_block_ids = needs_kv_cache_zeroing and type(kv_cache_spec) in (
87
- FullAttentionSpec,
88
- TQFullAttentionSpec,
89
- MLAAttentionSpec,
90
- HiddenStateCacheSpec,
87
+ self._record_new_block_ids = needs_kv_cache_zeroing and isinstance(
88
+ kv_cache_spec, AttentionSpec
89
)
90
self.new_block_ids: list[int] = []
91
vllm/v1/worker/utils.py
+1
-2
@@ -29,7 +29,6 @@ from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
29
from vllm.v1.kv_cache_interface import (
30
AttentionSpec,
31
EncoderOnlyAttentionSpec,
32
- FullAttentionSpec,
32
KVCacheConfig,
33
KVCacheGroupSpec,
34
KVCacheSpec,
@@ -147,7 +146,7 @@ class KVBlockZeroer:
146
147
for group in attn_groups_iter:
148
spec = group.kv_cache_spec
150
- if not isinstance(spec, FullAttentionSpec):
149
+ if not isinstance(spec, AttentionSpec):
150
continue
151
if group.kv_cache_group_id >= len(kernel_block_sizes):
152
continue