fix: preserve HiCache load-back allocator headroom

python/sglang/srt/managers/schedule_policy.py

@@ -499,6 +499,15 @@ class PrefillAdder:
     def ceil_paged_tokens(self, tokens: int) -> int:
         return -(-tokens // self.page_size) * self.page_size
 
+    def _get_available_device_tokens_for_load_back(self) -> int:
+        if self.is_hybrid_swa:
+            return self.token_to_kv_pool_allocator.full_available_size()
+        return self.token_to_kv_pool_allocator.available_size()
+
+    def _get_load_back_mem_quota(self, real_input_tokens: int) -> int:
+        reserve_tokens = real_input_tokens + self.page_size
+        return max(self._get_available_device_tokens_for_load_back() - reserve_tokens, 0)
+
     def budget_state(self):
         if self.rem_total_tokens <= 0 or self.cur_rem_tokens <= 0:
             return AddReqResult.NO_TOKEN
@@ -777,6 +786,7 @@ class PrefillAdder:
                     InitLoadBackParams(
                         last_host_node=req.last_host_node,
                         host_hit_length=req.host_hit_length,
+                        mem_quota=self._get_load_back_mem_quota(real_input_tokens),
                     )
                 )
                 req.prefix_indices = torch.cat([req.prefix_indices, new_indices])