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

Refactor cache init logic (#13800)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-11-24 11:41:46 +08:00
committed by GitHub
parent 75222bfed9
commit b2f7b08c49
17 changed files with 307 additions and 419 deletions
Download Patch File Download Diff File Expand all files Collapse all files

49
test/manual/test_schedule_policy.py
View File

@@ -14,7 +14,7 @@ from sglang.test.test_utils import CustomTestCase
class TestSchedulePolicy(CustomTestCase):
def setUp(self):
self.tree_cache = RadixCache(None, None, False)
self.tree_cache = RadixCache.create_simulated()
def test_init_with_cache_aware_policy(self):
policy = SchedulePolicy(
@@ -47,10 +47,10 @@ class TestSchedulePolicy(CustomTestCase):
)
def test_init_with_disabled_cache(self):
disabled_tree_cache = RadixCache(None, None, disable=True, page_size=1)
tree_cache = RadixCache.create_simulated(disable=True)
policy = SchedulePolicy(
policy="lpm",
tree_cache=disabled_tree_cache,
tree_cache=tree_cache,
enable_hierarchical_cache=True,
enable_priority_scheduling=False,
schedule_low_priority_values_first=False,
@@ -58,7 +58,7 @@ class TestSchedulePolicy(CustomTestCase):
self.assertEqual(policy.policy, CacheAgnosticPolicy.FCFS)
def test_calc_priority_fcfs(self):
tree_cache = RadixCache(None, None, False)
tree_cache = RadixCache.create_simulated()
waiting_queue = [
Req(1, "a b", [1, 2], SamplingParams()),
Req(3, "a b c", [1, 2, 3], SamplingParams()),
@@ -79,16 +79,15 @@ class TestSchedulePolicy(CustomTestCase):
self.assertEqual(waiting_queue[2].rid, 2)
def test_calc_priority_priority_enabled_fcfs_scheduling(self):
tree_cache = RadixCache(None, None, False)
tree_cache = RadixCache.create_simulated()
r1 = Req(1, "a b", [1, 2], SamplingParams())
r2 = Req(3, "a b c", [1, 2, 3], SamplingParams())
r3 = Req(2, "a", [1], SamplingParams())
r1.priority, r1.time_stats.wait_queue_entry_time = 1, 1
r2.priority, r2.time_stats.wait_queue_entry_time = 0, 1
r3.priority, r3.time_stats.wait_queue_entry_time = 0, 0
waiting_queue = [
Req(1, "a b", [1, 2], SamplingParams()),
Req(3, "a b c", [1, 2, 3], SamplingParams()),
Req(2, "a", [1], SamplingParams()),
]
waiting_queue[0].priority, waiting_queue[0].queue_time_start = 1, 1
waiting_queue[1].priority, waiting_queue[1].queue_time_start = 0, 1
waiting_queue[2].priority, waiting_queue[2].queue_time_start = 0, 0
waiting_queue = [r1, r2, r3]
policy = SchedulePolicy(
policy="fcfs",
@@ -98,6 +97,7 @@ class TestSchedulePolicy(CustomTestCase):
schedule_low_priority_values_first=False,
)
policy.calc_priority(waiting_queue)
# Check if priority enabled fcfs ordering is applied.
self.assertEqual(waiting_queue[0].rid, 1)
self.assertEqual(waiting_queue[1].rid, 2)
@@ -106,16 +106,15 @@ class TestSchedulePolicy(CustomTestCase):
def test_calc_priority_priority_enabled_fcfs_scheduling_with_low_priority_values_first(
self,
):
tree_cache = RadixCache(None, None, False)
tree_cache = RadixCache.create_simulated()
r1 = Req(1, "a b", [1, 2], SamplingParams())
r2 = Req(3, "a b c", [1, 2, 3], SamplingParams())
r3 = Req(2, "a", [1], SamplingParams())
r1.priority, r1.time_stats.wait_queue_entry_time = -1, 1
r2.priority, r2.time_stats.wait_queue_entry_time = 0, 1
r3.priority, r3.time_stats.wait_queue_entry_time = 0, 0
waiting_queue = [
Req(1, "a b", [1, 2], SamplingParams()),
Req(3, "a b c", [1, 2, 3], SamplingParams()),
Req(2, "a", [1], SamplingParams()),
]
waiting_queue[0].priority, waiting_queue[0].queue_time_start = -1, 0
waiting_queue[1].priority, waiting_queue[1].queue_time_start = 0, 1
waiting_queue[2].priority, waiting_queue[2].queue_time_start = 0, 0
waiting_queue = [r1, r2, r3]
policy = SchedulePolicy(
policy="fcfs",
@@ -131,7 +130,7 @@ class TestSchedulePolicy(CustomTestCase):
self.assertEqual(waiting_queue[2].rid, 3)
def test_calc_priority_longest_output_first_scheduling(self):
tree_cache = RadixCache(None, None, False)
tree_cache = RadixCache.create_simulated()
waiting_queue = [
Req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1000)),
@@ -153,7 +152,7 @@ class TestSchedulePolicy(CustomTestCase):
self.assertEqual(waiting_queue[2].rid, 3)
def test_calc_priority_priority_enabled_longest_output_first_scheduling(self):
tree_cache = RadixCache(None, None, False)
tree_cache = RadixCache.create_simulated()
waiting_queue = [
Req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1), priority=1),
@@ -177,7 +176,7 @@ class TestSchedulePolicy(CustomTestCase):
def test_calc_priority_priority_enabled_longest_output_first_scheduling_with_low_priority_values_first(
self,
):
tree_cache = RadixCache(None, None, False)
tree_cache = RadixCache.create_simulated()
waiting_queue = [
Req(1, "a b", [1, 2], SamplingParams(max_new_tokens=1), priority=0),

6
test/srt/test_mamba_unittest.py
View File

@@ -6,6 +6,7 @@ import torch
from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, HybridReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixKey
@@ -186,13 +187,14 @@ class TestMamba(unittest.TestCase):
kvcache=pool,
need_sort=False,
)
# setup radix cache
tree = MambaRadixCache(
params = CacheInitParams(
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=allocator,
page_size=1,
disable=False,
)
# setup radix cache
tree = MambaRadixCache(params=params)
def make_dummy_req():
sampling_params = SamplingParams(

73
test/srt/test_radix_cache_unit.py
View File

@@ -240,11 +240,9 @@ class TestRadixCache(unittest.TestCase):
with self.subTest(
page_size=page_size, disable=disable, enable_events=enable_events
):
cache = RadixCache(
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
page_size=page_size,
cache = RadixCache.create_simulated(
disable=disable,
page_size=page_size,
enable_kv_cache_events=enable_events,
)
@@ -257,9 +255,7 @@ class TestRadixCache(unittest.TestCase):
def test_reset(self):
"""Test reset method."""
cache = RadixCache(
req_to_token_pool=None, token_to_kv_pool_allocator=None, page_size=1
)
cache = RadixCache.create_simulated()
# Insert some data
cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 20, 30], dtype=torch.int64))
@@ -275,12 +271,7 @@ class TestRadixCache(unittest.TestCase):
"""Test basic insert and match operations."""
for disable_cache in [False, True]:
with self.subTest(disable_cache=disable_cache):
cache = RadixCache(
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
page_size=1,
disable=disable_cache,
)
cache = RadixCache.create_simulated(disable=disable_cache)
key = RadixKey([1, 2, 3])
value = torch.tensor([10, 20, 30], dtype=torch.int64)
@@ -309,9 +300,7 @@ class TestRadixCache(unittest.TestCase):
def test_insert_with_none_value(self):
"""Test insert with None value (should use token_ids as list)."""
cache = RadixCache(
req_to_token_pool=None, token_to_kv_pool_allocator=None, page_size=1
)
cache = RadixCache.create_simulated()
key = RadixKey([1, 2, 3])
prefix_len = cache.insert(key, None)
@@ -322,9 +311,7 @@ class TestRadixCache(unittest.TestCase):
def test_total_size(self):
"""Test total_size calculation."""
cache = RadixCache(
req_to_token_pool=None, token_to_kv_pool_allocator=None, page_size=1
)
cache = RadixCache.create_simulated()
self.assertEqual(cache.total_size(), 0)
@@ -344,11 +331,8 @@ class TestRadixCache(unittest.TestCase):
for page_size, enable_events in test_cases:
with self.subTest(page_size=page_size, enable_events=enable_events):
cache = RadixCache(
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
page_size=page_size,
enable_kv_cache_events=enable_events,
cache = RadixCache.create_simulated(
page_size=page_size, enable_kv_cache_events=enable_events
)
# Insert data
@@ -374,11 +358,8 @@ class TestRadixCache(unittest.TestCase):
mock_allocator = unittest.mock.Mock()
mock_allocator.device = torch.device("cpu")
cache = RadixCache(
req_to_token_pool=None,
token_to_kv_pool_allocator=mock_allocator,
page_size=1,
enable_kv_cache_events=True,
cache = RadixCache.create_simulated(
mock_allocator=mock_allocator, enable_kv_cache_events=True
)
# Insert and then evict data
@@ -400,9 +381,7 @@ class TestRadixCache(unittest.TestCase):
def test_extra_key_isolation(self):
"""Test that keys with different extra_key values are isolated."""
cache = RadixCache(
req_to_token_pool=None, token_to_kv_pool_allocator=None, page_size=1
)
cache = RadixCache.create_simulated()
# Insert same token sequence with different extra keys
cache.insert(
@@ -442,9 +421,7 @@ class TestRadixCache(unittest.TestCase):
def test_lock_ref_operations(self):
"""Test lock reference counting operations."""
cache = RadixCache(
req_to_token_pool=None, token_to_kv_pool_allocator=None, page_size=1
)
cache = RadixCache.create_simulated()
# Insert sequence
cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 20, 30], dtype=torch.int64))
@@ -471,11 +448,7 @@ class TestRadixCache(unittest.TestCase):
mock_allocator = unittest.mock.Mock()
mock_allocator.device = torch.device("cpu")
cache = RadixCache(
req_to_token_pool=None,
token_to_kv_pool_allocator=mock_allocator,
page_size=1,
)
cache = RadixCache.create_simulated(mock_allocator=mock_allocator)
# Insert sequences
cache.insert(RadixKey([1, 2]), torch.tensor([10, 20], dtype=torch.int64))
@@ -500,11 +473,7 @@ class TestRadixCache(unittest.TestCase):
for page_size, sequence_length in test_cases:
with self.subTest(page_size=page_size, sequence_length=sequence_length):
cache = RadixCache(
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
page_size=page_size,
)
cache = RadixCache.create_simulated(page_size=page_size)
tokens = list(range(sequence_length))
cache.insert(RadixKey(tokens), torch.tensor(tokens, dtype=torch.int64))
@@ -518,9 +487,7 @@ class TestRadixCache(unittest.TestCase):
def test_pretty_print_basic(self):
"""Test pretty_print produces output."""
cache = RadixCache(
req_to_token_pool=None, token_to_kv_pool_allocator=None, page_size=1
)
cache = RadixCache.create_simulated()
cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 20, 30], dtype=torch.int64))
@@ -532,9 +499,7 @@ class TestRadixCache(unittest.TestCase):
def test_all_values_flatten(self):
"""Test all_values_flatten method."""
cache = RadixCache(
req_to_token_pool=None, token_to_kv_pool_allocator=None, page_size=1
)
cache = RadixCache.create_simulated()
cache.insert(RadixKey([1, 2]), torch.tensor([10, 20], dtype=torch.int64))
cache.insert(RadixKey([3, 4]), torch.tensor([30, 40], dtype=torch.int64))
@@ -549,11 +514,7 @@ class TestRadixCache(unittest.TestCase):
"""Advanced prefix matching: splits inside nodes and across pages."""
for page_size in [1, 2]:
with self.subTest(page_size=page_size):
cache = RadixCache(
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
page_size=page_size,
)
cache = RadixCache.create_simulated(page_size=page_size)
# Insert a long sequence that will be split later.
seq1 = [1, 2, 3, 4, 5, 6, 7, 8]

23
test/srt/test_swa_unittest.py
View File

@@ -3,6 +3,7 @@ import unittest
import torch
from sglang.srt.mem_cache.allocator import SWAKVPool, SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
@@ -110,11 +111,13 @@ class TestSWA(unittest.TestCase):
)
# setup radix cache
tree = SWARadixCache(
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=allocator,
params=CacheInitParams(
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=allocator,
disable=False,
page_size=1,
),
sliding_window_size=sliding_window_size,
page_size=1,
disable=False,
)
# test
@@ -241,12 +244,14 @@ class TestSWA(unittest.TestCase):
)
# setup radix cache
tree = SWARadixCache(
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=allocator,
params=CacheInitParams(
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=allocator,
page_size=1,
disable=False,
is_eagle=True,
),
sliding_window_size=sliding_window_size,
page_size=1,
disable=False,
is_eagle=True,
)
# test