From 18cfa2a23d18e59abb887702626bb5dca8124105 Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Fri, 25 Sep 2026 13:30:46 -0700 Subject: [PATCH 1/3] Unify cache eviction cursors and lock receipts by component Keep request/session ownership in component-keyed cursor maps and distinguish missing lock receipts from root boundaries in Python and Rust caches. Co-authored-by: Lu Fang <30275821+houseroad@users.noreply.github.com> --- python/sglang/srt/disaggregation/decode.py | 5 +- python/sglang/srt/dllm/mixin/scheduler.py | 2 +- python/sglang/srt/managers/schedule_batch.py | 30 ++- .../scheduler_components/invariant_checker.py | 4 +- .../sglang/srt/mem_cache/base_prefix_cache.py | 67 +++-- .../srt/mem_cache/buffer_mode/pipeline.py | 4 +- python/sglang/srt/mem_cache/common.py | 19 +- .../srt/mem_cache/rust_tree_core/adapter.py | 22 +- .../unified_cache/components/README.md | 12 +- .../unified_cache/components/base.py | 3 +- .../mem_cache/unified_cache/components/swa.py | 39 +-- .../unified_cache/unified_cache_linker.py | 17 +- .../unified_cache/unified_tree_core.py | 13 +- .../sglang/srt/session/streaming_session.py | 12 +- rust/sglang-radix-tree/src/components/mod.rs | 2 +- rust/sglang-radix-tree/src/components/swa.rs | 21 +- rust/sglang-radix-tree/src/python_bindings.rs | 26 +- .../src/tests/components/mamba.rs | 18 +- .../src/tests/components/swa.rs | 251 +++++++----------- .../src/tests/unified_tree_core.rs | 28 +- .../src/unified_tree_core.rs | 65 +++-- .../dllm/test_gemma4_uniform_lifecycle.py | 1 - ...test_priority_scheduling_disaggregation.py | 1 - .../mem_cache/test_decode_radix_lock_ref.py | 19 +- .../unit/mem_cache/test_hisparse_allocator.py | 5 +- .../mem_cache/test_pure_swa_chunk_cache.py | 3 +- .../mem_cache/test_pure_swa_radix_cache.py | 7 +- .../test_rust_tree_core_integration.py | 98 +++++-- .../mem_cache/test_streaming_session_unit.py | 97 ++++--- .../unit/mem_cache/test_swa_unittest.py | 11 +- .../mem_cache/test_unified_cache_linker.py | 13 +- .../test_unified_radix_cache_unittest.py | 130 +++++---- 32 files changed, 596 insertions(+), 449 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index a7f85d5873a8..955e8c26f396 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -105,6 +105,7 @@ ReqToTokenPool, ) from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType from sglang.srt.observability.req_time_stats import ( set_schedule_time_batch, set_time_batch, @@ -2240,7 +2241,7 @@ def alloc_for_decode_prealloc_hisparse( ) swa_evicted_seqlen = fill_len - swa_tail_len assert swa_evicted_seqlen >= 0 and swa_evicted_seqlen % allocator.page_size == 0 - req.kv.swa_evicted_seqlen = swa_evicted_seqlen + req.kv.set_evicted_seqlen(ComponentType.SWA, swa_evicted_seqlen) else: kv_loc = allocator.alloc_logical_only( prefix_lens=prefix_lens, @@ -2312,7 +2313,7 @@ def alloc_for_decode_prealloc( swa_evicted_seqlen >= 0 and swa_evicted_seqlen % allocator.page_size == 0 ) - req.kv.swa_evicted_seqlen = swa_evicted_seqlen + req.kv.set_evicted_seqlen(ComponentType.SWA, swa_evicted_seqlen) else: kv_loc = allocator.alloc_extend( prefix_lens=torch.tensor( diff --git a/python/sglang/srt/dllm/mixin/scheduler.py b/python/sglang/srt/dllm/mixin/scheduler.py index 924d08a21b04..7de817bb321d 100644 --- a/python/sglang/srt/dllm/mixin/scheduler.py +++ b/python/sglang/srt/dllm/mixin/scheduler.py @@ -200,7 +200,7 @@ def _stash_dllm_context(self: Scheduler, req: Req) -> None: req.kv.kv_committed_len = context_len req.kv.kv_allocated_len = context_len - assert req.kv.swa_evicted_seqlen <= context_len + assert req.kv.max_evicted_seqlen <= context_len req.set_extend_range(context_len, context_len) self.stash_chunked_request(req) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 06bea0e6a83f..3394c7751e14 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -120,6 +120,7 @@ ) from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool from sglang.srt.mem_cache.radix_cache import RadixKey +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, @@ -936,9 +937,11 @@ class ReqKvInfo: kv_committed_len: int = 0 # KV content committed up to here, <= kv_allocated_len kv_allocated_len: int = 0 - # SWA slots in [swa_dead_lo(page_size), swa_evicted_seqlen) are already freed. + # SWA slots in [swa_dead_lo(page_size), get_evicted_seqlen(SWA)) are freed. swa_evict_floor: int = 0 # [0, here) never window-evicted (prefill-aware SWA) - swa_evicted_seqlen: int = 0 # SWA eviction cursor + component_evicted_seqlens: dict[ComponentType, int] = dataclasses.field( + default_factory=dict + ) # Host-side KV backup the request holds across a retraction (unified cache). retraction_backup: Optional[RetractionBackup] = None @@ -973,11 +976,25 @@ def holds_mamba(self) -> bool: @property def is_kv_released(self) -> bool: - return self.kv_allocated_len == 0 and self.swa_evicted_seqlen == 0 + return self.kv_allocated_len == 0 and self.max_evicted_seqlen == 0 + + @property + def max_evicted_seqlen(self) -> int: + return max(self.component_evicted_seqlens.values(), default=0) + + def get_evicted_seqlen(self, component_type: ComponentType) -> int: + return self.component_evicted_seqlens.get(component_type, 0) + + def set_evicted_seqlen(self, component_type: ComponentType, length: int) -> None: + self.component_evicted_seqlens[component_type] = length + + def clamp_evicted_seqlens(self, length: int) -> None: + for component, cursor in self.component_evicted_seqlens.items(): + self.component_evicted_seqlens[component] = min(cursor, length) def mark_kv_released(self) -> None: self.kv_allocated_len = 0 - self.swa_evicted_seqlen = 0 + self.component_evicted_seqlens.clear() class Req(ReqDllmMixin): @@ -3859,7 +3876,7 @@ def _swa_eviction_trigger(self, sliding_window_size: int, eviction_interval: int return lambda req: due_this_forward return lambda req: ( req.seqlen - 1 - sliding_window_size - >= req.kv.swa_evicted_seqlen + eviction_interval + >= req.kv.get_evicted_seqlen(ComponentType.SWA) + eviction_interval ) def maybe_evict_swa(self): @@ -3895,7 +3912,8 @@ def maybe_evict_swa(self): if ( release_leaf_lock and not req.swa_prefix_lock_released - and req.lock_receipt.swa_uuid_for_lock is not None + and req.lock_receipt.component_lock_uuids.get(ComponentType.SWA) + is not None and req.last_node is not None and req.decode_batch_idx >= sliding_window_size ): diff --git a/python/sglang/srt/managers/scheduler_components/invariant_checker.py b/python/sglang/srt/managers/scheduler_components/invariant_checker.py index 5f2a020012d4..9c902651c8d3 100644 --- a/python/sglang/srt/managers/scheduler_components/invariant_checker.py +++ b/python/sglang/srt/managers/scheduler_components/invariant_checker.py @@ -25,6 +25,7 @@ from sglang.srt.mem_cache.allocator.swa import is_swa_req_ring from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType from sglang.srt.observability.scheduler_stage_metrics import ( SCHEDULER_STAGE_SANITY_CHECK_CACHE, SchedulerStageMetricsRecorder, @@ -320,7 +321,8 @@ def _get_total_uncached_sizes( full_uncached += allocated_len - req.kv.cache_protected_len if self.is_hybrid_swa: swa_uncached += allocated_len - max( - req.kv.cache_protected_len, req.kv.swa_evicted_seqlen + req.kv.cache_protected_len, + req.kv.get_evicted_seqlen(ComponentType.SWA), ) if req.beam_group is not None: diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index def8d775b650..60f8fea13e79 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -87,10 +87,12 @@ class InsertParams: # SWA specific prev_prefix_len: int = 0 - swa_evicted_seqlen: int = 0 swa_branching_seqlen: Optional[int] = None # General + component_evicted_seqlens: dict[ComponentType, int] = dataclasses.field( + default_factory=dict, kw_only=True + ) chunked: bool = False priority: int = 0 session_id: Optional[str] = None @@ -101,6 +103,12 @@ class InsertParams: # off). See UnifiedTreeNode.rotation_base. rotation_base: Optional[int] = None + def get_evicted_seqlen(self, component_type: ComponentType) -> int: + return self.component_evicted_seqlens.get(component_type, 0) + + def set_evicted_seqlen(self, component_type: ComponentType, length: int) -> None: + self.component_evicted_seqlens[component_type] = length + @dataclasses.dataclass class InsertResult: @@ -162,24 +170,41 @@ class IncLockRefResult: """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. + receipt on that node only. A recorded UUID marks a segment boundary; + ``None`` means root, while an absent entry means no receipt. + ``skipped_lock_components`` records the components the acquire left + untaken, so the release leaves them untouched. """ delta: Optional[int] = None node_id: Optional[int] = None - swa_uuid_for_lock: Optional[int] = None - swa_uuid_for_host_lock: Optional[int] = None skipped_lock_components: tuple[ComponentType, ...] = () + component_lock_uuids: dict[ComponentType, Optional[int]] = dataclasses.field( + default_factory=dict + ) + component_host_lock_uuids: dict[ComponentType, Optional[int]] = dataclasses.field( + default_factory=dict + ) + + def set_lock_uuid( + self, + component_type: ComponentType, + uuid: Optional[int], + *, + lock_host: bool = False, + ) -> None: + uuids = ( + self.component_host_lock_uuids if lock_host else self.component_lock_uuids + ) + uuids[component_type] = uuid def to_dec_params(self) -> DecLockRefParams: """Convert to the corresponding DecLockRefParams for dec_lock_ref.""" return DecLockRefParams( node_id=self.node_id, - swa_uuid_for_lock=self.swa_uuid_for_lock, - swa_uuid_for_host_lock=self.swa_uuid_for_host_lock, skipped_lock_components=tuple(self.skipped_lock_components), + component_lock_uuids=dict(self.component_lock_uuids), + component_host_lock_uuids=dict(self.component_host_lock_uuids), ) @@ -187,16 +212,28 @@ def to_dec_params(self) -> DecLockRefParams: class DecLockRefParams: """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). + A segment release requires its component's boundary entry; a missing + entry must not be treated as a lock reaching the root. ``node_id`` is + ``None`` only for receipts that never came from a unified-tree acquire + (legacy caches, session sentinels). """ node_id: Optional[int] = None - swa_uuid_for_lock: Optional[int] = None - swa_uuid_for_host_lock: Optional[int] = None skipped_lock_components: tuple[ComponentType, ...] = () + component_lock_uuids: dict[ComponentType, Optional[int]] = dataclasses.field( + default_factory=dict + ) + component_host_lock_uuids: dict[ComponentType, Optional[int]] = dataclasses.field( + default_factory=dict + ) + + def get_lock_uuid( + self, component_type: ComponentType, *, lock_host: bool = False + ) -> Optional[int]: + uuids = ( + self.component_host_lock_uuids if lock_host else self.component_lock_uuids + ) + return uuids[component_type] @dataclasses.dataclass @@ -442,7 +479,7 @@ def free_kv_row(self, kv: Any, ranges: list[tuple[int, int]]) -> None: free_kv_row_segments( allocator, [(row[start:end], start) for start, end in coalesce_ranges(ranges)], - swa_evicted_seqlen=kv.swa_evicted_seqlen, + swa_evicted_seqlen=kv.get_evicted_seqlen(ComponentType.SWA), swa_dead_lo=kv.swa_dead_lo(allocator.page_size), ) diff --git a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py index 9066c28dacce..fc672c8672cc 100644 --- a/python/sglang/srt/mem_cache/buffer_mode/pipeline.py +++ b/python/sglang/srt/mem_cache/buffer_mode/pipeline.py @@ -1235,7 +1235,9 @@ def _defer_for_capacity(pool: str) -> None: key=key, value=torch.cat([req.prefix_indices, device_indices]), prev_prefix_len=splice_base, - swa_evicted_seqlen=(span_end - staged_swa) if staged_swa else 0, + component_evicted_seqlens={ + ComponentType.SWA: (span_end - staged_swa) if staged_swa else 0 + }, ) ) self.ongoing_buffer_load_back[load_back_id] = _OngoingBufferLoadBack( diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index a34f9d4d7c36..158506db3468 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -14,6 +14,7 @@ from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams from sglang.srt.mem_cache.hicache_storage import PoolTransfer from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType from sglang.srt.runtime_context import get_serving, get_spec from sglang.srt.utils.common import ceil_align @@ -69,9 +70,10 @@ def free_swa_out_of_window_slots( assert req.kv.cache_protected_len % page_size == 0, ( "cache_protected_len must be page aligned" ) - req.kv.swa_evicted_seqlen = max( - req.kv.swa_evicted_seqlen, req.kv.swa_dead_lo(page_size) + swa_evicted_seqlen = max( + req.kv.get_evicted_seqlen(ComponentType.SWA), req.kv.swa_dead_lo(page_size) ) + req.kv.set_evicted_seqlen(ComponentType.SWA, swa_evicted_seqlen) if is_chunk_cache: # Chunk cache builds no radix tree, so no tombstone-leaf concern; evict @@ -89,22 +91,19 @@ def free_swa_out_of_window_slots( # retained checkpoint could never be matched and holding it is pure cost. evict_threshold = min(evict_threshold, retain_floor) - new_swa_evicted_seqlen = max( - req.kv.swa_evicted_seqlen, - evict_threshold, - ) + new_swa_evicted_seqlen = max(swa_evicted_seqlen, evict_threshold) if page_size > 1: new_swa_evicted_seqlen = (new_swa_evicted_seqlen // page_size) * page_size - if new_swa_evicted_seqlen > req.kv.swa_evicted_seqlen: + if new_swa_evicted_seqlen > swa_evicted_seqlen: free_slots = req_to_token_pool.req_to_token[ - req.kv.req_pool_idx, req.kv.swa_evicted_seqlen : new_swa_evicted_seqlen + req.kv.req_pool_idx, swa_evicted_seqlen:new_swa_evicted_seqlen ] token_to_kv_pool_allocator.free_swa_segment( - free_slots, start_pos=req.kv.swa_evicted_seqlen + free_slots, start_pos=swa_evicted_seqlen ) - req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen + req.kv.set_evicted_seqlen(ComponentType.SWA, new_swa_evicted_seqlen) def coalesce_ranges(ranges: list[tuple[int, int]]) -> list[tuple[int, int]]: diff --git a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py index d02b84e09da6..95b52960b6b7 100644 --- a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py +++ b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py @@ -146,22 +146,30 @@ def _inc_lock_ref_result_from_binding(result) -> IncLockRefResult: return IncLockRefResult( delta=result.delta, node_id=result.node_id, - swa_uuid_for_lock=result.swa_uuid_for_lock, - swa_uuid_for_host_lock=result.swa_uuid_for_host_lock, skipped_lock_components=tuple( ComponentType(ct) for ct in result.skipped_lock_components ), + component_lock_uuids={ + ComponentType(ct): uuid for ct, uuid in result.component_lock_uuids.items() + }, + component_host_lock_uuids={ + ComponentType(ct): uuid + for ct, uuid in result.component_host_lock_uuids.items() + }, ) 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).""" + """Use the owning module's type for both production and inspection bindings.""" 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], + component_lock_uuids={ + int(ct): uuid for ct, uuid in params.component_lock_uuids.items() + }, + component_host_lock_uuids={ + int(ct): uuid for ct, uuid in params.component_host_lock_uuids.items() + }, ) @@ -619,7 +627,7 @@ def begin_insert(self, params: InsertParams) -> InsertStepResult: session_id=params.session_id, mamba_value=params.mamba_value, prev_prefix_len=params.prev_prefix_len, - swa_evicted_seqlen=params.swa_evicted_seqlen, + swa_evicted_seqlen=params.get_evicted_seqlen(ComponentType.SWA), swa_branching_seqlen=params.swa_branching_seqlen, chunked=params.chunked, priority=0 if params.priority is None else params.priority, diff --git a/python/sglang/srt/mem_cache/unified_cache/components/README.md b/python/sglang/srt/mem_cache/unified_cache/components/README.md index 26f524f9e109..3144ebdd1da8 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/README.md +++ b/python/sglang/srt/mem_cache/unified_cache/components/README.md @@ -122,7 +122,7 @@ Insert a key-value pair into the tree. | Aspect | Detail | |--------|--------| | **Purpose** | Insert token sequence + KV indices, reusing existing prefix and freeing duplicate KV slots | -| **Inputs** | `params.key: RadixKey`, `params.value: Tensor` (KV pool indices), plus component-specific fields (`mamba_value`, `swa_evicted_seqlen`, `prev_prefix_len`) | +| **Inputs** | `params.key: RadixKey`, `params.value: Tensor` (KV pool indices), plus component-specific fields (`mamba_value`, `component_evicted_seqlens`, `prev_prefix_len`) | | **Output** | `InsertResult(prefix_len, mamba_exist)` — `prefix_len` = length of reused prefix | | **Mutation** | Creates new leaf nodes; updates component data on overlapping nodes; frees duplicate KV indices; may split nodes; updates LRU lists and evictable sizes | | **Complexity** | **O(K + D·C)** | @@ -132,7 +132,7 @@ Insert a key-value pair into the tree. 2. If key diverges mid-node, calls `_split_node` → `redistribute_on_node_split()` per component 3. For each overlapping node, calls `update_component_on_insert_overlap()` per component — returns `consumed_from` index; the tree frees `value[dup_start:consumed_from]` as duplicate pool indices - Full: returns `prefix_len` (no consumption, default behavior) - - SWA: checks if the overlapping node is a tombstone (SWA value = None) within the SWA window boundary (`swa_evicted_seqlen`): + - SWA: checks if the overlapping node is a tombstone (SWA value = None) within the SWA window boundary (`params.get_evicted_seqlen(ComponentType.SWA)`): - If entirely within window: **recovers tombstone** — frees old `full_value`, clones `value_slice`, translates to SWA indices, inserts into SWA LRU (returns `0` = all consumed) - If partially within window: **splits node** at boundary, recovers SWA on the window portion (returns `start_idx`) - If entirely outside window: returns `prefix_len` (no consumption) @@ -181,7 +181,7 @@ Lock a node to protect it (and its ancestors) from eviction. |--------|--------| | **Purpose** | Called when a request begins using a cached prefix — prevents eviction of nodes it depends on | | **Inputs** | `node` — the last matched node (deepest); `skip_lock_components` names components to leave untaken (the decode hold passes `(MAMBA,)`) | -| **Output** | `IncLockRefResult(node_id, swa_uuid_for_lock, skipped_lock_components)` — the receipt the matching release must replay: the anchor node, the SWA boundary, and the skipped set | +| **Output** | `IncLockRefResult(node_id, skipped_lock_components, component_lock_uuids)` — the receipt the matching release must replay: the anchor node, the skipped set, and each component's device boundary. Host acquires use `component_host_lock_uuids`. | | **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).| @@ -204,7 +204,7 @@ Unlock a previously locked node path by replaying the acquire's receipt. | Aspect | Detail | |--------|--------| | **Purpose** | Called when a request finishes — releases eviction protection | -| **Inputs** | `node`; 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. | +| **Inputs** | `node`; required `params` receipt (`node_id` anchor, `component_lock_uuids` boundaries, `skipped_lock_components`); `skip_swa=True` after an earlier `dec_swa_lock_only`. A receipt whose anchor is not `node` is a protocol violation (assert): a mispaired release would otherwise walk another holder's segment. | | **Output** | `DecLockRefResult()` | | **Mutation** | Decrements `lock_ref` per component 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` | @@ -241,7 +241,7 @@ Cache a completed request's KV data into the tree. | **Complexity** | **O(K + D·C)** — insert O(K + D·C) + lock release O(D). Simplifies to **O(K)**. | **Algorithm detail:** -1. `prepare_for_caching_req()` per component — sets component-specific insert params, returns effective cache length (SWA: sets `swa_evicted_seqlen`; Mamba: prepares `mamba_value` from ping-pong buffer, returns `mamba_last_track_seqlen` as truncation hint) +1. `prepare_for_caching_req()` per component — sets component-specific insert params, returns effective cache length (SWA: copies its cursor with `set_evicted_seqlen`; Mamba: prepares `mamba_value` from ping-pong buffer, returns `mamba_last_track_seqlen` as truncation hint) 2. Truncates if `effective_cache_len < len(token_ids)`: frees excess pool indices 3. Converts token IDs (bigram if EAGLE), page-aligns keys, then calls `insert()` 4. Frees unaligned tail KV indices beyond page boundary @@ -320,7 +320,7 @@ Each component implements these hooks. See `base.py` for the ABC and docstrings. | Hook | Purpose | Called By | Default | |------|---------|-----------|----------| -| `prepare_for_caching_req()` | Prepare component-specific data before insert, fill fields in `InsertParams`, return effective cache length. Full: no-op. SWA: sets `swa_evicted_seqlen`. Mamba: prepares `mamba_value` from ping-pong buffer, returns `mamba_last_track_seqlen`. | `cache_finished/unfinished_req` | returns `None` | +| `prepare_for_caching_req()` | Prepare component-specific data before insert, fill fields in `InsertParams`, return effective cache length. Full: no-op. SWA: copies its cursor with `set_evicted_seqlen`. Mamba: prepares `mamba_value` from ping-pong buffer, returns `mamba_last_track_seqlen`. | `cache_finished/unfinished_req` | returns `None` | | `cleanup_after_caching_req()` | Post-cache cleanup. Full/SWA: no-op. Mamba: frees forked `mamba_value` based on `mamba_exist`, handles ping-pong buffer `keep_idx`, resets `mamba_last_track_seqlen` on unfinished. | `cache_finished/unfinished_req` | no-op | ### Utility diff --git a/python/sglang/srt/mem_cache/unified_cache/components/base.py b/python/sglang/srt/mem_cache/unified_cache/components/base.py index 571f564600b7..d29099d4244b 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/base.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/base.py @@ -615,7 +615,8 @@ def prepare_for_caching_req( Return None for no truncation opinion (use full length); return int >= 0 for effective cache length. - Full: no-op, returns None. - - SWA: sets insert_params.swa_evicted_seqlen on finished; returns None. + - SWA: copies its eviction cursor into insert_params for finished and + unfinished requests; may return a branching boundary. - Mamba: prepares mamba_value (finished from ping-pong buffer, unfinished fork from req); returns mamba_last_track_seqlen.""" return None diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa.py b/python/sglang/srt/mem_cache/unified_cache/components/swa.py index 054fb9422341..6e5709db459d 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa.py @@ -403,7 +403,7 @@ def update_component_on_insert_overlap( return prefix_len full_cd = node.component_data[BASE_COMPONENT_TYPE] - swa_evicted_seqlen = params.swa_evicted_seqlen + swa_evicted_seqlen = params.get_evicted_seqlen(self.component_type) # A locked tombstone is legal (segment locks count every node); the # full-value swap below is safe because full lock_ref >= swa # lock_ref, so a locked-SWA node always takes the Recover branch. @@ -480,7 +480,7 @@ def recover_after_unevict( ct = self.component_type if node.component_data[ct].value is not None: return - swa_evicted_seqlen = params.swa_evicted_seqlen + swa_evicted_seqlen = params.get_evicted_seqlen(self.component_type) assert swa_evicted_seqlen % self.tree_core.page_size == 0, ( f"{ct}: swa_evicted_seqlen must be page-aligned, {swa_evicted_seqlen=}" ) @@ -524,13 +524,14 @@ def commit_insert_component_data( node_start = result.prefix_len node_end = node_start + len(node.key) - split_pos = params.swa_evicted_seqlen - node_start + swa_evicted_seqlen = params.get_evicted_seqlen(self.component_type) + split_pos = swa_evicted_seqlen - node_start if split_pos >= len(node.key): # Entire leaf is outside the SWA window — left as a tombstone. return result.record_adopted_range( self.component_type, - max(node_start, params.swa_evicted_seqlen), + max(node_start, swa_evicted_seqlen), node_end, ) if split_pos > 0: @@ -864,10 +865,7 @@ def acquire_component_lock( swa_uuid = comp.metadata[uuid_key] cur = cur.parent - if lock_host: - result.swa_uuid_for_host_lock = swa_uuid - else: - result.swa_uuid_for_lock = swa_uuid + result.set_lock_uuid(ct, swa_uuid, lock_host=lock_host) return result def release_component_lock( @@ -878,9 +876,9 @@ def release_component_lock( ) -> None: ct = self.component_type root = self.tree_core.root_node - swa_uuid_for_lock = ( - params.swa_uuid_for_host_lock if lock_host else params.swa_uuid_for_lock - ) + if node is root: + return + swa_uuid_for_lock = params.get_lock_uuid(ct, lock_host=lock_host) dec_swa = True uuid_key = "host_uuid" if lock_host else "uuid" @@ -972,7 +970,9 @@ def prepare_for_caching_req( ) -> Optional[int]: # Unfinished requests can already have an SWA-evicted prefix; preserve # that boundary so insertion creates a tombstone instead of live SWA KV. - insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen + insert_params.set_evicted_seqlen( + self.component_type, req.kv.get_evicted_seqlen(self.component_type) + ) # A recurrent checkpoint must stay attached to its exact token prefix. # Let MambaComponent select the insertion length for hybrid caches. @@ -1009,7 +1009,9 @@ def free_out_of_window_slots( self, req: Req, pre_len: int, insert_params: InsertParams ) -> None: self._free_out_of_window_slots(req, pre_len) - insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen + insert_params.set_evicted_seqlen( + self.component_type, req.kv.get_evicted_seqlen(self.component_type) + ) def cleanup_after_caching_req( self, @@ -1258,12 +1260,11 @@ def update_external_linker_load( if req.kv is None: from sglang.srt.managers.schedule_batch import ReqKvInfo - req.kv = ReqKvInfo( - kv_allocated_len=prefix_len, - swa_evicted_seqlen=boundary, - ) - else: - req.kv.swa_evicted_seqlen = max(req.kv.swa_evicted_seqlen, boundary) + req.kv = ReqKvInfo(kv_allocated_len=prefix_len) + req.kv.set_evicted_seqlen( + self.component_type, + max(req.kv.get_evicted_seqlen(self.component_type), boundary), + ) return transfer assert phase == ExternalLinkerLoadPhase.COMMIT diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_cache_linker.py b/python/sglang/srt/mem_cache/unified_cache/unified_cache_linker.py index ff2acee154f5..9d322580fd8b 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_cache_linker.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_cache_linker.py @@ -330,12 +330,11 @@ def load_back(self, req: Req) -> tuple[torch.Tensor, NodeId]: if req.kv is None: from sglang.srt.managers.schedule_batch import ReqKvInfo - req.kv = ReqKvInfo( - kv_allocated_len=prefix_len, - swa_evicted_seqlen=prefix_len, - ) - else: - req.kv.swa_evicted_seqlen = max(req.kv.swa_evicted_seqlen, prefix_len) + req.kv = ReqKvInfo(kv_allocated_len=prefix_len) + req.kv.set_evicted_seqlen( + ComponentType.SWA, + max(req.kv.get_evicted_seqlen(ComponentType.SWA), prefix_len), + ) # Insert the newly loaded tail into the tree. prefix_indices = torch.cat( @@ -359,8 +358,10 @@ def load_back(self, req: Req) -> tuple[torch.Tensor, NodeId]: else None ), prev_prefix_len=device_hit_len, - swa_evicted_seqlen=( - req.kv.swa_evicted_seqlen if req.kv is not None else 0 + component_evicted_seqlens=( + req.kv.component_evicted_seqlens.copy() + if req.kv is not None + else {} ), chunked=True, priority=req.priority or 0, diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 5ac4a88b217e..6cccf443391b 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -749,11 +749,16 @@ def dec_swa_lock_only( result = DecSwaLockOnlyResult() node = self.node_by_id(node_id) self._assert_receipt_anchor(node, params) + if node is self.root_node: + return result swa_component = self.components_by_type.get(ComponentType.SWA) if swa_component is None: return result swa_component.release_window_lock( - node, params.swa_uuid_for_lock, result.device_frees, result.host_frees + node, + params.get_lock_uuid(ComponentType.SWA), + result.device_frees, + result.host_frees, ) # Drop strictly-lower-priority locks co-located on the node, skipping @@ -1264,7 +1269,11 @@ def _insert_walk_step(self, state: _InsertWalkState) -> None: dup = value_slice[dup_start:consumed_from] abs_start = state.total_prefix_length + dup_start swa_already_freed = min( - max(state.params.swa_evicted_seqlen - abs_start, 0), dup.numel() + max( + state.params.get_evicted_seqlen(ComponentType.SWA) - abs_start, + 0, + ), + dup.numel(), ) if swa_already_freed > 0: step_actions.append(FreeDeviceKVFullOnly([dup[:swa_already_freed]])) diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index 1ef6fe28dc01..038744af33bc 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -20,6 +20,7 @@ MatchPrefixParams, MatchResult, ) +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType from sglang.srt.utils.common import ceil_align, is_npu if TYPE_CHECKING: @@ -241,7 +242,7 @@ def try_match_prefix(self, params: MatchPrefixParams) -> Optional[MatchResult]: # NPU requires page-aligned KV reuse; a rewind below the SWA eviction # cursor must also land on a page boundary -- free_kv_row_segments # splits dead/alive at the cursor, and a mid-page cut frees a page twice. - if self.page_size > 1 and (is_npu() or req.kv.swa_evicted_seqlen > prefix_len): + if self.page_size > 1 and (is_npu() or req.kv.max_evicted_seqlen > prefix_len): prefix_len = (prefix_len // self.page_size) * self.page_size req.kv.kv_committed_len = min(req.kv.kv_committed_len, prefix_len) @@ -458,7 +459,8 @@ def session_held_swa_tokens(self, active_pool_idxs: Optional[set] = None) -> int if slot.kv.holds_kv and not in_batch: allocated = ceil_align(slot.kv.kv_allocated_len, self.page_size) total += allocated - max( - slot.kv.cache_protected_len, slot.kv.swa_evicted_seqlen + slot.kv.cache_protected_len, + slot.kv.get_evicted_seqlen(ComponentType.SWA), ) return total @@ -517,7 +519,7 @@ def _free_tail(self, kv: ReqKvInfo, prefix_len: int) -> None: self._free_kv_aligned(kv, prefix_len, kv.kv_allocated_len) kv.kv_allocated_len = prefix_len kv.kv_committed_len = min(kv.kv_committed_len, prefix_len) - kv.swa_evicted_seqlen = min(kv.swa_evicted_seqlen, prefix_len) + kv.clamp_evicted_seqlens(prefix_len) def _trim_overshoot(self, req: Req, finished_len: int) -> None: """Trim slot KV to finished_len boundary. Spec v2 may overshoot @@ -526,14 +528,14 @@ def _trim_overshoot(self, req: Req, finished_len: int) -> None: be released to avoid token/KV mismatch. """ target = len(req.origin_input_ids) + finished_len - if self.page_size > 1 and req.kv.swa_evicted_seqlen > target: + if self.page_size > 1 and req.kv.max_evicted_seqlen > target: # Same hazard as the match-path rewind: the cursor must stay # page-aligned; the partial page is re-prefilled next turn. target = (target // self.page_size) * self.page_size self._free_kv_aligned(req.kv, target, req.kv.kv_allocated_len) req.kv.kv_allocated_len = min(req.kv.kv_allocated_len, target) req.kv.kv_committed_len = min(req.kv.kv_committed_len, target) - req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, target) + req.kv.clamp_evicted_seqlens(target) req.output_ids = req.output_ids[:finished_len] def _free_kv_aligned(self, kv: ReqKvInfo, target: int, end: int) -> None: diff --git a/rust/sglang-radix-tree/src/components/mod.rs b/rust/sglang-radix-tree/src/components/mod.rs index 2337903c45ae..1d8c165c8222 100644 --- a/rust/sglang-radix-tree/src/components/mod.rs +++ b/rust/sglang-radix-tree/src/components/mod.rs @@ -354,7 +354,7 @@ pub trait TreeComponent { &self, _tree_core: &mut UnifiedTreeCore, _node_id: NodeIdx_, - _swa_uuid_for_lock: Option, + _params: &DecLockRefParams, _device_frees: &mut HashMap>, _host_frees: &mut HashMap>, ) { diff --git a/rust/sglang-radix-tree/src/components/swa.rs b/rust/sglang-radix-tree/src/components/swa.rs index 49bb05f94908..acd7c19d9114 100644 --- a/rust/sglang-radix-tree/src/components/swa.rs +++ b/rust/sglang-radix-tree/src/components/swa.rs @@ -1232,11 +1232,7 @@ impl TreeComponent for SwaComponent { cur = parent; } - if lock_host { - result.swa_uuid_for_host_lock = swa_uuid; - } else { - result.swa_uuid_for_lock = swa_uuid; - } + result.set_lock_uuid(SWA.idx() as u8, swa_uuid, lock_host); result } @@ -1247,11 +1243,10 @@ impl TreeComponent for SwaComponent { params: &DecLockRefParams, lock_host: bool, ) { - let swa_uuid_for_lock = if lock_host { - params.swa_uuid_for_host_lock - } else { - params.swa_uuid_for_lock - }; + if tree_core.arena.node(node_id).is_root() { + return; + } + let swa_uuid_for_lock = params.get_lock_uuid(SWA.idx() as u8, lock_host); let mut cur = node_id; loop { @@ -1311,10 +1306,14 @@ impl TreeComponent for SwaComponent { &self, tree_core: &mut UnifiedTreeCore, node_id: NodeIdx_, - swa_uuid_for_lock: Option, + params: &DecLockRefParams, device_frees: &mut HashMap>, host_frees: &mut HashMap>, ) { + if tree_core.arena.node(node_id).is_root() { + return; + } + let swa_uuid_for_lock = params.get_lock_uuid(SWA.idx() as u8, false); let ct = SWA; let mut cur = node_id; loop { diff --git a/rust/sglang-radix-tree/src/python_bindings.rs b/rust/sglang-radix-tree/src/python_bindings.rs index 7a871b0ee905..3be75b21ebc1 100644 --- a/rust/sglang-radix-tree/src/python_bindings.rs +++ b/rust/sglang-radix-tree/src/python_bindings.rs @@ -677,26 +677,26 @@ impl InsertResultBinding { #[derive(Clone, Default)] pub struct DecLockRefParamsBinding { pub node_id: Option, - pub swa_uuid_for_lock: Option, - pub swa_uuid_for_host_lock: Option, pub skipped_lock_components: Vec, + pub component_lock_uuids: HashMap>, + pub component_host_lock_uuids: HashMap>, } #[pymethods] impl DecLockRefParamsBinding { #[new] - #[pyo3(signature = (node_id = None, swa_uuid_for_lock = None, swa_uuid_for_host_lock = None, skipped_lock_components = Vec::new()))] + #[pyo3(signature = (node_id = None, skipped_lock_components = Vec::new(), *, component_lock_uuids = None, component_host_lock_uuids = None))] fn new( node_id: Option, - swa_uuid_for_lock: Option, - swa_uuid_for_host_lock: Option, skipped_lock_components: Vec, + component_lock_uuids: Option>>, + component_host_lock_uuids: Option>>, ) -> Self { DecLockRefParamsBinding { node_id, - swa_uuid_for_lock, - swa_uuid_for_host_lock, skipped_lock_components, + component_lock_uuids: component_lock_uuids.unwrap_or_default(), + component_host_lock_uuids: component_host_lock_uuids.unwrap_or_default(), } } } @@ -706,9 +706,9 @@ impl DecLockRefParamsBinding { fn to_dec_lock_ref_params(&self) -> PyResult { Ok(DecLockRefParams { node_id: self.node_id, - swa_uuid_for_lock: self.swa_uuid_for_lock, - swa_uuid_for_host_lock: self.swa_uuid_for_host_lock, skipped_lock_components: component_set_from_py(&self.skipped_lock_components)?, + component_lock_uuids: self.component_lock_uuids.clone(), + component_host_lock_uuids: self.component_host_lock_uuids.clone(), }) } } @@ -719,9 +719,9 @@ impl DecLockRefParamsBinding { pub struct IncLockRefResultBinding { delta: Option, node_id: Option, - swa_uuid_for_lock: Option, - swa_uuid_for_host_lock: Option, skipped_lock_components: Vec, + component_lock_uuids: HashMap>, + component_host_lock_uuids: HashMap>, } impl IncLockRefResultBinding { @@ -729,8 +729,8 @@ impl IncLockRefResultBinding { Self { delta: result.delta, node_id: result.node_id, - swa_uuid_for_lock: result.swa_uuid_for_lock, - swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, + component_lock_uuids: result.component_lock_uuids, + component_host_lock_uuids: result.component_host_lock_uuids, skipped_lock_components: result .skipped_lock_components .iter() diff --git a/rust/sglang-radix-tree/src/tests/components/mamba.rs b/rust/sglang-radix-tree/src/tests/components/mamba.rs index 666f11d3b73b..6e59aebd3c86 100644 --- a/rust/sglang-radix-tree/src/tests/components/mamba.rs +++ b/rust/sglang-radix-tree/src/tests/components/mamba.rs @@ -340,11 +340,7 @@ fn lock_without_mamba_records_the_receipt_and_leaves_mamba_evictable() { assert_eq!(tc.arena.node(leaf).device_lock_ref(FULL), 1); // The receipt replays exactly what was taken: FULL only. - let params = DecLockRefParams { - swa_uuid_for_lock: result.swa_uuid_for_lock, - skipped_lock_components: result.skipped_lock_components, - ..Default::default() - }; + let params = result.to_dec_params(); tc.dec_lock_ref(leaf_handle, ¶ms, /* skip_swa = */ false) .expect("live test node"); assert_eq!(tc.arena.node(parent).device_lock_ref(FULL), 0); @@ -368,11 +364,7 @@ fn swa_only_release_spares_another_holders_mamba_lock() { 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 holder_params = holder.to_dec_params(); let mut device_frees = HashMap::new(); let mut host_frees = HashMap::new(); tc.dec_swa_lock_only( @@ -391,11 +383,7 @@ fn swa_only_release_spares_another_holders_mamba_lock() { tc.dec_lock_ref(leaf_handle, &holder_params, /* skip_swa = */ true) .expect("live test node"); assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 1); - let owner_params = DecLockRefParams { - swa_uuid_for_lock: owner.swa_uuid_for_lock, - skipped_lock_components: owner.skipped_lock_components, - ..Default::default() - }; + let owner_params = owner.to_dec_params(); tc.dec_lock_ref(leaf_handle, &owner_params, /* skip_swa = */ false) .expect("live test node"); assert_eq!(tc.arena.node(leaf).device_lock_ref(MAMBA), 0); diff --git a/rust/sglang-radix-tree/src/tests/components/swa.rs b/rust/sglang-radix-tree/src/tests/components/swa.rs index bfac6bc2f5e9..2c3738b7ae6d 100644 --- a/rust/sglang-radix-tree/src/tests/components/swa.rs +++ b/rust/sglang-radix-tree/src/tests/components/swa.rs @@ -1504,8 +1504,8 @@ fn acquire_lock_walks_until_the_window_fills_and_stamps_the_crossing_node() { assert_eq!(tc.arena.device_lock_ref(b, SWA), 1); assert_eq!(tc.arena.device_lock_ref(a, SWA), 0); // The walk fills the 2-atom window at b; b carries the first minted uuid. - assert_eq!(result.swa_uuid_for_lock, Some(2)); - assert_eq!(result.swa_uuid_for_host_lock, None); + assert_eq!(result.component_lock_uuids[&(SWA.idx() as u8)], Some(2)); + assert!(result.component_host_lock_uuids.is_empty()); assert_eq!(node_swa_uuid(&tc, b), Some(2)); assert_eq!(node_swa_uuid(&tc, c), None); assert_eq!(tc.swa_evictable_size(), 1); @@ -1557,7 +1557,7 @@ fn acquire_lock_overshooting_the_window_stops_at_the_crossing_node() { 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(a, SWA), 0); - assert_eq!(result.swa_uuid_for_lock, Some(2)); + assert_eq!(result.component_lock_uuids[&(SWA.idx() as u8)], Some(2)); assert_eq!(node_swa_uuid(&tc, b), Some(2)); assert_eq!(node_swa_uuid(&tc, a), None); assert_eq!(tc.swa_evictable_size(), 2); @@ -1584,7 +1584,10 @@ fn acquire_lock_reuses_the_stamped_uuid_and_shifts_sizes_once() { IncLockRefResult::default(), /* lock_host = */ false, ); - assert_eq!(second.swa_uuid_for_lock, first.swa_uuid_for_lock); + assert_eq!( + second.component_lock_uuids[&(SWA.idx() as u8)], + first.component_lock_uuids[&(SWA.idx() as u8)] + ); assert_eq!(tc.arena.device_lock_ref(c, SWA), 2); assert_eq!(tc.arena.device_lock_ref(b, SWA), 2); assert_eq!(tc.swa_evictable_size(), 1); @@ -1608,8 +1611,11 @@ fn acquire_lock_counts_tombstones_toward_the_window() { 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(a, SWA), 0); - assert!(result.swa_uuid_for_lock.is_some()); - assert_eq!(node_swa_uuid(&tc, b), result.swa_uuid_for_lock); + assert!(result.component_lock_uuids[&(SWA.idx() as u8)].is_some()); + assert_eq!( + node_swa_uuid(&tc, b), + result.component_lock_uuids[&(SWA.idx() as u8)] + ); } #[test] @@ -1624,7 +1630,7 @@ fn acquire_lock_under_the_window_reaches_the_root_without_a_uuid() { IncLockRefResult::default(), /* lock_host = */ false, ); - assert_eq!(result.swa_uuid_for_lock, None); + assert_eq!(result.component_lock_uuids[&(SWA.idx() as u8)], None); assert_eq!(tc.arena.device_lock_ref(a, SWA), 1); assert_eq!(tc.arena.device_lock_ref(b, SWA), 1); assert_eq!(tc.swa_evictable_size(), 0); @@ -1647,7 +1653,7 @@ fn inc_lock_ref_runs_full_and_swa_walks_together() { 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.component_lock_uuids[&(SWA.idx() as u8)].is_some()); 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(a, SWA), 0); @@ -1671,13 +1677,9 @@ fn inc_host_lock_ref_runs_full_and_swa_host_arms_together() { assert_eq!(tc.arena.host_lock_ref(c, SWA), 1); assert_eq!(tc.arena.host_lock_ref(b, SWA), 1); assert_eq!(tc.arena.host_lock_ref(a, SWA), 0); - assert!(result.swa_uuid_for_host_lock.is_some()); + assert!(result.component_host_lock_uuids[&(SWA.idx() as u8)].is_some()); // The release replays the acquire's uuid and unwinds both arms. - let params = DecLockRefParams { - swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, - skipped_lock_components: result.skipped_lock_components, - ..Default::default() - }; + let params = result.to_dec_params(); tc.dec_host_lock_ref(tc.arena.node(c).id, ¶ms) .expect("live test node"); assert_eq!(tc.arena.host_lock_ref(c, FULL), 0); @@ -1700,17 +1702,13 @@ fn dec_host_lock_ref_with_the_inner_uuid_leaves_an_outer_window_pinned() { .expect("live test node"); tc.inc_host_lock_ref(tc.arena.node(b).id) .expect("live test node"); - assert!(inner.swa_uuid_for_host_lock.is_some()); + assert!(inner.component_host_lock_uuids[&(SWA.idx() as u8)].is_some()); assert_eq!(tc.arena.host_lock_ref(c, SWA), 1); assert_eq!(tc.arena.host_lock_ref(b, SWA), 2); assert_eq!(tc.arena.host_lock_ref(a, SWA), 1); // Releasing the inner window with its own uuid stops at b; the outer // window's lock above the boundary survives. - let params = DecLockRefParams { - swa_uuid_for_host_lock: inner.swa_uuid_for_host_lock, - skipped_lock_components: inner.skipped_lock_components, - ..Default::default() - }; + let params = inner.to_dec_params(); tc.dec_host_lock_ref(tc.arena.node(c).id, ¶ms) .expect("live test node"); assert_eq!(tc.arena.host_lock_ref(c, SWA), 0); @@ -1736,8 +1734,11 @@ fn acquire_host_lock_walks_until_the_window_fills_and_stamps_the_host_uuid() { assert_eq!(tc.arena.host_lock_ref(b, SWA), 1); assert_eq!(tc.arena.host_lock_ref(a, SWA), 0); // The window fills at b; b carries the host uuid and leaves the host LRU. - assert_eq!(result.swa_uuid_for_host_lock, Some(2)); - assert_eq!(result.swa_uuid_for_lock, None); + assert_eq!( + result.component_host_lock_uuids[&(SWA.idx() as u8)], + Some(2) + ); + assert!(result.component_lock_uuids.is_empty()); assert_eq!(node_swa_host_uuid(&tc, b), Some(2)); assert_eq!(node_swa_host_uuid(&tc, c), None); assert!(!tc.host_lru_list(SWA).in_list(Some(c))); @@ -1769,7 +1770,10 @@ fn acquire_host_lock_reuses_the_stamped_uuid_and_skips_unlisted_nodes() { IncLockRefResult::default(), /* lock_host = */ true, ); - assert_eq!(second.swa_uuid_for_host_lock, first.swa_uuid_for_host_lock); + assert_eq!( + second.component_host_lock_uuids[&(SWA.idx() as u8)], + first.component_host_lock_uuids[&(SWA.idx() as u8)] + ); assert_eq!(tc.arena.host_lock_ref(c, SWA), 2); assert_eq!(tc.arena.host_lock_ref(b, SWA), 2); assert_eq!(tc.arena.host_lock_ref(a, SWA), 0); @@ -1792,8 +1796,11 @@ fn acquire_host_lock_counts_host_tombstones_toward_the_window() { assert_eq!(tc.arena.host_lock_ref(c, SWA), 1); assert_eq!(tc.arena.host_lock_ref(b, SWA), 1); assert_eq!(tc.arena.host_lock_ref(a, SWA), 0); - assert_eq!(node_swa_host_uuid(&tc, b), result.swa_uuid_for_host_lock); - assert!(result.swa_uuid_for_host_lock.is_some()); + assert_eq!( + node_swa_host_uuid(&tc, b), + result.component_host_lock_uuids[&(SWA.idx() as u8)] + ); + assert!(result.component_host_lock_uuids[&(SWA.idx() as u8)].is_some()); } #[test] @@ -1808,7 +1815,7 @@ fn acquire_host_lock_under_the_window_reaches_the_root_without_a_uuid() { IncLockRefResult::default(), /* lock_host = */ true, ); - assert_eq!(result.swa_uuid_for_host_lock, None); + assert_eq!(result.component_host_lock_uuids[&(SWA.idx() as u8)], None); assert_eq!(tc.arena.host_lock_ref(a, SWA), 1); assert_eq!(tc.arena.host_lock_ref(b, SWA), 1); } @@ -1866,7 +1873,10 @@ fn acquire_host_lock_stamps_the_host_tier_uuid_field_only() { /* lock_host = */ true, ); // The boundary uuid lands on the host-tier field; the device field stays clear. - assert_eq!(result.swa_uuid_for_host_lock, Some(2)); + assert_eq!( + result.component_host_lock_uuids[&(SWA.idx() as u8)], + Some(2) + ); assert_eq!(node_swa_host_uuid(&tc, b), Some(2)); assert_eq!(node_swa_uuid(&tc, b), None); assert_eq!(node_swa_host_uuid(&tc, c), None); @@ -1895,8 +1905,8 @@ fn device_and_host_lock_walks_mint_independent_uuids() { IncLockRefResult::default(), /* lock_host = */ true, ); - assert_eq!(device.swa_uuid_for_lock, Some(2)); - assert_eq!(host.swa_uuid_for_host_lock, Some(3)); + assert_eq!(device.component_lock_uuids[&(SWA.idx() as u8)], Some(2)); + assert_eq!(host.component_host_lock_uuids[&(SWA.idx() as u8)], Some(3)); assert_eq!(node_swa_uuid(&tc, b), Some(2)); assert_eq!(node_swa_host_uuid(&tc, b), Some(3)); } @@ -2013,12 +2023,7 @@ fn release_lock_returns_the_window_to_evictable() { IncLockRefResult::default(), /* lock_host = */ false, ); - let params = DecLockRefParams { - node_id: None, - swa_uuid_for_lock: result.swa_uuid_for_lock, - swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, - skipped_lock_components: result.skipped_lock_components, - }; + let params = result.to_dec_params(); swa.release_component_lock(&mut tc, c, ¶ms, /* lock_host = */ false); assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); @@ -2047,12 +2052,7 @@ fn release_lock_keeps_sizes_while_other_locks_remain() { IncLockRefResult::default(), /* lock_host = */ false, ); - let params = DecLockRefParams { - node_id: None, - swa_uuid_for_lock: first.swa_uuid_for_lock, - swa_uuid_for_host_lock: first.swa_uuid_for_host_lock, - skipped_lock_components: first.skipped_lock_components, - }; + let params = first.to_dec_params(); swa.release_component_lock(&mut tc, c, ¶ms, /* lock_host = */ false); assert_eq!(tc.arena.device_lock_ref(c, SWA), 1); assert_eq!(tc.arena.device_lock_ref(b, SWA), 1); @@ -2075,12 +2075,7 @@ fn release_lock_replays_the_tombstone_skips() { ); // b gained a device value AFTER the acquire recorded it as a tombstone. store_swa_device(&mut tc, b); - let params = DecLockRefParams { - node_id: None, - swa_uuid_for_lock: result.swa_uuid_for_lock, - swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, - skipped_lock_components: result.skipped_lock_components, - }; + let params = result.to_dec_params(); swa.release_component_lock(&mut tc, c, ¶ms, /* lock_host = */ false); assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); @@ -2107,12 +2102,7 @@ fn release_lock_stops_at_the_window_uuid() { tc.arena .node_mut(a) .set_lock_ref_(ValueSlotIdx::device(SWA), 1); - let params = DecLockRefParams { - node_id: None, - swa_uuid_for_lock: result.swa_uuid_for_lock, - swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, - skipped_lock_components: result.skipped_lock_components, - }; + let params = result.to_dec_params(); swa.release_component_lock(&mut tc, c, ¶ms, /* lock_host = */ false); assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); assert_eq!(tc.arena.device_lock_ref(b, SWA), 0); @@ -2141,12 +2131,7 @@ fn release_host_lock_stops_at_the_host_uuid_boundary() { IncLockRefResult::default(), /* lock_host = */ true, ); - let params = DecLockRefParams { - node_id: None, - swa_uuid_for_lock: result.swa_uuid_for_lock, - swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, - skipped_lock_components: result.skipped_lock_components, - }; + let params = result.to_dec_params(); swa.release_component_lock(&mut tc, c, ¶ms, /* lock_host = */ true); assert_eq!(tc.arena.host_lock_ref(c, SWA), 0); assert_eq!(tc.arena.host_lock_ref(b, SWA), 0); @@ -2168,12 +2153,15 @@ fn release_lock_without_the_boundary_uuid_dies_at_the_segment_edge() { IncLockRefResult::default(), /* lock_host = */ false, ); - // A receipt-less release overshoots the boundary into unlocked territory + // An incorrect root boundary overshoots into unlocked territory // and dies there instead of silently stealing whatever it crosses. swa.release_component_lock( &mut tc, c, - &DecLockRefParams::default(), + &DecLockRefParams { + component_lock_uuids: HashMap::from([(SWA.idx() as u8, None)]), + ..Default::default() + }, /* lock_host = */ false, ); } @@ -2193,12 +2181,7 @@ fn release_host_lock_reparks_tombstoned_host_nodes() { IncLockRefResult::default(), /* lock_host = */ true, ); - let params = DecLockRefParams { - node_id: None, - swa_uuid_for_lock: result.swa_uuid_for_lock, - swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, - skipped_lock_components: result.skipped_lock_components, - }; + let params = result.to_dec_params(); swa.release_component_lock(&mut tc, c, ¶ms, /* lock_host = */ true); assert_eq!(tc.arena.host_lock_ref(c, SWA), 0); assert_eq!(tc.arena.host_lock_ref(b, SWA), 0); @@ -2217,12 +2200,7 @@ fn inc_then_dec_lock_ref_roundtrips_with_dec_params() { let result = tc .inc_lock_ref(tc.arena.node(c).id, ComponentSet::EMPTY) .expect("live test node"); - let params = DecLockRefParams { - node_id: None, - swa_uuid_for_lock: result.swa_uuid_for_lock, - swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, - skipped_lock_components: result.skipped_lock_components, - }; + let params = result.to_dec_params(); tc.dec_lock_ref(tc.arena.node(c).id, ¶ms, /* skip_swa = */ false) .expect("live test node"); assert_eq!(tc.arena.device_lock_ref(c, SWA), 0); @@ -2250,11 +2228,7 @@ fn dec_swa_lock_only_releases_swa_while_full_stays_locked() { let mut host_frees = HashMap::new(); tc.dec_swa_lock_only( tc.arena.node(c).id, - &DecLockRefParams { - swa_uuid_for_lock: result.swa_uuid_for_lock, - skipped_lock_components: ComponentSet::EMPTY, - ..Default::default() - }, + &result.to_dec_params(), &mut device_frees, &mut host_frees, ) @@ -2291,11 +2265,7 @@ fn dec_swa_lock_only_evicts_a_fully_unlocked_device_leaf() { let mut host_frees = HashMap::new(); tc.dec_swa_lock_only( tc.arena.node(c).id, - &DecLockRefParams { - swa_uuid_for_lock: result.swa_uuid_for_lock, - skipped_lock_components: ComponentSet::EMPTY, - ..Default::default() - }, + &result.to_dec_params(), &mut device_frees, &mut host_frees, ) @@ -2319,7 +2289,6 @@ fn dec_swa_lock_only_is_a_noop_without_the_swa_component() { tc.dec_swa_lock_only( tc.arena.node(root).id, &DecLockRefParams { - swa_uuid_for_lock: None, skipped_lock_components: ComponentSet::EMPTY, ..Default::default() }, @@ -2351,7 +2320,10 @@ fn release_window_lock_breaks_on_a_tombstone_carrying_the_uuid() { &swa, &mut tc, c, - Some(99), + &DecLockRefParams { + component_lock_uuids: HashMap::from([(SWA.idx() as u8, Some(99))]), + ..Default::default() + }, &mut device_frees, &mut host_frees, ); @@ -2371,7 +2343,7 @@ fn release_window_lock_panics_on_a_non_swa_component() { &FullComponent, &mut tc, root, - None, + &DecLockRefParams::default(), &mut device_frees, &mut host_frees, ); @@ -2399,12 +2371,7 @@ fn release_lock_skip_set_leaves_a_relocked_tombstone_credited() { IncLockRefResult::default(), /* lock_host = */ false, ); - let params = DecLockRefParams { - node_id: None, - swa_uuid_for_lock: first.swa_uuid_for_lock, - swa_uuid_for_host_lock: first.swa_uuid_for_host_lock, - skipped_lock_components: first.skipped_lock_components, - }; + let params = first.to_dec_params(); swa.release_component_lock(&mut tc, c, ¶ms, /* lock_host = */ false); assert_eq!(tc.arena.device_lock_ref(c, SWA), 1); assert_eq!(tc.arena.device_lock_ref(b, SWA), 1); @@ -2428,10 +2395,7 @@ fn double_release_with_one_receipt_dies_loud() { IncLockRefResult::default(), /* lock_host = */ false, ); - let params = DecLockRefParams { - swa_uuid_for_lock: result.swa_uuid_for_lock, - ..Default::default() - }; + let params = result.to_dec_params(); swa.release_component_lock(&mut tc, c, ¶ms, /* lock_host = */ false); // Consuming the same receipt twice dies at the first unlocked node. swa.release_component_lock(&mut tc, c, ¶ms, /* lock_host = */ false); @@ -2461,11 +2425,7 @@ fn dec_swa_lock_only_releases_the_window_exactly_once() { let mut host_frees = HashMap::new(); tc.dec_swa_lock_only( tc.arena.node(c).id, - &DecLockRefParams { - swa_uuid_for_lock: first.swa_uuid_for_lock, - skipped_lock_components: ComponentSet::EMPTY, - ..Default::default() - }, + &first.to_dec_params(), &mut device_frees, &mut host_frees, ) @@ -2477,11 +2437,7 @@ fn dec_swa_lock_only_releases_the_window_exactly_once() { assert_eq!(tc.swa_protected_size(), 2); tc.dec_swa_lock_only( tc.arena.node(c).id, - &DecLockRefParams { - swa_uuid_for_lock: first.swa_uuid_for_lock, - skipped_lock_components: ComponentSet::EMPTY, - ..Default::default() - }, + &first.to_dec_params(), &mut device_frees, &mut host_frees, ) @@ -2517,11 +2473,7 @@ fn dec_swa_lock_only_leaves_out_of_window_swa_locks_alone() { let mut host_frees = HashMap::new(); tc.dec_swa_lock_only( tc.arena.node(c).id, - &DecLockRefParams { - swa_uuid_for_lock: result.swa_uuid_for_lock, - skipped_lock_components: ComponentSet::EMPTY, - ..Default::default() - }, + &result.to_dec_params(), &mut device_frees, &mut host_frees, ) @@ -2552,7 +2504,16 @@ fn release_window_lock_without_the_uuid_dies_past_the_boundary() { let mut host_frees = HashMap::new(); // 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, + &DecLockRefParams { + component_lock_uuids: HashMap::from([(SWA.idx() as u8, None)]), + ..Default::default() + }, + &mut device_frees, + &mut host_frees, + ); } #[test] @@ -2562,7 +2523,7 @@ fn release_window_lock_passes_over_a_mid_chain_tombstone_without_a_uuid() { store_swa_device(&mut tc, a); store_swa_device(&mut tc, c); let swa = swa_component(100); - let _ = swa.acquire_component_lock( + let receipt = swa.acquire_component_lock( &mut tc, c, IncLockRefResult::default(), @@ -2572,7 +2533,13 @@ fn release_window_lock_passes_over_a_mid_chain_tombstone_without_a_uuid() { let mut host_frees = HashMap::new(); // 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, + &receipt.to_dec_params(), + &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); @@ -2595,12 +2562,7 @@ fn release_host_lock_does_not_repark_a_node_whose_host_value_was_taken() { // The host value moved out while the lock was held; the node has no // device value either, so the release has nothing to park. let _ = tc.arena.take_host_value(a, SWA); - let params = DecLockRefParams { - node_id: None, - swa_uuid_for_lock: result.swa_uuid_for_lock, - swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, - skipped_lock_components: result.skipped_lock_components, - }; + let params = result.to_dec_params(); swa.release_component_lock(&mut tc, a, ¶ms, /* lock_host = */ true); assert_eq!(tc.arena.host_lock_ref(a, SWA), 0); assert!(!tc.host_lru_list(SWA).in_list(Some(a))); @@ -2622,12 +2584,7 @@ fn release_host_lock_skips_reparking_device_valued_nodes() { IncLockRefResult::default(), /* lock_host = */ true, ); - let params = DecLockRefParams { - node_id: None, - swa_uuid_for_lock: result.swa_uuid_for_lock, - swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, - skipped_lock_components: result.skipped_lock_components, - }; + let params = result.to_dec_params(); swa.release_component_lock(&mut tc, c, ¶ms, /* lock_host = */ true); // Device-valued nodes never re-park in the host LRU on host release. assert!(!tc.host_lru_list(SWA).in_list(Some(c))); @@ -2651,12 +2608,7 @@ 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). tc.host_lru_list_mut(SWA).insert_mru(b); - let params = DecLockRefParams { - node_id: None, - swa_uuid_for_lock: result.swa_uuid_for_lock, - swa_uuid_for_host_lock: result.swa_uuid_for_host_lock, - skipped_lock_components: result.skipped_lock_components, - }; + let params = result.to_dec_params(); swa.release_component_lock(&mut tc, c, ¶ms, /* lock_host = */ true); assert!(tc.host_lru_list(SWA).in_list(Some(b))); assert!(tc.host_lru_list(SWA).in_list(Some(c))); @@ -4967,12 +4919,7 @@ fn deep_swa_tree_survives_backup_evict_and_load_back_rounds() { let lock = tc .inc_lock_ref(anchor, ComponentSet::EMPTY) .expect("live test node"); - let params = DecLockRefParams { - node_id: None, - swa_uuid_for_lock: lock.swa_uuid_for_lock, - swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock, - skipped_lock_components: lock.skipped_lock_components, - }; + let params = lock.to_dec_params(); tc.dec_lock_ref(anchor, ¶ms, /* skip_swa = */ false) .expect("live test node"); tc.finish_load_back(anchor).expect("live test node"); @@ -5083,12 +5030,7 @@ fn aux_release_refreshes_the_leaf_set_whatever_the_release_order() { .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, - }; + let params = result.to_dec_params(); // Full first: its walk still sees the SWA lock, so the leaf stays out. crate::components::FullComponent .release_component_lock(&mut tc, leaf, ¶ms, /* lock_host = */ false); @@ -5107,12 +5049,7 @@ fn dec_lock_ref_rejects_a_receipt_from_another_node() { 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, - }; + let params = result.to_dec_params(); // 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, ¶ms, /* skip_swa = */ false) @@ -5129,12 +5066,7 @@ fn dec_host_lock_ref_rejects_a_receipt_from_another_node() { 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, - }; + let params = result.to_dec_params(); tc.dec_host_lock_ref(tc.arena.node(a).id, ¶ms) .expect("live test node"); } @@ -5152,12 +5084,7 @@ fn receipt_anchor_follows_the_locked_node_through_a_split() { 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, - }; + let params = result.to_dec_params(); tc.dec_lock_ref(leaf_id, ¶ms, /* skip_swa = */ false) .expect("live test node"); assert_eq!( diff --git a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs index 8bd7700ff616..82ecc752e703 100644 --- a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs @@ -297,7 +297,7 @@ impl TreeComponent> for LowPriorityComponentForTest { lock_host: bool, ) { assert!(!lock_host); - assert!(params.swa_uuid_for_lock.is_some()); + assert!(params.component_lock_uuids[&(SWA.idx() as u8)].is_some()); panic!("low-priority release dispatched"); } } @@ -639,7 +639,7 @@ fn dec_swa_lock_only_dispatches_lower_priority_releases() { let _ = tc.dec_swa_lock_only( tc.arena.node(root).id, &DecLockRefParams { - swa_uuid_for_lock: Some(7), + component_lock_uuids: HashMap::from([(SWA.idx() as u8, Some(7))]), skipped_lock_components: ComponentSet::EMPTY, ..Default::default() }, @@ -683,11 +683,7 @@ fn dec_swa_lock_only_returns_device_frees_in_the_device_dict() { let mut host_frees = HashMap::new(); tc.dec_swa_lock_only( tc.arena.node(a).id, - &DecLockRefParams { - swa_uuid_for_lock: result.swa_uuid_for_lock, - skipped_lock_components: ComponentSet::EMPTY, - ..Default::default() - }, + &result.to_dec_params(), &mut device_frees, &mut host_frees, ) @@ -710,8 +706,8 @@ fn next_swa_uuid_counts_up_from_two() { #[test] fn inc_lock_ref_result_defaults_carry_no_uuids() { let result = IncLockRefResult::default(); - assert_eq!(result.swa_uuid_for_lock, None); - assert_eq!(result.swa_uuid_for_host_lock, None); + assert!(result.component_lock_uuids.is_empty()); + assert!(result.component_host_lock_uuids.is_empty()); } // A tree with the Swa stub registered and a node carrying an SWA device value. @@ -8527,12 +8523,7 @@ fn run_random_op_sequence(mut tc: UnifiedTreeCore>, page: usize, mamba: let lock = tc .inc_lock_ref(anchor, ComponentSet::EMPTY) .expect("live match anchor"); - let params = DecLockRefParams { - node_id: None, - swa_uuid_for_lock: lock.swa_uuid_for_lock, - swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock, - skipped_lock_components: lock.skipped_lock_components, - }; + let params = lock.to_dec_params(); tc.dec_lock_ref(anchor, ¶ms, /* skip_swa = */ false) .expect("live match anchor"); } @@ -8551,12 +8542,7 @@ fn run_random_op_sequence(mut tc: UnifiedTreeCore>, page: usize, mamba: &mut mamba_next, mamba, )); - let params = DecLockRefParams { - node_id: None, - swa_uuid_for_lock: lock.swa_uuid_for_lock, - swa_uuid_for_host_lock: lock.swa_uuid_for_host_lock, - skipped_lock_components: lock.skipped_lock_components, - }; + let params = lock.to_dec_params(); tc.dec_lock_ref(anchor, ¶ms, /* skip_swa = */ false) .expect("live match anchor"); } diff --git a/rust/sglang-radix-tree/src/unified_tree_core.rs b/rust/sglang-radix-tree/src/unified_tree_core.rs index 742f9c190260..74f9938a3651 100644 --- a/rust/sglang-radix-tree/src/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/unified_tree_core.rs @@ -35,42 +35,61 @@ fn next_coexist_reclaim_digest(current: i64, node_id: NodeId, component_idx: usi // ---- interface types ---- -/// 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. +/// Receipt returned by `inc_lock_ref` and replayed by `dec_lock_ref`. #[derive(Default)] pub struct IncLockRefResult { /// Tokens newly protected (moved out of evictable) by this lock. pub delta: Option, /// The node the lock was taken on; a release replays the receipt there only. pub node_id: Option, - /// SWA lock-window uuid minted/reused by the device lock walk. - pub swa_uuid_for_lock: Option, - /// SWA lock-window uuid minted/reused by the host lock walk. - pub swa_uuid_for_host_lock: Option, /// Components the acquire left untaken; the release skips them too. pub skipped_lock_components: ComponentSet, + /// A recorded None means the segment reaches the root; absence means no receipt. + pub component_lock_uuids: HashMap>, + pub component_host_lock_uuids: HashMap>, } -/// 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. +impl IncLockRefResult { + pub fn set_lock_uuid(&mut self, component: u8, uuid: Option, lock_host: bool) { + let uuids = if lock_host { + &mut self.component_host_lock_uuids + } else { + &mut self.component_lock_uuids + }; + uuids.insert(component, uuid); + } + + pub fn to_dec_params(&self) -> DecLockRefParams { + DecLockRefParams { + node_id: self.node_id, + skipped_lock_components: self.skipped_lock_components, + component_lock_uuids: self.component_lock_uuids.clone(), + component_host_lock_uuids: self.component_host_lock_uuids.clone(), + } + } +} + +/// Receipt required by `dec_lock_ref`. #[derive(Default)] 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, - /// SWA lock-window uuid the device unlock stops at, from the matching acquire. - pub swa_uuid_for_lock: Option, - /// SWA lock-window uuid the host unlock stops at, from the matching acquire. - pub swa_uuid_for_host_lock: Option, /// Components the matching acquire left untaken. pub skipped_lock_components: ComponentSet, + pub component_lock_uuids: HashMap>, + pub component_host_lock_uuids: HashMap>, +} + +impl DecLockRefParams { + pub fn get_lock_uuid(&self, component: u8, lock_host: bool) -> Option { + let uuids = if lock_host { + &self.component_host_lock_uuids + } else { + &self.component_lock_uuids + }; + uuids[&component] + } } /// Result of `dec_lock_ref`. @@ -962,13 +981,7 @@ impl UnifiedTreeCore { let Some(swa) = self.try_component_by_type_(SWA) else { return Ok(()); }; - swa.release_window_lock( - self, - node_idx, - params.swa_uuid_for_lock, - device_frees, - host_frees, - ); + swa.release_window_lock(self, node_idx, params, device_frees, host_frees); // Drop strictly-lower-priority locks co-located on the node, skipping // any the paired inc never took. diff --git a/test/registered/unit/dllm/test_gemma4_uniform_lifecycle.py b/test/registered/unit/dllm/test_gemma4_uniform_lifecycle.py index 247cd4487ac4..7d35a91609fe 100644 --- a/test/registered/unit/dllm/test_gemma4_uniform_lifecycle.py +++ b/test/registered/unit/dllm/test_gemma4_uniform_lifecycle.py @@ -39,7 +39,6 @@ def __init__(self, context_len=8, block_size=4, *, prefill=False): req_pool_idx=1, kv_allocated_len=context_len + block_size, kv_committed_len=context_len + block_size, - swa_evicted_seqlen=0, ) self.finished_reason = None self.finish_on_update = False diff --git a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py index c772003c357b..cd271086b796 100644 --- a/test/registered/unit/managers/test_priority_scheduling_disaggregation.py +++ b/test/registered/unit/managers/test_priority_scheduling_disaggregation.py @@ -188,7 +188,6 @@ def _new_unpublished_swa_state(self): kv=SimpleNamespace( req_pool_idx=1, cache_protected_len=2, - swa_evicted_seqlen=1, ), prefix_indices=torch.tensor([8, 9], dtype=torch.int64), priority=3, diff --git a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py index c606f36aa7b0..e970ea3c06f7 100644 --- a/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py +++ b/test/registered/unit/mem_cache/test_decode_radix_lock_ref.py @@ -455,7 +455,9 @@ def test_pop_preallocated_rechecks_budget_after_lock(self): req.last_node = object() req.finished_reason = None req.kv.cache_protected_len = 0 - req.lock_receipt = DecLockRefParams(swa_uuid_for_lock=123) + req.lock_receipt = DecLockRefParams( + component_lock_uuids={ComponentType.SWA: 123} + ) req.swa_prefix_lock_released = False req.pd_rebootstrap_in_progress = False req.sampling_params.max_new_tokens = 16 @@ -538,11 +540,11 @@ def test_pop_preallocated_rechecks_budget_after_lock(self): queue._pre_alloc.assert_not_called() queue.tree_cache.dec_swa_lock_only.assert_called_once_with( req.last_node, - DecLockRefParams(swa_uuid_for_lock=123), + DecLockRefParams(component_lock_uuids={ComponentType.SWA: 123}), ) queue.tree_cache.dec_lock_ref.assert_called_once_with( req.last_node, - DecLockRefParams(swa_uuid_for_lock=123), + DecLockRefParams(component_lock_uuids={ComponentType.SWA: 123}), skip_swa=True, ) self.assertFalse(req.swa_prefix_lock_released) @@ -561,7 +563,9 @@ def test_hicache_restore_commit_hands_over_lock_with_receipt(self): req = MagicMock() req.req_pool_idx = 0 - req.lock_receipt = DecLockRefParams(swa_uuid_for_lock=123) + req.lock_receipt = DecLockRefParams( + component_lock_uuids={ComponentType.SWA: 123} + ) req.swa_prefix_lock_released = True # SWA tail-prealloc released early prealloc_node = object() @@ -576,7 +580,8 @@ def test_hicache_restore_commit_hands_over_lock_with_receipt(self): ) decode_req.hicache_restored_node = restored_node decode_req.hicache_restore_lock_receipt = DecLockRefParams( - swa_uuid_for_lock=456, skipped_lock_components=(ComponentType.MAMBA,) + component_lock_uuids={ComponentType.SWA: 456}, + skipped_lock_components=(ComponentType.MAMBA,), ) decode_req.hicache_restored_kv_indices = torch.arange(4, 8, dtype=torch.int64) @@ -584,11 +589,11 @@ def test_hicache_restore_commit_hands_over_lock_with_receipt(self): q.tree_cache.dec_lock_ref.assert_called_once_with( prealloc_node, - DecLockRefParams(swa_uuid_for_lock=123), + DecLockRefParams(component_lock_uuids={ComponentType.SWA: 123}), skip_swa=True, ) self.assertIs(req.last_node, restored_node) - self.assertEqual(req.lock_receipt.swa_uuid_for_lock, 456) + self.assertEqual(req.lock_receipt.component_lock_uuids[ComponentType.SWA], 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) diff --git a/test/registered/unit/mem_cache/test_hisparse_allocator.py b/test/registered/unit/mem_cache/test_hisparse_allocator.py index 837adf53cb56..b234d93987eb 100644 --- a/test/registered/unit/mem_cache/test_hisparse_allocator.py +++ b/test/registered/unit/mem_cache/test_hisparse_allocator.py @@ -10,6 +10,7 @@ from sglang.srt.mem_cache.allocator.hisparse import ( DeepSeekV4HiSparseTokenToKVPoolAllocator, ) +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType from sglang.srt.runtime_context import get_context, publish, reset_context from sglang.srt.server_args import ServerArgs from sglang.test.ci.ci_register import register_cpu_ci @@ -242,7 +243,9 @@ def write(self, indices, values): _, kwargs = allocator.alloc_extend_swa_tail.call_args self.assertEqual(kwargs["extend_num_tokens"], fill_len) self.assertEqual(kwargs["swa_tail_len"], swa_tail_len) - self.assertEqual(req.kv.swa_evicted_seqlen, fill_len - swa_tail_len) + self.assertEqual( + req.kv.get_evicted_seqlen(ComponentType.SWA), fill_len - swa_tail_len + ) self.assertEqual(req.kv.kv_allocated_len, fill_len) self.assertEqual(req.kv.kv_committed_len, fill_len) self.assertEqual(req.extend_range.length, fill_len) diff --git a/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py b/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py index cd81a35872c7..b74469e52d44 100644 --- a/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py +++ b/test/registered/unit/mem_cache/test_pure_swa_chunk_cache.py @@ -7,6 +7,7 @@ from sglang.srt.managers.schedule_batch import ReqKvInfo from sglang.srt.mem_cache.chunk_cache import PureSWAChunkCache +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -41,7 +42,7 @@ def _make_req(): req_pool_idx=0, cache_protected_len=2, swa_evict_floor=3, - swa_evicted_seqlen=6, + component_evicted_seqlens={ComponentType.SWA: 6}, ), ) diff --git a/test/registered/unit/mem_cache/test_pure_swa_radix_cache.py b/test/registered/unit/mem_cache/test_pure_swa_radix_cache.py index 80368a9a53bb..d8595defd233 100644 --- a/test/registered/unit/mem_cache/test_pure_swa_radix_cache.py +++ b/test/registered/unit/mem_cache/test_pure_swa_radix_cache.py @@ -11,6 +11,7 @@ from sglang.srt.mem_cache.cache_init_params import CacheInitParams from sglang.srt.mem_cache.pure_swa_radix_cache import PureSWARadixCache from sglang.srt.mem_cache.radix_cache import RadixKey +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -61,7 +62,11 @@ def test_finish_inserts_up_to_the_evict_floor_and_frees_the_rest(self): cache_salt=None, last_node=None, priority=0, - kv=ReqKvInfo(req_pool_idx=0, swa_evict_floor=4, swa_evicted_seqlen=6), + kv=ReqKvInfo( + req_pool_idx=0, + swa_evict_floor=4, + component_evicted_seqlens={ComponentType.SWA: 6}, + ), ) cache.cache_finished_req(req, owned_kv_len=8) diff --git a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py index 6e710ed370f0..8748e56c9719 100644 --- a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py +++ b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py @@ -745,8 +745,9 @@ def test_insert_value_none_materializes_the_token_ids(): assert matched.device_indices.tolist() == [1, 2, 3] -def test_empty_match_result_is_root_anchored(): - core = _tree_core() +@pytest.mark.parametrize("swa", [False, True]) +def test_empty_match_result_is_root_anchored(swa): + core = _swa_tree_core() if swa else _tree_core() empty = core.empty_match_result assert empty.device_indices.numel() == 0 assert empty.host_hit_length == 0 @@ -754,6 +755,12 @@ def test_empty_match_result_is_root_anchored(): assert empty.best_match_node == probe.best_match_node assert empty.last_device_node == probe.last_device_node assert empty.last_host_node == probe.last_host_node + core.dec_lock_ref(empty.best_match_node, DecLockRefParams()) + core.dec_host_lock_ref(empty.best_match_node, DecLockRefParams()) + released = core.dec_swa_lock_only(empty.best_match_node, DecLockRefParams()) + assert not released.device_frees + assert not released.host_frees + core.sanity_check([], []) def test_set_hicache_enabled_marks_the_tree(): @@ -893,14 +900,46 @@ def test_insert_host_reports_a_dropped_write_through_suffix(): assert result.host_insert_dropped -def test_host_lock_refs_round_trip(): - core = _tree_core() +@pytest.mark.parametrize( + "swa, missing_receipt", [(False, False), (True, False), (True, True)] +) +def test_host_lock_refs_round_trip(swa, missing_receipt): + core = _swa_tree_core() if swa else _tree_core() core.set_hicache_enabled() _insert(core, [1], [10]) leaf = core.match_prefix(MatchPrefixParams(key=_key([1]))).best_match_node - core.commit_backup(leaf, torch.tensor([100], dtype=torch.int64), {}) + component_transfers = {} + if swa: + core.has_swa_host_pool = True + core.set_component_device_value( + leaf, ComponentType.SWA, torch.tensor([20], dtype=torch.int64) + ) + component_transfers[ComponentType.SWA] = [ + PoolTransfer( + name=PoolName.SWA, + host_indices=torch.tensor([200], dtype=torch.int64), + ) + ] + core.commit_backup( + leaf, torch.tensor([100], dtype=torch.int64), component_transfers + ) + device_lock = core.inc_lock_ref(leaf) host_lock = core.inc_host_lock_ref(leaf) - core.dec_host_lock_ref(leaf, host_lock.to_dec_params()) + assert not host_lock.component_lock_uuids + assert not device_lock.component_host_lock_uuids + if swa: + assert host_lock.component_host_lock_uuids == {ComponentType.SWA: None} + params = host_lock.to_dec_params() + if missing_receipt: + params.component_host_lock_uuids.clear() + with pytest.raises(BaseException, match="no entry found for key") as error: + core.dec_host_lock_ref(leaf, params) + assert error.type.__name__ == "PanicException" + return # A Rust ownership violation poisons the core. + core.dec_host_lock_ref(leaf, params) + if swa: + assert core.swa_protected_size() == 1 + core.dec_lock_ref(leaf, device_lock.to_dec_params()) core.sanity_check([], []) @@ -1498,7 +1537,7 @@ def test_swa_straddling_insert_crosses_the_boundary_actions(): InsertParams( key=_key([1, 2, 3, 4]), value=torch.tensor([20, 21, 22, 23], dtype=torch.int64), - swa_evicted_seqlen=2, + component_evicted_seqlens={ComponentType.SWA: 2}, ), ) free_tail, rebuild, free_duplicates = result.cache_actions @@ -1943,24 +1982,37 @@ def test_lock_uuid_round_trips_through_dec_lock_ref(): ) node = first.cache_actions[-1].node_id result = core.inc_lock_ref(node) - assert result.swa_uuid_for_lock is not None - assert result.swa_uuid_for_host_lock is None + assert result.component_lock_uuids[ComponentType.SWA] is not None + assert not result.component_host_lock_uuids # The locked window is protected SWA accounting, visible through the binding. assert core.swa_protected_size() == 2 assert core.swa_evictable_size() == 1 core.dec_lock_ref( node, - DecLockRefParams(swa_uuid_for_lock=result.swa_uuid_for_lock), + result.to_dec_params(), ) # The uuid-bounded release returned the window to evictable. assert core.swa_protected_size() == 0 assert core.swa_evictable_size() == 3 # A repeat acquire reuses the stamped uuid. again = core.inc_lock_ref(node) - assert again.swa_uuid_for_lock == result.swa_uuid_for_lock + assert ( + again.component_lock_uuids[ComponentType.SWA] + == result.component_lock_uuids[ComponentType.SWA] + ) + skipped = core.inc_lock_ref(node, skip_lock_components=(ComponentType.SWA,)) + assert ComponentType.SWA not in skipped.component_lock_uuids + core.dec_lock_ref(node, skipped.to_dec_params()) + assert core.swa_protected_size() == 2 + core.dec_lock_ref(node, again.to_dec_params()) + core.sanity_check([], []) -def test_swa_tombstones_cross_the_binding_and_release_balanced(): +@pytest.mark.parametrize("swa_only", [False, True]) +@pytest.mark.parametrize("missing_receipt", [False, True]) +def test_swa_tombstones_cross_the_binding_and_release_balanced( + swa_only, missing_receipt +): core = _swa_tree_core(window=8) _insert(core, [1, 2], [10, 11]) second = _insert(core, [1, 2, 3, 4], [10, 11, 12, 13]) @@ -1971,13 +2023,23 @@ def test_swa_tombstones_cross_the_binding_and_release_balanced(): leaf, ComponentType.SWA, torch.tensor([52, 53], dtype=torch.int64) ) result = core.inc_lock_ref(leaf) - assert result.swa_uuid_for_lock is None + assert result.component_lock_uuids[ComponentType.SWA] is None assert core.swa_protected_size() == 2 - core.dec_lock_ref( - leaf, - DecLockRefParams(swa_uuid_for_lock=result.swa_uuid_for_lock), - ) + release = core.dec_swa_lock_only if swa_only else core.dec_lock_ref + if missing_receipt: + missing = DecLockRefParams(node_id=leaf) + if swa_only: + # Early release always walks SWA, even if the receipt marks it skipped. + missing.skipped_lock_components = (ComponentType.SWA,) + with pytest.raises(BaseException, match="no entry found for key") as error: + release(leaf, missing) + assert error.type.__name__ == "PanicException" + return # A Rust ownership violation poisons the core. + release(leaf, result.to_dec_params()) assert core.swa_protected_size() == 0 + if swa_only: + core.dec_lock_ref(leaf, DecLockRefParams(node_id=leaf), skip_swa=True) + core.sanity_check([], []) def test_dec_swa_lock_only_frees_flow_after_the_full_release(): @@ -1989,7 +2051,7 @@ def test_dec_swa_lock_only_frees_flow_after_the_full_release(): ) 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 + assert result.component_lock_uuids[ComponentType.SWA] is not None # The FULL lock releases first (skip_swa), then the early window release # finds a fully unlocked device leaf and evicts it in place. core.dec_lock_ref(node, result.to_dec_params(), skip_swa=True) diff --git a/test/registered/unit/mem_cache/test_streaming_session_unit.py b/test/registered/unit/mem_cache/test_streaming_session_unit.py index 335117a7ed5d..a41d58d65cb3 100644 --- a/test/registered/unit/mem_cache/test_streaming_session_unit.py +++ b/test/registered/unit/mem_cache/test_streaming_session_unit.py @@ -1,11 +1,17 @@ from types import SimpleNamespace +import pytest import torch from sglang.srt.managers.schedule_batch import FINISH_ABORT, ReqKvInfo from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator -from sglang.srt.mem_cache.base_prefix_cache import DecLockRefParams, MatchResult +from sglang.srt.mem_cache.base_prefix_cache import ( + DecLockRefParams, + IncLockRefResult, + MatchResult, +) from sglang.srt.mem_cache.common import release_kv_cache +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType from sglang.srt.session.streaming_session import SessionSlot, StreamingSession from sglang.test.ci.ci_register import register_cpu_ci @@ -92,7 +98,6 @@ def __init__( req_pool_idx=req_pool_idx, kv_committed_len=committed, kv_allocated_len=allocated, - swa_evicted_seqlen=0, cache_protected_len=0, ) self.origin_input_ids = list(range(committed)) @@ -112,13 +117,14 @@ def detach_kv(self): return kv -def test_session_slot_round_trip_preserves_mamba_state(): +def test_session_slot_round_trip_preserves_component_state(): # The mamba state rides in the shared ReqKvInfo record. mamba_branching_seqlen # is a per-turn match observation on the Req and is not preserved by the slot. req = _FakeReq("session-a", req_pool_idx=0, committed=4, allocated=4) req.kv.mamba_next_track_idx = 1 req.kv.mamba_last_track_idx = 0 req.kv.mamba_last_track_seqlen = 3 + req.kv.set_evicted_seqlen(ComponentType.SWA, 2) slot = SessionSlot() slot.save_from_req(req, is_first=True) @@ -129,6 +135,11 @@ def test_session_slot_round_trip_preserves_mamba_state(): assert next_req.kv.mamba_next_track_idx == 1 assert next_req.kv.mamba_last_track_idx == 0 assert next_req.kv.mamba_last_track_seqlen == 3 + assert next_req.kv.component_evicted_seqlens == {ComponentType.SWA: 2} + assert req.kv.component_evicted_seqlens == {} + next_req.kv.mark_kv_released() + assert next_req.kv.is_kv_released + assert next_req.kv.component_evicted_seqlens == {} def test_preabort_detaches_session_and_preserves_slot(): @@ -156,7 +167,6 @@ def test_preabort_detaches_session_and_preserves_slot(): req_pool_idx=0, kv_committed_len=48, kv_allocated_len=48, - swa_evicted_seqlen=0, cache_protected_len=16, ), ) @@ -235,11 +245,9 @@ def test_nth_mid_abort_nukes_session_slot(): assert req.kv.req_pool_idx is None -def test_release_session_threads_mamba_lock_receipt(): - """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 - 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.""" +@pytest.mark.parametrize("uuid", [None, 17]) +def test_release_session_preserves_component_lock_receipt(uuid): + """Closing a session releases only the component locks it acquired.""" req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128) req_to_token_pool = _FakeReqToTokenPool(req_to_token) allocator = _FakeAllocator() @@ -247,23 +255,38 @@ def test_release_session_threads_mamba_lock_receipt(): tree_cache = StreamingSession(inner) lock_node = SimpleNamespace(id=42) + acquired = IncLockRefResult( + node_id=42, + skipped_lock_components=(ComponentType.MAMBA,), + ) + acquired.set_lock_uuid(ComponentType.SWA, uuid) + acquired.set_lock_uuid(ComponentType.SWA, 19, lock_host=True) 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=acquired.to_dec_params(), ) + acquired.set_lock_uuid(ComponentType.SWA, 99) + acquired.set_lock_uuid(ComponentType.SWA, 99, lock_host=True) tree_cache.release_session("session-a") assert inner.dec_lock_ref_calls == [lock_node] params = inner.dec_lock_ref_params[0] assert params is not None - assert params.skipped_lock_components == () + assert params.skipped_lock_components == (ComponentType.MAMBA,) + assert params.get_lock_uuid(ComponentType.SWA) == uuid + assert params.get_lock_uuid(ComponentType.SWA, lock_host=True) == 19 + for lock_host in (False, True): + with pytest.raises(KeyError): + params.get_lock_uuid(ComponentType.MAMBA, lock_host=lock_host) + with pytest.raises(KeyError): + DecLockRefParams().get_lock_uuid(ComponentType.SWA, lock_host=lock_host) assert inner.dec_lock_ref_skip_swa == [False] @@ -283,18 +306,19 @@ def test_release_session_skips_swa_after_early_release(): 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), + lock_receipt=DecLockRefParams( + node_id=42, component_lock_uuids={ComponentType.SWA: 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_params[0].component_lock_uuids[ComponentType.SWA] == 7 assert inner.dec_lock_ref_skip_swa == [True] @@ -339,23 +363,24 @@ def test_trim_overshoot_postcondition(): req = _FakeReq("session-a", req_pool_idx=0, committed=40, allocated=44) req.origin_input_ids = list(range(26)) req.output_ids = list(range(14)) - req.kv.swa_evicted_seqlen = 42 + req.kv.set_evicted_seqlen(ComponentType.SWA, 42) tree_cache._trim_overshoot(req, finished_len=12) target = 38 assert req.kv.kv_committed_len == target assert req.kv.kv_allocated_len == target - assert req.kv.swa_evicted_seqlen == target + assert req.kv.get_evicted_seqlen(ComponentType.SWA) == target assert len(req.output_ids) == 12 # Tail [38, 44) freed by _free_kv_aligned, split at the pre-trim eviction # floor 42: [38, 42) gave its SWA peers back already, so it goes back full-only. assert [t.tolist() for t in allocator.freed] == [[38, 39, 40, 41], [42, 43]] -def test_trim_overshoot_keeps_cursor_page_aligned_on_paged(): - """A mid-page trim target must not become the SWA eviction cursor (the - dead/alive split there frees the shared page twice); rewind to the boundary.""" +@pytest.mark.parametrize("operation", ["trim", "match"]) +@pytest.mark.parametrize("component", [ComponentType.SWA, ComponentType.FULL]) +def test_session_rewind_keeps_component_cursors_page_aligned(operation, component): + """Rewinding below any component cursor must free whole pages.""" page_size = 16 req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128) req_to_token_pool = _FakeReqToTokenPool(req_to_token) @@ -368,21 +393,33 @@ def test_trim_overshoot_keeps_cursor_page_aligned_on_paged(): req = _FakeReq("session-a", req_pool_idx=0, committed=52, allocated=64) req.origin_input_ids = list(range(26)) req.output_ids = list(range(14)) - req.kv.swa_evicted_seqlen = 48 - - tree_cache._trim_overshoot(req, finished_len=12) + req.kv.set_evicted_seqlen(ComponentType.SWA, 16) + req.kv.set_evicted_seqlen(component, 48) + + if operation == "trim": + tree_cache._trim_overshoot(req, finished_len=12) + assert len(req.output_ids) == 12 + else: + slot = SessionSlot() + slot.save_from_req(req, is_first=True) + tree_cache.slots["session-a"] = slot + req = _FakeReq("session-a", req_pool_idx=0, committed=0, allocated=0) + result = tree_cache.match_prefix(SimpleNamespace(req=req, key=list(range(38)))) + assert result.device_indices.tolist() == list(range(32)) # Rewound to floor_align(38) = 32; every cursor lands page-aligned. assert req.kv.kv_allocated_len == 32 assert req.kv.kv_committed_len == 32 - assert req.kv.swa_evicted_seqlen == 32 - assert len(req.output_ids) == 12 - # Freed [32, 64): [32, 48) below the old cursor goes back full-only, - # [48, 64) both halves. - assert [t.tolist() for t in allocator.freed] == [ - list(range(32, 48)), - list(range(48, 64)), - ] + assert req.kv.get_evicted_seqlen(component) == 32 + assert req.kv.max_evicted_seqlen == 32 + assert req.kv.get_evicted_seqlen(ComponentType.SWA) == ( + 32 if component == ComponentType.SWA else 16 + ) + assert [t.tolist() for t in allocator.freed] == ( + [list(range(32, 48)), list(range(48, 64))] + if component == ComponentType.SWA + else [list(range(32, 64))] + ) if __name__ == "__main__": diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index 7e6b9dd45e08..410cbec7f619 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -20,6 +20,7 @@ ) from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool +from sglang.srt.mem_cache.unified_cache.component_type import ComponentType from sglang.srt.utils import get_device from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import CustomTestCase @@ -484,7 +485,7 @@ def test_grouped_full_side_frees_defer_and_skip_the_unique_path(self): def test_free_kv_row_reads_the_record_row_and_its_floor(self): indices = _swa_alloc(self.allocator, 8) cache = _RowCache(self.allocator, indices) - kv = ReqKvInfo(req_pool_idx=0, swa_evicted_seqlen=3) + kv = ReqKvInfo(req_pool_idx=0, component_evicted_seqlens={ComponentType.SWA: 3}) self.allocator.free_swa(indices[:3]) cache.free_kv_row(kv, [(1, 5)]) @@ -496,7 +497,7 @@ def test_free_kv_row_reads_the_record_row_and_its_floor(self): def test_single_pool_free_kv_row_still_frees_the_whole_range(self): allocator = _SinglePoolAllocator() cache = _RowCache(allocator, torch.arange(16, dtype=torch.int64)) - kv = ReqKvInfo(req_pool_idx=0, swa_evicted_seqlen=4) + kv = ReqKvInfo(req_pool_idx=0, component_evicted_seqlens={ComponentType.SWA: 4}) cache.free_kv_row(kv, [(2, 6)]) @@ -516,7 +517,11 @@ def test_rows_below_the_dead_floor_are_still_alive(self): self.assertEqual(allocator.swa_available_size(), baseline - 6) cache = _RowCache(allocator, indices) - kv = ReqKvInfo(req_pool_idx=0, swa_evicted_seqlen=6, swa_evict_floor=4) + kv = ReqKvInfo( + req_pool_idx=0, + component_evicted_seqlens={ComponentType.SWA: 6}, + swa_evict_floor=4, + ) cache.free_kv_row(kv, [(2, 8)]) self.assertEqual(allocator.swa_available_size(), baseline - 2) diff --git a/test/registered/unit/mem_cache/test_unified_cache_linker.py b/test/registered/unit/mem_cache/test_unified_cache_linker.py index fe6ba8ce9b0b..01760d66ed84 100644 --- a/test/registered/unit/mem_cache/test_unified_cache_linker.py +++ b/test/registered/unit/mem_cache/test_unified_cache_linker.py @@ -1263,7 +1263,8 @@ def test_linker_load_preserves_swa_boundaries( def prepare(phase, req, full_transfer, transfer, prefix_len, **kwargs): if phase == ExternalLinkerLoadPhase.PREPARE: req.kv = ReqKvInfo( - kv_allocated_len=prefix_len, swa_evicted_seqlen=prefix_len - 2 + kv_allocated_len=prefix_len, + component_evicted_seqlens={ComponentType.SWA: prefix_len - 2}, ) return transfer @@ -1307,7 +1308,8 @@ def prepare(phase, req, full_transfer, transfer, prefix_len, **kwargs): None if previous_boundary is None else ReqKvInfo( - kv_allocated_len=previous_boundary, swa_evicted_seqlen=previous_boundary + kv_allocated_len=previous_boundary, + component_evicted_seqlens={ComponentType.SWA: previous_boundary}, ) ) req = SimpleNamespace( @@ -1321,9 +1323,12 @@ def prepare(phase, req, full_transfer, transfer, prefix_len, **kwargs): assert restored.tolist() == full_indices.tolist() assert last_node == 0 - assert req.kv.swa_evicted_seqlen == expected_boundary + assert req.kv.get_evicted_seqlen(ComponentType.SWA) == expected_boundary assert req.kv.kv_allocated_len == (previous_boundary or 4) - assert cache.insert.call_args.args[0].swa_evicted_seqlen == expected_boundary + assert ( + cache.insert.call_args.args[0].get_evicted_seqlen(ComponentType.SWA) + == expected_boundary + ) cache.tree_core.mark_external_cache_stored_path.assert_called_once_with(0, 0) assert [c.args[0] for c in full.build_external_linker_transfer.call_args_list] == [ LinkerTransferPhase.LOAD diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index 79cdf03ca84d..9975a0522111 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -1924,7 +1924,7 @@ def test_swa_unfinished_req_preserves_existing_eviction_boundary(self): req.kv.cache_protected_len = 0 req.lock_receipt = DecLockRefParams() req.extra_key = None - req.kv.swa_evicted_seqlen = evicted_len + req.kv.set_evicted_seqlen(ComponentType.SWA, evicted_len) cache.cache_unfinished_req(req) @@ -2149,7 +2149,7 @@ def test_swa_unfinished_recovery_preserves_locked_full_value(self): req.kv.cache_protected_len = 0 req.lock_receipt = DecLockRefParams() req.extra_key = None - req.kv.swa_evicted_seqlen = 0 + req.kv.set_evicted_seqlen(ComponentType.SWA, 0) full_available_before_insert = allocator.full_attn_allocator.available_size() @@ -2203,7 +2203,7 @@ def test_swa_insert_keeps_full_leaf_when_entire_span_is_outside_window(self): key=RadixKey(array("q", tokens)), value=value, prev_prefix_len=0, - swa_evicted_seqlen=len(tokens), + component_evicted_seqlens={ComponentType.SWA: len(tokens)}, ) ) @@ -2293,9 +2293,7 @@ def test_leaf_transition_swa_evict_spares_locked_full(self): ) cache.sanity_check() - cache.dec_lock_ref( - node_a, DecLockRefParams(swa_uuid_for_lock=None), skip_swa=True - ) + cache.dec_lock_ref(node_a, DecLockRefParams(), skip_swa=True) cache.sanity_check() def test_mamba_opt_out_holder_cannot_release_another_holders_mamba_lock(self): @@ -2603,9 +2601,7 @@ def _collect_descendants(node_id): "SWA stays in LRU for drive_eviction to pick later", ) - cache.dec_lock_ref( - node_a, DecLockRefParams(swa_uuid_for_lock=None), skip_swa=True - ) + cache.dec_lock_ref(node_a, DecLockRefParams(), skip_swa=True) self.assertTrue(cache.tree_core.is_device_leaf(node_a)) cache.sanity_check() @@ -2851,10 +2847,10 @@ def test_swa_eager_eviction_on_unfinished_req(self): cushion = max(self.cfg.sliding_window_size, self.cfg.page_size) expected_evicted = (pre_len - 1) - cushion self.assertEqual( - req.kv.swa_evicted_seqlen, + req.kv.get_evicted_seqlen(ComponentType.SWA), expected_evicted, f"swa_evicted_seqlen should advance to (pre_len-1) - cushion = " - f"{expected_evicted}, got {req.kv.swa_evicted_seqlen}", + f"{expected_evicted}, got {req.kv.get_evicted_seqlen(ComponentType.SWA)}", ) swa_avail_after = allocator.swa_attn_allocator.available_size() @@ -2936,7 +2932,7 @@ def test_swa_eager_eviction_noop_when_within_window(self): cache.cache_unfinished_req(req) self.assertEqual( - req.kv.swa_evicted_seqlen, + req.kv.get_evicted_seqlen(ComponentType.SWA), 0, "Nothing should be evicted when prefill fits inside the cushion", ) @@ -4678,7 +4674,7 @@ def _assert_masked_full_is_reused(self, storage_dir, seq, live_tail: int): InsertParams( key=RadixKey(array("q", seq)), value=value[: len(seq)], - swa_evicted_seqlen=live_from, + component_evicted_seqlens={ComponentType.SWA: live_from}, ) ) masked = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) @@ -4806,7 +4802,7 @@ def test_buffer_only_load_back_reuses_partial_masked_full(self): InsertParams( key=RadixKey(array("q", sibling)), value=value[: len(sibling)], - swa_evicted_seqlen=len(sibling), + component_evicted_seqlens={ComponentType.SWA: len(sibling)}, ) ) live = cons.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq)))) @@ -4886,7 +4882,7 @@ def test_buffer_only_masked_head_is_not_evicted_for_tail_load(self): InsertParams( key=RadixKey(array("q", seq[:split_at])), value=value[:split_at], - swa_evicted_seqlen=split_at, + component_evicted_seqlens={ComponentType.SWA: split_at}, ) ) @@ -6055,7 +6051,7 @@ def test_hicache_write_through_offloads_swa_split_leaf(self): InsertParams( key=RadixKey(array("q", seq)), value=value, - swa_evicted_seqlen=ps, + component_evicted_seqlens={ComponentType.SWA: ps}, ) ) self.assertEqual(result.prefix_len, 0) @@ -6096,7 +6092,7 @@ def insert_swa(tokens, swa_ev): InsertParams( key=RadixKey(array("q", tokens)), value=value, - swa_evicted_seqlen=swa_ev, + component_evicted_seqlens={ComponentType.SWA: swa_ev}, ) ) cache.writing_check() @@ -6972,15 +6968,21 @@ def test_swa_branching_seqlen_caps_insert_after_forward(self): swa = cache.components[ComponentType.SWA] req = mock.Mock( swa_branching_seqlen=8, - kv=mock.Mock(cache_protected_len=4, swa_evicted_seqlen=0), + kv=ReqKvInfo( + cache_protected_len=4, + component_evicted_seqlens={ComponentType.SWA: 3}, + ), ) for is_finished in (False, True): - params = InsertParams() + params = InsertParams(component_evicted_seqlens={ComponentType.MAMBA: 6}) self.assertEqual( swa.prepare_for_caching_req(req, params, 12, is_finished), 8 ) - self.assertEqual(params.swa_evicted_seqlen, 0) + self.assertEqual( + params.component_evicted_seqlens, + {ComponentType.SWA: 3, ComponentType.MAMBA: 6}, + ) self.assertIsNone(swa.prepare_for_caching_req(req, InsertParams(), 7, False)) req.kv.cache_protected_len = 8 @@ -7045,7 +7047,7 @@ def test_swa_branch_insert_releases_forward_overshoot(self): ) expected_evicted = forward_len - 1 - self.cfg.sliding_window_size - self.assertEqual(req.kv.swa_evicted_seqlen, expected_evicted) + self.assertEqual(req.kv.get_evicted_seqlen(ComponentType.SWA), expected_evicted) mapping = allocator.full_to_swa_index_mapping self.assertEqual( @@ -9550,7 +9552,13 @@ def _recover_swa_tombstone_with_tracking(self, *, lock_full: bool): evicted = self._alloc(allocator, len(seq)) # Window eviction already released the peers below the floor. allocator.free_swa(evicted) - cache.insert(InsertParams(key=key, value=evicted, swa_evicted_seqlen=len(seq))) + cache.insert( + InsertParams( + key=key, + value=evicted, + component_evicted_seqlens={ComponentType.SWA: len(seq)}, + ) + ) (leaf,) = _node_children(cache, cache.root_node_handle()) lock_result = cache.inc_lock_ref(leaf) if lock_full else None try: @@ -9558,7 +9566,6 @@ def _recover_swa_tombstone_with_tracking(self, *, lock_full: bool): InsertParams( key=key, value=self._alloc(allocator, len(seq)), - swa_evicted_seqlen=0, track_adopted_ranges=True, ) ) @@ -9584,11 +9591,17 @@ def test_swa_tombstone_recovery_frees_full_only(self): key = RadixKey(array("q", seq)) evicted = self._alloc(allocator, len(seq)) allocator.free_swa(evicted[:sw]) - cache.insert(InsertParams(key=key, value=evicted, swa_evicted_seqlen=sw)) + cache.insert( + InsertParams( + key=key, + value=evicted, + component_evicted_seqlens={ComponentType.SWA: sw}, + ) + ) value = self._alloc(allocator, len(seq)) full_available = allocator.full_attn_allocator.available_size() swa_available = allocator.swa_attn_allocator.available_size() - cache.insert(InsertParams(key=key, value=value, swa_evicted_seqlen=0)) + cache.insert(InsertParams(key=key, value=value)) self.assertEqual( allocator.full_attn_allocator.available_size(), @@ -9610,17 +9623,19 @@ def test_swa_recovery_keeps_recovered_node_below_window_nodes(self): key = RadixKey(array("q", seq)) evicted = self._alloc(allocator, len(seq)) allocator.free_swa(evicted[:sw]) - cache.insert(InsertParams(key=key, value=evicted, swa_evicted_seqlen=sw)) + cache.insert( + InsertParams( + key=key, + value=evicted, + component_evicted_seqlens={ComponentType.SWA: sw}, + ) + ) (prefix_node,) = _node_children(cache, cache.root_node_handle()) (window_node,) = _node_children(cache, prefix_node) self.assertIsNone(_device_value(cache, prefix_node, ComponentType.SWA)) # Re-inserting fully in-window recovers the prefix span's SWA data. - cache.insert( - InsertParams( - key=key, value=self._alloc(allocator, len(seq)), swa_evicted_seqlen=0 - ) - ) + cache.insert(InsertParams(key=key, value=self._alloc(allocator, len(seq)))) self.assertIsNotNone(_device_value(cache, prefix_node, ComponentType.SWA)) # SWA eviction takes the recovered span and keeps the window leaf. @@ -9645,7 +9660,13 @@ def test_dup_slice_below_eviction_floor_frees_full_only(self): with mock.patch.object( cache, "_apply_cache_action", wraps=cache._apply_cache_action ) as spy: - cache.insert(InsertParams(key=key, value=value, swa_evicted_seqlen=sw)) + cache.insert( + InsertParams( + key=key, + value=value, + component_evicted_seqlens={ComponentType.SWA: sw}, + ) + ) actions = [c.args[0] for c in spy.call_args_list] full_only = [ @@ -9671,9 +9692,7 @@ def test_dec_swa_lock_only_early_release_keeps_full_lock(self): self.assertEqual(_device_lock_ref(cache, node, ComponentType.SWA), 0) self.assertGreaterEqual(_device_lock_ref(cache, node, ComponentType.FULL), 1) - cache.dec_lock_ref( - node, DecLockRefParams(swa_uuid_for_lock=None), skip_swa=True - ) + cache.dec_lock_ref(node, DecLockRefParams(), skip_swa=True) cache.sanity_check() @@ -10340,10 +10359,10 @@ def test_match_after_insert_reaches_the_new_leaf(self): boundary = (seq_len - 1) // page_size * page_size self.assertGreaterEqual( - boundary - req.kv.swa_evicted_seqlen, + boundary - req.kv.get_evicted_seqlen(ComponentType.SWA), self.cfg.sliding_window_size, f"leaf ending at {boundary} keeps only " - f"{boundary - req.kv.swa_evicted_seqlen} live SWA tokens against a " + f"{boundary - req.kv.get_evicted_seqlen(ComponentType.SWA)} live SWA tokens against a " f"{self.cfg.sliding_window_size} window", ) self.assertEqual( @@ -10689,7 +10708,7 @@ def test_rebuilt_tombstone_relock_release_no_theft(self): InsertParams( key=RadixKey(array("q", seq)), value=value, - swa_evicted_seqlen=swa_evicted, + component_evicted_seqlens={ComponentType.SWA: swa_evicted}, ) ) leaf = self._deepest(cache) @@ -10702,7 +10721,7 @@ def test_rebuilt_tombstone_relock_release_no_theft(self): 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.assertIsNotNone(lock_a.component_lock_uuids[ComponentType.SWA]) self.assertEqual(lock_a.node_id, leaf) for n in segment: self.assertEqual(self._swa_ref(cache, n), 1) @@ -10714,7 +10733,6 @@ def test_rebuilt_tombstone_relock_release_no_theft(self): InsertParams( key=RadixKey(array("q", seq)), value=self._alloc(allocator, len(seq)), - swa_evicted_seqlen=0, ) ) cache.sanity_check() @@ -10738,20 +10756,31 @@ def test_rebuilt_tombstone_relock_release_no_theft(self): 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.""" + def test_root_release_needs_no_boundary(self): + cache, _, _ = build_fixture(self.cfg) + core = cache.tree_core + root = cache.root_node_handle() + core.dec_lock_ref(root, DecLockRefParams()) + core.dec_host_lock_ref(root, DecLockRefParams()) + released = core.dec_swa_lock_only(root, DecLockRefParams()) + self.assertFalse(released.device_frees) + self.assertFalse(released.host_frees) + cache.sanity_check() + + def test_release_with_incorrect_root_boundary_fails_loud(self): + """An incorrect root boundary must fail at the acquired segment's edge.""" 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.assertIsNotNone(lock.component_lock_uuids[ComponentType.SWA]) self._assert_protocol_violation( - lambda: cache.dec_lock_ref(leaf, DecLockRefParams(swa_uuid_for_lock=None)), + lambda: cache.dec_lock_ref( + leaf, DecLockRefParams(component_lock_uuids={ComponentType.SWA: None}) + ), "lock_ref=0", ) @@ -10914,7 +10943,7 @@ def _run_seed(self, seed: int, steps: int = 120): params = InsertParams( key=RadixKey(array("q", seq)), value=value, - swa_evicted_seqlen=swa_evict, + component_evicted_seqlens={ComponentType.SWA: swa_evict}, ) if self.cfg.has_mamba: req = self._make_req(req_to_token_pool) @@ -10945,7 +10974,10 @@ def _run_seed(self, seed: int, steps: int = 120): # 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: + if ( + released + or receipt.component_lock_uuids[ComponentType.SWA] is None + ): continue cache.dec_swa_lock_only( node_id, @@ -11014,7 +11046,7 @@ def _lock_and_early_release(self, cache, allocator): 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) + self.assertIsNotNone(lock.component_lock_uuids[ComponentType.SWA]) cache.dec_swa_lock_only(node, lock.to_dec_params()) return node, lock From f89a22e839c2e1714ec4291b7ff98262d4e801fb Mon Sep 17 00:00:00 2001 From: lmzheng Date: Fri, 25 Sep 2026 15:23:49 -0700 Subject: [PATCH 2/3] [MemCache] Update SWA lock fixture for component receipts --- .../unit/mem_cache/test_mamba_donated_alloc_ratio.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py b/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py index b628d176b5a3..91f91f27ae28 100644 --- a/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py +++ b/test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py @@ -195,6 +195,7 @@ def _run(self, skipped_lock_components): mamba = _RecordingComp(ComponentType.MAMBA, 0) node = SimpleNamespace(id=7) tree_core = SimpleNamespace( + root_node=object(), components=(full, swa, mamba), components_by_type={ComponentType.SWA: swa}, node_by_id=lambda node_id: node, @@ -203,7 +204,11 @@ def _run(self, skipped_lock_components): UnifiedTreeCore.dec_swa_lock_only( tree_core, node.id, - DecLockRefParams(skipped_lock_components=skipped_lock_components), + DecLockRefParams( + node_id=node.id, + skipped_lock_components=skipped_lock_components, + component_lock_uuids={ComponentType.SWA: None}, + ), ) return full, mamba From 1f651d63a528caffc6f14a8e1b265e4251f006f3 Mon Sep 17 00:00:00 2001 From: lmzheng Date: Fri, 25 Sep 2026 15:39:20 -0700 Subject: [PATCH 3/3] [Test] Initialize LMCache flag in scheduler flush fixture --- .../unit/managers/test_scheduler_flush_cache_after_retract.py | 1 + 1 file changed, 1 insertion(+) diff --git a/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py b/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py index 2cc85d60f956..2f9ef4ae853e 100644 --- a/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py +++ b/test/registered/unit/managers/test_scheduler_flush_cache_after_retract.py @@ -48,6 +48,7 @@ def _new_scheduler(self) -> Scheduler: scheduler.grammar_manager = MagicMock() scheduler.grammar_manager.grammar_queue = [] scheduler.enable_hierarchical_cache = False + scheduler.enable_lmcache = False scheduler.enable_hisparse = False scheduler.tree_cache = MagicMock() scheduler.req_to_token_pool = MagicMock()