[Bugfix][Core] Accept needs_kv_cache_zeroing in SinkFullAttentionManager

tests/v1/core/test_single_type_kv_cache_manager.py

@@ -18,6 +18,7 @@
     FullAttentionManager,
     MambaManager,
     RSWAManager,
+    SinkFullAttentionManager,
     SlidingWindowManager,
 )
 from vllm.v1.kv_cache_interface import (
@@ -26,6 +27,7 @@
     FullAttentionSpec,
     MambaSpec,
     RSWASpec,
+    SinkFullAttentionSpec,
     SlidingWindowSpec,
 )
 
@@ -360,6 +362,35 @@ def test_sliding_window_records_new_blocks_for_zeroing():
     assert manager.take_new_block_ids() == []
 
 
+def test_sink_full_attention_records_new_blocks_for_zeroing():
+    block_size = 2
+    spec = SinkFullAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+        sink_len=block_size,
+    )
+    block_pool = BlockPool(
+        num_gpu_blocks=10, enable_caching=False, hash_block_size=block_size
+    )
+    manager = SinkFullAttentionManager(
+        spec,
+        block_pool=block_pool,
+        enable_caching=False,
+        kv_cache_group_id=0,
+        scheduler_block_size=block_size,
+        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]
+
+
 def test_chunked_local_attention_records_new_blocks_for_zeroing():
     block_size = 2
     spec = ChunkedLocalAttentionSpec(

vllm/v1/core/single_type_kv_cache_manager.py

@@ -2213,25 +2213,8 @@ def find_longest_cache_hit(
 
 
 class SinkFullAttentionManager(FullAttentionManager):
-    def __init__(
-        self,
-        kv_cache_spec: SinkFullAttentionSpec,
-        block_pool: BlockPool,
-        enable_caching: bool,
-        kv_cache_group_id: int,
-        scheduler_block_size: int,
-        dcp_world_size: int = 1,
-        pcp_world_size: int = 1,
-    ):
-        super().__init__(
-            kv_cache_spec=kv_cache_spec,
-            block_pool=block_pool,
-            enable_caching=enable_caching,
-            kv_cache_group_id=kv_cache_group_id,
-            scheduler_block_size=scheduler_block_size,
-            dcp_world_size=dcp_world_size,
-            pcp_world_size=pcp_world_size,
-        )
+    def __init__(self, kv_cache_spec: SinkFullAttentionSpec, **kwargs) -> None:
+        super().__init__(kv_cache_spec, **kwargs)
         sink_len = kv_cache_spec.sink_len
         assert sink_len is not None and sink_len > 0 and sink_len % self.block_size == 0