From bede776c2af38d6fb84d6fa365bba8cc033af6fd Mon Sep 17 00:00:00 2001 From: SuperSong Date: Thu, 10 Sep 2026 05:46:51 +0900 Subject: [PATCH] fix(unified-memory): evict Full KV for Mamba byte shortfalls (#36713) Co-authored-by: Yangmin Li Co-authored-by: YAMY <74099316+YAMY1234@users.noreply.github.com> --- python/sglang/srt/mem_cache/allocator/base.py | 20 +- .../mem_cache/allocator/unified_hybrid_swa.py | 26 ++ .../srt/mem_cache/allocator/unified_mamba.py | 24 + .../mem_cache/allocator/unified_sub_pool.py | 42 +- .../srt/mem_cache/unified_radix_cache.py | 112 ++++- .../mem_cache/test_multi_ended_allocator.py | 218 ++++++++- .../test_unified_radix_allocation_eviction.py | 434 ++++++++++++++++++ .../unit/mem_cache/test_unified_tri_pool.py | 115 +++++ 8 files changed, 969 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/mem_cache/allocator/base.py b/python/sglang/srt/mem_cache/allocator/base.py index ffbc4b96d..fa640b650 100644 --- a/python/sglang/srt/mem_cache/allocator/base.py +++ b/python/sglang/srt/mem_cache/allocator/base.py @@ -16,7 +16,7 @@ limitations under the License. from __future__ import annotations import abc -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Protocol import torch @@ -24,6 +24,20 @@ if TYPE_CHECKING: from sglang.srt.mem_cache.memory_pool import KVCache +class MambaFullCacheDonor(Protocol): + """Allocator capability for reclaiming Full KV on Mamba byte pressure.""" + + def flush_deferred_full_frees(self) -> None: ... + + def full_tokens_before_mamba_recheck(self, target_size: int) -> int: + """Lower bound on new Full tokens before preparation can help.""" + ... + + def prepare_mamba_allocation(self, target_size: int) -> None: + """Expose layout-specific reclaim so Mamba capacity is queryable.""" + ... + + class BaseTokenToKVPoolAllocator(abc.ABC): @abc.abstractmethod def __init__( @@ -74,6 +88,10 @@ class BaseTokenToKVPoolAllocator(abc.ABC): violation strings, empty when healthy. Static pools have no byte model.""" return [] + def mamba_full_cache_donor(self) -> MambaFullCacheDonor | None: + """Return the shared-pool donor capability, if this allocator has one.""" + return None + def debug_print(self) -> str: return "" diff --git a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py index 1f199fc69..775e0f778 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py +++ b/python/sglang/srt/mem_cache/allocator/unified_hybrid_swa.py @@ -23,6 +23,7 @@ import torch from torch.profiler import record_function from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator +from sglang.srt.mem_cache.allocator.base import MambaFullCacheDonor from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator from sglang.srt.mem_cache.allocator.unified_sub_pool import ( FloatMultiEndedAllocator, @@ -30,6 +31,8 @@ from sglang.srt.mem_cache.allocator.unified_sub_pool import ( _chain_byte_accounting_violations, _end_pair_chain, _float_open_short_side, + _flush_deferred_free_group, + _full_tokens_before_mamba_recheck, _relieve_for_alloc, ) from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool @@ -845,6 +848,29 @@ class UnifiedMambaSWATokenToKVPoolAllocator(UnifiedSWATokenToKVPoolAllocator): flt = b _float_open_short_side(flt, demand) + def mamba_full_cache_donor(self) -> MambaFullCacheDonor: + return self + + def flush_deferred_full_frees(self) -> None: + """Expose grouped composite frees while preserving the group scope.""" + _flush_deferred_free_group( + self, + (self.free_group, self.free_page_reps_group, self.full_free_group), + ) + + def full_tokens_before_mamba_recheck(self, target_size: int) -> int: + return _full_tokens_before_mamba_recheck( + self.full_attn_allocator, self.mamba_allocator, target_size + ) + + def prepare_mamba_allocation(self, target_size: int) -> None: + """Expose Full reclaim, then move the SWA float away from Mamba.""" + self.flush_deferred_full_frees() + if target_size <= self.mamba_allocator.available_size(): + return + self.full_attn_allocator.flush_for_allocation() + _relieve_for_alloc(self.mamba_allocator, target_size) + def mamba_slot_full_token_cost(self) -> int: """Full-token-equivalents one mamba/conv slot removes from the shared buffer: a tri-pool token costs e_f + e_s bytes, and the quotient is rounded UP.""" diff --git a/python/sglang/srt/mem_cache/allocator/unified_mamba.py b/python/sglang/srt/mem_cache/allocator/unified_mamba.py index bc2ac7c51..4ecaf63b7 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_mamba.py +++ b/python/sglang/srt/mem_cache/allocator/unified_mamba.py @@ -23,10 +23,14 @@ import torch from torch.profiler import record_function from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator +from sglang.srt.mem_cache.allocator.base import MambaFullCacheDonor from sglang.srt.mem_cache.allocator.unified_sub_pool import ( MultiEndedAllocator, _chain_byte_accounting_violations, _end_pair_chain, + _flush_deferred_free_group, + _full_tokens_before_mamba_recheck, + _relieve_for_alloc, ) from sglang.srt.mem_cache.unified_memory_pool import UnifiedKVPool from sglang.srt.runtime_context import get_parallel @@ -146,6 +150,26 @@ class UnifiedMambaTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): def full_available_size(self) -> int: return self.full_attn_allocator.schedulable_available_size() + def mamba_full_cache_donor(self) -> MambaFullCacheDonor: + return self + + def flush_deferred_full_frees(self) -> None: + """Expose grouped Full frees without ending the caller's free group.""" + _flush_deferred_free_group(self, (self.free_group, self.free_page_reps_group)) + + def full_tokens_before_mamba_recheck(self, target_size: int) -> int: + return _full_tokens_before_mamba_recheck( + self.full_attn_allocator, self.mamba_allocator, target_size + ) + + def prepare_mamba_allocation(self, target_size: int) -> None: + """Make deferred Full reclaim visible to the Mamba capacity view.""" + self.flush_deferred_full_frees() + if target_size > self.mamba_allocator.schedulable_available_size(): + return + if target_size > self.mamba_allocator.available_size(): + _relieve_for_alloc(self.mamba_allocator, target_size) + def mamba_slot_full_token_cost(self) -> int: """Full-token-equivalents of shared-gap bytes ONE mamba state consumes; the prefill planner reserves this so admission stays inside the JOINT budget, diff --git a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py index a5286fd58..7857f4196 100644 --- a/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py +++ b/python/sglang/srt/mem_cache/allocator/unified_sub_pool.py @@ -31,6 +31,7 @@ from typing import ( Generic, List, Optional, + Sequence, Set, Tuple, TypeVar, @@ -194,18 +195,50 @@ def _float_open_short_side(flt, demand) -> None: def _relieve_for_alloc(short_pool, need_tokens: int) -> bool: """THE shortfall ladder: every allocation shortfall in the unified pool runs exactly this, whether a single band's own alloc or a composite's coupled - multi-band alloc. `_flush` is called unconditionally -- an eager END no-ops + multi-band alloc. The urgent flush is called unconditionally -- an eager END no-ops and a FLOAT always has boundary absorption to do -- so the ladder never branches on lazy mode, member kind, or layout. """ for m in short_pool._flush_targets(): - m._flush(urgent=True) + m.flush_for_allocation() if need_tokens <= short_pool.available_size(): return True short_pool._ask_float_for_room(need_tokens) return need_tokens <= short_pool.available_size() +def _flush_deferred_free_group( + allocator: BaseTokenToKVPoolAllocator, + pending_groups: Sequence[Optional[Sequence[torch.Tensor]]], +) -> None: + """Apply queued frees and reopen the caller's free-group scope.""" + if allocator.free_group is None or not any(pending_groups): + return + allocator.free_group_end() + allocator.free_group_begin() + + +def _full_tokens_before_mamba_recheck( + full_allocator: MultiEndedAllocator, + mamba_allocator: MultiEndedAllocator, + target_size: int, +) -> int: + """Conservative Full-token lower bound for the next Mamba capacity check. + + The current Mamba slot count can hide at most one slot minus one byte of + residual room. Subtract that possible residue so this estimate only skips + checks that cannot succeed from Full bytes alone. Allocator capacity remains + the stop condition after the bound is crossed. + """ + missing_slots = max(0, target_size - mamba_allocator.schedulable_available_size()) + if missing_slots == 0: + return 0 + mamba_page_bytes = mamba_allocator.entry_bytes_per_page + minimum_missing_bytes = (missing_slots - 1) * mamba_page_bytes + 1 + dcp_size = get_parallel().attn_dcp_size if full_allocator.shards_under_dcp else 1 + return -(-minimum_missing_bytes * dcp_size // full_allocator.entry_bytes) + + class MultiEndedAllocator(BaseTokenToKVPoolAllocator): """Allocator for one sub-pool over a `UnifiedKVPool`.""" @@ -1919,6 +1952,11 @@ class MultiEndedAllocator(BaseTokenToKVPoolAllocator): self._pending_reuse[latest_event] = (srcs_copy, src_pages_t) self._pending_reuse_pages_cpu.update(srcs_copy) + def flush_for_allocation(self) -> int: + """Public urgent flush used by peer allocation-pressure recovery.""" + with record_function("MultiEndedAlloc.flush_for_allocation"): + return self._flush(urgent=True) + def flush_opportunistic(self) -> int: """Public, non-urgent flush at quiescent points; never blocks `schedule_stream`. Fast-path the empty state: the scheduler triggers this diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 482df86df..f83cd9e79 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -577,11 +577,62 @@ class UnifiedRadixCache(BasePrefixCache): request_by_type = self._evict_request_by_type(params) available_size_targets = { - ct: self._component_available_size(ct) + request_cnt + ct: (ct, self._component_available_size(ct) + request_cnt) for ct, request_cnt in request_by_type.items() if request_cnt > 0 } - return self._evict(params, available_size_targets) + allocator = self.token_to_kv_pool_allocator + mamba_full_donor = allocator.mamba_full_cache_donor() + mamba_target = available_size_targets.get(ComponentType.MAMBA) + initial_params = params + if mamba_target is not None and mamba_full_donor is not None: + # Full KV can supply bytes but cannot recycle Mamba virtual IDs. + mamba_id_shortfall = max( + 0, + mamba_target[1] + - self.req_to_token_pool.mamba_allocator.available_size(), + ) + initial_params = EvictParams( + num_tokens=params.num_tokens, + swa_num_tokens=params.swa_num_tokens, + mamba_num=mamba_id_shortfall, + ) + result = self._evict(initial_params, available_size_targets) + + if mamba_target is not None and mamba_full_donor is not None: + mamba_full_donor.prepare_mamba_allocation(mamba_target[1]) + mamba_free_ids = self.req_to_token_pool.mamba_allocator.available_size() + mamba_capacity = self._component_available_size(ComponentType.MAMBA) + + if mamba_free_ids >= mamba_target[1] and mamba_capacity < mamba_target[1]: + full_evictable = self.full_evictable_size() + if full_evictable > 0: + donor_result = self._evict( + EvictParams(num_tokens=full_evictable), + {ComponentType.FULL: mamba_target}, + ) + result.num_tokens_evicted += donor_result.num_tokens_evicted + result.swa_num_tokens_evicted += donor_result.swa_num_tokens_evicted + result.mamba_num_evicted += donor_result.mamba_num_evicted + + # Preserve Mamba-victim recovery if Full cannot fund the target. + if ( + self._component_available_size(ComponentType.MAMBA) + < mamba_target[1] + ): + mamba_evictable = self.mamba_evictable_size() + if mamba_evictable > 0: + fallback_result = self._evict( + EvictParams(mamba_num=mamba_evictable), + {ComponentType.MAMBA: mamba_target}, + ) + result.num_tokens_evicted += fallback_result.num_tokens_evicted + result.swa_num_tokens_evicted += ( + fallback_result.swa_num_tokens_evicted + ) + result.mamba_num_evicted += fallback_result.mamba_num_evicted + + return result @staticmethod def _evict_request_by_type(params: EvictParams) -> dict[ComponentType, int]: @@ -611,7 +662,9 @@ class UnifiedRadixCache(BasePrefixCache): def _evict( self, params: EvictParams, - available_size_targets: Optional[dict[ComponentType, int]] = None, + available_size_targets: Optional[ + dict[ComponentType, tuple[ComponentType, int]] + ] = None, ) -> EvictResult: if self.disable: return EvictResult() @@ -712,22 +765,45 @@ class UnifiedRadixCache(BasePrefixCache): self, request_by_type: dict[ComponentType, int], tracker: dict[ComponentType, int], - available_size_targets: Optional[dict[ComponentType, int]] = None, + available_size_targets: Optional[ + dict[ComponentType, tuple[ComponentType, int]] + ] = None, ) -> None: # Buffer mode: eviction always wins over queued backup intents — a # destroyed victim's intent is stale-swept and the content rewrites # after its recompute. + last_mamba_donor_check = 0 + mamba_donor_prepared = False def target_reached(component_type: ComponentType) -> bool: + nonlocal last_mamba_donor_check, mamba_donor_prepared if available_size_targets is None: return False target = available_size_targets.get(component_type) - # Do not compact on every eviction step. Shared allocators include - # drainable peer holes here and flush the peer once in alloc(). - return ( - target is not None - and self._component_available_size(component_type) >= target - ) + if target is None: + return False + target_component, target_size = target + # A Full-leaf cascade can release Mamba or SWA state directly. + if self._component_available_size(target_component) >= target_size: + return True + if ( + component_type == ComponentType.FULL + and target_component == ComponentType.MAMBA + ): + donor = self.token_to_kv_pool_allocator.mamba_full_cache_donor() + assert donor is not None, "Mamba target requires a Full donor" + recheck_after = ( + 1 + if mamba_donor_prepared + else donor.full_tokens_before_mamba_recheck(target_size) + ) + if tracker[component_type] - last_mamba_donor_check < recheck_after: + return False + donor.prepare_mamba_allocation(target_size) + last_mamba_donor_check = tracker[component_type] + mamba_donor_prepared = True + # Schedulable capacity includes donor holes that allocation can compact. + return self._component_available_size(target_component) >= target_size for ct in self.tree_components: request_cnt = request_by_type[ct] @@ -737,16 +813,14 @@ class UnifiedRadixCache(BasePrefixCache): continue self.tree_core.evict_device_start(ct, request_cnt) try: - while not target_reached(ct): + while True: node_id, made_progress = self._evict_device_next_node(ct, tracker) if node_id is None: - if made_progress: - # Internal tombstone frees are now allocator-visible; - # recheck the allocation target before walking again. - continue - break - backup_kv = self._evict_device_leaf(node_id, tracker) - if backup_kv is not None: + if not made_progress: + break + else: + backup_kv = self._evict_device_leaf(node_id, tracker) + if node_id is not None and backup_kv is not None: # Deferred demote: run the D->H backup, demote only on success. written = self._execute_and_commit_kv_backup( backup_kv, write_back=True @@ -768,6 +842,8 @@ class UnifiedRadixCache(BasePrefixCache): "until host space frees", node_id, ) + if target_reached(ct): + break finally: self.tree_core.evict_device_end(ct) diff --git a/test/registered/unit/mem_cache/test_multi_ended_allocator.py b/test/registered/unit/mem_cache/test_multi_ended_allocator.py index b7cac8288..4c02d2d10 100644 --- a/test/registered/unit/mem_cache/test_multi_ended_allocator.py +++ b/test/registered/unit/mem_cache/test_multi_ended_allocator.py @@ -25,6 +25,7 @@ register_cpu_ci(est_time=11, suite="base-a-test-cpu") import contextlib import random import unittest +from unittest.mock import MagicMock import torch @@ -38,12 +39,16 @@ from sglang.srt.mem_cache.allocator.unified_sub_pool import ( FloatMultiEndedAllocator, MultiEndedAllocator, ) +from sglang.srt.mem_cache.base_prefix_cache import EvictParams +from sglang.srt.mem_cache.unified_cache.components import ComponentType from sglang.srt.mem_cache.unified_memory_pool import ( MambaSubPoolSpec, MHASubPoolSpec, MLASubPoolSpec, UnifiedKVPool, + UnifiedMambaSlotAllocator, ) +from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache from sglang.srt.runtime_context import get_parallel, publish, reset_context from sglang.srt.server_args import ServerArgs @@ -3040,7 +3045,7 @@ class TestDcpWidening(unittest.TestCase): self.assertTrue(bool((written[~owned] == 0).all())) self.assertTrue(bool((written[owned] > 0).all())) - def _build_composite(self, *, page_size): + def _build_composite(self, *, page_size, lazy_compaction=False): from sglang.srt.mem_cache.allocator.unified_mamba import ( UnifiedMambaTokenToKVPoolAllocator, ) @@ -3070,8 +3075,51 @@ class TestDcpWidening(unittest.TestCase): page_size=page_size, need_sort=False, forward_stream=None, + lazy_compaction=lazy_compaction, ) + def _build_donor_cache(self, allocator, mamba_slot_allocator, full_leaves): + 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.token_to_kv_pool_allocator = allocator + cache.req_to_token_pool = MagicMock(mamba_allocator=mamba_slot_allocator) + + tree_core = MagicMock() + full_evictable = sum(int(indices.numel()) for indices in full_leaves) + tree_core.full_evictable_size.return_value = full_evictable + tree_core.mamba_evictable_size.return_value = 0 + walk = {"request_cnt": 0, "freed_leaves": 0} + + def start(component_type, request_cnt): + self.assertEqual(component_type, ComponentType.FULL) + walk["request_cnt"] = request_cnt + + def next_node(component_type, tracker): + self.assertEqual(component_type, ComponentType.FULL) + if tracker[ComponentType.FULL] >= walk["request_cnt"] or walk[ + "freed_leaves" + ] >= len(full_leaves): + return None, False + return walk["freed_leaves"] + 1, True + + def evict_leaf(node_id, tracker): + self.assertEqual(node_id, walk["freed_leaves"] + 1) + indices = full_leaves[-node_id] + tracker[ComponentType.FULL] += int(indices.numel()) + allocator.free_segment(indices, start_pos=0) + walk["freed_leaves"] += 1 + return None + + tree_core.evict_device_start.side_effect = start + cache.tree_core = tree_core + cache._evict_device_next_node = MagicMock(side_effect=next_node) + cache._evict_device_leaf = MagicMock(side_effect=evict_leaf) + return cache, walk + def test_mamba_slot_cost_is_in_the_same_units_as_available_size(self): """The planner charges `mamba_slot_full_token_cost()` against a budget fed by `available_size()`. Both are bytes/entry_bytes conversions, so @@ -3099,6 +3147,174 @@ class TestDcpWidening(unittest.TestCase): # The un-scaled cost -- the bug -- would not have covered it. self.assertLess(base_cost * full_entry, mamba_bytes * dcp_size) + def test_mamba_donor_recheck_bound_is_aggregate_and_dcp_widened(self): + missing_slots = 7 + for dcp_size in (1, 2, 4): + with self.subTest(dcp_size=dcp_size), self._dcp(dcp_size): + allocator = self._build_composite(page_size=1) + allocator.mamba_allocator.schedulable_available_size = MagicMock( + return_value=3 + ) + + bound = allocator.full_tokens_before_mamba_recheck(3 + missing_slots) + + mamba_bytes = allocator.mamba_allocator.entry_bytes_per_page + full_bytes = allocator.full_attn_allocator.entry_bytes + minimum_missing_bytes = (missing_slots - 1) * mamba_bytes + 1 + expected = -(-minimum_missing_bytes * dcp_size // full_bytes) + self.assertEqual(bound, expected) + self.assertLessEqual( + bound, + missing_slots * allocator.mamba_slot_full_token_cost(), + ) + + def test_allocation_flush_public_wrapper_is_urgent(self): + with self._dcp(1): + allocator = self._build_composite(page_size=1) + full = allocator.full_attn_allocator + full._flush = MagicMock(return_value=3) + + self.assertEqual(full.flush_for_allocation(), 3) + full._flush.assert_called_once_with(urgent=True) + + def test_full_donor_flushes_paged_free_group_without_closing_it(self): + with self._dcp(2): + allocator = self._build_composite(page_size=2) + full_indices = allocator.alloc(8) + self.assertIsNotNone(full_indices) + allocated_before = allocator.full_attn_allocator.allocated_count() + + allocator.free_group_begin() + allocator.free_segment(full_indices, start_pos=0) + self.assertTrue(allocator.free_page_reps_group) + + donor = allocator.mamba_full_cache_donor() + self.assertIsNotNone(donor) + donor.flush_deferred_full_frees() + + self.assertEqual(allocator.free_group, []) + self.assertEqual(allocator.free_page_reps_group, []) + self.assertLess( + allocator.full_attn_allocator.allocated_count(), allocated_before + ) + self.assertEqual(allocator.verify_byte_accounting(), []) + allocator.free_group_end() + + def test_full_donor_reaches_real_capacity_across_supported_layouts(self): + cases = ( + # Baseline eager layout. + (1, 1, False, False), + # Paged lazy layout with grouped frees. + (1, 4, True, True), + # DCP widens Full ids while Mamba remains one slot per request. + (2, 1, True, True), + (4, 1, False, False), + ) + for dcp_size, page_size, lazy_compaction, grouped_free in cases: + with ( + self.subTest( + dcp_size=dcp_size, + page_size=page_size, + lazy_compaction=lazy_compaction, + grouped_free=grouped_free, + ), + self._dcp(dcp_size), + ): + allocator = self._build_composite( + page_size=page_size, + lazy_compaction=lazy_compaction, + ) + mamba_slots = UnifiedMambaSlotAllocator( + allocator.mamba_allocator, + max_size=allocator.size_mamba, + device=_DEV, + ) + + # Fill Full to its page-granular frontier, then consume any + # sub-Full-page byte residue with Mamba states. Mamba still + # owns unused virtual ids, but one more state has no backing + # bytes. + full_leaves = [] + while True: + indices = allocator.alloc(allocator.page_size) + if indices is None: + break + full_leaves.append(indices) + self.assertTrue(full_leaves) + residual_mamba = mamba_slots.schedulable_available_size() + if residual_mamba: + self.assertIsNotNone(mamba_slots.alloc(residual_mamba)) + self.assertGreater(mamba_slots.available_size(), 0) + self.assertEqual(mamba_slots.schedulable_available_size(), 0) + self.assertIsNone(mamba_slots.alloc(1)) + + gap_before = allocator.mamba_allocator._current_gap_bytes() + mamba_bytes = allocator.mamba_allocator.entry_bytes_per_page + full_page_bytes = allocator.full_attn_allocator.entry_bytes_per_page + expected_leaves = -(-(mamba_bytes - gap_before) // full_page_bytes) + self.assertGreater(expected_leaves, 0) + + cache, walk = self._build_donor_cache( + allocator, mamba_slots, full_leaves + ) + if grouped_free: + allocator.free_group_begin() + + result = cache.evict_for_alloc(EvictParams(mamba_num=1)) + + self.assertEqual(walk["freed_leaves"], expected_leaves) + self.assertEqual( + result.num_tokens_evicted, + expected_leaves * allocator.page_size, + ) + self.assertEqual(result.mamba_num_evicted, 0) + self.assertGreaterEqual(mamba_slots.schedulable_available_size(), 1) + allocated_mamba = mamba_slots.alloc(1) + self.assertIsNotNone(allocated_mamba) + self.assertEqual(allocator.verify_byte_accounting(), []) + if grouped_free: + self.assertEqual(allocator.free_group, []) + self.assertEqual(allocator.free_page_reps_group, []) + allocator.free_group_end() + + def test_full_donor_batches_urgent_compaction_for_large_shortfall(self): + with self._dcp(1): + allocator = self._build_composite( + page_size=1, + lazy_compaction=True, + ) + mamba_slots = UnifiedMambaSlotAllocator( + allocator.mamba_allocator, + max_size=allocator.size_mamba, + device=_DEV, + ) + + full_leaves = [] + while True: + indices = allocator.alloc(1) + if indices is None: + break + full_leaves.append(indices) + residual_mamba = mamba_slots.schedulable_available_size() + if residual_mamba: + self.assertIsNotNone(mamba_slots.alloc(residual_mamba)) + + cache, walk = self._build_donor_cache(allocator, mamba_slots, full_leaves) + full = allocator.full_attn_allocator + original_flush = full.flush_for_allocation + full.flush_for_allocation = MagicMock(wraps=original_flush) + allocator.free_group_begin() + + result = cache.evict_for_alloc(EvictParams(mamba_num=4)) + + self.assertGreater(walk["freed_leaves"], 1) + self.assertEqual(result.num_tokens_evicted, walk["freed_leaves"]) + self.assertEqual(full.flush_for_allocation.call_count, 1) + self.assertGreaterEqual(mamba_slots.schedulable_available_size(), 4) + self.assertIsNotNone(mamba_slots.alloc(4)) + self.assertEqual(allocator.verify_byte_accounting(), []) + allocator.free_group_end() + class TestFusedWriteLocTranslate(unittest.TestCase): """`write_loc_to_kernel_ids` must equal the arithmetic it stands for. diff --git a/test/registered/unit/mem_cache/test_unified_radix_allocation_eviction.py b/test/registered/unit/mem_cache/test_unified_radix_allocation_eviction.py index 9820849b6..68c8cf126 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_allocation_eviction.py +++ b/test/registered/unit/mem_cache/test_unified_radix_allocation_eviction.py @@ -3,6 +3,14 @@ import unittest from unittest.mock import MagicMock +import torch + +from sglang.srt.mem_cache.allocator.unified_hybrid_swa import ( + UnifiedMambaSWATokenToKVPoolAllocator, +) +from sglang.srt.mem_cache.allocator.unified_mamba import ( + UnifiedMambaTokenToKVPoolAllocator, +) 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 @@ -27,6 +35,7 @@ class TestUnifiedRadixAllocationEviction(CustomTestCase): capacity = {"available": 30} allocator = MagicMock() allocator.available_size.side_effect = lambda: capacity["available"] + allocator.mamba_full_cache_donor.return_value = None cache.token_to_kv_pool_allocator = allocator cache.req_to_token_pool = MagicMock() @@ -50,6 +59,49 @@ class TestUnifiedRadixAllocationEviction(CustomTestCase): cache._evict_device_leaf = MagicMock(side_effect=evict_leaf) return cache, capacity, leaf_count + @staticmethod + def _build_unified_mamba_donor_cache( + *, free_ids: int, byte_slots: int, tri_pool: bool = False + ): + cache = object.__new__(UnifiedRadixCache) + cache.disable = False + cache.tree_components = ( + (ComponentType.FULL, ComponentType.SWA, ComponentType.MAMBA) + if tri_pool + else (ComponentType.FULL, ComponentType.MAMBA) + ) + cache.is_swa_enabled = tri_pool + cache.cache_controller = None + cache.metrics_collector = None + cache.tree_core = MagicMock() + cache.tree_core.full_evictable_size.return_value = 16 + cache.tree_core.mamba_evictable_size.return_value = 8 + + capacity = {"free_ids": free_ids, "byte_slots": byte_slots} + allocator_cls = ( + UnifiedMambaSWATokenToKVPoolAllocator + if tri_pool + else UnifiedMambaTokenToKVPoolAllocator + ) + allocator = object.__new__(allocator_cls) + allocator.free_group = None + allocator.free_page_reps_group = None + if tri_pool: + allocator.full_free_group = [] + allocator.full_attn_allocator = MagicMock() + allocator.full_attn_allocator.schedulable_available_size.return_value = 100 + allocator.mamba_allocator = MagicMock() + allocator.mamba_allocator.available_size.side_effect = lambda: capacity[ + "free_ids" + ] + allocator.mamba_allocator.schedulable_available_size.side_effect = lambda: min( + capacity["free_ids"], capacity["byte_slots"] + ) + allocator.full_tokens_before_mamba_recheck = MagicMock(return_value=1) + cache.token_to_kv_pool_allocator = allocator + cache.req_to_token_pool = MagicMock(mamba_allocator=allocator.mamba_allocator) + return cache, capacity, allocator + def test_allocation_eviction_stops_when_shared_capacity_is_sufficient(self): cache, capacity, leaf_count = self._build_cache(collateral_capacity_gain=70) @@ -91,6 +143,7 @@ class TestUnifiedRadixAllocationEviction(CustomTestCase): cache.metrics_collector = None cache.tree_core = MagicMock() cache.token_to_kv_pool_allocator = MagicMock() + cache.token_to_kv_pool_allocator.mamba_full_cache_donor.return_value = None capacity = {"available": 0} mamba_allocator = MagicMock() @@ -117,6 +170,387 @@ class TestUnifiedRadixAllocationEviction(CustomTestCase): self.assertEqual(result.num_tokens_evicted, 20) self.assertEqual(result.mamba_num_evicted, 1) + def test_mamba_allocation_uses_full_as_donor_for_byte_shortfall(self): + cache, capacity, _ = self._build_unified_mamba_donor_cache( + free_ids=2, byte_slots=1 + ) + + def next_node(component_type, tracker): + if ( + component_type == ComponentType.FULL + and tracker[ComponentType.FULL] < 16 + ): + return 1, True + return None, False + + def evict_leaf(_node_id, tracker): + tracker[ComponentType.FULL] += 4 + capacity["byte_slots"] += 1 + 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=1)) + + self.assertEqual(capacity, {"free_ids": 2, "byte_slots": 2}) + self.assertEqual(result.num_tokens_evicted, 4) + self.assertEqual(result.mamba_num_evicted, 0) + self.assertEqual( + cache.tree_core.evict_device_start.call_args_list, + [ + unittest.mock.call(ComponentType.FULL, 16), + ], + ) + + def test_mamba_allocation_recycles_mamba_for_id_shortfall(self): + cache, capacity, _ = self._build_unified_mamba_donor_cache( + free_ids=0, byte_slots=1 + ) + + def next_node(component_type, tracker): + if ( + component_type == ComponentType.MAMBA + and tracker[ComponentType.MAMBA] < 1 + ): + return 1, True + return None, False + + def evict_leaf(_node_id, tracker): + tracker[ComponentType.MAMBA] += 1 + capacity["free_ids"] += 1 + capacity["byte_slots"] += 1 + 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=1)) + + self.assertEqual(capacity, {"free_ids": 1, "byte_slots": 2}) + self.assertEqual(result.num_tokens_evicted, 0) + self.assertEqual(result.mamba_num_evicted, 1) + cache.tree_core.evict_device_start.assert_called_once_with( + ComponentType.MAMBA, 1 + ) + + def test_mamba_allocation_does_not_evict_full_while_id_bound(self): + cache, capacity, _ = self._build_unified_mamba_donor_cache( + free_ids=0, byte_slots=10 + ) + + def next_node(component_type, tracker): + if component_type == ComponentType.FULL and tracker[ComponentType.FULL] < 4: + return 1, True + return None, False + + def evict_leaf(_node_id, tracker): + tracker[ComponentType.FULL] += 4 + capacity["byte_slots"] += 1 + 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=1)) + + self.assertEqual(capacity, {"free_ids": 0, "byte_slots": 10}) + self.assertEqual(result.num_tokens_evicted, 0) + self.assertEqual(result.mamba_num_evicted, 0) + cache.tree_core.evict_device_start.assert_called_once_with( + ComponentType.MAMBA, 1 + ) + + def test_mamba_allocation_does_not_evict_full_after_partial_id_recovery(self): + cache, capacity, _ = self._build_unified_mamba_donor_cache( + free_ids=0, byte_slots=10 + ) + + def next_node(component_type, tracker): + if ( + component_type == ComponentType.MAMBA + and tracker[ComponentType.MAMBA] < 1 + ): + return 1, True + if component_type == ComponentType.FULL and tracker[ComponentType.FULL] < 4: + return 2, True + return None, False + + def evict_leaf(node_id, tracker): + if node_id == 1: + tracker[ComponentType.MAMBA] += 1 + capacity["free_ids"] += 1 + capacity["byte_slots"] += 1 + else: + tracker[ComponentType.FULL] += 4 + capacity["byte_slots"] += 1 + 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=2)) + + self.assertEqual(capacity, {"free_ids": 1, "byte_slots": 11}) + self.assertEqual(result.num_tokens_evicted, 0) + self.assertEqual(result.mamba_num_evicted, 1) + cache.tree_core.evict_device_start.assert_called_once_with( + ComponentType.MAMBA, 2 + ) + + def test_mamba_allocation_splits_mixed_id_and_byte_pressure(self): + cache, capacity, _ = self._build_unified_mamba_donor_cache( + free_ids=3, byte_slots=2 + ) + + def next_node(component_type, tracker): + if ( + component_type == ComponentType.MAMBA + and tracker[ComponentType.MAMBA] < 1 + ): + return 1, True + if component_type == ComponentType.FULL and tracker[ComponentType.FULL] < 4: + return 2, True + return None, False + + def evict_leaf(node_id, tracker): + if node_id == 1: + tracker[ComponentType.MAMBA] += 1 + capacity["free_ids"] += 1 + capacity["byte_slots"] += 1 + else: + tracker[ComponentType.FULL] += 4 + capacity["byte_slots"] += 1 + 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=2)) + + self.assertEqual(capacity, {"free_ids": 4, "byte_slots": 4}) + self.assertEqual(result.num_tokens_evicted, 4) + self.assertEqual(result.mamba_num_evicted, 1) + self.assertEqual( + cache.tree_core.evict_device_start.call_args_list, + [ + unittest.mock.call(ComponentType.MAMBA, 1), + unittest.mock.call(ComponentType.FULL, 16), + ], + ) + + def test_full_donor_walk_continues_until_mamba_capacity_is_visible(self): + cache, capacity, _ = self._build_unified_mamba_donor_cache( + free_ids=2, byte_slots=1 + ) + leaf_count = 0 + + def next_node(component_type, tracker): + if ( + component_type == ComponentType.FULL + and tracker[ComponentType.FULL] < 16 + ): + return tracker[ComponentType.FULL] // 4 + 1, True + return None, False + + def evict_leaf(_node_id, tracker): + nonlocal leaf_count + leaf_count += 1 + tracker[ComponentType.FULL] += 4 + if leaf_count == 2: + capacity["byte_slots"] += 1 + 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=1)) + + self.assertEqual(leaf_count, 2) + self.assertEqual(result.num_tokens_evicted, 8) + cache.tree_core.evict_device_start.assert_called_once_with( + ComponentType.FULL, 16 + ) + + def test_full_donor_defers_preparation_until_safe_lower_bound(self): + cache, capacity, allocator = self._build_unified_mamba_donor_cache( + free_ids=2, byte_slots=1 + ) + leaf_count = 0 + + allocator.full_tokens_before_mamba_recheck.return_value = 8 + + def prepare(_target_size): + if leaf_count >= 2: + capacity["byte_slots"] = 2 + + allocator.prepare_mamba_allocation = MagicMock(side_effect=prepare) + cache._evict_device_next_node = MagicMock(return_value=(1, True)) + + def evict_leaf(_node_id, tracker): + nonlocal leaf_count + leaf_count += 1 + tracker[ComponentType.FULL] += 4 + return None + + cache._evict_device_leaf = MagicMock(side_effect=evict_leaf) + + result = cache.evict_for_alloc(EvictParams(mamba_num=1)) + + self.assertEqual(leaf_count, 2) + self.assertEqual(result.num_tokens_evicted, 8) + # One initial layout preparation plus one check at the lower bound. + self.assertEqual(allocator.prepare_mamba_allocation.call_count, 2) + + def test_full_donor_observes_cascade_before_lower_bound(self): + cache, capacity, allocator = self._build_unified_mamba_donor_cache( + free_ids=2, byte_slots=1 + ) + allocator.full_tokens_before_mamba_recheck.return_value = 100 + allocator.prepare_mamba_allocation = MagicMock() + cache._evict_device_next_node = MagicMock(return_value=(1, True)) + + def evict_leaf(_node_id, tracker): + tracker[ComponentType.FULL] += 4 + tracker[ComponentType.MAMBA] += 1 + capacity["byte_slots"] = 2 + return None + + cache._evict_device_leaf = MagicMock(side_effect=evict_leaf) + + result = cache.evict_for_alloc(EvictParams(mamba_num=1)) + + self.assertEqual(result.num_tokens_evicted, 4) + self.assertEqual(result.mamba_num_evicted, 1) + # Only the pre-donor preparation runs; the cheap capacity check stops + # before the Full-byte lower bound after the Mamba cascade is visible. + allocator.prepare_mamba_allocation.assert_called_once() + + def test_full_donor_rechecks_each_leaf_after_lower_bound(self): + cache, capacity, allocator = self._build_unified_mamba_donor_cache( + free_ids=2, byte_slots=1 + ) + leaf_count = 0 + + allocator.full_tokens_before_mamba_recheck.return_value = 8 + + def prepare(_target_size): + if leaf_count >= 3: + capacity["byte_slots"] = 2 + + allocator.prepare_mamba_allocation = MagicMock(side_effect=prepare) + cache._evict_device_next_node = MagicMock(return_value=(1, True)) + + def evict_leaf(_node_id, tracker): + nonlocal leaf_count + leaf_count += 1 + tracker[ComponentType.FULL] += 4 + return None + + cache._evict_device_leaf = MagicMock(side_effect=evict_leaf) + + result = cache.evict_for_alloc(EvictParams(mamba_num=1)) + + self.assertEqual(leaf_count, 3) + self.assertEqual(result.num_tokens_evicted, 12) + # The failed lower-bound check falls back to leaf-granular checks. + self.assertEqual(allocator.prepare_mamba_allocation.call_count, 3) + + def test_mamba_cache_is_last_resort_when_full_donor_is_exhausted(self): + cache, capacity, _ = self._build_unified_mamba_donor_cache( + free_ids=2, byte_slots=0 + ) + cache.tree_core.full_evictable_size.return_value = 4 + + def next_node(component_type, tracker): + if component_type == ComponentType.FULL and tracker[component_type] < 4: + return 1, True + if component_type == ComponentType.MAMBA and tracker[component_type] < 8: + return 2, True + return None, False + + def evict_leaf(node_id, tracker): + if node_id == 1: + tracker[ComponentType.FULL] += 4 + else: + tracker[ComponentType.MAMBA] += 1 + capacity["free_ids"] += 1 + capacity["byte_slots"] += 1 + 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=1)) + + self.assertEqual(capacity, {"free_ids": 3, "byte_slots": 1}) + self.assertEqual(result.num_tokens_evicted, 4) + self.assertEqual(result.mamba_num_evicted, 1) + self.assertEqual( + cache.tree_core.evict_device_start.call_args_list, + [ + unittest.mock.call(ComponentType.FULL, 4), + unittest.mock.call(ComponentType.MAMBA, 8), + ], + ) + + def test_full_donor_flushes_grouped_frees_before_capacity_recheck(self): + cache, capacity, allocator = self._build_unified_mamba_donor_cache( + free_ids=2, byte_slots=1 + ) + allocator.free_group_begin() + allocator.full_attn_allocator.free.side_effect = lambda _indices: ( + capacity.update(byte_slots=capacity["byte_slots"] + 1) + ) + allocator.mamba_allocator.alloc.side_effect = lambda need_size: ( + torch.arange(need_size) + if min(capacity["free_ids"], capacity["byte_slots"]) >= need_size + else None + ) + + cache._evict_device_next_node = MagicMock(return_value=(1, True)) + + def evict_leaf(_node_id, tracker): + tracker[ComponentType.FULL] += 4 + allocator.free(torch.tensor([1], dtype=torch.int64)) + return None + + cache._evict_device_leaf = MagicMock(side_effect=evict_leaf) + + result = cache.evict_for_alloc(EvictParams(mamba_num=1)) + + self.assertEqual(capacity["byte_slots"], 2) + self.assertEqual(result.num_tokens_evicted, 4) + self.assertEqual(allocator.free_group, []) + self.assertEqual(allocator.free_page_reps_group, []) + allocator.full_attn_allocator.free.assert_called_once() + self.assertIsNotNone(allocator.mamba_allocator.alloc(2)) + allocator.free_group_end() + + def test_tri_pool_uses_the_same_full_donor_capability(self): + cache, capacity, allocator = self._build_unified_mamba_donor_cache( + free_ids=2, byte_slots=1, tri_pool=True + ) + + cache._evict_device_next_node = MagicMock(return_value=(1, True)) + + def evict_leaf(_node_id, tracker): + tracker[ComponentType.FULL] += 4 + capacity["byte_slots"] += 1 + return None + + cache._evict_device_leaf = MagicMock(side_effect=evict_leaf) + + result = cache.evict_for_alloc(EvictParams(mamba_num=1)) + + self.assertIs(allocator.mamba_full_cache_donor(), allocator) + self.assertEqual(capacity["byte_slots"], 2) + self.assertEqual(result.num_tokens_evicted, 4) + self.assertEqual(result.mamba_num_evicted, 0) + cache.tree_core.evict_device_start.assert_called_once_with( + ComponentType.FULL, 16 + ) + def test_common_helper_uses_allocation_aware_entry_point(self): tree_cache = MagicMock() tree_cache.is_chunk_cache.return_value = False diff --git a/test/registered/unit/mem_cache/test_unified_tri_pool.py b/test/registered/unit/mem_cache/test_unified_tri_pool.py index 371c752bc..755d96848 100644 --- a/test/registered/unit/mem_cache/test_unified_tri_pool.py +++ b/test/registered/unit/mem_cache/test_unified_tri_pool.py @@ -20,6 +20,7 @@ Pure CPU; fakes stand in for the KV pools (data markers verify moves). import inspect import unittest +from unittest.mock import MagicMock import torch @@ -30,6 +31,8 @@ from sglang.srt.mem_cache.allocator.unified_sub_pool import ( FloatMultiEndedAllocator, MultiEndedAllocator, ) +from sglang.srt.mem_cache.base_prefix_cache import EvictParams +from sglang.srt.mem_cache.unified_cache.components import ComponentType from sglang.srt.mem_cache.unified_memory_pool import ( MambaSubPoolSpec, MHASubPoolSpec, @@ -37,6 +40,7 @@ from sglang.srt.mem_cache.unified_memory_pool import ( UnifiedMambaSlotAllocator, init_unified_mamba_swa_pools, ) +from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache from sglang.test.ci.ci_register import register_cpu_ci # Hermetic convention of this directory's pool tests: plain unittest.TestCase, @@ -113,6 +117,7 @@ class TestUnifiedTriPool(unittest.TestCase): n_full=32, n_swa=16, n_state=8, + page_size=1, lazy_compaction=False, ): full, swa, mamba = _tri_specs() @@ -126,6 +131,7 @@ class TestUnifiedTriPool(unittest.TestCase): sub_pool_specs=[full, swa, mamba], device=_DEV, enable_memory_saver=False, + page_size=page_size, ) kvcache = _FakeUnifiedSWAKVPool(pool) mamba_kv = _FakeKVCache(pool.max_slots("mamba")) @@ -136,6 +142,7 @@ class TestUnifiedTriPool(unittest.TestCase): device=_DEV, full_max_total_num_tokens=n_full, swa_max_total_num_tokens=n_swa, + page_size=page_size, need_sort=False, forward_stream=None, lazy_compaction=lazy_compaction, @@ -313,6 +320,90 @@ class TestUnifiedTriPool(unittest.TestCase): _relieve_for_alloc(allocator, 1) self.assertEqual(sa._hole_pages(), holes) # holes are assets, not backlog + def test_full_donor_stops_after_float_exposes_mamba_capacity(self): + for page_size, lazy_compaction in ((1, False), (4, True)): + with self.subTest(page_size=page_size, lazy_compaction=lazy_compaction): + _, allocator, _, _ = self._build( + page_size=page_size, lazy_compaction=lazy_compaction + ) + mamba_slots = UnifiedMambaSlotAllocator( + allocator.mamba_allocator, + max_size=allocator.mamba_allocator.max_slots - 1, + device=_DEV, + ) + + full_leaves = [] + while True: + indices = allocator.alloc(allocator.page_size) + if indices is None: + break + full_leaves.append(indices) + self.assertTrue(full_leaves) + residual_mamba = mamba_slots.schedulable_available_size() + if residual_mamba: + self.assertIsNotNone(mamba_slots.alloc(residual_mamba)) + while mamba_slots.alloc(1) is not None: + pass + self.assertGreater(mamba_slots.available_size(), 0) + self.assertEqual(mamba_slots.schedulable_available_size(), 0) + + cache = object.__new__(UnifiedRadixCache) + cache.disable = False + cache.tree_components = ( + ComponentType.FULL, + ComponentType.SWA, + ComponentType.MAMBA, + ) + cache.is_swa_enabled = True + cache.cache_controller = None + cache.metrics_collector = None + cache.token_to_kv_pool_allocator = allocator + cache.req_to_token_pool = MagicMock(mamba_allocator=mamba_slots) + + tree_core = MagicMock() + tree_core.full_evictable_size.return_value = len(full_leaves) + tree_core.mamba_evictable_size.return_value = 0 + walk = {"request_cnt": 0, "freed_leaves": 0} + + def start(component_type, request_cnt): + self.assertEqual(component_type, ComponentType.FULL) + walk["request_cnt"] = request_cnt + + def next_node(component_type, tracker): + self.assertEqual(component_type, ComponentType.FULL) + if tracker[ComponentType.FULL] >= walk["request_cnt"] or walk[ + "freed_leaves" + ] >= len(full_leaves): + return None, False + return walk["freed_leaves"] + 1, True + + def evict_leaf(node_id, tracker): + self.assertEqual(node_id, walk["freed_leaves"] + 1) + indices = full_leaves[-node_id] + tracker[ComponentType.FULL] += int(indices.numel()) + allocator.full_attn_allocator.free(indices) + walk["freed_leaves"] += 1 + return None + + tree_core.evict_device_start.side_effect = start + cache.tree_core = tree_core + cache._evict_device_next_node = MagicMock(side_effect=next_node) + cache._evict_device_leaf = MagicMock(side_effect=evict_leaf) + + swa_live_before = allocator.swa_attn_allocator._live_pages() + result = cache.evict_for_alloc(EvictParams(mamba_num=1)) + + self.assertEqual(walk["freed_leaves"], 1) + self.assertEqual(result.num_tokens_evicted, page_size) + self.assertEqual(result.swa_num_tokens_evicted, 0) + self.assertEqual(result.mamba_num_evicted, 0) + self.assertEqual( + allocator.swa_attn_allocator._live_pages(), swa_live_before + ) + self.assertGreaterEqual(mamba_slots.schedulable_available_size(), 1) + self.assertIsNotNone(mamba_slots.alloc(1)) + self.assertEqual(allocator.verify_byte_accounting(), []) + class TestTriPagedFreeGroup(unittest.TestCase): """The tri composite at PAGE SIZE > 1, driven through the production free @@ -370,6 +461,30 @@ class TestTriPagedFreeGroup(unittest.TestCase): # Capacity fully recovered: the float parked, both ends rewound. self.assertTrue(allocator.swa_attn_allocator._is_frontier_transparent()) + def test_mamba_donor_flushes_full_only_group_without_closing_it(self): + _, allocator = self._build_paged(page_size=1) + full_indices = allocator.alloc(8) + self.assertIsNotNone(full_indices) + allocator.free_swa(full_indices) + allocated_before = allocator.full_attn_allocator.allocated_count() + + allocator.free_group_begin() + allocator.free_full_segment(full_indices, start_pos=0) + self.assertTrue(allocator.full_free_group) + + donor = allocator.mamba_full_cache_donor() + self.assertIsNotNone(donor) + donor.flush_deferred_full_frees() + + self.assertEqual(allocator.free_group, []) + self.assertEqual(allocator.free_page_reps_group, []) + self.assertEqual(allocator.full_free_group, []) + self.assertLess( + allocator.full_attn_allocator.allocated_count(), allocated_before + ) + self.assertEqual(allocator.verify_byte_accounting(), []) + allocator.free_group_end() + def test_ungrouped_segment_free_also_reaches_the_float(self): pool, allocator = self._build_paged() v = allocator.alloc(8)