minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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 prefix_match: Optional[DecodePrefixMatch] = None
hicache_restored_kv_indices: Optional[torch.Tensor] = None hicache_restored_kv_indices: Optional[torch.Tensor] = None
hicache_restored_node: Any = 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_load_consumer_index: int = -1
hicache_restore_status: HiCacheRestoreResult = HiCacheRestoreResult.PENDING hicache_restore_status: HiCacheRestoreResult = HiCacheRestoreResult.PENDING
@@ -426,12 +428,11 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
) )
def _release_matched_prefix_lock(self, req: Req) -> None: 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: 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 req.swa_prefix_lock_released = False
else: 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( def _reclaim_swa_tail_capacity(
self, swa_tail_len: int, req_id: str self, swa_tail_len: int, req_id: str
@@ -676,9 +677,11 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
include_req=True, include_req=True,
) )
# Keep aggregated scheduling semantics while preserving the SWA lock # Keep aggregated scheduling semantics while preserving the SWA lock
# boundary needed for the matching dec_lock_ref. # boundary needed for the matching dec_lock_ref; the full receipt
lock_result = self.tree_cache.inc_lock_ref(result.last_device_node) # travels on the req so every later release mirrors this acquire.
req.swa_uuid_for_lock = lock_result.swa_uuid_for_lock req.lock_receipt = self.tree_cache.inc_lock_ref(
result.last_device_node
).to_dec_params()
return self._build_decode_prefix_match(req, result) return self._build_decode_prefix_match(req, result)
def _resolve_prefill_dp_rank(self, req: Req) -> Optional[int]: 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") and hasattr(self.tree_cache, "dec_swa_lock_only")
): ):
self.tree_cache.dec_swa_lock_only( self.tree_cache.dec_swa_lock_only(
decode_req.req.last_node, decode_req.req.last_node, decode_req.req.lock_receipt
decode_req.req.swa_uuid_for_lock,
) )
decode_req.req.swa_prefix_lock_released = True 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.disaggregation.base import KVPoll
from sglang.srt.managers.schedule_policy import match_prefix_for_req 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: if TYPE_CHECKING:
from sglang.srt.disaggregation.decode import DecodeRequest from sglang.srt.disaggregation.decode import DecodeRequest
@@ -184,8 +186,12 @@ class DecodeHiCacheTransferMixin:
): ):
self.tree_cache.release_aborted_request(decode_req.req.rid) self.tree_cache.release_aborted_request(decode_req.req.rid)
if decode_req.hicache_restored_node is not None: 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_restored_node = None
decode_req.hicache_restore_lock_receipt = None
def _try_hicache_queue_load_back(self, dr: DecodeRequest) -> bool: 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. """Queue one L2->L1 load_back op for ``dr``; True iff a DMA was queued.
@@ -218,6 +224,12 @@ class DecodeHiCacheTransferMixin:
req=dr.req, 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. # Failback: total coverage < required prefix means device alloc likely failed.
if len(rematch.device_indices) + len(new_indices) < pm.decode_prefix_len: if len(rematch.device_indices) + len(new_indices) < pm.decode_prefix_len:
logger.warning( logger.warning(
@@ -238,7 +250,9 @@ class DecodeHiCacheTransferMixin:
[rematch.device_indices[pm.l1_prefix_len :], new_indices] [rematch.device_indices[pm.l1_prefix_len :], new_indices]
) )
dr.hicache_restored_node = restored_node 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: if len(new_indices) == 0:
# Whole prefix already on device; no DMA needed. # 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: if prefix_match is None or not prefix_match.needs_local_restore:
return 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( self.tree_cache.req_to_token_pool.write(
( (
@@ -312,7 +336,12 @@ class DecodeHiCacheTransferMixin:
), ),
decode_req.hicache_restored_kv_indices, 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] [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.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import ( from sglang.srt.mem_cache.base_prefix_cache import (
BasePrefixCache, BasePrefixCache,
DecLockRefParams,
MatchPrefixParams, MatchPrefixParams,
zero_match_result, zero_match_result,
) )
@@ -1125,13 +1126,11 @@ class Req(ReqDllmMixin):
self.storage_prefetch_retry_pending = False self.storage_prefetch_retry_pending = False
self.storage_prefetch_retry_wait_polls = 0 self.storage_prefetch_retry_wait_polls = 0
self.storage_prefetch_retry_attempts = 0 self.storage_prefetch_retry_attempts = 0
# The node to lock until for swa radix tree lock ref # Receipt of the tree lock held on last_node (anchor, SWA boundary,
self.swa_uuid_for_lock: Optional[int] = None # skipped components); every release replays it unchanged.
self.lock_receipt: DecLockRefParams = DecLockRefParams()
# Whether the prefill-time SWA tree lock has been released early # Whether the prefill-time SWA tree lock has been released early
self.swa_prefix_lock_released: bool = False 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 # Whether or not if it is chunked. It increments whenever
# it is chunked, and decrement whenever chunked request is # it is chunked, and decrement whenever chunked request is
@@ -1823,10 +1822,9 @@ class Req(ReqDllmMixin):
self.last_node = None self.last_node = None
self.kv.cache_protected_len = 0 self.kv.cache_protected_len = 0
self.num_matched_prefix_tokens = 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_prefix_lock_released = False
self.swa_branching_seqlen = None self.swa_branching_seqlen = None
self.skip_lock_node_ids = {}
self.extend_range = None self.extend_range = None
self.dllm_initialized = False self.dllm_initialized = False
self.is_retracted = True self.is_retracted = True
@@ -3675,14 +3673,12 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
if ( if (
release_leaf_lock release_leaf_lock
and not req.swa_prefix_lock_released 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.last_node is not None
and req.decode_batch_idx >= sliding_window_size and req.decode_batch_idx >= sliding_window_size
): ):
self.tree_cache.dec_swa_lock_only( self.tree_cache.dec_swa_lock_only(
req.last_node, req.last_node, req.lock_receipt
req.swa_uuid_for_lock,
skip_lock_node_ids=req.skip_lock_node_ids,
) )
req.swa_prefix_lock_released = True req.swa_prefix_lock_released = True
elif self.forward_mode.is_extend() and self.tree_cache.is_chunk_cache(): 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) self._account_prefill_cache_admission(req, prefix_len)
def _req_inc_lock_ref(self, req: Req): def _req_inc_lock_ref(self, req: Req):
result = self.tree_cache.inc_lock_ref(req.last_node) # Persist the release receipt.
if self.is_hybrid_swa: req.lock_receipt = self.tree_cache.inc_lock_ref(req.last_node).to_dec_params()
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): def add_dllm_staging_req(self, req: Req):
assert self.dllm_config is not None assert self.dllm_config is not None
@@ -1121,9 +1117,8 @@ class PrefillAdder:
try: try:
result = self.tree_cache.inc_lock_ref(last_node) result = self.tree_cache.inc_lock_ref(last_node)
if self.tree_cache.is_tree_cache(): if self.tree_cache.is_tree_cache():
# init_load_back may revive SWA/Mamba tombstones while this # Replay the acquire's receipt (SWA boundary uuid, mamba flag)
# temporary admission lock is held. Release must mirror the # so release takes back exactly what this temporary lock took.
# exact nodes skipped at acquire time.
dec_lock_params = result.to_dec_params() dec_lock_params = result.to_dec_params()
yield None yield None
finally: 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 # Walk the same match -> lock -> alloc lifecycle as a scheduled
# request so release_kv_cache can release it symmetrically. # request so release_kv_cache can release it symmetrically.
req.init_next_round_input(self.tree_cache) req.init_next_round_input(self.tree_cache)
lock = self.tree_cache.inc_lock_ref(req.last_node) req.lock_receipt = self.tree_cache.inc_lock_ref(
req.swa_uuid_for_lock = lock.swa_uuid_for_lock req.last_node
).to_dec_params()
req.set_extend_range( req.set_extend_range(
len(req.prefix_indices), len(req.full_untruncated_fill_ids) 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 @dataclasses.dataclass
class IncLockRefResult: 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 delta: Optional[int] = None
node_id: Optional[int] = None
swa_uuid_for_lock: Optional[int] = None swa_uuid_for_lock: Optional[int] = None
swa_uuid_for_host_lock: Optional[int] = None swa_uuid_for_host_lock: Optional[int] = None
# Component nodes that were tombstones at acquire time. Replaying this set skipped_lock_components: tuple[ComponentType, ...] = ()
# 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
)
def to_dec_params(self) -> DecLockRefParams: def to_dec_params(self) -> DecLockRefParams:
"""Convert to the corresponding DecLockRefParams for dec_lock_ref.""" """Convert to the corresponding DecLockRefParams for dec_lock_ref."""
return DecLockRefParams( return DecLockRefParams(
node_id=self.node_id,
swa_uuid_for_lock=self.swa_uuid_for_lock, swa_uuid_for_lock=self.swa_uuid_for_lock,
swa_uuid_for_host_lock=self.swa_uuid_for_host_lock, swa_uuid_for_host_lock=self.swa_uuid_for_host_lock,
skip_lock_node_ids={ skipped_lock_components=tuple(self.skipped_lock_components),
component_type: set(node_ids)
for component_type, node_ids in self.skip_lock_node_ids.items()
},
) )
@dataclasses.dataclass @dataclasses.dataclass
class DecLockRefParams: 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_lock: Optional[int] = None
swa_uuid_for_host_lock: Optional[int] = None swa_uuid_for_host_lock: Optional[int] = None
skip_lock_node_ids: dict[ComponentType, set[int]] = dataclasses.field( skipped_lock_components: tuple[ComponentType, ...] = ()
default_factory=dict
)
@dataclasses.dataclass @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: def _inc_lock_ref_result_from_binding(result) -> IncLockRefResult:
return IncLockRefResult( return IncLockRefResult(
delta=result.delta, delta=result.delta,
node_id=result.node_id,
swa_uuid_for_lock=result.swa_uuid_for_lock, swa_uuid_for_lock=result.swa_uuid_for_lock,
swa_uuid_for_host_lock=result.swa_uuid_for_host_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]: def _tracker_to_binding(tracker: dict[ComponentType, int]) -> dict[int, int]:
"""Rekey a ComponentType tracker by the binding's component values.""" """Rekey a ComponentType tracker by the binding's component values."""
return {int(component): freed for component, freed in tracker.items()} return {int(component): freed for component, freed in tracker.items()}
@@ -333,6 +327,12 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
raise ValueError( raise ValueError(
"Rust TreeCore does not support --radix-eviction-policy-config" "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._page_size = params.page_size
self.is_eagle = ( self.is_eagle = (
@@ -415,45 +415,29 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
skip_lock_components: Sequence[ComponentType] = (), skip_lock_components: Sequence[ComponentType] = (),
) -> IncLockRefResult: ) -> IncLockRefResult:
result = self._binding.inc_lock_ref( 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) return _inc_lock_ref_result_from_binding(result)
def dec_lock_ref( def dec_lock_ref(
self, self,
node_id: NodeId, node_id: NodeId,
params: Optional[DecLockRefParams] = None, params: DecLockRefParams,
skip_swa: bool = False, skip_swa: bool = False,
) -> DecLockRefResult: ) -> DecLockRefResult:
binding_params = ( self._binding.dec_lock_ref(
self._bindings.DecLockRefParamsBinding( node_id, _dec_lock_ref_params_to_binding(self._bindings, params), skip_swa
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, binding_params, skip_swa)
return DecLockRefResult() return DecLockRefResult()
def dec_swa_lock_only( def dec_swa_lock_only(
self, self,
node_id: NodeId, node_id: NodeId,
swa_uuid_for_lock: Optional[int], params: DecLockRefParams,
skip_lock_node_ids: Optional[dict] = None,
) -> DecSwaLockOnlyResult: ) -> DecSwaLockOnlyResult:
result = DecSwaLockOnlyResult() result = DecSwaLockOnlyResult()
new_device_frees, new_host_frees = self._binding.dec_swa_lock_only( new_device_frees, new_host_frees = self._binding.dec_swa_lock_only(
node_id, node_id, _dec_lock_ref_params_to_binding(self._bindings, params)
swa_uuid_for_lock,
(
_skip_lock_node_ids_to_binding(skip_lock_node_ids)
if skip_lock_node_ids
else None
),
) )
for component, tensors in new_device_frees.items(): for component, tensors in new_device_frees.items():
result.device_frees[ComponentType(component)].extend(tensors) result.device_frees[ComponentType(component)].extend(tensors)
@@ -503,30 +487,14 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult: def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult:
result = self._binding.inc_host_lock_ref(node_id) result = self._binding.inc_host_lock_ref(node_id)
return IncLockRefResult( return _inc_lock_ref_result_from_binding(result)
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
),
)
def dec_host_lock_ref( def dec_host_lock_ref(
self, node_id: NodeId, params: Optional[DecLockRefParams] = None self, node_id: NodeId, params: DecLockRefParams
) -> DecLockRefResult: ) -> DecLockRefResult:
binding_params = ( self._binding.dec_host_lock_ref(
self._bindings.DecLockRefParamsBinding( node_id, _dec_lock_ref_params_to_binding(self._bindings, params)
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, binding_params)
return DecLockRefResult() return DecLockRefResult()
def evictable_size(self) -> int: 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 # Remove req slot release the cache lock
self.dec_lock_ref( self.dec_lock_ref(
req.last_node, req.last_node, req.lock_receipt, skip_swa=req.swa_prefix_lock_released
DecLockRefParams(swa_uuid_for_lock=req.swa_uuid_for_lock),
skip_swa=req.swa_prefix_lock_released,
) )
req.swa_prefix_lock_released = False req.swa_prefix_lock_released = False
@@ -563,13 +561,10 @@ class SWARadixCache(BasePrefixCache):
req.kv.cache_protected_len = len(new_indices) req.kv.cache_protected_len = len(new_indices)
self.dec_lock_ref( self.dec_lock_ref(
req.last_node, req.last_node, req.lock_receipt, skip_swa=req.swa_prefix_lock_released
DecLockRefParams(swa_uuid_for_lock=req.swa_uuid_for_lock),
skip_swa=req.swa_prefix_lock_released,
) )
req.swa_prefix_lock_released = False req.swa_prefix_lock_released = False
result = self.inc_lock_ref(new_last_node) lock_receipt = self.inc_lock_ref(new_last_node).to_dec_params()
swa_uuid_for_lock = result.swa_uuid_for_lock
# `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
if len(new_indices) < len(kv_indices): if len(new_indices) < len(kv_indices):
@@ -579,7 +574,7 @@ class SWARadixCache(BasePrefixCache):
else: else:
req.prefix_indices = new_indices req.prefix_indices = new_indices
req.last_node = new_last_node 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: def pretty_print(self) -> None:
self._print_helper(self.root_node, 0) self._print_helper(self.root_node, 0)
@@ -806,13 +801,14 @@ class SWARadixCache(BasePrefixCache):
def dec_swa_lock_only( def dec_swa_lock_only(
self, self,
node: TreeNode, node: TreeNode,
swa_uuid_for_lock: Optional[int] = None, params: DecLockRefParams,
skip_lock_node_ids: Optional[dict] = None, # unused, signature parity only
): ):
""" """
Decrement only the swa_lock_ref (and swa_protected_size_) along the chain 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 [node, receipt boundary uuid], inclusive. The full_lock_ref is left
so the caller's full-cache protection is preserved. 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 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 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 as `swa_tombstone=True`. The full kv stays alive until the full-side
lock drops; future prefix-matches stop before this tombstoned leaf. 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 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` finally releases its full lock via `dec_lock_ref`, pass `skip_swa=True`
to avoid touching SWA state again. to avoid touching SWA state again.
""" """
if self.disable: if self.disable:
return return
swa_uuid_for_lock = params.swa_uuid_for_lock
while node != self.root_node: while node != self.root_node:
assert not node.swa_tombstone, ( 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 | | Aspect | Detail |
|--------|--------| |--------|--------|
| **Purpose** | Called when a request begins using a cached prefix — prevents eviction of nodes it depends on | | **Purpose** | Called when a request begins using a cached prefix — prevents eviction of nodes it depends on |
| **Inputs** | `node` — the last matched node (deepest) | | **Inputs** | `node` — the last matched node (deepest); `skip_lock_components` names components to leave untaken (the decode hold passes `(MAMBA,)`) |
| **Output** | `IncLockRefResult(swa_uuid_for_lock)` | | **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 the path; moves tokens from evictable to protected size counters | | **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).| | **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 | | 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_`. | | 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. | | 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). | | 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 | | Aspect | Detail |
|--------|--------| |--------|--------|
| **Purpose** | Called when a request finishes — releases eviction protection | | **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()` | | **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` | | **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
+14 -13
View File
@@ -281,9 +281,11 @@ class FullComponent(TreeComponent):
root = self.tree_core.root_node root = self.tree_core.root_node
cur = 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: 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 cur = cur.parent
# Lock the device-on segment up to root # Lock the device-on segment up to root
@@ -307,7 +309,7 @@ class FullComponent(TreeComponent):
def release_component_lock( def release_component_lock(
self, self,
node: UnifiedTreeNode, node: UnifiedTreeNode,
params: Optional[DecLockRefParams], params: DecLockRefParams,
lock_host: bool = False, lock_host: bool = False,
) -> None: ) -> None:
ct = self.component_type ct = self.component_type
@@ -315,7 +317,6 @@ class FullComponent(TreeComponent):
cd = node.component_data[ct] cd = node.component_data[ct]
if cd.host_lock_ref == 0: if cd.host_lock_ref == 0:
return return
# Mirror of `acquire`. write_back uses a pure counter.
if cd.host_value is None and not self.tree_core.is_write_back: if cd.host_value is None and not self.tree_core.is_write_back:
return return
cd.host_lock_ref -= 1 cd.host_lock_ref -= 1
@@ -323,17 +324,13 @@ class FullComponent(TreeComponent):
return return
root = self.tree_core.root_node root = self.tree_core.root_node
skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
cur = node cur = node
while cur != root: while cur != root:
if cur.id in skip_lock_node_ids:
cur = cur.parent
continue
cd = cur.component_data[ct] cd = cur.component_data[ct]
assert cd.value is not None assert cd.lock_ref > 0, (
assert cd.lock_ref > 0 f"FULL segment release hit lock_ref=0 on node {cur.id}"
)
if cd.lock_ref == 1: if cd.lock_ref == 1 and cd.value is not None:
key_len = len(cd.value) key_len = len(cd.value)
self.tree_core.component_evictable_size_[ct] += key_len self.tree_core.component_evictable_size_[ct] += key_len
self.tree_core.component_protected_size_[ct] -= key_len self.tree_core.component_protected_size_[ct] -= key_len
@@ -422,7 +419,11 @@ class FullComponent(TreeComponent):
n_len = len(cd.host_value) n_len = len(cd.host_value)
cd.value = device_indices[offset : offset + n_len].clone() cd.value = device_indices[offset : offset + n_len].clone()
offset += n_len offset += n_len
# Full uses leaf sets, not LRU # 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.component_evictable_size_[ct] += n_len
self.tree_core._update_evictable_leaf_sets(n) self.tree_core._update_evictable_leaf_sets(n)
python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py
+23 -31
View File
@@ -233,14 +233,8 @@ class MambaComponent(TreeComponent):
self._emit_excess_path_states_eviction(node, cache_actions) self._emit_excess_path_states_eviction(node, cache_actions)
return return
if node.component_data[self.component_type].value is None: if node.component_data[self.component_type].value is None:
node.component_data[self.component_type].value = params.mamba_value self.tree_core.set_component_device_value(
# move from host LRU to device LRU node.id, self.component_type, params.mamba_value
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
) )
node.last_access_time = get_and_increase_time_counter() node.last_access_time = get_and_increase_time_counter()
self._emit_excess_path_states_eviction(node, cache_actions) self._emit_excess_path_states_eviction(node, cache_actions)
@@ -441,19 +435,17 @@ class MambaComponent(TreeComponent):
return result return result
cd = node.component_data[ct] cd = node.component_data[ct]
value = cd.host_value if lock_host else cd.value 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. # Tombstones are counted too; ledger/LRU track only data-bearing
if value is None: # nodes (a value materialized under lock is credited to protected
result.skip_lock_node_ids.setdefault(ct, set()).add(node.id) # at the materialization site).
return result
if lock_host: 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] host_lru = self.tree_core.host_lru_lists[ct]
if host_lru.in_list(node): if host_lru.in_list(node):
host_lru.remove_node(node) host_lru.remove_node(node)
cd.host_lock_ref += 1 cd.host_lock_ref += 1
else: else:
if cd.lock_ref == 0: if cd.lock_ref == 0 and value is not None:
vlen = len(value) vlen = len(value)
self.tree_core.component_evictable_size_[ct] -= vlen self.tree_core.component_evictable_size_[ct] -= vlen
self.tree_core.component_protected_size_[ct] += vlen self.tree_core.component_protected_size_[ct] += vlen
@@ -463,32 +455,36 @@ class MambaComponent(TreeComponent):
def release_component_lock( def release_component_lock(
self, self,
node: UnifiedTreeNode, node: UnifiedTreeNode,
params: Optional[DecLockRefParams], params: DecLockRefParams,
lock_host: bool = False, lock_host: bool = False,
) -> None: ) -> None:
ct = self.component_type ct = self.component_type
if node is self.tree_core.root_node: if node is self.tree_core.root_node:
return return
cd = node.component_data[ct] 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 value = cd.host_value if lock_host else cd.value
if lock_host: 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 cd.host_lock_ref -= 1
if cd.host_lock_ref == 0 and cd.value is None and cd.host_value is not None: 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] host_lru = self.tree_core.host_lru_lists[ct]
if not host_lru.in_list(node): if not host_lru.in_list(node):
host_lru.insert_mru(node) host_lru.insert_mru(node)
self.tree_core._update_evictable_leaf_sets(node)
return return
if cd.lock_ref > 0: assert cd.lock_ref > 0, f"Mamba release hit lock_ref=0 on node {node.id}"
if cd.lock_ref == 1: if cd.lock_ref == 1 and value is not None:
vlen = len(value) vlen = len(value)
self.tree_core.component_evictable_size_[ct] += vlen self.tree_core.component_evictable_size_[ct] += vlen
self.tree_core.component_protected_size_[ct] -= vlen self.tree_core.component_protected_size_[ct] -= vlen
cd.lock_ref -= 1 cd.lock_ref -= 1
if cd.lock_ref == 0:
self.tree_core._update_evictable_leaf_sets(node)
def _alloc_mamba_slot(self) -> torch.Tensor: def _alloc_mamba_slot(self) -> torch.Tensor:
"""Allocate one mamba pool slot, evicting if necessary.""" """Allocate one mamba pool slot, evicting if necessary."""
@@ -809,15 +805,11 @@ class MambaComponent(TreeComponent):
return return
transfer = transfers[0] transfer = transfers[0]
if transfer.device_indices is not None: if transfer.device_indices is not None:
cd = node.component_data[ct] # The materialization primitive owns the ledger/LRU moves,
cd.value = transfer.device_indices.clone() # including crediting protected when restored under lock.
count = len(cd.value) self.tree_core.set_component_device_value(
# Move from host LRU to device LRU node.id, ct, transfer.device_indices.clone()
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
elif phase == CacheTransferPhase.PREFETCH: elif phase == CacheTransferPhase.PREFETCH:
if not transfers: 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)}" 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) super().__init__(cache, params)
self._session_leaf_covered_len: dict[str, dict[UnifiedTreeNode, int]] = {} self._session_leaf_covered_len: dict[str, dict[UnifiedTreeNode, int]] = {}
self.sliding_window_size = params.sliding_window_size self.sliding_window_size = params.sliding_window_size
@@ -388,9 +391,9 @@ class SWAComponent(TreeComponent):
full_cd = node.component_data[BASE_COMPONENT_TYPE] full_cd = node.component_data[BASE_COMPONENT_TYPE]
swa_evicted_seqlen = params.swa_evicted_seqlen swa_evicted_seqlen = params.swa_evicted_seqlen
assert node.component_data[self.component_type].lock_ref == 0, ( # A locked tombstone is legal (segment locks count every node); the
f"tombstone {self.component_type} lock_ref should be 0, node {node.id}" # 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, ( assert swa_evicted_seqlen % self.tree_core.page_size == 0, (
f"{self.component_type}: swa_evicted_seqlen must be page-aligned, {swa_evicted_seqlen=}" 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 ct = self.component_type
if node.component_data[ct].value is not None: if node.component_data[ct].value is not None:
return 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 swa_evicted_seqlen = params.swa_evicted_seqlen
assert swa_evicted_seqlen % self.tree_core.page_size == 0, ( assert swa_evicted_seqlen % self.tree_core.page_size == 0, (
f"{ct}: swa_evicted_seqlen must be page-aligned, {swa_evicted_seqlen=}" 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[ new_parent.component_data[self.component_type].lock_ref = child.component_data[
self.component_type self.component_type
].lock_ref ].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[ new_parent.component_data[
self.component_type self.component_type
].session_ref = child.component_data[self.component_type].session_ref ].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 parent_swa_data.metadata["host_uuid"] = host_uuid
host_lru = self.tree_core.host_lru_lists[self.component_type] 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 ( if (
new_parent.component_data[self.component_type].value is None new_parent.component_data[self.component_type].value is None
and parent_swa_data.host_lock_ref == 0 and parent_swa_data.host_lock_ref == 0
@@ -622,11 +627,16 @@ class SWAComponent(TreeComponent):
): ):
host_lru.insert_mru(child) 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"] = ( 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.get("uuid")
) )
child.component_data[self.component_type].metadata.pop("uuid", None) 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( def evict_component(
self, self,
@@ -754,10 +764,20 @@ class SWAComponent(TreeComponent):
result: IncLockRefResult, result: IncLockRefResult,
lock_host: bool = False, lock_host: bool = False,
) -> IncLockRefResult: ) -> 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 ct = self.component_type
root = self.tree_core.root_node root = self.tree_core.root_node
sliding_window_size = self.sliding_window_size sliding_window_size = self.sliding_window_size
swa_lock_size = 0 covered = 0
swa_uuid = None swa_uuid = None
uuid_key = "host_uuid" if lock_host else "uuid" uuid_key = "host_uuid" if lock_host else "uuid"
lru = ( lru = (
@@ -766,33 +786,25 @@ class SWAComponent(TreeComponent):
else self.tree_core.lru_lists[ct] 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 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] comp = cur.component_data[ct]
value = comp.host_value if lock_host else comp.value 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 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 lock_host:
if lru.in_list(cur): if lru.in_list(cur):
lru.remove_node(cur) lru.remove_node(cur)
else: else:
key_len = len(cur.key) key_len = len(value)
self.tree_core.component_evictable_size_[ct] -= key_len self.tree_core.component_evictable_size_[ct] -= key_len
self.tree_core.component_protected_size_[ct] += key_len self.tree_core.component_protected_size_[ct] += key_len
if lock_host: if lock_host:
comp.host_lock_ref = ref + 1 comp.host_lock_ref = ref + 1
else: else:
comp.lock_ref = ref + 1 comp.lock_ref = ref + 1
swa_lock_size += len(value) covered += len(cur.key)
if swa_lock_size >= sliding_window_size: if covered >= sliding_window_size:
if comp.metadata.get(uuid_key) is None: if comp.metadata.get(uuid_key) is None:
comp.metadata[uuid_key] = next_component_uuid() comp.metadata[uuid_key] = next_component_uuid()
swa_uuid = comp.metadata[uuid_key] swa_uuid = comp.metadata[uuid_key]
@@ -807,45 +819,47 @@ class SWAComponent(TreeComponent):
def release_component_lock( def release_component_lock(
self, self,
node: UnifiedTreeNode, node: UnifiedTreeNode,
params: Optional[DecLockRefParams], params: DecLockRefParams,
lock_host: bool = False, lock_host: bool = False,
) -> None: ) -> None:
ct = self.component_type ct = self.component_type
root = self.tree_core.root_node root = self.tree_core.root_node
swa_uuid_for_lock = ( swa_uuid_for_lock = (
(params.swa_uuid_for_host_lock if lock_host else params.swa_uuid_for_lock) params.swa_uuid_for_host_lock if lock_host else params.swa_uuid_for_lock
if params
else None
) )
skip_lock_node_ids = params.skip_lock_node_ids.get(ct, ()) if params else ()
dec_swa = True dec_swa = True
uuid_key = "host_uuid" if lock_host else "uuid" 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 cur = node
while cur != root and dec_swa: while cur != root and dec_swa:
comp = cur.component_data[ct] 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 ref = comp.host_lock_ref if lock_host else comp.lock_ref
if ref == 0: # Acquire counted every segment node and splits copy refs, so a
cur = cur.parent # zero here means the release does not mirror its acquire.
continue assert ref > 0, (
if ref == 1: 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 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] host_lru = self.tree_core.host_lru_lists[ct]
if not host_lru.in_list(cur): if not host_lru.in_list(cur):
host_lru.insert_mru(cur) host_lru.insert_mru(cur)
else: else:
key_len = len(comp.value) key_len = len(value)
self.tree_core.component_evictable_size_[ct] += key_len self.tree_core.component_evictable_size_[ct] += key_len
self.tree_core.component_protected_size_[ct] -= key_len self.tree_core.component_protected_size_[ct] -= key_len
if lock_host: if lock_host:
comp.host_lock_ref = ref - 1 comp.host_lock_ref = ref - 1
else: else:
comp.lock_ref = ref - 1 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: if swa_uuid_for_lock and comp.metadata.get(uuid_key) == swa_uuid_for_lock:
dec_swa = False dec_swa = False
cur = cur.parent cur = cur.parent
@@ -857,15 +871,14 @@ class SWAComponent(TreeComponent):
device_frees: dict[ComponentType, list[torch.Tensor]], device_frees: dict[ComponentType, list[torch.Tensor]],
host_frees: dict[ComponentType, list[torch.Tensor]], host_frees: dict[ComponentType, list[torch.Tensor]],
) -> None: ) -> None:
"""Early-release the SWA lock along [node, swa_uuid_for_lock] while """Early-release the SWA lock along [node, swa_uuid_for_lock]; this
leaving Full and Mamba locks intact. 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 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 window. The caller must invoke this at most once per
Full lock must stay so the request's prefix is protected. (node, swa_uuid_for_lock) pair.
Caller (UnifiedRadixCache.dec_swa_lock_only) must ensure this is
invoked at most once per (node, swa_uuid_for_lock) pair.
""" """
ct = self.component_type ct = self.component_type
root = self.tree_core.root_node root = self.tree_core.root_node
@@ -873,19 +886,16 @@ class SWAComponent(TreeComponent):
cur = node cur = node
while cur is not root: while cur is not root:
cd = cur.component_data[ct] cd = cur.component_data[ct]
# Acquire skips tombstoned nodes; release must skip them too. Same assert cd.lock_ref > 0, (
# for nodes with lock_ref == 0 — acquire never credited them. f"SWA window release hit lock_ref=0 on node {cur.id}"
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
cd.lock_ref -= 1 cd.lock_ref -= 1
if cd.lock_ref == 0: if cd.lock_ref == 0:
key_len = len(cur.key) self.tree_core._update_evictable_leaf_sets(cur)
self.tree_core.component_protected_size_[ct] -= key_len if cd.lock_ref == 0 and cd.value is not None:
self.tree_core.component_evictable_size_[ct] += key_len 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): if self.tree_core._is_device_leaf(cur):
self.tree_core._evict_component_and_detach_lru( self.tree_core._evict_component_and_detach_lru(
cur, 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( def release_component_lock(
self, self,
node: UnifiedTreeNode, node: UnifiedTreeNode,
params: Optional[DecLockRefParams], params: DecLockRefParams,
lock_host: bool = False, lock_host: bool = False,
) -> None: ) -> None:
"""Decrement component lock refs, un-protecting nodes. """Decrement component lock refs, un-protecting nodes.
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+79 -34
View File
@@ -622,32 +622,66 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = () self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = ()
) -> IncLockRefResult: ) -> IncLockRefResult:
node = self.node_by_id(node_id) 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: for component in self.components:
if component.component_type in skip_lock_components: if component.component_type in skipped:
# 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 continue
result = component.acquire_component_lock(node=node, result=result) result = component.acquire_component_lock(node=node, result=result)
self._update_evictable_leaf_sets(node) self._update_evictable_leaf_sets(node)
return result 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( def dec_lock_ref(
self, self,
node_id: NodeId, node_id: NodeId,
params: Optional[DecLockRefParams] = None, params: DecLockRefParams,
skip_swa: bool = False, skip_swa: bool = False,
) -> DecLockRefResult: ) -> DecLockRefResult:
node = self.node_by_id(node_id) node = self.node_by_id(node_id)
for component in self.components: self._assert_receipt_anchor(node, params)
if skip_swa and component.component_type == ComponentType.SWA: # After an SWA early release (dec_swa_lock_only), SWA and the
continue # lower-priority components it dropped are already released.
component.release_component_lock(node=node, params=params) self._release_components(node, params, skip_swa_and_below=skip_swa)
self._update_evictable_leaf_sets(node) self._update_evictable_leaf_sets(node)
# TODO: delta is not aggregated from components; no caller uses it yet. # TODO: delta is not aggregated from components; no caller uses it yet.
return DecLockRefResult() return DecLockRefResult()
@@ -655,36 +689,33 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
def dec_swa_lock_only( def dec_swa_lock_only(
self, self,
node_id: NodeId, node_id: NodeId,
swa_uuid_for_lock: Optional[int], params: DecLockRefParams,
skip_lock_node_ids: Optional[dict] = None,
) -> DecSwaLockOnlyResult: ) -> DecSwaLockOnlyResult:
"""Early-release the SWA portion of a request's tree lock, plus any """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.""" strictly-lower-priority locks (e.g. Mamba) co-located on the node."""
result = DecSwaLockOnlyResult() result = DecSwaLockOnlyResult()
node = self.node_by_id(node_id) node = self.node_by_id(node_id)
self._assert_receipt_anchor(node, params)
swa_component = self.components_by_type.get(ComponentType.SWA) swa_component = self.components_by_type.get(ComponentType.SWA)
if swa_component is None: if swa_component is None:
return result return result
swa_component.release_window_lock( 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, # Drop strictly-lower-priority locks co-located on the node, skipping
# honoring skip ids so we don't drop a lock a partial inc never took # any the paired inc never took (matters for FULL+SWA+MAMBA models).
# (matters for FULL+SWA+MAMBA models, e.g. Inkling).
swa_priority = swa_component.eviction_priority(is_leaf=False) swa_priority = swa_component.eviction_priority(is_leaf=False)
dec_params = DecLockRefParams( for comp in reversed(self.components):
swa_uuid_for_lock=swa_uuid_for_lock, if comp.component_type in params.skipped_lock_components:
skip_lock_node_ids=skip_lock_node_ids or {}, continue
)
for comp in self.components:
if comp.eviction_priority(is_leaf=False) < swa_priority: if comp.eviction_priority(is_leaf=False) < swa_priority:
comp.release_component_lock(node, dec_params) comp.release_component_lock(node, params)
return result return result
def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult: def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult:
node = self.node_by_id(node_id) node = self.node_by_id(node_id)
result = IncLockRefResult() result = IncLockRefResult(node_id=node.id)
for component in self.components: for component in self.components:
result = component.acquire_component_lock( result = component.acquire_component_lock(
node=node, result=result, lock_host=True node=node, result=result, lock_host=True
@@ -693,11 +724,11 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
return result return result
def dec_host_lock_ref( def dec_host_lock_ref(
self, node_id: NodeId, params: Optional[DecLockRefParams] = None self, node_id: NodeId, params: DecLockRefParams
) -> DecLockRefResult: ) -> DecLockRefResult:
node = self.node_by_id(node_id) node = self.node_by_id(node_id)
for component in self.components: self._assert_receipt_anchor(node, params)
component.release_component_lock(node=node, params=params, lock_host=True) self._release_components(node, params, lock_host=True)
self._update_evictable_leaf_sets(node) self._update_evictable_leaf_sets(node)
return DecLockRefResult() return DecLockRefResult()
@@ -1260,6 +1291,9 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
assert cd.value is None assert cd.value is None
n = len(fresh_value) n = len(fresh_value)
cd.value = fresh_value.clone() cd.value = fresh_value.clone()
if cd.lock_ref > 0:
self.component_protected_size_[ct] += n
else:
self.component_evictable_size_[ct] += n self.component_evictable_size_[ct] += n
self._update_evictable_leaf_sets(node) self._update_evictable_leaf_sets(node)
# A backuped node restored from fresh KV is a duplicate right away. # A backuped node restored from fresh KV is a duplicate right away.
@@ -1887,7 +1921,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
return True return True
def _is_host_leaf(self, node: UnifiedTreeNode) -> bool: 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 Only the Full (base) component host_value is required; auxiliary
components are not mandatory for H-leaf membership. In-flight DMA components are not mandatory for H-leaf membership. In-flight DMA
@@ -1898,6 +1933,10 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
return False return False
if any(cd.host_lock_ref > 0 for cd in node.component_data): if any(cd.host_lock_ref > 0 for cd in node.component_data):
return False 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: if len(node.children) > 0:
return False return False
return True return True
@@ -2258,11 +2297,17 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
# Full uses leaf sets, not LRU; its stores go through the insert paths. # Full uses leaf sets, not LRU; its stores go through the insert paths.
assert component_type != BASE_COMPONENT_TYPE assert component_type != BASE_COMPONENT_TYPE
node = self.node_by_id(node_id) 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] host_lru = self.host_lru_lists[component_type]
if host_lru.in_list(node): if host_lru.in_list(node):
host_lru.remove_node(node) host_lru.remove_node(node)
self.lru_lists[component_type].insert_mru(node) self.lru_lists[component_type].insert_mru(node)
# 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) self.component_evictable_size_[component_type] += len(value)
def get_component_device_value( def get_component_device_value(
@@ -2361,8 +2406,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
E(f"node {nid} {ct} host_lock_ref={cd.host_lock_ref}") E(f"node {nid} {ct} host_lock_ref={cd.host_lock_ref}")
if ct != FCT and fl < cd.lock_ref: if ct != FCT and fl < cd.lock_ref:
E(f"node {nid} full_lock={fl} < {ct}_lock={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: # Locked tombstones are legal: segment locks count every
E(f"node {nid} {ct} evicted but lock_ref={cd.lock_ref}") # node in [start, boundary], data-bearing or not.
# Collect expected leaf qualification (single pass) # Collect expected leaf qualification (single pass)
if self._is_device_leaf(node): 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( def inc_lock_ref(
self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = () self, node_id: NodeId, skip_lock_components: Sequence[ComponentType] = ()
) -> IncLockRefResult: ) -> IncLockRefResult:
"""Bump the reference count on a node's component locks, leaving any """Bump the reference count on a node's component locks. Components in
component in skip_lock_components evictable and recorded in the result.""" ``skip_lock_components`` are left untaken; the receipt records the
anchor node and the skipped set so the paired release mirrors them."""
... ...
@abstractmethod @abstractmethod
def dec_lock_ref( def dec_lock_ref(
self, self,
node_id: NodeId, node_id: NodeId,
params: Optional[DecLockRefParams] = None, params: DecLockRefParams,
skip_swa: bool = False, skip_swa: bool = False,
) -> DecLockRefResult: ) -> 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 @abstractmethod
def dec_swa_lock_only( def dec_swa_lock_only(
self, self,
node_id: NodeId, node_id: NodeId,
swa_uuid_for_lock: Optional[int], params: DecLockRefParams,
skip_lock_node_ids: Optional[dict] = None,
) -> DecSwaLockOnlyResult: ) -> DecSwaLockOnlyResult:
"""Decrease only the SWA (and lower-priority co-located) reference """Decrease only the SWA (and lower-priority co-located) reference
counts; the result carries the freed slots.""" counts; the result carries the freed slots."""
@@ -313,7 +314,7 @@ class UnifiedTreeCoreInterface(ABC):
@abstractmethod @abstractmethod
def dec_host_lock_ref( def dec_host_lock_ref(
self, node_id: NodeId, params: Optional[DecLockRefParams] = None self, node_id: NodeId, params: DecLockRefParams
) -> DecLockRefResult: ) -> DecLockRefResult:
"""Decrease the reference count on a node's host-side component locks.""" """Decrease the reference count on a node's host-side component locks."""
... ...
python/sglang/srt/mem_cache/unified_radix_cache.py
+14 -25
View File
@@ -878,7 +878,7 @@ class UnifiedRadixCache(BasePrefixCache):
def dec_lock_ref( def dec_lock_ref(
self, self,
node_id: NodeId, node_id: NodeId,
params: Optional[DecLockRefParams] = None, params: DecLockRefParams,
skip_swa: bool = False, skip_swa: bool = False,
) -> DecLockRefResult: ) -> DecLockRefResult:
result = self.session.try_dec_lock_ref(node_id, params) 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) return self.tree_core.dec_lock_ref(node_id, params, skip_swa)
def _dec_req_lock(self, req: Req, *, skip_swa: bool = False) -> None: 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 """Release the tree lock a request holds on its last_node with the
components it skipped locking so it never drops a lock it never took.""" receipt its acquire returned, so it never drops a lock it never took."""
self.dec_lock_ref( self.dec_lock_ref(req.last_node, req.lock_receipt, skip_swa=skip_swa)
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( def dec_swa_lock_only(
self, self,
node_id: NodeId, node_id: NodeId,
swa_uuid_for_lock: Optional[int] = None, params: DecLockRefParams,
skip_lock_node_ids: Optional[dict] = None,
) -> None: ) -> None:
if self.disable: if self.disable:
return return
result = self.tree_core.dec_swa_lock_only( result = self.tree_core.dec_swa_lock_only(node_id, params)
node_id, swa_uuid_for_lock, skip_lock_node_ids
)
self._free_values(result.device_frees, result.host_frees) self._free_values(result.device_frees, result.host_frees)
def inc_host_lock_ref(self, node_id: NodeId) -> IncLockRefResult: 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) return self.tree_core.inc_host_lock_ref(node_id)
def dec_host_lock_ref( def dec_host_lock_ref(
self, node_id: NodeId, params: Optional[DecLockRefParams] = None self, node_id: NodeId, params: DecLockRefParams
) -> DecLockRefResult: ) -> DecLockRefResult:
if self.disable: if self.disable:
return DecLockRefResult() return DecLockRefResult()
@@ -1099,20 +1089,20 @@ class UnifiedRadixCache(BasePrefixCache):
new_indices[req.kv.cache_protected_len :], 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 # 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). # 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 # Safe only because any future COW source is the COWing request's own
# admission-locked last_node (recorded only if still present, locked before # admission-locked last_node (recorded only if still present, locked before
# the next alloc) -- not this evictable node. A scheduler that matched a # the next alloc) -- not this evictable node. A scheduler that matched a
# whole batch before locking would break that. Off = original full lock. # whole batch before locking would break that. Off = original full lock.
skip_lock_components = ( lock_result = self.inc_lock_ref(
new_last_node,
skip_lock_components=(
(ComponentType.MAMBA,) (ComponentType.MAMBA,)
if envs.SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK.get() if envs.SGLANG_OPT_MAMBA_SKIP_DECODE_LOCK.get()
else () else ()
) ),
lock_result = self.inc_lock_ref(
new_last_node, skip_lock_components=skip_lock_components
) )
# Update req fields # Update req fields
@@ -1124,9 +1114,8 @@ class UnifiedRadixCache(BasePrefixCache):
req.prefix_indices = new_indices req.prefix_indices = new_indices
req.kv.cache_protected_len = len(new_indices) req.kv.cache_protected_len = len(new_indices)
req.last_node = new_last_node req.last_node = new_last_node
req.swa_uuid_for_lock = lock_result.swa_uuid_for_lock # Carry the receipt so this node's dec releases only what we locked.
# carry the skip set so this node's dec releases only what we locked req.lock_receipt = lock_result.to_dec_params()
req.skip_lock_node_ids = lock_result.skip_lock_node_ids
# The rematch acquired a new SWA prefix lock. # The rematch acquired a new SWA prefix lock.
req.swa_prefix_lock_released = False 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) # First req's radix tree node (for dec_lock_ref on session close)
last_node: Any = None last_node: Any = None
swa_uuid_for_lock: Optional[str] = None # Receipt of the first request's tree lock on last_node.
# components the first req skipped locking on last_node, so release dec lock_receipt: DecLockRefParams = field(default_factory=DecLockRefParams)
# releases only what it took (may share the node with another req). # Whether the first request already released its SWA lock.
skip_lock_node_ids: dict = field(default_factory=dict) swa_prefix_lock_released: bool = False
def save_from_req(self, req: Req, is_first: bool): def save_from_req(self, req: Req, is_first: bool):
"""Save KV state from a finishing request into this slot.""" """Save KV state from a finishing request into this slot."""
kv = req.detach_kv() kv = req.detach_kv()
if is_first: if is_first:
self.last_node = req.last_node self.last_node = req.last_node
self.swa_uuid_for_lock = req.swa_uuid_for_lock self.lock_receipt = req.lock_receipt
self.skip_lock_node_ids = req.skip_lock_node_ids self.swa_prefix_lock_released = req.swa_prefix_lock_released
# The slot takes over the request's KV record. # The slot takes over the request's KV record.
self.kv = kv self.kv = kv
else: else:
@@ -71,8 +71,8 @@ class SessionSlot:
def restore_to_req(self, req: Req): def restore_to_req(self, req: Req):
"""Restore KV state from this slot into an incoming request.""" """Restore KV state from this slot into an incoming request."""
req.kv = self.kv req.kv = self.kv
req.swa_uuid_for_lock = self.swa_uuid_for_lock req.lock_receipt = self.lock_receipt
req.skip_lock_node_ids = self.skip_lock_node_ids req.swa_prefix_lock_released = self.swa_prefix_lock_released
# NOTE: the slot keeps sharing the record it just handed out. During # NOTE: the slot keeps sharing the record it just handed out. During
# chunked prefill, a request may be rejected by # chunked prefill, a request may be rejected by
@@ -290,8 +290,8 @@ class StreamingSession(BasePrefixCache):
slot = SessionSlot( slot = SessionSlot(
kv=kv, kv=kv,
last_node=req.last_node, last_node=req.last_node,
swa_uuid_for_lock=req.swa_uuid_for_lock, lock_receipt=req.lock_receipt,
skip_lock_node_ids=req.skip_lock_node_ids, swa_prefix_lock_released=req.swa_prefix_lock_released,
) )
self.slots[session_id] = slot self.slots[session_id] = slot
else: else:
@@ -396,13 +396,10 @@ class StreamingSession(BasePrefixCache):
) )
if lock_node is not None: if lock_node is not None:
self.inner.dec_lock_ref( # skip_swa is an SWA-cache extension kwarg; a slot can only have
lock_node, # early-released when the inner cache supports SWA locks.
DecLockRefParams( skip = {"skip_swa": True} if slot.swa_prefix_lock_released else {}
swa_uuid_for_lock=slot.swa_uuid_for_lock, self.inner.dec_lock_ref(lock_node, slot.lock_receipt, **skip)
skip_lock_node_ids=slot.skip_lock_node_ids,
),
)
if slot.kv.holds_kv: if slot.kv.holds_kv:
self.free_kv_row(slot.kv, [(protected_len, slot.kv.kv_allocated_len)]) 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: class ScriptedLockRefExhauster:
def __init__(self, scheduler: Scheduler) -> None: def __init__(self, scheduler: Scheduler) -> None:
self.scheduler = scheduler 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: def exhaust(self, *, leave_refs: int) -> None:
tree_cache = self.scheduler.tree_cache tree_cache = self.scheduler.tree_cache
@@ -24,17 +25,17 @@ class ScriptedLockRefExhauster:
return return
target = evictable[0] 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] newly_locked = [node for node in evictable if _node_lock_ref(node) > 0]
if not newly_locked: if not newly_locked:
return return
self._locked.append(target) self._locked.append((target, result.to_dec_params()))
def release(self) -> None: def release(self) -> None:
tree_cache = self.scheduler.tree_cache tree_cache = self.scheduler.tree_cache
for node in self._locked: for node, dec_params in self._locked:
tree_cache.dec_lock_ref(to_node_handle(tree_cache, node)) tree_cache.dec_lock_ref(to_node_handle(tree_cache, node), dec_params)
self._locked.clear() self._locked.clear()
def _evictable_nodes(self) -> List[Any]: def _evictable_nodes(self) -> List[Any]:
rust/sglang-radix-tree/src/components/full.rs
+21 -31
View File
@@ -2,7 +2,7 @@
//! the rest from the `TreeComponent` defaults. //! the rest from the `TreeComponent` defaults.
use std::cmp::Reverse; use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap, HashSet}; use std::collections::{BinaryHeap, HashMap};
use tch::{Kind, Tensor}; use tch::{Kind, Tensor};
@@ -277,8 +277,6 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
mut result: IncLockRefResult, mut result: IncLockRefResult,
lock_host: bool, lock_host: bool,
) -> IncLockRefResult { ) -> IncLockRefResult {
let ct = FULL;
// Only the last host node needs to be protected. // Only the last host node needs to be protected.
if lock_host { if lock_host {
let node = tree_core.arena.node_mut(node_id); let node = tree_core.arena.node_mut(node_id);
@@ -291,20 +289,17 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
return result; return result;
} }
// Skip the bottom evicted segment, recording it for the matching release. // The bottom device-evicted segment is locked too (no ledger move —
let on_boundary = |node: &Node<K>| node.is_root() || node.has_device_value(FULL); // nothing is on device); a load-back that materializes a value under
// lock credits protected directly.
let mut cur = node_id; let mut cur = node_id;
let mut node = tree_core.arena.node(cur);
if !on_boundary(node) {
let skip_lock_node_ids = result.skip_lock_node_ids.entry(ct).or_default();
loop { loop {
skip_lock_node_ids.insert(node.id); let node = tree_core.arena.node_mut(cur);
cur = node.parent(); if node.is_root() || node.has_device_value(FULL) {
node = tree_core.arena.node(cur);
if on_boundary(node) {
break; break;
} }
} node.inc_device_lock_ref(FULL);
cur = node.parent();
} }
// Lock the device-on segment up to the root. // Lock the device-on segment up to the root.
@@ -341,11 +336,9 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
&self, &self,
tree_core: &mut UnifiedTreeCore<K>, tree_core: &mut UnifiedTreeCore<K>,
node_id: NodeIdx_, node_id: NodeIdx_,
params: Option<&DecLockRefParams>, _params: &DecLockRefParams,
lock_host: bool, lock_host: bool,
) { ) {
let ct = FULL;
if lock_host { if lock_host {
let node = tree_core.arena.node_mut(node_id); let node = tree_core.arena.node_mut(node_id);
if node.host_lock_ref(FULL) == 0 { if node.host_lock_ref(FULL) == 0 {
@@ -360,10 +353,6 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
return; return;
} }
let empty = HashSet::new();
let skip_lock_node_ids = params
.and_then(|p| p.skip_lock_node_ids.get(&ct))
.unwrap_or(&empty);
let mut cur = node_id; let mut cur = node_id;
loop { loop {
let node = tree_core.arena.node_mut(cur); let node = tree_core.arena.node_mut(cur);
@@ -371,20 +360,12 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
break; break;
} }
let parent = node.parent(); let parent = node.parent();
if skip_lock_node_ids.contains(&node.id) {
cur = parent;
continue;
}
assert!(
node.has_device_value(FULL),
"release_component_lock: node {cur} has no FULL device value"
);
let old_lock_ref = node.device_lock_ref(FULL); let old_lock_ref = node.device_lock_ref(FULL);
assert!( assert!(
old_lock_ref > 0, old_lock_ref > 0,
"release_component_lock: node {cur} is not locked" "FULL segment release hit lock_ref=0 on node {cur}"
); );
let newly_unlocked_len = if old_lock_ref == 1 { let newly_unlocked_len = if old_lock_ref == 1 && node.has_device_value(FULL) {
Some(node.device_value_len(FULL)) Some(node.device_value_len(FULL))
} else { } else {
None None
@@ -393,6 +374,8 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
if let Some(key_len) = newly_unlocked_len { if let Some(key_len) = newly_unlocked_len {
tree_core.dec_protected_size(FULL, key_len); tree_core.dec_protected_size(FULL, key_len);
tree_core.inc_evictable_size(FULL, key_len); tree_core.inc_evictable_size(FULL, key_len);
}
if old_lock_ref == 1 {
tree_core.update_evictable_leaf_sets_(cur); tree_core.update_evictable_leaf_sets_(cur);
} }
cur = parent; cur = parent;
@@ -478,9 +461,16 @@ impl<K: ChildKeyType> TreeComponent<K> for FullComponent {
let n_len = loaded.host_value_len(FULL) as i64; let n_len = loaded.host_value_len(FULL) as i64;
loaded loaded
.set_device_value(FULL, device_indices.narrow(0, offset, n_len).copy()); .set_device_value(FULL, device_indices.narrow(0, offset, n_len).copy());
let locked = loaded.device_lock_ref(FULL) > 0;
offset += n_len; offset += n_len;
// Full uses leaf sets, not LRU. // Full uses leaf sets, not LRU. A value materialized
// under lock is protected; the last release moves it
// to evictable.
if locked {
tree_core.inc_protected_size(FULL, n_len as usize);
} else {
tree_core.inc_evictable_size(FULL, n_len as usize); tree_core.inc_evictable_size(FULL, n_len as usize);
}
tree_core.update_evictable_leaf_sets_(loaded_idx); tree_core.update_evictable_leaf_sets_(loaded_idx);
} }
} }
rust/sglang-radix-tree/src/components/mamba.rs
+27 -47
View File
@@ -178,16 +178,7 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
return; return;
} }
if !tree_core.arena.has_device_value(node_id, MAMBA) { if !tree_core.arena.has_device_value(node_id, MAMBA) {
// Tombstone refill: the node moves from the host LRU to the device LRU. tree_core.set_component_device_value_(node_id, MAMBA, mamba_value.shallow_clone());
tree_core
.arena
.set_device_value(node_id, MAMBA, mamba_value.shallow_clone());
let host_lru = tree_core.host_lru_list_mut(MAMBA);
if host_lru.in_list(Some(node_id)) {
host_lru.remove_node(node_id);
}
tree_core.device_lru_list_mut(MAMBA).insert_mru(node_id);
tree_core.inc_evictable_size(MAMBA, slot_len);
let tick = tree_core.arena.get_and_bump_access_counter(); let tick = tree_core.arena.get_and_bump_access_counter();
tree_core.arena.node_mut(node_id).last_access_counter = tick; tree_core.arena.node_mut(node_id).last_access_counter = tick;
self.emit_excess_path_states_eviction_(tree_core.arena.node(node_id).id, cache_actions); self.emit_excess_path_states_eviction_(tree_core.arena.node(node_id).id, cache_actions);
@@ -417,24 +408,19 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
&self, &self,
tree_core: &mut UnifiedTreeCore<K>, tree_core: &mut UnifiedTreeCore<K>,
node_id: NodeIdx_, node_id: NodeIdx_,
mut result: IncLockRefResult, result: IncLockRefResult,
lock_host: bool, lock_host: bool,
) -> IncLockRefResult { ) -> IncLockRefResult {
let node = tree_core.arena.node(node_id); let node = tree_core.arena.node(node_id);
if node.is_root() { if node.is_root() {
return result; return result;
} }
// A node in skip_lock_node_ids was a tombstone when this lock was acquired. // Tombstones are counted too; ledger/LRU track only data-bearing
if !Self::has_value(node, lock_host) { // nodes (a value materialized under lock is credited to protected
result // at the materialization site).
.skip_lock_node_ids let has_value = Self::has_value(node, lock_host);
.entry(MAMBA)
.or_default()
.insert(node.id);
return result;
}
if lock_host { if lock_host {
if node.host_lock_ref(MAMBA) == 0 { if node.host_lock_ref(MAMBA) == 0 && has_value {
let host_lru = tree_core.host_lru_list_mut(MAMBA); let host_lru = tree_core.host_lru_list_mut(MAMBA);
if host_lru.in_list(Some(node_id)) { if host_lru.in_list(Some(node_id)) {
host_lru.remove_node(node_id); host_lru.remove_node(node_id);
@@ -443,7 +429,7 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
tree_core.arena.inc_host_lock_ref(node_id, MAMBA); tree_core.arena.inc_host_lock_ref(node_id, MAMBA);
} else { } else {
let value_len = node.device_value_len(MAMBA); let value_len = node.device_value_len(MAMBA);
if node.device_lock_ref(MAMBA) == 0 { if node.device_lock_ref(MAMBA) == 0 && has_value {
tree_core.dec_evictable_size(MAMBA, value_len); tree_core.dec_evictable_size(MAMBA, value_len);
tree_core.inc_protected_size(MAMBA, value_len); tree_core.inc_protected_size(MAMBA, value_len);
} }
@@ -457,43 +443,44 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
&self, &self,
tree_core: &mut UnifiedTreeCore<K>, tree_core: &mut UnifiedTreeCore<K>,
node_id: NodeIdx_, node_id: NodeIdx_,
params: Option<&DecLockRefParams>, _params: &DecLockRefParams,
lock_host: bool, lock_host: bool,
) { ) {
if tree_core.arena.node(node_id).is_root() { if tree_core.arena.node(node_id).is_root() {
return; return;
} }
if let Some(params) = params
&& params
.skip_lock_node_ids
.get(&MAMBA)
.is_some_and(|ids| ids.contains(&tree_core.arena.node(node_id).id))
{
return;
}
if lock_host { if lock_host {
let node = tree_core.arena.node_mut(node_id); let node = tree_core.arena.node_mut(node_id);
assert!(
node.host_lock_ref(MAMBA) > 0,
"Mamba release hit host_lock_ref=0 on node {node_id}"
);
node.dec_host_lock_ref(MAMBA); node.dec_host_lock_ref(MAMBA);
if node.host_lock_ref(MAMBA) == 0 if node.host_lock_ref(MAMBA) == 0 {
&& !node.has_device_value(MAMBA) if !node.has_device_value(MAMBA) && node.has_host_value(MAMBA) {
&& node.has_host_value(MAMBA)
{
let host_lru = tree_core.host_lru_list_mut(MAMBA); let host_lru = tree_core.host_lru_list_mut(MAMBA);
if !host_lru.in_list(Some(node_id)) { if !host_lru.in_list(Some(node_id)) {
host_lru.insert_mru(node_id); host_lru.insert_mru(node_id);
} }
} }
tree_core.update_evictable_leaf_sets_(node_id);
}
return; return;
} }
let node = tree_core.arena.node(node_id); let node = tree_core.arena.node(node_id);
let device_lock_ref = node.device_lock_ref(MAMBA); let device_lock_ref = node.device_lock_ref(MAMBA);
if device_lock_ref > 0 { assert!(
if device_lock_ref == 1 { device_lock_ref > 0,
"Mamba release hit lock_ref=0 on node {node_id}"
);
if device_lock_ref == 1 && node.has_device_value(MAMBA) {
let value_len = node.device_value_len(MAMBA); let value_len = node.device_value_len(MAMBA);
tree_core.inc_evictable_size(MAMBA, value_len); tree_core.inc_evictable_size(MAMBA, value_len);
tree_core.dec_protected_size(MAMBA, value_len); tree_core.dec_protected_size(MAMBA, value_len);
} }
tree_core.arena.dec_device_lock_ref(node_id, MAMBA); tree_core.arena.dec_device_lock_ref(node_id, MAMBA);
if device_lock_ref == 1 {
tree_core.update_evictable_leaf_sets_(node_id);
} }
} }
@@ -613,16 +600,9 @@ impl<K: ChildKeyType> TreeComponent<K> for MambaComponent {
return; return;
}; };
if let Some(device_indices) = &transfer.device_indices { if let Some(device_indices) = &transfer.device_indices {
let node = tree_core.arena.node_mut(node_id); // The materialization primitive owns the ledger/LRU moves,
node.set_device_value(MAMBA, device_indices.copy()); // including crediting protected when restored under lock.
let count = node.device_value_len(MAMBA); tree_core.set_component_device_value_(node_id, MAMBA, device_indices.copy());
// Move from host LRU to device LRU
let host_lru = tree_core.host_lru_list_mut(MAMBA);
if host_lru.in_list(Some(node_id)) {
host_lru.remove_node(node_id);
}
tree_core.device_lru_list_mut(MAMBA).insert_mru(node_id);
tree_core.inc_evictable_size(MAMBA, count);
} }
} }
// The python elif chain has no BACKUP_STORAGE arm. // The python elif chain has no BACKUP_STORAGE arm.
rust/sglang-radix-tree/src/components/mod.rs
+44 -1
View File
@@ -344,7 +344,7 @@ pub trait TreeComponent<K: ChildKeyType> {
&self, &self,
tree_core: &mut UnifiedTreeCore<K>, tree_core: &mut UnifiedTreeCore<K>,
node_id: NodeIdx_, node_id: NodeIdx_,
params: Option<&DecLockRefParams>, params: &DecLockRefParams,
lock_host: bool, lock_host: bool,
); );
@@ -466,6 +466,49 @@ pub const BASE_COMPONENT_TYPE: ComponentType = ComponentType::Full;
/// Slots per tier — the arrays are sized to this, not the enabled subset. /// Slots per tier — the arrays are sized to this, not the enabled subset.
pub const NUM_COMPONENT_TYPES: usize = ComponentType::Mamba as usize + 1; pub const NUM_COMPONENT_TYPES: usize = ComponentType::Mamba as usize + 1;
/// A set of component types (bitmask over `ComponentType::idx`), e.g. the
/// components an `inc_lock_ref` left untaken.
#[derive(Copy, Clone, Default, PartialEq, Eq, Debug)]
pub struct ComponentSet(u8);
impl ComponentSet {
pub const EMPTY: ComponentSet = ComponentSet(0);
/// The set holding exactly one component.
pub const fn of(component_type: ComponentType) -> ComponentSet {
ComponentSet(1 << component_type.idx())
}
pub fn insert(&mut self, component_type: ComponentType) {
self.0 |= 1 << component_type.idx();
}
pub const fn contains(self, component_type: ComponentType) -> bool {
self.0 & (1 << component_type.idx()) != 0
}
pub const fn is_empty(self) -> bool {
self.0 == 0
}
/// The members, in component-index order.
pub fn iter(self) -> impl Iterator<Item = ComponentType> {
(0..NUM_COMPONENT_TYPES)
.filter(move |idx| self.0 & (1 << idx) != 0)
.map(ComponentType::from_idx)
}
}
impl FromIterator<ComponentType> for ComponentSet {
fn from_iter<I: IntoIterator<Item = ComponentType>>(iter: I) -> Self {
let mut set = ComponentSet::EMPTY;
for component_type in iter {
set.insert(component_type);
}
set
}
}
impl ComponentType { impl ComponentType {
/// Index into a per-component array. /// Index into a per-component array.
pub const fn idx(self) -> usize { pub const fn idx(self) -> usize {
rust/sglang-radix-tree/src/components/swa.rs
+86 -87
View File
@@ -3,7 +3,7 @@
//! SWA values arrive pool-resolved; the full->SWA index translation happens at //! SWA values arrive pool-resolved; the full->SWA index translation happens at
//! the cache boundary. //! the cache boundary.
use std::collections::{HashMap, HashSet}; use std::collections::HashMap;
use tch::{Kind, Tensor}; use tch::{Kind, Tensor};
@@ -34,10 +34,15 @@ impl SwaComponent {
impl SwaComponent { impl SwaComponent {
/// Build the driver from the tree's init params. /// Build the driver from the tree's init params.
pub fn new(params: &CacheInitParams) -> Self { pub fn new(params: &CacheInitParams) -> Self {
SwaComponent { let sliding_window_size = params
sliding_window_size: params
.swa_sliding_window_size .swa_sliding_window_size
.expect("the SWA component requires swa_sliding_window_size"), .expect("the SWA component requires swa_sliding_window_size");
assert!(
sliding_window_size > 0,
"swa_sliding_window_size must be positive"
);
SwaComponent {
sliding_window_size,
} }
} }
@@ -408,11 +413,9 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
} }
let swa_evicted_seqlen = params.swa_evicted_seqlen; let swa_evicted_seqlen = params.swa_evicted_seqlen;
assert_eq!( // A locked tombstone is legal (segment locks count every node); the
node.device_lock_ref(SWA), // full-value swap below is safe because full lock_ref >= swa
0, // lock_ref, so a locked-SWA node always takes the Recover branch.
"tombstone Swa lock_ref should be 0, node {node_id}"
);
assert_eq!( assert_eq!(
swa_evicted_seqlen % tree_core.page_size, swa_evicted_seqlen % tree_core.page_size,
0, 0,
@@ -495,11 +498,6 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
if node.has_device_value(SWA) { if node.has_device_value(SWA) {
return; return;
} }
assert_eq!(
node.device_lock_ref(SWA),
0,
"tombstone Swa lock_ref should be 0 on unevict, node {node_id}"
);
let swa_evicted_seqlen = params.swa_evicted_seqlen; let swa_evicted_seqlen = params.swa_evicted_seqlen;
assert_eq!( assert_eq!(
swa_evicted_seqlen % tree_core.page_size, swa_evicted_seqlen % tree_core.page_size,
@@ -587,26 +585,33 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
let (new_parent, child) = tree_core.arena.node_pair_mut(new_parent_id, child_id); let (new_parent, child) = tree_core.arena.node_pair_mut(new_parent_id, child_id);
let split_len = new_parent.key.atom_len() as i64; let split_len = new_parent.key.atom_len() as i64;
new_parent.copy_device_lock_ref(SWA, child); new_parent.copy_device_lock_ref(SWA, child);
new_parent.copy_host_lock_ref(SWA, child);
if child.has_device_value(SWA) { if child.has_device_value(SWA) {
Node::redistribute_child_device_value(new_parent, child, SWA, split_len); Node::redistribute_child_device_value(new_parent, child, SWA, split_len);
} }
if child.has_host_value(SWA) { if child.has_host_value(SWA) {
Node::redistribute_child_host_value(new_parent, child, SWA, split_len); Node::redistribute_child_host_value(new_parent, child, SWA, split_len);
// Device-tombstoned sides park in the host LRU. // Device-tombstoned sides park in the host LRU. Host-locked
let parent_is_tombstone = !new_parent.has_device_value(SWA); // halves stay out of it: in-flight IO holds them, and host
let child_is_tombstone = !child.has_device_value(SWA); // acquire removed the node at 0->1.
let parent_parks =
!new_parent.has_device_value(SWA) && new_parent.host_lock_ref(SWA) == 0;
let child_parks = !child.has_device_value(SWA) && child.host_lock_ref(SWA) == 0;
let host_lru = tree_core.host_lru_list_mut(SWA); let host_lru = tree_core.host_lru_list_mut(SWA);
if parent_is_tombstone { if parent_parks {
host_lru.insert_mru(new_parent_id); host_lru.insert_mru(new_parent_id);
} }
if child_is_tombstone && !host_lru.in_list(Some(child_id)) { if child_parks && !host_lru.in_list(Some(child_id)) {
host_lru.insert_mru(child_id); host_lru.insert_mru(child_id);
} }
} }
// 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.
let swa_uuid = tree_core.arena.node_mut(child_id).swa_uuid.take(); let swa_uuid = tree_core.arena.node_mut(child_id).swa_uuid.take();
tree_core.arena.node_mut(new_parent_id).swa_uuid = swa_uuid; tree_core.arena.node_mut(new_parent_id).swa_uuid = swa_uuid;
let swa_host_uuid = tree_core.arena.node_mut(child_id).swa_host_uuid.take();
tree_core.arena.node_mut(new_parent_id).swa_host_uuid = swa_host_uuid;
} }
fn evict_component( fn evict_component(
@@ -1033,46 +1038,44 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
mut result: IncLockRefResult, mut result: IncLockRefResult,
lock_host: bool, lock_host: bool,
) -> IncLockRefResult { ) -> IncLockRefResult {
let ct = SWA; // 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 (key length); 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.
let sliding_window_size = self.sliding_window_size; let sliding_window_size = self.sliding_window_size;
let mut swa_lock_size = 0; let mut covered = 0;
let mut swa_uuid = None; let mut swa_uuid = None;
// 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.
let mut cur = node_id; let mut cur = node_id;
loop { loop {
let node = tree_core.arena.node_mut(cur); let node = tree_core.arena.node_mut(cur);
if node.is_root() || swa_lock_size >= sliding_window_size { if node.is_root() || covered >= sliding_window_size {
break; break;
} }
let parent = node.parent(); let parent = node.parent();
if !Self::has_value(node, lock_host) {
result
.skip_lock_node_ids
.entry(ct)
.or_default()
.insert(node.id);
cur = parent;
continue;
}
let key_len = node.key.atom_len(); let key_len = node.key.atom_len();
let has_value = Self::has_value(node, lock_host);
let value_len = Self::value_len(node, lock_host);
let newly_locked = Self::lock_ref(node, lock_host) == 0; let newly_locked = Self::lock_ref(node, lock_host) == 0;
Self::inc_lock_ref(node, lock_host); Self::inc_lock_ref(node, lock_host);
swa_lock_size += Self::value_len(node, lock_host); if newly_locked && has_value {
if newly_locked {
if lock_host { if lock_host {
let host_lru = tree_core.host_lru_list_mut(SWA); let host_lru = tree_core.host_lru_list_mut(SWA);
if host_lru.in_list(Some(cur)) { if host_lru.in_list(Some(cur)) {
host_lru.remove_node(cur); host_lru.remove_node(cur);
} }
} else { } else {
tree_core.dec_evictable_size(SWA, key_len); tree_core.dec_evictable_size(SWA, value_len);
tree_core.inc_protected_size(SWA, key_len); tree_core.inc_protected_size(SWA, value_len);
} }
} }
if swa_lock_size >= sliding_window_size { covered += key_len;
if covered >= sliding_window_size {
swa_uuid = Some(Self::ensure_swa_uuid(tree_core, cur, lock_host)); swa_uuid = Some(Self::ensure_swa_uuid(tree_core, cur, lock_host));
} }
cur = parent; cur = parent;
@@ -1090,23 +1093,15 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
&self, &self,
tree_core: &mut UnifiedTreeCore<K>, tree_core: &mut UnifiedTreeCore<K>,
node_id: NodeIdx_, node_id: NodeIdx_,
params: Option<&DecLockRefParams>, params: &DecLockRefParams,
lock_host: bool, lock_host: bool,
) { ) {
let ct = SWA; let swa_uuid_for_lock = if lock_host {
let swa_uuid_for_lock = params.and_then(|p| { params.swa_uuid_for_host_lock
if lock_host {
p.swa_uuid_for_host_lock
} else { } else {
p.swa_uuid_for_lock params.swa_uuid_for_lock
} };
});
let empty = HashSet::new();
let skip_lock_node_ids = params
.and_then(|p| p.skip_lock_node_ids.get(&ct))
.unwrap_or(&empty);
// A node in skip_lock_node_ids was a tombstone when this lock was acquired.
let mut cur = node_id; let mut cur = node_id;
loop { loop {
let node = tree_core.arena.node_mut(cur); let node = tree_core.arena.node_mut(cur);
@@ -1114,30 +1109,36 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
break; break;
} }
let parent = node.parent(); let parent = node.parent();
if skip_lock_node_ids.contains(&node.id) {
cur = parent;
continue;
}
let lock_ref = Self::lock_ref(node, lock_host); let lock_ref = Self::lock_ref(node, lock_host);
if lock_ref == 0 { // Acquire counted every segment node and splits copy refs, so a
cur = parent; // zero here means the release does not mirror its acquire.
continue; assert!(
} lock_ref > 0,
if lock_ref == 1 { "SWA segment release hit {}lock_ref=0 on node {cur}",
if lock_host { "host_" } else { "" }
);
let has_value = Self::has_value(node, lock_host);
let value_len = Self::value_len(node, lock_host);
if lock_ref == 1 && has_value {
if lock_host { if lock_host {
if !node.has_device_value(SWA) && node.has_host_value(SWA) { if !node.has_device_value(SWA) {
let host_lru = tree_core.host_lru_list_mut(SWA); let host_lru = tree_core.host_lru_list_mut(SWA);
if !host_lru.in_list(Some(cur)) { if !host_lru.in_list(Some(cur)) {
host_lru.insert_mru(cur); host_lru.insert_mru(cur);
} }
} }
} else { } else {
let key_len = node.device_value_len(SWA); tree_core.inc_evictable_size(SWA, value_len);
tree_core.inc_evictable_size(SWA, key_len); tree_core.dec_protected_size(SWA, value_len);
tree_core.dec_protected_size(SWA, key_len);
} }
} }
Self::dec_lock_ref(tree_core.arena.node_mut(cur), lock_host); Self::dec_lock_ref(tree_core.arena.node_mut(cur), lock_host);
if lock_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.
tree_core.update_evictable_leaf_sets_(cur);
}
if swa_uuid_for_lock.is_some() if swa_uuid_for_lock.is_some()
&& Self::swa_uuid(tree_core.arena.node(cur), lock_host) == swa_uuid_for_lock && Self::swa_uuid(tree_core.arena.node(cur), lock_host) == swa_uuid_for_lock
{ {
@@ -1147,15 +1148,14 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
} }
} }
/// Early-release the SWA lock along [node, swa_uuid_for_lock] while /// Early-release the SWA lock along [node, swa_uuid_for_lock]; this
/// leaving Full and Mamba locks intact. /// 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 /// 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 /// window. The caller must invoke this at most once per
/// Full lock must stay so the request's prefix is protected. /// (node, swa_uuid_for_lock) pair.
///
/// Caller (UnifiedRadixCache.dec_swa_lock_only) must ensure this is
/// invoked at most once per (node, swa_uuid_for_lock) pair.
fn release_window_lock( fn release_window_lock(
&self, &self,
tree_core: &mut UnifiedTreeCore<K>, tree_core: &mut UnifiedTreeCore<K>,
@@ -1172,21 +1172,20 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
break; break;
} }
let parent = node.parent(); let parent = node.parent();
// Acquire skips tombstoned nodes; release must skip them too. Same assert!(
// for nodes with lock_ref == 0 — acquire never credited them. node.device_lock_ref(SWA) > 0,
if !node.has_device_value(SWA) || node.device_lock_ref(SWA) == 0 { "SWA window release hit lock_ref=0 on node {cur}"
if swa_uuid_for_lock.is_some() && node.swa_uuid == swa_uuid_for_lock { );
break; let has_value = node.has_device_value(SWA);
} let value_len = node.device_value_len(SWA);
cur = parent;
continue;
}
node.dec_device_lock_ref(SWA); node.dec_device_lock_ref(SWA);
if node.device_lock_ref(SWA) == 0 { let now_unlocked = node.device_lock_ref(SWA) == 0;
let key_len = node.key.atom_len(); if now_unlocked {
tree_core.dec_protected_size(SWA, key_len); tree_core.update_evictable_leaf_sets_(cur);
tree_core.inc_evictable_size(SWA, key_len); }
if now_unlocked && has_value {
tree_core.dec_protected_size(SWA, value_len);
tree_core.inc_evictable_size(SWA, value_len);
if tree_core.is_evictable_device_leaf_(tree_core.arena.node(cur)) { if tree_core.is_evictable_device_leaf_(tree_core.arena.node(cur)) {
tree_core.evict_component_and_detach_lru_( tree_core.evict_component_and_detach_lru_(
cur, cur,
rust/sglang-radix-tree/src/node.rs
+6
View File
@@ -283,6 +283,12 @@ impl<K: ChildKeyType> Node<K> {
self.set_lock_ref_(slot, src_node.lock_ref_(slot)); self.set_lock_ref_(slot, src_node.lock_ref_(slot));
} }
/// Copy the component's host lock refcount from `src_node`.
pub fn copy_host_lock_ref(&mut self, component_type: ComponentType, src_node: &Node<K>) {
let slot = ValueSlotIdx::host(component_type);
self.set_lock_ref_(slot, src_node.lock_ref_(slot));
}
/// Split the component's device value between a new parent and the child. /// Split the component's device value between a new parent and the child.
pub fn redistribute_child_device_value( pub fn redistribute_child_device_value(
parent_node: &mut Node<K>, parent_node: &mut Node<K>,
rust/sglang-radix-tree/src/python_bindings.rs
+55 -73
View File
@@ -1,7 +1,7 @@
//! Python bindings: the `mem_cache` extension module and its TreeCore adapter. //! Python bindings: the `mem_cache` extension module and its TreeCore adapter.
use std::borrow::Cow; use std::borrow::Cow;
use std::collections::{HashMap, HashSet}; use std::collections::HashMap;
use std::sync::Mutex; use std::sync::Mutex;
use pyo3::buffer::PyBuffer; use pyo3::buffer::PyBuffer;
@@ -10,7 +10,7 @@ use pyo3::prelude::*;
use pyo3::types::{PyBytes, PyDict, PyList}; use pyo3::types::{PyBytes, PyDict, PyList};
use tch::{Device, Kind, Tensor}; use tch::{Device, Kind, Tensor};
use crate::components::{ComponentType, FULL, MAMBA, SWA}; use crate::components::{ComponentSet, ComponentType, FULL, MAMBA, SWA};
use crate::node::ChildKeyType; use crate::node::ChildKeyType;
use crate::node::{KeyNamespaceRef, NodeAccessError, NodeId, TreeCoreRuntimeError}; use crate::node::{KeyNamespaceRef, NodeAccessError, NodeId, TreeCoreRuntimeError};
use crate::unified_tree_core::KvCacheEvent; use crate::unified_tree_core::KvCacheEvent;
@@ -648,24 +648,27 @@ impl InsertResultBinding {
#[pyclass(get_all, set_all)] #[pyclass(get_all, set_all)]
#[derive(Clone, Default)] #[derive(Clone, Default)]
pub struct DecLockRefParamsBinding { pub struct DecLockRefParamsBinding {
pub node_id: Option<NodeId>,
pub swa_uuid_for_lock: Option<i64>, pub swa_uuid_for_lock: Option<i64>,
pub swa_uuid_for_host_lock: Option<i64>, pub swa_uuid_for_host_lock: Option<i64>,
pub skip_lock_node_ids: HashMap<u8, HashSet<NodeId>>, pub skipped_lock_components: Vec<u8>,
} }
#[pymethods] #[pymethods]
impl DecLockRefParamsBinding { impl DecLockRefParamsBinding {
#[new] #[new]
#[pyo3(signature = (swa_uuid_for_lock = None, swa_uuid_for_host_lock = None, skip_lock_node_ids = None))] #[pyo3(signature = (node_id = None, swa_uuid_for_lock = None, swa_uuid_for_host_lock = None, skipped_lock_components = Vec::new()))]
fn new( fn new(
node_id: Option<NodeId>,
swa_uuid_for_lock: Option<i64>, swa_uuid_for_lock: Option<i64>,
swa_uuid_for_host_lock: Option<i64>, swa_uuid_for_host_lock: Option<i64>,
skip_lock_node_ids: Option<HashMap<u8, HashSet<NodeId>>>, skipped_lock_components: Vec<u8>,
) -> Self { ) -> Self {
DecLockRefParamsBinding { DecLockRefParamsBinding {
node_id,
swa_uuid_for_lock, swa_uuid_for_lock,
swa_uuid_for_host_lock, swa_uuid_for_host_lock,
skip_lock_node_ids: skip_lock_node_ids.unwrap_or_default(), skipped_lock_components,
} }
} }
} }
@@ -674,44 +677,49 @@ impl DecLockRefParamsBinding {
/// Convert into the tree core's dec-lock params. /// Convert into the tree core's dec-lock params.
fn to_dec_lock_ref_params(&self) -> PyResult<DecLockRefParams> { fn to_dec_lock_ref_params(&self) -> PyResult<DecLockRefParams> {
Ok(DecLockRefParams { Ok(DecLockRefParams {
node_id: self.node_id,
swa_uuid_for_lock: self.swa_uuid_for_lock, swa_uuid_for_lock: self.swa_uuid_for_lock,
swa_uuid_for_host_lock: self.swa_uuid_for_host_lock, swa_uuid_for_host_lock: self.swa_uuid_for_host_lock,
skip_lock_node_ids: self skipped_lock_components: component_set_from_py(&self.skipped_lock_components)?,
.skip_lock_node_ids
.iter()
.map(|(ct, node_ids)| {
Ok::<_, PyErr>((parse_component_type(*ct)?, node_ids.clone()))
})
.collect::<PyResult<_>>()?,
}) })
} }
} }
/// Python-visible inc_lock_ref result; hand skip_lock_node_ids back to the /// Python-visible inc_lock_ref result; the receipt (anchor node, boundary
/// matching dec_lock_ref. /// uuids, skipped components) is handed back to the matching dec_lock_ref.
#[pyclass(get_all)] #[pyclass(get_all)]
pub struct IncLockRefResultBinding { pub struct IncLockRefResultBinding {
delta: Option<usize>, delta: Option<usize>,
node_id: Option<NodeId>,
swa_uuid_for_lock: Option<i64>, swa_uuid_for_lock: Option<i64>,
swa_uuid_for_host_lock: Option<i64>, swa_uuid_for_host_lock: Option<i64>,
skip_lock_node_ids: HashMap<u8, HashSet<NodeId>>, skipped_lock_components: Vec<u8>,
} }
impl IncLockRefResultBinding { impl IncLockRefResultBinding {
fn from_result(result: crate::unified_tree_core::IncLockRefResult) -> Self { fn from_result(result: crate::unified_tree_core::IncLockRefResult) -> Self {
Self { Self {
delta: result.delta, delta: result.delta,
node_id: result.node_id,
swa_uuid_for_lock: result.swa_uuid_for_lock, swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result skipped_lock_components: result
.skip_lock_node_ids .skipped_lock_components
.into_iter() .iter()
.map(|(ct, node_ids)| (component_type_to_u8(ct), node_ids)) .map(|ct| ct.idx() as u8)
.collect(), .collect(),
} }
} }
} }
/// Parse Python component-type ids into a component set.
fn component_set_from_py(component_types: &[u8]) -> PyResult<ComponentSet> {
component_types
.iter()
.map(|ct| parse_component_type(*ct))
.collect()
}
/// Convert a Python component-keyed tracker into the core's counts. /// Convert a Python component-keyed tracker into the core's counts.
fn tracker_from_py(tracker: HashMap<u8, usize>) -> PyResult<HashMap<ComponentType, usize>> { fn tracker_from_py(tracker: HashMap<u8, usize>) -> PyResult<HashMap<ComponentType, usize>> {
tracker tracker
@@ -1074,23 +1082,17 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
cache_actions_to_py(py, actions) cache_actions_to_py(py, actions)
} }
/// Bump the reference count on a node's component locks. /// Bump the reference count on a node's component locks; the listed
/// components are left untaken and recorded in the receipt.
fn inc_lock_ref( fn inc_lock_ref(
&self, &self,
py: Python<'_>, py: Python<'_>,
node_id: NodeId, node_id: NodeId,
skip_lock_components: Option<Vec<u8>>, skip_lock_components: Vec<u8>,
) -> PyResult<IncLockRefResultBinding> { ) -> PyResult<IncLockRefResultBinding> {
let skip_lock_components = skip_lock_components let skip = component_set_from_py(&skip_lock_components)?;
.unwrap_or_default()
.into_iter()
.map(parse_component_type)
.collect::<PyResult<Vec<_>>>()?;
let result = py let result = py
.allow_threads(|| { .allow_threads(|| self.core().inc_lock_ref(node_id, skip))
self.core()
.inc_lock_ref_with_skip(node_id, &skip_lock_components)
})
.map_err(node_access_error)?; .map_err(node_access_error)?;
Ok(IncLockRefResultBinding::from_result(result)) Ok(IncLockRefResultBinding::from_result(result))
} }
@@ -1100,11 +1102,11 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
&self, &self,
py: Python<'_>, py: Python<'_>,
node_id: NodeId, node_id: NodeId,
params: Option<&DecLockRefParamsBinding>, params: &DecLockRefParamsBinding,
skip_swa: bool, skip_swa: bool,
) -> PyResult<()> { ) -> PyResult<()> {
let params = params.map(|p| p.to_dec_lock_ref_params()).transpose()?; let params = params.to_dec_lock_ref_params()?;
py.allow_threads(|| self.core().dec_lock_ref(node_id, params.as_ref(), skip_swa)) py.allow_threads(|| self.core().dec_lock_ref(node_id, &params, skip_swa))
.map_err(node_access_error)?; .map_err(node_access_error)?;
Ok(()) Ok(())
} }
@@ -1115,22 +1117,16 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
&self, &self,
py: Python<'_>, py: Python<'_>,
node_id: NodeId, node_id: NodeId,
swa_uuid_for_lock: Option<i64>, params: &DecLockRefParamsBinding,
skip_lock_node_ids: Option<HashMap<u8, HashSet<NodeId>>>,
) -> PyResult<(Py<PyDict>, Py<PyDict>)> { ) -> PyResult<(Py<PyDict>, Py<PyDict>)> {
let skip_lock_node_ids = skip_lock_node_ids let params = params.to_dec_lock_ref_params()?;
.unwrap_or_default()
.into_iter()
.map(|(ct, node_ids)| Ok((parse_component_type(ct)?, node_ids)))
.collect::<PyResult<HashMap<_, _>>>()?;
let (device_frees, host_frees) = py let (device_frees, host_frees) = py
.allow_threads(|| { .allow_threads(|| {
let mut device_frees = HashMap::new(); let mut device_frees = HashMap::new();
let mut host_frees = HashMap::new(); let mut host_frees = HashMap::new();
self.core().dec_swa_lock_only_with_skip( self.core().dec_swa_lock_only(
node_id, node_id,
swa_uuid_for_lock, &params,
Some(&skip_lock_node_ids),
&mut device_frees, &mut device_frees,
&mut host_frees, &mut host_frees,
)?; )?;
@@ -1733,16 +1729,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
let result = py let result = py
.allow_threads(|| self.core().inc_host_lock_ref(node_id)) .allow_threads(|| self.core().inc_host_lock_ref(node_id))
.map_err(node_access_error)?; .map_err(node_access_error)?;
Ok(IncLockRefResultBinding { Ok(IncLockRefResultBinding::from_result(result))
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: result
.skip_lock_node_ids
.into_iter()
.map(|(ct, node_ids)| (component_type_to_u8(ct), node_ids))
.collect(),
})
} }
/// Decrease the reference count on a node's host-side component locks. /// Decrease the reference count on a node's host-side component locks.
@@ -1750,10 +1737,10 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
&self, &self,
py: Python<'_>, py: Python<'_>,
node_id: NodeId, node_id: NodeId,
params: Option<&DecLockRefParamsBinding>, params: &DecLockRefParamsBinding,
) -> PyResult<()> { ) -> PyResult<()> {
let params = params.map(|p| p.to_dec_lock_ref_params()).transpose()?; let params = params.to_dec_lock_ref_params()?;
py.allow_threads(|| self.core().dec_host_lock_ref(node_id, params.as_ref())) py.allow_threads(|| self.core().dec_host_lock_ref(node_id, &params))
.map_err(node_access_error)?; .map_err(node_access_error)?;
Ok(()) Ok(())
} }
@@ -2371,23 +2358,24 @@ macro_rules! tree_core_binding {
} }
/// Bump the reference count on a node's component locks. /// Bump the reference count on a node's component locks.
#[pyo3(signature = (node_id, skip_lock_components = None))] #[pyo3(signature = (node_id, skip_lock_components = Vec::new()))]
fn inc_lock_ref( fn inc_lock_ref(
&self, &self,
py: Python<'_>, py: Python<'_>,
node_id: NodeId, node_id: NodeId,
skip_lock_components: Option<Vec<u8>>, skip_lock_components: Vec<u8>,
) -> PyResult<IncLockRefResultBinding> { ) -> PyResult<IncLockRefResultBinding> {
self.inner.inc_lock_ref(py, node_id, skip_lock_components) self.inner.inc_lock_ref(py, node_id, skip_lock_components)
} }
/// Decrease the reference count on a node's component locks. /// Decrease the reference count on a node's component locks. The
#[pyo3(signature = (node_id, params = None, skip_swa = false))] /// receipt is required: a release must replay its acquire's evidence.
#[pyo3(signature = (node_id, params, skip_swa = false))]
fn dec_lock_ref( fn dec_lock_ref(
&self, &self,
py: Python<'_>, py: Python<'_>,
node_id: NodeId, node_id: NodeId,
params: Option<&DecLockRefParamsBinding>, params: &DecLockRefParamsBinding,
skip_swa: bool, skip_swa: bool,
) -> PyResult<()> { ) -> PyResult<()> {
self.inner.dec_lock_ref(py, node_id, params, skip_swa) self.inner.dec_lock_ref(py, node_id, params, skip_swa)
@@ -2395,20 +2383,14 @@ macro_rules! tree_core_binding {
/// Early-release the SWA portion of a request's tree lock; returns this /// Early-release the SWA portion of a request's tree lock; returns this
/// release's per-component (device_frees, host_frees). /// release's per-component (device_frees, host_frees).
#[pyo3(signature = (node_id, swa_uuid_for_lock = None, skip_lock_node_ids = None))] #[pyo3(signature = (node_id, params))]
fn dec_swa_lock_only( fn dec_swa_lock_only(
&self, &self,
py: Python<'_>, py: Python<'_>,
node_id: NodeId, node_id: NodeId,
swa_uuid_for_lock: Option<i64>, params: &DecLockRefParamsBinding,
skip_lock_node_ids: Option<HashMap<u8, HashSet<NodeId>>>,
) -> PyResult<(Py<PyDict>, Py<PyDict>)> { ) -> PyResult<(Py<PyDict>, Py<PyDict>)> {
self.inner.dec_swa_lock_only( self.inner.dec_swa_lock_only(py, node_id, params)
py,
node_id,
swa_uuid_for_lock,
skip_lock_node_ids,
)
} }
/// Store a component's device value on a node (the SWA rebuild write-back). /// Store a component's device value on a node (the SWA rebuild write-back).
@@ -2817,12 +2799,12 @@ macro_rules! tree_core_binding {
} }
/// Decrease the reference count on a node's host-side component locks. /// Decrease the reference count on a node's host-side component locks.
#[pyo3(signature = (node_id, params = None))] /// The receipt is required, as for dec_lock_ref.
fn dec_host_lock_ref( fn dec_host_lock_ref(
&self, &self,
py: Python<'_>, py: Python<'_>,
node_id: NodeId, node_id: NodeId,
params: Option<&DecLockRefParamsBinding>, params: &DecLockRefParamsBinding,
) -> PyResult<()> { ) -> PyResult<()> {
self.inner.dec_host_lock_ref(py, node_id, params) self.inner.dec_host_lock_ref(py, node_id, params)
} }
rust/sglang-radix-tree/src/tests/components/base.rs
+1 -1
View File
@@ -70,7 +70,7 @@ impl TreeComponent<Vec<i64>> for DefaultComponentForTest {
&self, &self,
tree_core: &mut UnifiedTreeCore<Vec<i64>>, tree_core: &mut UnifiedTreeCore<Vec<i64>>,
node_id: NodeIdx_, node_id: NodeIdx_,
params: Option<&DecLockRefParams>, params: &DecLockRefParams,
lock_host: bool, lock_host: bool,
) { ) {
unimplemented!() unimplemented!()
rust/sglang-radix-tree/src/tests/components/full.rs
+119 -79
View File
@@ -1,5 +1,5 @@
use super::*; use super::*;
use crate::components::FULL; use crate::components::{ComponentSet, FULL};
use crate::node::NodeAccessError; use crate::node::NodeAccessError;
use crate::test_utils::accumulate_step; use crate::test_utils::accumulate_step;
use crate::unified_tree_core::CacheInitParams; use crate::unified_tree_core::CacheInitParams;
@@ -625,10 +625,9 @@ fn inc_lock_ref_locks_the_device_path() {
let mut tc = core(); let mut tc = core();
let (n1, n2) = lock_chain(&mut tc); let (n1, n2) = lock_chain(&mut tc);
let result = tc let result = tc
.inc_lock_ref(tc.arena.node(n2).id) .inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
assert_eq!(result.delta, Some(5)); assert_eq!(result.delta, Some(5));
assert!(result.skip_lock_node_ids.is_empty());
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1); assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(n2, FULL), 1); assert_eq!(tc.arena.device_lock_ref(n2, FULL), 1);
let state = tc.component_state(FULL); let state = tc.component_state(FULL);
@@ -641,10 +640,10 @@ fn inc_lock_ref_locks_the_device_path() {
fn inc_lock_ref_again_only_bumps_the_refs() { fn inc_lock_ref_again_only_bumps_the_refs() {
let mut tc = core(); let mut tc = core();
let (n1, n2) = lock_chain(&mut tc); let (n1, n2) = lock_chain(&mut tc);
tc.inc_lock_ref(tc.arena.node(n2).id) tc.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
let result = tc let result = tc
.inc_lock_ref(tc.arena.node(n2).id) .inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
assert_eq!(result.delta, Some(0)); assert_eq!(result.delta, Some(0));
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 2); assert_eq!(tc.arena.device_lock_ref(n1, FULL), 2);
@@ -659,10 +658,10 @@ fn inc_lock_ref_counts_only_newly_locked_nodes() {
// n1 is already locked via its own path; locking n2 moves only n2's tokens. // n1 is already locked via its own path; locking n2 moves only n2's tokens.
let mut tc = core(); let mut tc = core();
let (n1, n2) = lock_chain(&mut tc); let (n1, n2) = lock_chain(&mut tc);
tc.inc_lock_ref(tc.arena.node(n1).id) tc.inc_lock_ref(tc.arena.node(n1).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
let result = tc let result = tc
.inc_lock_ref(tc.arena.node(n2).id) .inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
assert_eq!(result.delta, Some(3)); assert_eq!(result.delta, Some(3));
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 2); assert_eq!(tc.arena.device_lock_ref(n1, FULL), 2);
@@ -673,8 +672,9 @@ fn inc_lock_ref_counts_only_newly_locked_nodes() {
} }
#[test] #[test]
fn inc_lock_ref_collects_the_evicted_bottom_segment() { fn inc_lock_ref_counts_the_evicted_bottom_segment() {
// n2 and n3 are evicted (no device value): the walk records both and locks only n1. // n2 and n3 are evicted (no device value): counted in the segment with no
// ledger move; only n1's tokens turn protected.
let mut tc = core(); let mut tc = core();
let root = tc.arena.root(); let root = tc.arena.root();
let n1 = tc let n1 = tc
@@ -709,16 +709,12 @@ fn inc_lock_ref_collects_the_evicted_bottom_segment() {
tc.component_state_mut(FULL).evictable_size = 2; tc.component_state_mut(FULL).evictable_size = 2;
tc.evictable_device_leaves.add(n1); tc.evictable_device_leaves.add(n1);
let result = tc let result = tc
.inc_lock_ref(tc.arena.node(n3).id) .inc_lock_ref(tc.arena.node(n3).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
assert_eq!(result.delta, Some(2)); assert_eq!(result.delta, Some(2));
assert_eq!(
result.skip_lock_node_ids[&FULL],
HashSet::from([tc.arena.node(n2).id, tc.arena.node(n3).id])
);
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1); assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(n2, FULL), 0); assert_eq!(tc.arena.device_lock_ref(n2, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(n3, FULL), 0); assert_eq!(tc.arena.device_lock_ref(n3, FULL), 1);
// The locked ancestor leaves the D-leaf set. // The locked ancestor leaves the D-leaf set.
assert!(!tc.evictable_device_leaves.contains(n1)); assert!(!tc.evictable_device_leaves.contains(n1));
} }
@@ -728,15 +724,18 @@ fn lock_round_trips_on_a_root_anchor_are_noops() {
let mut tc = core(); let mut tc = core();
let root = tc.arena.root(); let root = tc.arena.root();
let result = tc let result = tc
.inc_lock_ref(tc.arena.node(root).id) .inc_lock_ref(tc.arena.node(root).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
assert_eq!(result.delta, Some(0)); assert_eq!(result.delta, Some(0));
assert!(result.skip_lock_node_ids.is_empty());
// The protected root keeps its construction-time lock through the pair. // The protected root keeps its construction-time lock through the pair.
assert_eq!(tc.arena.device_lock_ref(root, FULL), 1); assert_eq!(tc.arena.device_lock_ref(root, FULL), 1);
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena.node(root).id, tc.arena.node(root).id,
/* params = */ None, /* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false, /* skip_swa = */ false,
) )
.expect("live test node"); .expect("live test node");
@@ -760,7 +759,7 @@ fn lock_walks_stop_at_the_root_of_a_salted_chain() {
.set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1])); .set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1]));
tc.component_state_mut(FULL).evictable_size = 2; tc.component_state_mut(FULL).evictable_size = 2;
let result = tc let result = tc
.inc_lock_ref(tc.arena.node(n1).id) .inc_lock_ref(tc.arena.node(n1).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
assert_eq!(result.delta, Some(2)); assert_eq!(result.delta, Some(2));
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1); assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1);
@@ -769,7 +768,11 @@ fn lock_walks_stop_at_the_root_of_a_salted_chain() {
// The release walk stops at the same boundary. // The release walk stops at the same boundary.
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena.node(n1).id, tc.arena.node(n1).id,
/* params = */ None, /* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false, /* skip_swa = */ false,
) )
.expect("live test node"); .expect("live test node");
@@ -798,17 +801,20 @@ fn lock_walks_treat_a_present_but_empty_value_as_device_on() {
Tensor::from_slice(&empty), Tensor::from_slice(&empty),
); );
let result = tc let result = tc
.inc_lock_ref(tc.arena.node(n1).id) .inc_lock_ref(tc.arena.node(n1).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
// A present-but-empty value is device-on (Python `value is not None`): // A present-but-empty value is device-on (Python `value is not None`):
// locked, zero tokens moved. // locked, zero tokens moved.
assert_eq!(result.delta, Some(0)); assert_eq!(result.delta, Some(0));
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1); assert_eq!(tc.arena.device_lock_ref(n1, FULL), 1);
assert!(result.skip_lock_node_ids.is_empty());
// The release side moves the same zero tokens back. // The release side moves the same zero tokens back.
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena.node(n1).id, tc.arena.node(n1).id,
/* params = */ None, /* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false, /* skip_swa = */ false,
) )
.expect("live test node"); .expect("live test node");
@@ -822,11 +828,15 @@ fn lock_walks_treat_a_present_but_empty_value_as_device_on() {
fn dec_lock_ref_unlocks_and_restores_sizes() { fn dec_lock_ref_unlocks_and_restores_sizes() {
let mut tc = core(); let mut tc = core();
let (n1, n2) = lock_chain(&mut tc); let (n1, n2) = lock_chain(&mut tc);
tc.inc_lock_ref(tc.arena.node(n2).id) tc.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena.node(n2).id, tc.arena.node(n2).id,
/* params = */ None, /* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false, /* skip_swa = */ false,
) )
.expect("live test node"); .expect("live test node");
@@ -875,18 +885,14 @@ fn dec_lock_ref_replays_the_skip_set() {
.set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1])); .set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1]));
tc.component_state_mut(FULL).evictable_size = 2; tc.component_state_mut(FULL).evictable_size = 2;
let result = tc let result = tc
.inc_lock_ref(tc.arena.node(n3).id) .inc_lock_ref(tc.arena.node(n3).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
let params = DecLockRefParams { let params = DecLockRefParams {
skip_lock_node_ids: result.skip_lock_node_ids, skipped_lock_components: result.skipped_lock_components,
..Default::default() ..Default::default()
}; };
// The still-evicted n2 and n3 are skipped instead of tripping the lock asserts. // The still-evicted n2 and n3 are skipped instead of tripping the lock asserts.
tc.dec_lock_ref( tc.dec_lock_ref(tc.arena.node(n3).id, &params, /* skip_swa = */ false)
tc.arena.node(n3).id,
Some(&params),
/* skip_swa = */ false,
)
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.device_lock_ref(n1, FULL), 0); assert_eq!(tc.arena.device_lock_ref(n1, FULL), 0);
assert_eq!(tc.arena.device_lock_ref(n2, FULL), 0); assert_eq!(tc.arena.device_lock_ref(n2, FULL), 0);
@@ -897,7 +903,7 @@ fn dec_lock_ref_replays_the_skip_set() {
} }
#[test] #[test]
fn temp_lock_skips_the_evicted_anchor_and_mirrors_on_release() { fn temp_lock_counts_the_evicted_anchor_and_mirrors_on_release() {
// Chain root -> a -> y -> anchor with FULL device values; the anchor is evicted. // Chain root -> a -> y -> anchor with FULL device values; the anchor is evicted.
let mut tc = core(); let mut tc = core();
let root = tc.arena.root(); let root = tc.arena.root();
@@ -933,34 +939,32 @@ fn temp_lock_skips_the_evicted_anchor_and_mirrors_on_release() {
tc.arena tc.arena
.set_device_value(y, FULL, Tensor::from_slice(&[0i64])); .set_device_value(y, FULL, Tensor::from_slice(&[0i64]));
tc.component_state_mut(FULL).evictable_size = 3; tc.component_state_mut(FULL).evictable_size = 3;
// The temp lock records the evicted anchor and locks only its ancestors. // The temp lock counts the evicted anchor too (no ledger move).
let temp_lock = tc let temp_lock = tc
.inc_lock_ref(tc.arena.node(anchor).id) .inc_lock_ref(tc.arena.node(anchor).id, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(anchor, FULL), 0);
assert_eq!(tc.arena.device_lock_ref(y, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(a, FULL), 1);
assert_eq!(
temp_lock.skip_lock_node_ids[&FULL],
HashSet::from([tc.arena.node(anchor).id])
);
// A load-back restores the anchor; the second acquire covers it.
tc.arena
.set_device_value(anchor, FULL, Tensor::from_slice(&[0i64]));
let second_lock = tc
.inc_lock_ref(tc.arena.node(anchor).id)
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.device_lock_ref(anchor, FULL), 1); assert_eq!(tc.arena.device_lock_ref(anchor, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(y, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(a, FULL), 1);
// A load-back restores the anchor under the held lock: credited to
// protected, as the production commit does. The second acquire stacks.
tc.arena
.set_device_value(anchor, FULL, Tensor::from_slice(&[0i64]));
tc.inc_protected_size(FULL, 1);
let second_lock = tc
.inc_lock_ref(tc.arena.node(anchor).id, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(anchor, FULL), 2);
assert_eq!(tc.arena.device_lock_ref(y, FULL), 2); assert_eq!(tc.arena.device_lock_ref(y, FULL), 2);
assert_eq!(tc.arena.device_lock_ref(a, FULL), 2); assert_eq!(tc.arena.device_lock_ref(a, FULL), 2);
// Releasing the temp lock mirrors its skip set: the anchor keeps its lock. // Each release takes back exactly its own refs.
let temp_params = DecLockRefParams { let temp_params = DecLockRefParams {
skip_lock_node_ids: temp_lock.skip_lock_node_ids, skipped_lock_components: temp_lock.skipped_lock_components,
..Default::default() ..Default::default()
}; };
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena.node(anchor).id, tc.arena.node(anchor).id,
Some(&temp_params), &temp_params,
/* skip_swa = */ false, /* skip_swa = */ false,
) )
.expect("live test node"); .expect("live test node");
@@ -968,12 +972,12 @@ fn temp_lock_skips_the_evicted_anchor_and_mirrors_on_release() {
assert_eq!(tc.arena.device_lock_ref(y, FULL), 1); assert_eq!(tc.arena.device_lock_ref(y, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(a, FULL), 1); assert_eq!(tc.arena.device_lock_ref(a, FULL), 1);
let second_params = DecLockRefParams { let second_params = DecLockRefParams {
skip_lock_node_ids: second_lock.skip_lock_node_ids, skipped_lock_components: second_lock.skipped_lock_components,
..Default::default() ..Default::default()
}; };
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena.node(anchor).id, tc.arena.node(anchor).id,
Some(&second_params), &second_params,
/* skip_swa = */ false, /* skip_swa = */ false,
) )
.expect("live test node"); .expect("live test node");
@@ -983,9 +987,9 @@ fn temp_lock_skips_the_evicted_anchor_and_mirrors_on_release() {
} }
#[test] #[test]
#[should_panic(expected = "has no FULL device value")] #[should_panic(expected = "FULL segment release hit lock_ref=0")]
fn dec_lock_ref_panics_without_replaying_the_skip_set() { fn dec_lock_ref_panics_on_double_release() {
// Dropping the acquire's skip set makes the release walk hit the tombstone. // The second, unpaired release hits the already-unlocked segment.
let mut tc = core(); let mut tc = core();
let root = tc.arena.root(); let root = tc.arena.root();
let n1 = tc let n1 = tc
@@ -1009,11 +1013,25 @@ fn dec_lock_ref_panics_without_replaying_the_skip_set() {
tc.arena tc.arena
.set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1])); .set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1]));
tc.component_state_mut(FULL).evictable_size = 2; tc.component_state_mut(FULL).evictable_size = 2;
tc.inc_lock_ref(tc.arena.node(n2).id) tc.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena.node(n2).id, tc.arena.node(n2).id,
/* params = */ None, /* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
)
.expect("live test node");
tc.dec_lock_ref(
tc.arena.node(n2).id,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false, /* skip_swa = */ false,
) )
.expect("live test node"); .expect("live test node");
@@ -1023,11 +1041,15 @@ fn dec_lock_ref_panics_without_replaying_the_skip_set() {
fn dec_lock_ref_with_skip_swa_still_releases_full() { fn dec_lock_ref_with_skip_swa_still_releases_full() {
let mut tc = core(); let mut tc = core();
let (_n1, n2) = lock_chain(&mut tc); let (_n1, n2) = lock_chain(&mut tc);
tc.inc_lock_ref(tc.arena.node(n2).id) tc.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena.node(n2).id, tc.arena.node(n2).id,
/* params = */ None, /* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ true, /* skip_swa = */ true,
) )
.expect("live test node"); .expect("live test node");
@@ -1039,13 +1061,17 @@ fn nested_locks_release_pairwise() {
// Two acquires then two releases: sizes move only on the outermost pair. // Two acquires then two releases: sizes move only on the outermost pair.
let mut tc = core(); let mut tc = core();
let (_n1, n2) = lock_chain(&mut tc); let (_n1, n2) = lock_chain(&mut tc);
tc.inc_lock_ref(tc.arena.node(n2).id) tc.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
tc.inc_lock_ref(tc.arena.node(n2).id) tc.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena.node(n2).id, tc.arena.node(n2).id,
/* params = */ None, /* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false, /* skip_swa = */ false,
) )
.expect("live test node"); .expect("live test node");
@@ -1055,7 +1081,11 @@ fn nested_locks_release_pairwise() {
assert!(!tc.evictable_device_leaves.contains(n2)); assert!(!tc.evictable_device_leaves.contains(n2));
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena.node(n2).id, tc.arena.node(n2).id,
/* params = */ None, /* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false, /* skip_swa = */ false,
) )
.expect("live test node"); .expect("live test node");
@@ -1066,13 +1096,17 @@ fn nested_locks_release_pairwise() {
} }
#[test] #[test]
#[should_panic(expected = "is not locked")] #[should_panic(expected = "FULL segment release hit lock_ref=0")]
fn dec_lock_ref_panics_on_an_unlocked_node() { fn dec_lock_ref_panics_on_an_unlocked_node() {
let mut tc = core(); let mut tc = core();
let (_n1, n2) = lock_chain(&mut tc); let (_n1, n2) = lock_chain(&mut tc);
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena.node(n2).id, tc.arena.node(n2).id,
/* params = */ None, /* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false, /* skip_swa = */ false,
) )
.expect("live test node"); .expect("live test node");
@@ -1105,7 +1139,7 @@ fn inc_lock_ref_panics_on_an_evicted_ancestor() {
tc.arena tc.arena
.set_device_value(n2, FULL, Tensor::from_slice(&[0i64, 1, 2])); .set_device_value(n2, FULL, Tensor::from_slice(&[0i64, 1, 2]));
tc.component_state_mut(FULL).evictable_size = 3; tc.component_state_mut(FULL).evictable_size = 3;
tc.inc_lock_ref(tc.arena.node(n2).id) tc.inc_lock_ref(tc.arena.node(n2).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
} }
@@ -1125,7 +1159,7 @@ fn inc_lock_ref_panics_when_evictable_size_is_unaccounted() {
.unwrap(); .unwrap();
tc.arena tc.arena
.set_device_value(n1, FULL, Tensor::from_slice(&[0i64])); .set_device_value(n1, FULL, Tensor::from_slice(&[0i64]));
tc.inc_lock_ref(tc.arena.node(n1).id) tc.inc_lock_ref(tc.arena.node(n1).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
} }
@@ -1140,7 +1174,11 @@ fn dec_lock_ref_panics_on_protected_underflow() {
.set_lock_ref_(ValueSlotIdx::device(FULL), 1); .set_lock_ref_(ValueSlotIdx::device(FULL), 1);
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena.node(n2).id, tc.arena.node(n2).id,
/* params = */ None, /* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false, /* skip_swa = */ false,
) )
.expect("live test node"); .expect("live test node");
@@ -1183,7 +1221,6 @@ fn inc_host_lock_ref_pins_the_backuped_anchor() {
.inc_host_lock_ref(tc.arena.node(node).id) .inc_host_lock_ref(tc.arena.node(node).id)
.expect("live test node"); .expect("live test node");
assert_eq!(result.delta, None); assert_eq!(result.delta, None);
assert!(result.skip_lock_node_ids.is_empty());
assert_eq!(tc.arena.host_lock_ref(node, FULL), 1); assert_eq!(tc.arena.host_lock_ref(node, FULL), 1);
// The pinned anchor leaves the H-leaf set; the device tier is untouched. // The pinned anchor leaves the H-leaf set; the device tier is untouched.
assert!(!tc.evictable_host_leaves.contains(node)); assert!(!tc.evictable_host_leaves.contains(node));
@@ -1256,7 +1293,7 @@ fn host_lock_round_trips_on_a_root_anchor_are_noops() {
.expect("live test node"); .expect("live test node");
assert_eq!(result.delta, None); assert_eq!(result.delta, None);
assert_eq!(tc.arena.host_lock_ref(root, FULL), 0); assert_eq!(tc.arena.host_lock_ref(root, FULL), 0);
tc.dec_host_lock_ref(tc.arena.node(root).id, /* params = */ None) tc.dec_host_lock_ref(tc.arena.node(root).id, &DecLockRefParams::default())
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.host_lock_ref(root, FULL), 0); assert_eq!(tc.arena.host_lock_ref(root, FULL), 0);
} }
@@ -1281,7 +1318,7 @@ fn dec_host_lock_ref_unpins_and_restores_the_h_leaf_set() {
tc.component_state_mut(FULL).evictable_size = 7; tc.component_state_mut(FULL).evictable_size = 7;
tc.inc_host_lock_ref(tc.arena.node(node).id) tc.inc_host_lock_ref(tc.arena.node(node).id)
.expect("live test node"); .expect("live test node");
tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None) tc.dec_host_lock_ref(tc.arena.node(node).id, &DecLockRefParams::default())
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 0); assert_eq!(tc.arena.host_lock_ref(node, FULL), 0);
assert!(tc.evictable_host_leaves.contains(node)); assert!(tc.evictable_host_leaves.contains(node));
@@ -1294,7 +1331,7 @@ fn dec_host_lock_ref_unpins_and_restores_the_h_leaf_set() {
fn dec_host_lock_ref_on_an_unlocked_anchor_is_a_noop() { fn dec_host_lock_ref_on_an_unlocked_anchor_is_a_noop() {
let mut tc = core(); let mut tc = core();
let node = host_lock_anchor(&mut tc); let node = host_lock_anchor(&mut tc);
tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None) tc.dec_host_lock_ref(tc.arena.node(node).id, &DecLockRefParams::default())
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 0); assert_eq!(tc.arena.host_lock_ref(node, FULL), 0);
} }
@@ -1307,7 +1344,7 @@ fn dec_host_lock_ref_keeps_the_counter_when_the_host_value_is_gone() {
tc.inc_host_lock_ref(tc.arena.node(node).id) tc.inc_host_lock_ref(tc.arena.node(node).id)
.expect("live test node"); .expect("live test node");
let _ = tc.arena.take_host_value(node, FULL); let _ = tc.arena.take_host_value(node, FULL);
tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None) tc.dec_host_lock_ref(tc.arena.node(node).id, &DecLockRefParams::default())
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 1); assert_eq!(tc.arena.host_lock_ref(node, FULL), 1);
} }
@@ -1318,7 +1355,7 @@ fn host_lock_round_trip_under_write_back_is_a_pure_counter() {
let (_n1, n2) = lock_chain(&mut tc); let (_n1, n2) = lock_chain(&mut tc);
tc.inc_host_lock_ref(tc.arena.node(n2).id) tc.inc_host_lock_ref(tc.arena.node(n2).id)
.expect("live test node"); .expect("live test node");
tc.dec_host_lock_ref(tc.arena.node(n2).id, /* params = */ None) tc.dec_host_lock_ref(tc.arena.node(n2).id, &DecLockRefParams::default())
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.host_lock_ref(n2, FULL), 0); assert_eq!(tc.arena.host_lock_ref(n2, FULL), 0);
let state = tc.component_state(FULL); let state = tc.component_state(FULL);
@@ -1346,7 +1383,10 @@ fn release_host_arm_updates_the_h_leaf_set_without_the_dispatcher() {
tc.inc_host_lock_ref(tc.arena.node(node).id) tc.inc_host_lock_ref(tc.arena.node(node).id)
.expect("live test node"); .expect("live test node");
FullComponent.release_component_lock( FullComponent.release_component_lock(
&mut tc, node, /* params = */ None, /* lock_host = */ true, &mut tc,
node,
&DecLockRefParams::default(),
/* lock_host = */ true,
); );
assert!(tc.evictable_host_leaves.contains(node)); assert!(tc.evictable_host_leaves.contains(node));
} }
@@ -1359,11 +1399,11 @@ fn nested_host_locks_release_pairwise() {
.expect("live test node"); .expect("live test node");
tc.inc_host_lock_ref(tc.arena.node(node).id) tc.inc_host_lock_ref(tc.arena.node(node).id)
.expect("live test node"); .expect("live test node");
tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None) tc.dec_host_lock_ref(tc.arena.node(node).id, &DecLockRefParams::default())
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 1); assert_eq!(tc.arena.host_lock_ref(node, FULL), 1);
assert!(!tc.evictable_host_leaves.contains(node)); assert!(!tc.evictable_host_leaves.contains(node));
tc.dec_host_lock_ref(tc.arena.node(node).id, /* params = */ None) tc.dec_host_lock_ref(tc.arena.node(node).id, &DecLockRefParams::default())
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.host_lock_ref(node, FULL), 0); assert_eq!(tc.arena.host_lock_ref(node, FULL), 0);
assert!(tc.evictable_host_leaves.contains(node)); assert!(tc.evictable_host_leaves.contains(node));
rust/sglang-radix-tree/src/tests/components/mamba.rs
+84 -58
View File
@@ -1,5 +1,5 @@
use super::*; use super::*;
use crate::components::{FULL, MAMBA, SWA}; use crate::components::{ComponentSet, FULL, MAMBA, SWA};
use crate::test_utils::{accumulate_step, action_kinds}; use crate::test_utils::{accumulate_step, action_kinds};
use crate::unified_lru_list::UnifiedLRUList; use crate::unified_lru_list::UnifiedLRUList;
@@ -292,7 +292,6 @@ fn device_lock_moves_the_slot_between_evictable_and_protected_once() {
IncLockRefResult::default(), IncLockRefResult::default(),
/* lock_host = */ false, /* lock_host = */ false,
); );
assert!(result.skip_lock_node_ids.is_empty());
assert_eq!(tc.evictable_size_(MAMBA), 0); assert_eq!(tc.evictable_size_(MAMBA), 0);
assert_eq!(tc.protected_size_(MAMBA), 1); assert_eq!(tc.protected_size_(MAMBA), 1);
mamba.acquire_component_lock( mamba.acquire_component_lock(
@@ -303,25 +302,34 @@ fn device_lock_moves_the_slot_between_evictable_and_protected_once() {
); );
assert_eq!(tc.protected_size_(MAMBA), 1); assert_eq!(tc.protected_size_(MAMBA), 1);
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 2); assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 2);
mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ false); mamba.release_component_lock(
&mut tc,
a,
&DecLockRefParams::default(),
/* lock_host = */ false,
);
assert_eq!(tc.protected_size_(MAMBA), 1); assert_eq!(tc.protected_size_(MAMBA), 1);
mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ false); mamba.release_component_lock(
&mut tc,
a,
&DecLockRefParams::default(),
/* lock_host = */ false,
);
assert_eq!(tc.evictable_size_(MAMBA), 1); assert_eq!(tc.evictable_size_(MAMBA), 1);
assert_eq!(tc.protected_size_(MAMBA), 0); assert_eq!(tc.protected_size_(MAMBA), 0);
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0); assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0);
} }
#[test] #[test]
fn skip_aware_lock_records_only_the_mamba_target() { fn lock_without_mamba_records_the_receipt_and_leaves_mamba_evictable() {
let (mut tc, parent, leaf) = hybrid_lock_core(); let (mut tc, parent, leaf) = hybrid_lock_core();
let leaf_handle = tc.arena.node(leaf).id; let leaf_handle = tc.arena.node(leaf).id;
let result = tc let result = tc
.inc_lock_ref_with_skip(leaf_handle, &[MAMBA]) .inc_lock_ref(leaf_handle, ComponentSet::of(MAMBA))
.expect("live test node"); .expect("live test node");
assert_eq!(result.skip_lock_node_ids[&MAMBA].len(), 1); assert!(result.skipped_lock_components.contains(MAMBA));
assert!(result.skip_lock_node_ids[&MAMBA].contains(&leaf_handle));
assert_eq!(tc.arena.node(parent).device_lock_ref(MAMBA), 0); assert_eq!(tc.arena.node(parent).device_lock_ref(MAMBA), 0);
assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 0); assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 0);
assert_eq!(tc.evictable_size_(MAMBA), 2); assert_eq!(tc.evictable_size_(MAMBA), 2);
@@ -329,36 +337,45 @@ fn skip_aware_lock_records_only_the_mamba_target() {
assert_eq!(tc.arena.node(parent).device_lock_ref(FULL), 1); assert_eq!(tc.arena.node(parent).device_lock_ref(FULL), 1);
assert_eq!(tc.arena.node(leaf).device_lock_ref(FULL), 1); assert_eq!(tc.arena.node(leaf).device_lock_ref(FULL), 1);
tc.dec_lock_ref( // The receipt replays exactly what was taken: FULL only.
leaf_handle, let params = DecLockRefParams {
Some(&DecLockRefParams {
swa_uuid_for_lock: result.swa_uuid_for_lock, swa_uuid_for_lock: result.swa_uuid_for_lock,
skip_lock_node_ids: result.skip_lock_node_ids, skipped_lock_components: result.skipped_lock_components,
..Default::default() ..Default::default()
}), };
/* skip_swa = */ false, tc.dec_lock_ref(leaf_handle, &params, /* skip_swa = */ false)
)
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.node(parent).device_lock_ref(FULL), 0); assert_eq!(tc.arena.node(parent).device_lock_ref(FULL), 0);
assert_eq!(tc.arena.node(leaf).device_lock_ref(FULL), 0); assert_eq!(tc.arena.node(leaf).device_lock_ref(FULL), 0);
assert_eq!(tc.evictable_size_(MAMBA), 2);
assert_eq!(tc.protected_size_(MAMBA), 0);
} }
#[test] #[test]
fn swa_only_release_honors_a_skipped_mamba_target() { fn swa_only_release_spares_another_holders_mamba_lock() {
let (mut tc, _parent, leaf) = hybrid_lock_core(); let (mut tc, _parent, leaf) = hybrid_lock_core();
let leaf_handle = tc.arena.node(leaf).id; let leaf_handle = tc.arena.node(leaf).id;
let owner = tc.inc_lock_ref(leaf_handle).expect("live test node"); let owner = tc
let skipped = tc .inc_lock_ref(leaf_handle, ComponentSet::EMPTY)
.inc_lock_ref_with_skip(leaf_handle, &[MAMBA])
.expect("live test node"); .expect("live test node");
let holder = tc
.inc_lock_ref(leaf_handle, ComponentSet::of(MAMBA))
.expect("live test node");
assert!(!owner.skipped_lock_components.contains(MAMBA));
assert!(holder.skipped_lock_components.contains(MAMBA));
assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 1); assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 1);
// The holder's early SWA release must not drop the owner's mamba lock.
let holder_params = DecLockRefParams {
swa_uuid_for_lock: holder.swa_uuid_for_lock,
skipped_lock_components: holder.skipped_lock_components,
..Default::default()
};
let mut device_frees = HashMap::new(); let mut device_frees = HashMap::new();
let mut host_frees = HashMap::new(); let mut host_frees = HashMap::new();
tc.dec_swa_lock_only_with_skip( tc.dec_swa_lock_only(
leaf_handle, leaf_handle,
skipped.swa_uuid_for_lock, &holder_params,
Some(&skipped.skip_lock_node_ids),
&mut device_frees, &mut device_frees,
&mut host_frees, &mut host_frees,
) )
@@ -369,33 +386,22 @@ fn swa_only_release_honors_a_skipped_mamba_target() {
assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 1); assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 1);
assert_eq!(tc.protected_size_(MAMBA), 1); assert_eq!(tc.protected_size_(MAMBA), 1);
let skipped_params = DecLockRefParams { tc.dec_lock_ref(leaf_handle, &holder_params, /* skip_swa = */ true)
swa_uuid_for_lock: skipped.swa_uuid_for_lock,
skip_lock_node_ids: skipped.skip_lock_node_ids,
..Default::default()
};
tc.dec_lock_ref(
leaf_handle,
Some(&skipped_params),
/* skip_swa = */ true,
)
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 1);
let owner_params = DecLockRefParams { let owner_params = DecLockRefParams {
swa_uuid_for_lock: owner.swa_uuid_for_lock, swa_uuid_for_lock: owner.swa_uuid_for_lock,
skip_lock_node_ids: owner.skip_lock_node_ids, skipped_lock_components: owner.skipped_lock_components,
..Default::default() ..Default::default()
}; };
tc.dec_lock_ref( tc.dec_lock_ref(leaf_handle, &owner_params, /* skip_swa = */ false)
leaf_handle,
Some(&owner_params),
/* skip_swa = */ false,
)
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 0);
assert_eq!(tc.protected_size_(MAMBA), 0); assert_eq!(tc.protected_size_(MAMBA), 0);
} }
#[test] #[test]
fn tombstone_lock_is_recorded_and_replayed_at_release() { fn tombstone_lock_is_counted_with_no_ledger_move() {
let mut tc = mamba_core(/* page_size = */ 1); let mut tc = mamba_core(/* page_size = */ 1);
let [a] = chain::<1>(&mut tc); let [a] = chain::<1>(&mut tc);
let mamba = mamba_component(); let mamba = mamba_component();
@@ -405,14 +411,15 @@ fn tombstone_lock_is_recorded_and_replayed_at_release() {
IncLockRefResult::default(), IncLockRefResult::default(),
/* lock_host = */ false, /* lock_host = */ false,
); );
assert!(result.skip_lock_node_ids[&MAMBA].contains(&tc.arena.node(a).id)); assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 1);
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0); assert_eq!(tc.evictable_size_(MAMBA), 0);
// The replayed skip set keeps the release from touching the node. assert_eq!(tc.protected_size_(MAMBA), 0);
// The paired release decrements the counted tombstone, ledger untouched.
let params = DecLockRefParams { let params = DecLockRefParams {
skip_lock_node_ids: result.skip_lock_node_ids.clone(), skipped_lock_components: result.skipped_lock_components,
..DecLockRefParams::default() ..DecLockRefParams::default()
}; };
mamba.release_component_lock(&mut tc, a, Some(&params), /* lock_host = */ false); mamba.release_component_lock(&mut tc, a, &params, /* lock_host = */ false);
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0); assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0);
assert_eq!(tc.evictable_size_(MAMBA), 0); assert_eq!(tc.evictable_size_(MAMBA), 0);
} }
@@ -428,8 +435,12 @@ fn root_locks_are_noops() {
IncLockRefResult::default(), IncLockRefResult::default(),
/* lock_host = */ false, /* lock_host = */ false,
); );
assert!(result.skip_lock_node_ids.is_empty()); mamba.release_component_lock(
mamba.release_component_lock(&mut tc, root, None, /* lock_host = */ false); &mut tc,
root,
&DecLockRefParams::default(),
/* lock_host = */ false,
);
assert_eq!(tc.evictable_size_(MAMBA), 0); assert_eq!(tc.evictable_size_(MAMBA), 0);
} }
@@ -448,7 +459,12 @@ fn host_lock_detaches_and_reattaches_the_host_lru() {
); );
assert!(!tc.host_lru_list(MAMBA).in_list(Some(a))); assert!(!tc.host_lru_list(MAMBA).in_list(Some(a)));
assert_eq!(tc.arena.node(a).host_lock_ref(MAMBA), 1); assert_eq!(tc.arena.node(a).host_lock_ref(MAMBA), 1);
mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ true); mamba.release_component_lock(
&mut tc,
a,
&DecLockRefParams::default(),
/* lock_host = */ true,
);
assert!(tc.host_lru_list(MAMBA).in_list(Some(a))); assert!(tc.host_lru_list(MAMBA).in_list(Some(a)));
assert_eq!(tc.arena.node(a).host_lock_ref(MAMBA), 0); assert_eq!(tc.arena.node(a).host_lock_ref(MAMBA), 0);
} }
@@ -466,7 +482,12 @@ fn host_unlock_skips_the_lru_for_device_backed_nodes() {
IncLockRefResult::default(), IncLockRefResult::default(),
/* lock_host = */ true, /* lock_host = */ true,
); );
mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ true); mamba.release_component_lock(
&mut tc,
a,
&DecLockRefParams::default(),
/* lock_host = */ true,
);
assert!(!tc.host_lru_list(MAMBA).in_list(Some(a))); assert!(!tc.host_lru_list(MAMBA).in_list(Some(a)));
} }
@@ -1820,7 +1841,7 @@ fn branching_from_a_host_full_hit_is_reusable_after_insert() {
} }
#[test] #[test]
fn skip_set_release_after_a_restore_and_relock_keeps_the_new_lock() { fn release_after_a_restore_and_relock_keeps_the_other_lock() {
let mut tc = mamba_core(/* page_size = */ 1); let mut tc = mamba_core(/* page_size = */ 1);
let [a] = chain::<1>(&mut tc); let [a] = chain::<1>(&mut tc);
let mamba = mamba_component(); let mamba = mamba_component();
@@ -1830,28 +1851,33 @@ fn skip_set_release_after_a_restore_and_relock_keeps_the_new_lock() {
IncLockRefResult::default(), IncLockRefResult::default(),
/* lock_host = */ false, /* lock_host = */ false,
); );
assert!(first.skip_lock_node_ids[&MAMBA].contains(&tc.arena.node(a).id)); assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 1);
// The tombstone is restored and a second request locks it before the // The tombstone is restored under the held lock (credited to protected)
// first release replays its skip set. // and a second request stacks its own lock on it.
set_mamba_device(&mut tc, a, 7); tc.set_component_device_value_(a, MAMBA, Tensor::from_slice(&[7i64]));
let _ = mamba.acquire_component_lock( let _ = mamba.acquire_component_lock(
&mut tc, &mut tc,
a, a,
IncLockRefResult::default(), IncLockRefResult::default(),
/* lock_host = */ false, /* lock_host = */ false,
); );
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 1); assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 2);
assert_eq!(tc.evictable_size_(MAMBA), 0); assert_eq!(tc.evictable_size_(MAMBA), 0);
assert_eq!(tc.protected_size_(MAMBA), 1); assert_eq!(tc.protected_size_(MAMBA), 1);
let params = DecLockRefParams { let params = DecLockRefParams {
skip_lock_node_ids: first.skip_lock_node_ids.clone(), skipped_lock_components: first.skipped_lock_components,
..DecLockRefParams::default() ..DecLockRefParams::default()
}; };
mamba.release_component_lock(&mut tc, a, Some(&params), /* lock_host = */ false); mamba.release_component_lock(&mut tc, a, &params, /* lock_host = */ false);
// The replayed skip keeps the restored node's fresh lock intact. // The first release takes back exactly its own ref.
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 1); assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 1);
assert_eq!(tc.protected_size_(MAMBA), 1); assert_eq!(tc.protected_size_(MAMBA), 1);
mamba.release_component_lock(&mut tc, a, None, /* lock_host = */ false); mamba.release_component_lock(
&mut tc,
a,
&DecLockRefParams::default(),
/* lock_host = */ false,
);
assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0); assert_eq!(tc.arena.node(a).device_lock_ref(MAMBA), 0);
assert_eq!(tc.evictable_size_(MAMBA), 1); assert_eq!(tc.evictable_size_(MAMBA), 1);
assert_eq!(tc.protected_size_(MAMBA), 0); assert_eq!(tc.protected_size_(MAMBA), 0);
rust/sglang-radix-tree/src/tests/components/swa.rs
+272 -99
View File
@@ -1,5 +1,5 @@
use super::*; use super::*;
use crate::components::{FULL, MAMBA, SWA}; use crate::components::{ComponentSet, FULL, MAMBA, SWA};
use crate::test_utils::{accumulate_step, action_kinds}; use crate::test_utils::{accumulate_step, action_kinds};
use crate::unified_tree_core::CacheInitParams; use crate::unified_tree_core::CacheInitParams;
@@ -766,18 +766,32 @@ fn insert_overlap_recovers_a_tombstone_inside_the_window() {
} }
#[test] #[test]
#[should_panic(expected = "tombstone Swa lock_ref should be 0, node")] fn insert_overlap_recovers_a_locked_swa_tombstone() {
fn insert_overlap_panics_on_a_locked_swa_tombstone() {
let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1);
tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)); tc.insert(&insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0));
let root = tc.arena.root(); let root = tc.arena.root();
let leaf = child_of(&tc, root, &[1]); let leaf = child_of(&tc, root, &[1]);
// The rebuild is deferred, so the leaf is still an SWA tombstone; a raw // A segment lock may hold an SWA tombstone; the co-held FULL lock (the
// lock on it breaks the tombstones-are-unlocked contract. // full >= swa protocol invariant) forces the Recover branch, so the
// locked full stays on the node.
tc.arena tc.arena
.node_mut(leaf) .node_mut(leaf)
.set_lock_ref_(ValueSlotIdx::device(SWA), 1); .set_lock_ref_(ValueSlotIdx::device(SWA), 1);
tc.insert(&insert_params_swa(&vec![1, 2, 3], &[20, 21, 22], 0, 0)); tc.arena
.node_mut(leaf)
.set_lock_ref_(ValueSlotIdx::device(FULL), 1);
let result = tc.insert(&insert_params_swa(&vec![1, 2, 3], &[20, 21, 22], 0, 0));
assert!(
tc.arena
.device_value(leaf, FULL)
.equal(&Tensor::from_slice(&[10i64, 11, 12]))
);
assert!(
result
.cache_actions
.iter()
.any(|action| matches!(action, CacheAction::RecoverSwaWithLockedFull { .. }))
);
} }
#[test] #[test]
@@ -1576,7 +1590,7 @@ fn acquire_lock_reuses_the_stamped_uuid_and_shifts_sizes_once() {
} }
#[test] #[test]
fn acquire_lock_skips_tombstones_and_records_them() { fn acquire_lock_counts_tombstones_toward_the_window() {
let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1);
let [a, b, c] = chain(&mut tc); let [a, b, c] = chain(&mut tc);
store_swa_device(&mut tc, a); store_swa_device(&mut tc, a);
@@ -1587,14 +1601,13 @@ fn acquire_lock_skips_tombstones_and_records_them() {
IncLockRefResult::default(), IncLockRefResult::default(),
/* lock_host = */ false, /* lock_host = */ false,
); );
// The valueless b is recorded and skipped; the window fills at a. // The valueless b is counted too (no ledger move); position-based
// coverage fills the window at b, so a stays outside the segment.
assert_eq!(tc.arena.device_lock_ref(c, SWA), 1); assert_eq!(tc.arena.device_lock_ref(c, SWA), 1);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); assert_eq!(tc.arena.device_lock_ref(b, SWA), 1);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 1); assert_eq!(tc.arena.device_lock_ref(a, SWA), 0);
assert_eq!(result.skip_lock_node_ids[&SWA].len(), 1);
assert!(result.skip_lock_node_ids[&SWA].contains(&tc.arena.node(b).id));
assert!(result.swa_uuid_for_lock.is_some()); assert!(result.swa_uuid_for_lock.is_some());
assert_eq!(node_swa_uuid(&tc, a), result.swa_uuid_for_lock); assert_eq!(node_swa_uuid(&tc, b), result.swa_uuid_for_lock);
} }
#[test] #[test]
@@ -1624,11 +1637,14 @@ fn inc_lock_ref_runs_full_and_swa_walks_together() {
store_swa_device(&mut tc, b); store_swa_device(&mut tc, b);
store_swa_device(&mut tc, c); store_swa_device(&mut tc, c);
let result = tc let result = tc
.inc_lock_ref(tc.arena.node(c).id) .inc_lock_ref(tc.arena.node(c).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
// FULL sees a valueless path (skip segment only); SWA locks its window. // FULL counts its valueless bottom segment (no ledger move); SWA locks
// its window.
assert_eq!(result.delta, Some(0)); assert_eq!(result.delta, Some(0));
assert_eq!(result.skip_lock_node_ids[&FULL].len(), 3); assert_eq!(tc.arena.device_lock_ref(c, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(b, FULL), 1);
assert_eq!(tc.arena.device_lock_ref(a, FULL), 1);
assert!(result.swa_uuid_for_lock.is_some()); assert!(result.swa_uuid_for_lock.is_some());
assert_eq!(tc.arena.device_lock_ref(c, SWA), 1); assert_eq!(tc.arena.device_lock_ref(c, SWA), 1);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 1); assert_eq!(tc.arena.device_lock_ref(b, SWA), 1);
@@ -1657,10 +1673,10 @@ fn inc_host_lock_ref_runs_full_and_swa_host_arms_together() {
// The release replays the acquire's uuid and unwinds both arms. // The release replays the acquire's uuid and unwinds both arms.
let params = DecLockRefParams { let params = DecLockRefParams {
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids, skipped_lock_components: result.skipped_lock_components,
..Default::default() ..Default::default()
}; };
tc.dec_host_lock_ref(tc.arena.node(c).id, Some(&params)) tc.dec_host_lock_ref(tc.arena.node(c).id, &params)
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.host_lock_ref(c, FULL), 0); assert_eq!(tc.arena.host_lock_ref(c, FULL), 0);
assert_eq!(tc.arena.host_lock_ref(c, SWA), 0); assert_eq!(tc.arena.host_lock_ref(c, SWA), 0);
@@ -1690,10 +1706,10 @@ fn dec_host_lock_ref_with_the_inner_uuid_leaves_an_outer_window_pinned() {
// window's lock above the boundary survives. // window's lock above the boundary survives.
let params = DecLockRefParams { let params = DecLockRefParams {
swa_uuid_for_host_lock: inner.swa_uuid_for_host_lock, swa_uuid_for_host_lock: inner.swa_uuid_for_host_lock,
skip_lock_node_ids: inner.skip_lock_node_ids, skipped_lock_components: inner.skipped_lock_components,
..Default::default() ..Default::default()
}; };
tc.dec_host_lock_ref(tc.arena.node(c).id, Some(&params)) tc.dec_host_lock_ref(tc.arena.node(c).id, &params)
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.host_lock_ref(c, SWA), 0); assert_eq!(tc.arena.host_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.host_lock_ref(b, SWA), 1); assert_eq!(tc.arena.host_lock_ref(b, SWA), 1);
@@ -1758,7 +1774,7 @@ fn acquire_host_lock_reuses_the_stamped_uuid_and_skips_unlisted_nodes() {
} }
#[test] #[test]
fn acquire_host_lock_skips_host_tombstones_and_records_them() { fn acquire_host_lock_counts_host_tombstones_toward_the_window() {
let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1);
let [a, b, c] = chain(&mut tc); let [a, b, c] = chain(&mut tc);
set_swa_host(&mut tc, a); set_swa_host(&mut tc, a);
@@ -1769,12 +1785,12 @@ fn acquire_host_lock_skips_host_tombstones_and_records_them() {
IncLockRefResult::default(), IncLockRefResult::default(),
/* lock_host = */ true, /* lock_host = */ true,
); );
// The host-valueless b is counted too; position-based coverage fills
// the window at b, so a stays outside the segment.
assert_eq!(tc.arena.host_lock_ref(c, SWA), 1); assert_eq!(tc.arena.host_lock_ref(c, SWA), 1);
assert_eq!(tc.arena.host_lock_ref(b, SWA), 0); assert_eq!(tc.arena.host_lock_ref(b, SWA), 1);
assert_eq!(tc.arena.host_lock_ref(a, SWA), 1); assert_eq!(tc.arena.host_lock_ref(a, SWA), 0);
assert_eq!(result.skip_lock_node_ids[&SWA].len(), 1); assert_eq!(node_swa_host_uuid(&tc, b), result.swa_uuid_for_host_lock);
assert!(result.skip_lock_node_ids[&SWA].contains(&tc.arena.node(b).id));
assert_eq!(node_swa_host_uuid(&tc, a), result.swa_uuid_for_host_lock);
assert!(result.swa_uuid_for_host_lock.is_some()); assert!(result.swa_uuid_for_host_lock.is_some());
} }
@@ -1996,11 +2012,12 @@ fn release_lock_returns_the_window_to_evictable() {
/* lock_host = */ false, /* lock_host = */ false,
); );
let params = DecLockRefParams { let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock, swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids, skipped_lock_components: result.skipped_lock_components,
}; };
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ false); swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ false);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); assert_eq!(tc.arena.device_lock_ref(a, SWA), 0);
@@ -2029,11 +2046,12 @@ fn release_lock_keeps_sizes_while_other_locks_remain() {
/* lock_host = */ false, /* lock_host = */ false,
); );
let params = DecLockRefParams { let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: first.swa_uuid_for_lock, swa_uuid_for_lock: first.swa_uuid_for_lock,
swa_uuid_for_host_lock: first.swa_uuid_for_host_lock, swa_uuid_for_host_lock: first.swa_uuid_for_host_lock,
skip_lock_node_ids: first.skip_lock_node_ids, skipped_lock_components: first.skipped_lock_components,
}; };
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ false); swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ false);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 1); assert_eq!(tc.arena.device_lock_ref(c, SWA), 1);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 1); assert_eq!(tc.arena.device_lock_ref(b, SWA), 1);
assert_eq!(tc.swa_evictable_size(), 1); assert_eq!(tc.swa_evictable_size(), 1);
@@ -2056,11 +2074,12 @@ fn release_lock_replays_the_tombstone_skips() {
// b gained a device value AFTER the acquire recorded it as a tombstone. // b gained a device value AFTER the acquire recorded it as a tombstone.
store_swa_device(&mut tc, b); store_swa_device(&mut tc, b);
let params = DecLockRefParams { let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock, swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids, skipped_lock_components: result.skipped_lock_components,
}; };
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ false); swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ false);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); assert_eq!(tc.arena.device_lock_ref(a, SWA), 0);
@@ -2087,11 +2106,12 @@ fn release_lock_stops_at_the_window_uuid() {
.node_mut(a) .node_mut(a)
.set_lock_ref_(ValueSlotIdx::device(SWA), 1); .set_lock_ref_(ValueSlotIdx::device(SWA), 1);
let params = DecLockRefParams { let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock, swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids, skipped_lock_components: result.skipped_lock_components,
}; };
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ false); swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ false);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 1); assert_eq!(tc.arena.device_lock_ref(a, SWA), 1);
@@ -2120,11 +2140,12 @@ fn release_host_lock_stops_at_the_host_uuid_boundary() {
/* lock_host = */ true, /* lock_host = */ true,
); );
let params = DecLockRefParams { let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock, swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids, skipped_lock_components: result.skipped_lock_components,
}; };
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ true); swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ true);
assert_eq!(tc.arena.host_lock_ref(c, SWA), 0); assert_eq!(tc.arena.host_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.host_lock_ref(b, SWA), 0); assert_eq!(tc.arena.host_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.host_lock_ref(a, SWA), 1); assert_eq!(tc.arena.host_lock_ref(a, SWA), 1);
@@ -2132,34 +2153,27 @@ fn release_host_lock_stops_at_the_host_uuid_boundary() {
} }
#[test] #[test]
fn release_lock_without_params_passes_over_an_unlocked_middle_node() { #[should_panic(expected = "SWA segment release hit lock_ref=0")]
fn release_lock_without_the_boundary_uuid_dies_at_the_segment_edge() {
let mut tc = swa_core(/* window = */ 1, /* page_size = */ 1); let mut tc = swa_core(/* window = */ 1, /* page_size = */ 1);
let [a, b, c] = chain(&mut tc); let [_a, _b, c] = chain(&mut tc);
store_swa_device(&mut tc, a);
store_swa_device(&mut tc, b);
store_swa_device(&mut tc, c); store_swa_device(&mut tc, c);
let swa = swa_component(1); let swa = swa_component(1);
// The 1-atom window locks only the acquired node: c and a, never b. // The 1-atom window locks only the acquired node c.
let _ = swa.acquire_component_lock( let _ = swa.acquire_component_lock(
&mut tc, &mut tc,
c, c,
IncLockRefResult::default(), IncLockRefResult::default(),
/* lock_host = */ false, /* lock_host = */ false,
); );
let _ = swa.acquire_component_lock( // A receipt-less release overshoots the boundary into unlocked territory
// and dies there instead of silently stealing whatever it crosses.
swa.release_component_lock(
&mut tc, &mut tc,
a, c,
IncLockRefResult::default(), &DecLockRefParams::default(),
/* lock_host = */ false, /* lock_host = */ false,
); );
swa.release_component_lock(
&mut tc, c, /* params = */ None, /* lock_host = */ false,
);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 0);
assert_eq!(tc.swa_evictable_size(), 3);
assert_eq!(tc.swa_protected_size(), 0);
} }
#[test] #[test]
@@ -2178,11 +2192,12 @@ fn release_host_lock_reparks_tombstoned_host_nodes() {
/* lock_host = */ true, /* lock_host = */ true,
); );
let params = DecLockRefParams { let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock, swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids, skipped_lock_components: result.skipped_lock_components,
}; };
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ true); swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ true);
assert_eq!(tc.arena.host_lock_ref(c, SWA), 0); assert_eq!(tc.arena.host_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.host_lock_ref(b, SWA), 0); assert_eq!(tc.arena.host_lock_ref(b, SWA), 0);
assert!(tc.host_lru_list(SWA).in_list(Some(c))); assert!(tc.host_lru_list(SWA).in_list(Some(c)));
@@ -2198,18 +2213,15 @@ fn inc_then_dec_lock_ref_roundtrips_with_dec_params() {
store_swa_device(&mut tc, b); store_swa_device(&mut tc, b);
store_swa_device(&mut tc, c); store_swa_device(&mut tc, c);
let result = tc let result = tc
.inc_lock_ref(tc.arena.node(c).id) .inc_lock_ref(tc.arena.node(c).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
let params = DecLockRefParams { let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock, swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids, skipped_lock_components: result.skipped_lock_components,
}; };
tc.dec_lock_ref( tc.dec_lock_ref(tc.arena.node(c).id, &params, /* skip_swa = */ false)
tc.arena.node(c).id,
Some(&params),
/* skip_swa = */ false,
)
.expect("live test node"); .expect("live test node");
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
@@ -2230,13 +2242,17 @@ fn dec_swa_lock_only_releases_swa_while_full_stays_locked() {
// Fund FULL's evictable counter for its lock walk (raw slot sets skip it). // Fund FULL's evictable counter for its lock walk (raw slot sets skip it).
tc.component_state_mut(FULL).evictable_size = 3; tc.component_state_mut(FULL).evictable_size = 3;
let result = tc let result = tc
.inc_lock_ref(tc.arena.node(c).id) .inc_lock_ref(tc.arena.node(c).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
let mut device_frees = HashMap::new(); let mut device_frees = HashMap::new();
let mut host_frees = HashMap::new(); let mut host_frees = HashMap::new();
tc.dec_swa_lock_only( tc.dec_swa_lock_only(
tc.arena.node(c).id, tc.arena.node(c).id,
result.swa_uuid_for_lock, &DecLockRefParams {
swa_uuid_for_lock: result.swa_uuid_for_lock,
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees, &mut device_frees,
&mut host_frees, &mut host_frees,
) )
@@ -2273,7 +2289,11 @@ fn dec_swa_lock_only_evicts_a_fully_unlocked_device_leaf() {
let mut host_frees = HashMap::new(); let mut host_frees = HashMap::new();
tc.dec_swa_lock_only( tc.dec_swa_lock_only(
tc.arena.node(c).id, tc.arena.node(c).id,
result.swa_uuid_for_lock, &DecLockRefParams {
swa_uuid_for_lock: result.swa_uuid_for_lock,
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees, &mut device_frees,
&mut host_frees, &mut host_frees,
) )
@@ -2296,7 +2316,11 @@ fn dec_swa_lock_only_is_a_noop_without_the_swa_component() {
let mut host_frees = HashMap::new(); let mut host_frees = HashMap::new();
tc.dec_swa_lock_only( tc.dec_swa_lock_only(
tc.arena.node(root).id, tc.arena.node(root).id,
None, &DecLockRefParams {
swa_uuid_for_lock: None,
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees, &mut device_frees,
&mut host_frees, &mut host_frees,
) )
@@ -2374,11 +2398,12 @@ fn release_lock_skip_set_leaves_a_relocked_tombstone_credited() {
/* lock_host = */ false, /* lock_host = */ false,
); );
let params = DecLockRefParams { let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: first.swa_uuid_for_lock, swa_uuid_for_lock: first.swa_uuid_for_lock,
swa_uuid_for_host_lock: first.swa_uuid_for_host_lock, swa_uuid_for_host_lock: first.swa_uuid_for_host_lock,
skip_lock_node_ids: first.skip_lock_node_ids, skipped_lock_components: first.skipped_lock_components,
}; };
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ false); swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ false);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 1); assert_eq!(tc.arena.device_lock_ref(c, SWA), 1);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 1); assert_eq!(tc.arena.device_lock_ref(b, SWA), 1);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); assert_eq!(tc.arena.device_lock_ref(a, SWA), 0);
@@ -2387,28 +2412,27 @@ fn release_lock_skip_set_leaves_a_relocked_tombstone_credited() {
} }
#[test] #[test]
fn release_lock_passes_over_uncredited_nodes_without_params() { #[should_panic(expected = "SWA segment release hit lock_ref=0")]
fn double_release_with_one_receipt_dies_loud() {
let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1);
let [a, b, c] = chain(&mut tc); let [a, b, c] = chain(&mut tc);
store_swa_device(&mut tc, a); store_swa_device(&mut tc, a);
store_swa_device(&mut tc, b); store_swa_device(&mut tc, b);
store_swa_device(&mut tc, c); store_swa_device(&mut tc, c);
let swa = swa_component(2); let swa = swa_component(2);
let _ = swa.acquire_component_lock( let result = swa.acquire_component_lock(
&mut tc, &mut tc,
c, c,
IncLockRefResult::default(), IncLockRefResult::default(),
/* lock_host = */ false, /* lock_host = */ false,
); );
// No params: the walk crosses the never-credited a up to the root. let params = DecLockRefParams {
swa.release_component_lock( swa_uuid_for_lock: result.swa_uuid_for_lock,
&mut tc, c, /* params = */ None, /* lock_host = */ false, ..Default::default()
); };
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ false);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); // Consuming the same receipt twice dies at the first unlocked node.
assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ false);
assert_eq!(tc.swa_evictable_size(), 3);
assert_eq!(tc.swa_protected_size(), 0);
} }
#[test] #[test]
@@ -2435,7 +2459,11 @@ fn dec_swa_lock_only_releases_the_window_exactly_once() {
let mut host_frees = HashMap::new(); let mut host_frees = HashMap::new();
tc.dec_swa_lock_only( tc.dec_swa_lock_only(
tc.arena.node(c).id, tc.arena.node(c).id,
first.swa_uuid_for_lock, &DecLockRefParams {
swa_uuid_for_lock: first.swa_uuid_for_lock,
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees, &mut device_frees,
&mut host_frees, &mut host_frees,
) )
@@ -2447,7 +2475,11 @@ fn dec_swa_lock_only_releases_the_window_exactly_once() {
assert_eq!(tc.swa_protected_size(), 2); assert_eq!(tc.swa_protected_size(), 2);
tc.dec_swa_lock_only( tc.dec_swa_lock_only(
tc.arena.node(c).id, tc.arena.node(c).id,
first.swa_uuid_for_lock, &DecLockRefParams {
swa_uuid_for_lock: first.swa_uuid_for_lock,
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees, &mut device_frees,
&mut host_frees, &mut host_frees,
) )
@@ -2483,7 +2515,11 @@ fn dec_swa_lock_only_leaves_out_of_window_swa_locks_alone() {
let mut host_frees = HashMap::new(); let mut host_frees = HashMap::new();
tc.dec_swa_lock_only( tc.dec_swa_lock_only(
tc.arena.node(c).id, tc.arena.node(c).id,
result.swa_uuid_for_lock, &DecLockRefParams {
swa_uuid_for_lock: result.swa_uuid_for_lock,
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees, &mut device_frees,
&mut host_frees, &mut host_frees,
) )
@@ -2496,7 +2532,8 @@ fn dec_swa_lock_only_leaves_out_of_window_swa_locks_alone() {
} }
#[test] #[test]
fn release_window_lock_passes_over_an_unlocked_valued_node() { #[should_panic(expected = "SWA window release hit lock_ref=0")]
fn release_window_lock_without_the_uuid_dies_past_the_boundary() {
let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1); let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1);
let [a, b, c] = chain(&mut tc); let [a, b, c] = chain(&mut tc);
store_swa_device(&mut tc, a); store_swa_device(&mut tc, a);
@@ -2511,11 +2548,9 @@ fn release_window_lock_passes_over_an_unlocked_valued_node() {
); );
let mut device_frees = HashMap::new(); let mut device_frees = HashMap::new();
let mut host_frees = HashMap::new(); let mut host_frees = HashMap::new();
// No uuid bound: the walk crosses the valued-but-unlocked a to the root. // Without the boundary uuid the walk crosses the segment edge into the
// unlocked a and dies there instead of stealing.
swa.release_window_lock(&mut tc, c, None, &mut device_frees, &mut host_frees); swa.release_window_lock(&mut tc, c, None, &mut device_frees, &mut host_frees);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(a, SWA), 0);
} }
#[test] #[test]
@@ -2533,7 +2568,8 @@ fn release_window_lock_passes_over_a_mid_chain_tombstone_without_a_uuid() {
); );
let mut device_frees = HashMap::new(); let mut device_frees = HashMap::new();
let mut host_frees = HashMap::new(); let mut host_frees = HashMap::new();
// No uuid bound: the walk crosses the mid-chain tombstone b and releases a. // Walked-to-root acquire (window > chain): the uuid-less release counts
// back through the mid-chain tombstone b and releases a.
swa.release_window_lock(&mut tc, c, None, &mut device_frees, &mut host_frees); swa.release_window_lock(&mut tc, c, None, &mut device_frees, &mut host_frees);
assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); assert_eq!(tc.arena.device_lock_ref(c, SWA), 0);
assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); assert_eq!(tc.arena.device_lock_ref(b, SWA), 0);
@@ -2558,11 +2594,12 @@ fn release_host_lock_does_not_repark_a_node_whose_host_value_was_taken() {
// device value either, so the release has nothing to park. // device value either, so the release has nothing to park.
let _ = tc.arena.take_host_value(a, SWA); let _ = tc.arena.take_host_value(a, SWA);
let params = DecLockRefParams { let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock, swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids, skipped_lock_components: result.skipped_lock_components,
}; };
swa.release_component_lock(&mut tc, a, Some(&params), /* lock_host = */ true); swa.release_component_lock(&mut tc, a, &params, /* lock_host = */ true);
assert_eq!(tc.arena.host_lock_ref(a, SWA), 0); assert_eq!(tc.arena.host_lock_ref(a, SWA), 0);
assert!(!tc.host_lru_list(SWA).in_list(Some(a))); assert!(!tc.host_lru_list(SWA).in_list(Some(a)));
} }
@@ -2584,11 +2621,12 @@ fn release_host_lock_skips_reparking_device_valued_nodes() {
/* lock_host = */ true, /* lock_host = */ true,
); );
let params = DecLockRefParams { let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock, swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids, skipped_lock_components: result.skipped_lock_components,
}; };
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ true); swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ true);
// Device-valued nodes never re-park in the host LRU on host release. // Device-valued nodes never re-park in the host LRU on host release.
assert!(!tc.host_lru_list(SWA).in_list(Some(c))); assert!(!tc.host_lru_list(SWA).in_list(Some(c)));
assert!(!tc.host_lru_list(SWA).in_list(Some(b))); assert!(!tc.host_lru_list(SWA).in_list(Some(b)));
@@ -2612,11 +2650,12 @@ fn release_host_lock_leaves_an_already_listed_node_listed() {
// Something re-listed b while the lock was held (e.g. a split re-park). // Something re-listed b while the lock was held (e.g. a split re-park).
tc.host_lru_list_mut(SWA).insert_mru(b); tc.host_lru_list_mut(SWA).insert_mru(b);
let params = DecLockRefParams { let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: result.swa_uuid_for_lock, swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skip_lock_node_ids: result.skip_lock_node_ids, skipped_lock_components: result.skipped_lock_components,
}; };
swa.release_component_lock(&mut tc, c, Some(&params), /* lock_host = */ true); swa.release_component_lock(&mut tc, c, &params, /* lock_host = */ true);
assert!(tc.host_lru_list(SWA).in_list(Some(b))); assert!(tc.host_lru_list(SWA).in_list(Some(b)));
assert!(tc.host_lru_list(SWA).in_list(Some(c))); assert!(tc.host_lru_list(SWA).in_list(Some(c)));
let _ = a; let _ = a;
@@ -2890,6 +2929,29 @@ fn redistribute_on_node_split_moves_the_swa_uuid_to_the_parent() {
assert_eq!(node_swa_uuid(&tc, node), None); assert_eq!(node_swa_uuid(&tc, node), None);
} }
#[test]
fn redistribute_on_node_split_preserves_host_lock_state() {
let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1);
tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0));
let root = tc.arena.root();
let node = child_of(&tc, root, &[1]);
// A host-locked tombstone mid-IO: both halves must stay pinned and out
// of the host LRU, and the host boundary uuid moves to the parent.
tc.arena
.set_host_value(node, SWA, Tensor::from_slice(&[70i64, 71]));
tc.arena
.node_mut(node)
.set_lock_ref_(ValueSlotIdx::host(SWA), 1);
tc.arena.node_mut(node).swa_host_uuid = Some(9);
let (parent, _) = tc.split_node_(node, /* split_len = */ 1);
assert_eq!(tc.arena.host_lock_ref(parent, SWA), 1);
assert_eq!(tc.arena.host_lock_ref(node, SWA), 1);
assert_eq!(node_swa_host_uuid(&tc, parent), Some(9));
assert_eq!(node_swa_host_uuid(&tc, node), None);
assert!(!tc.host_lru_list(SWA).in_list(Some(parent)));
assert!(!tc.host_lru_list(SWA).in_list(Some(node)));
}
#[test] #[test]
fn finalize_window_arithmetic_at_page_boundaries() { fn finalize_window_arithmetic_at_page_boundaries() {
let mut tc = swa_core(/* window = */ 4, /* page_size = */ 2); let mut tc = swa_core(/* window = */ 4, /* page_size = */ 2);
@@ -3277,17 +3339,24 @@ fn reinsert_rejects_a_page_misaligned_boundary() {
} }
#[test] #[test]
#[should_panic(expected = "tombstone Swa lock_ref should be 0 on unevict")] fn reinsert_rebuilds_a_locked_tombstone() {
fn reinsert_rejects_a_locked_tombstone() {
let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1); let mut tc = swa_core(/* window = */ 8, /* page_size = */ 1);
tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0)); tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0));
let root = tc.arena.root(); let root = tc.arena.root();
let node = child_of(&tc, root, &[1]); let node = child_of(&tc, root, &[1]);
evict_full(&mut tc, node, /* remaining_size = */ 0); evict_full(&mut tc, node, /* remaining_size = */ 0);
// Segment locks count evicted nodes, so a locked tombstone is legal and
// the re-insert rebuilds its SWA from the fresh KV.
tc.arena tc.arena
.node_mut(node) .node_mut(node)
.set_lock_ref_(ValueSlotIdx::device(SWA), 1); .set_lock_ref_(ValueSlotIdx::device(SWA), 1);
tc.insert(&insert_params_swa(&vec![1, 2], &[20, 21], 0, 0)); let result = tc.insert(&insert_params_swa(&vec![1, 2], &[20, 21], 0, 0));
assert!(
result
.cache_actions
.iter()
.any(|action| matches!(action, CacheAction::SwaRebuild { .. }))
);
} }
fn set_full_host(tc: &mut UnifiedTreeCore<Vec<i64>>, node: NodeIdx_) { fn set_full_host(tc: &mut UnifiedTreeCore<Vec<i64>>, node: NodeIdx_) {
@@ -4856,13 +4925,16 @@ fn deep_swa_tree_survives_backup_evict_and_load_back_rounds() {
) )
.expect("live test node"); .expect("live test node");
assert!(actions.is_empty()); assert!(actions.is_empty());
let lock = tc.inc_lock_ref(anchor).expect("live test node"); let lock = tc
.inc_lock_ref(anchor, ComponentSet::EMPTY)
.expect("live test node");
let params = DecLockRefParams { let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: lock.swa_uuid_for_lock, swa_uuid_for_lock: lock.swa_uuid_for_lock,
swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock, swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock,
skip_lock_node_ids: lock.skip_lock_node_ids, skipped_lock_components: lock.skipped_lock_components,
}; };
tc.dec_lock_ref(anchor, Some(&params), /* skip_swa = */ false) tc.dec_lock_ref(anchor, &params, /* skip_swa = */ false)
.expect("live test node"); .expect("live test node");
tc.finish_load_back(anchor).expect("live test node"); tc.finish_load_back(anchor).expect("live test node");
} }
@@ -4960,3 +5032,104 @@ fn recovered_swa_span_evicts_before_the_window_leaf() {
assert!(tc.arena.has_device_value(leaf, FULL)); assert!(tc.arena.has_device_value(leaf, FULL));
tc.sanity_check(&[], &[]); tc.sanity_check(&[], &[]);
} }
#[test]
fn aux_release_refreshes_the_leaf_set_whatever_the_release_order() {
let mut tc = swa_core(/* window = */ 1, /* page_size = */ 1);
tc.insert(&insert_params_swa(&vec![1], &[10], 0, 0));
let leaf = child_of(&tc, tc.arena.root(), &[1]);
store_swa_device(&mut tc, leaf);
assert!(tc.evictable_device_leaves.contains(leaf));
let result = tc
.inc_lock_ref(tc.arena.node(leaf).id, ComponentSet::EMPTY)
.expect("live test node");
assert!(!tc.evictable_device_leaves.contains(leaf));
let params = DecLockRefParams {
node_id: result.node_id,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: None,
skipped_lock_components: ComponentSet::EMPTY,
};
// Full first: its walk still sees the SWA lock, so the leaf stays out.
crate::components::FullComponent
.release_component_lock(&mut tc, leaf, &params, /* lock_host = */ false);
assert!(!tc.evictable_device_leaves.contains(leaf));
// The SWA release drops the last lock and must readmit the leaf itself.
swa_component(1).release_component_lock(&mut tc, leaf, &params, /* lock_host = */ false);
assert!(tc.evictable_device_leaves.contains(leaf));
}
#[test]
#[should_panic(expected = "lock receipt anchored on node")]
fn dec_lock_ref_rejects_a_receipt_from_another_node() {
let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1);
let [a, _b, c] = chain(&mut tc);
store_swa_device(&mut tc, c);
let result = tc
.inc_lock_ref(tc.arena.node(c).id, ComponentSet::EMPTY)
.expect("live test node");
let params = DecLockRefParams {
node_id: result.node_id,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: None,
skipped_lock_components: result.skipped_lock_components,
};
// Same receipt, wrong anchor: the walk would otherwise release a's
// segment, which this holder never locked.
tc.dec_lock_ref(tc.arena.node(a).id, &params, /* skip_swa = */ false)
.expect("live test node");
}
#[test]
#[should_panic(expected = "lock receipt anchored on node")]
fn dec_host_lock_ref_rejects_a_receipt_from_another_node() {
let mut tc = swa_core(/* window = */ 2, /* page_size = */ 1);
let [a, _b, c] = chain(&mut tc);
set_swa_host(&mut tc, c);
tc.host_lru_list_mut(SWA).insert_mru(c);
let result = tc
.inc_host_lock_ref(tc.arena.node(c).id)
.expect("live test node");
let params = DecLockRefParams {
node_id: result.node_id,
swa_uuid_for_lock: None,
swa_uuid_for_host_lock: result.swa_uuid_for_host_lock,
skipped_lock_components: ComponentSet::EMPTY,
};
tc.dec_host_lock_ref(tc.arena.node(a).id, &params)
.expect("live test node");
}
#[test]
fn receipt_anchor_follows_the_locked_node_through_a_split() {
let mut tc = swa_core(/* window = */ 1, /* page_size = */ 1);
tc.insert(&insert_params_swa(&vec![1, 2], &[10, 11], 0, 0));
let leaf = child_of(&tc, tc.arena.root(), &[1]);
store_swa_device(&mut tc, leaf);
let leaf_id = tc.arena.node(leaf).id;
let result = tc
.inc_lock_ref(leaf_id, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(result.node_id, Some(leaf_id));
// Diverge inside the node: the split keeps the id on the deeper half.
tc.insert(&insert_params_swa(&vec![1, 3], &[12, 13], 0, 0));
let params = DecLockRefParams {
node_id: result.node_id,
swa_uuid_for_lock: result.swa_uuid_for_lock,
swa_uuid_for_host_lock: None,
skipped_lock_components: ComponentSet::EMPTY,
};
tc.dec_lock_ref(leaf_id, &params, /* skip_swa = */ false)
.expect("live test node");
assert_eq!(
tc.arena
.device_lock_ref(tc.arena.resolve(leaf_id).expect("live test node"), SWA),
0
);
}
#[test]
#[should_panic(expected = "swa_sliding_window_size must be positive")]
fn new_panics_on_a_zero_sliding_window_size() {
SwaComponent::new(&swa_params_with_window(0));
}
rust/sglang-radix-tree/src/tests/unified_tree_core.rs
+85 -35
View File
@@ -3,7 +3,7 @@ use std::sync::Mutex;
use tch::Tensor; use tch::Tensor;
use super::*; use super::*;
use crate::components::{FULL, MAMBA, SWA}; use crate::components::{ComponentSet, FULL, MAMBA, SWA};
use crate::node::{NodeAccessError, ValueSlotIdx}; use crate::node::{NodeAccessError, ValueSlotIdx};
use crate::test_utils::{accumulate_step, action_kinds}; use crate::test_utils::{accumulate_step, action_kinds};
@@ -91,7 +91,7 @@ impl TreeComponent<Vec<i64>> for RecordingComponentForTest {
&self, &self,
_tree_core: &mut UnifiedTreeCore<Vec<i64>>, _tree_core: &mut UnifiedTreeCore<Vec<i64>>,
_node_id: NodeIdx_, _node_id: NodeIdx_,
_params: Option<&DecLockRefParams>, _params: &DecLockRefParams,
_lock_host: bool, _lock_host: bool,
) { ) {
unimplemented!() unimplemented!()
@@ -213,7 +213,7 @@ impl TreeComponent<Vec<i64>> for CountingComponentForTest {
&self, &self,
_tree_core: &mut UnifiedTreeCore<Vec<i64>>, _tree_core: &mut UnifiedTreeCore<Vec<i64>>,
_node_id: NodeIdx_, _node_id: NodeIdx_,
_params: Option<&DecLockRefParams>, _params: &DecLockRefParams,
_lock_host: bool, _lock_host: bool,
) { ) {
unimplemented!() unimplemented!()
@@ -293,11 +293,11 @@ impl TreeComponent<Vec<i64>> for LowPriorityComponentForTest {
&self, &self,
_tree_core: &mut UnifiedTreeCore<Vec<i64>>, _tree_core: &mut UnifiedTreeCore<Vec<i64>>,
_node_id: NodeIdx_, _node_id: NodeIdx_,
params: Option<&DecLockRefParams>, params: &DecLockRefParams,
lock_host: bool, lock_host: bool,
) { ) {
assert!(!lock_host); assert!(!lock_host);
assert!(params.is_some_and(|p| p.swa_uuid_for_lock.is_some())); assert!(params.swa_uuid_for_lock.is_some());
panic!("low-priority release dispatched"); panic!("low-priority release dispatched");
} }
} }
@@ -368,7 +368,7 @@ impl TreeComponent<Vec<i64>> for SwaComponentForTest {
&self, &self,
_tree_core: &mut UnifiedTreeCore<Vec<i64>>, _tree_core: &mut UnifiedTreeCore<Vec<i64>>,
_node_id: NodeIdx_, _node_id: NodeIdx_,
_params: Option<&DecLockRefParams>, _params: &DecLockRefParams,
_lock_host: bool, _lock_host: bool,
) { ) {
unimplemented!() unimplemented!()
@@ -481,7 +481,7 @@ impl TreeComponent<Vec<i64>> for SwaEvictionComponentForTest {
&self, &self,
_tree_core: &mut UnifiedTreeCore<Vec<i64>>, _tree_core: &mut UnifiedTreeCore<Vec<i64>>,
_node_id: NodeIdx_, _node_id: NodeIdx_,
_params: Option<&DecLockRefParams>, _params: &DecLockRefParams,
_lock_host: bool, _lock_host: bool,
) { ) {
unimplemented!() unimplemented!()
@@ -503,7 +503,7 @@ fn locked_anchor_for_dispatch(tc: &mut UnifiedTreeCore<Vec<i64>>) -> NodeIdx_ {
tc.arena tc.arena
.set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1])); .set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1]));
tc.component_state_mut(FULL).evictable_size = 2; tc.component_state_mut(FULL).evictable_size = 2;
tc.inc_lock_ref(tc.arena.node(n1).id) tc.inc_lock_ref(tc.arena.node(n1).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
n1 n1
} }
@@ -516,7 +516,11 @@ fn dec_lock_ref_skip_swa_skips_the_swa_component() {
// The skipped Swa driver is never dispatched, so its stub cannot panic. // The skipped Swa driver is never dispatched, so its stub cannot panic.
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena.node(n1).id, tc.arena.node(n1).id,
/* params = */ None, /* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ true, /* skip_swa = */ true,
) )
.expect("live test node"); .expect("live test node");
@@ -541,7 +545,7 @@ fn inc_lock_ref_reaches_every_component() {
tc.arena tc.arena
.set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1])); .set_device_value(n1, FULL, Tensor::from_slice(&[0i64, 1]));
tc.component_state_mut(FULL).evictable_size = 2; tc.component_state_mut(FULL).evictable_size = 2;
let _ = tc.inc_lock_ref(tc.arena.node(n1).id); let _ = tc.inc_lock_ref(tc.arena.node(n1).id, ComponentSet::EMPTY);
} }
#[test] #[test]
@@ -552,7 +556,11 @@ fn dec_lock_ref_without_skip_swa_reaches_every_component() {
tc.register_component_(Arc::new(SwaComponentForTest)); tc.register_component_(Arc::new(SwaComponentForTest));
let _ = tc.dec_lock_ref( let _ = tc.dec_lock_ref(
tc.arena.node(n1).id, tc.arena.node(n1).id,
/* params = */ None, /* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false, /* skip_swa = */ false,
); );
} }
@@ -630,7 +638,11 @@ fn dec_swa_lock_only_dispatches_lower_priority_releases() {
let mut host_frees = HashMap::new(); let mut host_frees = HashMap::new();
let _ = tc.dec_swa_lock_only( let _ = tc.dec_swa_lock_only(
tc.arena.node(root).id, tc.arena.node(root).id,
Some(7), &DecLockRefParams {
swa_uuid_for_lock: Some(7),
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees, &mut device_frees,
&mut host_frees, &mut host_frees,
); );
@@ -671,7 +683,11 @@ fn dec_swa_lock_only_returns_device_frees_in_the_device_dict() {
let mut host_frees = HashMap::new(); let mut host_frees = HashMap::new();
tc.dec_swa_lock_only( tc.dec_swa_lock_only(
tc.arena.node(a).id, tc.arena.node(a).id,
result.swa_uuid_for_lock, &DecLockRefParams {
swa_uuid_for_lock: result.swa_uuid_for_lock,
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
&mut device_frees, &mut device_frees,
&mut host_frees, &mut host_frees,
) )
@@ -3755,11 +3771,15 @@ fn commit_load_back_reattaches_device_slices_and_restores_the_match() {
assert_eq!(tc.full_evictable_size(), 4); assert_eq!(tc.full_evictable_size(), 4);
// The orchestrator re-locks the loaded path right after commit; that lock walk // The orchestrator re-locks the loaded path right after commit; that lock walk
// also re-evaluates the parent's transient D-leaf membership. // also re-evaluates the parent's transient D-leaf membership.
tc.inc_lock_ref(tc.arena.node(child).id) tc.inc_lock_ref(tc.arena.node(child).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena.node(child).id, tc.arena.node(child).id,
/* params = */ None, /* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false, /* skip_swa = */ false,
) )
.expect("live test node"); .expect("live test node");
@@ -5261,7 +5281,7 @@ fn stale_handle_returns_err_after_its_node_is_freed() {
tc.evict_device_leaf(leaf, /* is_write_back = */ false) tc.evict_device_leaf(leaf, /* is_write_back = */ false)
.expect("live test node"); .expect("live test node");
assert!(matches!( assert!(matches!(
tc.inc_lock_ref(leaf), tc.inc_lock_ref(leaf, ComponentSet::EMPTY),
Err(NodeAccessError { node_id }) if node_id == leaf Err(NodeAccessError { node_id }) if node_id == leaf
)); ));
} }
@@ -6164,7 +6184,7 @@ fn reset_restores_a_fresh_tree() {
..insert_params(&vec![7, 8], &[20, 21]) ..insert_params(&vec![7, 8], &[20, 21])
}); });
let matched = tc.match_prefix(&match_params(&vec![1, 2, 3])); let matched = tc.match_prefix(&match_params(&vec![1, 2, 3]));
tc.inc_lock_ref(matched.best_match_node_id) tc.inc_lock_ref(matched.best_match_node_id, ComponentSet::EMPTY)
.expect("live match node"); .expect("live match node");
assert_eq!(tc.protected_size(), 3); assert_eq!(tc.protected_size(), 3);
// Seed aux LRU, host LRU, and host-leaf state so the reset must clear each. // Seed aux LRU, host LRU, and host-leaf state so the reset must clear each.
@@ -6216,7 +6236,7 @@ fn size_accessors_mirror_the_full_component_state() {
assert_eq!(tc.protected_size(), 0); assert_eq!(tc.protected_size(), 0);
assert_eq!(tc.component_evictable_size(FULL), 3); assert_eq!(tc.component_evictable_size(FULL), 3);
let matched = tc.match_prefix(&match_params(&vec![1, 2, 3])); let matched = tc.match_prefix(&match_params(&vec![1, 2, 3]));
tc.inc_lock_ref(matched.best_match_node_id) tc.inc_lock_ref(matched.best_match_node_id, ComponentSet::EMPTY)
.expect("live match node"); .expect("live match node");
assert_eq!(tc.protected_size(), 3); assert_eq!(tc.protected_size(), 3);
assert_eq!(tc.full_protected_size(), 3); assert_eq!(tc.full_protected_size(), 3);
@@ -6306,7 +6326,7 @@ fn walk_for_kv_canary_chains_slots_across_namespaces() {
fn walk_for_kv_canary_unlocked_only_skips_locked_nodes_but_keeps_the_chain() { fn walk_for_kv_canary_unlocked_only_skips_locked_nodes_but_keeps_the_chain() {
let mut tc = core(); let mut tc = core();
let (a, _b) = matched_chain(&mut tc); let (a, _b) = matched_chain(&mut tc);
tc.inc_lock_ref(tc.arena.node(a).id) tc.inc_lock_ref(tc.arena.node(a).id, ComponentSet::EMPTY)
.expect("live test node"); .expect("live test node");
assert_eq!( assert_eq!(
sorted_canary_rows(tc.walk_for_kv_canary(true, false)), sorted_canary_rows(tc.walk_for_kv_canary(true, false)),
@@ -6659,13 +6679,18 @@ fn sanity_check_passes_on_a_healthy_tree() {
let leaf = tc let leaf = tc
.match_prefix(&match_params(&vec![1, 2, 9])) .match_prefix(&match_params(&vec![1, 2, 9]))
.best_match_node_id; .best_match_node_id;
tc.inc_lock_ref(leaf).expect("live test node"); tc.inc_lock_ref(leaf, ComponentSet::EMPTY)
.expect("live test node");
tc.sanity_check(&[(1, leaf)], &[(2, leaf)]); tc.sanity_check(&[(1, leaf)], &[(2, leaf)]);
tc.dec_lock_ref( tc.dec_lock_ref(
tc.arena tc.arena
.node(tc.arena.resolve(leaf).expect("live test node")) .node(tc.arena.resolve(leaf).expect("live test node"))
.id, .id,
/* params = */ None, /* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false, /* skip_swa = */ false,
) )
.expect("live test node"); .expect("live test node");
@@ -6774,16 +6799,28 @@ fn sanity_check_detects_an_evicted_parent_prefix() {
} }
#[test] #[test]
#[should_panic(expected = "evicted but lock_ref")] fn sanity_check_accepts_a_locked_tombstone() {
fn sanity_check_detects_a_locked_tombstone() { // Segment locks count evicted nodes, so a device-locked tombstone is a
// legal state the checker must not flag.
let mut tc = sane_tree(); let mut tc = sane_tree();
// write_back spares the tombstone's ancestors the backup-chain rule.
tc.is_write_back = true;
let leaf = tc let leaf = tc
.match_prefix(&match_params(&vec![1, 2, 9])) .match_prefix(&match_params(&vec![1, 2, 9]))
.best_match_node_id; .best_match_node_id;
tc.inc_lock_ref(leaf).expect("live test node"); let leaf_idx = tc.arena.resolve(leaf).expect("live test node");
let _ = tc // Tombstone the leaf consistently first (host copy, ledger, leaf sets),
.arena // then lock through it: the bottom segment counts the tombstone.
.take_device_value(tc.arena.resolve(leaf).expect("live test node"), FULL); tc.arena
.set_host_value(leaf_idx, FULL, Tensor::from_slice(&[9i64]));
let taken = tc.arena.take_device_value(leaf_idx, FULL);
tc.dec_evictable_size(FULL, taken.size()[0] as usize);
tc.update_evictable_leaf_sets_(leaf_idx);
let parent_idx = tc.arena.node(leaf_idx).parent();
tc.update_evictable_leaf_sets_(parent_idx);
tc.inc_lock_ref(leaf, ComponentSet::EMPTY)
.expect("live test node");
assert_eq!(tc.arena.device_lock_ref(leaf_idx, FULL), 1);
tc.sanity_check(&[], &[]); tc.sanity_check(&[], &[]);
} }
@@ -8049,13 +8086,16 @@ fn run_random_op_sequence(mut tc: UnifiedTreeCore<Vec<i64>>, page: usize, mamba:
2 => { 2 => {
// Balanced lock round trip on whatever the key matches. // Balanced lock round trip on whatever the key matches.
let anchor = tc.match_prefix(&match_params(&key)).best_match_node_id; let anchor = tc.match_prefix(&match_params(&key)).best_match_node_id;
let lock = tc.inc_lock_ref(anchor).expect("live match anchor"); let lock = tc
.inc_lock_ref(anchor, ComponentSet::EMPTY)
.expect("live match anchor");
let params = DecLockRefParams { let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: lock.swa_uuid_for_lock, swa_uuid_for_lock: lock.swa_uuid_for_lock,
swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock, swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock,
skip_lock_node_ids: lock.skip_lock_node_ids, skipped_lock_components: lock.skipped_lock_components,
}; };
tc.dec_lock_ref(anchor, Some(&params), /* skip_swa = */ false) tc.dec_lock_ref(anchor, &params, /* skip_swa = */ false)
.expect("live match anchor"); .expect("live match anchor");
} }
_ => { _ => {
@@ -8063,7 +8103,9 @@ fn run_random_op_sequence(mut tc: UnifiedTreeCore<Vec<i64>>, page: usize, mamba:
let matched = tc.match_prefix(&match_params(&key)); let matched = tc.match_prefix(&match_params(&key));
let anchor = matched.best_match_node_id; let anchor = matched.best_match_node_id;
let matched_len = matched.device_indices.numel() as usize; let matched_len = matched.device_indices.numel() as usize;
let lock = tc.inc_lock_ref(anchor).expect("live match anchor"); let lock = tc
.inc_lock_ref(anchor, ComponentSet::EMPTY)
.expect("live match anchor");
tc.insert(&sequence_insert_params( tc.insert(&sequence_insert_params(
&key, &key,
matched_len, matched_len,
@@ -8072,11 +8114,12 @@ fn run_random_op_sequence(mut tc: UnifiedTreeCore<Vec<i64>>, page: usize, mamba:
mamba, mamba,
)); ));
let params = DecLockRefParams { let params = DecLockRefParams {
node_id: None,
swa_uuid_for_lock: lock.swa_uuid_for_lock, swa_uuid_for_lock: lock.swa_uuid_for_lock,
swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock, swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock,
skip_lock_node_ids: lock.skip_lock_node_ids, skipped_lock_components: lock.skipped_lock_components,
}; };
tc.dec_lock_ref(anchor, Some(&params), /* skip_swa = */ false) tc.dec_lock_ref(anchor, &params, /* skip_swa = */ false)
.expect("live match anchor"); .expect("live match anchor");
} }
} }
@@ -8224,10 +8267,17 @@ fn a_zero_length_match_anchors_at_the_root() {
.best_match_node_id; .best_match_node_id;
assert_eq!(anchor, tc.root_node_handle(Some("salted"))); assert_eq!(anchor, tc.root_node_handle(Some("salted")));
// The root handle stays valid across a full namespace eviction. // The root handle stays valid across a full namespace eviction.
tc.inc_lock_ref(anchor).expect("live root"); tc.inc_lock_ref(anchor, ComponentSet::EMPTY)
.expect("live root");
drain_full_device(&mut tc); drain_full_device(&mut tc);
tc.dec_lock_ref( tc.dec_lock_ref(
anchor, /* params = */ None, /* skip_swa = */ false, anchor,
/* params = */
&DecLockRefParams {
skipped_lock_components: ComponentSet::EMPTY,
..Default::default()
},
/* skip_swa = */ false,
) )
.expect("live root"); .expect("live root");
assert!(tc.arena.resolve(anchor).is_ok()); assert!(tc.arena.resolve(anchor).is_ok());
rust/sglang-radix-tree/src/unified_tree_core.rs
+147 -97
View File
@@ -8,7 +8,9 @@ use std::sync::Arc;
use sha2::{Digest, Sha256}; use sha2::{Digest, Sha256};
use tch::{Device, Kind, Tensor}; use tch::{Device, Kind, Tensor};
use crate::components::{self, FullComponent, MambaComponent, SwaComponent, TreeComponent}; use crate::components::{
self, ComponentSet, FullComponent, MambaComponent, SwaComponent, TreeComponent,
};
use crate::components::{ use crate::components::{
BASE_COMPONENT_TYPE, ComponentType, FULL, MAMBA, NUM_COMPONENT_TYPES, SWA, BASE_COMPONENT_TYPE, ComponentType, FULL, MAMBA, NUM_COMPONENT_TYPES, SWA,
}; };
@@ -34,28 +36,41 @@ fn next_coexist_reclaim_digest(current: i64, node_id: NodeId, component_idx: usi
// ---- interface types ---- // ---- interface types ----
/// Result of `inc_lock_ref`, handed back to the matching `dec_lock_ref`. /// Result of `inc_lock_ref`, handed back to the matching `dec_lock_ref`.
///
/// The receipt a release needs is per-component lock evidence: the SWA
/// segment boundary uuid (None means the segment reached the root) and
/// whether the single-node Mamba lock was taken (the decode hold opts
/// out). Locks count every node in their contiguous segment, so no
/// per-node skip state exists. Receipt fields default to nothing-acquired;
/// `inc_lock_ref` stamps what it actually took.
#[derive(Default)] #[derive(Default)]
pub struct IncLockRefResult { pub struct IncLockRefResult {
/// Tokens newly protected (moved out of evictable) by this lock. /// Tokens newly protected (moved out of evictable) by this lock.
pub delta: Option<usize>, pub delta: Option<usize>,
/// The node the lock was taken on; a release replays the receipt there only.
pub node_id: Option<NodeId>,
/// SWA lock-window uuid minted/reused by the device lock walk. /// SWA lock-window uuid minted/reused by the device lock walk.
pub swa_uuid_for_lock: Option<i64>, pub swa_uuid_for_lock: Option<i64>,
/// SWA lock-window uuid minted/reused by the host lock walk. /// SWA lock-window uuid minted/reused by the host lock walk.
pub swa_uuid_for_host_lock: Option<i64>, pub swa_uuid_for_host_lock: Option<i64>,
/// Per-component nodes that were tombstones at acquire time; replayed at /// Components the acquire left untaken; the release skips them too.
/// release so the unlock skips them. pub skipped_lock_components: ComponentSet,
pub skip_lock_node_ids: HashMap<ComponentType, HashSet<NodeId>>,
} }
/// Params for `dec_lock_ref`. /// Params for `dec_lock_ref`. Receipt fields default to nothing-acquired so
/// a lost receipt under-releases (a leak sanity checks report) instead of
/// releasing a lock another holder owns.
#[derive(Default)] #[derive(Default)]
pub struct DecLockRefParams { pub struct DecLockRefParams {
/// The node the matching acquire locked; None only for receipts that did
/// not come from this core (a mispaired anchor is a protocol violation).
pub node_id: Option<NodeId>,
/// SWA lock-window uuid the device unlock stops at, from the matching acquire. /// SWA lock-window uuid the device unlock stops at, from the matching acquire.
pub swa_uuid_for_lock: Option<i64>, pub swa_uuid_for_lock: Option<i64>,
/// SWA lock-window uuid the host unlock stops at, from the matching acquire. /// SWA lock-window uuid the host unlock stops at, from the matching acquire.
pub swa_uuid_for_host_lock: Option<i64>, pub swa_uuid_for_host_lock: Option<i64>,
/// Per-component nodes the unlock walk skips (from the matching acquire). /// Components the matching acquire left untaken.
pub skip_lock_node_ids: HashMap<ComponentType, HashSet<NodeId>>, pub skipped_lock_components: ComponentSet,
} }
/// Result of `dec_lock_ref`. /// Result of `dec_lock_ref`.
@@ -792,58 +807,90 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
self.swa_uuid_counter self.swa_uuid_counter
} }
/// Bump the reference count on a node's component locks. /// Bump the reference count on a node's component locks. Components in
pub fn inc_lock_ref(&mut self, node_id: NodeId) -> Result<IncLockRefResult, NodeAccessError> { /// `skip_lock_components` are left untaken; the receipt records the anchor
self.inc_lock_ref_with_skip(node_id, &[]) /// node and the skipped set so the paired release mirrors them.
} pub fn inc_lock_ref(
/// Bump component locks, leaving explicitly skipped target components evictable.
pub fn inc_lock_ref_with_skip(
&mut self, &mut self,
node_id: NodeId, node_id: NodeId,
skip_lock_components: &[ComponentType], skip_lock_components: ComponentSet,
) -> Result<IncLockRefResult, NodeAccessError> { ) -> Result<IncLockRefResult, NodeAccessError> {
let node_id = self.arena.resolve(node_id)?; let node_idx = self.arena.resolve(node_id)?;
let node = self.arena.node(node_id); let mut result = IncLockRefResult {
let node_handle = node.id; node_id: Some(self.arena.node(node_idx).id),
let is_root = node.is_root(); skipped_lock_components: skip_lock_components,
let mut result = IncLockRefResult::default(); ..Default::default()
};
for i in 0..self.components.len() { for i in 0..self.components.len() {
let component_type = self.components[i].component_type(); let component = Arc::clone(&self.components[i]);
if skip_lock_components.contains(&component_type) { if skip_lock_components.contains(component.component_type()) {
if !is_root {
result
.skip_lock_node_ids
.entry(component_type)
.or_default()
.insert(node_handle);
}
continue; continue;
} }
let component = Arc::clone(&self.components[i]);
result = component result = component
.acquire_component_lock(self, node_id, result, /* lock_host = */ false); .acquire_component_lock(self, node_idx, result, /* lock_host = */ false);
} }
self.update_evictable_leaf_sets_(node_id); self.update_evictable_leaf_sets_(node_idx);
Ok(result) Ok(result)
} }
/// Decrease the reference count on a node's component locks. /// A receipt releases only the node its acquire returned; a mispaired
/// node would silently release (or steal) another holder's segment.
fn assert_receipt_anchor_(&self, node_idx: NodeIdx_, params: &DecLockRefParams) {
if let Some(anchor) = params.node_id {
let node_handle = self.arena.node(node_idx).id;
assert!(
anchor == node_handle,
"lock receipt anchored on node {anchor} released on node {node_handle}"
);
}
}
/// 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.
fn release_components_(
&mut self,
node_idx: NodeIdx_,
params: &DecLockRefParams,
lock_host: bool,
skip_swa_and_below: bool,
) {
let swa_priority = if skip_swa_and_below {
self.try_component_by_type_(SWA)
.map(|swa| swa.eviction_priority(/* is_leaf = */ false))
} else {
None
};
for i in (0..self.components.len()).rev() {
let component = Arc::clone(&self.components[i]);
let ct = component.component_type();
if params.skipped_lock_components.contains(ct) {
continue;
}
if let Some(swa_priority) = swa_priority
&& (ct == SWA || component.eviction_priority(/* is_leaf = */ false) < swa_priority)
{
continue;
}
component.release_component_lock(self, node_idx, params, lock_host);
}
}
/// Decrease the reference count on a node's component locks. The receipt
/// is required: a release must replay its acquire's evidence. After an SWA
/// early release (`dec_swa_lock_only`), `skip_swa` leaves SWA and the
/// lower-priority components it already dropped alone.
pub fn dec_lock_ref( pub fn dec_lock_ref(
&mut self, &mut self,
node_id: NodeId, node_id: NodeId,
params: Option<&DecLockRefParams>, params: &DecLockRefParams,
skip_swa: bool, skip_swa: bool,
) -> Result<DecLockRefResult, NodeAccessError> { ) -> Result<DecLockRefResult, NodeAccessError> {
let node_id = self.arena.resolve(node_id)?; let node_idx = self.arena.resolve(node_id)?;
for i in 0..self.components.len() { self.assert_receipt_anchor_(node_idx, params);
if skip_swa && self.components[i].component_type() == SWA { self.release_components_(node_idx, params, /* lock_host = */ false, skip_swa);
continue; self.update_evictable_leaf_sets_(node_idx);
}
let component = Arc::clone(&self.components[i]);
component.release_component_lock(self, node_id, params, /* lock_host = */ false);
}
self.update_evictable_leaf_sets_(node_id);
// TODO: delta is not aggregated from components; no caller uses it yet. // TODO: delta is not aggregated from components; no caller uses it yet.
Ok(DecLockRefResult::default()) Ok(DecLockRefResult::default())
} }
@@ -853,50 +900,37 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
pub fn dec_swa_lock_only( pub fn dec_swa_lock_only(
&mut self, &mut self,
node_id: NodeId, node_id: NodeId,
swa_uuid_for_lock: Option<i64>, params: &DecLockRefParams,
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>, device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>, host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
) -> Result<(), NodeAccessError> { ) -> Result<(), NodeAccessError> {
self.dec_swa_lock_only_with_skip( let node_idx = self.arena.resolve(node_id)?;
node_id, self.assert_receipt_anchor_(node_idx, params);
swa_uuid_for_lock,
/* skip_lock_node_ids = */ None,
device_frees,
host_frees,
)
}
/// Skip-aware variant used when an acquire deliberately omitted a component.
pub fn dec_swa_lock_only_with_skip(
&mut self,
node_id: NodeId,
swa_uuid_for_lock: Option<i64>,
skip_lock_node_ids: Option<&HashMap<ComponentType, HashSet<NodeId>>>,
device_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
host_frees: &mut HashMap<ComponentType, Vec<Tensor>>,
) -> Result<(), NodeAccessError> {
let node_id = self.arena.resolve(node_id)?;
let Some(swa) = self.try_component_by_type_(SWA) else { let Some(swa) = self.try_component_by_type_(SWA) else {
return Ok(()); return Ok(());
}; };
swa.release_window_lock(self, node_id, swa_uuid_for_lock, device_frees, host_frees); swa.release_window_lock(
// Drop strictly-lower-priority locks (e.g. Mamba) co-located on the node.
let swa_priority = swa.eviction_priority(/* is_leaf = */ false);
let dec_params = DecLockRefParams {
swa_uuid_for_lock,
skip_lock_node_ids: skip_lock_node_ids.cloned().unwrap_or_default(),
..Default::default()
};
for i in 0..self.components.len() {
let component = Arc::clone(&self.components[i]);
if component.eviction_priority(/* is_leaf = */ false) < swa_priority {
component.release_component_lock(
self, self,
node_id, node_idx,
Some(&dec_params), params.swa_uuid_for_lock,
/* lock_host = */ false, device_frees,
host_frees,
); );
// Drop strictly-lower-priority locks co-located on the node, skipping
// any the paired inc never took.
let swa_priority = swa.eviction_priority(/* is_leaf = */ false);
for i in (0..self.components.len()).rev() {
let component = Arc::clone(&self.components[i]);
if params
.skipped_lock_components
.contains(component.component_type())
{
continue;
}
if component.eviction_priority(/* is_leaf = */ false) < swa_priority {
component
.release_component_lock(self, node_idx, params, /* lock_host = */ false);
} }
} }
Ok(()) Ok(())
@@ -925,29 +959,31 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
&mut self, &mut self,
node_id: NodeId, node_id: NodeId,
) -> Result<IncLockRefResult, NodeAccessError> { ) -> Result<IncLockRefResult, NodeAccessError> {
let node_id = self.arena.resolve(node_id)?; let node_idx = self.arena.resolve(node_id)?;
let mut result = IncLockRefResult::default(); let mut result = IncLockRefResult {
node_id: Some(self.arena.node(node_idx).id),
..Default::default()
};
for i in 0..self.components.len() { for i in 0..self.components.len() {
let component = Arc::clone(&self.components[i]); let component = Arc::clone(&self.components[i]);
result = component result = component
.acquire_component_lock(self, node_id, result, /* lock_host = */ true); .acquire_component_lock(self, node_idx, result, /* lock_host = */ true);
} }
self.update_evictable_leaf_sets_(node_id); self.update_evictable_leaf_sets_(node_idx);
Ok(result) Ok(result)
} }
/// Decrease the reference count on a node's host-side component locks. /// Decrease the reference count on a node's host-side component locks.
/// The receipt is required, as for `dec_lock_ref`.
pub fn dec_host_lock_ref( pub fn dec_host_lock_ref(
&mut self, &mut self,
node_id: NodeId, node_id: NodeId,
params: Option<&DecLockRefParams>, params: &DecLockRefParams,
) -> Result<DecLockRefResult, NodeAccessError> { ) -> Result<DecLockRefResult, NodeAccessError> {
let node_id = self.arena.resolve(node_id)?; let node_idx = self.arena.resolve(node_id)?;
for i in 0..self.components.len() { self.assert_receipt_anchor_(node_idx, params);
let component = Arc::clone(&self.components[i]); self.release_components_(node_idx, params, /* lock_host = */ true, false);
component.release_component_lock(self, node_id, params, /* lock_host = */ true); self.update_evictable_leaf_sets_(node_idx);
}
self.update_evictable_leaf_sets_(node_id);
Ok(DecLockRefResult::default()) Ok(DecLockRefResult::default())
} }
@@ -1877,7 +1913,14 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
pub fn unevict_node_on_insert_(&mut self, node_id: NodeIdx_, fresh_value: &Tensor) { pub fn unevict_node_on_insert_(&mut self, node_id: NodeIdx_, fresh_value: &Tensor) {
self.arena self.arena
.set_device_value(node_id, FULL, fresh_value.copy()); .set_device_value(node_id, FULL, fresh_value.copy());
self.inc_evictable_size(FULL, fresh_value.size()[0] as usize); let tokens = fresh_value.size()[0] as usize;
// A value materialized under lock is protected; the last release
// moves it to evictable.
if self.arena.device_lock_ref(node_id, FULL) > 0 {
self.inc_protected_size(FULL, tokens);
} else {
self.inc_evictable_size(FULL, tokens);
}
self.update_evictable_leaf_sets_(node_id); self.update_evictable_leaf_sets_(node_id);
self.update_full_coexisting_host_tracking_(node_id); self.update_full_coexisting_host_tracking_(node_id);
if let Some(parent_id) = self.arena.node(node_id).try_parent() { if let Some(parent_id) = self.arena.node(node_id).try_parent() {
@@ -2644,6 +2687,11 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
if node.is_host_locked() { if node.is_host_locked() {
return false; 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 node.is_device_locked() {
return false;
}
if !node.children.is_empty() { if !node.children.is_empty() {
return false; return false;
} }
@@ -3702,8 +3750,14 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
host_lru.remove_node(node_id); host_lru.remove_node(node_id);
} }
self.device_lru_list_mut(component_type).insert_mru(node_id); self.device_lru_list_mut(component_type).insert_mru(node_id);
// A value materialized under lock is protected; the last release
// moves it to evictable.
if self.arena.device_lock_ref(node_id, component_type) > 0 {
self.inc_protected_size(component_type, tokens);
} else {
self.inc_evictable_size(component_type, tokens); self.inc_evictable_size(component_type, tokens);
} }
}
/// The component's device value on the node, or None if evicted. /// The component's device value on the node, or None if evicted.
pub fn get_component_device_value( pub fn get_component_device_value(
@@ -3931,12 +3985,8 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
device_state.lock_ref device_state.lock_ref
)); ));
} }
if device_state.value.is_none() && device_state.lock_ref > 0 { // Locked tombstones are legal: segment locks count every
errors.push(format!( // node in [start, boundary], data-bearing or not.
"node {node_id} {ct:?} evicted but lock_ref={}",
device_state.lock_ref
));
}
} }
// Collect expected leaf qualification (single pass) // Collect expected leaf qualification (single pass)
test/registered/unit/managers/test_scheduler_chunked_req_gate.py
+2 -1
View File
@@ -7,6 +7,7 @@ from unittest.mock import MagicMock
import torch import torch
from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams
from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel from sglang.test.test_utils import CustomTestCase, maybe_stub_sgl_kernel
@@ -40,7 +41,7 @@ def _make_req(
req.kv = ReqKvInfo(req_pool_idx=req_pool_idx) req.kv = ReqKvInfo(req_pool_idx=req_pool_idx)
req.skip_radix_cache_insert = False req.skip_radix_cache_insert = False
req.last_node = None req.last_node = None
req.swa_uuid_for_lock = None req.lock_receipt = DecLockRefParams()
req.session = None req.session = None
req.return_logprob = False req.return_logprob = False
req.logprob_start_len = -1 req.logprob_start_len = -1
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+55 -3
View File
@@ -35,13 +35,17 @@ from unittest.mock import MagicMock
import torch import torch
from sglang.srt.disaggregation.decode import DecodePreallocQueue from sglang.srt.disaggregation.decode import DecodePreallocQueue
from sglang.srt.disaggregation.decode_hicache_mixin import DecodePrefixMatch from sglang.srt.disaggregation.decode_hicache_mixin import (
DecodeHiCacheTransferMixin,
DecodePrefixMatch,
)
from sglang.srt.mem_cache.base_prefix_cache import ( from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams, DecLockRefParams,
InsertParams, InsertParams,
MatchPrefixParams, MatchPrefixParams,
) )
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey
from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
from sglang.srt.utils.common import Range from sglang.srt.utils.common import Range
@@ -387,7 +391,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
req.last_node = object() req.last_node = object()
req.finished_reason = None req.finished_reason = None
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
req.swa_uuid_for_lock = 123 req.lock_receipt = DecLockRefParams(swa_uuid_for_lock=123)
req.swa_prefix_lock_released = False req.swa_prefix_lock_released = False
req.pd_rebootstrap_in_progress = False req.pd_rebootstrap_in_progress = False
req.sampling_params.max_new_tokens = 16 req.sampling_params.max_new_tokens = 16
@@ -462,7 +466,10 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
self.assertEqual(preallocated, []) self.assertEqual(preallocated, [])
self.assertEqual(failed, []) self.assertEqual(failed, [])
queue._pre_alloc.assert_not_called() queue._pre_alloc.assert_not_called()
queue.tree_cache.dec_swa_lock_only.assert_called_once_with(req.last_node, 123) queue.tree_cache.dec_swa_lock_only.assert_called_once_with(
req.last_node,
DecLockRefParams(swa_uuid_for_lock=123),
)
queue.tree_cache.dec_lock_ref.assert_called_once_with( queue.tree_cache.dec_lock_ref.assert_called_once_with(
req.last_node, req.last_node,
DecLockRefParams(swa_uuid_for_lock=123), DecLockRefParams(swa_uuid_for_lock=123),
@@ -472,6 +479,51 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
queue._swa_tail_len.assert_called_once_with(8) queue._swa_tail_len.assert_called_once_with(8)
queue._allocatable_token_budgets.assert_called_once() queue._allocatable_token_budgets.assert_called_once()
def test_hicache_restore_commit_hands_over_lock_with_receipt(self):
"""The hicache-restore commit must release the prealloc lock with the
req's receipt, honoring a prior early SWA release (skip_swa), and hand
the restored node's lock to the req atomically: receipt fields move
with last_node, the early-release flag resets (the restored lock is
fresh), and the decode_req drops ownership so a post-commit abort
cannot release the restored lock a second time."""
q = DecodeHiCacheTransferMixin.__new__(DecodeHiCacheTransferMixin)
q.tree_cache = MagicMock()
req = MagicMock()
req.req_pool_idx = 0
req.lock_receipt = DecLockRefParams(swa_uuid_for_lock=123)
req.swa_prefix_lock_released = True # SWA tail-prealloc released early
prealloc_node = object()
restored_node = object()
decode_req = MagicMock()
decode_req.req = req
decode_req.prefix_match = DecodePrefixMatch(
prefix_indices=torch.arange(4, dtype=torch.int64),
l2_host_hit_length=4,
l3_storage_hit_length=0,
last_device_node=prealloc_node,
)
decode_req.hicache_restored_node = restored_node
decode_req.hicache_restore_lock_receipt = DecLockRefParams(
swa_uuid_for_lock=456, skipped_lock_components=(ComponentType.MAMBA,)
)
decode_req.hicache_restored_kv_indices = torch.arange(4, 8, dtype=torch.int64)
q._commit_hicache_local_restore_to_req(decode_req)
q.tree_cache.dec_lock_ref.assert_called_once_with(
prealloc_node,
DecLockRefParams(swa_uuid_for_lock=123),
skip_swa=True,
)
self.assertIs(req.last_node, restored_node)
self.assertEqual(req.lock_receipt.swa_uuid_for_lock, 456)
self.assertIn(ComponentType.MAMBA, req.lock_receipt.skipped_lock_components)
self.assertFalse(req.swa_prefix_lock_released)
self.assertIsNone(decode_req.hicache_restored_node)
self.assertIsNone(decode_req.hicache_restore_lock_receipt)
def test_repeated_incremental_no_leak(self): def test_repeated_incremental_no_leak(self):
"""Multiple incremental transfers shouldn't leak lock_refs.""" """Multiple incremental transfers shouldn't leak lock_refs."""
cache, req_to_token = _make_cache_with_pools() cache, req_to_token = _make_cache_with_pools()
test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py
+18 -13
View File
@@ -16,6 +16,7 @@ from types import SimpleNamespace
import torch import torch
from sglang.srt.mem_cache.base_prefix_cache import ( from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
EvictParams, EvictParams,
IncLockRefResult, IncLockRefResult,
) )
@@ -182,12 +183,12 @@ class _RecordingComp:
class TestDecSwaLockSkip(unittest.TestCase): class TestDecSwaLockSkip(unittest.TestCase):
"""dec_swa_lock_only early-releases SWA plus co-located lower-tier (Mamba) """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 locks. On a node whose acquire skipped Mamba (decode hold), the release
into that lower-tier release, else it drops a mamba lock it never took -- must skip it too, else it drops a mamba lock it never took -- another
another request's, on a shared FULL+SWA+MAMBA node (Inkling). Guards the request's, on a shared FULL+SWA+MAMBA node (Inkling). Guards the contract
contract without booting a 3-component model.""" without booting a 3-component model."""
def test_threads_skip_ids_into_lower_tier_release(self): def _run(self, skipped_lock_components):
# internal-node priority: full=2 > swa=1 > mamba=0 # internal-node priority: full=2 > swa=1 > mamba=0
full = _RecordingComp(ComponentType.FULL, 2) full = _RecordingComp(ComponentType.FULL, 2)
swa = _RecordingComp(ComponentType.SWA, 1) swa = _RecordingComp(ComponentType.SWA, 1)
@@ -197,23 +198,27 @@ class TestDecSwaLockSkip(unittest.TestCase):
components=(full, swa, mamba), components=(full, swa, mamba),
components_by_type={ComponentType.SWA: swa}, components_by_type={ComponentType.SWA: swa},
node_by_id=lambda node_id: node, node_by_id=lambda node_id: node,
_assert_receipt_anchor=UnifiedTreeCore._assert_receipt_anchor,
) )
UnifiedTreeCore.dec_swa_lock_only( UnifiedTreeCore.dec_swa_lock_only(
tree_core, tree_core,
node.id, node.id,
swa_uuid_for_lock=None, DecLockRefParams(skipped_lock_components=skipped_lock_components),
skip_lock_node_ids={ComponentType.MAMBA: {7}},
) )
return full, mamba
# mamba (below swa) is released, honoring the skip set def test_unlocked_mamba_is_not_released(self):
self.assertEqual(len(mamba.released), 1) full, mamba = self._run(skipped_lock_components=(ComponentType.MAMBA,))
self.assertEqual( # mamba took no lock at acquire, so the early release skips it too
mamba.released[0].skip_lock_node_ids.get(ComponentType.MAMBA), {7} self.assertEqual(mamba.released, [])
)
# full (above swa) is never touched # full (above swa) is never touched
self.assertEqual(full.released, []) self.assertEqual(full.released, [])
def test_lower_tier_released_when_locked(self):
full, mamba = self._run(skipped_lock_components=())
self.assertEqual(len(mamba.released), 1)
self.assertEqual(full.released, [])
class TestMambaDonatedAllocRatio(unittest.TestCase): class TestMambaDonatedAllocRatio(unittest.TestCase):
def test_prefill_peak_ratio2_exhausts_pool(self): def test_prefill_peak_ratio2_exhausts_pool(self):
test/registered/unit/mem_cache/test_rust_tree_core.py
+2 -2
View File
@@ -76,10 +76,10 @@ def test_lock_moves_tokens_between_evictable_and_protected():
InsertParams(key=_key([1, 2]), value=torch.tensor([10, 11], dtype=torch.int64)), InsertParams(key=_key([1, 2]), value=torch.tensor([10, 11], dtype=torch.int64)),
) )
matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))) matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2])))
core.inc_lock_ref(matched.best_match_node) lock = core.inc_lock_ref(matched.best_match_node)
assert core.protected_size() == 2 assert core.protected_size() == 2
assert core.evictable_size() == 0 assert core.evictable_size() == 0
core.dec_lock_ref(matched.best_match_node) core.dec_lock_ref(matched.best_match_node, lock.to_dec_params())
assert core.evictable_size() == 2 assert core.evictable_size() == 2
test/registered/unit/mem_cache/test_rust_tree_core_integration.py
+50 -58
View File
@@ -26,6 +26,7 @@ from sglang.srt.disaggregation.kv_events import (
) )
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.mem_cache.base_prefix_cache import ( from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
InsertParams, InsertParams,
InsertResult, InsertResult,
MatchPrefixParams, MatchPrefixParams,
@@ -230,8 +231,10 @@ def test_stale_handle_operations_raise_key_error_without_poisoning_the_core():
operations = { operations = {
"inc_lock_ref": lambda: core.inc_lock_ref(stale_root), "inc_lock_ref": lambda: core.inc_lock_ref(stale_root),
"dec_lock_ref": lambda: core.dec_lock_ref(stale_root), "dec_lock_ref": lambda: core.dec_lock_ref(stale_root, DecLockRefParams()),
"dec_swa_lock_only": lambda: core.dec_swa_lock_only(stale_root, None), "dec_swa_lock_only": lambda: core.dec_swa_lock_only(
stale_root, DecLockRefParams()
),
"evict_device_leaf": lambda: core.evict_device_leaf(stale_root, False), "evict_device_leaf": lambda: core.evict_device_leaf(stale_root, False),
"drop_subtree_no_host": lambda: core.drop_subtree_no_host(stale_root), "drop_subtree_no_host": lambda: core.drop_subtree_no_host(stale_root),
"demote": lambda: core.demote(stale_root), "demote": lambda: core.demote(stale_root),
@@ -267,7 +270,9 @@ def test_stale_handle_operations_raise_key_error_without_poisoning_the_core():
stale_root, {}, {} stale_root, {}, {}
), ),
"inc_host_lock_ref": lambda: core.inc_host_lock_ref(stale_root), "inc_host_lock_ref": lambda: core.inc_host_lock_ref(stale_root),
"dec_host_lock_ref": lambda: core.dec_host_lock_ref(stale_root), "dec_host_lock_ref": lambda: core.dec_host_lock_ref(
stale_root, DecLockRefParams()
),
"mark_write_through_pending": lambda: core.mark_write_through_pending( "mark_write_through_pending": lambda: core.mark_write_through_pending(
[stale_root], stale_root [stale_root], stale_root
), ),
@@ -416,10 +421,10 @@ def test_lock_and_unlock_move_tokens_between_protected_and_evictable():
_insert(core, [1, 2, 3], [10, 11, 12]) _insert(core, [1, 2, 3], [10, 11, 12])
_insert(core, [1, 2, 3, 4, 5], [20, 21, 22, 13, 14]) _insert(core, [1, 2, 3, 4, 5], [20, 21, 22, 13, 14])
matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3, 4, 5]))) matched = core.match_prefix(MatchPrefixParams(key=_key([1, 2, 3, 4, 5])))
core.inc_lock_ref(matched.best_match_node) lock = core.inc_lock_ref(matched.best_match_node)
assert core.protected_size() == 5 assert core.protected_size() == 5
assert core.evictable_size() == 0 assert core.evictable_size() == 0
core.dec_lock_ref(matched.best_match_node) core.dec_lock_ref(matched.best_match_node, lock.to_dec_params())
assert core.protected_size() == 0 assert core.protected_size() == 0
assert core.evictable_size() == 5 assert core.evictable_size() == 5
@@ -874,8 +879,8 @@ def test_host_lock_refs_round_trip():
_insert(core, [1], [10]) _insert(core, [1], [10])
leaf = core.match_prefix(MatchPrefixParams(key=_key([1]))).best_match_node leaf = core.match_prefix(MatchPrefixParams(key=_key([1]))).best_match_node
core.commit_backup(leaf, torch.tensor([100], dtype=torch.int64), {}) core.commit_backup(leaf, torch.tensor([100], dtype=torch.int64), {})
core.inc_host_lock_ref(leaf) host_lock = core.inc_host_lock_ref(leaf)
core.dec_host_lock_ref(leaf) core.dec_host_lock_ref(leaf, host_lock.to_dec_params())
core.sanity_check([], []) core.sanity_check([], [])
@@ -1204,11 +1209,6 @@ def test_swa_requires_the_sliding_window_size():
) )
def test_swa_without_a_window_is_rejected_through_the_adapter():
with pytest.raises(ValueError, match="requires swa_sliding_window_size"):
_tree_core(tree_components=(ComponentType.FULL, ComponentType.SWA))
def test_enable_hicache_constructs(): def test_enable_hicache_constructs():
mem_cache.RustUnifiedTreeCoreBinding( mem_cache.RustUnifiedTreeCoreBinding(
mem_cache.TreeCoreInitParamsBinding(enable_hicache=True), mem_cache.TreeCoreInitParamsBinding(enable_hicache=True),
@@ -1260,6 +1260,14 @@ def _swa_tree_core(window: int = 8, **params_overrides) -> RustUnifiedTreeCore:
) )
def test_swa_core_rejects_a_missing_or_non_positive_window():
"""A zero window can never fill, so no boundary uuid would ever be stamped;
the adapter refuses it up front instead of letting the core misbehave later."""
for window in (None, 0, -1):
with pytest.raises(ValueError, match="positive sliding_window_size"):
_swa_tree_core(window=window)
def test_write_back_load_back_ignores_auxiliary_nodes_for_pending_ownership(): def test_write_back_load_back_ignores_auxiliary_nodes_for_pending_ownership():
core = _swa_tree_core(window=4) core = _swa_tree_core(window=4)
core.set_hicache_enabled() core.set_hicache_enabled()
@@ -1583,20 +1591,20 @@ def test_skipped_mamba_lock_survives_swa_only_release_through_the_adapter():
node = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node node = core.match_prefix(MatchPrefixParams(key=_key([1, 2]))).best_match_node
owner = core.inc_lock_ref(node) owner = core.inc_lock_ref(node)
skipped = core.inc_lock_ref(node, skip_lock_components=(ComponentType.MAMBA,)) holder = core.inc_lock_ref(node, skip_lock_components=(ComponentType.MAMBA,))
assert skipped.skip_lock_node_ids == {ComponentType.MAMBA: {node}} assert ComponentType.MAMBA not in owner.skipped_lock_components
assert ComponentType.MAMBA in holder.skipped_lock_components
assert core.mamba_protected_size() == 1 assert core.mamba_protected_size() == 1
released = core.dec_swa_lock_only( # The holder's receipt says it never took mamba: its early SWA release
node, # must leave the owner's mamba lock alone.
skipped.swa_uuid_for_lock, released = core.dec_swa_lock_only(node, holder.to_dec_params())
skip_lock_node_ids=skipped.skip_lock_node_ids,
)
assert dict(released.device_frees) == {} assert dict(released.device_frees) == {}
assert dict(released.host_frees) == {} assert dict(released.host_frees) == {}
assert core.mamba_protected_size() == 1 assert core.mamba_protected_size() == 1
core.dec_lock_ref(node, skipped.to_dec_params(), skip_swa=True) core.dec_lock_ref(node, holder.to_dec_params(), skip_swa=True)
assert core.mamba_protected_size() == 1
core.dec_lock_ref(node, owner.to_dec_params()) core.dec_lock_ref(node, owner.to_dec_params())
assert core.protected_size() == 0 assert core.protected_size() == 0
assert core.swa_protected_size() == 0 assert core.swa_protected_size() == 0
@@ -1674,10 +1682,9 @@ def test_mamba_eviction_walk_frees_slots_through_the_adapter():
assert torch.cat(device_frees[ComponentType.MAMBA]).tolist() == [7] assert torch.cat(device_frees[ComponentType.MAMBA]).tolist() == [7]
assert core.mamba_evictable_size() == 1 assert core.mamba_evictable_size() == 1
# A pre-eviction node handle locked after the tombstoning lands in the # A pre-eviction node handle still lock-round-trips: the segment lock
# skip map, and the replay keeps the release off it. # counts the tombstone and the paired release takes it back exactly.
lock = core.inc_lock_ref(internal) lock = core.inc_lock_ref(internal)
assert internal in lock.skip_lock_node_ids[ComponentType.MAMBA]
core.dec_lock_ref(internal, lock.to_dec_params()) core.dec_lock_ref(internal, lock.to_dec_params())
core.sanity_check([], []) core.sanity_check([], [])
@@ -1892,8 +1899,6 @@ def test_component_device_value_round_trips():
def test_lock_uuid_round_trips_through_dec_lock_ref(): def test_lock_uuid_round_trips_through_dec_lock_ref():
from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams
core = _swa_tree_core(window=2) core = _swa_tree_core(window=2)
first = _insert(core, [1, 2, 3], [10, 11, 12]) first = _insert(core, [1, 2, 3], [10, 11, 12])
# The window cap split the leaf: rebuild the in-window nodes' SWA values. # The window cap split the leaf: rebuild the in-window nodes' SWA values.
@@ -1912,10 +1917,7 @@ def test_lock_uuid_round_trips_through_dec_lock_ref():
assert core.swa_evictable_size() == 1 assert core.swa_evictable_size() == 1
core.dec_lock_ref( core.dec_lock_ref(
node, node,
DecLockRefParams( DecLockRefParams(swa_uuid_for_lock=result.swa_uuid_for_lock),
swa_uuid_for_lock=result.swa_uuid_for_lock,
skip_lock_node_ids=result.skip_lock_node_ids,
),
) )
# The uuid-bounded release returned the window to evictable. # The uuid-bounded release returned the window to evictable.
assert core.swa_protected_size() == 0 assert core.swa_protected_size() == 0
@@ -1925,32 +1927,27 @@ def test_lock_uuid_round_trips_through_dec_lock_ref():
assert again.swa_uuid_for_lock == result.swa_uuid_for_lock assert again.swa_uuid_for_lock == result.swa_uuid_for_lock
def test_swa_skip_map_crosses_the_binding_and_replays(): def test_swa_tombstones_cross_the_binding_and_release_balanced():
from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams
core = _swa_tree_core(window=8) core = _swa_tree_core(window=8)
_insert(core, [1, 2], [10, 11]) _insert(core, [1, 2], [10, 11])
second = _insert(core, [1, 2, 3, 4], [10, 11, 12, 13]) second = _insert(core, [1, 2, 3, 4], [10, 11, 12, 13])
leaf = second.cache_actions[-1].node_id leaf = second.cache_actions[-1].node_id
# Only the leaf carries SWA; its ancestor is recorded as a tombstone skip. # Only the leaf carries SWA; the ancestor tombstone is counted too, and
# the under-window walk reaches the root without stamping a uuid.
core.set_component_device_value( core.set_component_device_value(
leaf, ComponentType.SWA, torch.tensor([52, 53], dtype=torch.int64) leaf, ComponentType.SWA, torch.tensor([52, 53], dtype=torch.int64)
) )
result = core.inc_lock_ref(leaf) result = core.inc_lock_ref(leaf)
assert result.skip_lock_node_ids[ComponentType.SWA] assert result.swa_uuid_for_lock is None
assert core.swa_protected_size() == 2
core.dec_lock_ref( core.dec_lock_ref(
leaf, leaf,
DecLockRefParams( DecLockRefParams(swa_uuid_for_lock=result.swa_uuid_for_lock),
swa_uuid_for_lock=result.swa_uuid_for_lock,
skip_lock_node_ids=result.skip_lock_node_ids,
),
) )
assert core.swa_protected_size() == 0 assert core.swa_protected_size() == 0
def test_dec_swa_lock_only_frees_flow_after_the_full_release(): def test_dec_swa_lock_only_frees_flow_after_the_full_release():
from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams
core = _swa_tree_core(window=2) core = _swa_tree_core(window=2)
first = _insert(core, [1, 2], [10, 11]) first = _insert(core, [1, 2], [10, 11])
node = first.cache_actions[0].node_id node = first.cache_actions[0].node_id
@@ -1958,17 +1955,15 @@ def test_dec_swa_lock_only_frees_flow_after_the_full_release():
node, ComponentType.SWA, torch.tensor([50, 51], dtype=torch.int64) node, ComponentType.SWA, torch.tensor([50, 51], dtype=torch.int64)
) )
result = core.inc_lock_ref(node) result = core.inc_lock_ref(node)
# A non-None boundary: the window fills at the locked node itself.
assert result.swa_uuid_for_lock is not None
# The FULL lock releases first (skip_swa), then the early window release # The FULL lock releases first (skip_swa), then the early window release
# finds a fully unlocked device leaf and evicts it in place. # finds a fully unlocked device leaf and evicts it in place.
core.dec_lock_ref( core.dec_lock_ref(node, result.to_dec_params(), skip_swa=True)
node,
DecLockRefParams(skip_lock_node_ids=result.skip_lock_node_ids),
skip_swa=True,
)
device_frees: dict = {} device_frees: dict = {}
host_frees: dict = {} host_frees: dict = {}
_accumulate_step( _accumulate_step(
core.dec_swa_lock_only(node, result.swa_uuid_for_lock), core.dec_swa_lock_only(node, result.to_dec_params()),
{}, {},
device_frees, device_frees,
host_frees, host_frees,
@@ -1977,7 +1972,7 @@ def test_dec_swa_lock_only_frees_flow_after_the_full_release():
assert core.get_component_device_value(node, ComponentType.SWA) is None assert core.get_component_device_value(node, ComponentType.SWA) is None
def test_dec_swa_lock_only_returns_the_window_frees(): def test_dec_swa_lock_only_releases_once_and_a_repeat_dies_loud():
core = _swa_tree_core(window=2) core = _swa_tree_core(window=2)
first = _insert(core, [1, 2, 3], [10, 11, 12]) first = _insert(core, [1, 2, 3], [10, 11, 12])
for action in first.cache_actions: for action in first.cache_actions:
@@ -1991,22 +1986,19 @@ def test_dec_swa_lock_only_returns_the_window_frees():
device_frees: dict = {} device_frees: dict = {}
host_frees: dict = {} host_frees: dict = {}
_accumulate_step( _accumulate_step(
core.dec_swa_lock_only(node, result.swa_uuid_for_lock), core.dec_swa_lock_only(node, result.to_dec_params()),
{}, {},
device_frees, device_frees,
host_frees, host_frees,
) )
# The FULL lock still protects the path: the SWA release frees nothing and # The FULL lock still protects the path: the SWA release frees nothing
# the rebuilt values survive; a repeat release is a no-op. # and the rebuilt values survive.
assert device_frees == {} assert device_frees == {}
assert core.get_component_device_value(node, ComponentType.SWA) is not None assert core.get_component_device_value(node, ComponentType.SWA) is not None
_accumulate_step( # A repeat release of the same window is a protocol violation and dies
core.dec_swa_lock_only(node, result.swa_uuid_for_lock), # at the segment instead of silently walking it.
{}, with pytest.raises(BaseException, match="SWA window release hit lock_ref=0"):
device_frees, core.dec_swa_lock_only(node, result.to_dec_params())
host_frees,
)
assert device_frees == {}
def test_swa_rebuild_applies_through_the_python_allocator(): def test_swa_rebuild_applies_through_the_python_allocator():
@@ -2035,7 +2027,7 @@ def test_recover_with_locked_full_applies_through_the_python_allocator():
# The decode advanced past the window: the SWA lock releases early, then # The decode advanced past the window: the SWA lock releases early, then
# window eviction tombstones the SWA slot under the FULL lock (the state a # window eviction tombstones the SWA slot under the FULL lock (the state a
# locked-full overlap recovers from); its frees return to the allocator. # locked-full overlap recovers from); its frees return to the allocator.
cache.dec_swa_lock_only(node, lock.swa_uuid_for_lock) cache.dec_swa_lock_only(node, lock.to_dec_params())
tracker = {ComponentType.FULL: 0, ComponentType.SWA: 0} tracker = {ComponentType.FULL: 0, ComponentType.SWA: 0}
device_frees: dict = {} device_frees: dict = {}
host_frees: dict = {} host_frees: dict = {}
test/registered/unit/mem_cache/test_streaming_session_unit.py
+41 -10
View File
@@ -4,7 +4,7 @@ import torch
from sglang.srt.managers.schedule_batch import FINISH_ABORT, ReqKvInfo from sglang.srt.managers.schedule_batch import FINISH_ABORT, ReqKvInfo
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import MatchResult from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams, MatchResult
from sglang.srt.session.streaming_session import SessionSlot, StreamingSession from sglang.srt.session.streaming_session import SessionSlot, StreamingSession
from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.ci.ci_register import register_cpu_ci
@@ -54,6 +54,7 @@ class _FakeInnerCache:
self.match_results = list(match_results or []) self.match_results = list(match_results or [])
self.dec_lock_ref_calls = [] self.dec_lock_ref_calls = []
self.dec_lock_ref_params = [] self.dec_lock_ref_params = []
self.dec_lock_ref_skip_swa = []
def cache_finished_req(self, *args, **kwargs): def cache_finished_req(self, *args, **kwargs):
raise AssertionError("Streaming requests should not delegate to inner cache") raise AssertionError("Streaming requests should not delegate to inner cache")
@@ -66,6 +67,7 @@ class _FakeInnerCache:
def dec_lock_ref(self, node, *args, **kwargs): def dec_lock_ref(self, node, *args, **kwargs):
self.dec_lock_ref_calls.append(node) self.dec_lock_ref_calls.append(node)
self.dec_lock_ref_params.append(args[0] if args else kwargs.get("params")) self.dec_lock_ref_params.append(args[0] if args else kwargs.get("params"))
self.dec_lock_ref_skip_swa.append(kwargs.get("skip_swa", False))
def supports_mamba(self): def supports_mamba(self):
return False return False
@@ -97,9 +99,9 @@ class _FakeReq:
self.extra_key = None self.extra_key = None
self.cache_salt = None self.cache_salt = None
self.last_node = None self.last_node = None
self.swa_uuid_for_lock = None
self.skip_lock_node_ids = {}
self.swa_branching_seqlen = None self.swa_branching_seqlen = None
self.lock_receipt = DecLockRefParams()
self.swa_prefix_lock_released = False
self.to_finish = None self.to_finish = None
self.finished_reason = None self.finished_reason = None
self.finished_len = None self.finished_len = None
@@ -232,13 +234,11 @@ def test_nth_mid_abort_nukes_session_slot():
assert req.kv.req_pool_idx is None assert req.kv.req_pool_idx is None
def test_release_session_threads_mamba_skip_ids(): def test_release_session_threads_mamba_lock_receipt():
"""release_session must forward the slot's skip_lock_node_ids to """release_session must forward the slot's mamba lock receipt to
dec_lock_ref. The first req's last_node may be full-only-locked (mamba 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 not taken at inc), so without the receipt the release would drop a mamba
lock the session never took -- another request's, on a shared node.""" 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 = torch.arange(256, dtype=torch.int32).reshape(2, 128)
req_to_token_pool = _FakeReqToTokenPool(req_to_token) req_to_token_pool = _FakeReqToTokenPool(req_to_token)
allocator = _FakeAllocator() allocator = _FakeAllocator()
@@ -255,7 +255,6 @@ def test_release_session_threads_mamba_skip_ids():
cache_protected_len=0, cache_protected_len=0,
), ),
last_node=lock_node, last_node=lock_node,
skip_lock_node_ids={ComponentType.MAMBA: {42}},
) )
tree_cache.release_session("session-a") tree_cache.release_session("session-a")
@@ -263,7 +262,39 @@ def test_release_session_threads_mamba_skip_ids():
assert inner.dec_lock_ref_calls == [lock_node] assert inner.dec_lock_ref_calls == [lock_node]
params = inner.dec_lock_ref_params[0] params = inner.dec_lock_ref_params[0]
assert params is not None assert params is not None
assert params.skip_lock_node_ids.get(ComponentType.MAMBA) == {42} assert params.skipped_lock_components == ()
assert inner.dec_lock_ref_skip_swa == [False]
def test_release_session_skips_swa_after_early_release():
"""A slot saved from a req that early-released its SWA lock
(swa_prefix_lock_released) must release with skip_swa, or the session
close double-releases the SWA segment."""
req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
req_to_token_pool = _FakeReqToTokenPool(req_to_token)
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(
kv=ReqKvInfo(
req_pool_idx=0,
kv_committed_len=50,
kv_allocated_len=50,
swa_evicted_seqlen=0,
cache_protected_len=0,
),
last_node=lock_node,
lock_receipt=DecLockRefParams(node_id=42, swa_uuid_for_lock=7),
swa_prefix_lock_released=True,
)
tree_cache.release_session("session-a")
assert inner.dec_lock_ref_calls == [lock_node]
assert inner.dec_lock_ref_params[0].swa_uuid_for_lock == 7
assert inner.dec_lock_ref_skip_swa == [True]
def test_session_slot_does_not_restore_swa_branching_seqlen(): def test_session_slot_does_not_restore_swa_branching_seqlen():
test/registered/unit/mem_cache/test_swa_eviction_boundary.py
+2 -1
View File
@@ -20,6 +20,7 @@ import torch
from sglang.srt.managers.schedule_batch import ReqKvInfo, ScheduleBatch from sglang.srt.managers.schedule_batch import ReqKvInfo, ScheduleBatch
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams
from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.common import free_swa_out_of_window_slots from sglang.srt.mem_cache.common import free_swa_out_of_window_slots
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
@@ -111,7 +112,7 @@ def _make_req(req_pool_idx, token_ids, cache_protected_len, tree):
extra_key=None, extra_key=None,
cache_salt=None, cache_salt=None,
last_node=tree.root_node, last_node=tree.root_node,
swa_uuid_for_lock=None, lock_receipt=DecLockRefParams(),
swa_prefix_lock_released=False, swa_prefix_lock_released=False,
prefix_indices=torch.tensor([], dtype=torch.int64, device=tree.device), prefix_indices=torch.tensor([], dtype=torch.int64, device=tree.device),
_kv_committed_len=len(token_ids), _kv_committed_len=len(token_ids),
test/registered/unit/mem_cache/test_swa_lock_release_lifecycle.py
+8 -6
View File
@@ -158,7 +158,7 @@ class TestSWALockReleaseLifecycle(CustomTestCase):
self.assertFalse(leaf.swa_tombstone) self.assertFalse(leaf.swa_tombstone)
self.assertTrue(tree.swa_lru_list.in_list(leaf)) self.assertTrue(tree.swa_lru_list.in_list(leaf))
tree.dec_swa_lock_only(leaf, swa_uuid_for_lock=swa_uuid) tree.dec_swa_lock_only(leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid))
self.assertTrue(leaf.swa_tombstone) self.assertTrue(leaf.swa_tombstone)
self.assertFalse(tree.swa_lru_list.in_list(leaf)) self.assertFalse(tree.swa_lru_list.in_list(leaf))
@@ -196,7 +196,7 @@ class TestSWALockReleaseLifecycle(CustomTestCase):
swa_evictable_before = tree.swa_evictable_size_ swa_evictable_before = tree.swa_evictable_size_
swa_avail_before = allocator.swa_available_size() swa_avail_before = allocator.swa_available_size()
tree.dec_swa_lock_only(leaf_a, swa_uuid_for_lock=swa_uuid) tree.dec_swa_lock_only(leaf_a, DecLockRefParams(swa_uuid_for_lock=swa_uuid))
self.assertFalse(internal.swa_tombstone) self.assertFalse(internal.swa_tombstone)
self.assertTrue(tree.swa_lru_list.in_list(internal)) self.assertTrue(tree.swa_lru_list.in_list(internal))
@@ -221,7 +221,7 @@ class TestSWALockReleaseLifecycle(CustomTestCase):
inc_res = tree.inc_lock_ref(leaf) inc_res = tree.inc_lock_ref(leaf)
swa_uuid = inc_res.swa_uuid_for_lock swa_uuid = inc_res.swa_uuid_for_lock
tree.dec_swa_lock_only(leaf, swa_uuid_for_lock=swa_uuid) tree.dec_swa_lock_only(leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid))
self.assertTrue(leaf.swa_tombstone) self.assertTrue(leaf.swa_tombstone)
self.assertEqual(leaf.full_lock_ref, 1) self.assertEqual(leaf.full_lock_ref, 1)
@@ -308,7 +308,7 @@ class TestSWALockReleaseLifecycle(CustomTestCase):
inc_res = tree.inc_lock_ref(leaf) inc_res = tree.inc_lock_ref(leaf)
swa_uuid = inc_res.swa_uuid_for_lock swa_uuid = inc_res.swa_uuid_for_lock
tree.dec_swa_lock_only(leaf, swa_uuid_for_lock=swa_uuid) tree.dec_swa_lock_only(leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid))
self.assertTrue(leaf.swa_tombstone) self.assertTrue(leaf.swa_tombstone)
swa_evictable_before_delete = tree.swa_evictable_size_ swa_evictable_before_delete = tree.swa_evictable_size_
@@ -354,7 +354,9 @@ class TestSWALockReleaseLifecycle(CustomTestCase):
swa_avail_before = allocator.swa_available_size() swa_avail_before = allocator.swa_available_size()
full_avail_before = allocator.full_available_size() full_avail_before = allocator.full_available_size()
tree.dec_swa_lock_only(leaf, swa_uuid_for_lock=swa_uuid) tree.dec_swa_lock_only(
leaf, DecLockRefParams(swa_uuid_for_lock=swa_uuid)
)
self.assertTrue(leaf.swa_tombstone) self.assertTrue(leaf.swa_tombstone)
self.assertFalse(tree.swa_lru_list.in_list(leaf)) self.assertFalse(tree.swa_lru_list.in_list(leaf))
@@ -406,7 +408,7 @@ class TestSWALockReleaseLifecycle(CustomTestCase):
swa_evictable_before = tree.swa_evictable_size_ swa_evictable_before = tree.swa_evictable_size_
swa_avail_before = allocator.swa_available_size() swa_avail_before = allocator.swa_available_size()
tree.dec_swa_lock_only(leaf_a, swa_uuid_for_lock=swa_uuid) tree.dec_swa_lock_only(leaf_a, DecLockRefParams(swa_uuid_for_lock=swa_uuid))
# Leaf side: tombstoned and pages freed. # Leaf side: tombstoned and pages freed.
self.assertTrue(leaf_a.swa_tombstone) self.assertTrue(leaf_a.swa_tombstone)
test/registered/unit/mem_cache/test_swa_unittest.py
+3 -3
View File
@@ -794,7 +794,7 @@ class TestSWA(unittest.TestCase):
req.extra_key = None req.extra_key = None
req.cache_salt = None req.cache_salt = None
req.last_node = tree.root_node req.last_node = tree.root_node
req.swa_uuid_for_lock = None req.lock_receipt = DecLockRefParams()
req.kv.swa_evicted_seqlen = 0 req.kv.swa_evicted_seqlen = 0
req.kv.cache_protected_len = 1 req.kv.cache_protected_len = 1
# Intentionally mismatch to ensure code does not use len(prefix_indices). # Intentionally mismatch to ensure code does not use len(prefix_indices).
@@ -832,7 +832,7 @@ class TestSWA(unittest.TestCase):
req2.extra_key = None req2.extra_key = None
req2.cache_salt = None req2.cache_salt = None
req2.last_node = tree.root_node req2.last_node = tree.root_node
req2.swa_uuid_for_lock = None req2.lock_receipt = DecLockRefParams()
req2.kv.swa_evicted_seqlen = 0 req2.kv.swa_evicted_seqlen = 0
req2.kv.cache_protected_len = 1 req2.kv.cache_protected_len = 1
req2.prefix_indices = torch.tensor([21, 22, 23, 24, 25], device=tree.device) req2.prefix_indices = torch.tensor([21, 22, 23, 24, 25], device=tree.device)
@@ -1322,7 +1322,7 @@ class TestCacheUnfinishedReqEvictedPrefix(CustomTestCase):
req.cache_salt = None req.cache_salt = None
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
req.last_node = tree.root_node req.last_node = tree.root_node
req.swa_uuid_for_lock = None req.lock_receipt = DecLockRefParams()
req.prefix_indices = torch.empty(0, dtype=torch.int64, device=tree.device) req.prefix_indices = torch.empty(0, dtype=torch.int64, device=tree.device)
req.kv.swa_evicted_seqlen = evicted req.kv.swa_evicted_seqlen = evicted
test/registered/unit/mem_cache/test_unified_radix_cache_bench.py
+4 -7
View File
@@ -25,7 +25,6 @@ from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import ( from sglang.srt.mem_cache.base_prefix_cache import (
DecLockRefParams,
EvictParams, EvictParams,
InsertParams, InsertParams,
MatchPrefixParams, MatchPrefixParams,
@@ -588,7 +587,7 @@ def bench_lock_unlock(
lr = env.tree.inc_lock_ref(node) lr = env.tree.inc_lock_ref(node)
env.tree.dec_lock_ref( env.tree.dec_lock_ref(
node, node,
DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)), lr.to_dec_params(),
) )
warmup = min(20, num_pairs // 10) warmup = min(20, num_pairs // 10)
@@ -633,9 +632,7 @@ def bench_cache_finished(
if v is None: if v is None:
env.tree.dec_lock_ref( env.tree.dec_lock_ref(
node, node,
DecLockRefParams( lr.to_dec_params(),
swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)
),
) )
continue continue
kv_indices = torch.cat([mr.device_indices, v]) kv_indices = torch.cat([mr.device_indices, v])
@@ -652,8 +649,8 @@ def bench_cache_finished(
req.last_node = node req.last_node = node
req.kv.cache_protected_len = matched_len req.kv.cache_protected_len = matched_len
req.kv.kv_committed_len = len(seq) req.kv.kv_committed_len = len(seq)
if hasattr(lr, "swa_uuid_for_lock"): if hasattr(lr, "to_dec_params"):
req.swa_uuid_for_lock = lr.swa_uuid_for_lock req.lock_receipt = lr.to_dec_params()
env.rtp.req_to_token[req.kv.req_pool_idx, : len(kv_indices)] = kv_indices env.rtp.req_to_token[req.kv.req_pool_idx, : len(kv_indices)] = kv_indices
req_items.append(req) req_items.append(req)
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+580 -76
View File
@@ -9,6 +9,7 @@ import unittest
from array import array from array import array
from collections import defaultdict from collections import defaultdict
from dataclasses import dataclass, replace from dataclasses import dataclass, replace
from types import SimpleNamespace
from typing import Optional from typing import Optional
from unittest import mock from unittest import mock
@@ -25,7 +26,7 @@ from sglang.srt.disaggregation.kv_events import (
StorageMedium, StorageMedium,
) )
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.managers.schedule_batch import Req from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req, ReqKvInfo
from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import ( from sglang.srt.mem_cache.base_prefix_cache import (
@@ -98,6 +99,7 @@ from sglang.srt.server_args import (
ServerArgs, ServerArgs,
set_global_server_args_for_scheduler, set_global_server_args_for_scheduler,
) )
from sglang.srt.session.streaming_session import SessionSlot
from sglang.srt.utils import get_device from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
from sglang.test.test_utils import CustomTestCase from sglang.test.test_utils import CustomTestCase
@@ -1446,9 +1448,7 @@ class UnifiedRadixCacheSuite:
# Unlock -> should now be evictable # Unlock -> should now be evictable
cache.dec_lock_ref( cache.dec_lock_ref(
m.last_device_node, m.last_device_node,
DecLockRefParams( lock_result.to_dec_params(),
swa_uuid_for_lock=getattr(lock_result, "swa_uuid_for_lock", None)
),
) )
result = cache.evict(EvictParams(num_tokens=len(seq_a))) result = cache.evict(EvictParams(num_tokens=len(seq_a)))
self.assertGreaterEqual(result.num_tokens_evicted, len(seq_a)) self.assertGreaterEqual(result.num_tokens_evicted, len(seq_a))
@@ -1539,7 +1539,7 @@ class UnifiedRadixCacheSuite:
req.kv.kv_committed_len = kv_len req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
req.swa_uuid_for_lock = None req.lock_receipt = DecLockRefParams()
req.extra_key = None req.extra_key = None
req.full_untruncated_fill_ids = array("q", input_ids + output_ids) req.full_untruncated_fill_ids = array("q", input_ids + output_ids)
req.set_extend_range( req.set_extend_range(
@@ -1580,7 +1580,7 @@ class UnifiedRadixCacheSuite:
req.kv.kv_allocated_len = kv_len req.kv.kv_allocated_len = kv_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
req.swa_uuid_for_lock = None req.lock_receipt = DecLockRefParams()
req.extra_key = None req.extra_key = None
if self.cfg.has_mamba: if self.cfg.has_mamba:
req.kv.mamba_last_track_seqlen = kv_len req.kv.mamba_last_track_seqlen = kv_len
@@ -1624,7 +1624,7 @@ class UnifiedRadixCacheSuite:
req.kv.kv_committed_len = kv_len req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
req.swa_uuid_for_lock = None req.lock_receipt = DecLockRefParams()
req.swa_prefix_lock_released = True req.swa_prefix_lock_released = True
req.extra_key = None req.extra_key = None
req.full_untruncated_fill_ids = array("q", tokens) req.full_untruncated_fill_ids = array("q", tokens)
@@ -1659,7 +1659,7 @@ class UnifiedRadixCacheSuite:
req.kv.kv_committed_len = kv_len req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
req.swa_uuid_for_lock = None req.lock_receipt = DecLockRefParams()
req.extra_key = None req.extra_key = None
if self.cfg.has_mamba: if self.cfg.has_mamba:
req.kv.mamba_last_track_seqlen = kv_len req.kv.mamba_last_track_seqlen = kv_len
@@ -1673,7 +1673,7 @@ class UnifiedRadixCacheSuite:
cache.dec_lock_ref( cache.dec_lock_ref(
req.last_node, req.last_node,
DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)), req.lock_receipt,
) )
cache.sanity_check() cache.sanity_check()
@@ -1698,7 +1698,7 @@ class UnifiedRadixCacheSuite:
req.kv.kv_committed_len = len(tokens) req.kv.kv_committed_len = len(tokens)
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
req.swa_uuid_for_lock = None req.lock_receipt = DecLockRefParams()
req.extra_key = None req.extra_key = None
req.kv.swa_evicted_seqlen = evicted_len req.kv.swa_evicted_seqlen = evicted_len
@@ -1716,7 +1716,7 @@ class UnifiedRadixCacheSuite:
cache.dec_lock_ref( cache.dec_lock_ref(
req.last_node, req.last_node,
DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)), req.lock_receipt,
) )
cache.sanity_check() cache.sanity_check()
@@ -1800,7 +1800,7 @@ class UnifiedRadixCacheSuite:
req.kv.kv_committed_len = kv_len req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
req.swa_uuid_for_lock = None req.lock_receipt = DecLockRefParams()
req.extra_key = None req.extra_key = None
req.full_untruncated_fill_ids = array("q", input_ids) req.full_untruncated_fill_ids = array("q", input_ids)
req.set_extend_range( req.set_extend_range(
@@ -1925,7 +1925,7 @@ class UnifiedRadixCacheSuite:
req.kv.kv_committed_len = kv_len req.kv.kv_committed_len = kv_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
req.swa_uuid_for_lock = None req.lock_receipt = DecLockRefParams()
req.extra_key = None req.extra_key = None
req.kv.swa_evicted_seqlen = 0 req.kv.swa_evicted_seqlen = 0
@@ -1955,7 +1955,7 @@ class UnifiedRadixCacheSuite:
cache.dec_lock_ref( cache.dec_lock_ref(
req.last_node, req.last_node,
DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)), req.lock_receipt,
) )
cache.dec_lock_ref(last_device_node, lock_result.to_dec_params()) cache.dec_lock_ref(last_device_node, lock_result.to_dec_params())
cache.sanity_check() cache.sanity_check()
@@ -2021,7 +2021,7 @@ class UnifiedRadixCacheSuite:
1, 1,
"Mamba locked before release", "Mamba locked before release",
) )
cache.dec_swa_lock_only(node_a, lock_result.swa_uuid_for_lock) cache.dec_swa_lock_only(node_a, lock_result.to_dec_params())
self.assertEqual(_device_lock_ref(cache, node_a, ComponentType.SWA), 0) self.assertEqual(_device_lock_ref(cache, node_a, ComponentType.SWA), 0)
self.assertEqual( self.assertEqual(
_device_lock_ref(cache, node_a, ComponentType.MAMBA), _device_lock_ref(cache, node_a, ComponentType.MAMBA),
@@ -2071,7 +2071,58 @@ class UnifiedRadixCacheSuite:
) )
cache.sanity_check() cache.sanity_check()
cache.dec_lock_ref(node_a, DecLockRefParams(swa_uuid_for_lock=None)) cache.dec_lock_ref(
node_a, DecLockRefParams(swa_uuid_for_lock=None), skip_swa=True
)
cache.sanity_check()
def test_mamba_opt_out_holder_cannot_release_another_holders_mamba_lock(self):
"""Holder A takes the mamba lock; holder B opts out (skip_lock_components=(ComponentType.MAMBA,))
on the same node. B's early SWA release and final release must leave
A's mamba lock intact -- a lost/defaulted receipt on B's side used to
decrement A's lock without tripping any assert."""
if not self.cfg.has_swa or not self.cfg.has_mamba:
self.skipTest("requires SWA and Mamba components")
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(
1, (self.cfg.sliding_window_size // self.cfg.page_size) + 4
)
self._insert(cache, allocator, req_to_token_pool, seq)
node = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", seq)))
).last_device_node
self.assertIsNotNone(_device_value(cache, node, ComponentType.MAMBA))
lock_a = cache.inc_lock_ref(node)
lock_b = cache.inc_lock_ref(node, skip_lock_components=(ComponentType.MAMBA,))
self.assertNotIn(ComponentType.MAMBA, lock_a.skipped_lock_components)
self.assertIn(ComponentType.MAMBA, lock_b.skipped_lock_components)
self.assertEqual(
_device_lock_ref(cache, node, ComponentType.MAMBA), 1, "only A holds mamba"
)
# B: early SWA release, then final release -- both replay B's receipt.
cache.dec_swa_lock_only(node, lock_b.to_dec_params())
self.assertEqual(
_device_lock_ref(cache, node, ComponentType.MAMBA),
1,
"B's early release spares A",
)
cache.dec_lock_ref(node, lock_b.to_dec_params(), skip_swa=True)
self.assertEqual(
_device_lock_ref(cache, node, ComponentType.MAMBA),
1,
"B's final release spares A",
)
cache.dec_swa_lock_only(node, lock_a.to_dec_params())
self.assertEqual(
_device_lock_ref(cache, node, ComponentType.MAMBA),
0,
"A's release drops its own lock",
)
cache.dec_lock_ref(node, lock_a.to_dec_params(), skip_swa=True)
cache.sanity_check() cache.sanity_check()
def test_swa_early_release_drops_co_located_mamba_lock(self): def test_swa_early_release_drops_co_located_mamba_lock(self):
@@ -2106,11 +2157,9 @@ class UnifiedRadixCacheSuite:
# Early SWA release (decode advanced past the window), via the public # Early SWA release (decode advanced past the window), via the public
# path the scheduler calls. The leaf's SWA is tombstoned and the # path the scheduler calls. The leaf's SWA is tombstoned and the
# co-located lower-tier Mamba lock must drop in the same release. # co-located lower-tier Mamba lock must drop in the same release.
cache.dec_swa_lock_only(node_a, lock_result.swa_uuid_for_lock) cache.dec_swa_lock_only(node_a, lock_result.to_dec_params())
self.assertEqual( self.assertEqual(
_device_lock_ref(cache, node_a, ComponentType.SWA), _device_lock_ref(cache, node_a, ComponentType.SWA), 0, "SWA early-released"
0,
"SWA early-released",
) )
self.assertEqual( self.assertEqual(
_device_lock_ref(cache, node_a, ComponentType.MAMBA), _device_lock_ref(cache, node_a, ComponentType.MAMBA),
@@ -2140,14 +2189,11 @@ class UnifiedRadixCacheSuite:
skipped = cache.inc_lock_ref( skipped = cache.inc_lock_ref(
node_a, skip_lock_components=(ComponentType.MAMBA,) node_a, skip_lock_components=(ComponentType.MAMBA,)
) )
self.assertEqual(skipped.skip_lock_node_ids, {ComponentType.MAMBA: {node_a}}) self.assertNotIn(ComponentType.MAMBA, owner.skipped_lock_components)
self.assertIn(ComponentType.MAMBA, skipped.skipped_lock_components)
self.assertEqual(_device_lock_ref(cache, node_a, ComponentType.MAMBA), 1) self.assertEqual(_device_lock_ref(cache, node_a, ComponentType.MAMBA), 1)
cache.dec_swa_lock_only( cache.dec_swa_lock_only(node_a, skipped.to_dec_params())
node_a,
skipped.swa_uuid_for_lock,
skip_lock_node_ids=skipped.skip_lock_node_ids,
)
self.assertEqual( self.assertEqual(
_device_lock_ref(cache, node_a, ComponentType.MAMBA), _device_lock_ref(cache, node_a, ComponentType.MAMBA),
1, 1,
@@ -2312,7 +2358,7 @@ class UnifiedRadixCacheSuite:
self.assertGreaterEqual(_device_lock_ref(cache, node_a, ComponentType.MAMBA), 1) self.assertGreaterEqual(_device_lock_ref(cache, node_a, ComponentType.MAMBA), 1)
self.assertGreaterEqual(_device_lock_ref(cache, node_a, ComponentType.FULL), 1) self.assertGreaterEqual(_device_lock_ref(cache, node_a, ComponentType.FULL), 1)
cache.dec_swa_lock_only(node_a, lock_result.swa_uuid_for_lock) cache.dec_swa_lock_only(node_a, lock_result.to_dec_params())
self.assertEqual( self.assertEqual(
_device_lock_ref(cache, node_a, ComponentType.SWA), 0, "SWA released" _device_lock_ref(cache, node_a, ComponentType.SWA), 0, "SWA released"
) )
@@ -2386,7 +2432,7 @@ class UnifiedRadixCacheSuite:
cache.sanity_check() cache.sanity_check()
cache.dec_lock_ref( cache.dec_lock_ref(
leaf, leaf,
DecLockRefParams(swa_uuid_for_lock=lock_result.swa_uuid_for_lock), lock_result.to_dec_params(),
) )
cache.sanity_check() cache.sanity_check()
@@ -2548,7 +2594,7 @@ class UnifiedRadixCacheSuite:
req.kv.kv_committed_len = pre_len req.kv.kv_committed_len = pre_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
req.swa_uuid_for_lock = None req.lock_receipt = DecLockRefParams()
req.extra_key = None req.extra_key = None
swa_avail_before = allocator.swa_attn_allocator.available_size() swa_avail_before = allocator.swa_attn_allocator.available_size()
@@ -2575,7 +2621,7 @@ class UnifiedRadixCacheSuite:
cache.dec_lock_ref( cache.dec_lock_ref(
req.last_node, req.last_node,
DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)), req.lock_receipt,
) )
cache.sanity_check() cache.sanity_check()
@@ -2637,7 +2683,7 @@ class UnifiedRadixCacheSuite:
req.kv.kv_committed_len = pre_len req.kv.kv_committed_len = pre_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
req.swa_uuid_for_lock = None req.lock_receipt = DecLockRefParams()
req.extra_key = None req.extra_key = None
with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True): with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True):
@@ -2651,7 +2697,7 @@ class UnifiedRadixCacheSuite:
cache.dec_lock_ref( cache.dec_lock_ref(
req.last_node, req.last_node,
DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)), req.lock_receipt,
) )
cache.sanity_check() cache.sanity_check()
@@ -2715,18 +2761,26 @@ class UnifiedRadixCacheSuite:
self.assertIsNotNone(_device_value(cache, node, ComponentType.FULL)) self.assertIsNotNone(_device_value(cache, node, ComponentType.FULL))
self.assertIsNotNone(_device_value(cache, node, aux)) self.assertIsNotNone(_device_value(cache, node, aux))
lock_result = cache.inc_lock_ref(node) # Reach the "FULL locked, aux unlocked" state the way production does:
self.assertGreater(_device_lock_ref(cache, node, ComponentType.FULL), 0) # mamba via the decode-hold opt-out (skip_lock_components=(ComponentType.MAMBA,)), SWA via the
self.assertGreater(_device_lock_ref(cache, node, aux), 0) # early window release (its own first-class op).
aux_len = len(_device_value(cache, node, aux)) aux_len = len(_device_value(cache, node, aux))
cache.tree_core.set_component_protected_size( if aux == ComponentType.MAMBA:
aux, cache.tree_core.component_protected_size(aux) - aux_len lock_result = cache.inc_lock_ref(
node, skip_lock_components=(ComponentType.MAMBA,)
) )
cache.tree_core.set_component_evictable_size( else:
aux, cache.tree_core.component_evictable_size(aux) + aux_len lock_result = cache.inc_lock_ref(node)
self.assertGreater(_device_lock_ref(cache, node, aux), 0)
cache.dec_swa_lock_only(
node,
lock_result.to_dec_params(),
) )
cache.tree_core.set_component_device_lock_ref(node, aux, 0) # FULL still locked -> not a device leaf -> no inline evict; the
# value stays evictable for the explicit aux eviction below.
self.assertIsNotNone(_device_value(cache, node, aux))
self.assertGreater(_device_lock_ref(cache, node, ComponentType.FULL), 0)
self.assertEqual(_device_lock_ref(cache, node, aux), 0)
self.assertFalse(cache.tree_core.is_device_evictable_leaf(node)) self.assertFalse(cache.tree_core.is_device_evictable_leaf(node))
evict_params = EvictParams(num_tokens=0) evict_params = EvictParams(num_tokens=0)
@@ -2747,7 +2801,8 @@ class UnifiedRadixCacheSuite:
cache.dec_lock_ref( cache.dec_lock_ref(
node, node,
DecLockRefParams(swa_uuid_for_lock=lock_result.swa_uuid_for_lock), lock_result.to_dec_params(),
skip_swa=(aux == ComponentType.SWA),
) )
cache.sanity_check() cache.sanity_check()
@@ -2789,9 +2844,7 @@ class UnifiedRadixCacheSuite:
cache.dec_lock_ref( cache.dec_lock_ref(
m_base.last_device_node, m_base.last_device_node,
DecLockRefParams( lock_result.to_dec_params(),
swa_uuid_for_lock=getattr(lock_result, "swa_uuid_for_lock", None)
),
) )
# After unlock, base should be in evictable_device_leaves # After unlock, base should be in evictable_device_leaves
self.assertTrue( self.assertTrue(
@@ -2910,7 +2963,7 @@ class UnifiedRadixCacheSuite:
cache.dec_lock_ref( cache.dec_lock_ref(
m.last_device_node, m.last_device_node,
DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)), lr.to_dec_params(),
) )
cache.sanity_check() cache.sanity_check()
@@ -2973,7 +3026,7 @@ class UnifiedRadixCacheSuite:
cache.dec_lock_ref( cache.dec_lock_ref(
m.last_device_node, m.last_device_node,
DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)), lr.to_dec_params(),
) )
cache.sanity_check() cache.sanity_check()
@@ -5094,7 +5147,7 @@ class UnifiedRadixCacheSuite:
cache.dec_lock_ref( cache.dec_lock_ref(
m.last_device_node, m.last_device_node,
DecLockRefParams(swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)), lr.to_dec_params(),
) )
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", leaf)))) m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", leaf))))
self.assertGreaterEqual(len(m.device_indices), len(base)) self.assertGreaterEqual(len(m.device_indices), len(base))
@@ -6070,9 +6123,7 @@ class UnifiedRadixCacheSuite:
finally: finally:
cache.dec_lock_ref( cache.dec_lock_ref(
parent, parent,
DecLockRefParams( lock_result.to_dec_params(),
swa_uuid_for_lock=getattr(lock_result, "swa_uuid_for_lock", None)
),
) )
self.assertTrue(cache.tree_core.is_full_device_evicted(leaf)) self.assertTrue(cache.tree_core.is_full_device_evicted(leaf))
self.assertTrue(cache.tree_core.is_backuped(leaf)) self.assertTrue(cache.tree_core.is_backuped(leaf))
@@ -7050,7 +7101,7 @@ class UnifiedRadixCacheSuite:
) )
temp_lock = cache.inc_lock_ref(leaf) temp_lock = cache.inc_lock_ref(leaf)
self.assertEqual(_device_lock_ref(cache, tombstone, ComponentType.SWA), 0) self.assertEqual(_device_lock_ref(cache, tombstone, ComponentType.SWA), 1)
xfer = cache.tree_core.build_hicache_transfers( xfer = cache.tree_core.build_hicache_transfers(
ComponentType.SWA, leaf, CacheTransferPhase.LOAD_BACK ComponentType.SWA, leaf, CacheTransferPhase.LOAD_BACK
@@ -7069,7 +7120,7 @@ class UnifiedRadixCacheSuite:
load_back_lock = cache.inc_lock_ref(leaf) load_back_lock = cache.inc_lock_ref(leaf)
request_lock = cache.inc_lock_ref(leaf) request_lock = cache.inc_lock_ref(leaf)
self.assertEqual(_device_lock_ref(cache, tombstone, ComponentType.SWA), 2) self.assertEqual(_device_lock_ref(cache, tombstone, ComponentType.SWA), 3)
cache.dec_lock_ref(leaf, temp_lock.to_dec_params()) cache.dec_lock_ref(leaf, temp_lock.to_dec_params())
self.assertEqual(_device_lock_ref(cache, tombstone, ComponentType.SWA), 2) self.assertEqual(_device_lock_ref(cache, tombstone, ComponentType.SWA), 2)
@@ -7167,14 +7218,12 @@ class UnifiedRadixCacheSuite:
self._release_ongoing_load_back_locks(cache) self._release_ongoing_load_back_locks(cache)
cache.sanity_check() cache.sanity_check()
def test_hicache_full_temp_lock_skips_evicted_anchor_and_mirrors_on_release( def test_hicache_full_temp_lock_covers_evicted_anchor_and_mirrors_on_release(
self, self,
): ):
"""Acquire records the evicted anchor in skip_lock_node_ids (phase 1) """Segment locks count the evicted anchor too (no skip receipts), so
and locks device-on ancestors only (phase 2). After load_back a value restored mid-hold stays correctly attributed: each release
restores the anchor, a second acquire covers it; releasing the takes back exactly its own ref regardless of interleaved holders.
first must mirror the skip so the anchor's lock_ref is not
decremented twice.
""" """
if self._skip_unsupported_hicache_test(): if self._skip_unsupported_hicache_test():
return return
@@ -7189,25 +7238,37 @@ class UnifiedRadixCacheSuite:
self._simulate_backup_tree(cache) self._simulate_backup_tree(cache)
anchor_value = _device_value(cache, anchor, ComponentType.FULL) anchor_value = _device_value(cache, anchor, ComponentType.FULL)
# Simulate the anchor's FULL device eviction: drop the value and take
# its tokens out of the evictable ledger, as a real evict would.
cache.tree_core.set_component_device_value_raw(anchor, ComponentType.FULL, None) cache.tree_core.set_component_device_value_raw(anchor, ComponentType.FULL, None)
cache.tree_core.set_component_evictable_size(
ComponentType.FULL,
cache.tree_core.component_evictable_size(ComponentType.FULL)
- len(anchor_value),
)
self.assertEqual(_device_lock_ref(cache, anchor, ComponentType.FULL), 0) self.assertEqual(_device_lock_ref(cache, anchor, ComponentType.FULL), 0)
self.assertEqual(_device_lock_ref(cache, y, ComponentType.FULL), 0) self.assertEqual(_device_lock_ref(cache, y, ComponentType.FULL), 0)
self.assertEqual(_device_lock_ref(cache, a, ComponentType.FULL), 0) self.assertEqual(_device_lock_ref(cache, a, ComponentType.FULL), 0)
temp_lock = cache.inc_lock_ref(anchor) temp_lock = cache.inc_lock_ref(anchor)
self.assertEqual(_device_lock_ref(cache, anchor, ComponentType.FULL), 0) self.assertEqual(_device_lock_ref(cache, anchor, ComponentType.FULL), 1)
self.assertEqual(_device_lock_ref(cache, y, ComponentType.FULL), 1) self.assertEqual(_device_lock_ref(cache, y, ComponentType.FULL), 1)
self.assertEqual(_device_lock_ref(cache, a, ComponentType.FULL), 1) self.assertEqual(_device_lock_ref(cache, a, ComponentType.FULL), 1)
self.assertIn(ComponentType.FULL, temp_lock.skip_lock_node_ids)
self.assertIn(anchor, temp_lock.skip_lock_node_ids[ComponentType.FULL])
# Restore the value mid-hold: a value materialized under lock is
# protected until the last release, exactly as a load-back credits it.
cache.tree_core.set_component_device_value_raw( cache.tree_core.set_component_device_value_raw(
anchor, ComponentType.FULL, anchor_value anchor, ComponentType.FULL, anchor_value
) )
cache.tree_core.set_component_protected_size(
ComponentType.FULL,
cache.tree_core.component_protected_size(ComponentType.FULL)
+ len(anchor_value),
)
second_lock = cache.inc_lock_ref(anchor) second_lock = cache.inc_lock_ref(anchor)
self.assertEqual(_device_lock_ref(cache, anchor, ComponentType.FULL), 1) self.assertEqual(_device_lock_ref(cache, anchor, ComponentType.FULL), 2)
self.assertEqual(_device_lock_ref(cache, y, ComponentType.FULL), 2) self.assertEqual(_device_lock_ref(cache, y, ComponentType.FULL), 2)
self.assertEqual(_device_lock_ref(cache, a, ComponentType.FULL), 2) self.assertEqual(_device_lock_ref(cache, a, ComponentType.FULL), 2)
@@ -7249,7 +7310,7 @@ class UnifiedRadixCacheSuite:
) )
temp_lock = cache.inc_lock_ref(node) temp_lock = cache.inc_lock_ref(node)
self.assertEqual(_device_lock_ref(cache, node, ComponentType.MAMBA), 0) self.assertEqual(_device_lock_ref(cache, node, ComponentType.MAMBA), 1)
xfer = cache.tree_core.build_hicache_transfers( xfer = cache.tree_core.build_hicache_transfers(
ComponentType.MAMBA, node, CacheTransferPhase.LOAD_BACK ComponentType.MAMBA, node, CacheTransferPhase.LOAD_BACK
@@ -7266,7 +7327,7 @@ class UnifiedRadixCacheSuite:
load_back_lock = cache.inc_lock_ref(node) load_back_lock = cache.inc_lock_ref(node)
request_lock = cache.inc_lock_ref(node) request_lock = cache.inc_lock_ref(node)
self.assertEqual(_device_lock_ref(cache, node, ComponentType.MAMBA), 2) self.assertEqual(_device_lock_ref(cache, node, ComponentType.MAMBA), 3)
cache.dec_lock_ref(node, temp_lock.to_dec_params()) cache.dec_lock_ref(node, temp_lock.to_dec_params())
self.assertEqual(_device_lock_ref(cache, node, ComponentType.MAMBA), 2) self.assertEqual(_device_lock_ref(cache, node, ComponentType.MAMBA), 2)
@@ -7274,6 +7335,9 @@ class UnifiedRadixCacheSuite:
cache.dec_lock_ref(node, load_back_lock.to_dec_params()) cache.dec_lock_ref(node, load_back_lock.to_dec_params())
cache.dec_lock_ref(node, request_lock.to_dec_params()) cache.dec_lock_ref(node, request_lock.to_dec_params())
self.assertEqual(_device_lock_ref(cache, node, ComponentType.MAMBA), 0) self.assertEqual(_device_lock_ref(cache, node, ComponentType.MAMBA), 0)
# The commit ran under held locks: the restored value must have been
# credited to protected, or the ledger drifts on the final release.
cache.sanity_check()
def test_hicache_mixed_backup_evict_insert(self): def test_hicache_mixed_backup_evict_insert(self):
"""Complex scenario: backup some, evict, insert new, verify invariants.""" """Complex scenario: backup some, evict, insert new, verify invariants."""
@@ -7337,9 +7401,7 @@ class UnifiedRadixCacheSuite:
finally: finally:
cache.dec_lock_ref( cache.dec_lock_ref(
parent, parent,
DecLockRefParams( lr.to_dec_params(),
swa_uuid_for_lock=getattr(lr, "swa_uuid_for_lock", None)
),
) )
self.assertTrue( self.assertTrue(
@@ -7627,7 +7689,7 @@ class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase):
req.kv.kv_committed_len = len(tokens) req.kv.kv_committed_len = len(tokens)
req.kv.kv_allocated_len = len(tokens) req.kv.kv_allocated_len = len(tokens)
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
req.swa_uuid_for_lock = None req.lock_receipt = DecLockRefParams()
req.extra_key = None req.extra_key = None
req.kv.mamba_last_track_seqlen = len(tokens) req.kv.mamba_last_track_seqlen = len(tokens)
return req return req
@@ -8140,7 +8202,7 @@ class TestResumableInsertWalk(_InsertWalkSuite):
# Fill the host pool below len(top) free, keeping the on-path H-leaf # Fill the host pool below len(top) free, keeping the on-path H-leaf
# the oldest host entry and pinning the unbacked path root. # the oldest host entry and pinning the unbacked path root.
cache.inc_lock_ref(top) top_lock = cache.inc_lock_ref(top)
host_pool = cache.cache_controller.mem_pool_host host_pool = cache.cache_controller.mem_pool_host
start = 1000 start = 1000
top_len = _node_key_length(cache, top) top_len = _node_key_length(cache, top)
@@ -8158,7 +8220,7 @@ class TestResumableInsertWalk(_InsertWalkSuite):
cache.writing_check(write_back=True) cache.writing_check(write_back=True)
cache.evict(EvictParams(num_tokens=count)) cache.evict(EvictParams(num_tokens=count))
self.assertTrue(cache.tree_core.is_full_device_evicted(filler)) self.assertTrue(cache.tree_core.is_full_device_evicted(filler))
cache.dec_lock_ref(top) cache.dec_lock_ref(top, top_lock.to_dec_params())
# The crossing backup evicts exactly the on-path H-leaf, then the # The crossing backup evicts exactly the on-path H-leaf, then the
# remaining suffix is recreated as a fresh leaf. # remaining suffix is recreated as a fresh leaf.
@@ -8527,11 +8589,13 @@ class TestResumableInsertWalkSWA(_InsertWalkSuite):
lock_result = cache.inc_lock_ref(node) lock_result = cache.inc_lock_ref(node)
self.assertGreaterEqual(_device_lock_ref(cache, node, ComponentType.SWA), 1) self.assertGreaterEqual(_device_lock_ref(cache, node, ComponentType.SWA), 1)
cache.dec_swa_lock_only(node, lock_result.swa_uuid_for_lock) cache.dec_swa_lock_only(node, lock_result.to_dec_params())
self.assertEqual(_device_lock_ref(cache, node, ComponentType.SWA), 0) self.assertEqual(_device_lock_ref(cache, node, ComponentType.SWA), 0)
self.assertGreaterEqual(_device_lock_ref(cache, node, ComponentType.FULL), 1) self.assertGreaterEqual(_device_lock_ref(cache, node, ComponentType.FULL), 1)
cache.dec_lock_ref(node, DecLockRefParams(swa_uuid_for_lock=None)) cache.dec_lock_ref(
node, DecLockRefParams(swa_uuid_for_lock=None), skip_swa=True
)
cache.sanity_check() cache.sanity_check()
@@ -8661,7 +8725,7 @@ class TestReturnedValuesDrain(_InsertWalkSuite):
( (
"dec_swa_lock_only", "dec_swa_lock_only",
lambda: make(DecSwaLockOnlyResult), lambda: make(DecSwaLockOnlyResult),
lambda: cache.dec_swa_lock_only(node), lambda: cache.dec_swa_lock_only(node, DecLockRefParams()),
None, None,
), ),
] ]
@@ -9182,7 +9246,7 @@ class TestSWAWindowUnderBigramKey(CustomTestCase):
req.kv.kv_committed_len = seq_len req.kv.kv_committed_len = seq_len
req.last_node = cache.root_node_handle() req.last_node = cache.root_node_handle()
req.kv.cache_protected_len = 0 req.kv.cache_protected_len = 0
req.swa_uuid_for_lock = None req.lock_receipt = DecLockRefParams()
req.extra_key = None req.extra_key = None
with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True): with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True):
@@ -9204,7 +9268,7 @@ class TestSWAWindowUnderBigramKey(CustomTestCase):
cache.dec_lock_ref( cache.dec_lock_ref(
req.last_node, req.last_node,
DecLockRefParams(swa_uuid_for_lock=getattr(req, "swa_uuid_for_lock", None)), req.lock_receipt,
) )
cache.sanity_check() cache.sanity_check()
@@ -9402,5 +9466,445 @@ class TestAnchorLockOutcomePolicy(CustomTestCase):
cache.match_prefix.assert_called_once() cache.match_prefix.assert_called_once()
@unittest.skipUnless(torch.cuda.is_available(), "cache fixtures need CUDA")
class TestSegmentLockProtocol(_InsertWalkSuite):
"""Segment-lock protocol regressions, replaying the production lock-theft
failure classes (F1/F2) and the split hazards.
The protocol: a lock covers the contiguous node segment
[start, boundary-uuid], counting every node (tombstones included), so a
release needs only the receipt (anchor node, boundary uuid, skipped
components) and any ref==0 met inside the segment is a hard protocol
violation. The replays read the tree through the inspection interface, so
they run unchanged against the Python and Rust cores.
"""
cfg = CacheConfig(
components=(ComponentType.FULL, ComponentType.SWA), sliding_window_size=8
)
@staticmethod
def _swa_ref(cache, node_id):
return _device_lock_ref(cache, node_id, ComponentType.SWA)
@staticmethod
def _segment(cache, leaf_id, window):
"""Node ids from leaf up to the position-based window boundary."""
nodes, covered, cur = [], 0, leaf_id
while not cache.tree_core.is_root(cur) and covered < window:
nodes.append(cur)
covered += cache.tree_core.get_node_key_length(cur)
cur = _node_parent(cache, cur)
return nodes
def _match_leaf(self, cache, seq):
m = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
return m.last_device_node
@staticmethod
def _deepest(cache):
"""Structurally deepest node — bypasses SWA match validation, which
never adopts a holed window (simulates the stale-relock drift case)."""
node_id = cache.root_node_handle()
while True:
children = _node_children(cache, node_id)
if not children:
return node_id
node_id = children[0]
def _assert_protocol_violation(self, fn, fragment):
"""The Python core asserts; the Rust core panics (a BaseException
subclass at the PyO3 boundary). Either way the message names the
violation and the operation never completes silently."""
try:
fn()
except (KeyboardInterrupt, SystemExit):
raise
except BaseException as exc: # pyo3 PanicException derives from BaseException
self.assertIn(fragment, str(exc))
else:
self.fail(f"protocol violation went unreported: {fragment}")
def test_rebuilt_tombstone_relock_release_no_theft(self):
"""F1 attribution replay: A locks a window containing a tombstone; the
tombstone is rebuilt and locked by B mid-hold; A's release must leave
B's refs intact (the old skip-set protocol decremented B's lock)."""
sw = self.cfg.sliding_window_size
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, 2 * sw)
# SWA data only for the last sw//2 positions: the window has a hole.
# As in production, the evicted prefix's SWA slots are released before
# the insert so a later FULL free finds no live SWA peer.
swa_evicted = len(seq) - sw // 2
value = self._alloc(allocator, len(seq))
allocator.free_swa(value[:swa_evicted])
cache.insert(
InsertParams(
key=RadixKey(array("q", seq)),
value=value,
swa_evicted_seqlen=swa_evicted,
)
)
leaf = self._deepest(cache)
segment = self._segment(cache, leaf, sw)
self.assertTrue(
any(_device_value(cache, n, ComponentType.SWA) is None for n in segment),
"fixture must place a tombstone inside the window",
)
lock_a = cache.inc_lock_ref(leaf)
# Count-everything: every segment node carries A's ref, tombstones
# included, and the boundary uuid is always stamped.
self.assertIsNotNone(lock_a.swa_uuid_for_lock)
self.assertEqual(lock_a.node_id, leaf)
for n in segment:
self.assertEqual(self._swa_ref(cache, n), 1)
cache.sanity_check()
# Rebuild the tombstones under A's lock (Recover path: FULL is
# locked); the rebuilt values must be credited to protected.
cache.insert(
InsertParams(
key=RadixKey(array("q", seq)),
value=self._alloc(allocator, len(seq)),
swa_evicted_seqlen=0,
)
)
cache.sanity_check()
leaf = self._deepest(cache)
segment = self._segment(cache, leaf, sw)
for n in segment:
self.assertIsNotNone(_device_value(cache, n, ComponentType.SWA))
lock_b = cache.inc_lock_ref(leaf)
for n in segment:
self.assertEqual(self._swa_ref(cache, n), 2)
# THE regression: A's release takes back exactly A's refs.
cache.dec_lock_ref(leaf, lock_a.to_dec_params())
for n in segment:
self.assertEqual(self._swa_ref(cache, n), 1)
cache.sanity_check()
cache.dec_lock_ref(leaf, lock_b.to_dec_params())
for n in segment:
self.assertEqual(self._swa_ref(cache, n), 0)
cache.sanity_check()
def test_release_without_receipt_fails_loud(self):
"""A release missing its boundary uuid must die at the segment edge
(ref==0 assert) instead of silently walking to root stealing other
holders' locks — the F1 failure made loud."""
sw = self.cfg.sliding_window_size
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, 3 * sw)
self._insert(cache, allocator, req_to_token_pool, seq)
leaf = self._match_leaf(cache, seq)
lock = cache.inc_lock_ref(leaf)
self.assertIsNotNone(lock.swa_uuid_for_lock)
self._assert_protocol_violation(
lambda: cache.dec_lock_ref(leaf, DecLockRefParams(swa_uuid_for_lock=None)),
"lock_ref=0",
)
def test_release_on_another_node_fails_loud(self):
"""The receipt anchors the lock on the node it was taken on; replaying
it on a different node (the rematch-clobbered ``req.last_node`` class
of bug) must assert instead of walking that node's segment."""
sw = self.cfg.sliding_window_size
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, 2 * sw)
self._insert(cache, allocator, req_to_token_pool, seq)
leaf = self._match_leaf(cache, seq)
parent = _node_parent(cache, leaf)
self.assertFalse(cache.tree_core.is_root(parent))
lock = cache.inc_lock_ref(leaf)
self.assertEqual(lock.node_id, leaf)
self._assert_protocol_violation(
lambda: cache.dec_lock_ref(parent, lock.to_dec_params()),
"lock receipt anchored on node",
)
def test_double_release_fails_loud(self):
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, 2 * self.cfg.sliding_window_size)
self._insert(cache, allocator, req_to_token_pool, seq)
leaf = self._match_leaf(cache, seq)
lock = cache.inc_lock_ref(leaf)
cache.dec_lock_ref(leaf, lock.to_dec_params())
self._assert_protocol_violation(
lambda: cache.dec_lock_ref(leaf, lock.to_dec_params()), "lock_ref=0"
)
def test_finish_after_early_release_without_skip_swa_fails_loud(self):
"""F2 replay: retraction-after-early-release used to run a second SWA
walk that stole ancestors' locks; now it dies at the first node."""
sw = self.cfg.sliding_window_size
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, 2 * sw)
self._insert(cache, allocator, req_to_token_pool, seq)
leaf = self._match_leaf(cache, seq)
lock = cache.inc_lock_ref(leaf)
cache.dec_swa_lock_only(leaf, lock.to_dec_params())
self._assert_protocol_violation(
lambda: cache.dec_lock_ref(leaf, lock.to_dec_params()), "lock_ref=0"
)
def test_finish_after_early_release_with_skip_swa(self):
"""The correct F2 flow: skip_swa honors the early release."""
sw = self.cfg.sliding_window_size
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, 2 * sw)
self._insert(cache, allocator, req_to_token_pool, seq)
leaf = self._match_leaf(cache, seq)
lock = cache.inc_lock_ref(leaf)
cache.dec_swa_lock_only(leaf, lock.to_dec_params())
cache.dec_lock_ref(leaf, lock.to_dec_params(), skip_swa=True)
self.assertEqual(self._swa_ref(cache, leaf), 0)
self.assertEqual(_device_lock_ref(cache, leaf, ComponentType.FULL), 0)
cache.sanity_check()
def test_split_under_lock_releases_balanced(self):
"""A mid-segment split mints a new node with copied refs and migrates
the boundary uuid; the original receipt (its anchor stays on the
deeper half) still releases exactly."""
sw = self.cfg.sliding_window_size
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, 2 * sw)
self._insert(cache, allocator, req_to_token_pool, seq)
leaf = self._match_leaf(cache, seq)
lock = cache.inc_lock_ref(leaf)
pre_segment = self._segment(cache, leaf, sw)
# Diverge inside the window to force a split of a locked node.
fork = seq[: len(seq) - sw // 2] + self._make_seq(9000, sw)
self._insert(cache, allocator, req_to_token_pool, fork)
post_segment = self._segment(cache, leaf, sw)
self.assertGreater(len(post_segment), len(pre_segment))
for n in post_segment:
self.assertEqual(self._swa_ref(cache, n), 1)
cache.dec_lock_ref(leaf, lock.to_dec_params())
for n in post_segment:
self.assertEqual(self._swa_ref(cache, n), 0)
cache.sanity_check()
def test_aux_release_readmits_the_leaf_whatever_the_release_order(self):
"""Each component's release refreshes the leaf sets of the nodes it
unlocks, so a leaf whose last lock is an auxiliary one is readmitted
even when Full released first. Component-level replay of the Python
core; the Rust crate covers its own order in its unit tests."""
if _selected_tree_core_test_backend() != "python":
self.skipTest("drives Python component objects directly")
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
seq = self._make_seq(1, self.cfg.sliding_window_size)
self._insert(cache, allocator, req_to_token_pool, seq)
leaf = self._match_leaf(cache, seq)
node = cache.tree_core.node_by_id(leaf)
self.assertIn(node, cache.tree_core.evictable_device_leaves)
params = cache.inc_lock_ref(leaf).to_dec_params()
self.assertNotIn(node, cache.tree_core.evictable_device_leaves)
# Full first: its walk still sees the SWA lock, so the leaf stays out.
cache.components[ComponentType.FULL].release_component_lock(node, params)
self.assertNotIn(node, cache.tree_core.evictable_device_leaves)
# The SWA release drops the last lock and must readmit the leaf itself.
cache.components[ComponentType.SWA].release_component_lock(node, params)
self.assertIn(node, cache.tree_core.evictable_device_leaves)
cache.sanity_check()
@unittest.skipUnless(torch.cuda.is_available(), "cache fixtures need CUDA")
class TestSegmentLockFuzz(_InsertWalkSuite):
"""Random lock/insert/evict interleavings with the tree's own ledger
recomputation (sanity_check) as the per-step oracle, on both cores."""
cfg = CacheConfig(
components=(ComponentType.FULL, ComponentType.SWA),
sliding_window_size=8,
kv_size=4096,
max_num_reqs=256,
)
def _lock_skips(self, rng):
"""The decode hold opts the Mamba lock out; exercise both receipts."""
if self.cfg.has_mamba and rng.random() < 0.5:
return (ComponentType.MAMBA,)
return ()
def _run_seed(self, seed: int, steps: int = 120):
import random as _random
rng = _random.Random(seed)
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
chains: list[list[int]] = []
held: list[list] = [] # [node_id, receipt, released] entries
for step in range(steps):
op = rng.random()
try:
if op < 0.35 or not chains:
# Insert: fresh chain or extend/diverge an existing one.
if chains and rng.random() < 0.6:
base = rng.choice(chains)
cut = rng.randrange(1, len(base) + 1)
seq = base[:cut] + self._make_seq(
1000 * (step + 1), rng.randrange(2, 12)
)
else:
seq = self._make_seq(1000 * (step + 1), rng.randrange(4, 20))
if allocator.available_size() < len(seq):
cache.evict(EvictParams(num_tokens=len(seq) * 2))
if allocator.available_size() < len(seq):
continue
swa_evict = rng.randrange(0, len(seq)) if rng.random() < 0.3 else 0
value = self._alloc(allocator, len(seq))
# Release the evicted prefix's SWA peers first, as the
# scheduler does before inserting a window-trimmed request.
allocator.free_swa(value[:swa_evict])
params = InsertParams(
key=RadixKey(array("q", seq)),
value=value,
swa_evicted_seqlen=swa_evict,
)
if self.cfg.has_mamba:
req = self._make_req(req_to_token_pool)
params.mamba_value = req.kv.mamba_pool_idx.unsqueeze(0)
cache.insert(params)
chains.append(seq)
elif op < 0.6:
# Lock a random chain's current deepest device node.
seq = rng.choice(chains)
m = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", seq)))
)
node_id = m.last_device_node
if cache.tree_core.is_root(node_id):
continue
receipt = cache.inc_lock_ref(
node_id, skip_lock_components=self._lock_skips(rng)
)
held.append([node_id, receipt, False])
elif op < 0.8 and held:
# Full release of a random held lock.
idx = rng.randrange(len(held))
node_id, receipt, released = held.pop(idx)
cache.dec_lock_ref(
node_id, receipt.to_dec_params(), skip_swa=released
)
elif op < 0.9 and held:
# Early SWA release of a random not-yet-released lock.
idx = rng.randrange(len(held))
node_id, receipt, released = held[idx]
if released or receipt.swa_uuid_for_lock is None:
continue
cache.dec_swa_lock_only(
node_id,
receipt.to_dec_params(),
)
held[idx][2] = True
else:
cache.evict(
EvictParams(
num_tokens=rng.randrange(0, 32),
swa_num_tokens=rng.randrange(0, 32),
mamba_num=rng.randrange(0, 4) if self.cfg.has_mamba else 0,
)
)
except AssertionError:
raise
cache.sanity_check()
# Drain remaining locks; the tree must come back exactly balanced.
for node_id, receipt, released in held:
cache.dec_lock_ref(node_id, receipt.to_dec_params(), skip_swa=released)
cache.sanity_check()
def test_fuzz_seed0(self):
self._run_seed(0)
def test_fuzz_seed1(self):
self._run_seed(1)
def test_fuzz_seed2(self):
self._run_seed(2)
@unittest.skipUnless(torch.cuda.is_available(), "cache fixtures need CUDA")
class TestSegmentLockFuzzWithMamba(TestSegmentLockFuzz):
"""The FULL+SWA+MAMBA (Inkling) shape: the Mamba opt-out receipt, the
lower-priority cascade on early SWA release, and Mamba evictions all join
the interleavings."""
cfg = CacheConfig(
components=(ComponentType.FULL, ComponentType.SWA, ComponentType.MAMBA),
sliding_window_size=8,
kv_size=4096,
max_num_reqs=256,
mamba_cache_size=512,
)
class TestStreamingSessionLockLifecycle(CustomTestCase):
"""A streaming session must persist swa_prefix_lock_released: closing or
aborting a session whose first turn early-released its SWA lock must not
release the SWA segment a second time."""
cfg = CacheConfig(
page_size=1,
components=(ComponentType.FULL, ComponentType.SWA),
sliding_window_size=4,
kv_size=64,
max_context_len=64,
)
def _lock_and_early_release(self, cache, allocator):
tokens = array("q", range(1, 9))
value = allocator.alloc(len(tokens))
cache.insert(InsertParams(key=RadixKey(tokens), value=value))
match = cache.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
node = match.last_device_node
lock = cache.inc_lock_ref(node)
self.assertIsNotNone(lock.swa_uuid_for_lock)
cache.dec_swa_lock_only(node, lock.to_dec_params())
return node, lock
def _streaming_req(self, node, lock, *, session):
# No KV row is held: the slot only carries the tree lock receipt.
kv = ReqKvInfo()
return SimpleNamespace(
kv=kv,
detach_kv=lambda: kv,
last_node=node,
lock_receipt=lock.to_dec_params(),
swa_prefix_lock_released=True,
session=session,
finished_reason=None,
)
def test_close_after_early_release_releases_swa_once(self):
cache, allocator, _ = build_fixture(self.cfg)
node, lock = self._lock_and_early_release(cache, allocator)
req = self._streaming_req(node, lock, session=None)
slot = SessionSlot()
cache.session.slots["s"] = slot
slot.save_from_req(req, is_first=True)
cache.session.release_session("s")
cache.sanity_check()
def test_first_req_mid_abort_after_early_release(self):
cache, allocator, pool = build_fixture(self.cfg)
node, lock = self._lock_and_early_release(cache, allocator)
session = SimpleNamespace(
session_id="s2", streaming=True, abort_req=lambda: None
)
req = self._streaming_req(node, lock, session=session)
req.finished_reason = FINISH_ABORT()
self.assertTrue(cache.session.try_cache_finished_req(req))
cache.sanity_check()
if __name__ == "__main__": if __name__ == "__main__":
unittest.main() unittest.main()