From 50b029257fb819b4673efb441cfdd048392571aa Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Wed, 29 Jul 2026 21:53:28 +0800 Subject: [PATCH] Skip mamba lock during decoding (#32228) --- python/sglang/srt/environ.py | 5 + python/sglang/srt/managers/schedule_batch.py | 8 +- python/sglang/srt/managers/schedule_policy.py | 3 + .../srt/mem_cache/kv_cache_configurator.py | 20 +- .../sglang/srt/mem_cache/swa_radix_cache.py | 5 +- .../unified_cache/unified_tree_core.py | 27 ++- .../unified_tree_core_interface.py | 12 +- .../srt/mem_cache/unified_radix_cache.py | 51 +++- .../sglang/srt/session/streaming_session.py | 20 +- .../test_mamba_donated_alloc_ratio.py | 227 ++++++++++++++++++ .../mem_cache/test_streaming_session_unit.py | 34 +++ 11 files changed, 381 insertions(+), 31 deletions(-) create mode 100644 test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index b125d9add..d0e89f3c2 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index f3e240c1a..4b3548de3 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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(): diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 21db31bc0..3cc7ce842 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 264cbc337..c30c474c7 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/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. diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index ef1a632c6..31cedca76 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index cf783205f..ff89de1ce 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/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) diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index a06b0517a..cbd9910ff 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/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.""" diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 234225c89..e10561f24 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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 diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index cb2a012e4..28cbd4cc6 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/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 diff --git a/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py b/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py new file mode 100644 index 000000000..cba3c2543 --- /dev/null +++ b/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() diff --git a/test/registered/unit/mem_cache/test_streaming_session_unit.py b/test/registered/unit/mem_cache/test_streaming_session_unit.py index dbaa775d2..46cd27433 100644 --- a/test/registered/unit/mem_cache/test_streaming_session_unit.py +++ b/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.