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 f20a96f5cc31..e25905ed367a 100644 --- a/tests/v1/core/test_single_type_kv_cache_manager.py +++ b/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( diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index c8b9c54435fc..24c1a997d463 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/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