[mem_cache] Split duplicate insert frees at the SWA eviction floor (#37481)

This commit is contained in:
Liangsheng Yin
2026-09-01 23:01:40 -07:00
committed by GitHub
parent a6a19f9290
commit 832d029870
5 changed files with 81 additions and 28 deletions
@@ -1292,7 +1292,9 @@ def test_swa_straddling_insert_crosses_the_boundary_actions():
assert free_tail.indices[0].tolist() == [12, 13]
assert isinstance(rebuild, SWARebuild)
assert rebuild.source_value.tolist() == [22, 23]
assert isinstance(free_duplicates, FreeDeviceKV)
# Below the floor the duplicate's SWA peers are gone: full side only.
assert isinstance(free_duplicates, FreeDeviceKVFullOnly)
assert free_duplicates.indices[0].tolist() == [20, 21]
def test_every_pool_name_crosses_the_prefetch_commit_boundary():
@@ -58,6 +58,7 @@ from sglang.srt.mem_cache.unified_cache.cache_action import (
FreeComponentDeviceSlot,
FreeComponentHostSlot,
FreeDeviceKV,
FreeDeviceKVFullOnly,
RebuildFullToSWAMapping,
RecoverSWAWithLockedFull,
ReplaceWriteThroughOnNodeSplit,
@@ -8096,6 +8097,30 @@ class TestResumableInsertWalkSWA(_InsertWalkSuite):
self.assertIsNotNone(_device_value(cache, window_node, ComponentType.FULL))
cache.sanity_check()
def test_dup_slice_below_eviction_floor_frees_full_only(self):
"""A re-insert whose duplicate slice starts below the request's eviction
floor gives back only the full side there: those SWA peers are gone."""
sw = self.cfg.sliding_window_size
cache, allocator, _ = build_fixture(self.cfg)
seq = list(range(1, 2 * sw + 1))
key = RadixKey(array("q", seq))
cache.insert(InsertParams(key=key, value=self._alloc(allocator, len(seq))))
value = self._alloc(allocator, len(seq))
with mock.patch.object(
cache, "_apply_cache_action", wraps=cache._apply_cache_action
) as spy:
cache.insert(InsertParams(key=key, value=value, swa_evicted_seqlen=sw))
actions = [c.args[0] for c in spy.call_args_list]
full_only = [
i for a in actions if isinstance(a, FreeDeviceKVFullOnly) for i in a.indices
]
both = [i for a in actions if isinstance(a, FreeDeviceKV) for i in a.indices]
torch.testing.assert_close(torch.cat(full_only), value[:sw])
torch.testing.assert_close(torch.cat(both), value[sw:])
cache.sanity_check()
def test_dec_swa_lock_only_early_release_keeps_full_lock(self):
"""The scheduler's early SWA release (decode past the window) drops
only the SWA lock; the Full path lock stays held until dec_lock_ref."""