@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 869f78732b

[Kimi Bug] Fix kimi k3 AssertionError assert 0 <= checkpoint_idx < len(blocks) (#55747)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

Wentao Ye committed Sep 7, 2026 at 22:34 UTC 869f78732b64454293d2ba42ae0008386fdaa6a6

2 files changed +17 -12

tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py

+1 -1

@@ -651,7 +651,7 @@ def test_hash_aligned_query_end_uses_regular_partial_tail():

651

652 assert new_blocks is not None

653 mamba_manager = manager.coordinator.single_type_managers[1]

654 - assert mamba_manager._checkpoint_positions[request.request_id] == 0

654 + assert request.request_id not in mamba_manager._checkpoints

655 mamba_blocks = manager.get_blocks(request.request_id).blocks[1]

656 assert sum(not block.is_null for block in mamba_blocks) == 1

657

vllm/v1/core/single_type_kv_cache_manager.py

+16 -11

@@ -1431,9 +1431,9 @@ class MambaManager(SingleTypeKVCacheManager):

1431 self.last_state_block_idx: dict[str, int] = {}

1432 # The set of the requests that have been allocated blocks

1433 self._allocated_block_reqs: set[str] = set()

1434 - # Absolute checkpoint position selected by the scheduler for the

1435 - # current allocation.

1436 - self._checkpoint_positions: dict[str, int] = {}

1434 + # checkpoint position and reserved block index for the current

1435 + # allocation.

1436 + self._checkpoints: dict[str, tuple[int, int]] = {}

1437 # Requests that registered their own last-prompt-boundary partial

1438 # tail (producers). A later CoW hands its private copy to the

1439 # connector; a request that finishes first hands off this table

@@ -1722,7 +1722,14 @@ class MambaManager(SingleTypeKVCacheManager):

1722 checkpoint_position = 0

1723 checkpoint_block = int(checkpoint_position > 0)

1724 if not apply_admission_cap:

1725 - self._checkpoint_positions[request_id] = checkpoint_position

1725 + if checkpoint_position > 0:

1726 + checkpoint_idx = cdiv(num_tokens, self.block_size) - 2

1727 + self._checkpoints[request_id] = (

1728 + checkpoint_position,

1729 + checkpoint_idx,

1730 + )

1731 + else:

1732 + self._checkpoints.pop(request_id, None)

1733 if num_new_blocks > 0:

1734 blocks_allocated = request_id in self._allocated_block_reqs

1735 if not (checkpoint_block and blocks_allocated):

@@ -1760,8 +1767,7 @@ class MambaManager(SingleTypeKVCacheManager):

1767 num_required_blocks = (

1768 cdiv(num_tokens, self.block_size) + self.num_speculative_blocks

1769 )

1763 - checkpoint_position = self._checkpoint_positions.get(request_id, 0)

1764 - checkpoint_block = int(checkpoint_position > 0)

1770 + checkpoint_block = int(request_id in self._checkpoints)

1771 partial_hit = self._partial_hit_reqs.get(request_id)

1772 has_partial_hit = partial_hit is not None

1773 # `num_required_blocks` might be less than `len(req_blocks)` if blocks are

@@ -1893,7 +1899,7 @@ class MambaManager(SingleTypeKVCacheManager):

1899 if self.mamba_cache_mode == "align":

1900 self._allocated_block_reqs.discard(request_id)

1901 self.last_state_block_idx.pop(request_id, None)

1896 - self._checkpoint_positions.pop(request_id, None)

1902 + self._checkpoints.pop(request_id, None)

1903 self._producer_partial_tail_reqs.pop(request_id, None)

1904 # An offer is only guaranteed to hold committed bytes until the end

1905 # of the pass that made it. This request's blocks are going back to

@@ -1963,10 +1969,9 @@ class MambaManager(SingleTypeKVCacheManager):

1969 ) -> BlockHashWithGroupId | None:

1970 hash_block_size = self.block_pool.hash_block_size

1971 # Re-key the reserved block at its exported checkpoint boundary.

1966 - checkpoint_position = self._checkpoint_positions.get(request.request_id, 0)

1967 - if checkpoint_position > 0:

1968 - # TODO: Store the reserved slot explicitly for multi-module MTP.

1969 - checkpoint_idx = cdiv(num_tokens, self.block_size) - 2

1972 + checkpoint = self._checkpoints.get(request.request_id)

1973 + if checkpoint is not None:

1974 + checkpoint_position, checkpoint_idx = checkpoint

1975 blocks = self.req_to_blocks[request.request_id]

1976 assert 0 <= checkpoint_idx < len(blocks)

1977 checkpoint_block = blocks[checkpoint_idx]