minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

Skip mamba lock during decoding (#32228)

Browse Source
This commit is contained in:
Ke Bao
2026-07-29 21:53:28 +08:00
committed by GitHub
parent d004a15a3e
commit 50b029257f
11 changed files with 381 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
+5
View File
@@ -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
+7 -1
View File
@@ -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
+3
View File
@@ -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
+18 -2
View File
@@ -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
+4 -1
View File
@@ -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
+23 -4
View File
@@ -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
+9 -3
View File
@@ -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
+38 -13
View File
@@ -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
+13 -7
View File
@@ -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