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