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
@@ -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