From e3fe212eaf0dc06da03d4df615f7d3d1e8bd5f20 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Tue, 11 Aug 2026 12:13:09 -0400 Subject: [PATCH] [Bugfix] Generalize KV block zeroing to `AttentionSpec` (#51749) Signed-off-by: mgoin Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../core/test_single_type_kv_cache_manager.py | 70 ++++++++++++++++++- tests/v1/worker/test_kv_block_zeroer.py | 63 ++++++++++++++++- vllm/v1/core/single_type_kv_cache_manager.py | 8 +-- vllm/v1/worker/utils.py | 3 +- 4 files changed, 134 insertions(+), 10 deletions(-) diff --git a/tests/v1/core/test_single_type_kv_cache_manager.py b/tests/v1/core/test_single_type_kv_cache_manager.py index 7300d45fbf..e20964ef76 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/tests/v1/core/test_single_type_kv_cache_manager.py @@ -26,7 +26,12 @@ from vllm.v1.kv_cache_interface import ( pytestmark = pytest.mark.cpu_test -def get_sliding_window_manager(sliding_window_spec, block_pool, enable_caching=True): +def get_sliding_window_manager( + sliding_window_spec, + block_pool, + enable_caching=True, + needs_kv_cache_zeroing=False, +): # Tests don't exercise admission gating; pass a large cap that is a no-op. return SlidingWindowManager( sliding_window_spec, @@ -34,12 +39,16 @@ def get_sliding_window_manager(sliding_window_spec, block_pool, enable_caching=T enable_caching=enable_caching, kv_cache_group_id=0, scheduler_block_size=sliding_window_spec.block_size, + needs_kv_cache_zeroing=needs_kv_cache_zeroing, max_admission_blocks_per_request=10**9, ) def get_chunked_local_attention_manager( - chunked_local_attention_spec, block_pool, enable_caching=True + chunked_local_attention_spec, + block_pool, + enable_caching=True, + needs_kv_cache_zeroing=False, ): return ChunkedLocalAttentionManager( chunked_local_attention_spec, @@ -47,10 +56,67 @@ def get_chunked_local_attention_manager( enable_caching=enable_caching, kv_cache_group_id=0, scheduler_block_size=chunked_local_attention_spec.block_size, + needs_kv_cache_zeroing=needs_kv_cache_zeroing, max_admission_blocks_per_request=10**9, ) +def test_sliding_window_records_new_blocks_for_zeroing(): + block_size = 2 + spec = SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=4, + ) + block_pool = BlockPool( + num_gpu_blocks=10, enable_caching=False, hash_block_size=block_size + ) + manager = get_sliding_window_manager( + spec, + block_pool, + enable_caching=False, + needs_kv_cache_zeroing=True, + ) + + blocks = manager.allocate_new_blocks( + "request", num_tokens=4, num_tokens_main_model=4 + ) + + assert manager.records_new_block_ids + assert manager.take_new_block_ids() == [block.block_id for block in blocks] + assert manager.take_new_block_ids() == [] + + +def test_chunked_local_attention_records_new_blocks_for_zeroing(): + block_size = 2 + spec = ChunkedLocalAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + attention_chunk_size=4, + ) + block_pool = BlockPool( + num_gpu_blocks=10, enable_caching=False, hash_block_size=block_size + ) + manager = get_chunked_local_attention_manager( + spec, + block_pool, + enable_caching=False, + needs_kv_cache_zeroing=True, + ) + + blocks = manager.allocate_new_blocks( + "request", num_tokens=4, num_tokens_main_model=4 + ) + + assert manager.records_new_block_ids + assert manager.take_new_block_ids() == [block.block_id for block in blocks] + assert manager.take_new_block_ids() == [] + + def test_chunked_local_attention_possible_cached_prefix(): block_size = 2 chunked_local_attention_spec = ChunkedLocalAttentionSpec( diff --git a/tests/v1/worker/test_kv_block_zeroer.py b/tests/v1/worker/test_kv_block_zeroer.py index 73d079ea24..58d1a2d3e7 100644 --- a/tests/v1/worker/test_kv_block_zeroer.py +++ b/tests/v1/worker/test_kv_block_zeroer.py @@ -1,10 +1,71 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch -from vllm.v1.worker.utils import KVBlockZeroer, _zero_kv_blocks_kernel +from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, + SlidingWindowSpec, +) +from vllm.v1.worker.utils import ( + AttentionGroup, + KVBlockZeroer, + _zero_kv_blocks_kernel, +) + + +class _BlockFirstBackend: + @staticmethod + def get_kv_cache_block_dim(*args, **kwargs): + return 0 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "spec", + [ + SlidingWindowSpec( + block_size=2, + num_kv_heads=1, + head_size=1, + dtype=torch.uint8, + sliding_window=4, + ), + ChunkedLocalAttentionSpec( + block_size=2, + num_kv_heads=1, + head_size=1, + dtype=torch.uint8, + attention_chunk_size=4, + ), + ], + ids=["sliding-window", "chunked-local"], +) +def test_attention_blocks_are_zeroed(spec): + device = torch.device("cuda") + storage = torch.ones((4, 1, 2, 2), dtype=torch.uint8, device=device) + layer_name = "draft.self_attn" + zeroer = KVBlockZeroer( + device, + attn_groups_iter=[ + AttentionGroup(_BlockFirstBackend, [layer_name], spec, 0) # type: ignore[arg-type] + ], + kernel_block_sizes=[2], + cache_dtype="fp8", + static_forward_context={ + layer_name: SimpleNamespace(kv_cache=storage), + }, + ) + + zeroer.zero_block_ids([1]) + torch.accelerator.synchronize() + + expected = torch.ones_like(storage) + expected[1] = 0 + assert torch.equal(storage, expected) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index dbf223ec9b..f5f854a5df 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -16,6 +16,7 @@ from vllm.v1.core.kv_cache_utils import ( resolve_block_hashes, ) from vllm.v1.kv_cache_interface import ( + AttentionSpec, ChunkedLocalAttentionSpec, CrossAttentionSpec, FullAttentionSpec, @@ -83,11 +84,8 @@ class SingleTypeKVCacheManager(ABC): self._max_admission_blocks_per_request = max_admission_blocks_per_request # Record newly allocated block ids only when worker-side zeroing will # consume them and this manager holds a spec type that gets zeroed. - self._record_new_block_ids = needs_kv_cache_zeroing and type(kv_cache_spec) in ( - FullAttentionSpec, - TQFullAttentionSpec, - MLAAttentionSpec, - HiddenStateCacheSpec, + self._record_new_block_ids = needs_kv_cache_zeroing and isinstance( + kv_cache_spec, AttentionSpec ) self.new_block_ids: list[int] = [] diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index b10a236ae0..3c8630a5e3 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -29,7 +29,6 @@ from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy from vllm.v1.kv_cache_interface import ( AttentionSpec, EncoderOnlyAttentionSpec, - FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, KVCacheSpec, @@ -147,7 +146,7 @@ class KVBlockZeroer: for group in attn_groups_iter: spec = group.kv_cache_spec - if not isinstance(spec, FullAttentionSpec): + if not isinstance(spec, AttentionSpec): continue if group.kv_cache_group_id >= len(kernel_block_sizes): continue