From 3a770da7566aa55a884896531cc7388898550e62 Mon Sep 17 00:00:00 2001 From: Jincong Chen Date: Sat, 5 Sep 2026 12:29:16 +0800 Subject: [PATCH] [Unified Tree] Support Branching-Point Caching for the SWA Component (#34565) Co-authored-by: alphabetc1 <2508695655@qq.com> Co-authored-by: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com> --- .../decode_schedule_batch_mixin.py | 3 + python/sglang/srt/managers/schedule_batch.py | 7 + python/sglang/srt/managers/schedule_policy.py | 1 + .../sglang/srt/mem_cache/base_prefix_cache.py | 7 + .../unified_cache/components/swa_component.py | 143 +++++++-- .../unified_cache/unified_tree_core.py | 10 +- .../sglang/srt/session/streaming_session.py | 2 + .../unit/mem_cache/test_radix_force_miss.py | 23 ++ .../mem_cache/test_streaming_session_unit.py | 15 + .../test_unified_radix_cache_unittest.py | 285 ++++++++++++++++++ 10 files changed, 467 insertions(+), 29 deletions(-) diff --git a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py index 07a786f76..5b5aada9f 100644 --- a/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py +++ b/python/sglang/srt/disaggregation/decode_schedule_batch_mixin.py @@ -117,6 +117,9 @@ class ScheduleBatchDisaggregationDecodeMixin: last_tokens: List[int] = [] for req in self.reqs: last_tokens.append(req.output_ids[-1]) + # PREBUILT does not materialize a local SWA branching window. + if req.swa_branching_seqlen is not None: + req.swa_branching_seqlen = None maybe_cache_unfinished_req(req, self.tree_cache) if req.grammar is not None: # FIXME: this try-except block is for handling unexpected xgrammar issue. diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 347d1483d..db3d02c29 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1004,6 +1004,10 @@ class Req(ReqDllmMixin): # For req-level memory management self.kv = ReqKvInfo() + # Full-KV-derived boundary whose SWA window should be inserted after + # the current prefill pass. + self.swa_branching_seqlen: Optional[int] = None + # for cross-encoder model self.token_type_ids = token_type_ids @@ -1515,6 +1519,7 @@ class Req(ReqDllmMixin): self.best_match_node, self.host_hit_length, self.swa_host_hit_length, + self.swa_branching_seqlen, self.mamba_host_hit_length, self.mamba_branching_seqlen, ) = ( @@ -1524,6 +1529,7 @@ class Req(ReqDllmMixin): match_result.best_match_node, match_result.host_hit_length, match_result.swa_host_hit_length, + match_result.swa_branching_seqlen, match_result.mamba_host_hit_length, match_result.mamba_branching_seqlen, ) @@ -1817,6 +1823,7 @@ class Req(ReqDllmMixin): self.num_matched_prefix_tokens = 0 self.swa_uuid_for_lock = None self.swa_prefix_lock_released = False + self.swa_branching_seqlen = None self.skip_lock_node_ids = {} self.extend_range = None self.dllm_initialized = False diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 169bcbbcf..023071a67 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -195,6 +195,7 @@ def match_prefix_for_req( req.num_matched_prefix_tokens = min( len(req.prefix_indices) + req.host_hit_length, max_len ) + req.swa_branching_seqlen = match_result.swa_branching_seqlen if match_result.mamba_branching_seqlen is not None: req.mamba_branching_seqlen = match_result.mamba_branching_seqlen if match_result.cache_protected_len 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 37891f6ea..530d630e8 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/python/sglang/srt/mem_cache/base_prefix_cache.py @@ -73,6 +73,7 @@ class InsertParams: # SWA specific prev_prefix_len: int = 0 swa_evicted_seqlen: int = 0 + swa_branching_seqlen: Optional[int] = None # General chunked: bool = False @@ -88,6 +89,7 @@ class InsertResult: total_len: int = 0 last_device_node: Any = None mamba_exist: bool = False + swa_branch_inserted: bool = False inserted_host_node: Any = None host_insert_dropped: bool = False adopted_ranges: Optional[dict[ComponentType, list[tuple[int, int]]]] = None @@ -201,6 +203,9 @@ class MatchResult(NamedTuple): loaded back to device. Pure-KV cache semantics; swa_host_hit_length : Number of SWA tokens that hit on host (within the sliding window) and will be load-back into the SWA device pool. + swa_branching_seqlen: The SWA radix cache branching point, which is the longest + page-aligned position that could've been cache hit if there + exists an SWA window. mamba_host_hit_length: Number of Mamba slots that hit on host and will be load-back into the Mamba device pool. Typically 0 or 1. mamba_branching_seqlen: The mamba radix cache branching point, which is the longest @@ -216,6 +221,7 @@ class MatchResult(NamedTuple): best_match_node: Any host_hit_length: int = 0 swa_host_hit_length: int = 0 + swa_branching_seqlen: Optional[int] = None mamba_host_hit_length: int = 0 mamba_branching_seqlen: Optional[int] = None cache_protected_len: Optional[int] = None @@ -240,6 +246,7 @@ def zero_match_result( best_match_node=root, host_hit_length=0, swa_host_hit_length=0, + swa_branching_seqlen=None, mamba_host_hit_length=0, full_kv_hit_length=0, ) diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py index 772699de0..9715fe343 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py @@ -4,6 +4,7 @@ from typing import TYPE_CHECKING, Callable, Optional, Sequence import torch +from sglang.srt.environ import envs from sglang.srt.mem_cache.base_prefix_cache import ( DecLockRefParams, EvictParams, @@ -84,8 +85,32 @@ class SWAComponent(TreeComponent): component_type = ComponentType.SWA + def _dirty_backup_window(self, node: UnifiedTreeNode) -> list[UnifiedTreeNode]: + if not self.tree_core.has_swa_host_pool: + return [] + + ct = self.component_type + covered = 0 + dirty: list[UnifiedTreeNode] = [] + cur = node + while ( + cur is not self.tree_core.root_node and covered < self.sliding_window_size + ): + if cur.write_through_pending_id is not None: + break + + cd = cur.component_data[ct] + value = cd.value if cd.value is not None else cd.host_value + if value is None: + break + covered += len(value) + if cd.value is not None and cd.host_value is None: + dirty.append(cur) + cur = cur.parent + return dirty + def needs_incremental_backup(self, node: UnifiedTreeNode) -> bool: - return False + return bool(self._dirty_backup_window(node)) def reset_session_state(self) -> None: super().reset_session_state() @@ -308,6 +333,16 @@ class SWAComponent(TreeComponent): best_value_len: int, ) -> MatchResult: ct = self.component_type + swa_boundary_len = len(result.device_indices) + result.host_hit_length + + # Full KV may extend beyond the latest reusable SWA window. The branching + # point is the last page-aligned position within the Full-KV hit that lies + # beyond the current SWA boundary. + aligned_seqlen = ( + result.full_kv_hit_length // self.tree_core.page_size + ) * self.tree_core.page_size + branching_seqlen = aligned_seqlen if aligned_seqlen > swa_boundary_len else None + n_swa = 0 swa_host_hit = 0 node = result.best_match_node @@ -328,11 +363,11 @@ class SWAComponent(TreeComponent): else: break node = node.parent - if swa_host_hit > 0: - return result._replace( - swa_host_hit_length=max(result.swa_host_hit_length, swa_host_hit) - ) - return result + + return result._replace( + swa_host_hit_length=max(result.swa_host_hit_length, swa_host_hit), + swa_branching_seqlen=branching_seqlen, + ) def update_component_on_insert_overlap( self, @@ -466,6 +501,11 @@ class SWAComponent(TreeComponent): result: InsertResult, cache_actions: list[CacheAction | ComponentAction], ) -> None: + branching_seqlen = params.swa_branching_seqlen + if branching_seqlen is not None: + assert params.key is not None + result.swa_branch_inserted = len(params.key) >= branching_seqlen + if not is_new_leaf: return @@ -857,23 +897,59 @@ class SWAComponent(TreeComponent): # 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 - return None + + branching_seqlen = req.swa_branching_seqlen + if branching_seqlen is None or branching_seqlen <= req.kv.cache_protected_len: + return None + + # An EAGLE key with N bigrams spans N + 1 raw tokens. + effective_cache_len = branching_seqlen + int(self.tree_core.is_eagle) + if effective_cache_len > token_ids_len: + return None + + # Record the logical SWA branch boundary for insertion. + insert_params.swa_branching_seqlen = branching_seqlen + return effective_cache_len + + def _free_out_of_window_slots(self, req: Req, pre_len: int) -> None: + if self.sliding_window_size is None: + return + free_swa_out_of_window_slots( + req, + pre_len, + sliding_window_size=self.sliding_window_size, + page_size=self.cache.page_size, + req_to_token_pool=self.cache.req_to_token_pool, + token_to_kv_pool_allocator=self.cache.token_to_kv_pool_allocator, + retain_floor=self.cache.swa_retain_floor(req), + ) def free_out_of_window_slots( self, req: Req, pre_len: int, insert_params: InsertParams ) -> None: - if self.sliding_window_size is not None: - free_swa_out_of_window_slots( - req, - pre_len, - sliding_window_size=self.sliding_window_size, - page_size=self.cache.page_size, - req_to_token_pool=self.cache.req_to_token_pool, - token_to_kv_pool_allocator=self.cache.token_to_kv_pool_allocator, - retain_floor=self.cache.swa_retain_floor(req), - ) + self._free_out_of_window_slots(req, pre_len) insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen + def cleanup_after_caching_req( + self, + req: Req, + is_finished: bool, + insert_result: Optional[InsertResult] = None, + insert_params: Optional[InsertParams] = None, + ) -> None: + if insert_result is not None and insert_result.swa_branch_inserted: + req.swa_branching_seqlen = None + + # Free unused SWA slots after inserting the branch. + if ( + not is_finished + and insert_result is not None + and insert_result.swa_branch_inserted + and envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.get() + ): + forward_key_len = len(req.get_fill_ids()) - int(self.tree_core.is_eagle) + self._free_out_of_window_slots(req, forward_key_len - 1) + # ---- HiCache Hooks ---- def prepare_prefetch( @@ -933,15 +1009,22 @@ class SWAComponent(TreeComponent): return None if phase == CacheTransferPhase.BACKUP_HOST: - cd = node.component_data[ct] - if cd.value is None: + if self.cache.host_memory_mode == "buffer_only": + # Buffer mode stages one node/hash span per FIFO backup intent. + cd = node.component_data[ct] + dirty = [node] if cd.value is not None else [] + else: + dirty = self._dirty_backup_window(node) + if not dirty: return None - # cd.value already holds SWA-pool indices (translated at insert time). - # Host pool indexing wants int64. + dirty.reverse() return [ PoolTransfer( name=PoolName.SWA, - device_indices=cd.value.to(torch.int64), + device_indices=torch.cat( + [n.component_data[ct].value for n in dirty] + ).to(torch.int64), + nodes_to_load=[n.id for n in dirty], ) ] @@ -1120,10 +1203,18 @@ class SWAComponent(TreeComponent): ct = self.component_type if phase == CacheTransferPhase.BACKUP_HOST: - if transfers and transfers[0].host_indices is not None: - cd = node.component_data[ct] - if cd.host_value is None: - cd.host_value = transfers[0].host_indices.clone() + if not transfers or transfers[0].host_indices is None: + return + xfer = transfers[0] + offset = 0 + for node_id in xfer.nodes_to_load or [node.id]: + target = self.tree_core.node_by_id(node_id) + cd = target.component_data[ct] + assert cd.value is not None and cd.host_value is None + size = len(cd.value) + cd.host_value = xfer.host_indices[offset : offset + size].clone() + offset += size + assert offset == len(xfer.host_indices) return if phase == CacheTransferPhase.LOAD_BACK: 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 37264a7d9..d97046d8c 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 @@ -1131,9 +1131,13 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): state.phase = _InsertPhase.TAIL def _needs_incremental_component_backup(self, node: UnifiedTreeNode) -> bool: + components = self.components + if self.is_write_back: + swa = self.components_by_type.get(ComponentType.SWA) + components = () if swa is None else (swa,) return any( component.needs_incremental_backup(node) - for component in self.components + for component in components if component.component_type != BASE_COMPONENT_TYPE ) @@ -1147,7 +1151,6 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): node = state.target_node return ( self.enable_hicache - and not self.is_write_back and node.backuped and node.write_through_pending_id is None and self._needs_incremental_component_backup(node) @@ -1979,7 +1982,8 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): for comp in self.components: if comp.component_type == BASE_COMPONENT_TYPE: continue - if node.component_data[comp.component_type].host_value is not None: + cd = node.component_data[comp.component_type] + if cd.host_value is not None and not comp.needs_incremental_backup(node): continue t = comp.build_hicache_transfers(node, CacheTransferPhase.BACKUP_HOST) if t: diff --git a/python/sglang/srt/session/streaming_session.py b/python/sglang/srt/session/streaming_session.py index d9b271c35..e661a82b9 100644 --- a/python/sglang/srt/session/streaming_session.py +++ b/python/sglang/srt/session/streaming_session.py @@ -66,6 +66,8 @@ class SessionSlot: # Later turns run on the slot's record (see restore_to_req). assert kv is self.kv + req.swa_branching_seqlen = None + def restore_to_req(self, req: Req): """Restore KV state from this slot into an incoming request.""" req.kv = self.kv diff --git a/test/registered/unit/mem_cache/test_radix_force_miss.py b/test/registered/unit/mem_cache/test_radix_force_miss.py index 5a06596b0..c2eacc9dd 100644 --- a/test/registered/unit/mem_cache/test_radix_force_miss.py +++ b/test/registered/unit/mem_cache/test_radix_force_miss.py @@ -39,6 +39,7 @@ class _StubReq: self.best_match_node = None self.host_hit_length = None self.num_matched_prefix_tokens = 0 + self.swa_branching_seqlen = None self.kv = SimpleNamespace(cache_protected_len=None) def _compute_max_prefix_len(self, input_len): @@ -79,6 +80,28 @@ class TestZeroMatchResult(unittest.TestCase): class TestMatchPrefixForReqForceMiss(unittest.TestCase): + def test_swa_branching_seqlen_is_cleared_without_new_branch(self): + class _StubTreeCache: + def swa_reprefill_tail_tokens(self): + return 0 + + def match_prefix(self, params): + return MatchResult( + device_indices=torch.empty((0,), dtype=torch.int64), + last_device_node=None, + last_host_node=None, + best_match_node=None, + host_hit_length=0, + swa_branching_seqlen=None, + ) + + req = _StubReq([1, 2, 3, 4]) + req.swa_branching_seqlen = 8 + + match_prefix_for_req(_StubTreeCache(), req) + + self.assertIsNone(req.swa_branching_seqlen) + def test_force_miss_zeros_req_prefix(self): tree = RadixCache.create_simulated() tree.insert( 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 47064be66..a200755e3 100644 --- a/test/registered/unit/mem_cache/test_streaming_session_unit.py +++ b/test/registered/unit/mem_cache/test_streaming_session_unit.py @@ -99,6 +99,7 @@ class _FakeReq: self.last_node = None self.swa_uuid_for_lock = None self.skip_lock_node_ids = {} + self.swa_branching_seqlen = None self.to_finish = None self.finished_reason = None self.finished_len = None @@ -265,6 +266,20 @@ def test_release_session_threads_mamba_skip_ids(): assert params.skip_lock_node_ids.get(ComponentType.MAMBA) == {42} +def test_session_slot_does_not_restore_swa_branching_seqlen(): + req = _FakeReq("session-a", req_pool_idx=0, committed=4, allocated=4) + req.swa_branching_seqlen = 8 + + slot = SessionSlot() + slot.save_from_req(req, is_first=True) + + next_req = _FakeReq("session-a", req_pool_idx=1, committed=0, allocated=0) + slot.restore_to_req(next_req) + + assert req.swa_branching_seqlen is None + assert next_req.swa_branching_seqlen is None + + # Shrink tests removed: streaming sessions are append-only after the # rollback fix in session_controller (rollback_aborted_req). The shrink # code path in cache_finished_req no longer exists. 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 527c6a423..645019745 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 @@ -5765,6 +5765,231 @@ class UnifiedRadixCacheSuite: self.assertEqual(result.host_hit_length, 0) self.assertEqual(result.swa_host_hit_length, _node_key_length(cache, leaf)) + def _skip_swa_branching_on_rust(self) -> None: + # TODO(alphabetc1): drop this gate once #37584 ports SWA branching-point + # caching to the Rust tree core. + if _selected_tree_core_test_backend() == "rust": + self.skipTest("SWA branching-point caching is Python-core only") + + def test_swa_branching_seqlen_uses_device_full_hit(self): + self._skip_swa_branching_on_rust() + if ( + not self.cfg.has_swa + or self.cfg.has_mamba + or self.cfg.page_size not in (1, 4) + or self.cfg.sliding_window_size != 4 + ): + self.skipTest("requires Full+SWA with window_size=4") + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + window = self.cfg.sliding_window_size + window_pages = (window + self.cfg.page_size - 1) // self.cfg.page_size + prefix = self._make_seq(1, window_pages) + tokens = prefix + self._make_seq(1000, 2) + self._insert(cache, allocator, req_to_token_pool, prefix) + self._insert(cache, allocator, req_to_token_pool, tokens) + + leaf = cache.resolve_node_handle( + cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ).last_device_node + ) + device_frees = defaultdict(list) + cache.tree_core._evict_component_and_detach_lru( + leaf, + cache.components[ComponentType.SWA], + device_frees=device_frees, + host_frees=defaultdict(list), + target=EvictLayer.DEVICE, + ) + cache._drain_device_frees(device_frees) + + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) + + self.assertEqual(result.full_kv_hit_length, len(tokens)) + self.assertEqual(result.swa_branching_seqlen, len(tokens)) + self.assertEqual(result.swa_branching_seqlen % self.cfg.page_size, 0) + + # Simulate forward producing fresh SWA KV at the branching point. + self._insert( + cache, + allocator, + req_to_token_pool, + tokens[: result.swa_branching_seqlen], + ) + + rematch = cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ) + self.assertEqual(len(rematch.device_indices), result.swa_branching_seqlen) + self.assertIsNone(rematch.swa_branching_seqlen) + + def test_swa_branching_seqlen_uses_host_full_hit(self): + self._skip_swa_branching_on_rust() + if ( + not self.cfg.has_swa + or self.cfg.has_mamba + or self.cfg.page_size != 1 + or self.cfg.sliding_window_size != 4 + ): + self.skipTest("requires page_size=1 Full+SWA with window_size=4") + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + window = self.cfg.sliding_window_size + prefix = self._make_seq(1, window) + tokens = prefix + self._make_seq(1000, window + 1) + self._insert(cache, allocator, req_to_token_pool, prefix) + self._insert(cache, allocator, req_to_token_pool, tokens) + + leaf = cache.resolve_node_handle( + cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ).last_device_node + ) + parent = leaf.parent + self._backup_node(cache, leaf.id) + lock_result = cache.inc_lock_ref(parent.id) + try: + cache.evict(EvictParams(num_tokens=len(leaf.key))) + finally: + cache.dec_lock_ref(parent.id, lock_result.to_dec_params()) + device_frees = defaultdict(list) + host_frees = defaultdict(list) + cache.components[ComponentType.SWA].evict_component( + leaf, device_frees, host_frees, target=EvictLayer.HOST + ) + cache._free_values(device_frees, host_frees) + full_host_pool = cache.cache_controller.mem_pool_host + swa_host_pool = cache.components[ComponentType.SWA]._swa_kv_pool_host + full_available_before = full_host_pool.available_size() + swa_available_before = swa_host_pool.available_size() + + result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens)))) + + self.assertEqual(result.full_kv_hit_length, len(tokens)) + self.assertEqual(result.swa_branching_seqlen, len(tokens)) + + self._insert( + cache, + allocator, + req_to_token_pool, + tokens[: result.swa_branching_seqlen], + ) + cache.writing_check(write_back=True) + # Full was already backed up, so only the SWA window is allocated. + self.assertEqual(full_host_pool.available_size(), full_available_before) + self.assertEqual( + swa_host_pool.available_size(), + swa_available_before - len(leaf.key), + ) + + rematch = cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", tokens))) + ) + self.assertEqual(len(rematch.device_indices), result.swa_branching_seqlen) + self.assertIsNone(rematch.swa_branching_seqlen) + + def test_swa_branching_seqlen_caps_insert_after_forward(self): + if ( + not self.cfg.has_swa + or self.cfg.has_mamba + or self.cfg.page_size != 1 + or self.cfg.sliding_window_size != 4 + ): + self.skipTest("requires page_size=1 Full+SWA with window_size=4") + cache, _, _ = build_fixture(self.cfg) + swa = cache.components[ComponentType.SWA] + req = mock.Mock( + swa_branching_seqlen=8, + kv=mock.Mock(cache_protected_len=4, swa_evicted_seqlen=0), + ) + + for is_finished in (False, True): + params = InsertParams() + self.assertEqual( + swa.prepare_for_caching_req(req, params, 12, is_finished), 8 + ) + self.assertEqual(params.swa_evicted_seqlen, 0) + + self.assertIsNone(swa.prepare_for_caching_req(req, InsertParams(), 7, False)) + req.kv.cache_protected_len = 8 + self.assertIsNone(swa.prepare_for_caching_req(req, InsertParams(), 12, False)) + + cache.tree_core.is_eagle = True + req.kv.cache_protected_len = 4 + params = InsertParams() + self.assertEqual(swa.prepare_for_caching_req(req, params, 12, False), 9) + params.key = RadixKey(array("q", range(9)), is_bigram=True) + result = mock.Mock(swa_branch_inserted=False) + swa.commit_insert_component_data(mock.Mock(), False, params, result, []) + self.assertTrue(result.swa_branch_inserted) + + def test_swa_branch_insert_releases_forward_overshoot(self): + if ( + not self.cfg.has_swa + or self.cfg.has_mamba + or self.cfg.page_size != 1 + or self.cfg.sliding_window_size != 4 + ): + self.skipTest("requires page_size=1 Full+SWA with window_size=4") + + cache, allocator, req_to_token_pool = build_fixture(self.cfg) + swa = cache.components[ComponentType.SWA] + + branching_seqlen = 8 + forward_len = 20 + req = self._make_req(req_to_token_pool) + tokens = self._make_seq(1, forward_len) + req.origin_input_ids = tokens + req.output_ids = [] + req.full_untruncated_fill_ids = array("q", tokens) + req.set_extend_range(0, forward_len) + req.kv.cache_protected_len = branching_seqlen + + kv_indices = self._alloc(allocator, forward_len) + req_to_token_pool.write( + (req.kv.req_pool_idx, slice(0, forward_len)), + kv_indices, + ) + + params = InsertParams( + key=RadixKey(array("q", tokens[:branching_seqlen])), + swa_branching_seqlen=branching_seqlen, + ) + result = mock.Mock(swa_branch_inserted=False) + swa.commit_insert_component_data( + mock.Mock(), + False, + params, + result, + [], + ) + self.assertTrue(result.swa_branch_inserted) + + with envs.SGLANG_OPT_UNIFIED_CACHE_FREE_OUT_OF_WINDOW_SLOTS.override(True): + swa.cleanup_after_caching_req( + req, + is_finished=False, + insert_result=result, + ) + + expected_evicted = forward_len - 1 - self.cfg.sliding_window_size + self.assertEqual(req.kv.swa_evicted_seqlen, expected_evicted) + + mapping = allocator.full_to_swa_index_mapping + self.assertEqual( + torch.count_nonzero(mapping[kv_indices[:branching_seqlen]]).item(), + branching_seqlen, + ) + self.assertEqual( + torch.count_nonzero( + mapping[kv_indices[branching_seqlen:expected_evicted]] + ).item(), + 0, + ) + self.assertEqual( + torch.count_nonzero(mapping[kv_indices[expected_evicted:]]).item(), + forward_len - expected_evicted, + ) + def test_mamba_branching_seqlen_disabled_under_hicache(self): if not self.cfg.has_mamba or self.cfg.has_swa or self.cfg.page_size != 1: self.skipTest("requires page_size=1 Full+Mamba") @@ -6467,6 +6692,66 @@ class UnifiedRadixCacheSuite: self.assertEqual(kv_xfer.nodes_to_load, [b]) self.assertEqual(comp_xfers[ComponentType.SWA][0].nodes_to_load, [a, b]) + def test_hicache_swa_backup_window_stops_at_pending_ancestor(self): + self._skip_swa_branching_on_rust() + if ( + not self.cfg.has_swa + or self.cfg.has_mamba + or self.cfg.page_size != 1 + or self.cfg.sliding_window_size != 4 + ): + self.skipTest("requires page_size=1 Full+SWA with window_size=4") + + cache, allocator, req_to_token_pool = self._build_hicache_fixture() + chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 3) + if len(chain) < 3: + self.skipTest("chain collapsed below the pending ancestor test") + + c = chain[-1] + c_swa = _device_value(cache, c, ComponentType.SWA).clone() + + # First transfer: publish Full for C only, leaving SWA dirty while the + # write-through ack is still pending. + cache.tree_core.set_component_device_value_raw(c, ComponentType.SWA, None) + self.assertGreater( + cache._execute_and_commit_kv_backup(BackupKV(node_ids=[c])), + 0, + ) + self.assertEqual( + cache.tree_core.node_by_id(c).write_through_pending_id, + c, + ) + self.assertIsNotNone(_host_value(cache, c, ComponentType.FULL)) + self.assertIsNone(_host_value(cache, c, ComponentType.SWA)) + + # Simulate SWA being reconstructed on device before the first ack. The + # next incremental SWA backup must treat C as the boundary and back up + # only the newly inserted descendant. + cache.tree_core.set_component_device_value_raw(c, ComponentType.SWA, c_swa) + tokens = self._match_tokens_for_chain(cache, chain) + next_tokens = tokens + self._make_seq(9000, 1) + cache.write_through_threshold = 1 + + self._insert(cache, allocator, req_to_token_pool, next_tokens) + + d = cache.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", next_tokens))) + ).last_device_node + self.assertEqual( + cache.tree_core.node_by_id(c).write_through_pending_id, + c, + ) + self.assertEqual( + cache.tree_core.node_by_id(d).write_through_pending_id, + d, + ) + self.assertIsNone(_host_value(cache, c, ComponentType.SWA)) + self.assertIsNotNone(_host_value(cache, d, ComponentType.SWA)) + + cache.writing_check(write_back=True) + self.assertIsNone(cache.tree_core.node_by_id(c).write_through_pending_id) + self.assertIsNone(cache.tree_core.node_by_id(d).write_through_pending_id) + def _swa_finalize_setup(self): """Build a SWA chain long enough to fill at least the window plus one extra page, and host-back every node so we can flip