392 lines
14 KiB
Python
392 lines
14 KiB
Python
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):
|
|
set_global_server_args_for_scheduler(
|
|
ServerArgs(model_path="dummy", page_size=1)
|
|
)
|
|
# kv cache
|
|
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
|
|
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,
|
|
)
|
|
# setup kv pool
|
|
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,
|
|
)
|
|
|
|
# setup token to kv pool allocator
|
|
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,
|
|
)
|
|
# setup radix cache
|
|
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
|
|
|
|
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()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|