[BugFix]Mamba pooling&mtp (#10565)

### What this PR does / why we need it?
- Fixed the bug where the KV cache usage statistics displayed a value
less than zero.
- Fixed the bug where the hybrid attention pointers registered to the
backend were not aligned to 2MB.

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

Author: zqs164

tests/ut/distributed/ascend_store/test_config_data.py

@@ -283,6 +283,74 @@ def test_update_invalid_type(self):
         with self.assertRaises(ValueError):
             tracker.update("invalid")  # type: ignore[arg-type]
 
+    def test_update_mamba_with_tuple(self):
+        tracker = RequestTracker(
+            req_id="r1", token_len=16, allocated_block_ids_by_group=[[1], [2], [3], [4]], block_sizes=[16] * 4
+        )
+        tracker.update(([5, 6], [0, 7], [0, 8], [0, 9]))
+        self.assertEqual(tracker.allocated_block_ids_by_group[0], [1, 5, 6])
+        self.assertEqual(tracker.allocated_block_ids_by_group[1], [2, 0, 7])
+        self.assertEqual(tracker.allocated_block_ids_by_group[2], [3, 0, 8])
+        self.assertEqual(tracker.allocated_block_ids_by_group[3], [4, 0, 9])
+
+    def test_update_mamba_mtp_with_tuple_chunk2(self):
+        tracker = RequestTracker(
+            req_id="r1",
+            token_len=32,
+            allocated_block_ids_by_group=[
+                [1, 2],
+                [0, 3, 4, 5, 6],
+                [0, 7, 8, 9, 10],
+                [0, 11, 12, 13, 14],
+            ],
+            mamba_group_ids=[1, 2, 3],
+            num_speculative_blocks=3,
+            block_sizes=[16] * 4,
+        )
+
+        tracker.update(([15, 16], [4, 17], [8, 18], [12, 19]), 32)
+        self.assertEqual(tracker.allocated_block_ids_by_group[0], [1, 2, 15, 16])
+        self.assertEqual(tracker.allocated_block_ids_by_group[1], [0, 3, 0, 5, 6, 4, 17])
+        self.assertEqual(tracker.allocated_block_ids_by_group[2], [0, 7, 0, 9, 10, 8, 18])
+        self.assertEqual(tracker.allocated_block_ids_by_group[3], [0, 11, 0, 13, 14, 12, 19])
+
+    def test_update_mamba_mtp_with_tuple_chunk8(self):
+        tracker = RequestTracker(
+            req_id="r1",
+            token_len=128,
+            allocated_block_ids_by_group=[
+                [1, 2, 3, 4, 5, 6, 7, 8],
+                [0, 0, 0, 0, 0, 0, 0, 9, 10, 11, 12],
+                [0, 0, 0, 0, 0, 0, 0, 13, 14, 15, 16],
+                [0, 0, 0, 0, 0, 0, 0, 17, 18, 19, 20],
+            ],
+            mamba_group_ids=[1, 2, 3],
+            num_speculative_blocks=3,
+            block_sizes=[16] * 4,
+        )
+
+        tracker.update(
+            (
+                [21, 22, 23, 24, 25, 26, 27, 28],
+                [0, 0, 0, 0, 10, 11, 12, 29],
+                [0, 0, 0, 0, 14, 15, 16, 30],
+                [0, 0, 0, 0, 18, 19, 20, 31],
+            ),
+            128,
+        )
+        self.assertEqual(
+            tracker.allocated_block_ids_by_group[0], [1, 2, 3, 4, 5, 6, 7, 8, 21, 22, 23, 24, 25, 26, 27, 28]
+        )
+        self.assertEqual(
+            tracker.allocated_block_ids_by_group[1], [0, 0, 0, 0, 0, 0, 0, 9, 0, 0, 0, 0, 0, 0, 0, 10, 11, 12, 29]
+        )
+        self.assertEqual(
+            tracker.allocated_block_ids_by_group[2], [0, 0, 0, 0, 0, 0, 0, 13, 0, 0, 0, 0, 0, 0, 0, 14, 15, 16, 30]
+        )
+        self.assertEqual(
+            tracker.allocated_block_ids_by_group[3], [0, 0, 0, 0, 0, 0, 0, 17, 0, 0, 0, 0, 0, 0, 0, 18, 19, 20, 31]
+        )
+
 
 class TestReqMeta(unittest.TestCase):
     def test_from_request_tracker_basic_save(self):

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

@@ -542,6 +542,13 @@ class RequestTracker:
 
     last_block_key: str | None = None
 
+    mamba_group_ids: list[int] | None = None
+
+    # spec blocks for mamba cache group
+    num_speculative_blocks: int = 0
+
+    block_sizes: list[int] | None = None
+
     def __init__(
         self,
         req_id: str,
@@ -558,9 +565,14 @@ def __init__(
         ends: list[int] | None = None,
         sizes_per_chunk: list[list[int]] | None = None,
         last_block_key: str | None = None,
+        mamba_group_ids: list[int] | None = None,
+        num_speculative_blocks: int = 0,
+        block_sizes: list[int] | None = None,
     ) -> None:
         self.req_id = req_id
         self.token_len = token_len
+        self.mamba_group_ids = mamba_group_ids
+        self.num_speculative_blocks = num_speculative_blocks
         block_ids = allocated_block_ids_by_group
         if block_ids is None:
             block_ids = normalize_block_ids_by_group(allocated_block_ids or [])
@@ -575,6 +587,7 @@ def __init__(
         self.ends = ends
         self.sizes_per_chunk = sizes_per_chunk
         self.last_block_key = last_block_key
+        self.block_sizes = block_sizes
 
     @property
     def allocated_block_ids(self) -> list[int]:
@@ -601,6 +614,7 @@ def from_new_request(
     def update(
         self,
         new_block_ids: tuple[list[int], ...] | list[int],
+        num_computed_tokens: int = 0,
     ) -> None:
         """Update the request tracker when a running request is scheduled again."""
         normalized = normalize_block_ids_by_group(new_block_ids)
@@ -609,8 +623,37 @@ def update(
                 [[] for _ in range(len(normalized) - len(self.allocated_block_ids_by_group))]
             )
         for group_id, ids in enumerate(normalized):
+            self.update_mamba_spec_blocks(ids, group_id, num_computed_tokens)
             self.allocated_block_ids_by_group[group_id].extend(ids)
 
+    def update_mamba_spec_blocks(self, block_ids: list[int], kv_cache_group_id: int, num_computed_tokens: int):
+        """
+        for mamba align groups, each step will:
+            - Firstly, remove some previous blocks and append some necessary null blocks
+            - Secondly, move the speculative blocks(maybe all or partially) to the last position for reuse
+            - Finally, allocate a new block
+        so, if a speculative block is moved to last position and replaced with null block,
+        we also need to update the previous allocated_block_ids to 0.
+        """
+        if self.mamba_group_ids and kv_cache_group_id in self.mamba_group_ids:
+            assert self.block_sizes is not None and len(self.block_sizes) > kv_cache_group_id
+            num_skipped_blocks = (
+                max(num_computed_tokens - self.num_speculative_blocks - 1, 0) // self.block_sizes[kv_cache_group_id]
+            )
+            num_skipped_blocks = min(len(self.allocated_block_ids_by_group[kv_cache_group_id]), num_skipped_blocks)
+            if num_skipped_blocks > 0:
+                self.allocated_block_ids_by_group[kv_cache_group_id][:num_skipped_blocks] = [0] * num_skipped_blocks
+            if not block_ids or self.num_speculative_blocks <= 0:
+                return
+            mask_spec_count = min(len(block_ids) - 1, self.num_speculative_blocks)
+            group_block_ids = self.allocated_block_ids_by_group[kv_cache_group_id]
+            if mask_spec_count >= self.num_speculative_blocks:
+                group_block_ids[-self.num_speculative_blocks :] = [0] * self.num_speculative_blocks
+            else:
+                group_block_ids[-self.num_speculative_blocks : mask_spec_count - self.num_speculative_blocks] = [
+                    0
+                ] * mask_spec_count
+
 
 @dataclass(init=False)
 class ReqMeta:

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

@@ -97,6 +97,9 @@ def __init__(
         self.dcp_size = getattr(vllm_config.parallel_config, "decode_context_parallel_size", 1)
 
         self.mamba_group_ids = self._infer_mamba_groups()
+        self.num_speculative_blocks = (
+            vllm_config.speculative_config.num_speculative_tokens if vllm_config.speculative_config else 0
+        )
         self.original_block_size = self._infer_group_block_sizes(vllm_config, kv_cache_config)
         cp_scale = self.pcp_size * self.dcp_size
         self.grouped_block_size = [block_size * cp_scale for block_size in self.original_block_size]
@@ -704,6 +707,9 @@ def _process_new_request(
             block_keys=(previous_tracker.block_keys.copy() if previous_tracker else []),
             block_gvas=(previous_tracker.block_gvas.copy() if previous_tracker else []),
             gva_block_offset=(previous_tracker.gva_block_offset if previous_tracker else 0),
+            mamba_group_ids=self.mamba_group_ids,
+            num_speculative_blocks=self.num_speculative_blocks,
+            block_sizes=self.grouped_block_size,
         )
         self._request_trackers[request.req_id] = request_tracker
         num_blocks = num_tokens_to_compute // self._block_size
@@ -751,6 +757,9 @@ def _process_preempted_cached_request(
             block_keys=(previous_tracker.block_keys.copy() if previous_tracker else []),
             block_gvas=(previous_tracker.block_gvas.copy() if previous_tracker else []),
             gva_block_offset=(previous_tracker.gva_block_offset if previous_tracker else 0),
+            mamba_group_ids=self.mamba_group_ids,
+            num_speculative_blocks=self.num_speculative_blocks,
+            block_sizes=self.grouped_block_size,
         )
         self._request_trackers[req_id] = request_tracker
         num_blocks = len(new_block_ids_by_group[0])
@@ -785,15 +794,15 @@ def _process_running_cached_request(
             raise ValueError(f"Request {req_id} is not in _request_trackers, but it is scheduled to be cached")
         num_new_tokens = scheduler_output.num_scheduled_tokens[req_id]
         req_tuple = self._unfinished_requests.get(req_id)
-        if req_tuple:
-            request = req_tuple[0]
-            num_current_tokens = request_tracker.token_len
-            new_token_ids = request.all_token_ids[num_current_tokens : num_current_tokens + num_new_tokens]
-            if request_tracker.token_ids is not None and new_token_ids:
-                request_tracker.token_ids.extend(new_token_ids)
-            request_tracker.token_len += num_new_tokens
-        else:
+        if not req_tuple:
             raise ValueError(f"Request {req_id} is not in _unfinished_requests, but it is scheduled to be cached")
+        request = req_tuple[0]
+        num_current_tokens = request_tracker.token_len
+        new_token_ids = request.all_token_ids[num_current_tokens : num_current_tokens + num_new_tokens]
+        if request_tracker.token_ids is not None and new_token_ids:
+            request_tracker.token_ids.extend(new_token_ids)
+        request_tracker.token_len += num_new_tokens
+
         prev_token_count = request_tracker.token_len - num_new_tokens
         prev_hash_count = prev_token_count // self._block_size
         current_hash_count = request_tracker.token_len // self._block_size
@@ -813,7 +822,7 @@ def _process_running_cached_request(
                     has_last_block=True,
                 )
         if new_block_ids is not None:
-            request_tracker.update(new_block_ids)
+            request_tracker.update(new_block_ids, request.num_computed_tokens)
         load_spec = None
         return self._build_req_meta(
             request_tracker,
@@ -846,6 +855,9 @@ def _process_async_load_request(
             block_keys=(previous_tracker.block_keys.copy() if previous_tracker else []),
             block_gvas=(previous_tracker.block_gvas.copy() if previous_tracker else []),
             gva_block_offset=(previous_tracker.gva_block_offset if previous_tracker else 0),
+            mamba_group_ids=self.mamba_group_ids,
+            num_speculative_blocks=self.num_speculative_blocks,
+            block_sizes=self.grouped_block_size,
         )
         self._request_trackers[request_id] = request_tracker
         num_blocks = num_tokens_to_compute // self._block_size
@@ -973,9 +985,9 @@ def update_connector_output(self, connector_output: KVConnectorOutput):
         hand the connector_output, free non-null mamba blocks and so on.
         """
         meta = connector_output.kv_connector_worker_meta
-        if not isinstance(meta, AscendStoreKVConnectorWorkerMetadata):
+        if not isinstance(meta, AscendStoreKVConnectorWorkerMetadata) or self._block_pool is None:
             return
-        to_free_block_ids: list[int] = []
+
         for event_id, count in meta.completed_events.items():
             logger.debug("event %s update with %s", event_id, count)
             total = self.sending_events.get(event_id, -1)
@@ -984,16 +996,14 @@ def update_connector_output(self, connector_output: KVConnectorOutput):
                 continue
             total = total + count
             if total >= self._expected_worker_count:
-                to_free_block_ids.extend(self.sending_blocks.pop(event_id, []))
+                to_free_block_ids = self.sending_blocks.pop(event_id, [])
                 self.sending_events.pop(event_id, None)
+                if to_free_block_ids:
+                    logger.debug("event %s free blocks: %s", event_id, to_free_block_ids)
+                    self._block_pool.free_blocks([self._block_pool.blocks[block_id] for block_id in to_free_block_ids])
             else:
                 self.sending_events[event_id] = total
 
-        if to_free_block_ids:
-            logger.debug("free blocks: %s", to_free_block_ids)
-            assert self._block_pool is not None
-            self._block_pool.free_blocks([self._block_pool.blocks[block_id] for block_id in to_free_block_ids])
-
     def request_finished(
         self,
         request: "Request",

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

@@ -500,6 +500,25 @@ def _infer_cache_group_metadata(self, group_id: int, layer_names: list[str]):
         self.group_block_stride[group_id] = group_block_strides
         self.group_num_layers[group_id] = len(layer_names)
 
+    def _align_kv_ptrs(self, registered_regions: dict[int, tuple[int, int]]):
+        """
+        In hybrid scenario, where a KVCacheTensor is shared by multiple layers,
+        but sometimes, layers cannot be evenly distributed among multiple groups,
+        the layers sharing the KVCacheTensor may not completely occupy all the space of the KVCacheTensor.
+        This results in the calculated start address not being the previously aligned address.
+        Therefore, we down-align the start address to meet the 2MB alignment requirement.
+        """
+        if not self.use_hybrid:
+            return
+        alignment = 2 * 1024 * 1024
+        for storage_key in registered_regions:
+            start, end = registered_regions[storage_key]
+            new_start = start // alignment * alignment
+            # Because the addresses of raw tensors are aligned to 2MB,
+            # all shared sub-tensors, when aligned downwards, should theoretically not exceed the address bounds.
+            assert new_start >= storage_key, "invalid kv cache tensor, raw tensor ptr must be align to 2MB"
+            registered_regions[storage_key] = (new_start, end)
+
     def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         _, first_kv_cache_tuple = next(iter(kv_caches.items()))
         first_kv_cache_tuple = self._as_cache_tuple(first_kv_cache_tuple)
@@ -553,6 +572,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                 else:
                     registered_regions[storage_key] = (start, end)
 
+        self._align_kv_ptrs(registered_regions)
         ptrs = [start for start, _ in registered_regions.values()]
         lengths = [end - start for start, end in registered_regions.values()]