[wip] support local chunked attentin for hybrid

tests/v1/core/test_specialized_manager.py

@@ -1,12 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 
 import torch
-
+import random
 from vllm.v1.core.block_pool import BlockPool
+
 from vllm.v1.core.kv_cache_utils import (BlockHashType, KVCacheBlock,
                                          KVCacheBlockBundle)
-from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager
-from vllm.v1.kv_cache_interface import SlidingWindowSpec
+from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager, ChunkedLocalAttentionManager
+from vllm.v1.kv_cache_interface import SlidingWindowSpec, ChunkedLocalAttentionSpec
 
 
 def get_sliding_window_manager(sliding_window_spec, block_pool):
@@ -17,6 +18,76 @@
                                 caching_hash_fn=lambda x: x,
                                 manager_id=0)
 
+def get_chunked_local_attention_manager(chunked_local_attention_spec, block_pool):
+    return ChunkedLocalAttentionManager(chunked_local_attention_spec,
+                                block_pool,
+                                caching_hash_fn=lambda x: x,
+                                kv_cache_group_id=0)
+
+
+def test_chunked_local_attention_possible_cached_prefix():
+    block_size = 2
+    chunked_local_attention_spec = ChunkedLocalAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+        attention_chunk_size=4,
+        use_mla=False,
+    )
+
+    block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
+    manager = get_chunked_local_attention_manager(chunked_local_attention_spec, block_pool)
+
+    def run_one_case(block_is_cached, tail_token, expect_length):
+        block_hash_list = [
+            BlockHash(i, ()) for i in range(len(block_is_cached))
+        ]
+
+        block_pool.cached_block_hash_to_block.clear()
+
+        # Mock the block pool with the cached blocks
+        for i, (block_hash,
+                is_cached) in enumerate(zip(block_hash_list, block_is_cached)):
+            if is_cached:
+                block_pool.cached_block_hash_to_block[BlockHashWithGroupId(
+                    block_hash, 0)] = {
+                        i: block_pool.blocks[i + 10],
+                    }
+
+        computed_blocks = manager.find_longest_cache_hit(
+            block_hashes=block_hash_list,
+            max_length=len(block_hash_list) * block_size + tail_token,
+            kv_cache_group_ids=[0],
+            block_pool=block_pool,
+            kv_cache_spec=chunked_local_attention_spec,
+            use_eagle=False)[0]
+        assert len(computed_blocks) == expect_length
+
+        assert all(block == block_pool.null_block
+                   for block in computed_blocks[:(expect_length-1) // 2])
+
+    run_one_case([True], 0,  1)
+    run_one_case([True], 1,  1)
+    run_one_case([True, False], 0, 1)
+    run_one_case([True, False], 1, 2)
+    run_one_case([True, True], 0, 2)
+    run_one_case([True, True], 1, 2)
+    run_one_case([True, True, False], 0, 2)
+    run_one_case([True, True, False], 1, 2)
+    run_one_case([True, True, True], 0, 3)
+    run_one_case([True, True, True], 1, 3)
+    run_one_case([True, True, True, False], 0, 3)
+    run_one_case([True, True, True, False], 1, 4)
+    run_one_case([random.choice([True, False])] * 8 + [True, True], 1, 10)
+    run_one_case([random.choice([True, False])] * 8 + [True, False], 0, 9)
+    run_one_case([random.choice([True, False])] * 8 + [True, False], 1, 10)
+    run_one_case([random.choice([True, False])] * 8 + [False, True], 0, 8)
+    run_one_case([random.choice([True, False])] * 8 + [False, True], 1, 10)
+    run_one_case([random.choice([True, False])] * 8 + [False, False], 0, 8)
+    run_one_case([random.choice([True, False])] * 8 + [False, False], 1, 10)
+
+
 
 def test_sliding_window_possible_cached_prefix():
     block_size = 2
@@ -66,22 +137,77 @@
     run_one_case([True], 1)
     run_one_case([True, False], 1)
     run_one_case([True, True], 2)
-    run_one_case([True, True, False], 2)
+    run_one_case([True, True, False], 0)
     run_one_case([True, True, True], 3)
     run_one_case([True, True, True, False], 3)
     run_one_case([
         True, True, False, True, False, False, True, True, False, True, True,
         True
     ], 12)
     run_one_case([
-        True, True, False, True, False, False, True, True, False, False, False
+        True, True, False, True, False, False, True, True, False, False, True
     ], 8)
     run_one_case([
         True, True, False, True, False, False, True, True, False, False, False,
         True
     ], 8)
 
 
+def test_chunked_local_attention_remove_skipped_blocks():
+    attention_spec = ChunkedLocalAttentionSpec(
+        block_size=2,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+        attention_chunk_size=4,
+        use_mla=False,
+    )
+
+    block_pool = BlockPool(num_gpu_blocks=2000, enable_caching=True)
+
+    manager = get_chunked_local_attention_manager(attention_spec, block_pool)
+
+    null_block_id = block_pool.null_block.block_id
+
+    def id_to_block_table(ids) -> list[KVCacheBlock]:
+        return [
+            KVCacheBlock(id_)
+            if id_ != null_block_id else block_pool.null_block for id_ in ids
+        ]
+
+    def assert_block_id(block_table: list[KVCacheBlock], ids: list[int]):
+        for block, id_ in zip(block_table, ids):
+            if id_ == null_block_id:
+                assert block == block_pool.null_block
+            else:
+                assert block.block_id == id_
+
+    original_block_ids = [
+        1000, 1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008, 1009, 1010
+    ]
+    block_table = id_to_block_table(original_block_ids)
+    manager.req_to_blocks["test"] = block_table
+
+    manager.remove_skipped_blocks("test", 0)
+    assert_block_id(block_table, original_block_ids)
+
+    # 4 tokens are computed. no token is out of the local attention window.
+    manager.remove_skipped_blocks("test", 4)
+    assert_block_id(block_table, original_block_ids)
+
+    # 5 tokens are computed. token 0 is out of the local attention window. no block can be removed.
+    manager.remove_skipped_blocks("test", 5)
+    assert_block_id(block_table, [null_block_id])
+
+    # 6 tokens are computed. token 4 - 5 are in local attention window, token 0 - 3 are out, 2 blocks can be removed.
+    manager.remove_skipped_blocks("test", 6)
+    assert_block_id(block_table, [null_block_id] * 2 + original_block_ids[2:])
+    # 11 tokens are computed. token 8 - 11 are in local attention window, token 0-7 are out, 4 block can be removed.
+    manager.remove_skipped_blocks("test", 11)
+    assert_block_id(block_table, [null_block_id] * 4 + original_block_ids[4:])
+
+
+
 def test_sliding_window_remove_skipped_blocks():
     sliding_window_spec = SlidingWindowSpec(
         block_size=2,
@@ -153,3 +279,47 @@
     assert_block_id(block_table, [null_block_id] * 4 + original_block_ids[4:])
 
     manager.free("test")
+
+def test_sliding_window_get_num_blocks_to_allocate():
+    block_size = 2
+    sliding_window_spec = SlidingWindowSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+        sliding_window=4,  # Placeholder value, not related to test result
+        use_mla=False,
+    )
+
+    block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
+    manager = get_sliding_window_manager(sliding_window_spec, block_pool)
+    cached_blocks_1 = [KVCacheBlock(i + 1) for i in range(10)]
+    cached_blocks_2 = [block_pool.null_block for _ in range(5)
+                       ] + [KVCacheBlock(i + 1) for i in range(5)]
+
+    assert manager.get_num_blocks_to_allocate("1", 20 * block_size,
+                                              cached_blocks_1) == 20
+    assert manager.get_num_blocks_to_allocate("2", 20 * block_size,
+                                              cached_blocks_2) == 15
+
+def test_chunked_local_attention_get_num_blocks_to_allocate():
+    block_size = 2
+    attention_spec = ChunkedLocalAttentionSpec(
+        block_size=block_size,
+        num_kv_heads=1,
+        head_size=1,
+        dtype=torch.float32,
+        attention_chunk_size=4,  # Placeholder value, not related to test result
+        use_mla=False,
+    )
+
+    block_pool = BlockPool(num_gpu_blocks=100, enable_caching=True)
+    manager = get_chunked_local_attention_manager(attention_spec, block_pool)
+    cached_blocks_1 = [KVCacheBlock(i + 1) for i in range(10)]
+    cached_blocks_2 = [block_pool.null_block for _ in range(5)
+                       ] + [KVCacheBlock(i + 1) for i in range(5)]
+
+    assert manager.get_num_blocks_to_allocate("1", 20 * block_size,
+                                              cached_blocks_1) == 20
+    assert manager.get_num_blocks_to_allocate("2", 20 * block_size,
+                                              cached_blocks_2) == 15

vllm/attention/layer.py

@@ -137,6 +137,7 @@ def __init__(
                              **extra_impl_args)
         self.backend = backend_name_to_enum(attn_backend.get_name())
         self.dtype = dtype
+        self.use_irope = extra_impl_args.get("use_irope", False)
 
         # For cuda-alike (CUDA and ROCM) and cpu platforms, we control how
         # torch.compile works by registering the attention as one giant

vllm/executor/executor_base.py

@@ -112,6 +112,7 @@ def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks) -> None:
         logger.info("# %s blocks: %d, # CPU blocks: %d",
                     vllm.platforms.current_platform.device_name,
                     num_gpu_blocks, num_cpu_blocks)
+        print(f"{self.cache_config.block_size=}, {num_gpu_blocks=}")
         max_concurrency = (num_gpu_blocks * self.cache_config.block_size /
                            self.model_config.max_model_len)
         logger.info("Maximum concurrency for %s tokens per request: %.2fx",

vllm/model_executor/models/gemma3.py

@@ -148,7 +148,7 @@ def __init__(self,
         layer_idx = extract_layer_index(prefix)
         self.is_sliding = (getattr(
             config, "interleaved_sliding_window", None) is not None and bool(
-                (layer_idx + 1) % config.sliding_window_pattern))
+                (layer_idx + 1) % getattr(config, "sliding_window_pattern", 2)))
         # Initialize the rotary embedding.
         if self.is_sliding:
             # Local attention. Override the values in config.json.

vllm/v1/core/kv_cache_utils.py

@@ -9,8 +9,9 @@
 
 from vllm.config import VllmConfig
 from vllm.logger import init_logger
-from vllm.utils import GiB_bytes, sha256
-from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,
+from vllm.utils import GiB_bytes, cdiv, sha256
+from vllm.v1.kv_cache_interface import (ChunkedLocalAttentionSpec,
+                                        FullAttentionSpec, KVCacheConfig,
                                         KVCacheGroupSpec, KVCacheSpec,
                                         KVCacheTensor, SlidingWindowSpec)
 from vllm.v1.metrics.stats import PrefixCacheStats
@@ -763,6 +764,7 @@ def _get_kv_cache_config_uniform_page_size(
     # full.0, sw.0, sw.2: share a Tensor with size=available_memory//2
     # full.1, sw.1: share another Tensor with size=available_memory//2
     page_size = get_uniform_page_size(kv_cache_spec)
+    # print(f"{page_size=}, {group_size=}")
     num_blocks = get_num_blocks(vllm_config, group_size, available_memory,
                                 page_size)
     per_memory_pool_size = page_size * num_blocks
@@ -818,11 +820,16 @@ def is_hybrid(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
     logger.warning("Hybrid KV cache manager is disabled for this hybrid model,"
                    "There can be some waste of KV cache memory.")
 
+    # TODO: fix this to be more generic
     has_full_attention = any(
         isinstance(spec, FullAttentionSpec) for spec in kv_cache_spec.values())
     has_sliding_window = any(
         isinstance(spec, SlidingWindowSpec) for spec in kv_cache_spec.values())
-    if has_full_attention and has_sliding_window:
+    has_chunked_local_attention = any(
+        isinstance(spec, ChunkedLocalAttentionSpec)
+        for spec in kv_cache_spec.values())
+    if has_full_attention and (has_sliding_window
+                               or has_chunked_local_attention):
         for layer_name, spec in kv_cache_spec.items():
             if isinstance(spec, SlidingWindowSpec):
                 kv_cache_spec[layer_name] = FullAttentionSpec(
@@ -833,6 +840,14 @@ def is_hybrid(kv_cache_spec: dict[str, KVCacheSpec]) -> bool:
                     use_mla=spec.use_mla,
                     sliding_window=spec.sliding_window,
                 )
+            elif isinstance(spec, ChunkedLocalAttentionSpec):
+                kv_cache_spec[layer_name] = FullAttentionSpec(
+                    block_size=spec.block_size,
+                    num_kv_heads=spec.num_kv_heads,
+                    head_size=spec.head_size,
+                    dtype=spec.dtype,
+                    use_mla=spec.use_mla,
+                )
 
     if is_hybrid(kv_cache_spec):
         raise ValueError("Hybrid KV cache manager is disabled but failed to "
@@ -856,7 +871,6 @@ def get_kv_cache_config(
         The generated KVCacheConfigs
     """
     check_enough_kv_cache_memory(vllm_config, kv_cache_spec, available_memory)
-
     if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:
         unify_hybrid_kv_cache_specs(kv_cache_spec)