[BugFix]Get mamba cached block failed (vllm-project#11131)

### What this PR does / why we need it?
- Fixes vllm-project#10998 
Currently issue:
block size: 1k
chunk size: 4k
condition 1：
```
full  0: 1k 2k 3k 4k(evict) -> hit 3k
mamba 1:          4k        -> hit 4k
mamba 2:          4k        -> hit 4k
mamba 3:          4k        -> hit 4k

finally hit: 3k(min)
problem: group 2, 3, 4 load 3k failed(-704)
```
condition 2：
```
full  0: 1k 2k 3k 4k 5k 6k 7k 8k 9k 10k 11k 12k(evict) -> hit 11k
mamba 1:          4k(evict)   8k            12k        -> hit 8k
mamba 2:          4k          8k(evict)     12k        -> hit 4k  
mamba 3:          4k          8k            12k        -> hit 4k

finally hit: 4k(min)
problem: group 1 load 4k failed(-704)
```
After Fix:
condition 1：
```
full  0: 1k 2k 3k 4k(evict) -> hit [1k 2k 3k]
mamba 1:          4k        -> hit [4k]
mamba 2:          4k        -> hit [4k]
mamba 3:          4k        -> hit [4k]

finally hit: 0(max(intersection))
```
condition 2：
```
full  0: 1k 2k 3k 4k 5k 6k 7k 8k 9k 10k 11k 12k(evict) -> hit [1k 2k 3k 4k 5k 6k 7k 8k 9k 10k 11k]
mamba 1:          4k(evict)   8k            12k        -> hit [8k]
mamba 2:          4k          8k(evict)     12k        -> hit [4k] 
mamba 3:          4k          8k            12k        -> hit [4k]

finally hit: 0(max(intersection))
```
condition 2：
```
full  0: 1k 2k 3k 4k 5k 6k 7k 8k 9k 10k 11k 12k(evict) -> hit [1k 2k 3k 4k 5k 6k 7k 8k 9k 10k 11k]
mamba 1:          4k          8k            12k        -> hit [4k 8k]
mamba 2:          4k          8k(evict)     12k        -> hit [4k] 
mamba 3:          4k          8k            12k        -> hit [4k]

finally hit: 4k(max(intersection))
```

- vLLM version: v0.23.0
- vLLM main:
vllm-project/vllm@dc68bd8
---------
Signed-off-by: Qingsong Zhang <1640410765@qq.com>

Author: zqs164

tests/ut/distributed/ascend_store/test_pool_worker.py

@@ -50,28 +50,61 @@ def test_check_all_layers_exists_none(self):
         result = cls.check_all_layers_exists(None, [0, 0, 0], 3)
         self.assertEqual(result, [0])
 
-    def test_find_max_hit_index_found(self):
+    def test_find_all_continuous_hit_positions_found(self):
         cls = self._make_worker_class()
         arr = [[1, 1, 0], [1, 0, 1]]
-        result = cls.find_max_hit_index(None, arr, 3)
-        self.assertEqual(result, 0)
+        result = cls.find_all_continuous_hit_positions(arr, [16, 32, 48], 3, 48, 16)
+        self.assertEqual(result, [16])
 
-    def test_find_max_hit_index_all_one(self):
+    def test_find_all_continuous_hit_positions_all_one(self):
         cls = self._make_worker_class()
         arr = [[1, 1, 1], [1, 1, 1]]
-        result = cls.find_max_hit_index(None, arr, 3)
-        self.assertEqual(result, 2)
+        result = cls.find_all_continuous_hit_positions(arr, [16, 32, 48], 3, 48, 16)
+        self.assertEqual(result, [16, 32, 48])
 
-    def test_find_max_hit_index_first_pos(self):
+    def test_find_all_continuous_hit_positions_first_pos(self):
         cls = self._make_worker_class()
         arr = [[0, 1], [1, 0]]
-        result = cls.find_max_hit_index(None, arr, 3)
-        self.assertEqual(result, -1)
+        result = cls.find_all_continuous_hit_positions(arr, [16, 32], 2, 48, 16)
+        self.assertEqual(result, [])
+
+    def test_find_all_continuous_hit_positions_empty(self):
+        cls = self._make_worker_class()
+        result = cls.find_all_continuous_hit_positions([], [], 0, 48, 16)
+        self.assertEqual(result, [])
+
+    def test_find_all_discontinuous_hit_positions_all_tp_hits(self):
+        cls = self._make_worker_class()
+        arr = [[0, 0, 1, 0, 0, 1], [0, 0, 1, 0, 0, 1]]
+        result = cls.find_all_discontinuous_hit_positions(arr, [16, 32, 48, 64, 80, 96], 6, 128, 16)
+        self.assertEqual(result, [48, 96])
+
+    def test_find_all_discontinuous_hit_positions_some_tp_hits(self):
+        cls = self._make_worker_class()
+        arr = [[0, 0, 1, 0, 0, 1], [0, 0, 1, 0, 0, 0]]
+        result = cls.find_all_discontinuous_hit_positions(arr, [16, 32, 48, 64, 80, 96], 6, 128, 16)
+        self.assertEqual(result, [48])
+
+    def test_find_all_discontinuous_hit_positions_all_tp_hits_with_limits(self):
+        cls = self._make_worker_class()
+        arr = [[0, 0, 1, 0, 0, 1], [0, 0, 1, 0, 0, 1]]
+        result = cls.find_all_discontinuous_hit_positions(arr, [16, 32, 48, 64, 80, 96], 6, 64, 16)
+        self.assertEqual(result, [48])
+
+    def test_max_intersection_hit_position_single_group(self):
+        cls = self._make_worker_class()
+        hits = [[16, 32, 48]]
+        self.assertEqual(48, cls._max_intersection_hit_position(hits))
+
+    def test_max_intersection_hit_position_empty_group(self):
+        cls = self._make_worker_class()
+        hits: list[list[int]] = []
+        self.assertEqual(0, cls._max_intersection_hit_position(hits))
 
-    def test_find_max_hit_index_empty(self):
+    def test_max_intersection_hit_position_multi_group(self):
         cls = self._make_worker_class()
-        result = cls.find_max_hit_index(None, [], 0)
-        self.assertEqual(result, -1)
+        hits = [[16, 32, 48], [32, 48], [16, 32], [32, 48, 64]]
+        self.assertEqual(32, cls._max_intersection_hit_position(hits))
 
 
 class TestKVPoolWorkerInit(unittest.TestCase):
@@ -84,6 +117,7 @@ def _make_vllm_config(self, kv_role="kv_producer", extra_config=None, block_size
         config.model_config.hf_text_config = MagicMock(spec=[])  # no index_topk
         config.model_config.get_num_layers.return_value = 32
         config.model_config.get_total_num_kv_heads.return_value = 8
+        config.model_config.max_model_len = 1024
         config.parallel_config.data_parallel_rank = 0
         config.parallel_config.rank = 0
         config.parallel_config.pipeline_parallel_size = 1
@@ -461,6 +495,7 @@ def _make_config(self, kv_role="kv_producer", extra_config=None, block_size=16):
         config.model_config.model = "org/llama-7b"
         config.model_config.use_mla = False
         config.model_config.hf_text_config = MagicMock(spec=[])
+        config.model_config.max_model_len = 1024
         config.model_config.get_num_layers.return_value = 2
         config.model_config.get_total_num_kv_heads.return_value = 1
         config.parallel_config.data_parallel_rank = 0

vllm_ascend/distributed/kv_transfer/kv_pool/ascend_store/pool_worker.py

@@ -78,6 +78,7 @@ def __init__(
         hf_config = getattr(model_config, "hf_config", hf_text_config)
         self.hf_config = hf_text_config or hf_config
         self.compress_ratios = getattr(hf_text_config, "compress_ratios", None)
+        self.max_model_len = model_config.max_model_len
         if self.compress_ratios is None:
             self.compress_ratios = getattr(hf_config, "compress_ratios", None)
         self.use_compress = self.compress_ratios is not None
@@ -1018,7 +1019,8 @@ def lookup_scheduler(
         :return: An int indicating how many prefix tokens are cached.
         """
         try:
-            hits = []
+            hits: list[list[int]] = []
+            max_hit_position = self.max_model_len
             kv_cache_group_ids = kv_cache_group_ids or [0]
             kv_cache_group_ids = self._get_lookup_gate_group_ids(kv_cache_group_ids)
             for group_id in kv_cache_group_ids:
@@ -1081,32 +1083,22 @@ def lookup_scheduler(
                     multi_tp_keys[:3],
                 )
                 if group_id < len(self.group_uses_align_state) and self.group_uses_align_state[group_id]:
-                    # mamba group with align mode will skip some null block, we must loop it in reverse order
-                    for i in range(num_block - 1, -1, -1):
-                        if (
-                            all(values[i] == 1 for values in multi_tp_values)
-                            and ends[i] % self.cache_transfer_granularity == 0
-                        ):
-                            hits.append(ends[i])
-                            break
-                    else:
-                        return 0
+                    group_hits = self.find_all_discontinuous_hit_positions(
+                        multi_tp_values, ends, num_block, max_hit_position, self.cache_transfer_granularity
+                    )
                 else:
-                    index = self.find_max_hit_index(multi_tp_values, num_block)
-                    if index == -1:
-                        return 0
-                    else:
-                        for hit_index in range(index, -1, -1):
-                            if ends[hit_index] % self.cache_transfer_granularity == 0:
-                                hits.append(ends[hit_index])
-                                break
-                        else:
-                            return 0
+                    group_hits = self.find_all_continuous_hit_positions(
+                        multi_tp_values, ends, num_block, max_hit_position, self.cache_transfer_granularity
+                    )
+                if not group_hits:
+                    return 0
+                max_hit_position = min(max_hit_position, group_hits[-1])
+                hits.append(group_hits)
                 logger.debug(
                     "KV pool scheduler lookup group=%d keys=%d hit=%d token_len=%d",
                     group_id,
                     len(keys),
-                    hits[-1],
+                    max_hit_position,
                     token_len,
                 )
         except Exception as e:
@@ -1116,14 +1108,26 @@ def lookup_scheduler(
                 e,
             )
             return 0
+        final_hits = self._max_intersection_hit_position(hits)
         logger.debug(
-            "KV pool scheduler lookup final token_len=%d groups=%s hits=%s result=%d",
+            "KV pool scheduler lookup final token_len=%d groups=%s hit=%d",
             token_len,
             kv_cache_group_ids,
-            hits,
-            min(hits) if hits else 0,
+            final_hits,
         )
-        return min(hits) if hits else 0
+        return final_hits
+
+    @staticmethod
+    def _max_intersection_hit_position(hits: list[list[int]]) -> int:
+        """
+        For all attention groups, treat the position of the maximum common hit as the final hit position
+        """
+        if not hits:
+            return 0
+        common_elements = set(hits[0]).intersection(*hits[1:])
+        if not common_elements:
+            return 0
+        return max(common_elements)
 
     def check_all_layers_exists(self, res: list[int], num_layers: int) -> list[int]:
         total_chunks = len(res) // num_layers
@@ -1137,13 +1141,37 @@ def check_all_layers_exists(self, res: list[int], num_layers: int) -> list[int]:
 
         return result
 
-    def find_max_hit_index(self, arr, num_blocks: int):
+    @staticmethod
+    def find_all_discontinuous_hit_positions(
+        arr, ends, num_blocks: int, max_hit_position: int, cache_transfer_granularity: int
+    ) -> list[int]:
+        """
+        For mamba attn, there will be some uncached null blocks, we just collect all hit positions,
+        and use the last position as final hit position
+        """
+        hits: list[int] = []
         for i in range(num_blocks):
-            if any(row[i] != 1 for row in arr):
-                return i - 1
-        else:
-            # if arr is not empty, all hits, else no hits
-            return len(arr[0]) - 1 if arr else -1
+            if ends[i] > max_hit_position:
+                break
+            if all(row[i] == 1 for row in arr):
+                if ends[i] % cache_transfer_granularity == 0:
+                    hits.append(ends[i])
+        return hits
+
+    @staticmethod
+    def find_all_continuous_hit_positions(
+        arr, ends, num_blocks: int, max_hit_position: int, cache_transfer_granularity: int
+    ) -> list[int]:
+        hits: list[int] = []
+        for i in range(num_blocks):
+            if ends[i] > max_hit_position:
+                break
+            if all(row[i] == 1 for row in arr):
+                if ends[i] % cache_transfer_granularity == 0:
+                    hits.append(ends[i])
+            else:
+                break
+        return hits
 
     def get_kv_events(self) -> list[BlockStored]:
         if self.enable_kv_events and self.kv_send_thread is not None: