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

[HiCache] Replace skip_lock_node_ids with a segment lock protocol (#36848)

Browse Source
This commit is contained in:
Zhiqiang Xie
2026-09-09 14:55:51 -07:00
committed by GitHub
parent a84ffd1326
commit beaf3d9252
41 changed files with 2116 additions and 1146 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/decode.py
+10 -8
View File
@@ -317,6 +317,8 @@ class DecodeRequest:
prefix_match: Optional[DecodePrefixMatch] = None
hicache_restored_kv_indices: Optional[torch.Tensor] = None
hicache_restored_node: Any = None
# Receipt for the inc_lock_ref held on hicache_restored_node.
hicache_restore_lock_receipt: Optional[DecLockRefParams] = None
hicache_load_consumer_index: int = -1
hicache_restore_status: HiCacheRestoreResult = HiCacheRestoreResult.PENDING
@@ -426,12 +428,11 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
)
def _release_matched_prefix_lock(self, req: Req) -> None:
params = DecLockRefParams(swa_uuid_for_lock=req.swa_uuid_for_lock)
if req.swa_prefix_lock_released:
self.tree_cache.dec_lock_ref(req.last_node, params, skip_swa=True)
self.tree_cache.dec_lock_ref(req.last_node, req.lock_receipt, skip_swa=True)
req.swa_prefix_lock_released = False
else:
self.tree_cache.dec_lock_ref(req.last_node, params)
self.tree_cache.dec_lock_ref(req.last_node, req.lock_receipt)
def _reclaim_swa_tail_capacity(
self, swa_tail_len: int, req_id: str
@@ -676,9 +677,11 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
include_req=True,
)
# Keep aggregated scheduling semantics while preserving the SWA lock
# boundary needed for the matching dec_lock_ref.
lock_result = self.tree_cache.inc_lock_ref(result.last_device_node)
req.swa_uuid_for_lock = lock_result.swa_uuid_for_lock
# boundary needed for the matching dec_lock_ref; the full receipt
# travels on the req so every later release mirrors this acquire.
req.lock_receipt = self.tree_cache.inc_lock_ref(
result.last_device_node
).to_dec_params()
return self._build_decode_prefix_match(req, result)
def _resolve_prefill_dp_rank(self, req: Req) -> Optional[int]:
@@ -1239,8 +1242,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
and hasattr(self.tree_cache, "dec_swa_lock_only")
):
self.tree_cache.dec_swa_lock_only(
decode_req.req.last_node,
decode_req.req.swa_uuid_for_lock,
decode_req.req.last_node, decode_req.req.lock_receipt
)
decode_req.req.swa_prefix_lock_released = True
python/sglang/srt/disaggregation/decode_hicache_mixin.py
+35 -6
View File
@@ -11,7 +11,9 @@ import torch
from sglang.srt.disaggregation.base import KVPoll
from sglang.srt.managers.schedule_policy import match_prefix_for_req
from sglang.srt.mem_cache.base_prefix_cache import InitLoadBackParams
from sglang.srt.mem_cache.base_prefix_cache import (
InitLoadBackParams,
)
if TYPE_CHECKING:
from sglang.srt.disaggregation.decode import DecodeRequest
@@ -184,8 +186,12 @@ class DecodeHiCacheTransferMixin:
):
self.tree_cache.release_aborted_request(decode_req.req.rid)
if decode_req.hicache_restored_node is not None:
self.tree_cache.dec_lock_ref(decode_req.hicache_restored_node)
self.tree_cache.dec_lock_ref(
decode_req.hicache_restored_node,
decode_req.hicache_restore_lock_receipt,
)
decode_req.hicache_restored_node = None
decode_req.hicache_restore_lock_receipt = None
def _try_hicache_queue_load_back(self, dr: DecodeRequest) -> bool:
"""Queue one L2->L1 load_back op for ``dr``; True iff a DMA was queued.
@@ -218,6 +224,12 @@ class DecodeHiCacheTransferMixin:
req=dr.req,
)
)
# The rematch repointed req.last_node to feed init_load_back's device
# boundary, but the prealloc lock and the receipt on the req still
# belong to pm.last_device_node; restore the pairing so any release
# before the commit hands over the restored lock hits the right node
# (the receipt's anchor makes a mispaired release assert).
dr.req.last_node = pm.last_device_node
# Failback: total coverage < required prefix means device alloc likely failed.
if len(rematch.device_indices) + len(new_indices) < pm.decode_prefix_len:
logger.warning(
@@ -238,7 +250,9 @@ class DecodeHiCacheTransferMixin:
[rematch.device_indices[pm.l1_prefix_len :], new_indices]
)
dr.hicache_restored_node = restored_node
self.tree_cache.inc_lock_ref(restored_node)
dr.hicache_restore_lock_receipt = self.tree_cache.inc_lock_ref(
restored_node
).to_dec_params()
if len(new_indices) == 0:
# Whole prefix already on device; no DMA needed.
@@ -303,7 +317,17 @@ class DecodeHiCacheTransferMixin:
if prefix_match is None or not prefix_match.needs_local_restore:
return
self.tree_cache.dec_lock_ref(prefix_match.last_device_node)
req = decode_req.req
restored_node = decode_req.hicache_restored_node
restored_lock_receipt = decode_req.hicache_restore_lock_receipt
assert restored_node is not None
assert restored_lock_receipt is not None
# Release preallocation before installing the restored lock receipt.
self.tree_cache.dec_lock_ref(
prefix_match.last_device_node,
req.lock_receipt,
skip_swa=req.swa_prefix_lock_released,
)
self.tree_cache.req_to_token_pool.write(
(
@@ -312,7 +336,12 @@ class DecodeHiCacheTransferMixin:
),
decode_req.hicache_restored_kv_indices,
)
decode_req.req.prefix_indices = torch.cat(
req.prefix_indices = torch.cat(
[prefix_match.prefix_indices, decode_req.hicache_restored_kv_indices]
)
decode_req.req.last_node = decode_req.hicache_restored_node
req.last_node = restored_node
req.lock_receipt = restored_lock_receipt
req.swa_prefix_lock_released = False
# Prevent abort cleanup from releasing the transferred lock.
decode_req.hicache_restored_node = None
decode_req.hicache_restore_lock_receipt = None
python/sglang/srt/managers/schedule_batch.py
+7 -11
View File
@@ -107,6 +107,7 @@ from sglang.srt.mem_cache.allocation_sizing import get_alloc_reserve_per_decode
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache,
DecLockRefParams,
MatchPrefixParams,
zero_match_result,
)
@@ -1125,13 +1126,11 @@ class Req(ReqDllmMixin):
self.storage_prefetch_retry_pending = False
self.storage_prefetch_retry_wait_polls = 0
self.storage_prefetch_retry_attempts = 0
# The node to lock until for swa radix tree lock ref
self.swa_uuid_for_lock: Optional[int] = None
# Receipt of the tree lock held on last_node (anchor, SWA boundary,
# skipped components); every release replays it unchanged.
self.lock_receipt: DecLockRefParams = DecLockRefParams()
# 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 = {}
# Whether or not if it is chunked. It increments whenever
# it is chunked, and decrement whenever chunked request is
@@ -1823,10 +1822,9 @@ class Req(ReqDllmMixin):
self.last_node = None
self.kv.cache_protected_len = 0
self.num_matched_prefix_tokens = 0
self.swa_uuid_for_lock = None
self.lock_receipt = DecLockRefParams()
self.swa_prefix_lock_released = False
self.swa_branching_seqlen = None
self.skip_lock_node_ids = {}
self.extend_range = None
self.dllm_initialized = False
self.is_retracted = True
@@ -3675,14 +3673,12 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
if (
release_leaf_lock
and not req.swa_prefix_lock_released
and req.swa_uuid_for_lock is not None
and req.lock_receipt.swa_uuid_for_lock is not None
and req.last_node is not None
and req.decode_batch_idx >= sliding_window_size
):
self.tree_cache.dec_swa_lock_only(
req.last_node,
req.swa_uuid_for_lock,
skip_lock_node_ids=req.skip_lock_node_ids,
req.last_node, req.lock_receipt
)
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
+4 -9
View File
@@ -1016,12 +1016,8 @@ class PrefillAdder:
self._account_prefill_cache_admission(req, prefix_len)
def _req_inc_lock_ref(self, req: Req):
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 = {}
# Persist the release receipt.
req.lock_receipt = self.tree_cache.inc_lock_ref(req.last_node).to_dec_params()
def add_dllm_staging_req(self, req: Req):
assert self.dllm_config is not None
@@ -1121,9 +1117,8 @@ class PrefillAdder:
try:
result = self.tree_cache.inc_lock_ref(last_node)
if self.tree_cache.is_tree_cache():
# init_load_back may revive SWA/Mamba tombstones while this
# temporary admission lock is held. Release must mirror the
# exact nodes skipped at acquire time.
# Replay the acquire's receipt (SWA boundary uuid, mamba flag)
# so release takes back exactly what this temporary lock took.
dec_lock_params = result.to_dec_params()
yield None
finally:
python/sglang/srt/managers/scheduler_components/dynamic_chunk_sizer.py
+3 -2
View File
@@ -174,8 +174,9 @@ class DynamicChunkSizer:
# Walk the same match -> lock -> alloc lifecycle as a scheduled
# request so release_kv_cache can release it symmetrically.
req.init_next_round_input(self.tree_cache)
lock = self.tree_cache.inc_lock_ref(req.last_node)
req.swa_uuid_for_lock = lock.swa_uuid_for_lock
req.lock_receipt = self.tree_cache.inc_lock_ref(
req.last_node
).to_dec_params()
req.set_extend_range(
len(req.prefix_indices), len(req.full_untruncated_fill_ids)
)
python/sglang/srt/mem_cache/base_prefix_cache.py
+20 -15
View File
@@ -131,39 +131,44 @@ class EvictResult:
@dataclasses.dataclass
class IncLockRefResult:
"""Result of an inc_lock_ref operation."""
"""Receipt returned by ``inc_lock_ref``.
``node_id`` is the anchor the lock was taken on; a release replays the
receipt on that node only. The SWA UUID marks the segment boundary;
``None`` means root. ``skipped_lock_components`` records the components
the acquire left untaken, so the release leaves them untouched.
"""
delta: Optional[int] = None
node_id: Optional[int] = None
swa_uuid_for_lock: Optional[int] = None
swa_uuid_for_host_lock: Optional[int] = None
# Component nodes that were tombstones at acquire time. Replaying this set
# at release prevents a short-lived lock from consuming a later load-back or
# request lock after that tombstone becomes a valid device value.
skip_lock_node_ids: dict[ComponentType, set[int]] = dataclasses.field(
default_factory=dict
)
skipped_lock_components: tuple[ComponentType, ...] = ()
def to_dec_params(self) -> DecLockRefParams:
"""Convert to the corresponding DecLockRefParams for dec_lock_ref."""
return DecLockRefParams(
node_id=self.node_id,
swa_uuid_for_lock=self.swa_uuid_for_lock,
swa_uuid_for_host_lock=self.swa_uuid_for_host_lock,
skip_lock_node_ids={
component_type: set(node_ids)
for component_type, node_ids in self.skip_lock_node_ids.items()
},
skipped_lock_components=tuple(self.skipped_lock_components),
)
@dataclasses.dataclass
class DecLockRefParams:
"""Parameters for dec_lock_ref operation."""
"""Receipt required by unified-tree ``dec_lock_ref``.
Fields default to nothing-acquired, so a lost receipt under-releases (a
leak the sanity checks report) instead of releasing another holder's
lock. ``node_id`` is ``None`` only for receipts that never came from a
unified-tree acquire (legacy caches, session sentinels).
"""
node_id: Optional[int] = None
swa_uuid_for_lock: Optional[int] = None
swa_uuid_for_host_lock: Optional[int] = None
skip_lock_node_ids: dict[ComponentType, set[int]] = dataclasses.field(
default_factory=dict
)
skipped_lock_components: tuple[ComponentType, ...] = ()
@dataclasses.dataclass
python/sglang/srt/mem_cache/rust_tree_core/adapter.py
+31 -63
View File
@@ -152,9 +152,23 @@ def _cache_actions_from_tagged(actions: Sequence[tuple]) -> list[CacheAction]:
def _inc_lock_ref_result_from_binding(result) -> IncLockRefResult:
return IncLockRefResult(
delta=result.delta,
node_id=result.node_id,
swa_uuid_for_lock=result.swa_uuid_for_lock,
swa_uuid_for_host_lock=result.swa_uuid_for_host_lock,
skip_lock_node_ids=_skip_lock_node_ids_from_binding(result.skip_lock_node_ids),
skipped_lock_components=tuple(
ComponentType(ct) for ct in result.skipped_lock_components
),
)
def _dec_lock_ref_params_to_binding(bindings_module, params: DecLockRefParams):
"""Build the binding's params from the module that owns the core's binding
(the inspection build is a distinct extension module with its own types)."""
return bindings_module.DecLockRefParamsBinding(
node_id=params.node_id,
swa_uuid_for_lock=params.swa_uuid_for_lock,
swa_uuid_for_host_lock=params.swa_uuid_for_host_lock,
skipped_lock_components=[int(ct) for ct in params.skipped_lock_components],
)
@@ -245,26 +259,6 @@ def _match_result_from_binding(result) -> MatchResult:
)
def _skip_lock_node_ids_from_binding(
skip_lock_node_ids: dict[int, set[int]],
) -> dict[ComponentType, set[int]]:
"""Rekey the binding's component-value skip map by ComponentType."""
return {
ComponentType(component): set(node_ids)
for component, node_ids in skip_lock_node_ids.items()
}
def _skip_lock_node_ids_to_binding(
skip_lock_node_ids: dict[ComponentType, set[int]],
) -> dict[int, set[int]]:
"""Rekey a ComponentType skip map by the binding's component values."""
return {
int(component): set(node_ids)
for component, node_ids in skip_lock_node_ids.items()
}
def _tracker_to_binding(tracker: dict[ComponentType, int]) -> dict[int, int]:
"""Rekey a ComponentType tracker by the binding's component values."""
return {int(component): freed for component, freed in tracker.items()}
@@ -333,6 +327,12 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
raise ValueError(
"Rust TreeCore does not support --radix-eviction-policy-config"
)
if ComponentType.SWA in self.tree_components and (
params.sliding_window_size is None or params.sliding_window_size <= 0
):
raise ValueError(
"the SWA tree component requires a positive sliding_window_size"
)
self._page_size = params.page_size
self.is_eagle = (
@@ -415,45 +415,29 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
skip_lock_components: Sequence[ComponentType] = (),
) -> IncLockRefResult:
result = self._binding.inc_lock_ref(
node_id, [int(component) for component in skip_lock_components]
node_id, [int(ct) for ct in skip_lock_components]
)
return _inc_lock_ref_result_from_binding(result)
def dec_lock_ref(
self,
node_id: NodeId,
params: Optional[DecLockRefParams] = None,
params: DecLockRefParams,
skip_swa: bool = False,
) -> DecLockRefResult:
binding_params = (
self._bindings.DecLockRefParamsBinding(
swa_uuid_for_lock=params.swa_uuid_for_lock,
swa_uuid_for_host_lock=params.swa_uuid_for_host_lock,
skip_lock_node_ids=_skip_lock_node_ids_to_binding(
params.skip_lock_node_ids
),
)
if params is not None
else None
self._binding.dec_lock_ref(
node_id, _dec_lock_ref_params_to_binding(self._bindings, params), skip_swa
)
self._binding.dec_lock_ref(node_id, binding_params, skip_swa)
return DecLockRefResult()
def dec_swa_lock_only(
self,
node_id: NodeId,
swa_uuid_for_lock: Optional[int],
skip_lock_node_ids: Optional[dict] = None,
params: DecLockRefParams,
) -> DecSwaLockOnlyResult:
result = DecSwaLockOnlyResult()
new_device_frees, new_host_frees = self._binding.dec_swa_lock_only(
node_id,
swa_uuid_for_lock,
(
_skip_lock_node_ids_to_binding(skip_lock_node_ids)
if skip_lock_node_ids
else None
),
node_id, _dec_lock_ref_params_to_binding(self._bindings, params)
)
for component, tensors in new_device_frees.items():
result.device_frees[ComponentType(component)].extend(tensors)
@@ -503,30 +487,14 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult:
result = self._binding.inc_host_lock_ref(node_id)
return IncLockRefResult(
delta=result.delta,
swa_uuid_for_lock=result.swa_uuid_for_lock,
swa_uuid_for_host_lock=result.swa_uuid_for_host_lock,
skip_lock_node_ids=_skip_lock_node_ids_from_binding(
result.skip_lock_node_ids
),
)
return _inc_lock_ref_result_from_binding(result)
def dec_host_lock_ref(
self, node_id: NodeId, params: Optional[DecLockRefParams] = None
self, node_id: NodeId, params: DecLockRefParams
) -> DecLockRefResult:
binding_params = (
self._bindings.DecLockRefParamsBinding(
swa_uuid_for_lock=params.swa_uuid_for_lock,
swa_uuid_for_host_lock=params.swa_uuid_for_host_lock,
skip_lock_node_ids=_skip_lock_node_ids_to_binding(
params.skip_lock_node_ids
),
)
if params is not None
else None
self._binding.dec_host_lock_ref(
node_id, _dec_lock_ref_params_to_binding(self._bindings, params)
)
self._binding.dec_host_lock_ref(node_id, binding_params)
return DecLockRefResult()
def evictable_size(self) -> int:
python/sglang/srt/mem_cache/swa_radix_cache.py
+11 -14
View File
@@ -501,9 +501,7 @@ class SWARadixCache(BasePrefixCache):
# Remove req slot release the cache lock
self.dec_lock_ref(
req.last_node,
DecLockRefParams(swa_uuid_for_lock=req.swa_uuid_for_lock),
skip_swa=req.swa_prefix_lock_released,
req.last_node, req.lock_receipt, skip_swa=req.swa_prefix_lock_released
)
req.swa_prefix_lock_released = False
@@ -563,13 +561,10 @@ class SWARadixCache(BasePrefixCache):
req.kv.cache_protected_len = len(new_indices)
self.dec_lock_ref(
req.last_node,
DecLockRefParams(swa_uuid_for_lock=req.swa_uuid_for_lock),
skip_swa=req.swa_prefix_lock_released,
req.last_node, req.lock_receipt, skip_swa=req.swa_prefix_lock_released
)
req.swa_prefix_lock_released = False
result = self.inc_lock_ref(new_last_node)
swa_uuid_for_lock = result.swa_uuid_for_lock
lock_receipt = self.inc_lock_ref(new_last_node).to_dec_params()
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
if len(new_indices) < len(kv_indices):
@@ -579,7 +574,7 @@ class SWARadixCache(BasePrefixCache):
else:
req.prefix_indices = new_indices
req.last_node = new_last_node
req.swa_uuid_for_lock = swa_uuid_for_lock
req.lock_receipt = lock_receipt
def pretty_print(self) -> None:
self._print_helper(self.root_node, 0)
@@ -806,13 +801,14 @@ class SWARadixCache(BasePrefixCache):
def dec_swa_lock_only(
self,
node: TreeNode,
swa_uuid_for_lock: Optional[int] = None,
skip_lock_node_ids: Optional[dict] = None, # unused, signature parity only
params: DecLockRefParams,
):
"""
Decrement only the swa_lock_ref (and swa_protected_size_) along the chain
[node, swa_uuid_for_lock], inclusive. The full_lock_ref is left untouched
so the caller's full-cache protection is preserved.
[node, receipt boundary uuid], inclusive. The full_lock_ref is left
untouched so the caller's full-cache protection is preserved. Of the
receipt this cache consumes only ``swa_uuid_for_lock``; it has no
lower-priority components to drop.
Used to early-release the SWA portion of a request's tree lock once the
request's decode position has advanced past the sliding window, so the
@@ -826,13 +822,14 @@ class SWARadixCache(BasePrefixCache):
as `swa_tombstone=True`. The full kv stays alive until the full-side
lock drops; future prefix-matches stop before this tombstoned leaf.
Caller must ensure this is invoked at most once per (node, swa_uuid_for_lock)
Caller must ensure this is invoked at most once per (node, boundary uuid)
pair (track via e.g. `Req.swa_prefix_lock_released`). When the request
finally releases its full lock via `dec_lock_ref`, pass `skip_swa=True`
to avoid touching SWA state again.
"""
if self.disable:
return
swa_uuid_for_lock = params.swa_uuid_for_lock
while node != self.root_node:
assert not node.swa_tombstone, (
python/sglang/srt/mem_cache/unified_cache/components/README.md
+28 -11
View File
@@ -180,34 +180,51 @@ Lock a node to protect it (and its ancestors) from eviction.
| Aspect | Detail |
|--------|--------|
| **Purpose** | Called when a request begins using a cached prefix — prevents eviction of nodes it depends on |
| **Inputs** | `node` — the last matched node (deepest) |
| **Output** | `IncLockRefResult(swa_uuid_for_lock)` |
| **Mutation** | Increments `lock_ref` per component along the path; moves tokens from evictable to protected size counters |
| **Inputs** | `node` — the last matched node (deepest); `skip_lock_components` names components to leave untaken (the decode hold passes `(MAMBA,)`) |
| **Output** | `IncLockRefResult(node_id, swa_uuid_for_lock, skipped_lock_components)` — the receipt the matching release must replay: the anchor node, the SWA boundary, and the skipped set |
| **Mutation** | Increments `lock_ref` per component along its contiguous segment; moves data-bearing tokens from evictable to protected size counters |
| **Complexity** | **O(D)** — Full: node to root; SWA: up to window boundary O(min(D, W)); Mamba: O(1).|
**Algorithm detail:** Calls `acquire_component_lock()` for each component.
**Algorithm detail:** Calls `acquire_component_lock()` for each component. A lock
covers a contiguous node segment and counts **every** node in it — tombstones
included (they carry no tokens, so sizes only move for data-bearing nodes).
| Component | Strategy |
|-----------|----------|
| Full | **Path-lock**: walks from node to root, `lock_ref += 1` on every ancestor. On first lock (`lock_ref: 0→1`), moves tokens from `component_evictable_size_` to `component_protected_size_`. |
| SWA | **Window-lock**: walks upward, accumulating SWA value lengths until `sliding_window_size` is filled. Records a `component_uuid` at the boundary node for `dec_lock_ref` to know where to stop. |
| Mamba | **Single-node lock**: only `lock_ref += 1` on the node itself (mamba state is per-leaf, not per-path). |
| SWA | **Segment-lock**: walks upward, `lock_ref += 1` on every node (tombstones included), accumulating position coverage (`len(key)`) until `sliding_window_size` is filled. Always stamps a boundary `component_uuid` at the last locked node; a `None` uuid in the receipt means the walk reached the root. |
| Mamba | **Single-node lock**: only `lock_ref += 1` on the node itself (mamba state is per-leaf, not per-path). Taken unless the acquire lists it in `skip_lock_components`; the receipt records the skipped set. The core names no component: it drives whatever the tree registered through the same interface. |
---
### `dec_lock_ref(node, params?) → DecLockRefResult`
### `dec_lock_ref(node, params, skip_swa=False) → DecLockRefResult`
Unlock a previously locked node path.
Unlock a previously locked node path by replaying the acquire's receipt.
| Aspect | Detail |
|--------|--------|
| **Purpose** | Called when a request finishes — releases eviction protection |
| **Inputs** | `node`, optional `params.swa_uuid_for_lock` for SWA boundary detection |
| **Inputs** | `node`; required `params` receipt (`node_id` anchor, `swa_uuid_for_lock` boundary, `skipped_lock_components`); `skip_swa=True` after an earlier `dec_swa_lock_only`. A receipt whose anchor is not `node` is a protocol violation (assert): a mispaired release would otherwise walk another holder's segment. |
| **Output** | `DecLockRefResult()` |
| **Mutation** | Decrements `lock_ref` per component; moves tokens from protected back to evictable when `lock_ref` reaches 0 |
| **Mutation** | Decrements `lock_ref` per component along the same segment the acquire counted; moves tokens from protected back to evictable when `lock_ref` reaches 0 |
| **Complexity** | **O(D)** — symmetric to `inc_lock_ref` |
**Algorithm detail:** Calls `release_component_lock()` for each component. Full walks to root; SWA walks up until matching `component_uuid`; Mamba decrements single node.
**Algorithm detail:** Releases auxiliary components before Full; every walk
refreshes the evictable-leaf membership of each node whose last lock it drops,
so the order is not load-bearing for the leaf sets. Full walks to root; SWA
stops at the receipt boundary; components in `skipped_lock_components` are
left alone. `skip_swa=True` also skips lower-priority components already
released by `dec_swa_lock_only`. The host-side `dec_host_lock_ref` takes the
same required receipt.
---
### `dec_swa_lock_only(node, params) → DecSwaLockOnlyResult`
Early-release only the SWA portion of a lock (decode advanced past the
window), plus strictly-lower-priority co-located locks (e.g. Mamba) the
receipt proves were taken. The eventual full release must pass
`skip_swa=True`. At most once per (node, boundary uuid) pair.
---
python/sglang/srt/mem_cache/unified_cache/components/full_component.py
+15 -14
View File
@@ -281,9 +281,11 @@ class FullComponent(TreeComponent):
root = self.tree_core.root_node
cur = node
# Skip the bottom evicted segment
# The bottom device-evicted segment is locked too (no ledger move —
# nothing is on device); a load-back that materializes a value under
# lock credits protected directly.
while cur is not root and cur.component_data[ct].value is None:
result.skip_lock_node_ids.setdefault(ct, set()).add(cur.id)
cur.component_data[ct].lock_ref += 1
cur = cur.parent
# Lock the device-on segment up to root
@@ -307,7 +309,7 @@ class FullComponent(TreeComponent):
def release_component_lock(
self,
node: UnifiedTreeNode,
params: Optional[DecLockRefParams],
params: DecLockRefParams,
lock_host: bool = False,
) -> None:
ct = self.component_type
@@ -315,7 +317,6 @@ class FullComponent(TreeComponent):
cd = node.component_data[ct]
if cd.host_lock_ref == 0:
return
# Mirror of `acquire`. write_back uses a pure counter.
if cd.host_value is None and not self.tree_core.is_write_back:
return
cd.host_lock_ref -= 1
@@ -323,17 +324,13 @@ class FullComponent(TreeComponent):
return
root = self.tree_core.root_node
skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
cur = node
while cur != root:
if cur.id in skip_lock_node_ids:
cur = cur.parent
continue
cd = cur.component_data[ct]
assert cd.value is not None
assert cd.lock_ref > 0
if cd.lock_ref == 1:
assert cd.lock_ref > 0, (
f"FULL segment release hit lock_ref=0 on node {cur.id}"
)
if cd.lock_ref == 1 and cd.value is not None:
key_len = len(cd.value)
self.tree_core.component_evictable_size_[ct] += key_len
self.tree_core.component_protected_size_[ct] -= key_len
@@ -422,8 +419,12 @@ class FullComponent(TreeComponent):
n_len = len(cd.host_value)
cd.value = device_indices[offset : offset + n_len].clone()
offset += n_len
# Full uses leaf sets, not LRU
self.tree_core.component_evictable_size_[ct] += n_len
# Full uses leaf sets, not LRU. A value materialized under
# lock is protected; the last release moves it to evictable.
if cd.lock_ref > 0:
self.tree_core.component_protected_size_[ct] += n_len
else:
self.tree_core.component_evictable_size_[ct] += n_len
self.tree_core._update_evictable_leaf_sets(n)
self.tree_core._update_evictable_leaf_sets(node)
python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py
+30 -38
View File
@@ -233,14 +233,8 @@ class MambaComponent(TreeComponent):
self._emit_excess_path_states_eviction(node, cache_actions)
return
if node.component_data[self.component_type].value is None:
node.component_data[self.component_type].value = params.mamba_value
# move from host LRU to device LRU
host_lru = self.tree_core.host_lru_lists[self.component_type]
if host_lru.in_list(node):
host_lru.remove_node(node)
self.tree_core.lru_lists[self.component_type].insert_mru(node)
self.tree_core.component_evictable_size_[self.component_type] += len(
params.mamba_value
self.tree_core.set_component_device_value(
node.id, self.component_type, params.mamba_value
)
node.last_access_time = get_and_increase_time_counter()
self._emit_excess_path_states_eviction(node, cache_actions)
@@ -441,19 +435,17 @@ class MambaComponent(TreeComponent):
return result
cd = node.component_data[ct]
value = cd.host_value if lock_host else cd.value
# A node in skip_lock_node_ids was a tombstone when this lock was acquired.
if value is None:
result.skip_lock_node_ids.setdefault(ct, set()).add(node.id)
return result
# Tombstones are counted too; ledger/LRU track only data-bearing
# nodes (a value materialized under lock is credited to protected
# at the materialization site).
if lock_host:
if cd.host_lock_ref == 0:
if cd.host_lock_ref == 0 and value is not None:
host_lru = self.tree_core.host_lru_lists[ct]
if host_lru.in_list(node):
host_lru.remove_node(node)
cd.host_lock_ref += 1
else:
if cd.lock_ref == 0:
if cd.lock_ref == 0 and value is not None:
vlen = len(value)
self.tree_core.component_evictable_size_[ct] -= vlen
self.tree_core.component_protected_size_[ct] += vlen
@@ -463,32 +455,36 @@ class MambaComponent(TreeComponent):
def release_component_lock(
self,
node: UnifiedTreeNode,
params: Optional[DecLockRefParams],
params: DecLockRefParams,
lock_host: bool = False,
) -> None:
ct = self.component_type
if node is self.tree_core.root_node:
return
cd = node.component_data[ct]
skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
if node.id in skip_lock_node_ids:
return
value = cd.host_value if lock_host else cd.value
if lock_host:
assert cd.host_lock_ref > 0, (
f"Mamba release hit host_lock_ref=0 on node {node.id}"
)
cd.host_lock_ref -= 1
if cd.host_lock_ref == 0 and cd.value is None and cd.host_value is not None:
host_lru = self.tree_core.host_lru_lists[ct]
if not host_lru.in_list(node):
host_lru.insert_mru(node)
if cd.host_lock_ref == 0:
if cd.value is None and cd.host_value is not None:
host_lru = self.tree_core.host_lru_lists[ct]
if not host_lru.in_list(node):
host_lru.insert_mru(node)
self.tree_core._update_evictable_leaf_sets(node)
return
if cd.lock_ref > 0:
if cd.lock_ref == 1:
vlen = len(value)
self.tree_core.component_evictable_size_[ct] += vlen
self.tree_core.component_protected_size_[ct] -= vlen
cd.lock_ref -= 1
assert cd.lock_ref > 0, f"Mamba release hit lock_ref=0 on node {node.id}"
if cd.lock_ref == 1 and value is not None:
vlen = len(value)
self.tree_core.component_evictable_size_[ct] += vlen
self.tree_core.component_protected_size_[ct] -= vlen
cd.lock_ref -= 1
if cd.lock_ref == 0:
self.tree_core._update_evictable_leaf_sets(node)
def _alloc_mamba_slot(self) -> torch.Tensor:
"""Allocate one mamba pool slot, evicting if necessary."""
@@ -809,15 +805,11 @@ class MambaComponent(TreeComponent):
return
transfer = transfers[0]
if transfer.device_indices is not None:
cd = node.component_data[ct]
cd.value = transfer.device_indices.clone()
count = len(cd.value)
# Move from host LRU to device LRU
host_lru = self.tree_core.host_lru_lists[ct]
if host_lru.in_list(node):
host_lru.remove_node(node)
self.tree_core.lru_lists[ct].insert_mru(node)
self.tree_core.component_evictable_size_[ct] += count
# The materialization primitive owns the ledger/LRU moves,
# including crediting protected when restored under lock.
self.tree_core.set_component_device_value(
node.id, ct, transfer.device_indices.clone()
)
elif phase == CacheTransferPhase.PREFETCH:
if not transfers:
python/sglang/srt/mem_cache/unified_cache/components/swa_component.py
+64 -54
View File
@@ -74,6 +74,9 @@ class SWAComponent(TreeComponent):
), (
f"SWAComponent requires SWATokenToKVPoolAllocator, got {type(params.token_to_kv_pool_allocator)}"
)
if params.sliding_window_size is None or params.sliding_window_size <= 0:
raise ValueError("SWAComponent requires a positive sliding_window_size")
super().__init__(cache, params)
self._session_leaf_covered_len: dict[str, dict[UnifiedTreeNode, int]] = {}
self.sliding_window_size = params.sliding_window_size
@@ -388,9 +391,9 @@ class SWAComponent(TreeComponent):
full_cd = node.component_data[BASE_COMPONENT_TYPE]
swa_evicted_seqlen = params.swa_evicted_seqlen
assert node.component_data[self.component_type].lock_ref == 0, (
f"tombstone {self.component_type} lock_ref should be 0, node {node.id}"
)
# A locked tombstone is legal (segment locks count every node); the
# full-value swap below is safe because full lock_ref >= swa
# lock_ref, so a locked-SWA node always takes the Recover branch.
assert swa_evicted_seqlen % self.tree_core.page_size == 0, (
f"{self.component_type}: swa_evicted_seqlen must be page-aligned, {swa_evicted_seqlen=}"
)
@@ -464,9 +467,6 @@ class SWAComponent(TreeComponent):
ct = self.component_type
if node.component_data[ct].value is not None:
return
assert node.component_data[ct].lock_ref == 0, (
f"tombstone {ct} lock_ref should be 0 on unevict, node {node.id}"
)
swa_evicted_seqlen = params.swa_evicted_seqlen
assert swa_evicted_seqlen % self.tree_core.page_size == 0, (
f"{ct}: swa_evicted_seqlen must be page-aligned, {swa_evicted_seqlen=}"
@@ -576,6 +576,9 @@ class SWAComponent(TreeComponent):
new_parent.component_data[self.component_type].lock_ref = child.component_data[
self.component_type
].lock_ref
new_parent.component_data[
self.component_type
].host_lock_ref = child.component_data[self.component_type].host_lock_ref
new_parent.component_data[
self.component_type
].session_ref = child.component_data[self.component_type].session_ref
@@ -610,6 +613,8 @@ class SWAComponent(TreeComponent):
parent_swa_data.metadata["host_uuid"] = host_uuid
host_lru = self.tree_core.host_lru_lists[self.component_type]
# Host-locked halves stay out of the host LRU: in-flight IO
# holds them, and host acquire removed the node at 0->1.
if (
new_parent.component_data[self.component_type].value is None
and parent_swa_data.host_lock_ref == 0
@@ -622,11 +627,16 @@ class SWAComponent(TreeComponent):
):
host_lru.insert_mru(child)
# parent inherits the swa_uuid from child for swa lock ref
# The window-boundary uuids mark the node's older edge, which the
# split moves to the parent — both tiers migrate with it.
new_parent.component_data[self.component_type].metadata["uuid"] = (
child.component_data[self.component_type].metadata.get("uuid")
)
child.component_data[self.component_type].metadata.pop("uuid", None)
new_parent.component_data[self.component_type].metadata["host_uuid"] = (
child.component_data[self.component_type].metadata.get("host_uuid")
)
child.component_data[self.component_type].metadata.pop("host_uuid", None)
def evict_component(
self,
@@ -754,10 +764,20 @@ class SWAComponent(TreeComponent):
result: IncLockRefResult,
lock_host: bool = False,
) -> IncLockRefResult:
"""Lock the contiguous segment covering the trailing window.
Every node in [node, boundary] is counted, tombstones included, so
the paired release decrements the same contiguous segment with no
carried skip state. Coverage is position-based (len(cur.key)); the
boundary node is always uuid-stamped, so a release without a uuid
means the segment reached the root. Ledger/LRU transitions track
only data-bearing nodes; a value materialized later under lock is
credited to protected by set_component_device_value.
"""
ct = self.component_type
root = self.tree_core.root_node
sliding_window_size = self.sliding_window_size
swa_lock_size = 0
covered = 0
swa_uuid = None
uuid_key = "host_uuid" if lock_host else "uuid"
lru = (
@@ -766,33 +786,25 @@ class SWAComponent(TreeComponent):
else self.tree_core.lru_lists[ct]
)
# Tombstoned nodes (cd.value is None) have no SWA chunk to protect
# skip them and keep walking up. This path is hit when HiCache
# backs up a FULL present internal node whose SWA was already evicted.
cur = node
while cur != root and swa_lock_size < sliding_window_size:
while cur != root and covered < sliding_window_size:
comp = cur.component_data[ct]
value = comp.host_value if lock_host else comp.value
if value is None:
result.skip_lock_node_ids.setdefault(ct, set()).add(cur.id)
cur = cur.parent
continue
ref = comp.host_lock_ref if lock_host else comp.lock_ref
if ref == 0:
if ref == 0 and value is not None:
if lock_host:
if lru.in_list(cur):
lru.remove_node(cur)
else:
key_len = len(cur.key)
key_len = len(value)
self.tree_core.component_evictable_size_[ct] -= key_len
self.tree_core.component_protected_size_[ct] += key_len
if lock_host:
comp.host_lock_ref = ref + 1
else:
comp.lock_ref = ref + 1
swa_lock_size += len(value)
if swa_lock_size >= sliding_window_size:
covered += len(cur.key)
if covered >= sliding_window_size:
if comp.metadata.get(uuid_key) is None:
comp.metadata[uuid_key] = next_component_uuid()
swa_uuid = comp.metadata[uuid_key]
@@ -807,45 +819,47 @@ class SWAComponent(TreeComponent):
def release_component_lock(
self,
node: UnifiedTreeNode,
params: Optional[DecLockRefParams],
params: DecLockRefParams,
lock_host: bool = False,
) -> None:
ct = self.component_type
root = self.tree_core.root_node
swa_uuid_for_lock = (
(params.swa_uuid_for_host_lock if lock_host else params.swa_uuid_for_lock)
if params
else None
params.swa_uuid_for_host_lock if lock_host else params.swa_uuid_for_lock
)
skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
dec_swa = True
uuid_key = "host_uuid" if lock_host else "uuid"
# A node in skip_lock_node_ids was a tombstone when this lock was acquired.
cur = node
while cur != root and dec_swa:
comp = cur.component_data[ct]
if cur.id in skip_lock_node_ids:
cur = cur.parent
continue
ref = comp.host_lock_ref if lock_host else comp.lock_ref
if ref == 0:
cur = cur.parent
continue
if ref == 1:
# Acquire counted every segment node and splits copy refs, so a
# zero here means the release does not mirror its acquire.
assert ref > 0, (
f"SWA segment release hit {'host_' if lock_host else ''}"
f"lock_ref=0 on node {cur.id}"
)
value = comp.host_value if lock_host else comp.value
if ref == 1 and value is not None:
if lock_host:
if comp.value is None and comp.host_value is not None:
if comp.value is None:
host_lru = self.tree_core.host_lru_lists[ct]
if not host_lru.in_list(cur):
host_lru.insert_mru(cur)
else:
key_len = len(comp.value)
key_len = len(value)
self.tree_core.component_evictable_size_[ct] += key_len
self.tree_core.component_protected_size_[ct] -= key_len
if lock_host:
comp.host_lock_ref = ref - 1
else:
comp.lock_ref = ref - 1
if ref == 1:
# This may have been the last lock holding the node out of
# the evictable-leaf sets; refresh it here rather than rely
# on the Full walk running after this one.
self.tree_core._update_evictable_leaf_sets(cur)
if swa_uuid_for_lock and comp.metadata.get(uuid_key) == swa_uuid_for_lock:
dec_swa = False
cur = cur.parent
@@ -857,15 +871,14 @@ class SWAComponent(TreeComponent):
device_frees: dict[ComponentType, list[torch.Tensor]],
host_frees: dict[ComponentType, list[torch.Tensor]],
) -> None:
"""Early-release the SWA lock along [node, swa_uuid_for_lock] while
leaving Full and Mamba locks intact.
"""Early-release the SWA lock along [node, swa_uuid_for_lock]; this
method touches only SWA state. The wrapping ``dec_swa_lock_only`` also
drops strictly-lower-priority co-located locks (e.g. Mamba) per the
receipt; the Full lock stays so the request's prefix is protected.
Called when a request's decode position has advanced past the sliding
window — the SWA portion of the tree lock is no longer needed but the
Full lock must stay so the request's prefix is protected.
Caller (UnifiedRadixCache.dec_swa_lock_only) must ensure this is
invoked at most once per (node, swa_uuid_for_lock) pair.
window. The caller must invoke this at most once per
(node, swa_uuid_for_lock) pair.
"""
ct = self.component_type
root = self.tree_core.root_node
@@ -873,19 +886,16 @@ class SWAComponent(TreeComponent):
cur = node
while cur is not root:
cd = cur.component_data[ct]
# Acquire skips tombstoned nodes; release must skip them too. Same
# for nodes with lock_ref == 0 — acquire never credited them.
if cd.value is None or cd.lock_ref == 0:
if swa_uuid_for_lock and cd.metadata.get("uuid") == swa_uuid_for_lock:
break
cur = cur.parent
continue
assert cd.lock_ref > 0, (
f"SWA window release hit lock_ref=0 on node {cur.id}"
)
cd.lock_ref -= 1
if cd.lock_ref == 0:
key_len = len(cur.key)
self.tree_core.component_protected_size_[ct] -= key_len
self.tree_core.component_evictable_size_[ct] += key_len
self.tree_core._update_evictable_leaf_sets(cur)
if cd.lock_ref == 0 and cd.value is not None:
value_len = len(cd.value)
self.tree_core.component_protected_size_[ct] -= value_len
self.tree_core.component_evictable_size_[ct] += value_len
if self.tree_core._is_device_leaf(cur):
self.tree_core._evict_component_and_detach_lru(
cur,
python/sglang/srt/mem_cache/unified_cache/components/tree_component.py
+1 -1
View File
@@ -584,7 +584,7 @@ class TreeComponent(ABC):
def release_component_lock(
self,
node: UnifiedTreeNode,
params: Optional[DecLockRefParams],
params: DecLockRefParams,
lock_host: bool = False,
) -> None:
"""Decrement component lock refs, un-protecting nodes.
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+81 -36
View File
@@ -622,32 +622,66 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = ()
) -> IncLockRefResult:
node = self.node_by_id(node_id)
result = IncLockRefResult()
skipped = tuple(skip_lock_components)
# The receipt records the anchor and what was locked; the paired dec
# replays exactly that.
result = IncLockRefResult(node_id=node.id, skipped_lock_components=skipped)
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)
if component.component_type in skipped:
continue
result = component.acquire_component_lock(node=node, result=result)
self._update_evictable_leaf_sets(node)
return result
@staticmethod
def _assert_receipt_anchor(node: UnifiedTreeNode, params: DecLockRefParams) -> None:
"""A receipt releases only the node its acquire returned; a mispaired
node would silently release (or steal) another holder's segment."""
assert params.node_id is None or params.node_id == node.id, (
f"lock receipt anchored on node {params.node_id} released on node {node.id}"
)
def _release_components(
self,
node: UnifiedTreeNode,
params: DecLockRefParams,
*,
lock_host: bool = False,
skip_swa_and_below: bool = False,
) -> None:
"""Release each component this receipt acquired. Auxiliaries go first
so Full, whose walk refreshes leaf membership on every node it
unlocks, sees their final refs; the auxiliary walks also refresh the
nodes they unlock, so the order is not load-bearing for the sets."""
swa_priority = None
if skip_swa_and_below:
swa_component = self.components_by_type.get(ComponentType.SWA)
if swa_component is not None:
swa_priority = swa_component.eviction_priority(is_leaf=False)
for component in reversed(self.components):
ct = component.component_type
if ct in params.skipped_lock_components:
continue
if swa_priority is not None and (
ct == ComponentType.SWA
or component.eviction_priority(is_leaf=False) < swa_priority
):
continue
component.release_component_lock(
node=node, params=params, lock_host=lock_host
)
def dec_lock_ref(
self,
node_id: NodeId,
params: Optional[DecLockRefParams] = None,
params: DecLockRefParams,
skip_swa: bool = False,
) -> DecLockRefResult:
node = self.node_by_id(node_id)
for component in self.components:
if skip_swa and component.component_type == ComponentType.SWA:
continue
component.release_component_lock(node=node, params=params)
self._assert_receipt_anchor(node, params)
# After an SWA early release (dec_swa_lock_only), SWA and the
# lower-priority components it dropped are already released.
self._release_components(node, params, skip_swa_and_below=skip_swa)
self._update_evictable_leaf_sets(node)
# TODO: delta is not aggregated from components; no caller uses it yet.
return DecLockRefResult()
@@ -655,36 +689,33 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
def dec_swa_lock_only(
self,
node_id: NodeId,
swa_uuid_for_lock: Optional[int],
skip_lock_node_ids: Optional[dict] = None,
params: DecLockRefParams,
) -> 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."""
result = DecSwaLockOnlyResult()
node = self.node_by_id(node_id)
self._assert_receipt_anchor(node, params)
swa_component = self.components_by_type.get(ComponentType.SWA)
if swa_component is None:
return result
swa_component.release_window_lock(
node, swa_uuid_for_lock, result.device_frees, result.host_frees
node, params.swa_uuid_for_lock, result.device_frees, result.host_frees
)
# 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).
# Drop strictly-lower-priority locks co-located on the node, skipping
# any the paired inc never took (matters for FULL+SWA+MAMBA models).
swa_priority = swa_component.eviction_priority(is_leaf=False)
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:
for comp in reversed(self.components):
if comp.component_type in params.skipped_lock_components:
continue
if comp.eviction_priority(is_leaf=False) < swa_priority:
comp.release_component_lock(node, dec_params)
comp.release_component_lock(node, params)
return result
def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult:
node = self.node_by_id(node_id)
result = IncLockRefResult()
result = IncLockRefResult(node_id=node.id)
for component in self.components:
result = component.acquire_component_lock(
node=node, result=result, lock_host=True
@@ -693,11 +724,11 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
return result
def dec_host_lock_ref(
self, node_id: NodeId, params: Optional[DecLockRefParams] = None
self, node_id: NodeId, params: DecLockRefParams
) -> DecLockRefResult:
node = self.node_by_id(node_id)
for component in self.components:
component.release_component_lock(node=node, params=params, lock_host=True)
self._assert_receipt_anchor(node, params)
self._release_components(node, params, lock_host=True)
self._update_evictable_leaf_sets(node)
return DecLockRefResult()
@@ -1260,7 +1291,10 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
assert cd.value is None
n = len(fresh_value)
cd.value = fresh_value.clone()
self.component_evictable_size_[ct] += n
if cd.lock_ref > 0:
self.component_protected_size_[ct] += n
else:
self.component_evictable_size_[ct] += n
self._update_evictable_leaf_sets(node)
# A backuped node restored from fresh KV is a duplicate right away.
self._update_duplicate_tracking(node)
@@ -1887,7 +1921,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
return True
def _is_host_leaf(self, node: UnifiedTreeNode) -> bool:
"""H-leaf: evicted, Full host value present, no children, unlocked, not root.
"""H-leaf: evicted, Full host value present, no children, unlocked on
both tiers, not root.
Only the Full (base) component host_value is required; auxiliary
components are not mandatory for H-leaf membership. In-flight DMA
@@ -1898,6 +1933,10 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
return False
if any(cd.host_lock_ref > 0 for cd in node.component_data):
return False
# Segment locks count evicted nodes too: a device-locked candidate is
# a live segment's anchor, and _evict_host_leaf would delete it.
if any(cd.lock_ref > 0 for cd in node.component_data):
return False
if len(node.children) > 0:
return False
return True
@@ -2258,12 +2297,18 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
# Full uses leaf sets, not LRU; its stores go through the insert paths.
assert component_type != BASE_COMPONENT_TYPE
node = self.node_by_id(node_id)
node.component_data[component_type].value = value
cd = node.component_data[component_type]
cd.value = value
host_lru = self.host_lru_lists[component_type]
if host_lru.in_list(node):
host_lru.remove_node(node)
self.lru_lists[component_type].insert_mru(node)
self.component_evictable_size_[component_type] += len(value)
# A value materialized under lock is protected; the last release
# moves it to evictable.
if cd.lock_ref > 0:
self.component_protected_size_[component_type] += len(value)
else:
self.component_evictable_size_[component_type] += len(value)
def get_component_device_value(
self, node_id: NodeId, component_type: ComponentType
@@ -2361,8 +2406,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
E(f"node {nid} {ct} host_lock_ref={cd.host_lock_ref}")
if ct != FCT and fl < cd.lock_ref:
E(f"node {nid} full_lock={fl} < {ct}_lock={cd.lock_ref}")
if cd.value is None and cd.lock_ref > 0:
E(f"node {nid} {ct} evicted but lock_ref={cd.lock_ref}")
# Locked tombstones are legal: segment locks count every
# node in [start, boundary], data-bearing or not.
# Collect expected leaf qualification (single pass)
if self._is_device_leaf(node):
python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py
+8 -7
View File
@@ -237,26 +237,27 @@ class UnifiedTreeCoreInterface(ABC):
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."""
"""Bump the reference count on a node's component locks. Components in
``skip_lock_components`` are left untaken; the receipt records the
anchor node and the skipped set so the paired release mirrors them."""
...
@abstractmethod
def dec_lock_ref(
self,
node_id: NodeId,
params: Optional[DecLockRefParams] = None,
params: DecLockRefParams,
skip_swa: bool = False,
) -> DecLockRefResult:
"""Decrease the reference count on a node's component locks."""
"""Decrease the reference count on a node's component locks. The
receipt is required: a release must replay its acquire's evidence."""
...
@abstractmethod
def dec_swa_lock_only(
self,
node_id: NodeId,
swa_uuid_for_lock: Optional[int],
skip_lock_node_ids: Optional[dict] = None,
params: DecLockRefParams,
) -> DecSwaLockOnlyResult:
"""Decrease only the SWA (and lower-priority co-located) reference
counts; the result carries the freed slots."""
@@ -313,7 +314,7 @@ class UnifiedTreeCoreInterface(ABC):
@abstractmethod
def dec_host_lock_ref(
self, node_id: NodeId, params: Optional[DecLockRefParams] = None
self, node_id: NodeId, params: DecLockRefParams
) -> DecLockRefResult:
"""Decrease the reference count on a node's host-side component locks."""
...
python/sglang/srt/mem_cache/unified_radix_cache.py
+16 -27
View File
@@ -878,7 +878,7 @@ class UnifiedRadixCache(BasePrefixCache):
def dec_lock_ref(
self,
node_id: NodeId,
params: Optional[DecLockRefParams] = None,
params: DecLockRefParams,
skip_swa: bool = False,
) -> DecLockRefResult:
result = self.session.try_dec_lock_ref(node_id, params)
@@ -889,28 +889,18 @@ class UnifiedRadixCache(BasePrefixCache):
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,
)
"""Release the tree lock a request holds on its last_node with the
receipt its acquire returned, so it never drops a lock it never took."""
self.dec_lock_ref(req.last_node, req.lock_receipt, 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,
params: DecLockRefParams,
) -> None:
if self.disable:
return
result = self.tree_core.dec_swa_lock_only(
node_id, swa_uuid_for_lock, skip_lock_node_ids
)
result = self.tree_core.dec_swa_lock_only(node_id, params)
self._free_values(result.device_frees, result.host_frees)
def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult:
@@ -919,7 +909,7 @@ class UnifiedRadixCache(BasePrefixCache):
return self.tree_core.inc_host_lock_ref(node_id)
def dec_host_lock_ref(
self, node_id: NodeId, params: Optional[DecLockRefParams] = None
self, node_id: NodeId, params: DecLockRefParams
) -> DecLockRefResult:
if self.disable:
return DecLockRefResult()
@@ -1099,20 +1089,20 @@ class UnifiedRadixCache(BasePrefixCache):
new_indices[req.kv.cache_protected_len :],
)
self._dec_req_lock(req)
self._dec_req_lock(req, skip_swa=req.swa_prefix_lock_released)
# 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
new_last_node,
skip_lock_components=(
(ComponentType.MAMBA,)
if envs.SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK.get()
else ()
),
)
# Update req fields
@@ -1124,9 +1114,8 @@ class UnifiedRadixCache(BasePrefixCache):
req.prefix_indices = new_indices
req.kv.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
# Carry the receipt so this node's dec releases only what we locked.
req.lock_receipt = lock_result.to_dec_params()
# The rematch acquired a new SWA prefix lock.
req.swa_prefix_lock_released = False
python/sglang/srt/session/streaming_session.py
+14 -17
View File
@@ -48,18 +48,18 @@ class SessionSlot:
# First req's radix tree node (for dec_lock_ref on session close)
last_node: Any = None
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)
# Receipt of the first request's tree lock on last_node.
lock_receipt: DecLockRefParams = field(default_factory=DecLockRefParams)
# Whether the first request already released its SWA lock.
swa_prefix_lock_released: bool = False
def save_from_req(self, req: Req, is_first: bool):
"""Save KV state from a finishing request into this slot."""
kv = req.detach_kv()
if is_first:
self.last_node = req.last_node
self.swa_uuid_for_lock = req.swa_uuid_for_lock
self.skip_lock_node_ids = req.skip_lock_node_ids
self.lock_receipt = req.lock_receipt
self.swa_prefix_lock_released = req.swa_prefix_lock_released
# The slot takes over the request's KV record.
self.kv = kv
else:
@@ -71,8 +71,8 @@ class SessionSlot:
def restore_to_req(self, req: Req):
"""Restore KV state from this slot into an incoming request."""
req.kv = self.kv
req.swa_uuid_for_lock = self.swa_uuid_for_lock
req.skip_lock_node_ids = self.skip_lock_node_ids
req.lock_receipt = self.lock_receipt
req.swa_prefix_lock_released = self.swa_prefix_lock_released
# NOTE: the slot keeps sharing the record it just handed out. During
# chunked prefill, a request may be rejected by
@@ -290,8 +290,8 @@ class StreamingSession(BasePrefixCache):
slot = SessionSlot(
kv=kv,
last_node=req.last_node,
swa_uuid_for_lock=req.swa_uuid_for_lock,
skip_lock_node_ids=req.skip_lock_node_ids,
lock_receipt=req.lock_receipt,
swa_prefix_lock_released=req.swa_prefix_lock_released,
)
self.slots[session_id] = slot
else:
@@ -396,13 +396,10 @@ class StreamingSession(BasePrefixCache):
)
if lock_node is not None:
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,
),
)
# skip_swa is an SWA-cache extension kwarg; a slot can only have
# early-released when the inner cache supports SWA locks.
skip = {"skip_swa": True} if slot.swa_prefix_lock_released else {}
self.inner.dec_lock_ref(lock_node, slot.lock_receipt, **skip)
if slot.kv.holds_kv:
self.free_kv_row(slot.kv, [(protected_len, slot.kv.kv_allocated_len)])
python/sglang/test/scripted_runtime/context/lock_ref_exhauster.py
+6 -5
View File
@@ -11,7 +11,8 @@ if TYPE_CHECKING:
class ScriptedLockRefExhauster:
def __init__(self, scheduler: Scheduler) -> None:
self.scheduler = scheduler
self._locked: List[Any] = []
# (node, dec receipt) pairs; the receipt bounds the release walk.
self._locked: List[tuple[Any, Any]] = []
def exhaust(self, *, leave_refs: int) -> None:
tree_cache = self.scheduler.tree_cache
@@ -24,17 +25,17 @@ class ScriptedLockRefExhauster:
return
target = evictable[0]
tree_cache.inc_lock_ref(to_node_handle(tree_cache, target))
result = tree_cache.inc_lock_ref(to_node_handle(tree_cache, target))
newly_locked = [node for node in evictable if _node_lock_ref(node) > 0]
if not newly_locked:
return
self._locked.append(target)
self._locked.append((target, result.to_dec_params()))
def release(self) -> None:
tree_cache = self.scheduler.tree_cache
for node in self._locked:
tree_cache.dec_lock_ref(to_node_handle(tree_cache, node))
for node, dec_params in self._locked:
tree_cache.dec_lock_ref(to_node_handle(tree_cache, node), dec_params)
self._locked.clear()
def _evictable_nodes(self) -> List[Any]: