From 722e25a621300b9f6d8f16d51ad19eadc67626a1 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 9 Apr 2026 22:09:43 -0700 Subject: [PATCH] Fix SWA eviction boundary and page-align chunked prefill (#22470) --- python/sglang/srt/managers/schedule_batch.py | 8 +- .../sglang/srt/mem_cache/swa_radix_cache.py | 20 + .../mem_cache/test_swa_eviction_boundary.py | 392 ++++++++++++++++++ 3 files changed, 419 insertions(+), 1 deletion(-) create mode 100644 test/registered/unit/mem_cache/test_swa_eviction_boundary.py diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 078d6ab91..ddddbcef1 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2515,8 +2515,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): ), "cache_protected_len must be page aligned" req.swa_evicted_seqlen = max(req.swa_evicted_seqlen, req.cache_protected_len) + # Subtract an extra page_size so the eviction frontier never reaches the + # radix tree insert boundary (page_floor(seq_len)). This keeps at least one + # page of non-evicted SWA KV for the tree to store as a non-tombstone node, + # preserving cache reuse in multi-turn scenarios. + # See also: _insert_helper case 3 in swa_radix_cache.py (defensive counterpart). new_swa_evicted_seqlen = max( - req.swa_evicted_seqlen, pre_len - sliding_window_size + req.swa_evicted_seqlen, + pre_len - sliding_window_size - self.tree_cache.page_size, ) if self.tree_cache.page_size > 1: diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 34df1617d..6e1b1c847 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/python/sglang/srt/mem_cache/swa_radix_cache.py @@ -1013,6 +1013,26 @@ class SWARadixCache(BasePrefixCache): child_key = self.get_child_key_fn(key) if len(key): + # Layout: |--- total_prefix_length ---|--- len(key) ---| + # ^ ^ ^ + # 0 total_prefix_length total_length + # + # Cases based on swa_evicted_seqlen position: + # 1. swa_evicted_seqlen <= total_prefix_length: + # Already handled in the while loop above. All of len(key) is non-tombstone. + # 2. total_prefix_length < swa_evicted_seqlen < total_length: + # Split: [total_prefix_length, swa_evicted_seqlen) as tombstone, + # [swa_evicted_seqlen, total_length) as non-tombstone. + # 3. swa_evicted_seqlen == total_length: + # All remaining tokens are evicted. Free value and return without + # creating a node (leaf nodes must not be tombstone). + # Note: the -page_size fix in _evict_swa prevents this case from + # occurring in normal operation. This check is a defensive guard + # against unexpected eviction states from other code paths. + if swa_evicted_seqlen == total_prefix_length + len(key): + self.token_to_kv_pool_allocator.free(value) + return total_prefix_length + if ( swa_evicted_seqlen > total_prefix_length and swa_evicted_seqlen < total_prefix_length + len(key) diff --git a/test/registered/unit/mem_cache/test_swa_eviction_boundary.py b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py new file mode 100644 index 000000000..b4c302261 --- /dev/null +++ b/test/registered/unit/mem_cache/test_swa_eviction_boundary.py @@ -0,0 +1,392 @@ +"""Unit tests for SWA eviction boundary fixes. + +Bug: when page_size > sliding_window_size, _evict_swa could advance the +eviction frontier to exactly page_floor(seq_len), making all tokens being +inserted into the radix tree fully evicted (case 3). _insert_helper had no +handling for this, creating an incorrect non-tombstone node that caused +inflated swa_evictable_size_, negative usage, and potential double-free. + +Two-sided fix: +1. _evict_swa subtracts extra page_size (preventive). +2. _insert_helper early-returns on case 3 (defensive). + +Tests use real tree/allocator/pool with mock Req/ScheduleBatch wrappers. +""" + +import unittest +from types import SimpleNamespace + +import torch + +from sglang.srt.managers.schedule_batch import ScheduleBatch +from sglang.srt.mem_cache.cache_init_params import CacheInitParams +from sglang.srt.mem_cache.memory_pool import ReqToTokenPool +from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator +from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache +from sglang.srt.utils import get_device +from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci + +register_cuda_ci(est_time=8, suite="stage-b-test-1-gpu-large") +register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd") + +# --------------------------------------------------------------------------- +# Infrastructure helpers (shared setup, not logic) +# --------------------------------------------------------------------------- + + +def _swa_alloc(allocator, need_size): + """Allocate from SWA allocator for any page_size. + + SWATokenToKVPoolAllocator.alloc() asserts page_size == 1. For page_size > 1, + allocate from the underlying paged allocators directly and set up the mapping. + """ + if allocator.page_size == 1: + return allocator.alloc(need_size) + + if need_size > allocator.full_attn_allocator.available_size(): + return None + if need_size > allocator.swa_attn_allocator.available_size(): + return None + + full_indices = allocator.full_attn_allocator.alloc(need_size) + swa_indices = allocator.swa_attn_allocator.alloc(need_size) + assert full_indices is not None and swa_indices is not None + allocator.full_to_swa_index_mapping[full_indices] = swa_indices + return full_indices + + +def _build_swa_tree(page_size, sliding_window_size, kv_size=1024, kv_size_swa=512): + head_num, head_dim, num_layers, global_interval = 8, 128, 24, 4 + dtype = torch.bfloat16 + device = get_device() + full_ids = list(range(0, num_layers, global_interval)) + swa_ids = [i for i in range(num_layers) if i not in set(full_ids)] + + pool = ReqToTokenPool( + size=8, max_context_len=2048, device=device, enable_memory_saver=False + ) + kv_pool = SWAKVPool( + size=kv_size, + size_swa=kv_size_swa, + page_size=page_size, + dtype=dtype, + head_num=head_num, + head_dim=head_dim, + swa_attention_layer_ids=swa_ids, + full_attention_layer_ids=full_ids, + enable_kvcache_transpose=False, + device=device, + ) + allocator = SWATokenToKVPoolAllocator( + size=kv_size, + size_swa=kv_size_swa, + page_size=page_size, + dtype=dtype, + device=device, + kvcache=kv_pool, + need_sort=False, + ) + tree = SWARadixCache( + params=CacheInitParams( + req_to_token_pool=pool, + token_to_kv_pool_allocator=allocator, + page_size=page_size, + disable=False, + is_eagle=False, + sliding_window_size=sliding_window_size, + ), + ) + return tree, allocator, pool + + +def _make_req(req_pool_idx, token_ids, cache_protected_len, tree): + """Mock Req with fields needed by _evict_swa and cache_finished_req.""" + req = SimpleNamespace( + req_pool_idx=req_pool_idx, + origin_input_ids=token_ids, + output_ids=[], + cache_protected_len=cache_protected_len, + swa_evicted_seqlen=0, + extra_key=None, + last_node=tree.root_node, + swa_uuid_for_lock=None, + prefix_indices=torch.tensor([], dtype=torch.int64, device=tree.device), + _kv_committed_len=len(token_ids), + ) + req.pop_committed_kv_cache = lambda: req._kv_committed_len + return req + + +def _make_batch(tree, allocator, pool): + """Mock ScheduleBatch with fields needed by _evict_swa.""" + return SimpleNamespace( + tree_cache=tree, + req_to_token_pool=pool, + token_to_kv_pool_allocator=allocator, + ) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +class TestSWAEvictionBoundary(unittest.TestCase): + + # -- Eviction formula: page_size > window -- + + def test_formula_page_gt_window_sweep(self): + """Sweep page_size > window combinations. The -page_size fix must + prevent eviction from reaching page_floor(seq_len).""" + for page_size in [4, 8, 16, 32, 64, 128, 256]: + for window in [1, 2, 4, 8]: + if page_size <= window: + continue + tree, allocator, pool = _build_swa_tree( + page_size=page_size, + sliding_window_size=window, + kv_size=max(4096, page_size * 20), + kv_size_swa=max(2048, page_size * 10), + ) + for seq_len in range(page_size + 1, page_size * 5): + alloc_size = (seq_len + page_size - 1) // page_size * page_size + kv = _swa_alloc(allocator, alloc_size) + if kv is None: + break + pool.write((0, slice(0, alloc_size)), kv) + + req = _make_req(0, list(range(seq_len)), 0, tree) + batch = _make_batch(tree, allocator, pool) + ScheduleBatch._evict_swa(batch, req, seq_len - 1) + + insert_len = seq_len // page_size * page_size + self.assertLess( + req.swa_evicted_seqlen, + insert_len, + f"page={page_size}, win={window}, seq={seq_len}", + ) + allocator.free(kv) + + # -- Eviction formula: page_size <= window -- + + def test_formula_page_leq_window(self): + """page_size <= window: -page_size fix causes no regression.""" + page_size, window = 4, 8 + tree, allocator, pool = _build_swa_tree( + page_size=page_size, sliding_window_size=window + ) + + for seq_len in [13, 17, 25, 33]: + alloc_size = (seq_len + page_size - 1) // page_size * page_size + kv = _swa_alloc(allocator, alloc_size) + pool.write((0, slice(0, alloc_size)), kv) + + req = _make_req(0, list(range(seq_len)), 0, tree) + batch = _make_batch(tree, allocator, pool) + ScheduleBatch._evict_swa(batch, req, seq_len - 1) + + insert_len = seq_len // page_size * page_size + self.assertLess(req.swa_evicted_seqlen, insert_len) + + tree.cache_finished_req(req, is_insert=True) + tree.sanity_check() + + # -- Eviction formula: page_size == 1 -- + + def test_formula_page_size_1(self): + """page_size=1: no floor alignment, -1 just means one less token evicted.""" + page_size, window = 1, 4 + tree, allocator, pool = _build_swa_tree( + page_size=page_size, sliding_window_size=window + ) + + for seq_len in range(window + 2, 30): + kv = _swa_alloc(allocator, seq_len) + pool.write((0, slice(0, seq_len)), kv) + + req = _make_req(0, list(range(seq_len)), 0, tree) + batch = _make_batch(tree, allocator, pool) + ScheduleBatch._evict_swa(batch, req, seq_len - 1) + + self.assertLess(req.swa_evicted_seqlen, seq_len) + self.assertEqual(req.swa_evicted_seqlen, max(0, seq_len - 1 - window - 1)) + + tree.cache_finished_req(req, is_insert=True) + tree.sanity_check() + + # -- Eviction formula: no-op when seq too short -- + + def test_formula_noop_short_sequence(self): + """pre_len - window - page_size < 0: eviction stays at 0.""" + page_size, window = 8, 4 + tree, allocator, pool = _build_swa_tree( + page_size=page_size, sliding_window_size=window + ) + + seq_len = page_size + window - 2 # = 10, formula gives 10-1-4-8 = -3 + alloc_size = (seq_len + page_size - 1) // page_size * page_size + kv = _swa_alloc(allocator, alloc_size) + pool.write((0, slice(0, alloc_size)), kv) + + req = _make_req(0, list(range(seq_len)), 0, tree) + batch = _make_batch(tree, allocator, pool) + ScheduleBatch._evict_swa(batch, req, seq_len - 1) + + self.assertEqual(req.swa_evicted_seqlen, 0) + + # -- Insert case 1: swa_evicted <= total_prefix_length -- + + def test_insert_case1_evicted_within_matched(self): + """Eviction within matched region. New tokens all non-tombstone.""" + page_size, window = 8, 2 + tree, allocator, pool = _build_swa_tree( + page_size=page_size, sliding_window_size=window + ) + + # First request: populate tree with 16 tokens (2 pages) + first_len = page_size * 2 + kv1 = _swa_alloc(allocator, first_len) + pool.write((0, slice(0, first_len)), kv1) + req1 = _make_req(0, list(range(first_len)), 0, tree) + tree.cache_finished_req(req1, is_insert=True) + tree.sanity_check() + + # Second request: 24 tokens, first 16 overlap with tree + second_len = page_size * 3 + kv2 = _swa_alloc(allocator, second_len) + pool.write((1, slice(0, second_len)), kv2) + + req2 = _make_req(1, list(range(second_len)), 0, tree) + batch = _make_batch(tree, allocator, pool) + + # pre_len=15: 15-2-8=5, floor to 8 -> 0. Eviction stays within matched. + ScheduleBatch._evict_swa(batch, req2, first_len - 1) + self.assertLessEqual(req2.swa_evicted_seqlen, first_len) + + swa_evictable_before = tree.swa_evictable_size_ + tree.cache_finished_req(req2, is_insert=True) + + # New tokens [16, 24) should all be non-tombstone + new_tokens = second_len // page_size * page_size - first_len + self.assertEqual(tree.swa_evictable_size_, swa_evictable_before + new_tokens) + tree.sanity_check() + + # -- Insert case 2: total_prefix_length < swa_evicted < total_length -- + + def test_insert_case2_partial_tombstone(self): + """Partial eviction: some tombstone, some non-tombstone.""" + page_size, window = 8, 2 + tree, allocator, pool = _build_swa_tree( + page_size=page_size, sliding_window_size=window + ) + + # seq_len=25: insert_length=24, evicted should be 8 (1 page) + seq_len = page_size * 3 + 1 + alloc_size = (seq_len + page_size - 1) // page_size * page_size + kv = _swa_alloc(allocator, alloc_size) + pool.write((0, slice(0, alloc_size)), kv) + + req = _make_req(0, list(range(seq_len)), 0, tree) + batch = _make_batch(tree, allocator, pool) + swa_evictable_before = tree.swa_evictable_size_ + + ScheduleBatch._evict_swa(batch, req, seq_len - 1) + insert_len = seq_len // page_size * page_size + self.assertGreater(req.swa_evicted_seqlen, 0, "Should have some eviction") + self.assertLess(req.swa_evicted_seqlen, insert_len, "Should be partial") + + tree.cache_finished_req(req, is_insert=True) + + non_tombstone = insert_len - req.swa_evicted_seqlen + self.assertEqual(tree.swa_evictable_size_, swa_evictable_before + non_tombstone) + self.assertGreater(tree.full_evictable_size_, 0) + tree.sanity_check() + + # -- Insert case 3: swa_evicted == total_length (defensive) -- + + def test_insert_case3_defensive_early_return(self): + """Simulate OLD formula to trigger case 3. Defensive early return + must prevent non-tombstone node creation.""" + page_size, window = 8, 2 + tree, allocator, pool = _build_swa_tree( + page_size=page_size, sliding_window_size=window + ) + + # seq_len=11: OLD formula -> evicted=8 == insert_length=8 + seq_len = page_size + window + 1 + alloc_size = (seq_len + page_size - 1) // page_size * page_size + kv = _swa_alloc(allocator, alloc_size) + pool.write((0, slice(0, alloc_size)), kv) + + # OLD formula (without -page_size) + pre_len = seq_len - 1 + old_evicted = max(0, (pre_len - window) // page_size * page_size) + insert_len = seq_len // page_size * page_size + self.assertEqual( + old_evicted, insert_len, "Precondition: old formula hits boundary" + ) + + # Manually free SWA as _evict_swa would + allocator.free_swa(pool.req_to_token[0, :old_evicted]) + + req = _make_req(0, list(range(seq_len)), 0, tree) + req.swa_evicted_seqlen = old_evicted + swa_evictable_before = tree.swa_evictable_size_ + + tree.cache_finished_req(req, is_insert=True) + + self.assertEqual(tree.swa_evictable_size_, swa_evictable_before) + + # -- Integration: multiple decode turns -- + + def test_multiple_decodes(self): + """Multiple decode turns with page_size > window. No over-eviction, + tree stays consistent throughout.""" + page_size, window = 8, 2 + tree, allocator, pool = _build_swa_tree( + page_size=page_size, sliding_window_size=window + ) + + for turn in range(4): + seq_len = page_size * (turn + 2) + 1 + idx = turn % pool.size + alloc_size = (seq_len + page_size - 1) // page_size * page_size + kv = _swa_alloc(allocator, alloc_size) + assert kv is not None, f"Alloc failed at turn {turn}" + pool.write((idx, slice(0, alloc_size)), kv) + + req = _make_req(idx, list(range(seq_len)), 0, tree) + batch = _make_batch(tree, allocator, pool) + ScheduleBatch._evict_swa(batch, req, seq_len - 1) + + insert_len = seq_len // page_size * page_size + self.assertLess(req.swa_evicted_seqlen, insert_len, f"turn {turn}") + + tree.cache_finished_req(req, is_insert=True) + tree.sanity_check() + + # -- Integration: page_size=1 full flow -- + + def test_page_size_1_full_flow(self): + """End-to-end with page_size=1. Fix is near no-op.""" + page_size, window = 1, 4 + tree, allocator, pool = _build_swa_tree( + page_size=page_size, sliding_window_size=window + ) + + for seq_len in [10, 20, 30]: + kv = _swa_alloc(allocator, seq_len) + pool.write((0, slice(0, seq_len)), kv) + + req = _make_req(0, list(range(seq_len)), 0, tree) + batch = _make_batch(tree, allocator, pool) + ScheduleBatch._evict_swa(batch, req, seq_len - 1) + + self.assertEqual(req.swa_evicted_seqlen, max(0, seq_len - 1 - window - 1)) + + tree.cache_finished_req(req, is_insert=True) + tree.sanity_check() + + +if __name__ == "__main__": + unittest.main()