[Bugfix] Fix batched KV free aliasing (#34067)

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)

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

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

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

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

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)

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

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

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

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):
+    def test_cache_unfinished_req_deferred_free_owns_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)
@@ -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)
+        tree_indices = allocator.alloc(3)
-        request_indices = torch.tensor([20, 21, 22], dtype=torch.int64)
+        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
         )

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],
+                allocator.full_to_swa_index_mapping[original_indices.to(torch.int64)],
-                torch.zeros_like(all_indices),
+                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):