[Fix] unified-memory: Bound joint reclaim and preserve cached prefixes

python/sglang/srt/mem_cache/allocation.py

@@ -188,8 +188,15 @@ def alloc_paged_token_slots_extend(
     # availability floor must cover every request's ceil(K_i/N) rounding,
     # matching the N*ps per-request admission reserve).
     allocator = tree_cache.token_to_kv_pool_allocator
-    num_tokens = extend_num_tokens + len(seq_lens_cpu) * (
-        allocator.page_size * page_interleave_shard_size(allocator)
+    shard_size = page_interleave_shard_size(allocator)
+    conservative_num_tokens = extend_num_tokens + len(seq_lens_cpu) * (
+        allocator.page_size * shard_size
+    )
+    num_tokens = allocator.get_extend_allocation_demand(
+        prefix_lens_cpu,
+        seq_lens_cpu,
+        conservative_num_tokens=conservative_num_tokens,
+        shard_size=shard_size,
     )
     evict_from_tree_cache(tree_cache, num_tokens)
 

python/sglang/srt/mem_cache/allocator/base.py

@@ -141,6 +141,17 @@ def prealloc_fits(
             swa_budget_tokens is None or swa_tokens <= swa_budget_tokens
         )
 
+    def get_extend_allocation_demand(
+        self,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        *,
+        conservative_num_tokens: int,
+        shard_size: int,
+    ) -> int:
+        """Token demand for eviction before paged extend."""
+        return conservative_num_tokens
+
     def evict_to_free_tokens(self, tree_cache, num_tokens: int) -> bool | None:
         """Evict unlocked prefix-cache entries until this allocator can serve
         ``num_tokens`` or nothing evictable remains.