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

python/sglang/srt/disaggregation/decode_kvcache_offload_manager.py

@@ -250,27 +250,20 @@ class DecodeKVCacheOffloadManager:
         # Prefill-aligned slots are freed only here, at request finish; freeing
         # them mid-decode races with concurrent admission over live slots.
         prefill_len = self._prefill_offloaded_len(req)
+        ranges = []
         if prefill_len > 0:
-            prefill_indices = self.req_to_token_pool.req_to_token[
-                req.kv.req_pool_idx, :prefill_len
-            ]
-            self.token_to_kv_pool_allocator.free(prefill_indices)
-        start = prefill_len
-        end = kv_committed_len
-        # Free the incremental part of the request (DSA-aware)
-        kv_indices = self.req_to_token_pool.req_to_token[req.kv.req_pool_idx, start:end]
-        self.token_to_kv_pool_allocator.free(kv_indices)
+            ranges.append((0, prefill_len))
+        # The incremental part of the request (DSA-aware)
+        ranges.append((prefill_len, kv_committed_len))
 
-        # Free over-allocated KV cache slots (e.g. from speculative decoding v2).
-        # Without spec v2, start_p == end_p so this is a no-op.
+        # Over-allocated KV cache slots (e.g. from speculative decoding v2).
+        # Without spec v2, start_p == end_p so this contributes nothing.
         start_p, end_p = kv_committed_len, req.kv.kv_allocated_len
         if self.page_size > 1:
             start_p = ceil_align(start_p, self.page_size)
         if start_p < end_p:
-            overalloc_indices = self.req_to_token_pool.req_to_token[
-                req.kv.req_pool_idx, start_p:end_p
-            ]
-            self.token_to_kv_pool_allocator.free(overalloc_indices)
+            ranges.append((start_p, end_p))
+        self.tree_cache.free_kv_row(req.kv, ranges)
 
         self.req_to_token_pool.free(req)
         req.kv.mark_kv_released()

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -369,6 +369,19 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
     def cache_unfinished_req(self, req: Req, **kwargs):
         pass
 
+    def free_kv_row(self, kv: Any, ranges: list[tuple[int, int]]) -> None:
+        """Give back ascending, disjoint, half-open row-position ranges
+        of the ``kv`` record's row; one call keeps a shared page freed once.
+        """
+        from sglang.srt.mem_cache.common import free_kv_row_segments
+
+        row = self.req_to_token_pool.req_to_token[kv.req_pool_idx]
+        free_kv_row_segments(
+            self.token_to_kv_pool_allocator,
+            [(row[start:end], start) for start, end in ranges],
+            swa_evicted_seqlen=kv.swa_evicted_seqlen,
+        )
+
     @abstractmethod
     def evict(self, params: EvictParams) -> EvictResult:
         pass

python/sglang/srt/mem_cache/chunk_cache.py

@@ -81,10 +81,7 @@ class ChunkCache(BasePrefixCache):
     ):
         # For decode server: if req.output_ids is empty, we want to free all req.origin_input_ids
         # The protected prefix is not this req's to free.
-        kv_indices = self.req_to_token_pool.req_to_token[
-            req.kv.req_pool_idx, req.kv.cache_protected_len : kv_len_to_handle
-        ]
-        self.token_to_kv_pool_allocator.free(kv_indices)
+        self.free_kv_row(req.kv, [(req.kv.cache_protected_len, kv_len_to_handle)])
 
     def cache_unfinished_req(self, req: Req, chunked=False):
         kv_indices = self.req_to_token_pool.req_to_token[

python/sglang/srt/mem_cache/common.py

@@ -119,6 +119,45 @@ def free_swa_out_of_window_slots(
         req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen
 
 
+def free_kv_row_segments(
+    allocator: BaseTokenToKVPoolAllocator,
+    segments: list[tuple[torch.Tensor, int]],
+    *,
+    swa_evicted_seqlen: int,
+) -> None:
+    """Free ascending disjoint ``(kv_indices, start_pos)`` segments of one
+    request's kv row, split at the SWA eviction floor."""
+    swa_dead: list[torch.Tensor] = []
+    swa_alive: list[tuple[torch.Tensor, int]] = []
+    for kv_indices, start_pos in segments:
+        num_indices = kv_indices.numel()
+        if num_indices == 0:
+            continue
+        # Below the floor the SWA peers are already gone -- window eviction, or
+        # the deliberately unmapped prefix of a PD decode SWA-tail prealloc.
+        num_dead = min(max(swa_evicted_seqlen - start_pos, 0), num_indices)
+        if num_dead > 0:
+            swa_dead.append(kv_indices[:num_dead])
+        if num_dead < num_indices:
+            swa_alive.append((kv_indices[num_dead:], start_pos + num_dead))
+
+    if swa_dead and swa_alive:
+        # A mid-page floor would send a page shared by the dead and alive
+        # sides back twice.
+        assert swa_evicted_seqlen % allocator.page_size == 0, (
+            f"SWA eviction floor {swa_evicted_seqlen} splits a page "
+            f"(page_size {allocator.page_size})"
+        )
+    if len(swa_dead) == 1:
+        allocator.free_full(swa_dead[0])
+    elif swa_dead:
+        # Two dead pieces can share a boundary page, and only free_full's own
+        # page dedup covers that -- free_segments trims the alive side alone.
+        allocator.free_full(torch.cat(swa_dead))
+    if swa_alive:
+        allocator.free_segments(swa_alive)
+
+
 def maybe_cache_unfinished_req(req: Req, tree_cache: BasePrefixCache, **kwargs):
     if getattr(req, "skip_radix_cache_insert", False):
         return
@@ -275,12 +314,9 @@ def _release_overallocated_kv_indices(
         start_p = ceil_align(start_p, page_size)
 
     if start_p < end_p:
-        indices_to_free = tree_cache.req_to_token_pool.req_to_token[
-            req.kv.req_pool_idx
-        ][start_p:end_p]
         # start_p is aligned to the allocator's physical page size above, so it
         # never shares a page with cache_finished_req's tail free in this group.
-        allocator.free_segment(indices_to_free, start_pos=start_p)
+        tree_cache.free_kv_row(req.kv, [(start_p, end_p)])
 
 
 def available_and_evictable_str(tree_cache: BasePrefixCache) -> str:

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -42,6 +42,7 @@ from sglang.srt.mem_cache.base_prefix_cache import (
     MatchResult,
 )
 from sglang.srt.mem_cache.cache_init_params import CacheInitParams
+from sglang.srt.mem_cache.common import free_kv_row_segments
 from sglang.srt.mem_cache.events import KVCacheEventRecorder
 from sglang.srt.mem_cache.radix_cache import RadixKey
 from sglang.srt.mem_cache.utils import split_node_hash_value
@@ -464,10 +465,7 @@ class SWARadixCache(BasePrefixCache):
     ) -> None:
         """Cache request when it finishes."""
         if self.disable:
-            kv_indices = self.req_to_token_pool.req_to_token[
-                req.kv.req_pool_idx, :kv_len_to_handle
-            ]
-            self.token_to_kv_pool_allocator.free(kv_indices)
+            self.free_kv_row(req.kv, [(0, kv_len_to_handle)])
             return
 
         token_ids = (req.origin_input_ids + req.output_ids)[:kv_len_to_handle]
@@ -497,12 +495,10 @@ class SWARadixCache(BasePrefixCache):
                 )
             )
         else:
-            self.token_to_kv_pool_allocator.free(
-                kv_indices[old_prefix_len:page_aligned_len]
-            )
+            self.free_kv_row(req.kv, [(old_prefix_len, page_aligned_len)])
 
         # free the unaligned tail
-        self.token_to_kv_pool_allocator.free(kv_indices[page_aligned_len:])
+        self.free_kv_row(req.kv, [(page_aligned_len, kv_len_to_handle)])
 
         # Remove req slot release the cache lock
         self.dec_lock_ref(
@@ -1207,7 +1203,7 @@ class SWARadixCache(BasePrefixCache):
                             )
                         else:
                             # Free full tokens in the original tree node.
-                            self.token_to_kv_pool_allocator.free(
+                            self.token_to_kv_pool_allocator.free_full(
                                 node.value[:prefix_len]
                             )
                             # Overwrite the new value in request to the tree node.
@@ -1225,18 +1221,18 @@ class SWARadixCache(BasePrefixCache):
                             self._recover_tombstone_keeping_locked_full(
                                 node, value[start_update_idx:prefix_len]
                             )
-                            self.token_to_kv_pool_allocator.free(
+                            self.token_to_kv_pool_allocator.free_full(
                                 value[:start_update_idx]
                             )
                         else:
-                            self.token_to_kv_pool_allocator.free(
+                            self.token_to_kv_pool_allocator.free_full(
                                 node.value[start_update_idx:prefix_len]
                             )
                             self._split_node(node.key, node, start_update_idx)
                             # Here node is the new node after split, so we can overwrite the value to the new node.
                             # The old node is still swa tombstone and the full token is not freed.
                             node.value = value[start_update_idx:prefix_len].clone()
-                            self.token_to_kv_pool_allocator.free(
+                            self.token_to_kv_pool_allocator.free_full(
                                 value[:start_update_idx]
                             )
                             node.swa_tombstone = False
@@ -1244,10 +1240,16 @@ class SWARadixCache(BasePrefixCache):
                             self.swa_evictable_size_ += len(node.value)
                     else:
                         # Branch 3: all swa tokens of value[:prefix_len] are evicted, so we don't need to update the node.
-                        self.token_to_kv_pool_allocator.free(value[:prefix_len])
+                        self.token_to_kv_pool_allocator.free_full(value[:prefix_len])
                 else:
                     # The node is not tombstone, so we don't need to update the node.
-                    self.token_to_kv_pool_allocator.free(value[:prefix_len])
+                    # The incoming slice can still straddle this request's own
+                    # eviction floor, so split it there.
+                    free_kv_row_segments(
+                        self.token_to_kv_pool_allocator,
+                        [(value[:prefix_len], total_prefix_length)],
+                        swa_evicted_seqlen=swa_evicted_seqlen,
+                    )
 
             total_prefix_length += prefix_len
             key = key[prefix_len:]
@@ -1274,7 +1276,7 @@ class SWARadixCache(BasePrefixCache):
             #    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)
+                self.token_to_kv_pool_allocator.free_full(value)
                 return total_prefix_length
 
             if (

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -842,10 +842,7 @@ class UnifiedRadixCache(BasePrefixCache):
             return
 
         if self.disable:
-            kv_indices = self.req_to_token_pool.req_to_token[
-                req.kv.req_pool_idx, :kv_len_to_handle
-            ]
-            self.token_to_kv_pool_allocator.free_segment(kv_indices, start_pos=0)
+            self.free_kv_row(req.kv, [(0, kv_len_to_handle)])
             for comp in self._components_tuple:
                 comp.cleanup_after_caching_req(req, is_finished=True)
             return
@@ -899,15 +896,12 @@ class UnifiedRadixCache(BasePrefixCache):
             result = self.insert(insert_params)
 
             # Free unaligned tail (+ deferred truncation tail)
-            segments = [(kv_indices[page_aligned_len:], page_aligned_len)]
+            ranges = [(page_aligned_len, len(kv_indices))]
             if tail_free_start is not None:
-                segments.append((kv_indices_full[tail_free_start:], tail_free_start))
-            self.token_to_kv_pool_allocator.free_segments(segments)
+                ranges.append((tail_free_start, len(kv_indices_full)))
+            self.free_kv_row(req.kv, ranges)
         else:
-            self.token_to_kv_pool_allocator.free_segment(
-                kv_indices[req.kv.cache_protected_len :],
-                start_pos=req.kv.cache_protected_len,
-            )
+            self.free_kv_row(req.kv, [(req.kv.cache_protected_len, kv_len_to_handle)])
 
         self._dec_req_lock(req, skip_swa=req.swa_prefix_lock_released)
 

python/sglang/srt/session/streaming_session.py

@@ -232,8 +232,10 @@ class StreamingSession(BasePrefixCache):
             f"{slot.kv.cache_protected_len=}"
         )
 
-        # Floor-align prefix_len to page boundary (NPU workaround).
-        if is_npu() and self.page_size > 1:
+        # NPU requires page-aligned KV reuse; a rewind below the SWA eviction
+        # cursor must also land on a page boundary -- free_kv_row_segments
+        # splits dead/alive at the cursor, and a mid-page cut frees a page twice.
+        if self.page_size > 1 and (is_npu() or req.kv.swa_evicted_seqlen > prefix_len):
             prefix_len = (prefix_len // self.page_size) * self.page_size
             req.kv.kv_committed_len = min(req.kv.kv_committed_len, prefix_len)
 
@@ -401,13 +403,7 @@ class StreamingSession(BasePrefixCache):
             )
 
         if slot.kv.holds_kv:
-            start = protected_len
-            end = slot.kv.kv_allocated_len
-            if start < end:
-                kv_indices = self.req_to_token_pool.req_to_token[
-                    slot.kv.req_pool_idx, start:end
-                ]
-                self.token_to_kv_pool_allocator.free(kv_indices)
+            self.free_kv_row(slot.kv, [(protected_len, slot.kv.kv_allocated_len)])
             self.req_to_token_pool.free(slot)
 
         self._free_slot_mamba(slot)
@@ -504,7 +500,7 @@ class StreamingSession(BasePrefixCache):
         decoding pushes allocated above committed, or when retract retry's
         logit-reserve pulls prefix_len below committed.
         """
-        self._free_kv_aligned(kv.req_pool_idx, prefix_len, kv.kv_allocated_len)
+        self._free_kv_aligned(kv, prefix_len, kv.kv_allocated_len)
         kv.kv_allocated_len = prefix_len
         kv.kv_committed_len = min(kv.kv_committed_len, prefix_len)
         kv.swa_evicted_seqlen = min(kv.swa_evicted_seqlen, prefix_len)
@@ -516,14 +512,18 @@ class StreamingSession(BasePrefixCache):
         be released to avoid token/KV mismatch.
         """
         target = len(req.origin_input_ids) + finished_len
-        self._free_kv_aligned(req.kv.req_pool_idx, target, req.kv.kv_allocated_len)
+        if self.page_size > 1 and req.kv.swa_evicted_seqlen > target:
+            # Same hazard as the match-path rewind: the cursor must stay
+            # page-aligned; the partial page is re-prefilled next turn.
+            target = (target // self.page_size) * self.page_size
+        self._free_kv_aligned(req.kv, target, req.kv.kv_allocated_len)
         req.kv.kv_allocated_len = min(req.kv.kv_allocated_len, target)
         req.kv.kv_committed_len = min(req.kv.kv_committed_len, target)
         req.kv.swa_evicted_seqlen = min(req.kv.swa_evicted_seqlen, target)
         req.output_ids = req.output_ids[:finished_len]
 
-    def _free_kv_aligned(self, pool_idx: int, target: int, end: int) -> None:
-        """Free req_to_token[pool_idx, ceil_align(target):end). Page-aligned
+    def _free_kv_aligned(self, kv: ReqKvInfo, target: int, end: int) -> None:
+        """Free the record's kv row over [ceil_align(target), end). Page-aligned
         because PagedTokenToKVPoolAllocator.free returns whole pages
         (free_index // page_size), so partial-page free would corrupt pages
         still holding committed tokens. The range [target, ceil_align(target))
@@ -534,9 +534,7 @@ class StreamingSession(BasePrefixCache):
         start = target
         if self.page_size > 1:
             start = ceil_align(start, self.page_size)
-        if start < end:
-            tail = self.req_to_token_pool.req_to_token[pool_idx, start:end]
-            self.token_to_kv_pool_allocator.free(tail)
+        self.free_kv_row(kv, [(start, end)])
 
     # -- Pass-through methods --
 

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."""