Skip mamba lock during decoding (#32228)

python/sglang/srt/environ.py

@@ -845,6 +845,11 @@ class Envs:
     # Kill-switch for the fused per-slot conv clear/copy kernel (MambaPool);
     # falls back to the per-conv-type Python loop.
     SGLANG_DISABLE_FUSED_MAMBA_SLOT_OPS = EnvBool(False)
+    # Opt-in: on the unified radix tree, leave the matched-prefix mamba evictable
+    # during decode (it is already COW'd to the request's own slot) and shrink the
+    # mamba pool ratio accordingly. Frees one resident slot per running request,
+    # raising max_running_requests. Off = original locking + ratio (escape hatch).
+    SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK = EnvBool(False)
 
     # Unified Radix Tree
     SGLANG_ENABLE_UNIFIED_RADIX_TREE = EnvBool(False)

python/sglang/srt/managers/schedule_batch.py

@@ -907,6 +907,9 @@ class Req(ReqDllmMixin):
         self.swa_uuid_for_lock: Optional[int] = None
         # Whether the prefill-time SWA tree lock has been released early
         self.swa_prefix_lock_released: bool = False
+        # per-component nodes this req skipped locking (e.g. mamba on the decode
+        # hold, already COW'd), so their dec releases only what it took.
+        self.skip_lock_node_ids: dict = {}
         # The prefix length that is inserted into the tree cache
         self.cache_protected_len: int = 0
 
@@ -1522,6 +1525,7 @@ class Req(ReqDllmMixin):
         self.num_matched_prefix_tokens = 0
         self.swa_uuid_for_lock = None
         self.swa_prefix_lock_released = False
+        self.skip_lock_node_ids = {}
         self.extend_range = None
         self.dllm_initialized = False
         self.is_retracted = True
@@ -3130,7 +3134,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
                         and req.decode_batch_idx >= sliding_window_size
                     ):
                         self.tree_cache.dec_swa_lock_only(
-                            req.last_node, req.swa_uuid_for_lock
+                            req.last_node,
+                            req.swa_uuid_for_lock,
+                            skip_lock_node_ids=req.skip_lock_node_ids,
                         )
                         req.swa_prefix_lock_released = True
                 elif self.forward_mode.is_extend() and self.tree_cache.is_chunk_cache():

python/sglang/srt/managers/schedule_policy.py

@@ -822,6 +822,9 @@ class PrefillAdder:
         result = self.tree_cache.inc_lock_ref(req.last_node)
         if self.is_hybrid_swa:
             req.swa_uuid_for_lock = result.swa_uuid_for_lock
+        # match locks this node's components, so clear any stale skip set
+        # carried from a previous scheduling of this req.
+        req.skip_lock_node_ids = {}
 
     def add_dllm_staging_req(self, req: Req):
         assert self.dllm_config is not None

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -97,11 +97,17 @@ def _should_enable_lazy_compaction() -> bool:
     return not envs.SGLANG_DISABLE_LAZY_COMPACTION.get()
 
 
-# the ratio of mamba cache pool size to max_running_requests
+# base ratio of mamba pool size to max_running_requests. Under
+# SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK the decode-time skip frees one resident slot
+# per running request, so the base drops by 1 (overlap 5->4, lazy 4->3). no_buffer
+# stays at effective 3 either way: its binding limit is the prefill->decode peak,
+# which the decode-time drop does not shrink.
 MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO = 3
+MAMBA_CACHE_BASE_RATIO_DROP_ON_SKIP = 1
 MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP = 2
 MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP_LAZY = 1
 MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP = 1
+MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_BUFFER = 1
 
 if TYPE_CHECKING:
     from sglang.srt.distributed.parallel_state_wrapper import ParallelState
@@ -1578,6 +1584,11 @@ class KVCacheConfigurator:
         if self.server_args.disable_radix_cache:
             return 1
 
+        skip_decode_lock = envs.SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK.get()
+        base = MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO - (
+            MAMBA_CACHE_BASE_RATIO_DROP_ON_SKIP if skip_decode_lock else 0
+        )
+
         additional_ratio = 0
         if self.server_args.enable_mamba_extra_buffer():
             # ping-pong buffer size is 2 when overlap schedule is on, 1 otherwise.
@@ -1592,8 +1603,13 @@ class KVCacheConfigurator:
                     not self.server_args.enable_mamba_extra_buffer_lazy()
                 ), "Lazy extra buffer requires overlap schedule (--disable-overlap-schedule is incompatible)"
                 additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP
+        elif skip_decode_lock:
+            # no_buffer under skip: add the base drop back so effective stays 3,
+            # the prefill->decode peak needs ~3 slots/req and this leaf-only mode
+            # has no ping-pong to absorb it.
+            additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_BUFFER
 
-        return MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO + additional_ratio
+        return base + additional_ratio
 
     def _apply_token_constraints(self, token_capacity: int) -> int:
         """Apply external constraints to token capacity: user cap, PP sync.

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -780,7 +780,10 @@ class SWARadixCache(KVCacheEventMixin, BasePrefixCache):
         return DecLockRefResult()
 
     def dec_swa_lock_only(
-        self, node: TreeNode, swa_uuid_for_lock: Optional[int] = None
+        self,
+        node: TreeNode,
+        swa_uuid_for_lock: Optional[int] = None,
+        skip_lock_node_ids: Optional[dict] = None,  # unused, signature parity only
     ):
         """
         Decrement only the swa_lock_ref (and swa_protected_size_) along the chain

python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py

@@ -442,10 +442,21 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
         """Drop a tree node from the arena."""
         self._node_arena.pop(node.id, None)
 
-    def inc_lock_ref(self, node_id: NodeId) -> IncLockRefResult:
+    def inc_lock_ref(
+        self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = ()
+    ) -> IncLockRefResult:
         node = self.node_by_id(node_id)
         result = IncLockRefResult()
         for component in self.components:
+            if component.component_type in skip_lock_components:
+                # Leave this component's value evictable and record every
+                # non-root node (incl tombstones) so the matching dec skips a
+                # lock we never took, which may be another req's on a shared node.
+                if node is not self.root_node:
+                    result.skip_lock_node_ids.setdefault(
+                        component.component_type, set()
+                    ).add(node.id)
+                continue
             result = component.acquire_component_lock(node=node, result=result)
         self._update_evictable_leaf_sets(node)
         return result
@@ -466,7 +477,10 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
         return DecLockRefResult()
 
     def dec_swa_lock_only(
-        self, node_id: NodeId, swa_uuid_for_lock: Optional[int]
+        self,
+        node_id: NodeId,
+        swa_uuid_for_lock: Optional[int],
+        skip_lock_node_ids: Optional[dict] = None,
     ) -> DecSwaLockOnlyResult:
         """Early-release the SWA portion of a request's tree lock, plus any
         strictly-lower-priority locks (e.g. Mamba) co-located on the node."""
@@ -479,9 +493,14 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
             node, swa_uuid_for_lock, result.device_frees, result.host_frees
         )
 
-        # Drop strictly-lower-priority locks (e.g. Mamba) co-located on the node.
+        # Drop strictly-lower-priority locks (e.g. Mamba) co-located on the node,
+        # honoring skip ids so we don't drop a lock a partial inc never took
+        # (matters for FULL+SWA+MAMBA models, e.g. Inkling).
         swa_priority = swa_component.eviction_priority(is_leaf=False)
-        dec_params = DecLockRefParams(swa_uuid_for_lock=swa_uuid_for_lock)
+        dec_params = DecLockRefParams(
+            swa_uuid_for_lock=swa_uuid_for_lock,
+            skip_lock_node_ids=skip_lock_node_ids or {},
+        )
         for comp in self.components:
             if comp.eviction_priority(is_leaf=False) < swa_priority:
                 comp.release_component_lock(node, dec_params)

python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py

@@ -166,8 +166,11 @@ class UnifiedTreeCoreInterface(KVCacheEventMixin, ABC):
         ...
 
     @abstractmethod
-    def inc_lock_ref(self, node_id: NodeId) -> IncLockRefResult:
-        """Bump the reference count on a node's component locks."""
+    def inc_lock_ref(
+        self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = ()
+    ) -> IncLockRefResult:
+        """Bump the reference count on a node's component locks, leaving any
+        component in skip_lock_components evictable and recorded in the result."""
         ...
 
     @abstractmethod
@@ -182,7 +185,10 @@ class UnifiedTreeCoreInterface(KVCacheEventMixin, ABC):
 
     @abstractmethod
     def dec_swa_lock_only(
-        self, node_id: NodeId, swa_uuid_for_lock: Optional[int]
+        self,
+        node_id: NodeId,
+        swa_uuid_for_lock: Optional[int],
+        skip_lock_node_ids: Optional[dict] = None,
     ) -> DecSwaLockOnlyResult:
         """Decrease only the SWA (and lower-priority co-located) reference
         counts; the result carries the freed slots."""

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -4,7 +4,7 @@ import logging
 import threading
 import time
 from queue import Empty, Queue
-from typing import TYPE_CHECKING, Iterator, NamedTuple, Optional, TypeVar
+from typing import TYPE_CHECKING, Iterator, NamedTuple, Optional, Sequence, TypeVar
 
 import torch
 
@@ -535,13 +535,15 @@ class UnifiedRadixCache(BasePrefixCache):
             finally:
                 self.tree_core.evict_device_end(ct)
 
-    def inc_lock_ref(self, node_id: NodeId) -> IncLockRefResult:
+    def inc_lock_ref(
+        self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = ()
+    ) -> IncLockRefResult:
         result = self.session.try_inc_lock_ref(node_id)
         if result is not None:
             return result
         if self.disable:
             return IncLockRefResult()
-        return self.tree_core.inc_lock_ref(node_id)
+        return self.tree_core.inc_lock_ref(node_id, skip_lock_components)
 
     def dec_lock_ref(
         self,
@@ -556,14 +558,29 @@ class UnifiedRadixCache(BasePrefixCache):
             return DecLockRefResult()
         return self.tree_core.dec_lock_ref(node_id, params, skip_swa)
 
+    def _dec_req_lock(self, req: Req, *, skip_swa: bool = False) -> None:
+        """Release the tree lock a request holds on its last_node, honoring the
+        components it skipped locking so it never drops a lock it never took."""
+        self.dec_lock_ref(
+            req.last_node,
+            DecLockRefParams(
+                swa_uuid_for_lock=req.swa_uuid_for_lock,
+                skip_lock_node_ids=req.skip_lock_node_ids,
+            ),
+            skip_swa=skip_swa,
+        )
+
     def dec_swa_lock_only(
         self,
         node_id: NodeId,
         swa_uuid_for_lock: Optional[int] = None,
+        skip_lock_node_ids: Optional[dict] = None,
     ) -> None:
         if self.disable:
             return
-        result = self.tree_core.dec_swa_lock_only(node_id, swa_uuid_for_lock)
+        result = self.tree_core.dec_swa_lock_only(
+            node_id, swa_uuid_for_lock, skip_lock_node_ids
+        )
         self._free_values(result.device_frees, result.host_frees)
 
     def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult:
@@ -649,11 +666,7 @@ class UnifiedRadixCache(BasePrefixCache):
                 start_pos=req.cache_protected_len,
             )
 
-        self.dec_lock_ref(
-            req.last_node,
-            DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)),
-            skip_swa=getattr(req, "swa_prefix_lock_released", False),
-        )
+        self._dec_req_lock(req, skip_swa=req.swa_prefix_lock_released)
 
         # cleanup
         for comp in self._components_tuple:
@@ -739,11 +752,21 @@ class UnifiedRadixCache(BasePrefixCache):
             new_indices[req.cache_protected_len :],
         )
 
-        self.dec_lock_ref(
-            req.last_node,
-            DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)),
+        self._dec_req_lock(req)
+        # Opt-in: leave the matched-prefix mamba evictable during decode (it is
+        # already COW'd to the request's own slot, never read from this node again).
+        # Safe only because any future COW source is the COWing request's own
+        # admission-locked last_node (recorded only if still present, locked before
+        # the next alloc) -- not this evictable node. A scheduler that matched a
+        # whole batch before locking would break that. Off = original full lock.
+        skip_lock_components = (
+            (ComponentType.MAMBA,)
+            if envs.SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK.get()
+            else ()
+        )
+        lock_result = self.inc_lock_ref(
+            new_last_node, skip_lock_components=skip_lock_components
         )
-        lock_result = self.inc_lock_ref(new_last_node)
 
         # Update req fields
         if len(new_indices) < len(kv_indices_orig):
@@ -755,6 +778,8 @@ class UnifiedRadixCache(BasePrefixCache):
         req.cache_protected_len = len(new_indices)
         req.last_node = new_last_node
         req.swa_uuid_for_lock = lock_result.swa_uuid_for_lock
+        # carry the skip set so this node's dec releases only what we locked
+        req.skip_lock_node_ids = lock_result.skip_lock_node_ids
         # The rematch acquired a new SWA prefix lock.
         req.swa_prefix_lock_released = False
 

python/sglang/srt/session/streaming_session.py

@@ -52,6 +52,9 @@ class SessionSlot:
     last_node: Any = None
     cache_protected_len: int = 0
     swa_uuid_for_lock: Optional[str] = None
+    # components the first req skipped locking on last_node, so release dec
+    # releases only what it took (may share the node with another req).
+    skip_lock_node_ids: dict = field(default_factory=dict)
 
     # Mamba states
     mamba_pool_idx: Any = None
@@ -75,6 +78,7 @@ class SessionSlot:
             self.last_node = req.last_node
             self.cache_protected_len = req.cache_protected_len
             self.swa_uuid_for_lock = req.swa_uuid_for_lock
+            self.skip_lock_node_ids = req.skip_lock_node_ids
 
         self.mamba_pool_idx = req.mamba_pool_idx
         self.mamba_ping_pong_track_buffer = req.mamba_ping_pong_track_buffer
@@ -104,6 +108,7 @@ class SessionSlot:
         req.kv_committed_len = self.kv_committed_len
         req.kv = copy.copy(self.kv)
         req.swa_uuid_for_lock = self.swa_uuid_for_lock
+        req.skip_lock_node_ids = self.skip_lock_node_ids
 
         req.mamba_pool_idx = self.mamba_pool_idx
         req.mamba_ping_pong_track_buffer = self.mamba_ping_pong_track_buffer
@@ -306,6 +311,7 @@ class StreamingSession(BasePrefixCache):
                     last_node=req.last_node,
                     cache_protected_len=req.cache_protected_len,
                     swa_uuid_for_lock=req.swa_uuid_for_lock,
+                    skip_lock_node_ids=req.skip_lock_node_ids,
                     mamba_pool_idx=req.mamba_pool_idx,
                     mamba_ping_pong_track_buffer=req.mamba_ping_pong_track_buffer,
                 )
@@ -418,13 +424,13 @@ class StreamingSession(BasePrefixCache):
         )
 
         if lock_node is not None:
-            if slot.swa_uuid_for_lock is not None:
-                self.inner.dec_lock_ref(
-                    lock_node,
-                    DecLockRefParams(swa_uuid_for_lock=slot.swa_uuid_for_lock),
-                )
-            else:
-                self.inner.dec_lock_ref(lock_node)
+            self.inner.dec_lock_ref(
+                lock_node,
+                DecLockRefParams(
+                    swa_uuid_for_lock=slot.swa_uuid_for_lock,
+                    skip_lock_node_ids=slot.skip_lock_node_ids,
+                ),
+            )
 
         if slot.is_holding_kv:
             start = protected_len

test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py

@@ -0,0 +1,227 @@
+"""CPU-only unit tests for the mamba pool ratio vs the prefill->decode peak.
+
+Pins the sizing invariant behind MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO:
+at the first cache_unfinished_req, a request still holds its admission-locked
+matched-prefix mamba (protected) plus its own COW slot, and then allocates a
+donated slot. With N distinct-prefix requests that peak is N own + N locked +
+1 donated. An effective ratio of 2 (pool = 2N) leaves no evictable victim and
+the donated alloc asserts; ratio 3 (pool = 3N) has headroom. Once decode's
+skip_mamba leaves the matched prefix evictable, even ratio 2 recovers via
+eviction -- which is why the peak, not the decode steady state, sets the floor.
+"""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.mem_cache.base_prefix_cache import (
+    EvictParams,
+    IncLockRefResult,
+)
+from sglang.srt.mem_cache.unified_cache.components.mamba_component import MambaComponent
+from sglang.srt.mem_cache.unified_cache.components.tree_component import ComponentType
+from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
+from sglang.srt.mem_cache.unified_radix_cache import UnifiedTreeNode
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+N = 4  # concurrent distinct-prefix requests
+
+
+class _BoundedMambaAllocator:
+    """Fixed-capacity slot allocator; alloc returns None once exhausted."""
+
+    def __init__(self, size: int):
+        self.free_ids = list(range(size))
+
+    def alloc(self, n: int):
+        if len(self.free_ids) < n:
+            return None
+        return torch.tensor([self.free_ids.pop() for _ in range(n)], dtype=torch.int64)
+
+    def free(self, value: torch.Tensor):
+        self.free_ids.extend(int(v) for v in value.tolist())
+
+
+class _RatioCache:
+    tree_components = (ComponentType.FULL, ComponentType.MAMBA)
+
+    def __init__(self, pool_size: int):
+        self.root_node = UnifiedTreeNode(self.tree_components)
+        self.allocator = _BoundedMambaAllocator(pool_size)
+        self.req_to_token_pool = SimpleNamespace(mamba_allocator=self.allocator)
+        self.component_evictable_size_ = {ComponentType.MAMBA: 0}
+        self.component_protected_size_ = {ComponentType.MAMBA: 0}
+        self.prefix_nodes = []
+
+    def evict(self, params: EvictParams):
+        # Reclaim up to mamba_num evictable (unlocked) prefix snapshots, mirroring
+        # what the real tree eviction can hand back under mamba pressure.
+        need = params.mamba_num
+        for node in list(self.prefix_nodes):
+            if need <= 0:
+                break
+            cd = node.component_data[ComponentType.MAMBA]
+            if cd.lock_ref == 0 and cd.value is not None:
+                self.allocator.free(cd.value)
+                self.component_evictable_size_[ComponentType.MAMBA] -= len(cd.value)
+                cd.value = None
+                self.prefix_nodes.remove(node)
+                need -= 1
+
+
+def _build_peak(pool_size: int, lock_prefixes: bool):
+    """N own slots + N matched-prefix snapshots, then return the component ready
+    to allocate one donated slot. Prefix snapshots are locked (protected,
+    prefill peak) or left evictable (decode steady state after skip_mamba)."""
+    cache = _RatioCache(pool_size)
+    component = object.__new__(MambaComponent)
+    component.cache = cache
+    # The TreeCore owns the tree member-var state the component reads through.
+    component.tree_core = cache
+    component.component_type = ComponentType.MAMBA
+
+    owned = [cache.allocator.alloc(1) for _ in range(N)]
+    assert all(s is not None for s in owned)
+
+    for _ in range(N):
+        node = UnifiedTreeNode(cache.tree_components)
+        slot = cache.allocator.alloc(1)
+        assert slot is not None
+        node.component_data[ComponentType.MAMBA].value = slot
+        cache.component_evictable_size_[ComponentType.MAMBA] += len(slot)
+        cache.prefix_nodes.append(node)
+        if lock_prefixes:
+            component.acquire_component_lock(node, IncLockRefResult())
+
+    return component, cache, owned
+
+
+class TestMambaRatioEnvGate(unittest.TestCase):
+    """SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK gates the pool ratio: off restores the
+    original base 3 (overlap 5, lazy 4, no_buffer 3), on drops the base to 2
+    (overlap 4, lazy 3) while no_buffer stays 3. Guards the flag wiring so the
+    ratio can never drift out of sync with whether the decode lock is skipped."""
+
+    @staticmethod
+    def _ratio(*, extra_buffer, lazy, disable_overlap, skip):
+        from sglang.srt.environ import envs
+        from sglang.srt.mem_cache.kv_cache_configurator import KVCacheConfigurator
+
+        server_args = SimpleNamespace(
+            disable_radix_cache=False,
+            disable_overlap_schedule=disable_overlap,
+            enable_mamba_extra_buffer=lambda: extra_buffer,
+            enable_mamba_extra_buffer_lazy=lambda: lazy,
+        )
+        fake = SimpleNamespace(server_args=server_args)
+        with envs.SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK.override(skip):
+            return KVCacheConfigurator._calculate_mamba_ratio(fake)
+
+    def test_flag_off_restores_original_ratios(self):
+        r = lambda **kw: self._ratio(skip=False, **kw)
+        self.assertEqual(
+            r(extra_buffer=False, lazy=False, disable_overlap=True), 3
+        )  # no_buffer
+        self.assertEqual(
+            r(extra_buffer=True, lazy=True, disable_overlap=False), 4
+        )  # lazy
+        self.assertEqual(
+            r(extra_buffer=True, lazy=False, disable_overlap=False), 5
+        )  # overlap
+
+    def test_flag_on_drops_base_but_keeps_no_buffer(self):
+        r = lambda **kw: self._ratio(skip=True, **kw)
+        self.assertEqual(
+            r(extra_buffer=False, lazy=False, disable_overlap=True), 3
+        )  # no_buffer
+        self.assertEqual(
+            r(extra_buffer=True, lazy=True, disable_overlap=False), 3
+        )  # lazy
+        self.assertEqual(
+            r(extra_buffer=True, lazy=False, disable_overlap=False), 4
+        )  # overlap
+
+
+class _RecordingComp:
+    """Fake tree component: records the dec params it is asked to release with."""
+
+    def __init__(self, component_type, priority):
+        self.component_type = component_type
+        self._priority = priority
+        self.released = []
+
+    def eviction_priority(self, is_leaf):
+        return self._priority
+
+    def release_component_lock(self, node, params):
+        self.released.append(params)
+
+    def release_window_lock(  # SWA only
+        self, node, swa_uuid_for_lock, device_frees, host_frees
+    ):
+        pass
+
+
+class TestDecSwaLockSkip(unittest.TestCase):
+    """dec_swa_lock_only early-releases SWA plus co-located lower-tier (Mamba)
+    locks. On a full-only-locked node (decode skip) it must thread the skip set
+    into that lower-tier release, else it drops a mamba lock it never took --
+    another request's, on a shared FULL+SWA+MAMBA node (Inkling). Guards the
+    contract without booting a 3-component model."""
+
+    def test_threads_skip_ids_into_lower_tier_release(self):
+        # internal-node priority: full=2 > swa=1 > mamba=0
+        full = _RecordingComp(ComponentType.FULL, 2)
+        swa = _RecordingComp(ComponentType.SWA, 1)
+        mamba = _RecordingComp(ComponentType.MAMBA, 0)
+        node = SimpleNamespace(id=7)
+        tree_core = SimpleNamespace(
+            components=(full, swa, mamba),
+            components_by_type={ComponentType.SWA: swa},
+            node_by_id=lambda node_id: node,
+        )
+
+        UnifiedTreeCore.dec_swa_lock_only(
+            tree_core,
+            node.id,
+            swa_uuid_for_lock=None,
+            skip_lock_node_ids={ComponentType.MAMBA: {7}},
+        )
+
+        # mamba (below swa) is released, honoring the skip set
+        self.assertEqual(len(mamba.released), 1)
+        self.assertEqual(
+            mamba.released[0].skip_lock_node_ids.get(ComponentType.MAMBA), {7}
+        )
+        # full (above swa) is never touched
+        self.assertEqual(full.released, [])
+
+
+class TestMambaDonatedAllocRatio(unittest.TestCase):
+    def test_prefill_peak_ratio2_exhausts_pool(self):
+        # pool = 2N, all N prefixes admission-locked: no evictable victim.
+        component, _, _ = _build_peak(pool_size=2 * N, lock_prefixes=True)
+        with self.assertRaisesRegex(AssertionError, "Can not alloc mamba cache"):
+            component._alloc_mamba_slot()
+
+    def test_prefill_peak_ratio3_has_headroom(self):
+        # pool = 3N: N free slots remain after own + locked prefix.
+        component, cache, _ = _build_peak(pool_size=3 * N, lock_prefixes=True)
+        slot = component._alloc_mamba_slot()
+        self.assertIsNotNone(slot)
+        self.assertEqual(cache.component_protected_size_[ComponentType.MAMBA], N)
+
+    def test_decode_steady_evictable_prefix_ratio2_ok(self):
+        # pool = 2N but the matched prefixes are evictable (skip_mamba on decode):
+        # eviction reclaims a victim, so even ratio 2 serves the donated alloc.
+        component, cache, _ = _build_peak(pool_size=2 * N, lock_prefixes=False)
+        slot = component._alloc_mamba_slot()
+        self.assertIsNotNone(slot)
+        self.assertEqual(len(cache.prefix_nodes), N - 1)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_streaming_session_unit.py

@@ -25,6 +25,7 @@ class _FakeInnerCache:
         self.page_size = page_size
         self.match_results = list(match_results or [])
         self.dec_lock_ref_calls = []
+        self.dec_lock_ref_params = []
 
     def cache_finished_req(self, *args, **kwargs):
         raise AssertionError("Streaming requests should not delegate to inner cache")
@@ -36,6 +37,7 @@ class _FakeInnerCache:
 
     def dec_lock_ref(self, node, *args, **kwargs):
         self.dec_lock_ref_calls.append(node)
+        self.dec_lock_ref_params.append(args[0] if args else kwargs.get("params"))
 
     def supports_mamba(self):
         return False
@@ -67,6 +69,7 @@ class _FakeReq:
         self.last_node = None
         self.cache_protected_len = 0
         self.swa_uuid_for_lock = None
+        self.skip_lock_node_ids = {}
         self.mamba_pool_idx = None
         self.mamba_ping_pong_track_buffer = None
         self.mamba_next_track_idx = None
@@ -185,6 +188,37 @@ def test_nth_mid_abort_nukes_session_slot():
     assert req.req_pool_idx is None
 
 
+def test_release_session_threads_mamba_skip_ids():
+    """release_session must forward the slot's skip_lock_node_ids to
+    dec_lock_ref. The first req's last_node may be full-only-locked (mamba
+    skipped at inc), so without the skip set the release would drop a mamba
+    lock the session never took -- another request's, on a shared node."""
+    from sglang.srt.mem_cache.unified_cache.components import ComponentType
+
+    req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
+    req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
+    allocator = _FakeAllocator()
+    inner = _FakeInnerCache(req_to_token_pool, allocator, page_size=1)
+    tree_cache = StreamingSession(inner)
+
+    lock_node = SimpleNamespace(id=42)
+    tree_cache.slots["session-a"] = SessionSlot(
+        req_pool_idx=0,
+        kv_committed_len=50,
+        kv=SimpleNamespace(kv_allocated_len=50, swa_evicted_seqlen=0),
+        last_node=lock_node,
+        cache_protected_len=0,
+        skip_lock_node_ids={ComponentType.MAMBA: {42}},
+    )
+
+    tree_cache.release_session("session-a")
+
+    assert inner.dec_lock_ref_calls == [lock_node]
+    params = inner.dec_lock_ref_params[0]
+    assert params is not None
+    assert params.skip_lock_node_ids.get(ComponentType.MAMBA) == {42}
+
+
 # Shrink tests removed: streaming sessions are append-only after the
 # rollback fix in session_controller (rollback_aborted_req).  The shrink
 # code path in cache_finished_req no longer exists.