From f59bbef841705ef42e07ea360ed9d9fd0b443906 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Mon, 1 Jun 2026 23:46:55 +0800 Subject: [PATCH] Split SWA leaf to one window on insert (#26919) --- .../unified_cache_components/swa_component.py | 26 ++++ .../test_unified_radix_cache_unittest.py | 111 +++++++++++++----- 2 files changed, 108 insertions(+), 29 deletions(-) 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 415be271f..9f70db025 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 @@ -261,6 +261,32 @@ class SWAComponent(TreeComponent): node.component_data[self.component_type].value = swa_value self.cache.lru_lists[self.component_type].insert_mru(node) self.cache.component_evictable_size_[self.component_type] += len(swa_value) + else: + # Entire leaf is outside the SWA window — left as a tombstone. + return + + self._maybe_split_leaf_for_swa_lock(node) + + def _maybe_split_leaf_for_swa_lock(self, leaf: UnifiedTreeNode) -> None: + """Cap a fresh SWA leaf at one page-aligned window so locking it pins + only one window of SWA pool, not the whole (long chunked-prefill) leaf. + """ + ct = self.component_type + cd = leaf.component_data[ct] + if leaf is self.cache.root_node or cd.value is None or cd.lock_ref > 0: + return + + page_size = self.cache.page_size + # Smallest page-aligned size that still covers the sliding window. + tail_size = (self.sliding_window_size + page_size - 1) // page_size * page_size + leaf_len = len(leaf.key) + if leaf_len <= tail_size: + return + split_at = leaf_len - tail_size + if page_size > 1 and (split_at % page_size != 0 or leaf_len % page_size != 0): + return + + self.cache._split_node(leaf.key, leaf, split_at) def redistribute_on_node_split( self, new_parent: UnifiedTreeNode, child: UnifiedTreeNode 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 9eb4cdc6c..dc0110242 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 @@ -1114,6 +1114,51 @@ class UnifiedRadixCacheSuite: ) tree.sanity_check() + def test_swa_leaf_capped_to_window_on_insert(self): + """A long SWA leaf is split so locking it protects one window of SWA + while full attention still protects the whole sequence.""" + if not self.cfg.has_swa: + self.skipTest("requires SWA component") + + ps = self.cfg.page_size + window = self.cfg.sliding_window_size + tail_size = ((window + ps - 1) // ps) * ps + tail_pages = tail_size // ps + + for case in ("long_splits", "short_keeps"): + with self.subTest(case=case): + tree, allocator, req_to_token_pool = build_fixture(self.cfg) + num_pages = tail_pages + 2 if case == "long_splits" else tail_pages + seq = self._make_seq(1, num_pages) + self._insert(tree, allocator, req_to_token_pool, seq) + tree.sanity_check() + + leaf = tree.match_prefix( + MatchPrefixParams(key=RadixKey(array("q", seq))) + ).last_device_node + swa_val = leaf.component_data[ComponentType.SWA].value + self.assertIsNotNone(swa_val) + + if case == "long_splits": + # Capped to one page-aligned window; prefix is a real ancestor. + self.assertEqual(len(swa_val), tail_size) + self.assertIsNot(leaf.parent, tree.root_node) + else: + # Already within one window — no split. + self.assertEqual(len(swa_val), len(seq)) + self.assertIs(leaf.parent, tree.root_node) + + lock_result = tree.inc_lock_ref(leaf) + # SWA pins one window; full attention pins everything. + self.assertEqual(tree.swa_protected_size(), len(swa_val)) + self.assertEqual(tree.full_protected_size(), len(seq)) + tree.sanity_check() + tree.dec_lock_ref( + leaf, + DecLockRefParams(swa_uuid_for_lock=lock_result.swa_uuid_for_lock), + ) + tree.sanity_check() + def _swa_lru_order(self, tree): lru = tree.lru_lists[ComponentType.SWA] pt = lru._pt @@ -1153,29 +1198,29 @@ class UnifiedRadixCacheSuite: self._insert(tree, allocator, req_to_token_pool, seq_side) pre = self._swa_lru_order(tree) - self.assertEqual(len(pre), 4) - side_node, c_node, b_node, a_node = pre + # Each 8-page segment is cap-split into [prefix, tail]; the tail leads + # the pair in MRU order, so segment tails sit at even indices. + self.assertEqual(len(pre), 8) + side_node, c_node, b_node, a_node = pre[0], pre[2], pre[4], pre[6] seq_abcd = seq_abc + self._make_seq(300, 8) self._insert(tree, allocator, req_to_token_pool, seq_abcd) post = self._swa_lru_order(tree) - # new leaf E exists now, length 5 - self.assertEqual(len(post), 5) + # New segment E adds two nodes (prefix + tail). + self.assertEqual(len(post), 10) # side branch must still appear BEFORE B and A in MRU->LRU order: - # bounded refresh on new leaf E (size=8 >= cushion=5) refreshes only E. + # walk-down on the new segment must not refresh old ancestors. side_pos = post.index(side_node) - b_pos = post.index(b_node) - a_pos = post.index(a_node) self.assertLess( side_pos, - b_pos, + post.index(b_node), f"side branch must remain ahead of B (no walk-down refresh); " f"post={[n.id for n in post]}, side={side_node.id}, B={b_node.id}", ) self.assertLess( side_pos, - a_pos, + post.index(a_node), f"side branch must remain ahead of A (no walk-down refresh); " f"post={[n.id for n in post]}, side={side_node.id}, A={a_node.id}", ) @@ -1197,26 +1242,30 @@ class UnifiedRadixCacheSuite: self._insert(tree, allocator, req_to_token_pool, seq_side) pre = self._swa_lru_order(tree) - self.assertEqual(len(pre), 4) - side_node, c_node, b_node, a_node = pre + # Each 8-page segment is cap-split into [prefix, tail]; tails lead in + # MRU order, so segment tails sit at even indices. + self.assertEqual(len(pre), 8) + side_node, c_node, b_node, a_node = pre[0], pre[2], pre[4], pre[6] m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_abc)))) self.assertEqual(len(m.device_indices), len(seq_abc)) post = self._swa_lru_order(tree) - self.assertIs(post[0], c_node, "C (last matched node) must be MRU") - self.assertIs( - post[-1], - a_node, - "Oldest out-of-cushion ancestor A must remain at LRU tail; " - f"got post={[n.id for n in post]}, " - f"pre={[n.id for n in pre]}", - ) + # Matching seq_abc refreshes only the window cushion (C's capped nodes) + # to the MRU side; out-of-cushion ancestors B and A keep their order. self.assertIn( - side_node, - post[:2], - "Side branch must NOT be pushed below ancestors after deep match; " - f"got post={[n.id for n in post]}", + c_node, post[:2], f"C must be refreshed to MRU; post={[n.id for n in post]}" + ) + self.assertLess( + post.index(side_node), + post.index(b_node), + "Side branch must NOT be pushed below ancestors after deep match", + ) + self.assertLess(post.index(b_node), post.index(a_node), "B must stay above A") + self.assertGreaterEqual( + post.index(a_node), + len(post) - 2, + "Oldest out-of-cushion ancestor A must stay at the LRU-tail end", ) tree.sanity_check() @@ -1272,20 +1321,24 @@ class UnifiedRadixCacheSuite: self._insert(tree, allocator, req_to_token_pool, seq_side) pre = self._swa_lru_order(tree) - self.assertEqual(len(pre), 4) - side_node, c_node, b_node, a_node = pre + self.assertEqual(len(pre), 8) + side_node, c_node, b_node, a_node = pre[0], pre[2], pre[4], pre[6] + c_prefix = pre[3] # C's prefix pairs with its tail (c_node) at pre[2:4] m = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq_abc)))) self.assertEqual(len(m.device_indices), len(seq_abc)) post = self._swa_lru_order(tree) cushion = self.cfg.sliding_window_size + self.cfg.page_size - self.assertGreaterEqual(len(c_node.key), cushion) - self.assertIs(post[0], c_node, "C alone exhausts cushion → only C refreshed") - # B and A: untouched ordering relative to each other AND to side_node + # Under leaf-cap no single node exceeds the cushion; it spans C's capped + # tail plus its prefix, so both of C's nodes are refreshed to the MRU + # side while B and A keep their relative order below. + self.assertLess(len(c_node.key), cushion) + self.assertIn(c_node, post[:2]) + self.assertIn(c_prefix, post[:2]) + side_pos = post.index(side_node) b_pos = post.index(b_node) a_pos = post.index(a_node) - side_pos = post.index(side_node) self.assertLess(side_pos, b_pos, "B was below side in pre, must stay below") self.assertLess(b_pos, a_pos, "A was below B in pre, must stay below") tree.sanity_check()