fix(unified-memory): evict Full KV for Mamba byte shortfalls (#36713)

Co-authored-by: Yangmin Li <yangminl@nvidia.com>
Co-authored-by: YAMY <74099316+YAMY1234@users.noreply.github.com>

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

@@ -16,7 +16,7 @@ limitations under the License.
 from __future__ import annotations
 
 import abc
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Protocol
 
 import torch
 
@@ -24,6 +24,20 @@ if TYPE_CHECKING:
     from sglang.srt.mem_cache.memory_pool import KVCache
 
 
+class MambaFullCacheDonor(Protocol):
+    """Allocator capability for reclaiming Full KV on Mamba byte pressure."""
+
+    def flush_deferred_full_frees(self) -> None: ...
+
+    def full_tokens_before_mamba_recheck(self, target_size: int) -> int:
+        """Lower bound on new Full tokens before preparation can help."""
+        ...
+
+    def prepare_mamba_allocation(self, target_size: int) -> None:
+        """Expose layout-specific reclaim so Mamba capacity is queryable."""
+        ...
+
+
 class BaseTokenToKVPoolAllocator(abc.ABC):
     @abc.abstractmethod
     def __init__(
@@ -74,6 +88,10 @@ class BaseTokenToKVPoolAllocator(abc.ABC):
         violation strings, empty when healthy. Static pools have no byte model."""
         return []
 
+    def mamba_full_cache_donor(self) -> MambaFullCacheDonor | None:
+        """Return the shared-pool donor capability, if this allocator has one."""
+        return None
+
     def debug_print(self) -> str:
         return ""
 

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

@@ -23,6 +23,7 @@ import torch
 from torch.profiler import record_function
 
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.allocator.base import MambaFullCacheDonor
 from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
 from sglang.srt.mem_cache.allocator.unified_sub_pool import (
     FloatMultiEndedAllocator,
@@ -30,6 +31,8 @@ from sglang.srt.mem_cache.allocator.unified_sub_pool import (
     _chain_byte_accounting_violations,
     _end_pair_chain,
     _float_open_short_side,
+    _flush_deferred_free_group,
+    _full_tokens_before_mamba_recheck,
     _relieve_for_alloc,
 )
 from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool
@@ -845,6 +848,29 @@ class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator):
                 flt = b
         _float_open_short_side(flt, demand)
 
+    def mamba_full_cache_donor(self) -> MambaFullCacheDonor:
+        return self
+
+    def flush_deferred_full_frees(self) -> None:
+        """Expose grouped composite frees while preserving the group scope."""
+        _flush_deferred_free_group(
+            self,
+            (self.free_group, self.free_page_reps_group, self.full_free_group),
+        )
+
+    def full_tokens_before_mamba_recheck(self, target_size: int) -> int:
+        return _full_tokens_before_mamba_recheck(
+            self.full_attn_allocator, self.mamba_allocator, target_size
+        )
+
+    def prepare_mamba_allocation(self, target_size: int) -> None:
+        """Expose Full reclaim, then move the SWA float away from Mamba."""
+        self.flush_deferred_full_frees()
+        if target_size <= self.mamba_allocator.available_size():
+            return
+        self.full_attn_allocator.flush_for_allocation()
+        _relieve_for_alloc(self.mamba_allocator, target_size)
+
     def mamba_slot_full_token_cost(self) -> int:
         """Full-token-equivalents one mamba/conv slot removes from the shared buffer:
         a tri-pool token costs e_f + e_s bytes, and the quotient is rounded UP."""

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

@@ -23,10 +23,14 @@ import torch
 from torch.profiler import record_function
 
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.allocator.base import MambaFullCacheDonor
 from sglang.srt.mem_cache.allocator.unified_sub_pool import (
     MultiEndedAllocator,
     _chain_byte_accounting_violations,
     _end_pair_chain,
+    _flush_deferred_free_group,
+    _full_tokens_before_mamba_recheck,
+    _relieve_for_alloc,
 )
 from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool
 from sglang.srt.runtime_context import get_parallel
@@ -146,6 +150,26 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
     def full_available_size(self) -> int:
         return self.full_attn_allocator.schedulable_available_size()
 
+    def mamba_full_cache_donor(self) -> MambaFullCacheDonor:
+        return self
+
+    def flush_deferred_full_frees(self) -> None:
+        """Expose grouped Full frees without ending the caller's free group."""
+        _flush_deferred_free_group(self, (self.free_group, self.free_page_reps_group))
+
+    def full_tokens_before_mamba_recheck(self, target_size: int) -> int:
+        return _full_tokens_before_mamba_recheck(
+            self.full_attn_allocator, self.mamba_allocator, target_size
+        )
+
+    def prepare_mamba_allocation(self, target_size: int) -> None:
+        """Make deferred Full reclaim visible to the Mamba capacity view."""
+        self.flush_deferred_full_frees()
+        if target_size > self.mamba_allocator.schedulable_available_size():
+            return
+        if target_size > self.mamba_allocator.available_size():
+            _relieve_for_alloc(self.mamba_allocator, target_size)
+
     def mamba_slot_full_token_cost(self) -> int:
         """Full-token-equivalents of shared-gap bytes ONE mamba state consumes; the
         prefill planner reserves this so admission stays inside the JOINT budget,

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

@@ -31,6 +31,7 @@ from typing import (
     Generic,
     List,
     Optional,
+    Sequence,
     Set,
     Tuple,
     TypeVar,
@@ -194,18 +195,50 @@ def _float_open_short_side(flt, demand) -> None:
 def _relieve_for_alloc(short_pool, need_tokens: int) -> bool:
     """THE shortfall ladder: every allocation shortfall in the unified pool runs
     exactly this, whether a single band's own alloc or a composite's coupled
-    multi-band alloc. `_flush` is called unconditionally -- an eager END no-ops
+    multi-band alloc. The urgent flush is called unconditionally -- an eager END no-ops
     and a FLOAT always has boundary absorption to do -- so the ladder never
     branches on lazy mode, member kind, or layout.
     """
     for m in short_pool._flush_targets():
-        m._flush(urgent=True)
+        m.flush_for_allocation()
     if need_tokens <= short_pool.available_size():
         return True
     short_pool._ask_float_for_room(need_tokens)
     return need_tokens <= short_pool.available_size()
 
 
+def _flush_deferred_free_group(
+    allocator: BaseTokenToKVPoolAllocator,
+    pending_groups: Sequence[Optional[Sequence[torch.Tensor]]],
+) -> None:
+    """Apply queued frees and reopen the caller's free-group scope."""
+    if allocator.free_group is None or not any(pending_groups):
+        return
+    allocator.free_group_end()
+    allocator.free_group_begin()
+
+
+def _full_tokens_before_mamba_recheck(
+    full_allocator: MultiEndedAllocator,
+    mamba_allocator: MultiEndedAllocator,
+    target_size: int,
+) -> int:
+    """Conservative Full-token lower bound for the next Mamba capacity check.
+
+    The current Mamba slot count can hide at most one slot minus one byte of
+    residual room. Subtract that possible residue so this estimate only skips
+    checks that cannot succeed from Full bytes alone. Allocator capacity remains
+    the stop condition after the bound is crossed.
+    """
+    missing_slots = max(0, target_size - mamba_allocator.schedulable_available_size())
+    if missing_slots == 0:
+        return 0
+    mamba_page_bytes = mamba_allocator.entry_bytes_per_page
+    minimum_missing_bytes = (missing_slots - 1) * mamba_page_bytes + 1
+    dcp_size = get_parallel().attn_dcp_size if full_allocator.shards_under_dcp else 1
+    return -(-minimum_missing_bytes * dcp_size // full_allocator.entry_bytes)
+
+
 class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
     """Allocator for one sub-pool over a `UnifiedKVPool`."""
 
@@ -1919,6 +1952,11 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator):
                 self._pending_reuse[latest_event] = (srcs_copy, src_pages_t)
                 self._pending_reuse_pages_cpu.update(srcs_copy)
 
+    def flush_for_allocation(self) -> int:
+        """Public urgent flush used by peer allocation-pressure recovery."""
+        with record_function("MultiEndedAlloc.flush_for_allocation"):
+            return self._flush(urgent=True)
+
     def flush_opportunistic(self) -> int:
         """Public, non-urgent flush at quiescent points; never blocks
         `schedule_stream`. Fast-path the empty state: the scheduler triggers this

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -577,11 +577,62 @@ class UnifiedRadixCache(BasePrefixCache):
 
         request_by_type = self._evict_request_by_type(params)
         available_size_targets = {
-            ct: self._component_available_size(ct) + request_cnt
+            ct: (ct, self._component_available_size(ct) + request_cnt)
             for ct, request_cnt in request_by_type.items()
             if request_cnt > 0
         }
-        return self._evict(params, available_size_targets)
+        allocator = self.token_to_kv_pool_allocator
+        mamba_full_donor = allocator.mamba_full_cache_donor()
+        mamba_target = available_size_targets.get(ComponentType.MAMBA)
+        initial_params = params
+        if mamba_target is not None and mamba_full_donor is not None:
+            # Full KV can supply bytes but cannot recycle Mamba virtual IDs.
+            mamba_id_shortfall = max(
+                0,
+                mamba_target[1]
+                - self.req_to_token_pool.mamba_allocator.available_size(),
+            )
+            initial_params = EvictParams(
+                num_tokens=params.num_tokens,
+                swa_num_tokens=params.swa_num_tokens,
+                mamba_num=mamba_id_shortfall,
+            )
+        result = self._evict(initial_params, available_size_targets)
+
+        if mamba_target is not None and mamba_full_donor is not None:
+            mamba_full_donor.prepare_mamba_allocation(mamba_target[1])
+            mamba_free_ids = self.req_to_token_pool.mamba_allocator.available_size()
+            mamba_capacity = self._component_available_size(ComponentType.MAMBA)
+
+            if mamba_free_ids >= mamba_target[1] and mamba_capacity < mamba_target[1]:
+                full_evictable = self.full_evictable_size()
+                if full_evictable > 0:
+                    donor_result = self._evict(
+                        EvictParams(num_tokens=full_evictable),
+                        {ComponentType.FULL: mamba_target},
+                    )
+                    result.num_tokens_evicted += donor_result.num_tokens_evicted
+                    result.swa_num_tokens_evicted += donor_result.swa_num_tokens_evicted
+                    result.mamba_num_evicted += donor_result.mamba_num_evicted
+
+                # Preserve Mamba-victim recovery if Full cannot fund the target.
+                if (
+                    self._component_available_size(ComponentType.MAMBA)
+                    < mamba_target[1]
+                ):
+                    mamba_evictable = self.mamba_evictable_size()
+                    if mamba_evictable > 0:
+                        fallback_result = self._evict(
+                            EvictParams(mamba_num=mamba_evictable),
+                            {ComponentType.MAMBA: mamba_target},
+                        )
+                        result.num_tokens_evicted += fallback_result.num_tokens_evicted
+                        result.swa_num_tokens_evicted += (
+                            fallback_result.swa_num_tokens_evicted
+                        )
+                        result.mamba_num_evicted += fallback_result.mamba_num_evicted
+
+        return result
 
     @staticmethod
     def _evict_request_by_type(params: EvictParams) -> dict[ComponentType, int]:
@@ -611,7 +662,9 @@ class UnifiedRadixCache(BasePrefixCache):
     def _evict(
         self,
         params: EvictParams,
-        available_size_targets: Optional[dict[ComponentType, int]] = None,
+        available_size_targets: Optional[
+            dict[ComponentType, tuple[ComponentType, int]]
+        ] = None,
     ) -> EvictResult:
         if self.disable:
             return EvictResult()
@@ -712,22 +765,45 @@ class UnifiedRadixCache(BasePrefixCache):
         self,
         request_by_type: dict[ComponentType, int],
         tracker: dict[ComponentType, int],
-        available_size_targets: Optional[dict[ComponentType, int]] = None,
+        available_size_targets: Optional[
+            dict[ComponentType, tuple[ComponentType, int]]
+        ] = None,
     ) -> None:
         # Buffer mode: eviction always wins over queued backup intents — a
         # destroyed victim's intent is stale-swept and the content rewrites
         # after its recompute.
+        last_mamba_donor_check = 0
+        mamba_donor_prepared = False
 
         def target_reached(component_type: ComponentType) -> bool:
+            nonlocal last_mamba_donor_check, mamba_donor_prepared
             if available_size_targets is None:
                 return False
             target = available_size_targets.get(component_type)
-            # Do not compact on every eviction step. Shared allocators include
-            # drainable peer holes here and flush the peer once in alloc().
-            return (
-                target is not None
-                and self._component_available_size(component_type) >= target
-            )
+            if target is None:
+                return False
+            target_component, target_size = target
+            # A Full-leaf cascade can release Mamba or SWA state directly.
+            if self._component_available_size(target_component) >= target_size:
+                return True
+            if (
+                component_type == ComponentType.FULL
+                and target_component == ComponentType.MAMBA
+            ):
+                donor = self.token_to_kv_pool_allocator.mamba_full_cache_donor()
+                assert donor is not None, "Mamba target requires a Full donor"
+                recheck_after = (
+                    1
+                    if mamba_donor_prepared
+                    else donor.full_tokens_before_mamba_recheck(target_size)
+                )
+                if tracker[component_type] - last_mamba_donor_check < recheck_after:
+                    return False
+                donor.prepare_mamba_allocation(target_size)
+                last_mamba_donor_check = tracker[component_type]
+                mamba_donor_prepared = True
+            # Schedulable capacity includes donor holes that allocation can compact.
+            return self._component_available_size(target_component) >= target_size
 
         for ct in self.tree_components:
             request_cnt = request_by_type[ct]
@@ -737,16 +813,14 @@ class UnifiedRadixCache(BasePrefixCache):
                 continue
             self.tree_core.evict_device_start(ct, request_cnt)
             try:
-                while not target_reached(ct):
+                while True:
                     node_id, made_progress = self._evict_device_next_node(ct, tracker)
                     if node_id is None:
-                        if made_progress:
-                            # Internal tombstone frees are now allocator-visible;
-                            # recheck the allocation target before walking again.
-                            continue
-                        break
-                    backup_kv = self._evict_device_leaf(node_id, tracker)
-                    if backup_kv is not None:
+                        if not made_progress:
+                            break
+                    else:
+                        backup_kv = self._evict_device_leaf(node_id, tracker)
+                    if node_id is not None and backup_kv is not None:
                         # Deferred demote: run the D->H backup, demote only on success.
                         written = self._execute_and_commit_kv_backup(
                             backup_kv, write_back=True
@@ -768,6 +842,8 @@ class UnifiedRadixCache(BasePrefixCache):
                                 "until host space frees",
                                 node_id,
                             )
+                    if target_reached(ct):
+                        break
             finally:
                 self.tree_core.evict_device_end(ct)