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

[unified-memory] Stop eviction when shared allocation capacity is sufficient (#33091)

Browse Source 
Co-authored-by: seokwoosong <seokwoosong@users.noreply.github.com>
This commit is contained in:
SuperSong
2026-08-26 02:06:32 -07:00
committed by GitHub
co-authored by seokwoosong
parent ffc431cd4c
commit 2511743bd7
22 changed files with 885 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/hardware_backend/mlx/test_attention_patching.py
+12
View File
@@ -877,6 +877,17 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
self.assertEqual(forked.tolist(), [3])
self.assertEqual(restored[0].state[0].tolist(), [1.0])
self.assertEqual(pool.available_size(), 3)
self.assertEqual(pool.schedulable_available_size(), 3)
def test_auxiliary_state_pool_returns_unused_group_slots(self):
pool = MlxAuxiliaryStatePool(size=4, device="cpu")
pool.alloc_group_begin(3)
allocated = pool.alloc(1)
pool.alloc_group_end()
self.assertEqual(allocated.tolist(), [1])
self.assertEqual(pool.available_size(), 3)
def test_auxiliary_state_pool_restores_instance_meta_state(self):
pool = MlxAuxiliaryStatePool(size=2, device="cpu")
@@ -916,6 +927,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
self.assertIsNotNone(auxiliary_state_idx)
self.assertIsNone(req.req_pool_idx)
self.assertIsNotNone(req.mamba_pool_idx)
self.assertIs(pool.mamba_allocator, pool.mamba_pool)
self.assertEqual(pool.auxiliary_state_pool.available_size(), 3)
pool.free_auxiliary_state_cache(req)
self.assertIsNone(req.mamba_pool_idx)
test/registered/unit/mem_cache/test_decode_radix_lock_ref.py
+1 -1
View File
@@ -144,7 +144,7 @@ class TestDecodeLockRefScenarios(unittest.TestCase):
error = queue._reclaim_swa_tail_capacity(129, "req-1")
self.assertIsNone(error)
params = queue.tree_cache.evict.call_args.args[0]
params = queue.tree_cache.evict_for_alloc.call_args.args[0]
self.assertEqual(params.num_tokens, 0)
self.assertEqual(params.swa_num_tokens, 128)
test/registered/unit/mem_cache/test_hybrid_pool_assembler.py
+38 -1
View File
@@ -2,9 +2,12 @@
import unittest
from types import SimpleNamespace
from unittest.mock import patch
from unittest.mock import MagicMock, patch
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
from sglang.srt.mem_cache.hybrid_cache.hybrid_pool_assembler import (
_evict_mamba_for_device_alloc,
_evict_swa_for_device_alloc,
_split_hicache_size,
build_full_draft_pools,
)
@@ -23,6 +26,40 @@ class _Pool:
return self._kv_bytes
class TestDeviceAllocEviction(CustomTestCase):
def test_swa_evicts_only_allocation_shortfall(self):
cache = MagicMock()
cache.token_to_kv_pool_allocator.swa_available_size.return_value = 8
_evict_swa_for_device_alloc(cache, required_size=10)
cache.evict_for_alloc.assert_called_once_with(EvictParams(swa_num_tokens=2))
cache.evict.assert_not_called()
def test_mamba_evicts_only_allocation_shortfall(self):
cache = MagicMock()
allocator = cache.req_to_token_pool.mamba_allocator
allocator.schedulable_available_size.return_value = 8
_evict_mamba_for_device_alloc(cache, required_size=10)
cache.evict_for_alloc.assert_called_once_with(EvictParams(mamba_num=2))
cache.evict.assert_not_called()
def test_sufficient_capacity_skips_eviction(self):
cache = MagicMock()
cache.token_to_kv_pool_allocator.swa_available_size.return_value = 10
cache.req_to_token_pool.mamba_allocator.schedulable_available_size.return_value = (
10
)
_evict_swa_for_device_alloc(cache, required_size=10)
_evict_mamba_for_device_alloc(cache, required_size=10)
cache.evict_for_alloc.assert_not_called()
cache.evict.assert_not_called()
class TestSplitHicacheSize(CustomTestCase):
def test_splits_total_budget_by_device_bytes(self):
# scalar and (k, v) tuple return shapes both supported
test/registered/unit/mem_cache/test_mamba_donated_alloc_ratio.py
+16 -6
View File
@@ -55,10 +55,10 @@ class _RatioCache:
self.component_evictable_size_ = {ComponentType.MAMBA: 0}
self.component_protected_size_ = {ComponentType.MAMBA: 0}
self.prefix_nodes = []
self.alloc_evict_params = []
def evict(self, params: EvictParams):
# Reclaim up to mamba_num evictable (unlocked) prefix snapshots, mirroring
# what the real tree eviction can hand back under mamba pressure.
def evict_for_alloc(self, params: EvictParams):
self.alloc_evict_params.append(params)
need = params.mamba_num
for node in list(self.prefix_nodes):
if need <= 0:
@@ -130,7 +130,9 @@ class TestMambaRatioEnvGate(unittest.TestCase):
return KVCacheConfigurator._calculate_mamba_ratio(fake)
def test_flag_off_restores_original_ratios(self):
r = lambda **kw: self._ratio(skip=False, **kw)
def r(**kwargs):
return self._ratio(skip=False, **kwargs)
self.assertEqual(
r(extra_buffer=False, lazy=False, disable_overlap=True), 3
) # no_buffer
@@ -142,7 +144,9 @@ class TestMambaRatioEnvGate(unittest.TestCase):
) # overlap
def test_flag_on_drops_base_but_keeps_no_buffer(self):
r = lambda **kw: self._ratio(skip=True, **kw)
def r(**kwargs):
return self._ratio(skip=True, **kwargs)
self.assertEqual(
r(extra_buffer=False, lazy=False, disable_overlap=True), 3
) # no_buffer
@@ -212,9 +216,12 @@ class TestDecSwaLockSkip(unittest.TestCase):
class TestMambaDonatedAllocRatio(unittest.TestCase):
def test_prefill_peak_ratio2_exhausts_pool(self):
# pool = 2N, all N prefixes admission-locked: no evictable victim.
component, _, _ = _build_peak(pool_size=2 * N, lock_prefixes=True)
component, cache, _ = _build_peak(pool_size=2 * N, lock_prefixes=True)
with self.assertRaisesRegex(AssertionError, "Can not alloc mamba cache"):
component._alloc_mamba_slot()
self.assertEqual(
cache.alloc_evict_params, [EvictParams(num_tokens=0, mamba_num=1)]
)
def test_prefill_peak_ratio3_has_headroom(self):
# pool = 3N: N free slots remain after own + locked prefix.
@@ -230,6 +237,9 @@ class TestMambaDonatedAllocRatio(unittest.TestCase):
slot = component._alloc_mamba_slot()
self.assertIsNotNone(slot)
self.assertEqual(len(cache.prefix_nodes), N - 1)
self.assertEqual(
cache.alloc_evict_params, [EvictParams(num_tokens=0, mamba_num=1)]
)
class TestPPMambaPoolSizing(unittest.TestCase):
test/registered/unit/mem_cache/test_unified_radix_allocation_eviction.py
+132
View File
@@ -0,0 +1,132 @@
"""CPU-only tests for allocation-aware UnifiedRadixCache eviction."""
import unittest
from unittest.mock import MagicMock
from sglang.srt.mem_cache.base_prefix_cache import EvictParams
from sglang.srt.mem_cache.common import evict_from_tree_cache
from sglang.srt.mem_cache.unified_cache.components import ComponentType
from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
class TestUnifiedRadixAllocationEviction(CustomTestCase):
@staticmethod
def _build_cache(*, collateral_capacity_gain: int):
cache = object.__new__(UnifiedRadixCache)
cache.disable = False
cache.tree_components = (ComponentType.FULL, ComponentType.MAMBA)
cache.is_swa_enabled = False
cache.cache_controller = None
cache.metrics_collector = None
cache.tree_core = MagicMock()
capacity = {"available": 30}
allocator = MagicMock()
allocator.available_size.side_effect = lambda: capacity["available"]
cache.token_to_kv_pool_allocator = allocator
cache.req_to_token_pool = MagicMock()
leaf_count = {"value": 0}
def next_node(component_type, tracker):
if tracker[component_type] >= 70:
return None, False
return leaf_count["value"] + 1, True
def evict_leaf(_node_id, tracker):
leaf_count["value"] += 1
tracker[ComponentType.FULL] += 20
tracker[ComponentType.MAMBA] += 1
capacity["available"] += (
collateral_capacity_gain if leaf_count["value"] == 1 else 20
)
return None
cache._evict_device_next_node = MagicMock(side_effect=next_node)
cache._evict_device_leaf = MagicMock(side_effect=evict_leaf)
return cache, capacity, leaf_count
def test_allocation_eviction_stops_when_shared_capacity_is_sufficient(self):
cache, capacity, leaf_count = self._build_cache(collateral_capacity_gain=70)
result = cache.evict_for_alloc(EvictParams(num_tokens=70))
self.assertEqual(capacity["available"], 100)
self.assertEqual(leaf_count["value"], 1)
self.assertEqual(result.num_tokens_evicted, 20)
self.assertEqual(result.mamba_num_evicted, 1)
def test_explicit_evict_preserves_component_count_semantics(self):
cache, _, leaf_count = self._build_cache(collateral_capacity_gain=70)
result = cache.evict(EvictParams(num_tokens=70))
self.assertEqual(leaf_count["value"], 4)
self.assertEqual(result.num_tokens_evicted, 80)
self.assertEqual(result.mamba_num_evicted, 4)
def test_c128_component_keeps_zero_quota(self):
cache, _, _ = self._build_cache(collateral_capacity_gain=70)
cache.tree_components = (ComponentType.FULL, ComponentType.C128)
cache._evict_device_next_node.side_effect = None
cache._evict_device_next_node.return_value = (None, False)
result = cache.evict(EvictParams(num_tokens=1))
self.assertEqual(result.num_tokens_evicted, 0)
cache.tree_core.evict_device_start.assert_called_once_with(
ComponentType.FULL, 1
)
def test_mamba_allocation_counts_collateral_full_capacity(self):
cache = object.__new__(UnifiedRadixCache)
cache.disable = False
cache.tree_components = (ComponentType.FULL, ComponentType.MAMBA)
cache.is_swa_enabled = False
cache.cache_controller = None
cache.metrics_collector = None
cache.tree_core = MagicMock()
cache.token_to_kv_pool_allocator = MagicMock()
capacity = {"available": 0}
mamba_allocator = MagicMock()
mamba_allocator.schedulable_available_size.side_effect = lambda: capacity[
"available"
]
cache.req_to_token_pool = MagicMock(mamba_allocator=mamba_allocator)
def next_node(component_type, tracker):
return (None, False) if tracker[component_type] >= 3 else (1, True)
def evict_leaf(_node_id, tracker):
tracker[ComponentType.FULL] += 20
tracker[ComponentType.MAMBA] += 1
capacity["available"] += 3
return None
cache._evict_device_next_node = MagicMock(side_effect=next_node)
cache._evict_device_leaf = MagicMock(side_effect=evict_leaf)
result = cache.evict_for_alloc(EvictParams(mamba_num=3))
self.assertEqual(capacity["available"], 3)
self.assertEqual(result.num_tokens_evicted, 20)
self.assertEqual(result.mamba_num_evicted, 1)
def test_common_helper_uses_allocation_aware_entry_point(self):
tree_cache = MagicMock()
tree_cache.is_chunk_cache.return_value = False
tree_cache.token_to_kv_pool_allocator.available_size.return_value = 30
evict_from_tree_cache(tree_cache, num_tokens=100)
tree_cache.evict_for_alloc.assert_called_once_with(EvictParams(num_tokens=70))
tree_cache.evict.assert_not_called()
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+173 -7
View File
@@ -343,6 +343,7 @@ def build_fixture(
cfg: CacheConfig,
*,
enable_kv_cache_events: bool = False,
enable_session_radix_cache: bool = False,
tree_page_size: Optional[int] = None,
mamba_cache_chunk_size: Optional[int] = None,
):
@@ -472,6 +473,7 @@ def build_fixture(
tree_components=cfg.components,
enable_mamba_extra_buffer=cfg.enable_mamba_extra_buffer,
enable_kv_cache_events=enable_kv_cache_events,
enable_session_radix_cache=enable_session_radix_cache,
eviction_policy=cfg.eviction_policy,
is_eagle=cfg.is_eagle,
)
@@ -481,6 +483,162 @@ def build_fixture(
return cache, allocator, req_to_token_pool
@unittest.skipUnless(torch.cuda.is_available(), "cache fixtures need CUDA")
class TestUnifiedRadixAllocationEvictionRealComponents(CustomTestCase):
"""Allocation targets are observed between real auxiliary-tree steps."""
_SHORTFALL = 100
def _insert(self, cache, allocator, req_to_token_pool, tokens) -> None:
value = allocator.alloc(len(tokens))
self.assertIsNotNone(value)
params = InsertParams(
key=RadixKey(array("q", tokens)),
value=value[: len(tokens)],
)
if cache.supports_mamba():
req = Req(
rid=f"mamba-{len(tokens)}",
origin_input_text="",
origin_input_ids=array("q"),
sampling_params=SamplingParams(temperature=0, max_new_tokens=1),
)
req_to_token_pool.alloc([req])
params.mamba_value = req.mamba_pool_idx.unsqueeze(0)
cache.insert(params)
def _build_internal_chain(self, component_type, enable_session_radix_cache):
cfg = (
CacheConfig(
components=(ComponentType.FULL, ComponentType.SWA),
sliding_window_size=128,
)
if component_type is ComponentType.SWA
else CacheConfig(
components=(ComponentType.FULL, ComponentType.MAMBA),
mamba_cache_size=8,
)
)
cache, allocator, req_to_token_pool = build_fixture(
cfg, enable_session_radix_cache=enable_session_radix_cache
)
for length in (2, 4, 6):
self._insert(
cache,
allocator,
req_to_token_pool,
list(range(1, length + 1)),
)
lru = cache.tree_core.lru_lists[component_type]
first = lru.get_lru_no_lock()
second = lru.get_prev_no_lock(first)
leaf = lru.get_prev_no_lock(second)
self.assertNotIn(first, cache.tree_core.evictable_device_leaves)
self.assertNotIn(second, cache.tree_core.evictable_device_leaves)
self.assertIn(leaf, cache.tree_core.evictable_device_leaves)
for node in (first, second, leaf):
self.assertIsNotNone(node.component_data[component_type].value)
self.assertIsNotNone(node.component_data[ComponentType.FULL].value)
return cache, first, second, leaf
def _evict_for_alloc_after_first_drain(self, cache, component_type):
capacity = {"available": 0}
auxiliary_drains = {"count": 0}
real_available_size = cache._component_available_size
real_free_values = cache._free_values
def available_size(requested_type):
if requested_type is component_type:
return capacity["available"]
return real_available_size(requested_type)
def free_values(device_frees, host_frees):
freed_auxiliary = bool(device_frees.get(component_type))
real_free_values(device_frees, host_frees)
if freed_auxiliary:
auxiliary_drains["count"] += 1
capacity["available"] = self._SHORTFALL
params = (
EvictParams(swa_num_tokens=self._SHORTFALL)
if component_type is ComponentType.SWA
else EvictParams(mamba_num=self._SHORTFALL)
)
with (
mock.patch.object(
cache, "_component_available_size", side_effect=available_size
),
mock.patch.object(cache, "_free_values", side_effect=free_values),
):
result = cache.evict_for_alloc(params)
return result, auxiliary_drains["count"]
def test_allocation_target_stops_after_one_internal_tombstone(self):
for component_type in (ComponentType.SWA, ComponentType.MAMBA):
for enable_session_radix_cache in (False, True):
with self.subTest(
component_type=component_type,
enable_session_radix_cache=enable_session_radix_cache,
):
cache, first, second, leaf = self._build_internal_chain(
component_type, enable_session_radix_cache
)
first_size = len(first.component_data[component_type].value)
result, drain_count = self._evict_for_alloc_after_first_drain(
cache, component_type
)
self.assertIsNone(first.component_data[component_type].value)
self.assertIsNotNone(second.component_data[component_type].value)
self.assertIsNotNone(leaf.component_data[component_type].value)
self.assertIsNotNone(leaf.component_data[ComponentType.FULL].value)
self.assertEqual(result.num_tokens_evicted, 0)
self.assertEqual(drain_count, 1)
evicted = (
result.swa_num_tokens_evicted
if component_type is ComponentType.SWA
else result.mamba_num_evicted
)
self.assertEqual(evicted, first_size)
cache.sanity_check()
def test_explicit_evict_continues_across_internal_steps(self):
for component_type in (ComponentType.SWA, ComponentType.MAMBA):
for enable_session_radix_cache in (False, True):
with self.subTest(
component_type=component_type,
enable_session_radix_cache=enable_session_radix_cache,
):
cache, first, second, leaf = self._build_internal_chain(
component_type, enable_session_radix_cache
)
request_count = sum(
len(node.component_data[component_type].value)
for node in (first, second)
)
params = (
EvictParams(swa_num_tokens=request_count)
if component_type is ComponentType.SWA
else EvictParams(mamba_num=request_count)
)
result = cache.evict(params)
self.assertIsNone(first.component_data[component_type].value)
self.assertIsNone(second.component_data[component_type].value)
self.assertIsNotNone(leaf.component_data[component_type].value)
self.assertIsNotNone(leaf.component_data[ComponentType.FULL].value)
evicted = (
result.swa_num_tokens_evicted
if component_type is ComponentType.SWA
else result.mamba_num_evicted
)
self.assertEqual(evicted, request_count)
cache.sanity_check()
class TestUnifiedRadixCacheEagleHiCacheStorageKey(CustomTestCase):
cfg = CacheConfig(
page_size=4,
@@ -5363,10 +5521,12 @@ class UnifiedRadixCacheSuite:
"alloc",
side_effect=[None, retry_slot],
),
mock.patch.object(cache, "evict", autospec=True) as evict,
mock.patch.object(
cache, "evict_for_alloc", autospec=True
) as evict_for_alloc,
):
prep = comp.prepare_load_back(leaf.id, req=req)
evict.assert_called_once_with(EvictParams(num_tokens=0, mamba_num=1))
evict_for_alloc.assert_called_once_with(EvictParams(num_tokens=0, mamba_num=1))
self.assertIs(prep.allocated_mamba_slot, retry_slot)
self.assertEqual(int(req.mamba_pool_idx), int(retry_slot[0]))
@@ -5790,13 +5950,15 @@ class UnifiedRadixCacheSuite:
int(swa_xfer.host_indices.numel()),
)
with mock.patch.object(cache, "evict", wraps=cache.evict) as evict_mock:
with mock.patch.object(
cache, "evict_for_alloc", wraps=cache.evict_for_alloc
) as evict_for_alloc_mock:
self.assertTrue(cache.load_back(leaf.id))
# Full pre-eviction must not be triggered by SWA pool pressure.
full_pre_evict_calls = [
call
for call in evict_mock.call_args_list
for call in evict_for_alloc_mock.call_args_list
if call.args and call.args[0].num_tokens > 0
]
self.assertEqual(full_pre_evict_calls, [])
@@ -5807,7 +5969,7 @@ class UnifiedRadixCacheSuite:
call.args
and call.args[0].num_tokens == 0
and call.args[0].swa_num_tokens > 0
for call in evict_mock.call_args_list
for call in evict_for_alloc_mock.call_args_list
)
)
@@ -7026,9 +7188,13 @@ class TestReturnedValuesDrain(_InsertWalkSuite):
cases = [
(
"evict_device_next_node",
lambda: make(EvictDeviceNextNodeResult, node_id=node.id),
lambda: make(
EvictDeviceNextNodeResult,
node_id=node.id,
made_progress=True,
),
lambda: cache._evict_device_next_node(ComponentType.FULL, tracker),
node.id,
(node.id, True),
),
(
"evict_device_leaf",