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

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

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