/test/registered/unit/mem_cache/test_swa_unittest.py

import unittest
from array import array

import torch

from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored
from sglang.srt.environ import envs
from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
    DecLockRefParams,
    EvictParams,
    EvictResult,
    InsertParams,
    MatchPrefixParams,
)
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.common import available_and_evictable_str
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
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
from sglang.test.test_utils import CustomTestCase

register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-large")
register_amd_ci(est_time=10, suite="stage-b-test-1-gpu-small-amd")


def _event_hashes(events):
    return [block_hash for event in events for block_hash in event.block_hashes]


class _DummyReq:
    def __init__(self):
        self._kv_committed_len = 0
        self.swa_prefix_lock_released = False

    def pop_committed_kv_cache(self):
        return self._kv_committed_len


def _build_swa_tree(
    is_eagle: bool,
    page_size: int = 1,
    req_size: int = 8,
    max_context_len: int = 64,
    kv_size: int = 64,
    kv_size_swa: int = 32,
    sliding_window_size: int = 4,
    enable_kv_cache_events: bool = False,
):
    head_num = 8
    head_dim = 128
    num_layers = 24
    global_interval = 4
    dtype = torch.bfloat16
    device = get_device()
    full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)]
    full_attention_layer_ids_set = set(full_attention_layer_ids)
    swa_attention_layer_ids = [
        i for i in range(num_layers) if i not in full_attention_layer_ids_set
    ]

    req_to_token_pool = ReqToTokenPool(
        size=req_size,
        max_context_len=max_context_len,
        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_attention_layer_ids,
        full_attention_layer_ids=full_attention_layer_ids,
        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=req_to_token_pool,
            token_to_kv_pool_allocator=allocator,
            page_size=page_size,
            disable=False,
            is_eagle=is_eagle,
            sliding_window_size=sliding_window_size,
            enable_kv_cache_events=enable_kv_cache_events,
        ),
    )
    return tree, allocator, req_to_token_pool


def _swa_alloc(allocator, need_size):
    """SWA-pool alloc that also works for page_size > 1 (built-in alloc asserts page_size == 1)."""
    if allocator.page_size == 1:
        return allocator.alloc(need_size)

    assert need_size % allocator.page_size == 0
    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 _insert(tree, allocator, token_ids):
    indices = _swa_alloc(allocator, len(token_ids))
    assert indices is not None
    tree.insert(InsertParams(key=RadixKey(array("q", token_ids)), value=indices))


def _insert_chain(tree, allocator, token_ids):
    _insert(tree, allocator, token_ids)
    match = tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", token_ids))))
    return match.last_device_node


def _expected_tail_size(window: int, page_size: int) -> int:
    """Mirror of _maybe_split_leaf_for_swa_lock's tail_size formula."""
    return (window + page_size - 1) // page_size * page_size


class TestSWA(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        pass

    @classmethod
    def tearDownClass(cls):
        pass

    def test_swa_radix_cache_kv_events(self):
        tree, allocator, _ = _build_swa_tree(
            is_eagle=False, enable_kv_cache_events=True
        )
        tree.take_events()  # Clear the reset event.

        _insert(tree, allocator, [1, 2, 3, 4])
        first_insert_events = [
            e for e in tree.take_events() if isinstance(e, BlockStored)
        ]
        self.assertEqual(len(first_insert_events), 4)
        self.assertEqual([e.token_ids[0] for e in first_insert_events], [1, 2, 3, 4])

        _insert(tree, allocator, [1, 2, 3, 4, 5, 6])
        second_insert_events = [
            e for e in tree.take_events() if isinstance(e, BlockStored)
        ]
        self.assertEqual(len(second_insert_events), 2)
        self.assertEqual([e.token_ids[0] for e in second_insert_events], [5, 6])

        stored_hashes = [
            e.block_hashes[0] for e in first_insert_events + second_insert_events
        ]

        # Evicting only SWA tokens tombstones nodes but keeps full KV blocks.
        result = tree.evict(EvictParams(num_tokens=0, swa_num_tokens=1))
        self.assertEqual(result.num_tokens_evicted, 0)
        self.assertGreaterEqual(result.swa_num_tokens_evicted, 1)
        self.assertEqual(
            [e for e in tree.take_events() if isinstance(e, BlockRemoved)], []
        )

        result = tree.evict(EvictParams(num_tokens=1, swa_num_tokens=0))
        self.assertGreaterEqual(result.num_tokens_evicted, 1)
        removed_hashes = _event_hashes(
            [e for e in tree.take_events() if isinstance(e, BlockRemoved)]
        )
        self.assertCountEqual(removed_hashes, stored_hashes)

    def test_swa_radix_cache_kv_events_split_hash(self):
        tree, allocator, _ = _build_swa_tree(
            is_eagle=False, enable_kv_cache_events=True
        )
        tree.take_events()  # Clear the reset event.

        _insert(tree, allocator, [1, 2, 3, 4])
        first_insert_events = [
            e for e in tree.take_events() if isinstance(e, BlockStored)
        ]
        self.assertEqual(len(first_insert_events), 4)
        split_parent_hash = first_insert_events[1].block_hashes[0]

        _insert(tree, allocator, [1, 2, 5, 6])
        second_insert_events = [
            e for e in tree.take_events() if isinstance(e, BlockStored)
        ]
        self.assertEqual(len(second_insert_events), 2)
        self.assertEqual(list(second_insert_events[0].token_ids), [5])
        self.assertEqual(second_insert_events[0].parent_block_hash, split_parent_hash)

    def test_swa_memory_pool_paged_free_clears_full_page_mapping(self):
        page_size = 4
        _, allocator, _ = _build_swa_tree(
            is_eagle=False,
            page_size=page_size,
            kv_size=16,
            kv_size_swa=16,
            sliding_window_size=page_size,
        )

        full_indices = _swa_alloc(allocator, page_size)
        self.assertEqual(allocator.swa_available_size(), 16 - page_size)

        allocator.free_swa(full_indices[:1])
        self.assertEqual(allocator.swa_available_size(), 16)
        self.assertTrue(
            torch.all(
                allocator.full_to_swa_index_mapping[full_indices.to(torch.int64)] == 0
            )
        )

        allocator.free_swa(full_indices[1:2])
        self.assertEqual(allocator.swa_available_size(), 16)

    def test_swa_radix_cache_1(self):
        # args
        req_size = 10
        max_context_len = 128
        kv_size = 128
        kv_size_swa = 64
        page_size = 1
        sliding_window_size = 4
        head_num = 8
        head_dim = 128
        num_layers = 48
        global_interval = 4
        dtype = torch.bfloat16
        device = get_device()
        full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)]
        full_attention_layer_ids_set = set(full_attention_layer_ids)
        swa_attention_layer_ids = [
            i for i in range(num_layers) if i not in full_attention_layer_ids_set
        ]
        # setup req to token pool
        req_to_token_pool = ReqToTokenPool(
            size=req_size,
            max_context_len=max_context_len,
            device=device,
            enable_memory_saver=False,
        )
        # setup kv pool
        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_attention_layer_ids,
            full_attention_layer_ids=full_attention_layer_ids,
            device=device,
        )
        # setup token to kv pool allocator
        allocator = SWATokenToKVPoolAllocator(
            size=kv_size,
            size_swa=kv_size_swa,
            page_size=page_size,
            dtype=dtype,
            device=device,
            kvcache=kv_pool,
            need_sort=False,
        )
        # setup radix cache
        tree = SWARadixCache(
            params=CacheInitParams(
                req_to_token_pool=req_to_token_pool,
                token_to_kv_pool_allocator=allocator,
                disable=False,
                page_size=page_size,
                sliding_window_size=sliding_window_size,
            ),
        )

        # test
        print(
            f"[Start] allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
        )
        req1_token_ids, req1_kv_indices = [1, 2, 3], allocator.alloc(3)
        self.assertEqual(len(req1_token_ids), len(req1_kv_indices))
        print(
            f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}"
        )
        key = RadixKey(array("q", req1_token_ids))
        result = tree.insert(InsertParams(key=key, value=req1_kv_indices[: len(key)]))
        prefix_len = result.prefix_len
        print(
            f"req1: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
        )
        req2_token_ids, req2_kv_indices = [1, 2, 3, 4, 5, 6, 7], allocator.alloc(7)
        self.assertEqual(len(req2_token_ids), len(req2_kv_indices))
        print(
            f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}"
        )
        key = RadixKey(array("q", req2_token_ids))
        result = tree.insert(InsertParams(key=key, value=req2_kv_indices[: len(key)]))
        prefix_len = result.prefix_len
        print(
            f"req2: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
        )
        req3_token_ids, req3_kv_indices = [10, 11, 12], allocator.alloc(3)
        self.assertEqual(len(req3_token_ids), len(req3_kv_indices))
        print(
            f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}"
        )
        key = RadixKey(array("q", req3_token_ids))
        result = tree.insert(InsertParams(key=key, value=req3_kv_indices[: len(key)]))
        prefix_len = result.prefix_len
        print(
            f"req3: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
        )
        req4_token_ids, req4_kv_indices = [1, 2, 3, 4, 5, 60, 70], allocator.alloc(7)
        self.assertEqual(len(req4_token_ids), len(req4_kv_indices))
        print(
            f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}"
        )
        key = RadixKey(array("q", req4_token_ids))
        result = tree.insert(InsertParams(key=key, value=req4_kv_indices[: len(key)]))
        prefix_len = result.prefix_len
        print(
            f"req4: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
        )

        tree.pretty_print()
        full_num_tokens, swa_num_tokens = 1, 0
        print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token")
        tree.evict(
            EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
        )
        tree.pretty_print()

        full_num_tokens, swa_num_tokens = 0, 1
        print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token")
        tree.evict(
            EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
        )
        tree.pretty_print()

        full_num_tokens, swa_num_tokens = 1, 2
        print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token")
        tree.evict(
            EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
        )
        tree.pretty_print()

        req5_token_ids = [1, 2, 3, 4, 5]
        result = tree.match_prefix(
            MatchPrefixParams(key=RadixKey(array("q", req5_token_ids)))
        )
        kv_indices, last_node = result.device_indices, result.last_device_node
        print(
            f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
        )
        self.assertEqual(len(kv_indices), 0)

        req6_token_ids = [1, 2, 3, 4, 5, 60, 70]
        result = tree.match_prefix(
            MatchPrefixParams(key=RadixKey(array("q", req6_token_ids)))
        )
        kv_indices, last_node = result.device_indices, result.last_device_node
        print(
            f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
        )
        self.assertEqual(len(kv_indices), 7)
        self.assertEqual(len(last_node.key), 2)
        self.assertEqual(last_node.key.token_ids[0], 60)
        self.assertEqual(last_node.key.token_ids[1], 70)

        print(tree.available_and_evictable_str())
        print(available_and_evictable_str(tree))
        tree.sanity_check()

    def test_swa_radix_cache_eagle(self):
        # args
        req_size = 10
        max_context_len = 128
        kv_size = 128
        kv_size_swa = 64
        page_size = 1
        sliding_window_size = 4
        head_num = 8
        head_dim = 128
        num_layers = 48
        global_interval = 4
        dtype = torch.bfloat16
        device = get_device()
        full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)]
        full_attention_layer_ids_set = set(full_attention_layer_ids)
        swa_attention_layer_ids = [
            i for i in range(num_layers) if i not in full_attention_layer_ids_set
        ]
        # setup req to token pool
        req_to_token_pool = ReqToTokenPool(
            size=req_size,
            max_context_len=max_context_len,
            device=device,
            enable_memory_saver=False,
        )
        # setup kv pool
        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_attention_layer_ids,
            full_attention_layer_ids=full_attention_layer_ids,
            device=device,
        )
        # setup token to kv pool allocator
        allocator = SWATokenToKVPoolAllocator(
            size=kv_size,
            size_swa=kv_size_swa,
            page_size=page_size,
            dtype=dtype,
            device=device,
            kvcache=kv_pool,
            need_sort=False,
        )
        # setup radix cache
        tree = SWARadixCache(
            params=CacheInitParams(
                req_to_token_pool=req_to_token_pool,
                token_to_kv_pool_allocator=allocator,
                page_size=page_size,
                disable=False,
                is_eagle=True,
                sliding_window_size=sliding_window_size,
            ),
        )

        # test
        print(
            f"[Start] allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
        )
        req1_token_ids, req1_kv_indices = [1, 2, 3], allocator.alloc(3)
        self.assertEqual(len(req1_token_ids), len(req1_kv_indices))
        print(
            f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}"
        )
        key = RadixKey(array("q", req1_token_ids))
        result = tree.insert(InsertParams(key=key, value=req1_kv_indices[: len(key)]))
        prefix_len = result.prefix_len
        self.assertEqual(prefix_len, 0)
        print(
            f"req1: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
        )
        req2_token_ids, req2_kv_indices = [1, 2, 3, 4, 5, 6, 7], allocator.alloc(7)
        self.assertEqual(len(req2_token_ids), len(req2_kv_indices))
        print(
            f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}"
        )
        key = RadixKey(array("q", req2_token_ids))
        result = tree.insert(InsertParams(key=key, value=req2_kv_indices[: len(key)]))
        prefix_len = result.prefix_len
        self.assertEqual(prefix_len, 2)
        print(
            f"req2: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
        )
        req3_token_ids, req3_kv_indices = [10, 11, 12], allocator.alloc(3)
        self.assertEqual(len(req3_token_ids), len(req3_kv_indices))
        print(
            f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}"
        )
        key = RadixKey(array("q", req3_token_ids))
        result = tree.insert(InsertParams(key=key, value=req3_kv_indices[: len(key)]))
        prefix_len = result.prefix_len
        self.assertEqual(prefix_len, 0)
        print(
            f"req3: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
        )
        req4_token_ids, req4_kv_indices = [1, 2, 3, 4, 5, 60, 70], allocator.alloc(7)
        self.assertEqual(len(req4_token_ids), len(req4_kv_indices))
        print(
            f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}"
        )
        key = RadixKey(array("q", req4_token_ids))
        result = tree.insert(InsertParams(key=key, value=req4_kv_indices[: len(key)]))
        prefix_len = result.prefix_len
        self.assertEqual(prefix_len, 4)
        print(
            f"req4: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
        )

        tree.pretty_print()
        full_num_tokens, swa_num_tokens = 1, 0
        print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token")
        evict_result = tree.evict(
            EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
        )
        assert isinstance(evict_result, EvictResult)
        assert (
            evict_result.num_tokens_evicted >= full_num_tokens
        )  # May evict more due to node granularity
        print(
            f"evicted {evict_result.num_tokens_evicted} full tokens, {evict_result.swa_num_tokens_evicted} swa tokens"
        )
        tree.pretty_print()

        full_num_tokens, swa_num_tokens = 0, 1
        print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token")
        evict_result = tree.evict(
            EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
        )
        assert isinstance(evict_result, EvictResult)
        assert (
            evict_result.swa_num_tokens_evicted >= swa_num_tokens
        ), f"evicted {evict_result.swa_num_tokens_evicted} swa tokens, expected {swa_num_tokens}"
        tree.pretty_print()

        full_num_tokens, swa_num_tokens = 1, 2
        print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token")
        evict_result = tree.evict(
            EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
        )
        assert isinstance(evict_result, EvictResult)
        assert (
            evict_result.num_tokens_evicted >= full_num_tokens
        ), f"evicted {evict_result.num_tokens_evicted} full tokens, expected {full_num_tokens}"
        assert (
            evict_result.swa_num_tokens_evicted >= swa_num_tokens
        ), f"evicted {evict_result.swa_num_tokens_evicted} swa tokens, expected {swa_num_tokens}"
        tree.pretty_print()

        req5_token_ids = [1, 2, 3, 4, 5]
        result = tree.match_prefix(
            MatchPrefixParams(key=RadixKey(array("q", req5_token_ids)))
        )
        kv_indices, last_node = result.device_indices, result.last_device_node
        print(
            f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
        )
        self.assertEqual(len(kv_indices), 0)  # no swa prefix matched

        req6_token_ids = [1, 2, 3, 4, 5, 60, 70]
        result = tree.match_prefix(
            MatchPrefixParams(key=RadixKey(array("q", req6_token_ids)))
        )
        kv_indices, last_node = result.device_indices, result.last_device_node
        print(
            f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
        )
        self.assertEqual(len(kv_indices), 6)
        self.assertEqual(len(last_node.key), 2)
        # Bigram view: token_ids holds raw tokens; iteration yields bigram tuples.
        self.assertTrue(last_node.key.is_bigram)
        self.assertEqual(list(last_node.key), [(5, 60), (60, 70)])

    def test_swa_cache_finished_req_eagle_uses_cache_protected_len_and_bigram_key(self):
        tree, allocator, req_to_token_pool = _build_swa_tree(is_eagle=True)

        # Case 1: is_insert=True should pass bigram key and use cache_protected_len.
        req = _DummyReq()
        req.req_pool_idx = 0
        req.origin_input_ids = array("q", [1, 2, 3, 4, 5, 6])
        req.output_ids = array("q")
        req._kv_committed_len = len(req.origin_input_ids)
        kv_indices = allocator.alloc(req._kv_committed_len)
        req_to_token_pool.write(
            (req.req_pool_idx, slice(0, req._kv_committed_len)), kv_indices
        )
        req.extra_key = None
        req.last_node = tree.root_node
        req.swa_uuid_for_lock = None
        req.swa_evicted_seqlen = 0
        req.cache_protected_len = 1
        # Intentionally mismatch to ensure code does not use len(prefix_indices).
        req.prefix_indices = torch.tensor([7, 8, 9, 10, 11], device=tree.device)

        captured = {}
        original_insert = tree.insert

        def wrapped_insert(params):
            captured["prev_prefix_len"] = params.prev_prefix_len
            captured["is_bigram"] = params.key.is_bigram
            captured["key_len"] = len(params.key)
            return original_insert(params)

        tree.insert = wrapped_insert
        tree.cache_finished_req(req, is_insert=True)

        self.assertEqual(captured["prev_prefix_len"], req.cache_protected_len)
        self.assertTrue(captured["is_bigram"])
        self.assertEqual(captured["key_len"], len(req.origin_input_ids) - 1)

        # Case 2: is_insert=False should free [cache_protected_len:page_aligned_len]
        # even when len(prefix_indices) is intentionally larger.
        req2 = _DummyReq()
        req2.req_pool_idx = 1
        req2.origin_input_ids = array("q", [11, 12, 13, 14, 15, 16])
        req2.output_ids = array("q")
        req2._kv_committed_len = len(req2.origin_input_ids)
        kv_indices2 = allocator.alloc(req2._kv_committed_len)
        req_to_token_pool.write(
            (req2.req_pool_idx, slice(0, req2._kv_committed_len)), kv_indices2
        )
        req2.extra_key = None
        req2.last_node = tree.root_node
        req2.swa_uuid_for_lock = None
        req2.swa_evicted_seqlen = 0
        req2.cache_protected_len = 1
        req2.prefix_indices = torch.tensor([21, 22, 23, 24, 25], device=tree.device)

        freed_lens = []
        original_free = allocator.free

        def wrapped_free(indices):
            freed_lens.append(int(indices.numel()))
            return original_free(indices)

        allocator.free = wrapped_free
        tree.cache_finished_req(req2, is_insert=False)

        # EAGLE + page_size=1 => page_aligned_len = committed_len - 1 = 5
        # Expected frees:
        #   overlap range [1:5] -> 4
        #   tail range [5:]     -> 1
        self.assertEqual(freed_lens, [4, 1])


# Optimization: SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT.
# Splits a freshly-inserted leaf at the (page-aligned) sliding-window
# boundary so a future inc_lock_ref protects only ~sliding_window_size SWA
# tokens instead of the whole chunked-prefill chain.
class TestSWASplitLeafOnInsert(CustomTestCase):
    def _insert_and_lock(self, *, window, page_size, leaf_len, flag_on):
        tree, allocator, _ = _build_swa_tree(
            is_eagle=False,
            kv_size=128,
            kv_size_swa=64,
            sliding_window_size=window,
            page_size=page_size,
        )
        token_ids = list(range(leaf_len))
        with envs.SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT.override(flag_on):
            leaf = _insert_chain(tree, allocator, token_ids)
        result = tree.inc_lock_ref(leaf)
        return tree, leaf, result

    def test_flag_off_protects_full_leaf(self):
        tree, leaf, _ = self._insert_and_lock(
            window=4, page_size=1, leaf_len=12, flag_on=False
        )
        self.assertEqual(len(leaf.value), 12)
        self.assertEqual(tree.swa_protected_size_, 12)

    def test_flag_on_caps_protection_at_window(self):
        # (window, page_size, leaf_len, expected_tail_size); leaf_len picked
        # > tail_size and page-aligned for page_size > 1.
        cases = [
            (4, 1, 12, 4),
            (4, 1, 5, 4),
            (1, 1, 5, 1),
            (4, 2, 12, 4),
            (8, 2, 12, 8),
            (4, 4, 12, 4),
            # window NOT page-aligned -> tail rounds up to page boundary.
            (3, 2, 12, 4),
            (5, 4, 12, 8),
            (3, 4, 12, 4),
        ]
        for window, page_size, leaf_len, expected_tail in cases:
            with self.subTest(window=window, page_size=page_size, leaf_len=leaf_len):
                self.assertEqual(_expected_tail_size(window, page_size), expected_tail)
                tree, leaf, _ = self._insert_and_lock(
                    window=window,
                    page_size=page_size,
                    leaf_len=leaf_len,
                    flag_on=True,
                )
                self.assertEqual(len(leaf.value), expected_tail)
                self.assertEqual(tree.swa_protected_size_, expected_tail)

    def test_flag_on_no_split_when_leaf_within_window(self):
        # leaf_len <= tail_size: split must no-op.
        cases = [
            (4, 1, 4),
            (4, 1, 3),
            (4, 2, 4),
            (3, 2, 4),
            (8, 2, 4),
            (4, 4, 4),
        ]
        for window, page_size, leaf_len in cases:
            with self.subTest(window=window, page_size=page_size, leaf_len=leaf_len):
                tree, leaf, _ = self._insert_and_lock(
                    window=window,
                    page_size=page_size,
                    leaf_len=leaf_len,
                    flag_on=True,
                )
                self.assertEqual(len(leaf.value), leaf_len)
                self.assertEqual(tree.swa_protected_size_, leaf_len)

    def test_match_prefix_returns_full_chain_after_split(self):
        tree, allocator, _ = _build_swa_tree(
            is_eagle=False,
            kv_size=128,
            kv_size_swa=64,
            sliding_window_size=4,
            page_size=1,
        )
        token_ids = list(range(12))
        with envs.SGLANG_OPT_SWA_SPLIT_LEAF_ON_INSERT.override(True):
            inserted_leaf = _insert_chain(tree, allocator, token_ids)
        self.assertEqual(len(inserted_leaf.value), 4)
        match = tree.match_prefix(
            MatchPrefixParams(key=RadixKey(array("q", token_ids)))
        )
        self.assertEqual(match.device_indices.shape[0], 12)
        self.assertIs(match.last_device_node, inserted_leaf)

    def test_dec_lock_ref_after_split_balances_to_zero(self):
        tree, leaf, result = self._insert_and_lock(
            window=4, page_size=1, leaf_len=12, flag_on=True
        )
        self.assertEqual(tree.swa_protected_size_, 4)
        self.assertEqual(tree.full_protected_size_, 12)

        tree.dec_lock_ref(
            leaf,
            params=DecLockRefParams(swa_uuid_for_lock=result.swa_uuid_for_lock),
        )

        self.assertEqual(tree.swa_protected_size_, 0)
        self.assertEqual(tree.full_protected_size_, 0)
        tree.sanity_check()


if __name__ == "__main__":
    unittest.main()