From efe24704b70d5a7876b523e985f95be5af94153a Mon Sep 17 00:00:00 2001
From: shuwenn <47200617+alphabetc1@users.noreply.github.com>
Date: Fri, 5 Jun 2026 03:01:37 +0800
Subject: [PATCH] [HiCache] feat: truncate mamba prefetch length to available
 host KV size (#26945)

---
 python/sglang/srt/mem_cache/hi_mamba_radix_cache.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
index c3e4c7a80..6b5b904f1 100644
--- a/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
+++ b/python/sglang/srt/mem_cache/hi_mamba_radix_cache.py
@@ -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