[None][feat] add commit min snapshot for SSM reuse

tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi

@@ -132,13 +132,6 @@ class BufferConfig:
     size: int
     tokens_per_block_override: int | None = None
 
-@dataclass(slots=True)
-class HelixConfig:
-    helix_group_size: int
-    helix_gpu_rank: int
-    helix_shard_size: int
-    shared_comm_port: int
-
 @dataclass(slots=True)
 class AttentionLayerConfig:
     layer_id: LayerId
@@ -180,10 +173,9 @@ class KVCacheManagerConfig:
     constraints: list[BatchDesc] = ...
     typical_step: BatchDesc | None = None
     initial_pool_ratio: list[float] | None = None
-    ssm_reuse_interval: int = 512
     swa_scratch_reuse: SwaScratchReuseConfig | None = None
+    commit_min_snapshot: bool = False
     enable_stats: bool = True
-    helix_config: HelixConfig | None = None
     @property
     def enable_swa_scratch_reuse(self) -> bool: ...
 
@@ -349,6 +341,7 @@ class _KVCache:
         self,
         accepted_input_tokens: Sequence[TokenIdExt],
         beam_search_indices: Sequence[int] | None = None,
+        is_end: bool = False,
     ) -> None: ...
     @property
     def num_committed_tokens(self) -> int: ...
@@ -493,4 +486,4 @@ class KVCacheManager:
     @property
     def need_adjustment(self) -> bool: ...
     @property
-    def ssm_reuse_interval(self) -> int: ...
+    def commit_min_snapshot(self) -> bool: ...

tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py

@@ -20,7 +20,16 @@
 from . import rawref
 from ._common import NDEBUG, BlockOrdinal, PageStatus, TokenId, TokenIdExt
 from ._life_cycle_registry import AttnLifeCycle, LifeCycle, LifeCycleId, LifeCycleRegistry
-from ._utils import TypedIndexList, chunked, div_up, filled_list, find_index, unwrap_rawref
+from ._utils import (
+    TypedIndexList,
+    chunked,
+    div_up,
+    expect_type,
+    filled_list,
+    find_index,
+    map_optional,
+    unwrap_rawref,
+)
 
 if TYPE_CHECKING:
     from ._event_manager import KVCacheEventManager
@@ -377,8 +386,13 @@ def num_life_cycles(self) -> LifeCycleId:
     def prev(self) -> "Block | RootBlock":
         return unwrap_rawref(self._prev)
 
-    def unset_page(self, lc_idx: LifeCycleId, lc: LifeCycle) -> None:
-        if self.storage[lc_idx] is None:
+    def unset_page(
+        self, lc_idx: LifeCycleId, lc: LifeCycle, expected_page: "CommittedPage | None" = None
+    ) -> None:
+        ref = self.storage[lc_idx]
+        if ref is None:
+            return
+        if expected_page is not None and ref() is not expected_page:
             return
         ordinal = self.ordinal
         self.storage[lc_idx] = None
@@ -556,24 +570,32 @@ def check_no_page_lc(b: tuple[Block, int]) -> bool:
             n = find_index(matched[: lc.num_sink_blocks], check_no_page_lc)
             if n < lc.num_sink_blocks:
                 matched = matched[:n]
-        # Check SWA window and SSM snapshot constraints together,
-        # since SSM truncation can invalidate SWA invariants.
-        # SSM is checked first (intervals are large, so it prunes more).
+        # Check SSM snapshot availability before SWA window constraints.
+        # Truncating to the last reusable SSM snapshot can change the matched
+        # length used by the SWA check.
         ssm_lc_id = life_cycles.ssm_life_cycle_id
+        if ssm_lc_id is not None:
+            from ._page import SsmCommittedPage
+
         while matched:
-            # SSM truncation: truncate to the last block with an SSM snapshot
             if ssm_lc_id is not None:
                 ssm_trunc = 0
+                ssm_match_len = 0
                 for i in reversed(range(len(matched))):
-                    if matched[i][0].storage[ssm_lc_id] is not None:
-                        assert NDEBUG or matched[i][1] == self._tokens_per_block, (
-                            "SSM reuse snapshot must only be selected from a fully matched block"
-                        )
+                    block = matched[i][0]
+                    page = map_optional(block.storage[ssm_lc_id], lambda f: f())
+                    if page is None:
+                        continue
+                    page = expect_type(SsmCommittedPage, page)
+                    snapshot_len = page.num_tokens_in_block
+                    if matched[i][1] >= snapshot_len:
                         ssm_trunc = i + 1
+                        ssm_match_len = snapshot_len
                         break
                 matched = matched[:ssm_trunc]
                 if not matched:
                     break
+                matched[-1] = (matched[-1][0], ssm_match_len)
             # SWA window check
             num_tokens = self._num_matched_tokens(matched)
             for lc_idx, lc in swa_life_cycles:

tensorrt_llm/runtime/kv_cache_manager_v2/_config.py

@@ -161,16 +161,6 @@ def __post_init__(self) -> None:
         assert self.system_prompt_length >= 0
 
 
-@dataclass(slots=True)
-class HelixConfig:
-    helix_group_size: int
-    helix_gpu_rank: int
-    # number of tokens in one helix shard
-    helix_shard_size: int
-    # must be the same for all ranks in the same helix group and different for different helix groups.
-    shared_comm_port: int
-
-
 @dataclass(slots=True)
 class SwaScratchReuseConfig:
     """
@@ -229,12 +219,6 @@ class KVCacheManagerConfig:
     takes precedence over typical_step and constraints for initial sizing.
     """
 
-    ssm_reuse_interval: int = 512
-    """
-    Interval (in tokens) at which SSM state is snapshotted for prefix reuse.
-    Must be a positive multiple of tokens_per_block. Only takes effect when SSM layers are present.
-    """
-
     swa_scratch_reuse: SwaScratchReuseConfig | None = None
     """
     When set, SWA layers reuse physical pages for out-of-window blocks during prefill.
@@ -249,14 +233,20 @@ class KVCacheManagerConfig:
     where the number of out-of-window blocks dominates memory usage.
     """
 
+    commit_min_snapshot: bool = False
+    """
+    If True, commit() records only the minimum cache snapshot reusable at the post-call
+    num_committed_tokens. Only the minimum amount of pages required for such reuse will
+    be preserved.
+
+    Required when SSM layers are present.
+    """
+
     enable_stats: bool = True
     """
     Collect V2 KV cache allocation, reuse, and transfer statistics.
     """
 
-    # unsupported yet
-    helix_config: HelixConfig | None = None
-
     @property
     def enable_swa_scratch_reuse(self) -> bool:
         return self.swa_scratch_reuse is not None
@@ -273,11 +263,6 @@ def __post_init__(self) -> None:
             for buffer in layer.buffers
         )
         if any(layer.type == LayerType.SSM for layer in self.layers):
-            assert self.ssm_reuse_interval > 0, "ssm_reuse_interval must be positive"
-            assert self.ssm_reuse_interval % self.tokens_per_block == 0, (
-                f"ssm_reuse_interval ({self.ssm_reuse_interval}) must be a multiple of "
-                f"tokens_per_block ({self.tokens_per_block})"
-            )
-            assert not self.enable_partial_reuse, (
-                "enable_partial_reuse must be False when SSM layers are present"
+            assert self.commit_min_snapshot, (
+                "commit_min_snapshot must be True when SSM layers are present"
             )