minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Leon Gao
2026-08-08 03:34:44 -07:00
committed by GitHub
parent a1ca76b24b
commit cfb354bcfc
11 changed files with 60 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/allocator_npu.py
+1 -1
View File
@@ -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
+5
View File
@@ -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
+2 -2
View File
@@ -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
+4 -2
View File
@@ -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
+4 -4
View File
@@ -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
+1 -1
View File
@@ -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
+3 -3
View File
@@ -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
+1 -3
View File
@@ -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
+13
View File
@@ -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
+19 -21
View File
@@ -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
+7 -7
View File
@@ -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):