diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 5034ad6a5..9695d1048 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/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, ) diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 9757fe57a..2e6b47f45 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/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]], diff --git a/test/registered/unit/mem_cache/test_free_kv_row_coalesce.py b/test/registered/unit/mem_cache/test_free_kv_row_coalesce.py new file mode 100644 index 000000000..74e9e2e0f --- /dev/null +++ b/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()