diff --git a/python/sglang/srt/hardware_backend/npu/allocator_npu.py b/python/sglang/srt/hardware_backend/npu/allocator_npu.py index 1a6ce9e6e..32ec7ffc8 100644 --- a/python/sglang/srt/hardware_backend/npu/allocator_npu.py +++ b/python/sglang/srt/hardware_backend/npu/allocator_npu.py @@ -149,7 +149,7 @@ class NPUPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator): else: self.free_pages = torch.cat((free_page_indices, self.free_pages)) else: - self.free_group.append(free_index) + self.free_group.append(self._copy_for_free_group(free_index)) if self.debug_mode: assert len(torch.unique(self.free_pages)) == len(self.free_pages) diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index 1a8b08560..532cee46a 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -69,6 +69,11 @@ class BaseTokenToKVPoolAllocator(abc.ABC): if self.free_group: self.free(torch.cat(self.free_group)) + @staticmethod + def _copy_for_free_group(free_index: torch.Tensor) -> torch.Tensor: + """Take ownership before a caller can mutate a deferred tensor view.""" + return free_index.clone() + def merge_and_sort_free(self): if len(self.release_pages) > 0: self.free_pages = torch.cat((self.free_pages, self.release_pages)) diff --git a/python/sglang/srt/mem_cache/allocator/hisparse.py b/python/sglang/srt/mem_cache/allocator/hisparse.py index 79245ac29..8ebfee249 100644 --- a/python/sglang/srt/mem_cache/allocator/hisparse.py +++ b/python/sglang/srt/mem_cache/allocator/hisparse.py @@ -262,7 +262,7 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.logical_attn_allocator.free(free_index) self.free_hisparse(free_index) else: - self.free_group.append(free_index) + self.free_group.append(self._copy_for_free_group(free_index)) assert ( self.logical_attn_allocator.available_size() <= self.logical_attn_allocator.size @@ -585,4 +585,4 @@ class DeepSeekV4HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): if self.is_not_in_free_group: self.logical_attn_allocator.free(free_index) else: - self.free_group.append(free_index) + self.free_group.append(self._copy_for_free_group(free_index)) diff --git a/python/sglang/srt/mem_cache/allocator/paged.py b/python/sglang/srt/mem_cache/allocator/paged.py index 9cfdfa069..b092911f2 100755 --- a/python/sglang/srt/mem_cache/allocator/paged.py +++ b/python/sglang/srt/mem_cache/allocator/paged.py @@ -265,7 +265,7 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): if self.is_not_in_free_group: self._release_page_ids(torch.unique(free_index // self.page_size)) else: - self.free_group.append(free_index) + self.free_group.append(self._copy_for_free_group(free_index)) if self.debug_mode: self._debug_check_no_duplicate_pages() @@ -298,7 +298,9 @@ class PagedTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): if self.debug_mode: self._debug_check_no_duplicate_pages() else: - self.free_page_reps_group.extend(pieces) + self.free_page_reps_group.extend( + self._copy_for_free_group(piece) for piece in pieces + ) def _debug_check_no_duplicate_pages(self): # span both containers: need_sort (PD disagg) routes frees into release_pages diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index b8f841294..b429b40d0 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -325,7 +325,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.full_attn_allocator.free(free_index) self.free_swa(free_index) else: - self.free_group.append(free_index) + self.free_group.append(self._copy_for_free_group(free_index)) assert ( self.full_attn_allocator.available_size() <= self.full_attn_allocator.size ) @@ -350,7 +350,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): return if not self.is_not_in_free_group: - self.swa_free_group.append(free_index) + self.swa_free_group.append(self._copy_for_free_group(free_index)) return if self.page_size == 1: @@ -500,7 +500,7 @@ class PureSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): if self.is_not_in_free_group: self.swa_attn_allocator.free(free_index[free_index > 0]) else: - self.free_group.append(free_index) + self.free_group.append(self._copy_for_free_group(free_index)) assert self.swa_attn_allocator.available_size() <= self.swa_attn_allocator.size def free_swa(self, free_index: torch.Tensor): @@ -509,7 +509,7 @@ class PureSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): if self.is_not_in_free_group: self.swa_attn_allocator.free(free_index[free_index > 0]) else: - self.free_group.append(free_index) + self.free_group.append(self._copy_for_free_group(free_index)) def free_group_begin(self): self.is_not_in_free_group = False diff --git a/python/sglang/srt/mem_cache/allocator/token.py b/python/sglang/srt/mem_cache/allocator/token.py index 990e468dd..06aceef12 100644 --- a/python/sglang/srt/mem_cache/allocator/token.py +++ b/python/sglang/srt/mem_cache/allocator/token.py @@ -73,7 +73,7 @@ class TokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): else: self.free_pages = torch.cat((self.free_pages, free_index)) else: - self.free_group.append(free_index) + self.free_group.append(self._copy_for_free_group(free_index)) def get_cpu_copy(self, indices, mamba_indices=None): return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices) diff --git a/python/sglang/srt/mem_cache/multi_ended_allocator.py b/python/sglang/srt/mem_cache/multi_ended_allocator.py index ec2358bc5..075c2e930 100644 --- a/python/sglang/srt/mem_cache/multi_ended_allocator.py +++ b/python/sglang/srt/mem_cache/multi_ended_allocator.py @@ -954,7 +954,7 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): if free_index is None or free_index.numel() == 0: return if not self.is_not_in_free_group: - self.free_group.append(free_index) + self.free_group.append(self._copy_for_free_group(free_index)) return if self.lazy_compaction: self._free_lazy(free_index) @@ -1928,7 +1928,7 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): if free_index is None or free_index.numel() == 0: return if not self.is_not_in_free_group: - self.free_group.append(free_index) + self.free_group.append(self._copy_for_free_group(free_index)) return self.full_attn_allocator.free(free_index) self.full_attn_allocator.clear_inverse_history() @@ -2397,7 +2397,7 @@ class UnifiedSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): if free_index is None or free_index.numel() == 0: return if not self.is_not_in_free_group: - self.free_group.append(free_index) + self.free_group.append(self._copy_for_free_group(free_index)) return # Free both peers; the per-sub-pool v2p IS the mapping, so order isn't # load-bearing. Filter the swa side to skip already-tombstoned virtuals diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 7609d44af..2916f8c64 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -511,10 +511,8 @@ class RadixCache(KVCacheEventMixin, BasePrefixCache): ) new_prefix_len = result.prefix_len - # Use the out-of-place values copy so the allocator can safely defer or group - # this free after req_to_token is overwritten below. self.token_to_kv_pool_allocator.free_segment( - values[req.cache_protected_len : new_prefix_len], + kv_indices[req.cache_protected_len : new_prefix_len], start_pos=req.cache_protected_len, ) diff --git a/test/registered/unit/mem_cache/test_paged_free_segment.py b/test/registered/unit/mem_cache/test_paged_free_segment.py index 77da04a4f..b2630cb8e 100644 --- a/test/registered/unit/mem_cache/test_paged_free_segment.py +++ b/test/registered/unit/mem_cache/test_paged_free_segment.py @@ -82,6 +82,19 @@ class TestFreeSegment(unittest.TestCase): alloc.free_group_end() self.assertEqual(len(alloc.free_pages), before + 2) + def test_group_owns_deferred_page_representatives(self): + alloc = _make_allocator() + row = _make_kv_row(alloc, 2 * PAGE_SIZE) + expected_pages = torch.unique(row // PAGE_SIZE) + + alloc.free_group_begin() + alloc.free_segment(row, start_pos=0) + row.zero_() + alloc.free_group_end() + + freed_pages = alloc.free_pages[: expected_pages.numel()] + self.assertTrue(torch.equal(torch.sort(freed_pages)[0], expected_pages)) + def test_group_end_debug_assert_catches_cross_call_double_free(self): # legacy free() + free_segment() on the same page in one group must # trip free_group_end's debug assert diff --git a/test/registered/unit/mem_cache/test_radix_cache_unit.py b/test/registered/unit/mem_cache/test_radix_cache_unit.py index d1beca13f..6e3782e87 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -31,6 +31,7 @@ from array import array import torch from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored +from sglang.srt.mem_cache.allocator.token import TokenToKVPoolAllocator from sglang.srt.mem_cache.base_prefix_cache import ( EvictParams, EvictResult, @@ -391,23 +392,7 @@ class TestRadixCache(unittest.TestCase): ) self.assertEqual(cache.total_size(), 5) - def test_cache_unfinished_req_deferred_free_keeps_original_indices(self): - class DeferredFreeAllocator: - device = torch.device("cpu") - - def __init__(self): - self.free_group = [] - self.freed = None - - def free_group_begin(self): - self.free_group = [] - - def free_segment(self, free_index, *, start_pos): - self.free_group.append(free_index) - - def free_group_end(self): - self.freed = torch.cat(self.free_group) - + def test_cache_unfinished_req_deferred_free_owns_original_indices(self): class ReqToTokenPool: def __init__(self, row): self.req_to_token = row.unsqueeze(0) @@ -415,11 +400,19 @@ class TestRadixCache(unittest.TestCase): def write(self, indices, values): self.req_to_token[indices] = values - allocator = DeferredFreeAllocator() + allocator = TokenToKVPoolAllocator( + size=16, + dtype=torch.float16, + device="cpu", + kvcache=None, + need_sort=False, + ) cache = RadixCache.create_simulated(mock_allocator=allocator) token_ids = array("q", [1, 2, 3]) - tree_indices = torch.tensor([10, 11, 12], dtype=torch.int64) - request_indices = torch.tensor([20, 21, 22], dtype=torch.int64) + tree_indices = allocator.alloc(3) + request_indices = allocator.alloc(3) + assert tree_indices is not None + assert request_indices is not None cache.insert( InsertParams( key=RadixKey(array("q", token_ids)), @@ -436,11 +429,16 @@ class TestRadixCache(unittest.TestCase): ) req.get_fill_ids.return_value = token_ids + available_before_free = allocator.available_size() allocator.free_group_begin() cache.cache_unfinished_req(req) allocator.free_group_end() - torch.testing.assert_close(allocator.freed, request_indices) + self.assertEqual( + allocator.available_size(), + available_before_free + request_indices.numel(), + ) + torch.testing.assert_close(allocator.free_pages[-3:], request_indices) torch.testing.assert_close( cache.req_to_token_pool.req_to_token[0], tree_indices ) diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index b36125b54..52822984d 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -224,7 +224,7 @@ class TestSWA(unittest.TestCase): allocator.free_swa(full_indices[1:2]) self.assertEqual(allocator.swa_available_size(), 16) - def test_free_swa_batches_with_free_group(self): + def test_free_swa_group_owns_deferred_indices(self): _, allocator, _ = _build_swa_tree( is_eagle=False, kv_size=32, @@ -235,6 +235,7 @@ class TestSWA(unittest.TestCase): indices = _swa_alloc(allocator, size) assert indices is not None index_batches.append(indices) + original_indices = torch.cat([indices.clone() for indices in index_batches]) available_before_free = allocator.swa_available_size() allocator.free_group_begin() @@ -243,20 +244,19 @@ class TestSWA(unittest.TestCase): self.assertEqual(len(allocator.swa_free_group), len(index_batches)) self.assertEqual(allocator.swa_available_size(), available_before_free) - + for indices in index_batches: + indices.zero_() allocator.free_group_end() - all_indices = torch.cat(index_batches).to(torch.int64) - self.assertEqual(allocator.swa_free_group, []) self.assertTrue( torch.equal( - allocator.full_to_swa_index_mapping[all_indices], - torch.zeros_like(all_indices), + allocator.full_to_swa_index_mapping[original_indices.to(torch.int64)], + torch.zeros_like(original_indices), ) ) self.assertEqual( allocator.swa_available_size(), - available_before_free + all_indices.numel(), + available_before_free + original_indices.numel(), ) def test_swa_radix_cache_1(self):