[srt] Batch scheduler cache frees (#33475)

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"

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
         )

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

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,
         )
 

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 = [

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