[Fix] Merge adjacent KV-row frees so a mid-page split under DCP cannot double-free (#38941)

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -415,12 +415,14 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
         """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
+        from sglang.srt.mem_cache.common import coalesce_ranges, free_kv_row_segments
 
         row = self.req_to_token_pool.req_to_token[kv.req_pool_idx]
+        # Adjacent pieces whose seam falls inside one (DCP-widened) page would
+        # free that page twice; the allocator rejects that, so merge them first.
         free_kv_row_segments(
             self.token_to_kv_pool_allocator,
-            [(row[start:end], start) for start, end in ranges],
+            [(row[start:end], start) for start, end in coalesce_ranges(ranges)],
             swa_evicted_seqlen=kv.swa_evicted_seqlen,
         )
 

python/sglang/srt/mem_cache/common.py

@@ -108,6 +108,17 @@ def free_swa_out_of_window_slots(
         req.kv.swa_evicted_seqlen = new_swa_evicted_seqlen
 
 
+def coalesce_ranges(ranges: list[tuple[int, int]]) -> list[tuple[int, int]]:
+    """Merge adjacent half-open ranges so a split that falls mid-page frees that page once."""
+    merged: list[tuple[int, int]] = []
+    for start, end in ranges:
+        if merged and start == merged[-1][1]:
+            merged[-1] = (merged[-1][0], end)
+        else:
+            merged.append((start, end))
+    return merged
+
+
 def free_kv_row_segments(
     allocator: BaseTokenToKVPoolAllocator,
     segments: list[tuple[torch.Tensor, int]],

test/registered/unit/mem_cache/test_free_kv_row_coalesce.py

@@ -0,0 +1,92 @@
+"""free_kv_row must not hand the allocator two frees that meet inside one page.
+
+Under decode context parallelism the allocator page is page_size * dcp_size,
+while cache-length caps (e.g. the Mamba track boundary) stay on the unwidened
+grid, so cache_finished_req can free a request's tail as two adjacent ranges
+split mid-page. free_segments rejects that as a double free.
+
+    python -m pytest test/registered/unit/mem_cache/test_free_kv_row_coalesce.py -v
+"""
+
+import unittest
+from types import SimpleNamespace
+
+import torch
+
+from sglang.srt.mem_cache.allocator.paged import PagedTokenToKVPoolAllocator
+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=10, suite="base-a-test-cpu")
+
+# A 64-token page widened by dcp_size 8 is 512; scaled down 64x: page 8, and
+# a Mamba track boundary at 576 tokens lands at position 9.
+PAGE_SIZE = 8
+NUM_PAGES = 16
+
+
+class _RowCache(BasePrefixCache):
+    """Just enough of a prefix cache to exercise free_kv_row."""
+
+    def __init__(self, allocator, req_to_token):
+        self.token_to_kv_pool_allocator = allocator
+        self.req_to_token_pool = SimpleNamespace(req_to_token=req_to_token)
+
+    def reset(self):
+        raise NotImplementedError
+
+    def match_prefix(self, params):
+        raise NotImplementedError
+
+    def cache_finished_req(self, req, is_insert=True, **kwargs):
+        raise NotImplementedError
+
+    def cache_unfinished_req(self, req, **kwargs):
+        raise NotImplementedError
+
+    def evict(self, params):
+        raise NotImplementedError
+
+    def inc_lock_ref(self, node):
+        raise NotImplementedError
+
+    def dec_lock_ref(self, node):
+        raise NotImplementedError
+
+
+def _make_cache_with_row(num_tokens):
+    alloc = PagedTokenToKVPoolAllocator(
+        size=NUM_PAGES * PAGE_SIZE,
+        page_size=PAGE_SIZE,
+        dtype=torch.float16,
+        device="cpu",
+        kvcache=None,
+        need_sort=False,
+    )
+    row = alloc.alloc(num_tokens)
+    req_to_token = torch.zeros((1, num_tokens), dtype=row.dtype)
+    req_to_token[0] = row
+    kv = SimpleNamespace(req_pool_idx=0, swa_evicted_seqlen=0)
+    return _RowCache(alloc, req_to_token), alloc, row, kv
+
+
+class TestFreeKvRowCoalesce(unittest.TestCase):
+    def test_adjacent_ranges_split_mid_page_free_every_page_once(self):
+        cache, alloc, row, kv = _make_cache_with_row(3 * PAGE_SIZE)
+        before = len(alloc.free_pages)
+        # Kept prefix [0, 8); the unaligned tail [8, 9) and the deferred
+        # truncation tail [9, 24) meet at 9, inside page 1.
+        cache.free_kv_row(kv, [(8, 9), (9, 3 * PAGE_SIZE)])
+        freed = alloc.free_pages[: len(alloc.free_pages) - before]
+        reference = torch.unique(row[8:] // PAGE_SIZE)
+        self.assertTrue(torch.equal(torch.sort(freed)[0], reference))
+
+    def test_non_adjacent_ranges_sharing_a_page_still_rejected(self):
+        cache, _, _, kv = _make_cache_with_row(3 * PAGE_SIZE)
+        # A gap between the ranges means the shared page really is freed twice.
+        with self.assertRaises(AssertionError):
+            cache.free_kv_row(kv, [(8, 9), (12, 3 * PAGE_SIZE)])
+
+
+if __name__ == "__main__":
+    unittest.main()