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

[CI] Move existing unit tests into unit directory (#20631)

Browse Source
This commit is contained in:
Ke Bao
2026-03-15 23:25:18 +08:00
committed by GitHub
parent e2be31824f
commit c3483e8e97
29 changed files with 16 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/mem_cache/test_mamba_unittest.py
+462
View File
@@ -0,0 +1,462 @@
import unittest
import torch
from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
from sglang.srt.environ import envs
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
EvictParams,
InsertParams,
MatchPrefixParams,
)
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.common import available_and_evictable_str
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
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=9, suite="stage-b-test-small-1-gpu")
register_amd_ci(est_time=9, suite="stage-b-test-small-1-gpu-amd")
class TestMamba(unittest.TestCase):
@classmethod
def setUpClass(cls):
pass
@classmethod
def tearDownClass(cls):
pass
def test_hybrid_linear_kv_pool(self):
size = 16
head_num = 2
head_dim = 256
num_layers = 48
global_interval = 4
dtype = torch.bfloat16
device = get_device()
full_attention_layer_ids = [
i for i in range(global_interval - 1, num_layers, global_interval)
]
pool = HybridLinearKVPool(
size=size,
dtype=dtype,
page_size=1,
head_num=head_num,
head_dim=head_dim,
full_attention_layer_ids=full_attention_layer_ids,
enable_kvcache_transpose=False,
device=device,
enable_memory_saver=False,
mamba_pool=None,
)
assert pool._transfer_full_attention_id(global_interval - 1) == 0
assert pool._transfer_full_attention_id(2 * global_interval - 1) == 1
with self.assertRaises(ValueError) as context:
pool._transfer_full_attention_id(1)
self.assertIn(
"layer_id=1 not in full attention layers:", str(context.exception)
)
def test_mamba_pool(self):
max_num_reqs = 10
mamba_cache_size = 20
max_context_len = 128
device = get_device()
global_interval = 4
num_layers = 48
full_attention_layer_ids = [
i for i in range(global_interval - 1, num_layers, global_interval)
]
mamba_layers = [
i for i in range(num_layers) if i not in full_attention_layer_ids
]
shape = Mamba2StateShape.create(
tp_world_size=1,
intermediate_size=4096,
n_groups=16,
num_heads=32,
head_dim=128,
state_size=128,
conv_kernel=4,
)
with envs.SGLANG_MAMBA_SSM_DTYPE.override("bfloat16"):
mamba2_cache_params = Mamba2CacheParams(shape=shape, layers=mamba_layers)
req_to_token_pool = HybridReqToTokenPool(
size=max_num_reqs,
mamba_size=mamba_cache_size,
mamba_spec_state_size=max_num_reqs,
max_context_len=max_context_len,
device=device,
enable_memory_saver=False,
cache_params=mamba2_cache_params,
enable_mamba_extra_buffer=False,
speculative_num_draft_tokens=3,
)
assert req_to_token_pool.available_size() == max_num_reqs
assert req_to_token_pool.mamba_pool.available_size() == mamba_cache_size
sampling_params = SamplingParams(
temperature=0,
max_new_tokens=1,
)
req = Req(
rid=0,
origin_input_text="",
origin_input_ids=[],
sampling_params=sampling_params,
)
# alloc req
req_to_token_pool.alloc([req])
assert req_to_token_pool.available_size() == max_num_reqs - 1
assert req_to_token_pool.mamba_pool.available_size() == mamba_cache_size - 1
# free req
req_to_token_pool.free_mamba_cache(req)
req_to_token_pool.free(req)
assert req_to_token_pool.available_size() == max_num_reqs
assert req_to_token_pool.mamba_pool.available_size() == mamba_cache_size
# alloc req without free mamba cache
req.mamba_pool_idx = None
req_to_token_pool.alloc([req])
req_to_token_pool.free(req)
assert req_to_token_pool.available_size() == max_num_reqs
assert req_to_token_pool.mamba_pool.available_size() == mamba_cache_size - 1
# alloc again
req_to_token_pool.alloc([req])
assert req_to_token_pool.available_size() == max_num_reqs - 1
assert req_to_token_pool.mamba_pool.available_size() == mamba_cache_size - 1
def test_mamba_radix_cache_1(self):
tree, allocator, req_to_token_pool, make_dummy_req = (
self._setup_tree_and_allocator()
)
mamba_pool = req_to_token_pool.mamba_pool
# test
print(
f"[Start] allocator mamba available size: {mamba_pool.available_size()}, full available size: {allocator.available_size()}"
)
req1 = make_dummy_req()
req1_token_ids, req1_kv_indices = [1, 2, 3], allocator.alloc(3)
assert len(req1_token_ids) == len(req1_kv_indices)
print(
f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}"
)
result = tree.insert(
InsertParams(
key=RadixKey(req1_token_ids),
value=req1_kv_indices,
mamba_value=req1.mamba_pool_idx.unsqueeze(0),
)
)
prefix_len = result.prefix_len
print(
f"req1: prefix_len: {prefix_len}, allocator mamba available size: {mamba_pool.available_size()}, full available size: {allocator.available_size()}"
)
req2 = make_dummy_req()
req2_token_ids, req2_kv_indices = [1, 2, 3, 4, 5, 6, 7], allocator.alloc(7)
assert len(req2_token_ids) == len(req2_kv_indices)
print(
f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}"
)
result = tree.insert(
InsertParams(
key=RadixKey(req2_token_ids),
value=req2_kv_indices,
mamba_value=req2.mamba_pool_idx.unsqueeze(0),
)
)
prefix_len = result.prefix_len
print(
f"req2: prefix_len: {prefix_len}, allocator mamba available size: {mamba_pool.available_size()}, full available size: {allocator.available_size()}"
)
req3 = make_dummy_req()
req3_token_ids, req3_kv_indices = [10, 11, 12], allocator.alloc(3)
assert len(req3_token_ids) == len(req3_kv_indices)
print(
f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}"
)
result = tree.insert(
InsertParams(
key=RadixKey(req3_token_ids),
value=req3_kv_indices,
mamba_value=req3.mamba_pool_idx.unsqueeze(0),
)
)
prefix_len = result.prefix_len
print(
f"req3: prefix_len: {prefix_len}, allocator mamba available size: {mamba_pool.available_size()}, full available size: {allocator.available_size()}"
)
req4 = make_dummy_req()
req4_token_ids, req4_kv_indices = [1, 2, 3, 4, 5, 60, 70], allocator.alloc(7)
assert len(req4_token_ids) == len(req4_kv_indices)
print(
f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}"
)
result = tree.insert(
InsertParams(
key=RadixKey(req4_token_ids),
value=req4_kv_indices,
mamba_value=req4.mamba_pool_idx.unsqueeze(0),
)
)
prefix_len = result.prefix_len
print(
f"req4: prefix_len: {prefix_len}, allocator mamba available size: {mamba_pool.available_size()}, full available size: {allocator.available_size()}"
)
tree.pretty_print()
full_num_tokens = 1
print(f"evicting {full_num_tokens} full token")
result = tree.evict(EvictParams(num_tokens=full_num_tokens))
assert (
result.num_tokens_evicted >= full_num_tokens
), f"evicted {result.num_tokens_evicted} full tokens, expected {full_num_tokens}"
tree.pretty_print()
mamba_num = 1
print(f"evicting {mamba_num} mamba")
result = tree.evict(EvictParams(num_tokens=0, mamba_num=mamba_num))
assert (
result.mamba_num_evicted >= mamba_num
), f"evicted {result.mamba_num_evicted} mamba states, expected {mamba_num}"
tree.pretty_print()
req5_token_ids = [1, 2, 3, 4, 5]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req5_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
)
assert len(kv_indices) == 0
req6_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req6_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
)
assert len(kv_indices) == 7
assert len(last_node.key) == 2
req7_token_ids = [1, 2, 3, 4, 5, 6, 7]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req7_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req7: token_ids: {req7_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
)
assert len(kv_indices) == 7
assert len(last_node.key) == 2
mamba_num = 1
print(f"evicting {mamba_num} mamba")
result = tree.evict(EvictParams(num_tokens=0, mamba_num=mamba_num))
assert (
result.mamba_num_evicted >= mamba_num
), f"evicted {result.mamba_num_evicted} mamba states, expected {mamba_num}"
tree.pretty_print()
req8_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req8_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req8: token_ids: {req8_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
)
assert len(kv_indices) == 0
assert len(last_node.key) == 0
req9_token_ids = [1, 2, 3, 4, 5, 6, 7]
req9 = make_dummy_req()
result = tree.match_prefix(
MatchPrefixParams(key=RadixKey(req9_token_ids), req=req9, cow_mamba=True)
)
kv_indices, last_node = result.device_indices, result.last_device_node
assert req9.mamba_pool_idx is not None
assert torch.all(
mamba_pool.mamba_cache.conv[0][:, req9.mamba_pool_idx]
== mamba_pool.mamba_cache.conv[0][:, last_node.mamba_value]
)
assert torch.all(
mamba_pool.mamba_cache.temporal[:, req9.mamba_pool_idx]
== mamba_pool.mamba_cache.temporal[:, last_node.mamba_value]
)
print(tree.available_and_evictable_str())
print(available_and_evictable_str(tree))
tree.sanity_check()
def _setup_tree_and_allocator(self):
"""Helper to create a MambaRadixCache with allocator for testing."""
set_global_server_args_for_scheduler(
ServerArgs(model_path="dummy", page_size=1)
)
size = 128
dtype = torch.bfloat16
head_num = 2
head_dim = 256
num_layers = 48
global_interval = 4
max_num_reqs = 10
mamba_cache_size = 20
max_context_len = 128
device = get_device()
full_attention_layer_ids = [
i for i in range(global_interval - 1, num_layers, global_interval)
]
mamba_layers = [
i for i in range(num_layers) if i not in full_attention_layer_ids
]
with envs.SGLANG_MAMBA_SSM_DTYPE.override("bfloat16"):
shape = Mamba2StateShape.create(
tp_world_size=1,
intermediate_size=4096,
n_groups=16,
num_heads=32,
head_dim=128,
state_size=128,
conv_kernel=4,
)
mamba2_cache_params = Mamba2CacheParams(shape=shape, layers=mamba_layers)
req_to_token_pool = HybridReqToTokenPool(
size=max_num_reqs,
mamba_size=mamba_cache_size,
mamba_spec_state_size=max_num_reqs,
max_context_len=max_context_len,
device=device,
enable_memory_saver=False,
cache_params=mamba2_cache_params,
enable_mamba_extra_buffer=False,
speculative_num_draft_tokens=3,
)
pool = HybridLinearKVPool(
size=size,
dtype=dtype,
page_size=1,
head_num=head_num,
head_dim=head_dim,
full_attention_layer_ids=full_attention_layer_ids,
enable_kvcache_transpose=False,
device=device,
enable_memory_saver=False,
mamba_pool=req_to_token_pool.mamba_pool,
)
allocator = TokenToKVPoolAllocator(
size=size,
dtype=dtype,
device=device,
kvcache=pool,
need_sort=False,
)
params = CacheInitParams(
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=allocator,
page_size=1,
disable=False,
)
tree = MambaRadixCache(params=params)
def make_dummy_req():
sampling_params = SamplingParams(
temperature=0,
max_new_tokens=1,
)
req = Req(
rid=0,
origin_input_text="",
origin_input_ids=[],
sampling_params=sampling_params,
)
req_to_token_pool.alloc([req])
return req
return tree, allocator, req_to_token_pool, make_dummy_req
def test_insert_prev_prefix_len(self):
"""Test that prev_prefix_len correctly controls which KV indices are freed
during insert, covering: full free, partial free across multi-node, and no free.
"""
tree, allocator, req_to_token_pool, make_dummy_req = (
self._setup_tree_and_allocator()
)
initial_avail = allocator.available_size()
# Step 1: Insert [1,2,3] to create first node
req1 = make_dummy_req()
tree.insert(
InsertParams(
key=RadixKey([1, 2, 3]),
value=allocator.alloc(3),
mamba_value=req1.mamba_pool_idx.unsqueeze(0),
)
)
assert allocator.available_size() == initial_avail - 3
# Step 2: Insert [1,2,3,4,5,6,7] with prev_prefix_len=0 (free all matched)
# Creates tree: [1,2,3] -> [4,5,6,7]
req2 = make_dummy_req()
result = tree.insert(
InsertParams(
key=RadixKey([1, 2, 3, 4, 5, 6, 7]),
value=allocator.alloc(7),
mamba_value=req2.mamba_pool_idx.unsqueeze(0),
prev_prefix_len=0,
)
)
assert result.prefix_len == 3
# alloc 7, freed 3 (dup prefix [0..2]), stored 4 in new node => net -4
assert allocator.available_size() == initial_avail - 3 - 4
avail_after_step2 = allocator.available_size()
# Step 3: Insert [1,2,3,4,5,6,7,8] with prev_prefix_len=2
# Matched prefix = 7 (across two nodes: [1,2,3] len=3, [4,5,6,7] len=4)
# Protected [0..1], freed [2..6] = 5 slots, new [7] = 1 slot stored
req3 = make_dummy_req()
result = tree.insert(
InsertParams(
key=RadixKey([1, 2, 3, 4, 5, 6, 7, 8]),
value=allocator.alloc(8),
mamba_value=req3.mamba_pool_idx.unsqueeze(0),
prev_prefix_len=2,
)
)
assert result.prefix_len == 7
# alloc 8, freed 5, stored 1 => net -3
assert allocator.available_size() == avail_after_step2 - 3
avail_after_step3 = allocator.available_size()
# Step 4: Insert [1,2,3,4,5,6,7,8,9] with prev_prefix_len=8 (covers all matched)
# Matched prefix = 8, prev_prefix_len=8 => nothing freed
req4 = make_dummy_req()
result = tree.insert(
InsertParams(
key=RadixKey([1, 2, 3, 4, 5, 6, 7, 8, 9]),
value=allocator.alloc(9),
mamba_value=req4.mamba_pool_idx.unsqueeze(0),
prev_prefix_len=8,
)
)
assert result.prefix_len == 8
# alloc 9, freed 0, stored 1 => net -9
assert allocator.available_size() == avail_after_step3 - 9
tree.sanity_check()
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_nsa_pool_host_unit.py
+131
View File
@@ -0,0 +1,131 @@
import unittest
import torch
from sglang.srt.mem_cache.memory_pool import NSATokenToKVPool
from sglang.srt.mem_cache.memory_pool_host import (
ALLOC_MEMORY_FUNCS,
NSATokenToKVPoolHost,
alloc_with_pin_memory,
)
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=3, suite="stage-b-test-small-1-gpu")
class TestNSAHiCacheTransfer(unittest.TestCase):
def setUp(self):
if not torch.cuda.is_available():
self.skipTest("CUDA is required for NSA host transfer tests.")
if is_npu() or is_xpu():
self.skipTest("NSA host transfer tests only support CUDA/ROCm.")
if not (is_cuda() or is_hip()):
self.skipTest("CUDA/ROCm not available.")
@staticmethod
def _token_indices_for_pages(pages: torch.Tensor, page_size: int, device: str):
parts = [
torch.arange(
int(page_id) * page_size,
(int(page_id) + 1) * page_size,
device=device,
dtype=torch.int64,
)
for page_id in pages.tolist()
]
return torch.cat(parts, dim=0)
def _run_device_to_host_indexer_copy(self, io_backend: str):
page_size = 1 if is_hip() else 64
layer_num = 2
size = page_size * 4
device_pool = NSATokenToKVPool(
size=size,
page_size=page_size,
kv_lora_rank=128,
dtype=torch.bfloat16,
qk_rope_head_dim=32,
layer_num=layer_num,
device="cuda",
enable_memory_saver=False,
kv_cache_dim=576,
index_head_dim=128,
)
pin_memory = io_backend == "kernel"
original_alloc = ALLOC_MEMORY_FUNCS["cuda"]
if pin_memory:
ALLOC_MEMORY_FUNCS["cuda"] = alloc_with_pin_memory
try:
host_pool = NSATokenToKVPoolHost(
device_pool=device_pool,
host_to_device_ratio=2.0,
host_size=0,
page_size=page_size,
layout="layer_first",
pin_memory=pin_memory,
device="cpu",
)
finally:
ALLOC_MEMORY_FUNCS["cuda"] = original_alloc
for layer_id in range(layer_num):
buf = device_pool.index_k_with_scale_buffer[layer_id]
data = torch.arange(
buf.numel(), device=buf.device, dtype=torch.uint8
).view_as(buf)
buf.copy_((data + layer_id) % 256)
kv_buf = device_pool.kv_buffer[layer_id]
kv_data = torch.arange(
kv_buf.numel(), device=kv_buf.device, dtype=kv_buf.dtype
).view_as(kv_buf)
kv_buf.copy_(kv_data + layer_id)
device_pages = torch.tensor([1, 2, 3], device="cuda", dtype=torch.int64)
host_pages = torch.tensor(
[0, 1, 2],
device="cuda" if io_backend == "kernel" else "cpu",
dtype=torch.int64,
)
device_indices = self._token_indices_for_pages(
device_pages, page_size, device="cuda"
)
host_indices = self._token_indices_for_pages(
host_pages,
page_size,
device="cuda" if io_backend == "kernel" else "cpu",
)
host_pool.backup_from_device_all_layer(
device_pool, host_indices, device_indices, io_backend
)
for layer_id in range(layer_num):
for host_page, device_page in zip(
host_pages.tolist(), device_pages.tolist()
):
got = host_pool.index_k_with_scale_buffer[layer_id][host_page].cpu()
expected = device_pool.index_k_with_scale_buffer[layer_id][
device_page
].cpu()
self.assertTrue(torch.equal(got, expected))
host_start = host_page * page_size
device_start = device_page * page_size
got_kv = host_pool.kv_buffer[layer_id][
host_start : host_start + page_size
].cpu()
expected_kv = device_pool.kv_buffer[layer_id][
device_start : device_start + page_size
].cpu()
self.assertTrue(torch.equal(got_kv, expected_kv))
def test_device_to_host_indexer_kernel(self):
self._run_device_to_host_indexer_copy(io_backend="kernel")
def test_device_to_host_indexer_direct(self):
self._run_device_to_host_indexer_copy(io_backend="direct")
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_radix_cache_slru_accuracy.py
+144
View File
@@ -0,0 +1,144 @@
import unittest
import torch
from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import (
EvictParams,
InsertParams,
MatchPrefixParams,
)
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey
from sglang.test.ci.ci_register import register_cuda_ci
register_cuda_ci(est_time=5, suite="stage-b-test-small-1-gpu")
class TestSLRUAccuracy(unittest.TestCase):
def setUp(self):
"""Setup minimal memory pools for testing"""
torch.set_default_device(None)
device = "cpu"
dtype = torch.float16
# Create smaller KV cache to ensure evictions occur
self.kv_cache = MHATokenToKVPool(
size=8, # Very small size to trigger evictions quickly
page_size=1,
dtype=dtype,
head_num=8,
head_dim=64,
layer_num=1,
device=device,
enable_memory_saver=False,
)
# Create token-to-KV pool allocator
self.token_to_kv_pool = TokenToKVPoolAllocator(
size=8, dtype=dtype, device=device, kvcache=self.kv_cache, need_sort=False
)
# Create req-to-token pool
self.req_to_token_pool = ReqToTokenPool(
size=8, max_context_len=1024, device=device, enable_memory_saver=False
)
# Create a cache with the memory pools
params = CacheInitParams(
disable=False,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool,
page_size=1,
eviction_policy="slru",
enable_kv_cache_events=False,
)
self.cache = RadixCache(params)
def test_eviction_mechanism(self):
"""Test that SLRU eviction mechanism works correctly"""
# Insert one key-value three times (high frequency access)
frequent_key = RadixKey(
token_ids=[1, 2], extra_key=None
) # High hit rate, should be retained
frequent_val = torch.tensor([10, 20], dtype=torch.int64)
# Insert the frequent key multiple times to increase its hit count
for _ in range(3):
self.cache.insert(InsertParams(key=frequent_key, value=frequent_val))
# Insert first low-frequency key-value pair that should be evicted
first_low_freq_key = RadixKey(
token_ids=[5, 6], extra_key=None
) # Low hit rate, should be evicted
first_low_freq_val = torch.tensor([50, 60], dtype=torch.int64)
self.cache.insert(
InsertParams(key=first_low_freq_key, value=first_low_freq_val)
)
# Insert other key-values once each (low frequency access) - fill up the cache
other_keys = []
for i in range(4): # Reduce the number to fit in our smaller cache
key = RadixKey(
token_ids=[i + 10], extra_key=None
) # Unique keys for low-frequency items
val = torch.tensor([i + 100], dtype=torch.int64)
self.cache.insert(InsertParams(key=key, value=val))
other_keys.append(key)
# Now insert more items to trigger evictions
for i in range(6, 10): # Add more items to definitely exceed capacity
key = RadixKey(
token_ids=[i * 2], extra_key=None
) # Different pattern to avoid conflicts
val = torch.tensor([i * 200], dtype=torch.int64)
self.cache.insert(InsertParams(key=key, value=val))
# Now trigger eviction explicitly to make space
evict_result = self.cache.evict(
EvictParams(num_tokens=4)
) # Try to evict 4 tokens worth of space
# Check if the frequently accessed key-value is still present
# The frequent key should have higher hit count and remain in cache due to SLRU policy
frequent_match_result = self.cache.match_prefix(
MatchPrefixParams(key=frequent_key)
)
# Check if the first low-frequency key-value has been evicted
# The first low-freq key should have lower hit count and be evicted due to SLRU policy
first_low_freq_match_result = self.cache.match_prefix(
MatchPrefixParams(key=first_low_freq_key)
)
# Verify the frequent key is still present in cache after evictions
self.assertIsNotNone(
frequent_match_result,
"Frequently accessed key should still be in cache after evictions",
)
# Check if the tensor is empty, which indicates the key was not found (evicted)
is_frequent_key_present = frequent_match_result.device_indices.numel() > 0
self.assertTrue(
is_frequent_key_present,
"Frequently accessed key should still be in cache after evictions",
)
# Verify the first low-frequency key has been evicted
# The device_indices tensor should be empty when the key is not found
is_first_low_freq_key_present = (
first_low_freq_match_result.device_indices.numel() > 0
)
self.assertFalse(
is_first_low_freq_key_present,
"First inserted low-frequency key should be evicted after evictions",
)
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_radix_cache_unit.py
+778
View File
@@ -0,0 +1,778 @@
"""
Unit tests for the RadixCache implementation.
This module tests the core functionality of RadixCache, RadixKey, and TreeNode
following SGLang testing patterns.
Test Coverage:
- RadixKey: token ID management, slicing, iteration, representation
- TreeNode: node properties, reference counting, hash values
- RadixCache: insert/match operations, eviction, page alignment, error handling
- Cache events and request handling
- Boundary conditions with parameterized testing
Usage:
python test_radix_cache_unit.py
python -m pytest test_radix_cache_unit.py -v
python -m pytest test_radix_cache_unit.py::TestRadixCache::test_insert_basic
"""
from sglang.srt.mem_cache.common import available_and_evictable_str
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
# CPU-based unit test, runs quickly on any GPU runner
register_cuda_ci(est_time=5, suite="stage-b-test-small-1-gpu")
register_amd_ci(est_time=5, suite="stage-b-test-small-1-gpu-amd")
import random
import time
import unittest
import unittest.mock
import torch
from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored
from sglang.srt.mem_cache.base_prefix_cache import (
EvictParams,
EvictResult,
InsertParams,
MatchPrefixParams,
)
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
# Test constants
DEFAULT_PAGE_SIZE = 4
class TestRadixKey(unittest.TestCase):
"""Test cases for RadixKey class."""
def test_init_basic(self):
"""Test basic initialization of RadixKey."""
token_ids = [1, 2, 3, 4]
key = RadixKey(token_ids)
self.assertEqual(key.token_ids, token_ids)
self.assertIsNone(key.extra_key)
def test_init_with_extra_key(self):
"""Test initialization with extra_key."""
token_ids = [1, 2, 3]
extra_key = "test_key"
key = RadixKey(token_ids, extra_key)
self.assertEqual(key.token_ids, token_ids)
self.assertEqual(key.extra_key, extra_key)
def test_len(self):
"""Test __len__ method."""
key = RadixKey([1, 2, 3])
self.assertEqual(len(key), 3)
empty_key = RadixKey([])
self.assertEqual(len(empty_key), 0)
def test_iter(self):
"""Test __iter__ method."""
token_ids = [1, 2, 3, 4]
key = RadixKey(token_ids)
self.assertEqual(list(key), token_ids)
def test_len_and_iter(self):
"""Test __len__ and __iter__ methods."""
test_cases = [
([1, 2, 3], 3),
([], 0),
([42], 1),
]
for tokens, expected in test_cases:
with self.subTest(tokens=tokens):
key = RadixKey(tokens)
self.assertEqual(len(key), expected)
self.assertEqual(list(key), tokens)
def test_getitem_int(self):
"""Test __getitem__ with int index."""
test_cases = [
([10, 20, 30], 0, [10]),
([10, 20, 30], -1, [30]),
([10, 20, 30], 2, [30]),
]
for tokens, index, expected in test_cases:
with self.subTest(tokens=tokens, index=index):
key = RadixKey(tokens)
result = key[index]
self.assertIsInstance(result, RadixKey)
self.assertEqual(result.token_ids, expected)
def test_getitem_slice(self):
"""Test __getitem__ with slice and edge cases."""
key = RadixKey([1, 2, 3, 4, 5], "extra")
# Basic slice
sliced = key[1:4]
self.assertIsInstance(sliced, RadixKey)
self.assertEqual(sliced.token_ids, [2, 3, 4])
self.assertEqual(sliced.extra_key, "extra")
# Edge cases
self.assertEqual(key[2:2].token_ids, []) # Empty slice
self.assertEqual(key[:].token_ids, [1, 2, 3, 4, 5]) # Full slice
def test_getitem_invalid_index(self):
"""Test __getitem__ with invalid indices."""
key = RadixKey([1, 2, 3])
with self.assertRaises(IndexError):
_ = key[10] # Out of bounds
def test_repr(self):
"""Test __repr__ method."""
key = RadixKey([1, 2, 3], "test")
repr_str = repr(key)
self.assertIn("RadixKey", repr_str)
self.assertIn("extra_key='test'", repr_str)
self.assertIn("[1, 2, 3]", repr_str)
def test_repr_long_token_ids(self):
"""Test __repr__ with long token_ids."""
long_tokens = list(range(15))
key = RadixKey(long_tokens)
repr_str = repr(key)
self.assertIn("...", repr_str) # Should be truncated
class TestTreeNode(unittest.TestCase):
"""Test cases for TreeNode class."""
def setUp(self):
"""Reset the counter before each test."""
TreeNode.counter = 0
def test_init_basic(self):
"""Test basic initialization of TreeNode."""
node = TreeNode()
self.assertEqual(node.id, 0)
self.assertEqual(len(node.children), 0)
self.assertIsNone(node.parent)
self.assertIsNone(node.key)
self.assertIsNone(node.value)
self.assertEqual(node.lock_ref, 0)
self.assertEqual(node.hit_count, 0)
self.assertEqual(node.host_ref_counter, 0)
self.assertIsNone(node.host_value)
self.assertIsNone(node.hash_value)
def test_init_with_id(self):
"""Test initialization with custom ID."""
node = TreeNode(id=42)
self.assertEqual(node.id, 42)
node2 = TreeNode()
self.assertEqual(node2.id, 1) # Counter was incremented
def test_counter_increment(self):
"""Test that counter increments properly."""
node1 = TreeNode()
node2 = TreeNode()
self.assertEqual(node1.id, 0)
self.assertEqual(node2.id, 1)
def test_evicted_backuped_properties(self):
"""Test evicted and backuped properties."""
test_cases = [
(False, False, True, False),
(True, False, False, False),
(True, True, False, True),
(False, True, True, True),
]
for (
has_value,
has_host_value,
expected_evicted,
expected_backuped,
) in test_cases:
with self.subTest(has_value=has_value, has_host_value=has_host_value):
node = TreeNode()
if has_value:
node.value = torch.tensor([1, 2, 3])
if has_host_value:
node.host_value = torch.tensor([4, 5, 6])
self.assertEqual(node.evicted, expected_evicted)
self.assertEqual(node.backuped, expected_backuped)
def test_protect_release_host(self):
"""Test protect_host and release_host methods."""
node = TreeNode()
self.assertEqual(node.host_ref_counter, 0)
node.protect_host()
self.assertEqual(node.host_ref_counter, 1)
node.release_host()
self.assertEqual(node.host_ref_counter, 0)
# Test error case
with self.assertRaises(RuntimeError):
node.release_host()
def test_get_last_hash_value(self):
"""Test get_last_hash_value method."""
node = TreeNode()
self.assertIsNone(node.get_last_hash_value())
node.hash_value = ["hash1", "hash2", "hash3"]
self.assertEqual(node.get_last_hash_value(), "hash3")
def test_lt_comparison(self):
"""Test less than comparison based on last_access_time."""
node1 = TreeNode()
time.sleep(0.001) # Small delay to ensure different timestamps
node2 = TreeNode()
self.assertTrue(node1 < node2)
self.assertFalse(node2 < node1)
class TestRadixCache(unittest.TestCase):
"""Test cases for RadixCache class."""
def setUp(self):
"""Set up test fixtures."""
TreeNode.counter = 0
def test_init_variations(self):
"""Test cache initialization with different parameters."""
test_cases = [
(1, False, False),
(4, False, True),
(1, True, False),
]
for page_size, disable, enable_events in test_cases:
with self.subTest(
page_size=page_size, disable=disable, enable_events=enable_events
):
cache = RadixCache.create_simulated(
disable=disable,
page_size=page_size,
enable_kv_cache_events=enable_events,
)
self.assertEqual(cache.page_size, page_size)
self.assertEqual(cache.disable, disable)
self.assertEqual(cache.enable_kv_cache_events, enable_events)
self.assertEqual(cache.device, torch.device("cpu"))
self.assertIsNotNone(cache.root_node)
self.assertEqual(len(cache.root_node.key), 0)
def test_reset(self):
"""Test reset method."""
cache = RadixCache.create_simulated()
# Insert some data
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3]),
value=torch.tensor([10, 20, 30], dtype=torch.int64),
)
)
self.assertGreater(cache.total_size(), 0)
# Reset
cache.reset()
self.assertEqual(cache.total_size(), 0)
self.assertEqual(cache.evictable_size(), 0)
self.assertEqual(cache.protected_size(), 0)
def test_insert_and_match_basic(self):
"""Test basic insert and match operations."""
for disable_cache in [False, True]:
with self.subTest(disable_cache=disable_cache):
cache = RadixCache.create_simulated(disable=disable_cache)
key = RadixKey([1, 2, 3])
value = torch.tensor([10, 20, 30], dtype=torch.int64)
result = cache.insert(InsertParams(key=key, value=value))
prefix_len = result.prefix_len
if disable_cache:
self.assertEqual(prefix_len, 0)
self.assertEqual(cache.total_size(), 0)
continue
self.assertEqual(prefix_len, 0) # No existing prefix
self.assertEqual(cache.total_size(), 3)
self.assertEqual(cache.evictable_size(), 3)
# Test match_prefix
result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3])))
self.assertEqual(len(result.device_indices), 3)
torch.testing.assert_close(result.device_indices, value)
# Test partial match
result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2])))
self.assertEqual(len(result.device_indices), 2)
torch.testing.assert_close(
result.device_indices, torch.tensor([10, 20], dtype=torch.int64)
)
def test_insert_with_none_value(self):
"""Test insert with None value (should use token_ids as list)."""
cache = RadixCache.create_simulated()
key = RadixKey([1, 2, 3])
result = cache.insert(InsertParams(key=key, value=None))
prefix_len = result.prefix_len
# When None is passed, it should create value from token_ids
self.assertEqual(prefix_len, 0)
self.assertEqual(cache.total_size(), 3)
def test_total_size(self):
"""Test total_size calculation."""
cache = RadixCache.create_simulated()
self.assertEqual(cache.total_size(), 0)
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3]),
value=torch.tensor([10, 20, 30], dtype=torch.int64),
)
)
self.assertEqual(cache.total_size(), 3)
cache.insert(
InsertParams(
key=RadixKey([4, 5]), value=torch.tensor([40, 50], dtype=torch.int64)
)
)
self.assertEqual(cache.total_size(), 5)
def test_kv_cache_events(self):
"""Test KV cache events functionality."""
test_cases = [
(1, True),
(2, True),
(1, False),
]
for page_size, enable_events in test_cases:
with self.subTest(page_size=page_size, enable_events=enable_events):
cache = RadixCache.create_simulated(
page_size=page_size, enable_kv_cache_events=enable_events
)
# Insert data
cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4, 5]), value=None))
# Take events
events = cache.take_events()
if enable_events:
self.assertGreater(len(events), 0)
# Verify events include BlockStored events (there might be other event types)
block_stored_events = [
e for e in events if isinstance(e, BlockStored)
]
self.assertGreater(len(block_stored_events), 0)
for event in block_stored_events:
self.assertLessEqual(len(event.token_ids), page_size)
else:
self.assertEqual(len(events), 0)
def test_kv_cache_events_with_eviction(self):
"""Test KV cache events include removal events."""
mock_allocator = unittest.mock.Mock()
mock_allocator.device = torch.device("cpu")
cache = RadixCache.create_simulated(
mock_allocator=mock_allocator, enable_kv_cache_events=True
)
# Insert and then evict data
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3]),
value=torch.tensor([10, 20, 30], dtype=torch.int64),
)
)
result = cache.evict(EvictParams(num_tokens=3))
self.assertIsInstance(result, EvictResult)
self.assertGreaterEqual(
result.num_tokens_evicted,
3,
f"evicted {result.num_tokens_evicted} tokens, expected at least 3",
)
# Take events - should include both store and remove events
events = cache.take_events()
self.assertGreater(len(events), 0)
# Check event types
event_types = [type(event).__name__ for event in events]
self.assertIn("BlockStored", event_types)
# Verify BlockRemoved event content
remove_events = [e for e in events if isinstance(e, BlockRemoved)]
for event in remove_events:
self.assertGreater(len(event.block_hashes), 0)
def test_extra_key_isolation(self):
"""Test that keys with different extra_key values are isolated."""
cache = RadixCache.create_simulated()
# Insert same token sequence with different extra keys
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3], "key1"),
value=torch.tensor([10, 20, 30], dtype=torch.int64),
)
)
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3], "key2"),
value=torch.tensor([40, 50, 60], dtype=torch.int64),
)
)
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3], None),
value=torch.tensor([70, 80, 90], dtype=torch.int64),
)
)
# Keys with different extra_key should not match each other
result1 = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3], "key1")))
result2 = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3], "key2")))
result3 = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3], None)))
result4 = cache.match_prefix(
MatchPrefixParams(key=RadixKey([1, 2, 3], "nonexistent"))
)
# Each should match only its own data
self.assertEqual(len(result1.device_indices), 3)
torch.testing.assert_close(
result1.device_indices, torch.tensor([10, 20, 30], dtype=torch.int64)
)
self.assertEqual(len(result2.device_indices), 3)
torch.testing.assert_close(
result2.device_indices, torch.tensor([40, 50, 60], dtype=torch.int64)
)
self.assertEqual(len(result3.device_indices), 3)
torch.testing.assert_close(
result3.device_indices, torch.tensor([70, 80, 90], dtype=torch.int64)
)
# Non-existent extra_key should not match
self.assertEqual(len(result4.device_indices), 0)
def test_lock_ref_operations(self):
"""Test lock reference counting operations."""
cache = RadixCache.create_simulated()
# Insert sequence
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3]),
value=torch.tensor([10, 20, 30], dtype=torch.int64),
)
)
# Get node
result = cache.match_prefix(MatchPrefixParams(key=RadixKey([1, 2, 3])))
node = result.last_device_node
initial_evictable = cache.evictable_size()
initial_protected = cache.protected_size()
# Lock the node
cache.inc_lock_ref(node)
self.assertEqual(cache.protected_size(), initial_protected + 3)
self.assertEqual(cache.evictable_size(), initial_evictable - 3)
# Unlock the node
cache.dec_lock_ref(node)
self.assertEqual(cache.protected_size(), initial_protected)
self.assertEqual(cache.evictable_size(), initial_evictable)
def test_evict_functionality(self):
"""Test eviction functionality."""
mock_allocator = unittest.mock.Mock()
mock_allocator.device = torch.device("cpu")
cache = RadixCache.create_simulated(mock_allocator=mock_allocator)
# Insert sequences
cache.insert(
InsertParams(
key=RadixKey([1, 2]), value=torch.tensor([10, 20], dtype=torch.int64)
)
)
cache.insert(
InsertParams(
key=RadixKey([3, 4]), value=torch.tensor([30, 40], dtype=torch.int64)
)
)
initial_size = cache.total_size()
# Evict some tokens
result = cache.evict(EvictParams(num_tokens=2))
self.assertIsInstance(result, EvictResult)
self.assertGreaterEqual(
result.num_tokens_evicted,
2,
f"evicted {result.num_tokens_evicted} tokens, expected at least 2",
)
# Should have called free and reduced size
mock_allocator.free.assert_called()
self.assertLess(cache.total_size(), initial_size)
def test_page_alignment_boundary(self):
"""Test page alignment with different sizes."""
test_cases = [
(1, 5),
(2, 5),
(4, 6),
]
for page_size, sequence_length in test_cases:
with self.subTest(page_size=page_size, sequence_length=sequence_length):
cache = RadixCache.create_simulated(page_size=page_size)
tokens = list(range(sequence_length))
cache.insert(
InsertParams(
key=RadixKey(tokens),
value=torch.tensor(tokens, dtype=torch.int64),
)
)
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(tokens)))
self.assertGreater(len(result.device_indices), 0)
# Match length should be page-aligned
match_len = len(result.device_indices)
self.assertEqual(match_len % page_size, 0)
def test_pretty_print_basic(self):
"""Test pretty_print produces output."""
cache = RadixCache.create_simulated()
cache.insert(
InsertParams(
key=RadixKey([1, 2, 3]),
value=torch.tensor([10, 20, 30], dtype=torch.int64),
)
)
# Just test that it doesn't crash
try:
cache.pretty_print()
except Exception as e:
self.fail(f"pretty_print raised an exception: {e}")
def test_all_values_flatten(self):
"""Test all_values_flatten method."""
cache = RadixCache.create_simulated()
cache.insert(
InsertParams(
key=RadixKey([1, 2]), value=torch.tensor([10, 20], dtype=torch.int64)
)
)
cache.insert(
InsertParams(
key=RadixKey([3, 4]), value=torch.tensor([30, 40], dtype=torch.int64)
)
)
all_values = cache.all_values_flatten()
self.assertEqual(len(all_values), 4)
# Values should contain all inserted values (order may vary)
values_set = set(all_values.tolist())
self.assertEqual(values_set, {10, 20, 30, 40})
def test_advanced_prefix_match_with_node_splits(self):
"""Advanced prefix matching: splits inside nodes and across pages."""
for page_size in [1, 2]:
with self.subTest(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]
val1 = torch.tensor([x * 10 for x in seq1], dtype=torch.int64)
cache.insert(InsertParams(key=RadixKey(seq1), value=val1))
# Insert a diverging branch to create an internal node on the path.
seq2 = [1, 2, 9, 10]
val2 = torch.tensor([x * 10 for x in seq2], dtype=torch.int64)
cache.insert(InsertParams(key=RadixKey(seq2), value=val2))
print(cache.pretty_print())
baseline_total = cache.total_size()
expected_total = 10 # 8 + 2
self.assertEqual(baseline_total, expected_total)
# Match that causes a split inside an existing node:
# take first 4 tokens of seq1, then diverge.
query1 = [1, 2, 3, 4, 999, 1000]
result1 = cache.match_prefix(MatchPrefixParams(key=RadixKey(query1)))
torch.testing.assert_close(result1.device_indices, val1[:4])
# No data change after structural split during matching.
self.assertEqual(cache.total_size(), baseline_total)
# Full match of the long sequence still returns the full indices.
result_full = cache.match_prefix(MatchPrefixParams(key=RadixKey(seq1)))
torch.testing.assert_close(result_full.device_indices, val1)
# Another split deeper on the path (after matching 6 tokens, then diverge).
query2 = [1, 2, 3, 4, 5, 6, 777, 888]
result2 = cache.match_prefix(MatchPrefixParams(key=RadixKey(query2)))
torch.testing.assert_close(result2.device_indices, val1[:6])
self.assertEqual(cache.total_size(), baseline_total)
# Matching the short diverging branch should return exactly its indices.
result_branch = cache.match_prefix(
MatchPrefixParams(key=RadixKey(seq2))
)
torch.testing.assert_close(result_branch.device_indices, val2)
def test_hash_value_storage(self):
"""Test that hash_value is stored correctly after insert operations."""
cache = RadixCache.create_simulated(
page_size=4,
enable_kv_cache_events=True,
)
# Insert a sequence
cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4, 5, 6, 7, 8]), value=None))
# Trigger event emission to compute hash_value lazily
cache.take_events()
# Find the inserted node (traverse from root)
node = cache.root_node
for i in range(0, 8, 4): # page_size=4, so 2 pages
child_key = tuple([1, 2, 3, 4][:4]) if i == 0 else tuple([5, 6, 7, 8][:4])
if child_key in node.children:
node = node.children[child_key]
break
# Verify hash_value is set (computed lazily during event emission)
self.assertIsNotNone(node.hash_value)
# Should have 2 pages (8 tokens / 4 page_size)
self.assertEqual(len(node.hash_value), 2)
def test_hash_value_repeating_tokens(self):
"""Test that repeating token patterns get different hash values."""
cache = RadixCache.create_simulated(
page_size=4,
enable_kv_cache_events=True,
)
# Insert a sequence with repeating token pattern: [1,2,3,4, 1,2,3,4]
cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4, 1, 2, 3, 4]), value=None))
events = cache.take_events()
block_stored_events = [e for e in events if isinstance(e, BlockStored)]
# Should have 2 blocks (2 pages of size 4)
self.assertEqual(len(block_stored_events), 2)
# Extract block hashes
block_hash_1 = block_stored_events[0].block_hashes[0]
block_hash_2 = block_stored_events[1].block_hashes[0]
# The two blocks should have DIFFERENT hashes despite same content
# because they are at different positions (sequence-aware hashing)
self.assertNotEqual(
block_hash_1,
block_hash_2,
"Repeating token patterns should get different sequence-aware hashes",
)
# First block should have no parent
self.assertIsNone(block_stored_events[0].parent_block_hash)
# Second block's parent should be the first block's hash
self.assertEqual(block_stored_events[1].parent_block_hash, block_hash_1)
def test_hash_value_split(self):
"""Test that hash_value is split correctly when nodes are split."""
cache = RadixCache.create_simulated(
page_size=2,
enable_kv_cache_events=True,
)
# Insert a sequence that will cause a split
cache.insert(InsertParams(key=RadixKey([1, 2, 3, 4]), value=None))
cache.take_events() # Clear events and compute hash_value for first node
# Insert a diverging sequence that will cause a split at page boundary
cache.insert(InsertParams(key=RadixKey([1, 2, 5, 6]), value=None))
cache.take_events() # Trigger event emission to compute hash_value
# Find the split node
node = cache.root_node
child_key = tuple([1, 2])
if child_key in node.children:
node = node.children[child_key]
# After split and event emission, hash_value should be computed
# Note: If hash_value wasn't set before split, it will be computed lazily
# during event emission. If it was set, it will be split.
# Either way, after events are emitted, it should be set.
self.assertIsNotNone(node.hash_value)
# Should have 1 page (split at page_size=2)
self.assertEqual(len(node.hash_value), 1)
def test_memory_allocated(self):
keys, values = [], []
num_seqs = 10000
vocab_size = 1000
base_prefix_len = 10000
suffix_len = 100
torch_allocated_before = torch.cuda.memory_allocated()
# build dataset with common prefix
common_prefix = [random.randint(1, vocab_size) for _ in range(base_prefix_len)]
for _ in range(num_seqs):
suffix = [random.randint(1, vocab_size) for _ in range(suffix_len)]
seq = common_prefix + suffix
keys.append(seq)
values.append(torch.zeros(len(seq), device="cuda", dtype=torch.int32))
cache: RadixCache = RadixCache.create_simulated()
for key, value in zip(keys, values):
cache.insert(InsertParams(key=RadixKey(key), value=value))
del values
torch_allocated = torch.cuda.memory_allocated() - torch_allocated_before
cache_size_bytes = cache.total_size() * 4
print(f"\nCache size (MB): {cache_size_bytes / (1024 * 1024)}")
print(f"Torch allocated (MB): {torch_allocated / (1024 * 1024)}")
# The cache size should be within reasonable bounds of the actual allocated memory.
self.assertLess(torch_allocated, cache_size_bytes * 2)
def test_available_and_evictable_str(self):
mock_allocator = unittest.mock.Mock()
mock_allocator.available_size.return_value = 10
cache: RadixCache = RadixCache.create_simulated(mock_allocator=mock_allocator)
print(cache.available_and_evictable_str())
print(available_and_evictable_str(cache))
if __name__ == "__main__":
unittest.main()
test/registered/unit/mem_cache/test_swa_unittest.py
+557
View File
@@ -0,0 +1,557 @@
import unittest
import torch
from sglang.srt.mem_cache.base_prefix_cache import (
EvictParams,
EvictResult,
InsertParams,
MatchPrefixParams,
)
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.common import available_and_evictable_str
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixKey
from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool, SWATokenToKVPoolAllocator
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
from sglang.srt.utils import get_device
from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
register_cuda_ci(est_time=8, suite="stage-b-test-large-1-gpu")
register_amd_ci(est_time=10, suite="stage-b-test-small-1-gpu-amd")
class TestSWA(unittest.TestCase):
class _DummyReq:
def __init__(self):
self._kv_committed_len = 0
def pop_committed_kv_cache(self):
return self._kv_committed_len
def _build_swa_tree(
self,
is_eagle: bool,
page_size: int = 1,
req_size: int = 8,
max_context_len: int = 64,
kv_size: int = 64,
kv_size_swa: int = 32,
sliding_window_size: int = 4,
):
head_num = 8
head_dim = 128
num_layers = 24
global_interval = 4
dtype = torch.bfloat16
device = get_device()
full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)]
full_attention_layer_ids_set = set(full_attention_layer_ids)
swa_attention_layer_ids = [
i for i in range(num_layers) if i not in full_attention_layer_ids_set
]
req_to_token_pool = ReqToTokenPool(
size=req_size,
max_context_len=max_context_len,
device=device,
enable_memory_saver=False,
)
kv_pool = SWAKVPool(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=dtype,
head_num=head_num,
head_dim=head_dim,
swa_attention_layer_ids=swa_attention_layer_ids,
full_attention_layer_ids=full_attention_layer_ids,
enable_kvcache_transpose=False,
device=device,
)
allocator = SWATokenToKVPoolAllocator(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=dtype,
device=device,
kvcache=kv_pool,
need_sort=False,
)
tree = SWARadixCache(
params=CacheInitParams(
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=allocator,
page_size=page_size,
disable=False,
is_eagle=is_eagle,
sliding_window_size=sliding_window_size,
),
)
return tree, allocator, req_to_token_pool
@classmethod
def setUpClass(cls):
pass
@classmethod
def tearDownClass(cls):
pass
def test_swa_memory_pool(self):
size = 16
size_swa = 16
page_size = 1
head_num = 8
head_dim = 128
num_layers = 48
global_interval = 4
dtype = torch.bfloat16
device = get_device()
full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)]
full_attention_layer_ids_set = set(full_attention_layer_ids)
swa_attention_layer_ids = [
i for i in range(num_layers) if i not in full_attention_layer_ids_set
]
pool = SWAKVPool(
size=size,
size_swa=size_swa,
page_size=page_size,
dtype=dtype,
head_num=head_num,
head_dim=head_dim,
swa_attention_layer_ids=swa_attention_layer_ids,
full_attention_layer_ids=full_attention_layer_ids,
enable_kvcache_transpose=False,
device=device,
)
alloc = SWATokenToKVPoolAllocator(
size=size,
size_swa=size_swa,
page_size=page_size,
dtype=dtype,
device=device,
kvcache=pool,
need_sort=False,
)
self.assertEqual(
alloc.full_available_size() + alloc.swa_available_size(), size + size_swa
)
index = alloc.alloc(1)
self.assertEqual(
alloc.full_available_size() + alloc.swa_available_size(),
size_swa + size_swa - 2,
)
alloc.free_swa(index)
result = alloc.translate_loc_from_full_to_swa(index)
print(result)
def test_swa_radix_cache_1(self):
# args
req_size = 10
max_context_len = 128
kv_size = 128
kv_size_swa = 64
page_size = 1
sliding_window_size = 4
head_num = 8
head_dim = 128
num_layers = 48
global_interval = 4
dtype = torch.bfloat16
device = get_device()
full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)]
full_attention_layer_ids_set = set(full_attention_layer_ids)
swa_attention_layer_ids = [
i for i in range(num_layers) if i not in full_attention_layer_ids_set
]
# setup req to token pool
req_to_token_pool = ReqToTokenPool(
size=req_size,
max_context_len=max_context_len,
device=device,
enable_memory_saver=False,
)
# setup kv pool
kv_pool = SWAKVPool(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=dtype,
head_num=head_num,
head_dim=head_dim,
swa_attention_layer_ids=swa_attention_layer_ids,
full_attention_layer_ids=full_attention_layer_ids,
enable_kvcache_transpose=False,
device=device,
)
# setup token to kv pool allocator
allocator = SWATokenToKVPoolAllocator(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=dtype,
device=device,
kvcache=kv_pool,
need_sort=False,
)
# setup radix cache
tree = SWARadixCache(
params=CacheInitParams(
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=allocator,
disable=False,
page_size=page_size,
sliding_window_size=sliding_window_size,
),
)
# test
print(
f"[Start] allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
)
req1_token_ids, req1_kv_indices = [1, 2, 3], allocator.alloc(3)
self.assertEqual(len(req1_token_ids), len(req1_kv_indices))
print(
f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}"
)
result = tree.insert(
InsertParams(key=RadixKey(req1_token_ids), value=req1_kv_indices)
)
prefix_len = result.prefix_len
print(
f"req1: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
)
req2_token_ids, req2_kv_indices = [1, 2, 3, 4, 5, 6, 7], allocator.alloc(7)
self.assertEqual(len(req2_token_ids), len(req2_kv_indices))
print(
f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}"
)
result = tree.insert(
InsertParams(key=RadixKey(req2_token_ids), value=req2_kv_indices)
)
prefix_len = result.prefix_len
print(
f"req2: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
)
req3_token_ids, req3_kv_indices = [10, 11, 12], allocator.alloc(3)
self.assertEqual(len(req3_token_ids), len(req3_kv_indices))
print(
f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}"
)
result = tree.insert(
InsertParams(key=RadixKey(req3_token_ids), value=req3_kv_indices)
)
prefix_len = result.prefix_len
print(
f"req3: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
)
req4_token_ids, req4_kv_indices = [1, 2, 3, 4, 5, 60, 70], allocator.alloc(7)
self.assertEqual(len(req4_token_ids), len(req4_kv_indices))
print(
f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}"
)
result = tree.insert(
InsertParams(key=RadixKey(req4_token_ids), value=req4_kv_indices)
)
prefix_len = result.prefix_len
print(
f"req4: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
)
tree.pretty_print()
full_num_tokens, swa_num_tokens = 1, 0
print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token")
tree.evict(
EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
)
tree.pretty_print()
full_num_tokens, swa_num_tokens = 0, 1
print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token")
tree.evict(
EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
)
tree.pretty_print()
full_num_tokens, swa_num_tokens = 1, 2
print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token")
tree.evict(
EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
)
tree.pretty_print()
req5_token_ids = [1, 2, 3, 4, 5]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req5_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
)
self.assertEqual(len(kv_indices), 0)
req6_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req6_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
)
self.assertEqual(len(kv_indices), 7)
self.assertEqual(len(last_node.key), 2)
self.assertEqual(last_node.key.token_ids[0], 60)
self.assertEqual(last_node.key.token_ids[1], 70)
print(tree.available_and_evictable_str())
print(available_and_evictable_str(tree))
tree.sanity_check()
def test_swa_radix_cache_eagle(self):
# args
req_size = 10
max_context_len = 128
kv_size = 128
kv_size_swa = 64
page_size = 1
sliding_window_size = 4
head_num = 8
head_dim = 128
num_layers = 48
global_interval = 4
dtype = torch.bfloat16
device = get_device()
full_attention_layer_ids = [i for i in range(0, num_layers, global_interval)]
full_attention_layer_ids_set = set(full_attention_layer_ids)
swa_attention_layer_ids = [
i for i in range(num_layers) if i not in full_attention_layer_ids_set
]
# setup req to token pool
req_to_token_pool = ReqToTokenPool(
size=req_size,
max_context_len=max_context_len,
device=device,
enable_memory_saver=False,
)
# setup kv pool
kv_pool = SWAKVPool(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=dtype,
head_num=head_num,
head_dim=head_dim,
swa_attention_layer_ids=swa_attention_layer_ids,
full_attention_layer_ids=full_attention_layer_ids,
enable_kvcache_transpose=False,
device=device,
)
# setup token to kv pool allocator
allocator = SWATokenToKVPoolAllocator(
size=kv_size,
size_swa=kv_size_swa,
page_size=page_size,
dtype=dtype,
device=device,
kvcache=kv_pool,
need_sort=False,
)
# setup radix cache
tree = SWARadixCache(
params=CacheInitParams(
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=allocator,
page_size=page_size,
disable=False,
is_eagle=True,
sliding_window_size=sliding_window_size,
),
)
# test
print(
f"[Start] allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
)
req1_token_ids, req1_kv_indices = [1, 2, 3], allocator.alloc(3)
self.assertEqual(len(req1_token_ids), len(req1_kv_indices))
print(
f"req1: inserting, req1_token_ids: {req1_token_ids}, req1_kv_indices: {req1_kv_indices}"
)
result = tree.insert(
InsertParams(key=RadixKey(req1_token_ids), value=req1_kv_indices)
)
prefix_len = result.prefix_len
self.assertEqual(prefix_len, 0)
print(
f"req1: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
)
req2_token_ids, req2_kv_indices = [1, 2, 3, 4, 5, 6, 7], allocator.alloc(7)
self.assertEqual(len(req2_token_ids), len(req2_kv_indices))
print(
f"req2: inserting, req2_token_ids: {req2_token_ids}, req2_kv_indices: {req2_kv_indices}"
)
result = tree.insert(
InsertParams(key=RadixKey(req2_token_ids), value=req2_kv_indices)
)
prefix_len = result.prefix_len
self.assertEqual(prefix_len, 2)
print(
f"req2: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
)
req3_token_ids, req3_kv_indices = [10, 11, 12], allocator.alloc(3)
self.assertEqual(len(req3_token_ids), len(req3_kv_indices))
print(
f"req3: inserting, req3_token_ids: {req3_token_ids}, req3_kv_indices: {req3_kv_indices}"
)
result = tree.insert(
InsertParams(key=RadixKey(req3_token_ids), value=req3_kv_indices)
)
prefix_len = result.prefix_len
self.assertEqual(prefix_len, 0)
print(
f"req3: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
)
req4_token_ids, req4_kv_indices = [1, 2, 3, 4, 5, 60, 70], allocator.alloc(7)
self.assertEqual(len(req4_token_ids), len(req4_kv_indices))
print(
f"req4: inserting, req4_token_ids: {req4_token_ids}, req4_kv_indices: {req4_kv_indices}"
)
result = tree.insert(
InsertParams(key=RadixKey(req4_token_ids), value=req4_kv_indices)
)
prefix_len = result.prefix_len
self.assertEqual(prefix_len, 4)
print(
f"req4: prefix_len: {prefix_len}, allocator swa available size: {allocator.swa_available_size()}, full available size: {allocator.full_available_size()}"
)
tree.pretty_print()
full_num_tokens, swa_num_tokens = 1, 0
print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token")
evict_result = tree.evict(
EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
)
assert isinstance(evict_result, EvictResult)
assert (
evict_result.num_tokens_evicted >= full_num_tokens
) # May evict more due to node granularity
print(
f"evicted {evict_result.num_tokens_evicted} full tokens, {evict_result.swa_num_tokens_evicted} swa tokens"
)
tree.pretty_print()
full_num_tokens, swa_num_tokens = 0, 1
print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token")
evict_result = tree.evict(
EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
)
assert isinstance(evict_result, EvictResult)
assert (
evict_result.swa_num_tokens_evicted >= swa_num_tokens
), f"evicted {evict_result.swa_num_tokens_evicted} swa tokens, expected {swa_num_tokens}"
tree.pretty_print()
full_num_tokens, swa_num_tokens = 1, 2
print(f"evicting {full_num_tokens} full token and {swa_num_tokens} swa token")
evict_result = tree.evict(
EvictParams(num_tokens=full_num_tokens, swa_num_tokens=swa_num_tokens)
)
assert isinstance(evict_result, EvictResult)
assert (
evict_result.num_tokens_evicted >= full_num_tokens
), f"evicted {evict_result.num_tokens_evicted} full tokens, expected {full_num_tokens}"
assert (
evict_result.swa_num_tokens_evicted >= swa_num_tokens
), f"evicted {evict_result.swa_num_tokens_evicted} swa tokens, expected {swa_num_tokens}"
tree.pretty_print()
req5_token_ids = [1, 2, 3, 4, 5]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req5_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req5: token_ids: {req5_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
)
self.assertEqual(len(kv_indices), 0) # no swa prefix matched
req6_token_ids = [1, 2, 3, 4, 5, 60, 70]
result = tree.match_prefix(MatchPrefixParams(key=RadixKey(req6_token_ids)))
kv_indices, last_node = result.device_indices, result.last_device_node
print(
f"req6: token_ids: {req6_token_ids}, matched kv_indices: {kv_indices}, last_node.key: {last_node.key}"
)
self.assertEqual(len(kv_indices), 6)
self.assertEqual(len(last_node.key), 2)
self.assertEqual(last_node.key.token_ids[0], (5, 60))
self.assertEqual(last_node.key.token_ids[1], (60, 70))
def test_swa_cache_finished_req_eagle_uses_cache_protected_len_and_bigram_key(self):
tree, allocator, req_to_token_pool = self._build_swa_tree(is_eagle=True)
# Case 1: is_insert=True should pass bigram key and use cache_protected_len.
req = self._DummyReq()
req.req_pool_idx = 0
req.origin_input_ids = [1, 2, 3, 4, 5, 6]
req.output_ids = []
req._kv_committed_len = len(req.origin_input_ids)
kv_indices = allocator.alloc(req._kv_committed_len)
req_to_token_pool.write(
(req.req_pool_idx, slice(0, req._kv_committed_len)), kv_indices
)
req.extra_key = None
req.last_node = tree.root_node
req.swa_uuid_for_lock = None
req.swa_evicted_seqlen = 0
req.cache_protected_len = 1
# Intentionally mismatch to ensure code does not use len(prefix_indices).
req.prefix_indices = torch.tensor([7, 8, 9, 10, 11], device=tree.device)
captured = {}
original_insert = tree.insert
def wrapped_insert(params):
captured["prev_prefix_len"] = params.prev_prefix_len
captured["is_bigram"] = params.key.is_bigram
captured["key_len"] = len(params.key)
return original_insert(params)
tree.insert = wrapped_insert
tree.cache_finished_req(req, is_insert=True)
self.assertEqual(captured["prev_prefix_len"], req.cache_protected_len)
self.assertTrue(captured["is_bigram"])
self.assertEqual(captured["key_len"], len(req.origin_input_ids) - 1)
# Case 2: is_insert=False should free [cache_protected_len:page_aligned_len]
# even when len(prefix_indices) is intentionally larger.
req2 = self._DummyReq()
req2.req_pool_idx = 1
req2.origin_input_ids = [11, 12, 13, 14, 15, 16]
req2.output_ids = []
req2._kv_committed_len = len(req2.origin_input_ids)
kv_indices2 = allocator.alloc(req2._kv_committed_len)
req_to_token_pool.write(
(req2.req_pool_idx, slice(0, req2._kv_committed_len)), kv_indices2
)
req2.extra_key = None
req2.last_node = tree.root_node
req2.swa_uuid_for_lock = None
req2.swa_evicted_seqlen = 0
req2.cache_protected_len = 1
req2.prefix_indices = torch.tensor([21, 22, 23, 24, 25], device=tree.device)
freed_lens = []
original_free = allocator.free
def wrapped_free(indices):
freed_lens.append(int(indices.numel()))
return original_free(indices)
allocator.free = wrapped_free
tree.cache_finished_req(req2, is_insert=False)
# EAGLE + page_size=1 => page_aligned_len = committed_len - 1 = 5
# Expected frees:
# overlap range [1:5] -> 4
# tail range [5:] -> 1
self.assertEqual(freed_lens, [4, 1])
if __name__ == "__main__":
unittest.main()