fix(unified-memory): evict Full KV for Mamba byte shortfalls (#36713)
Co-authored-by: Yangmin Li <yangminl@nvidia.com> Co-authored-by: YAMY <74099316+YAMY1234@users.noreply.github.com>
This commit is contained in:
SuperSong
co-authored by
Yangmin Li
YAMY
Yangmin Li
YAMY
parent
6c1d0b1b29
commit
bede776c2a
@@ -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 ""
|
||||
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user