[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>

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(

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")

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] = []
 

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