minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Ke Bao
2026-06-01 23:46:55 +08:00
committed by GitHub
parent d8a5a25c36
commit f59bbef841
2 changed files with 108 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/unified_cache_components/swa_component.py
+26
View File
@@ -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
+82 -29
View File
@@ -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()