diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 15e9f8805..6483c53a6 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -3232,6 +3232,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): eviction_interval = max(1, envs.SGLANG_SWA_EVICTION_INTERVAL.get()) swa_maintenance_step = (self.forward_iter or 0) % eviction_interval == 0 + self.token_to_kv_pool_allocator.free_group_begin() for idx, req in enumerate(self.reqs): if self.forward_mode.is_decode(): # We set evict_swa condition here with two reasons: @@ -3276,6 +3277,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self._evict_swa(req, pre_len) else: self._evict_swa(req, pre_len) + self.token_to_kv_pool_allocator.free_group_end() def _evict_swa(self, req: Req, pre_len: int): assert self.tree_cache.supports_swa(), "prefix cache must support swa" diff --git a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py index dbf83eea9..24f0e089c 100644 --- a/python/sglang/srt/managers/scheduler_components/batch_result_processor.py +++ b/python/sglang/srt/managers/scheduler_components/batch_result_processor.py @@ -195,6 +195,7 @@ class SchedulerBatchResultProcessor: result: Union[GenerationBatchResult, EmbeddingBatchResult], ): skip_stream_req = None + self.token_to_kv_pool_allocator.free_group_begin() if self.is_generation: if result.copy_done is not None: @@ -360,6 +361,7 @@ class SchedulerBatchResultProcessor: req.inflight_middle_chunks -= 1 req.time_stats.set_last_chunked_prefill_finish_time() + self.token_to_kv_pool_allocator.free_group_end() self.output_streamer.stream_output( batch.reqs, batch.return_logprob, skip_stream_req ) diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index 308066de3..b8f841294 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -95,6 +95,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.release_pages = None self.is_not_in_free_group = True self.free_group = [] + self.swa_free_group = [] self._kvcache = kvcache self.clear() @@ -348,6 +349,10 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): if free_index.numel() == 0: return + if not self.is_not_in_free_group: + self.swa_free_group.append(free_index) + return + if self.page_size == 1: mapping_indices = free_index else: @@ -358,6 +363,17 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.swa_attn_allocator.free(swa_indices) self.full_to_swa_index_mapping[mapping_indices] = 0 + def free_group_begin(self): + super().free_group_begin() + self.swa_free_group = [] + + def free_group_end(self): + super().free_group_end() + if self.swa_free_group: + swa_free_group = self.swa_free_group + self.swa_free_group = [] + self.free_swa(torch.cat(swa_free_group)) + def _expand_to_full_pages(self, indices: torch.Tensor) -> torch.Tensor: pages = torch.unique(indices // self.page_size) page_offsets = torch.arange( @@ -386,6 +402,7 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): self.full_to_swa_index_mapping[:-1].fill_(0) self.is_not_in_free_group = True self.free_group = [] + self.swa_free_group = [] def get_cpu_copy(self, indices, mamba_indices=None): return self._kvcache.get_cpu_copy(indices, mamba_indices=mamba_indices) @@ -489,7 +506,10 @@ class PureSWATokenToKVPoolAllocator(SWATokenToKVPoolAllocator): def free_swa(self, free_index: torch.Tensor): if free_index.numel() == 0: return - self.swa_attn_allocator.free(free_index[free_index > 0]) + if self.is_not_in_free_group: + self.swa_attn_allocator.free(free_index[free_index > 0]) + else: + self.free_group.append(free_index) def free_group_begin(self): self.is_not_in_free_group = False diff --git a/python/sglang/srt/mem_cache/radix_cache.py b/python/sglang/srt/mem_cache/radix_cache.py index 2916f8c64..7609d44af 100644 --- a/python/sglang/srt/mem_cache/radix_cache.py +++ b/python/sglang/srt/mem_cache/radix_cache.py @@ -511,8 +511,10 @@ 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( - kv_indices[req.cache_protected_len : new_prefix_len], + values[req.cache_protected_len : new_prefix_len], start_pos=req.cache_protected_len, ) 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 778afa06c..d1beca13f 100644 --- a/test/registered/unit/mem_cache/test_radix_cache_unit.py +++ b/test/registered/unit/mem_cache/test_radix_cache_unit.py @@ -391,6 +391,60 @@ 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) + + class ReqToTokenPool: + def __init__(self, row): + self.req_to_token = row.unsqueeze(0) + + def write(self, indices, values): + self.req_to_token[indices] = values + + allocator = DeferredFreeAllocator() + 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) + cache.insert( + InsertParams( + key=RadixKey(array("q", token_ids)), + value=tree_indices, + ) + ) + cache.req_to_token_pool = ReqToTokenPool(request_indices.clone()) + req = unittest.mock.Mock( + req_pool_idx=0, + cache_protected_len=0, + extra_key=None, + priority=0, + last_node=cache.root_node, + ) + req.get_fill_ids.return_value = token_ids + + allocator.free_group_begin() + cache.cache_unfinished_req(req) + allocator.free_group_end() + + torch.testing.assert_close(allocator.freed, request_indices) + torch.testing.assert_close( + cache.req_to_token_pool.req_to_token[0], tree_indices + ) + def test_kv_cache_events(self): """Test KV cache events functionality.""" test_cases = [ diff --git a/test/registered/unit/mem_cache/test_swa_unittest.py b/test/registered/unit/mem_cache/test_swa_unittest.py index 534303900..b36125b54 100644 --- a/test/registered/unit/mem_cache/test_swa_unittest.py +++ b/test/registered/unit/mem_cache/test_swa_unittest.py @@ -224,6 +224,41 @@ 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): + _, allocator, _ = _build_swa_tree( + is_eagle=False, + kv_size=32, + kv_size_swa=32, + ) + index_batches = [] + for size in (2, 3, 1, 4): + indices = _swa_alloc(allocator, size) + assert indices is not None + index_batches.append(indices) + + available_before_free = allocator.swa_available_size() + allocator.free_group_begin() + for indices in index_batches: + allocator.free_swa(indices) + + self.assertEqual(len(allocator.swa_free_group), len(index_batches)) + self.assertEqual(allocator.swa_available_size(), available_before_free) + + 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), + ) + ) + self.assertEqual( + allocator.swa_available_size(), + available_before_free + all_indices.numel(), + ) + def test_swa_radix_cache_1(self): # args req_size = 10