[MemCache] Unify component eviction cursors and lock receipts

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(

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)
 

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
                     ):

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:

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,41 +170,70 @@ 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),
         )
 
 
 @dataclasses.dataclass
 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),
         )
 

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(

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]]:

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,