[HiCache] feat: truncate mamba prefetch length to available host KV size (#26945)

2026-06-05 03:01:37 +08:00
parent 8e836e7dc9
commit efe24704b7
1 changed files with 10 additions and 0 deletions
@@ -1735,6 +1735,16 @@ class HiMambaRadixCache(MambaRadixCache):
            prefetch_length,
            self.evict_host,
        )
+        if host_indices is None:
+            # truncate the prefetch length to the page-aligned available host size
+            available_size = self.cache_controller.mem_pool_host.available_size()
+            prefetch_length = available_size - (available_size % self.page_size)
+            if prefetch_length < self.prefetch_threshold:
+                self._release_host_node(last_host_node, release_mamba=False)
+                return
+            new_input_tokens = new_input_tokens[:prefetch_length]
+            host_indices = self.cache_controller.mem_pool_host.alloc(prefetch_length)
+
        if host_indices is None:
            self._release_host_node(last_host_node, release_mamba=False)
            return