[Hybrid] Marconi-style admission policy for hybrid cache

vllm/v1/core/kv_cache_coordinator.py

@@ -482,6 +482,7 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:
 
         num_groups = len(self.kv_cache_config.kv_cache_groups)
         hit_length = max_cache_hit_length
+        longest_hit_length = 0
         hit_blocks_by_group: list[list[KVCacheBlock] | None] = [None] * num_groups
 
         # Simple hybrid (1 full attn + 1 other): one iteration suffices.
@@ -523,7 +524,12 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:
                     curr_hit_length = len(hit_blocks[0]) * spec.block_size
                     for group_id, blocks in zip(group_ids, hit_blocks):
                         hit_blocks_by_group[group_id] = blocks
-
+
+                # Collect information on the longest cached prefix overall
+                # (no matter the attention type) to allow for more complex
+                # caching policies
+                longest_hit_length = max(longest_hit_length, curr_hit_length)
+
             if curr_hit_length >= hit_length:
                 break
             hit_length = curr_hit_length
@@ -541,7 +547,7 @@ def _get_block_hashes(kv_cache_spec: KVCacheSpec) -> BlockHashList:
 
         return tuple(
             blocks if blocks is not None else [] for blocks in hit_blocks_by_group
-        ), hit_length
+        ), longest_hit_length # longest_hit_length >= hit_length
 
 
 def get_kv_cache_coordinator(

vllm/v1/core/sched/scheduler.py

@@ -291,6 +291,7 @@
         num_new_tokens: int,
         num_new_local_computed_tokens: int = 0,
         num_external_computed_tokens: int = 0,
+        mamba_tokens_lag: int = 0,
     ) -> int:
         assert num_external_computed_tokens == 0, (
             "External KV connector is not verified yet"
@@ -333,6 +334,21 @@
             else:
                 # prefill the last few tokens
                 pass
+
+            # Marconi cache admission optimization:
+            # Create cache entries at divergence points of common prefixes.
+            # 
+            # Implementation:
+            # If mamba cache "lags" behind the KVCache hits for >= block_size,
+            # there is a common shared prefix that wasn't cached.
+            if mamba_tokens_lag >= block_size:
+                # If num_new_tokens is longer than lag, 
+                # the prefix normally still wouldn't be cached
+                if num_new_tokens > mamba_tokens_lag:
+                    # So we force caching at mamba_tokens_lag
+                    num_new_tokens = mamba_tokens_lag
+                    assert mamba_tokens_lag % block_size == 0 #TODO?
+                    #num_new_tokens = num_new_tokens // block_size * block_size
         return num_new_tokens
 
     def schedule(self) -> SchedulerOutput:
@@ -602,6 +618,21 @@
                     new_computed_blocks, num_new_local_computed_tokens = (
                         self.kv_cache_manager.get_computed_blocks(request)
                     )
+
+                    # More proper check would be:
+                    # if isinstance(self.kv_cache_manager.coordinator, 
+                    #               HybridKVCacheCoordinator):
+                    # but this check is similar and avoids
+                    # importing HybridKVCacheCoordinator:
+                    if self.has_mamba_layers:
+                        # HybridKVCacheCoordinator returns the longest hit:
+                        longest_hit_length = num_new_local_computed_tokens
+                        # Obtain the shortest cached prefix from the blocks:
+                        num_new_local_computed_tokens = \
+                            len(new_computed_blocks.blocks[0]) * self.block_size
+                        # Mamba tokens "lag" - how far it's behind longest hit:
+                        mamba_tokens_lag = \
+                            longest_hit_length - num_new_local_computed_tokens
 
                     # Get externally-cached tokens if using a KVConnector.
                     if self.connector is not None:
@@ -694,6 +725,7 @@
                         num_new_tokens,
                         num_new_local_computed_tokens,
                         num_external_computed_tokens,
+                        mamba_tokens_lag,
                     )
                     if num_new_tokens == 0:
                         break