fix issue of local attention start idx based on num computed tokens

Signed-off-by: Lu Fang <[email protected]>

Author: Lu Fang

tests/v1/core/test_specialized_manager.py

@@ -74,22 +74,24 @@ def run_one_case(block_is_cached, tail_token, expect_length):
 
     run_one_case([True], 0, 1)
     run_one_case([True], 1, 1)
-    run_one_case([True, False], 0, 1)
+    run_one_case([True, False], 0, 2)
     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], 0, 4)
     run_one_case([True, True, True, False], 1, 4)
+    run_one_case([random.choice([True, False])] * 8 + [True], 1, 9)
+    run_one_case([random.choice([True, False])] * 8 + [False], 1, 8)
     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], 0, 10)
     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], 0, 10)
     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], 0, 10)
     run_one_case([random.choice([True, False])] * 8 + [False, False], 1, 10)
 
 
@@ -198,23 +200,18 @@ def assert_block_id(block_table: list[KVCacheBlock], ids: list[int]):
     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.
+    # For 4th token (0-indexed), token 0-3 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])
+    assert_block_id(block_table, [null_block_id] * 2)
 
-    # 6 tokens are computed. token 4 - 5 are in local attention window,
+    # For 6th token (0-indexed), token 4 - 6 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:])
+    # For 12th token (0-indexed),
+    # token 0-11 are out, 6 block can be removed.
+    manager.remove_skipped_blocks("test", 12)
+    assert_block_id(block_table, [null_block_id] * 6)
 
 
 def test_sliding_window_remove_skipped_blocks():

vllm/v1/core/single_type_kv_cache_manager.py

@@ -396,6 +396,7 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
     def __init__(self, kv_cache_spec: ChunkedLocalAttentionSpec,
                  block_pool: BlockPool, **kwargs) -> None:
         super().__init__(kv_cache_spec, block_pool, **kwargs)
+        print("local attetntion manager is init")
         self.attention_chunk_size = kv_cache_spec.attention_chunk_size
         self._null_block = block_pool.null_block
 
@@ -414,7 +415,7 @@ def find_longest_cache_hit(
             "chunked local attention groups")
         max_num_blocks = max_length // kv_cache_spec.block_size
         if max_length > 0:
-            local_attention_start_idx = ((max_length - 1) //
+            local_attention_start_idx = (max_length //
                                          kv_cache_spec.attention_chunk_size *
                                          kv_cache_spec.attention_chunk_size)
         else:
@@ -450,12 +451,12 @@ def remove_skipped_blocks(self, request_id: str,
         # chunked attention window and skipped
         # during the attention computation.
 
-        # (N-1) // chunk_size * chunk_size
+        # N // chunk_size * chunk_size
         # [chunk 0][chunk 1]local_attention_start_idx ... current
 
         local_attention_start_idx = (
-            num_computed_tokens -
-            1) // self.attention_chunk_size * self.attention_chunk_size
+            num_computed_tokens
+        ) // self.attention_chunk_size * self.attention_chunk_size
         # 1024-> 0, 1025-> 1024
         first_useful_block_idx = local_attention_start_idx // self.block_size
         # block size =128, 0 -> block 0, 1024 -> block 8, 372 -> block 2