Radix Cache Split: Spin off TreeCore (#29901)
This commit is contained in:
Jialin Ouyang
@@ -1,11 +1,16 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import unittest
|
||||
from array import array
|
||||
from unittest import mock
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.srt.kv_canary.radix_cache_walker import walk_radix_cache_for_canary
|
||||
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
|
||||
from sglang.srt.mem_cache.radix_cache import RadixKey
|
||||
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache, TreeNode
|
||||
from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
|
||||
from sglang.srt.mem_cache.unified_cache_components import (
|
||||
BASE_COMPONENT_TYPE,
|
||||
ComponentType,
|
||||
@@ -129,17 +134,37 @@ class TestSelfUnitRadixWalker(CustomTestCase):
|
||||
)
|
||||
self.assertEqual(result.slot_indices.tolist(), [3, 4])
|
||||
|
||||
def test_unified_swa_sweep_gates_on_swa_lock_not_full_lock(self):
|
||||
"""With unlocked_only + swa_resident_only, the sweep filters on the SWA
|
||||
component lock: a FULL-locked node whose SWA lock was already released
|
||||
(early dec_swa_lock_only) must still be swept, and a node whose SWA
|
||||
lock is still held must not."""
|
||||
cache = self._make_unified_cache((ComponentType.FULL, ComponentType.SWA))
|
||||
self._add_unified_child(cache, [1, 2], lock_ref=1, swa_value=[1, 2])
|
||||
held = self._add_unified_child(cache, [3, 4], swa_value=[3, 4])
|
||||
held.component_data[ComponentType.SWA].lock_ref = 1
|
||||
|
||||
result = walk_radix_cache_for_canary(
|
||||
radix_cache=cache, unlocked_only=True, swa_resident_only=True
|
||||
)
|
||||
self.assertEqual(result.slot_indices.tolist(), [1, 2])
|
||||
|
||||
def _make_unified_cache(
|
||||
self, tree_components: tuple[ComponentType, ...]
|
||||
) -> UnifiedRadixCache:
|
||||
cache = UnifiedRadixCache.__new__(UnifiedRadixCache)
|
||||
cache.tree_components = tree_components
|
||||
cache.components = {ct: None for ct in tree_components}
|
||||
root = UnifiedTreeNode(tree_components)
|
||||
root.component_data[BASE_COMPONENT_TYPE].value = torch.tensor(
|
||||
[], dtype=torch.int32, device=self.device
|
||||
cache.is_swa_enabled = ComponentType.SWA in tree_components
|
||||
cache.tree_core = UnifiedTreeCore(
|
||||
CacheInitParams(
|
||||
disable=False,
|
||||
req_to_token_pool=None,
|
||||
token_to_kv_pool_allocator=None,
|
||||
page_size=1,
|
||||
),
|
||||
{ct: mock.MagicMock() for ct in tree_components},
|
||||
)
|
||||
cache.root_node = root
|
||||
return cache
|
||||
|
||||
def _add_unified_child(
|
||||
@@ -190,6 +215,23 @@ class TestSelfUnitRadixWalker(CustomTestCase):
|
||||
)
|
||||
self.assertEqual(result.slot_indices.tolist(), [3, 4])
|
||||
|
||||
def test_unified_walk_spans_device_evicted_nodes_without_emitting_them(self):
|
||||
"""Verify device-evicted (host-only) nodes emit no slots but still advance
|
||||
positions by their key length and pass the prev-slot chain through."""
|
||||
cache = self._make_unified_cache((ComponentType.FULL,))
|
||||
evicted = UnifiedTreeNode(cache.tree_components)
|
||||
evicted.parent = cache.root_node
|
||||
evicted.key = RadixKey(array("q", [7, 8]), None)
|
||||
cache.root_node.children[evicted.id] = evicted
|
||||
grandchild = self._add_unified_child(cache, [5, 6])
|
||||
cache.root_node.children.pop(grandchild.id)
|
||||
grandchild.parent = evicted
|
||||
evicted.children[grandchild.id] = grandchild
|
||||
result = walk_radix_cache_for_canary(radix_cache=cache)
|
||||
self.assertEqual(result.slot_indices.tolist(), [5, 6])
|
||||
self.assertEqual(result.positions.tolist(), [2, 3])
|
||||
self.assertEqual(result.prev_slot_indices.tolist(), [-1, 5])
|
||||
|
||||
def test_unified_swa_resident_only_noop_without_swa_component(self):
|
||||
"""Verify swa_resident_only is a no-op when SWA is not enabled."""
|
||||
cache = self._make_unified_cache((ComponentType.FULL,))
|
||||
|
||||
@@ -1,44 +1,37 @@
|
||||
"""CPU-only unit tests for the per-path Mamba checkpoint cap."""
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.ci.ci_register import register_cpu_ci, register_cuda_ci
|
||||
|
||||
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
|
||||
register_cuda_ci(est_time=5, stage="base-b", runner_config="1-gpu-small")
|
||||
|
||||
import argparse
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from collections import defaultdict
|
||||
from unittest import mock
|
||||
|
||||
import torch
|
||||
from test_unified_radix_cache_unittest import CacheConfig, UnifiedRadixCacheSuite
|
||||
|
||||
from sglang.srt.mem_cache.unified_cache.cache_action import MambaEvictExcessPathStates
|
||||
from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
|
||||
from sglang.srt.mem_cache.unified_cache_components.mamba_component import (
|
||||
MambaComponent,
|
||||
)
|
||||
from sglang.srt.mem_cache.unified_cache_components.tree_component import (
|
||||
ComponentType,
|
||||
)
|
||||
from sglang.srt.mem_cache.unified_radix_cache import (
|
||||
UnifiedLRUList,
|
||||
UnifiedRadixCache,
|
||||
UnifiedTreeNode,
|
||||
)
|
||||
from sglang.srt.mem_cache.unified_radix_cache import UnifiedLRUList, UnifiedTreeNode
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
|
||||
class _RecordingAllocator:
|
||||
def __init__(self):
|
||||
self.freed = []
|
||||
|
||||
def free(self, value):
|
||||
self.freed.extend(value.tolist())
|
||||
|
||||
|
||||
class _FakeUnifiedCache:
|
||||
class _FakeTreeCore:
|
||||
tree_components = (ComponentType.FULL, ComponentType.MAMBA)
|
||||
|
||||
def __init__(self):
|
||||
self.root_node = UnifiedTreeNode(self.tree_components)
|
||||
self.evictable_device_leaves = set()
|
||||
self.req_to_token_pool = SimpleNamespace(mamba_allocator=_RecordingAllocator())
|
||||
self.component_evictable_size_ = {ComponentType.MAMBA: 0}
|
||||
self.component_protected_size_ = {ComponentType.MAMBA: 0}
|
||||
self.lru_lists = {
|
||||
@@ -48,41 +41,47 @@ class _FakeUnifiedCache:
|
||||
}
|
||||
self.host_lru_lists = {
|
||||
ComponentType.MAMBA: UnifiedLRUList(
|
||||
ComponentType.MAMBA, self.tree_components
|
||||
ComponentType.MAMBA, self.tree_components, use_host_ptr=True
|
||||
)
|
||||
}
|
||||
self.evicted = []
|
||||
self.cascaded = []
|
||||
|
||||
def _evict_component_and_detach_lru(self, node, component, **kwargs):
|
||||
def _evict_component_and_detach_lru(self, node, component, *args, **kwargs):
|
||||
self.evicted.append(node)
|
||||
return UnifiedRadixCache._evict_component_and_detach_lru(
|
||||
self, node, component, **kwargs
|
||||
return UnifiedTreeCore._evict_component_and_detach_lru(
|
||||
self, node, component, *args, **kwargs
|
||||
)
|
||||
|
||||
def _cascade_evict(self, node, component, tracker):
|
||||
def _cascade_evict(self, node, component, tracker, device_frees, host_frees):
|
||||
self.cascaded.append(node)
|
||||
|
||||
|
||||
class _FakeUnifiedCache:
|
||||
tree_components = _FakeTreeCore.tree_components
|
||||
|
||||
|
||||
def _build_unified_chain(cap, length=3):
|
||||
cache = _FakeUnifiedCache()
|
||||
core = _FakeTreeCore()
|
||||
component = object.__new__(MambaComponent)
|
||||
component.cache = cache
|
||||
component.tree_core = core
|
||||
component.mamba_max_states_per_path = cap
|
||||
|
||||
nodes = []
|
||||
parent = cache.root_node
|
||||
parent = core.root_node
|
||||
for index in range(length):
|
||||
node = UnifiedTreeNode(cache.tree_components)
|
||||
node = UnifiedTreeNode(core.tree_components)
|
||||
node.parent = parent
|
||||
node.component_data[ComponentType.FULL].value = torch.tensor([100 + index])
|
||||
node.component_data[ComponentType.MAMBA].value = torch.tensor([index])
|
||||
parent.children[index] = node
|
||||
cache.component_evictable_size_[ComponentType.MAMBA] += 1
|
||||
cache.lru_lists[ComponentType.MAMBA].insert_mru(node)
|
||||
core.component_evictable_size_[ComponentType.MAMBA] += 1
|
||||
core.lru_lists[ComponentType.MAMBA].insert_mru(node)
|
||||
nodes.append(node)
|
||||
parent = node
|
||||
return component, nodes, cache
|
||||
return component, nodes, core, cache
|
||||
|
||||
|
||||
class TestMambaPathStateCap(unittest.TestCase):
|
||||
@@ -114,20 +113,22 @@ class TestMambaPathStateCap(unittest.TestCase):
|
||||
)
|
||||
|
||||
def test_unified_cache_removes_only_shallow_mamba_state(self):
|
||||
component, nodes, cache = _build_unified_chain(cap=2)
|
||||
component, nodes, core, cache = _build_unified_chain(cap=2)
|
||||
|
||||
component._evict_excess_path_states(nodes[-1])
|
||||
device_frees = defaultdict(list)
|
||||
host_frees = defaultdict(list)
|
||||
component._evict_excess_path_states(nodes[-1], device_frees, host_frees)
|
||||
|
||||
self.assertEqual(cache.evicted, [nodes[0]])
|
||||
self.assertEqual(cache.cascaded, [nodes[0]])
|
||||
self.assertEqual(core.evicted, [nodes[0]])
|
||||
self.assertEqual(core.cascaded, [nodes[0]])
|
||||
self.assertIsNone(nodes[0].component_data[ComponentType.MAMBA].value)
|
||||
self.assertIsNotNone(nodes[-1].component_data[ComponentType.MAMBA].value)
|
||||
self.assertEqual(
|
||||
cache.req_to_token_pool.mamba_allocator.freed,
|
||||
[v.item() for v in device_frees[ComponentType.MAMBA]],
|
||||
[0],
|
||||
)
|
||||
self.assertEqual(cache.component_evictable_size_[ComponentType.MAMBA], 2)
|
||||
self.assertFalse(cache.lru_lists[ComponentType.MAMBA].in_list(nodes[0]))
|
||||
self.assertEqual(core.component_evictable_size_[ComponentType.MAMBA], 2)
|
||||
self.assertFalse(core.lru_lists[ComponentType.MAMBA].in_list(nodes[0]))
|
||||
self.assertTrue(
|
||||
all(
|
||||
node.component_data[ComponentType.FULL].value is not None
|
||||
@@ -136,37 +137,44 @@ class TestMambaPathStateCap(unittest.TestCase):
|
||||
)
|
||||
|
||||
def test_unified_cache_cap_is_soft_for_fork_and_locked_nodes(self):
|
||||
component, nodes, cache = _build_unified_chain(cap=1, length=4)
|
||||
fork_child = UnifiedTreeNode(cache.tree_components)
|
||||
component, nodes, core, cache = _build_unified_chain(cap=1, length=4)
|
||||
fork_child = UnifiedTreeNode(core.tree_components)
|
||||
fork_child.parent = nodes[0]
|
||||
nodes[0].children["fork"] = fork_child
|
||||
nodes[1].component_data[ComponentType.MAMBA].lock_ref = 1
|
||||
|
||||
component._evict_excess_path_states(nodes[-1])
|
||||
device_frees = defaultdict(list)
|
||||
host_frees = defaultdict(list)
|
||||
component._evict_excess_path_states(nodes[-1], device_frees, host_frees)
|
||||
|
||||
self.assertEqual(cache.evicted, [nodes[2]])
|
||||
self.assertEqual(core.evicted, [nodes[2]])
|
||||
self.assertIsNotNone(nodes[0].component_data[ComponentType.MAMBA].value)
|
||||
self.assertIsNotNone(nodes[1].component_data[ComponentType.MAMBA].value)
|
||||
self.assertIsNone(nodes[2].component_data[ComponentType.MAMBA].value)
|
||||
self.assertIsNotNone(nodes[3].component_data[ComponentType.MAMBA].value)
|
||||
|
||||
def test_unified_cache_preserves_existing_host_backup(self):
|
||||
component, nodes, cache = _build_unified_chain(cap=2)
|
||||
component, nodes, core, cache = _build_unified_chain(cap=2)
|
||||
mamba_data = nodes[0].component_data[ComponentType.MAMBA]
|
||||
mamba_data.host_value = torch.tensor([10])
|
||||
|
||||
component._evict_excess_path_states(nodes[-1])
|
||||
device_frees = defaultdict(list)
|
||||
host_frees = defaultdict(list)
|
||||
component._evict_excess_path_states(nodes[-1], device_frees, host_frees)
|
||||
|
||||
self.assertIsNone(mamba_data.value)
|
||||
self.assertIsNotNone(mamba_data.host_value)
|
||||
self.assertTrue(cache.host_lru_lists[ComponentType.MAMBA].in_list(nodes[0]))
|
||||
self.assertTrue(core.host_lru_lists[ComponentType.MAMBA].in_list(nodes[0]))
|
||||
|
||||
def test_unified_cache_negative_one_disables_cap(self):
|
||||
component, nodes, cache = _build_unified_chain(cap=-1)
|
||||
component, nodes, core, cache = _build_unified_chain(cap=-1)
|
||||
|
||||
component._evict_excess_path_states(nodes[-1])
|
||||
device_frees = defaultdict(list)
|
||||
host_frees = defaultdict(list)
|
||||
component._evict_excess_path_states(nodes[-1], device_frees, host_frees)
|
||||
|
||||
self.assertEqual(cache.evicted, [])
|
||||
self.assertEqual(dict(device_frees), {})
|
||||
self.assertEqual(core.evicted, [])
|
||||
self.assertTrue(
|
||||
all(
|
||||
node.component_data[ComponentType.MAMBA].value is not None
|
||||
@@ -175,5 +183,162 @@ class TestMambaPathStateCap(unittest.TestCase):
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(torch.cuda.is_available(), "mamba pool fixtures need CUDA")
|
||||
class TestMambaPathCapWriteThroughOrdering(CustomTestCase):
|
||||
"""CI-active write-through/path-cap ordering regressions (the unified radix
|
||||
cache unittest module is temporarily gated off on trunk)."""
|
||||
|
||||
cfg = CacheConfig(components=(ComponentType.FULL, ComponentType.MAMBA))
|
||||
_rid = 0
|
||||
# Borrow the fixture helpers without inheriting the full gated suite.
|
||||
_make_req = UnifiedRadixCacheSuite._make_req
|
||||
_alloc = UnifiedRadixCacheSuite._alloc
|
||||
_insert = UnifiedRadixCacheSuite._insert
|
||||
_init_hicache = UnifiedRadixCacheSuite._init_hicache
|
||||
_build_hicache_fixture = UnifiedRadixCacheSuite._build_hicache_fixture
|
||||
|
||||
def test_write_through_backup_survives_mamba_path_cap(self):
|
||||
cache, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||
cache.write_through_threshold = 2
|
||||
cache.components[ComponentType.MAMBA].mamba_max_states_per_path = 1
|
||||
|
||||
# The first insert creates the ancestor with a mamba state (hit_count 1).
|
||||
self._insert(cache, allocator, req_to_token_pool, [1, 2])
|
||||
ancestor = next(iter(cache.root_node.children.values()))
|
||||
self.assertIsNotNone(ancestor.component_data[ComponentType.MAMBA].value)
|
||||
|
||||
# The extending insert crosses the ancestor's write-through threshold in
|
||||
# the same walk whose commit runs the path-cap eviction; the cap must
|
||||
# leave the pending-backup node's device state for the deferred BackupKV.
|
||||
self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4])
|
||||
cache.writing_check(write_back=True)
|
||||
|
||||
self.assertTrue(ancestor.backuped)
|
||||
self.assertIsNotNone(ancestor.component_data[ComponentType.MAMBA].host_value)
|
||||
|
||||
def test_write_through_backup_chain_survives_mamba_path_cap(self):
|
||||
"""A failed backup leaves an unbacked ancestor inside a later deferred
|
||||
backup chain; the cap walk must spare the whole chain, not just its tip."""
|
||||
cache, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||
cache.write_through_threshold = 3
|
||||
mamba_comp = cache.components[ComponentType.MAMBA]
|
||||
|
||||
self._insert(cache, allocator, req_to_token_pool, [1, 2])
|
||||
ancestor = next(iter(cache.root_node.children.values()))
|
||||
self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4])
|
||||
middle = next(iter(ancestor.children.values()))
|
||||
|
||||
# The ancestor crosses the threshold here; a host-exhaustion failure
|
||||
# leaves it unbacked with hit_count past the bar and its state intact.
|
||||
with mock.patch.object(cache, "_execute_kv_backup", return_value=None):
|
||||
self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4, 5, 6])
|
||||
self.assertFalse(ancestor.backuped)
|
||||
self.assertIsNotNone(ancestor.component_data[ComponentType.MAMBA].value)
|
||||
|
||||
# The middle node crosses next, so the deferred chain is
|
||||
# [ancestor, middle]; the extending insert adopts a new leaf state,
|
||||
# firing the now-enabled cap walk before the chain executes — it must
|
||||
# not evict either chain node's device state.
|
||||
mamba_comp.mamba_max_states_per_path = 1
|
||||
self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4, 5, 6, 7, 8])
|
||||
cache.writing_check(write_back=True)
|
||||
|
||||
self.assertTrue(ancestor.backuped)
|
||||
self.assertIsNotNone(ancestor.component_data[ComponentType.MAMBA].host_value)
|
||||
self.assertTrue(middle.backuped)
|
||||
self.assertIsNotNone(middle.component_data[ComponentType.MAMBA].host_value)
|
||||
|
||||
def test_backup_retry_after_mamba_cap_skips_tombstoned_state(self):
|
||||
"""A backup that fails before the cap and retries via the leaf action
|
||||
rebuilds its spec post-cap: KV backs up, the tombstoned mamba arm stays gone."""
|
||||
cache, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||
cache.write_through_threshold = 1
|
||||
mamba_comp = cache.components[ComponentType.MAMBA]
|
||||
|
||||
# A failed write-through leaves the ancestor unbacked with device state.
|
||||
with mock.patch.object(cache, "_execute_kv_backup", return_value=None):
|
||||
self._insert(cache, allocator, req_to_token_pool, [1, 2])
|
||||
ancestor = next(iter(cache.root_node.children.values()))
|
||||
self.assertFalse(ancestor.backuped)
|
||||
self.assertIsNotNone(ancestor.component_data[ComponentType.MAMBA].value)
|
||||
|
||||
# The walk backup fails again, the cap tombstones the unlocked
|
||||
# ancestor's mamba state, then the leaf-action retry succeeds.
|
||||
mamba_comp.mamba_max_states_per_path = 1
|
||||
real_backup = cache._execute_kv_backup
|
||||
attempts = []
|
||||
|
||||
def fail_once(*args, **kwargs):
|
||||
attempts.append(args)
|
||||
if len(attempts) == 1:
|
||||
return None
|
||||
return real_backup(*args, **kwargs)
|
||||
|
||||
with mock.patch.object(cache, "_execute_kv_backup", side_effect=fail_once):
|
||||
self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4])
|
||||
leaf = next(iter(ancestor.children.values()))
|
||||
cache.writing_check(write_back=True)
|
||||
|
||||
# Post-cap spec rebuild: no resurrection of the tombstoned mamba state.
|
||||
self.assertTrue(ancestor.backuped)
|
||||
ancestor_cd = ancestor.component_data[ComponentType.MAMBA]
|
||||
self.assertIsNone(ancestor_cd.value)
|
||||
self.assertIsNone(ancestor_cd.host_value)
|
||||
self.assertTrue(leaf.backuped)
|
||||
self.assertIsNotNone(leaf.component_data[ComponentType.MAMBA].host_value)
|
||||
cache.sanity_check()
|
||||
|
||||
def test_walk_backup_excludes_same_insert_restamped_mamba(self):
|
||||
"""The walked target's backup executes before commit hooks, so a mamba
|
||||
value re-stamped by the same insert stays out of the host backup."""
|
||||
cache, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||
mamba_comp = cache.components[ComponentType.MAMBA]
|
||||
|
||||
self._insert(cache, allocator, req_to_token_pool, [1, 2])
|
||||
ancestor = next(iter(cache.root_node.children.values()))
|
||||
self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4])
|
||||
|
||||
# The cap walk tombstones the ancestor's mamba state.
|
||||
mamba_comp.mamba_max_states_per_path = 1
|
||||
self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4, 5, 6])
|
||||
self.assertIsNone(ancestor.component_data[ComponentType.MAMBA].value)
|
||||
|
||||
# Re-inserting [1, 2] crosses the threshold and re-stamps the tombstone
|
||||
# in the same insert; the backup must not carry the fresh mamba state.
|
||||
cache.write_through_threshold = ancestor.hit_count + 1
|
||||
self._insert(cache, allocator, req_to_token_pool, [1, 2])
|
||||
cache.writing_check(write_back=True)
|
||||
|
||||
self.assertTrue(ancestor.backuped)
|
||||
ancestor_cd = ancestor.component_data[ComponentType.MAMBA]
|
||||
self.assertIsNotNone(ancestor_cd.value)
|
||||
self.assertIsNone(ancestor_cd.host_value)
|
||||
|
||||
def test_cap_walk_failure_still_drains_collected_frees(self):
|
||||
"""A cap walk that raises mid-eviction must still free the tombstoned
|
||||
slots it already collected (the pre-split inline frees could not leak)."""
|
||||
cache, allocator, req_to_token_pool = self._build_hicache_fixture()
|
||||
mamba_comp = cache.components[ComponentType.MAMBA]
|
||||
|
||||
self._insert(cache, allocator, req_to_token_pool, [1, 2])
|
||||
ancestor = next(iter(cache.root_node.children.values()))
|
||||
self._insert(cache, allocator, req_to_token_pool, [1, 2, 3, 4])
|
||||
leaf = next(iter(ancestor.children.values()))
|
||||
self.assertIsNotNone(ancestor.component_data[ComponentType.MAMBA].value)
|
||||
|
||||
mamba_comp.mamba_max_states_per_path = 1
|
||||
available = req_to_token_pool.mamba_allocator.available_size()
|
||||
with mock.patch.object(
|
||||
cache.tree_core, "_cascade_evict", side_effect=RuntimeError("boom")
|
||||
):
|
||||
with self.assertRaises(RuntimeError):
|
||||
mamba_comp.apply_component_action(MambaEvictExcessPathStates(leaf.id))
|
||||
|
||||
self.assertIsNone(ancestor.component_data[ComponentType.MAMBA].value)
|
||||
self.assertEqual(
|
||||
req_to_token_pool.mamba_allocator.available_size(), available + 1
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,189 @@
|
||||
"""Unit tests for the tree-core backend registry."""
|
||||
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
|
||||
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
|
||||
from sglang.srt.mem_cache.unified_cache.component_type import ComponentType
|
||||
from sglang.srt.mem_cache.unified_cache.tree_core_registry import (
|
||||
_TREE_CORE_REGISTRY,
|
||||
create_tree_core,
|
||||
register_tree_core_backend,
|
||||
registered_tree_core_backends,
|
||||
)
|
||||
from sglang.srt.mem_cache.unified_cache.unified_tree_core import UnifiedTreeCore
|
||||
from sglang.srt.mem_cache.unified_cache_components.tree_component import (
|
||||
EvictLayer,
|
||||
TreeComponent,
|
||||
)
|
||||
from sglang.srt.mem_cache.unified_radix_cache import UnifiedRadixCache
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def _cache_init_params(**kwargs) -> CacheInitParams:
|
||||
return CacheInitParams(
|
||||
disable=False,
|
||||
req_to_token_pool=None,
|
||||
token_to_kv_pool_allocator=None,
|
||||
page_size=2,
|
||||
tree_components=(ComponentType.FULL,),
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
|
||||
class _StubFullComponent(TreeComponent):
|
||||
component_type = ComponentType.FULL
|
||||
|
||||
def create_match_validator(self, match_device_only: bool = False):
|
||||
return lambda node: True
|
||||
|
||||
def redistribute_on_node_split(self, new_parent, child):
|
||||
return None
|
||||
|
||||
def evict_component(
|
||||
self, node, device_frees, host_frees, target: EvictLayer = EvictLayer.DEVICE
|
||||
) -> tuple[int, int]:
|
||||
return 0, 0
|
||||
|
||||
def acquire_component_lock(self, node, result):
|
||||
return result
|
||||
|
||||
def release_component_lock(self, node, params):
|
||||
return None
|
||||
|
||||
def _evict_device_start(self, request_cnt) -> None:
|
||||
pass
|
||||
|
||||
def _evict_device_next_node(self, tracker, device_frees, host_frees):
|
||||
return None
|
||||
|
||||
def _evict_device_end(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
class _StubMambaComponent(_StubFullComponent):
|
||||
component_type = ComponentType.MAMBA
|
||||
|
||||
|
||||
class TreeCoreRegistryTest(CustomTestCase):
|
||||
def setUp(self):
|
||||
self._registry_snapshot = dict(_TREE_CORE_REGISTRY)
|
||||
|
||||
def tearDown(self):
|
||||
_TREE_CORE_REGISTRY.clear()
|
||||
_TREE_CORE_REGISTRY.update(self._registry_snapshot)
|
||||
|
||||
def test_registry_contains_python(self):
|
||||
self.assertIn("python", registered_tree_core_backends())
|
||||
|
||||
def test_python_backend_builds_the_python_tree(self):
|
||||
component = mock.MagicMock()
|
||||
core = create_tree_core(
|
||||
name="python",
|
||||
params=_cache_init_params(),
|
||||
components={ComponentType.FULL: component},
|
||||
)
|
||||
self.assertIsInstance(core, UnifiedTreeCore)
|
||||
self.assertIs(component.tree_core, core)
|
||||
|
||||
def test_unknown_backend_raises_naming_the_known_backends(self):
|
||||
with self.assertRaisesRegex(ValueError, "not registered") as cm:
|
||||
create_tree_core(
|
||||
name="not_a_real_backend",
|
||||
params=_cache_init_params(),
|
||||
components={},
|
||||
)
|
||||
self.assertIn("'python'", str(cm.exception))
|
||||
|
||||
def test_register_rejects_empty_name(self):
|
||||
with self.assertRaises(ValueError):
|
||||
register_tree_core_backend(" ", mock.MagicMock())
|
||||
|
||||
def test_register_rejects_duplicate_name(self):
|
||||
with self.assertRaises(ValueError):
|
||||
register_tree_core_backend("python", mock.MagicMock())
|
||||
|
||||
def test_create_dispatches_to_a_registered_factory(self):
|
||||
core = mock.MagicMock()
|
||||
factory = mock.MagicMock(return_value=core)
|
||||
register_tree_core_backend("custom", factory)
|
||||
params = _cache_init_params()
|
||||
components = {ComponentType.FULL: mock.MagicMock()}
|
||||
result = create_tree_core(name="custom", params=params, components=components)
|
||||
factory.assert_called_once_with(params, components)
|
||||
self.assertIs(result, core)
|
||||
|
||||
|
||||
class UnifiedRadixCacheTreeCoreSelectionTest(CustomTestCase):
|
||||
def setUp(self):
|
||||
self._registry_snapshot = dict(_TREE_CORE_REGISTRY)
|
||||
|
||||
def tearDown(self):
|
||||
_TREE_CORE_REGISTRY.clear()
|
||||
_TREE_CORE_REGISTRY.update(self._registry_snapshot)
|
||||
|
||||
def _cache_params(
|
||||
self,
|
||||
tree_components=(ComponentType.FULL,),
|
||||
component_registry_override={ComponentType.FULL: _StubFullComponent},
|
||||
**kwargs,
|
||||
) -> CacheInitParams:
|
||||
return CacheInitParams(
|
||||
disable=True,
|
||||
req_to_token_pool=ReqToTokenPool(
|
||||
size=2,
|
||||
max_context_len=8,
|
||||
device="cpu",
|
||||
enable_memory_saver=False,
|
||||
),
|
||||
token_to_kv_pool_allocator=None,
|
||||
page_size=1,
|
||||
tree_components=tree_components,
|
||||
component_registry_override=component_registry_override,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def test_init_downgrades_is_eagle_when_mamba_is_enabled(self):
|
||||
params = self._cache_params(
|
||||
is_eagle=True,
|
||||
tree_components=(ComponentType.FULL, ComponentType.MAMBA),
|
||||
component_registry_override={
|
||||
ComponentType.FULL: _StubFullComponent,
|
||||
ComponentType.MAMBA: _StubMambaComponent,
|
||||
},
|
||||
)
|
||||
cache = UnifiedRadixCache(params)
|
||||
self.assertFalse(cache.tree_core.is_eagle)
|
||||
|
||||
def test_init_keeps_is_eagle_without_mamba(self):
|
||||
params = self._cache_params(is_eagle=True)
|
||||
cache = UnifiedRadixCache(params)
|
||||
self.assertTrue(cache.tree_core.is_eagle)
|
||||
|
||||
def test_default_backend_builds_the_python_tree_core(self):
|
||||
cache = UnifiedRadixCache(params=self._cache_params())
|
||||
self.assertIsInstance(cache.tree_core, UnifiedTreeCore)
|
||||
component = cache.components[ComponentType.FULL]
|
||||
self.assertIs(component.tree_core, cache.tree_core)
|
||||
|
||||
def test_env_var_routes_construction_to_the_selected_backend(self):
|
||||
"""SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND selects the registered factory
|
||||
the cache constructs its tree through."""
|
||||
core = mock.MagicMock()
|
||||
factory = mock.MagicMock(return_value=core)
|
||||
register_tree_core_backend("custom_env_backend", factory)
|
||||
with envs.SGLANG_UNIFIED_RADIX_TREE_CORE_BACKEND.override("custom_env_backend"):
|
||||
cache = UnifiedRadixCache(params=self._cache_params())
|
||||
factory.assert_called_once()
|
||||
self.assertIs(cache.tree_core, core)
|
||||
component = cache.components[ComponentType.FULL]
|
||||
self.assertIs(component.tree_core, core)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -566,7 +566,7 @@ def bench_lock_unlock(
|
||||
nodes = []
|
||||
for seq in env.seqs[: num_seqs // 2]:
|
||||
r = env.tree.match_prefix(MatchPrefixParams(key=RadixKey(array("q", seq))))
|
||||
if r.last_device_node != env.tree.root_node:
|
||||
if r.last_device_node != env.tree.root_node_handle():
|
||||
nodes.append(r.last_device_node)
|
||||
if not nodes:
|
||||
return BenchResult("lock_unlock", 0, 0, 0, [])
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user