Split SWA leaf to one window on insert (#26919)

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

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)
+        # Each 8-page segment is cap-split into [prefix, tail]; the tail leads
-        side_node, c_node, b_node, a_node = pre
+        # 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
+        # New segment E adds two nodes (prefix + tail).
-        self.assertEqual(len(post), 5)
+        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)
+        # Each 8-page segment is cap-split into [prefix, tail]; tails lead in
-        side_node, c_node, b_node, a_node = pre
+        # 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")
+        # Matching seq_abc refreshes only the window cushion (C's capped nodes)
-        self.assertIs(
+        # to the MRU side; out-of-cushion ancestors B and A keep their order.
-            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]}",
-        )
         self.assertIn(
-            side_node,
+            c_node, post[:2], f"C must be refreshed to MRU; post={[n.id for n in post]}"
-            post[:2],
+        )
-            "Side branch must NOT be pushed below ancestors after deep match; "
+        self.assertLess(
-            f"got post={[n.id for n in post]}",
+            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)
+        self.assertEqual(len(pre), 8)
-        side_node, c_node, b_node, a_node = pre
+        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)
+        # Under leaf-cap no single node exceeds the cushion; it spans C's capped
-        self.assertIs(post[0], c_node, "C alone exhausts cushion → only C refreshed")
+        # tail plus its prefix, so both of C's nodes are refreshed to the MRU
-        # B and A: untouched ordering relative to each other AND to side_node
+        # 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()