[mem_cache] Add free_kv_row to release a request's kv row by row range (#36721)

test/registered/unit/disaggregation/test_specv2_kvcache_offloading.py

@@ -20,6 +20,8 @@ from sglang.srt.disaggregation.decode_kvcache_offload_manager import (
 from sglang.srt.disaggregation.kv_events import OffloadedState
 from sglang.srt.managers.cache_controller import HiCacheAck
 from sglang.srt.managers.schedule_batch import ReqKvInfo
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=8, suite="base-a-test-cpu")
@@ -47,6 +49,31 @@ def _make_mock_req(
     return req
 
 
+class _RecordingAllocator(BaseTokenToKVPoolAllocator):
+    """Single-pool double. Subclassing the base routes free_full / free_segment /
+    free_segments into free(), so a new free API cannot slip past the recorder."""
+
+    def __init__(self, page_size: int):
+        super().__init__(
+            size=1024,
+            page_size=page_size,
+            dtype=torch.bfloat16,
+            device="cpu",
+            kvcache=None,
+            need_sort=False,
+        )
+        self.freed = []
+
+    def clear(self):
+        self.freed = []
+
+    def alloc(self, need_size: int):
+        raise NotImplementedError
+
+    def free(self, free_index: torch.Tensor):
+        self.freed.append(free_index.clone())
+
+
 def _make_manager(pool_size: int, page_size: int = 1):
     """Create a DecodeKVCacheOffloadManager with mock pools for testing."""
     # Build a real req_to_token tensor so indexing works
@@ -55,15 +82,16 @@ def _make_manager(pool_size: int, page_size: int = 1):
     req_to_token_pool = MagicMock()
     req_to_token_pool.req_to_token = req_to_token
 
-    freed_indices = []
-
-    allocator = MagicMock()
-    allocator.free = MagicMock(
-        side_effect=lambda idx: freed_indices.append(idx.clone())
-    )
+    allocator = _RecordingAllocator(page_size)
+    freed_indices = allocator.freed
 
     tree_cache = MagicMock()
     tree_cache.protected_size_ = 0
+    tree_cache.req_to_token_pool = req_to_token_pool
+    tree_cache.token_to_kv_pool_allocator = allocator
+    tree_cache.free_kv_row = lambda owner, ranges: BasePrefixCache.free_kv_row(
+        tree_cache, owner, ranges
+    )
 
     # Bypass __init__ entirely and set attributes directly
     manager = object.__new__(DecodeKVCacheOffloadManager)

test/registered/unit/layers/test_minicpm_sparse_cache.py

@@ -14,6 +14,7 @@ from sglang.srt.managers.scheduler_components.invariant_checker import (
 from sglang.srt.managers.scheduler_components.pool_stats_observer import (
     SchedulerPoolStatsObserver,
 )
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
 from sglang.srt.session.streaming_session import SessionSlot, StreamingSession
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -21,17 +22,23 @@ from sglang.test.ci.ci_register import register_cpu_ci
 register_cpu_ci(est_time=5, suite="base-a-test-cpu")
 
 
-class RecordingAllocator:
+class RecordingAllocator(BaseTokenToKVPoolAllocator):
+    """Single-pool double. Subclassing the base routes free_full / free_segment /
+    free_segments into free(), so a new free API cannot slip past the recorder."""
+
     def __init__(self, capacity: int):
+        super().__init__(
+            size=capacity,
+            page_size=1,
+            dtype=torch.bfloat16,
+            device="cpu",
+            kvcache=None,
+            need_sort=False,
+        )
         self.capacity = capacity
-        self.page_size = 1
         self.next_slot = 1
         self.live: set[int] = set()
 
-    @property
-    def size(self):
-        return self.capacity
-
     def alloc(self, size: int):
         if size > self.available_size():
             return None

test/registered/unit/mem_cache/test_decode_radix_lock_ref.py

@@ -85,6 +85,7 @@ class MockReq:
             kv_committed_len=len(fill_ids),
             kv_allocated_len=len(fill_ids),
             cache_protected_len=cache_protected_len,
+            swa_evicted_seqlen=0,
         )
 
     def get_fill_ids(self):

test/registered/unit/mem_cache/test_streaming_session_unit.py

@@ -3,6 +3,7 @@ from types import SimpleNamespace
 import torch
 
 from sglang.srt.managers.schedule_batch import FINISH_ABORT, ReqKvInfo
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import MatchResult
 from sglang.srt.session.streaming_session import SessionSlot, StreamingSession
 from sglang.test.ci.ci_register import register_cpu_ci
@@ -10,10 +11,27 @@ from sglang.test.ci.ci_register import register_cpu_ci
 register_cpu_ci(est_time=12, suite="base-a-test-cpu")
 
 
-class _FakeAllocator:
-    def __init__(self):
+class _FakeAllocator(BaseTokenToKVPoolAllocator):
+    """Single-pool double. Subclassing the base routes free_full / free_segment /
+    free_segments into free(), so a new free API cannot slip past the recorder."""
+
+    def __init__(self, page_size: int = 1):
+        super().__init__(
+            size=1024,
+            page_size=page_size,
+            dtype=torch.bfloat16,
+            device="cpu",
+            kvcache=None,
+            need_sort=False,
+        )
         self.freed = []
 
+    def clear(self):
+        self.freed = []
+
+    def alloc(self, need_size: int):
+        raise NotImplementedError
+
     def free(self, free_index: torch.Tensor):
         self.freed.append(free_index.clone())
 
@@ -114,7 +132,7 @@ def test_preabort_detaches_session_and_preserves_slot():
     the session: session=None, abort_req() called. Slot stays intact."""
     req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
     req_to_token_pool = _FakeReqToTokenPool(req_to_token)
-    allocator = _FakeAllocator()
+    allocator = _FakeAllocator(page_size=16)
     inner = _FakeInnerCache(
         req_to_token_pool,
         allocator,
@@ -283,9 +301,41 @@ def test_trim_overshoot_postcondition():
     assert req.kv.kv_allocated_len == target
     assert req.kv.swa_evicted_seqlen == target
     assert len(req.output_ids) == 12
-    # Tail [38, 44) freed by _free_kv_aligned.
-    assert len(allocator.freed) == 1
-    assert allocator.freed[0].tolist() == list(range(38, 44))
+    # Tail [38, 44) freed by _free_kv_aligned, split at the pre-trim eviction
+    # floor 42: [38, 42) gave its SWA peers back already, so it goes back full-only.
+    assert [t.tolist() for t in allocator.freed] == [[38, 39, 40, 41], [42, 43]]
+
+
+def test_trim_overshoot_keeps_cursor_page_aligned_on_paged():
+    """A mid-page trim target must not become the SWA eviction cursor (the
+    dead/alive split there frees the shared page twice); rewind to the boundary."""
+    page_size = 16
+    req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128)
+    req_to_token_pool = _FakeReqToTokenPool(req_to_token)
+    allocator = _FakeAllocator(page_size=page_size)
+    tree_cache = StreamingSession(
+        _FakeInnerCache(req_to_token_pool, allocator, page_size)
+    )
+
+    # origin=26, finished=12 -> raw target 38 (mid-page); cursor 48 > target.
+    req = _FakeReq("session-a", req_pool_idx=0, committed=52, allocated=64)
+    req.origin_input_ids = list(range(26))
+    req.output_ids = list(range(14))
+    req.kv.swa_evicted_seqlen = 48
+
+    tree_cache._trim_overshoot(req, finished_len=12)
+
+    # Rewound to floor_align(38) = 32; every cursor lands page-aligned.
+    assert req.kv.kv_allocated_len == 32
+    assert req.kv.kv_committed_len == 32
+    assert req.kv.swa_evicted_seqlen == 32
+    assert len(req.output_ids) == 12
+    # Freed [32, 64): [32, 48) below the old cursor goes back full-only,
+    # [48, 64) both halves.
+    assert [t.tolist() for t in allocator.freed] == [
+        list(range(32, 48)),
+        list(range(48, 64)),
+    ]
 
 
 if __name__ == "__main__":

test/registered/unit/mem_cache/test_swa_unittest.py

@@ -6,8 +6,10 @@ import torch
 
 from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored
 from sglang.srt.environ import envs
+from sglang.srt.mem_cache.allocator.base import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import (
+    BasePrefixCache,
     DecLockRefParams,
     EvictParams,
     EvictResult,
@@ -15,7 +17,10 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     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.common import (
+    available_and_evictable_str,
+    free_kv_row_segments,
+)
 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
@@ -851,6 +856,31 @@ class TestSWASplitLeafOnInsert(CustomTestCase):
         tree.sanity_check()
 
 
+class _SinglePoolAllocator(BaseTokenToKVPoolAllocator):
+    """Minimal single-pool allocator: no SWA peer, so the whole range dies
+    together whatever the floor says."""
+
+    def __init__(self):
+        super().__init__(
+            size=16,
+            page_size=1,
+            dtype=torch.bfloat16,
+            device="cpu",
+            kvcache=None,
+            need_sort=False,
+        )
+        self.freed = []
+
+    def clear(self):
+        self.freed = []
+
+    def alloc(self, need_size: int):
+        raise NotImplementedError
+
+    def free(self, free_index: torch.Tensor):
+        self.freed.append(free_index)
+
+
 class TestFreeFullPartition(CustomTestCase):
     """`free_full` releases only the full side of a hybrid SWA allocator."""
 
@@ -894,6 +924,96 @@ class TestFreeFullPartition(CustomTestCase):
         self.assertEqual(self.allocator.full_available_size(), self.full_baseline)
 
 
+class _RowCache:
+    """Minimal PrefixCacheTrait host, so free_kv_row can be exercised without
+    standing up a whole tree."""
+
+    free_kv_row = BasePrefixCache.free_kv_row
+
+    def __init__(self, allocator, row):
+        self.req_to_token_pool = SimpleNamespace(req_to_token=row.unsqueeze(0))
+        self.token_to_kv_pool_allocator = allocator
+        self.page_size = allocator.page_size
+
+
+class TestFreeKvRow(CustomTestCase):
+    """A kv row is given back split at `swa_evicted_seqlen`: the full side
+    whole, the SWA side only from the floor up."""
+
+    def setUp(self):
+        _, self.allocator, _ = _build_swa_tree(is_eagle=False)
+        self.full_baseline = self.allocator.full_available_size()
+        self.swa_baseline = self.allocator.swa_available_size()
+
+    def _sizes(self):
+        return (
+            self.allocator.full_available_size(),
+            self.allocator.swa_available_size(),
+        )
+
+    def test_floor_decides_how_much_of_the_swa_side_stays_out(self):
+        # (start_pos, num_slots, floor, rows whose SWA peers are already gone)
+        cases = [
+            (0, 4, 4, 4),
+            (8, 4, 8, 0),
+            (8, 4, 10, 2),
+            (8, 4, 4, 0),
+        ]
+        for start_pos, num_slots, floor, num_dead in cases:
+            with self.subTest(start_pos=start_pos, floor=floor):
+                indices = _swa_alloc(self.allocator, num_slots)
+                free_kv_row_segments(
+                    self.allocator, [(indices, start_pos)], swa_evicted_seqlen=floor
+                )
+                self.assertEqual(
+                    self._sizes(),
+                    (self.full_baseline, self.swa_baseline - num_dead),
+                )
+                # Give the held-back SWA peers back, so the next case starts clean.
+                if num_dead:
+                    self.allocator.free_swa(indices[:num_dead])
+                self.assertEqual(self._sizes(), (self.full_baseline, self.swa_baseline))
+
+    def test_adjacent_below_floor_pieces_release_their_shared_page_once(self):
+        _, allocator, _ = _build_swa_tree(is_eagle=False, page_size=4)
+        indices = _swa_alloc(allocator, 8)
+        after_alloc = allocator.full_available_size()
+
+        # Rows [0, 6) and [6, 8) both sit below the floor and share page 1.
+        free_kv_row_segments(
+            allocator,
+            [(indices[:6], 0), (indices[6:], 6)],
+            swa_evicted_seqlen=8,
+        )
+
+        self.assertEqual(allocator.full_available_size(), after_alloc + 8)
+
+    def test_free_kv_row_reads_the_record_row_and_its_floor(self):
+        indices = _swa_alloc(self.allocator, 8)
+        cache = _RowCache(self.allocator, indices)
+        kv = SimpleNamespace(req_pool_idx=0, swa_evicted_seqlen=3)
+
+        cache.free_kv_row(kv, [(1, 5)])
+
+        # Rows [1, 5) go back on the full side; of those, [1, 3) lost their SWA
+        # peers already, so 6 of the 8 SWA slots are still out.
+        self.assertEqual(self._sizes(), (self.full_baseline - 4, self.swa_baseline - 6))
+
+    def test_single_pool_free_kv_row_still_frees_the_whole_range(self):
+        allocator = _SinglePoolAllocator()
+        cache = _RowCache(allocator, torch.arange(16, dtype=torch.int64))
+        kv = SimpleNamespace(req_pool_idx=0, swa_evicted_seqlen=4)
+
+        cache.free_kv_row(kv, [(2, 6)])
+
+        self.assertEqual([t.tolist() for t in allocator.freed], [[2, 3], [4, 5]])
+
+        # release_session and _free_kv_aligned dropped their own emptiness
+        # guards, so an empty range has to stay a no-op here.
+        cache.free_kv_row(kv, [(6, 6)])
+        self.assertEqual(len(allocator.freed), 2)
+
+
 class TestCacheUnfinishedReqEvictedPrefix(CustomTestCase):
     """An unfinished request whose SWA prefix is already gone must insert that
     prefix as a tombstone, not as live SWA KV."""