[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>

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(

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.

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