820 lines
31 KiB
Python
Executable File
820 lines
31 KiB
Python
Executable File
import unittest
|
|
from array import array
|
|
|
|
import torch
|
|
|
|
from sglang.kernels.ops.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE
|
|
from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
|
|
from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored
|
|
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.hi_mamba_radix_cache import HiMambaRadixCache
|
|
from sglang.srt.mem_cache.mamba_radix_cache import LRUList, MambaRadixCache, TreeNode
|
|
from sglang.srt.mem_cache.memory_pool import (
|
|
HybridLinearKVPool,
|
|
HybridReqToTokenPool,
|
|
MambaPool,
|
|
)
|
|
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=10, stage="base-b", runner_config="1-gpu-small")
|
|
register_amd_ci(est_time=9, suite="stage-b-test-1-gpu-small-amd")
|
|
|
|
|
|
def _event_hashes(events):
|
|
return [block_hash for event in events for block_hash in event.block_hashes]
|
|
|
|
|
|
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,
|
|
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,
|
|
mamba_layer_ids=mamba_layers,
|
|
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_allocator.available_size() == mamba_cache_size
|
|
|
|
sampling_params = SamplingParams(
|
|
temperature=0,
|
|
max_new_tokens=1,
|
|
)
|
|
req = Req(
|
|
rid=0,
|
|
origin_input_text="",
|
|
origin_input_ids=array("q"),
|
|
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_allocator.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_allocator.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_allocator.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_allocator.available_size() == mamba_cache_size - 1
|
|
)
|
|
|
|
def test_mamba_pool_deduplicated_conv_window_axis(self):
|
|
class WindowFirstMambaPool(MambaPool):
|
|
conv_window_axis = 0
|
|
|
|
num_mamba_layers = 2
|
|
spec_state_size = 3
|
|
speculative_num_draft_tokens = 4
|
|
window_size = 3
|
|
conv_dim = 5
|
|
|
|
pool = object.__new__(WindowFirstMambaPool)
|
|
physical, view = pool._allocate_deduplicated_conv_window(
|
|
conv_shape=(window_size, conv_dim),
|
|
num_mamba_layers=num_mamba_layers,
|
|
spec_state_size=spec_state_size,
|
|
speculative_num_draft_tokens=speculative_num_draft_tokens,
|
|
conv_dtype=torch.float32,
|
|
)
|
|
|
|
shared_window_size = speculative_num_draft_tokens + window_size - 1
|
|
self.assertEqual(
|
|
physical.shape,
|
|
(
|
|
num_mamba_layers,
|
|
spec_state_size + 1,
|
|
shared_window_size,
|
|
conv_dim,
|
|
),
|
|
)
|
|
self.assertEqual(
|
|
view.shape,
|
|
(
|
|
num_mamba_layers,
|
|
spec_state_size + 1,
|
|
speculative_num_draft_tokens,
|
|
window_size,
|
|
conv_dim,
|
|
),
|
|
)
|
|
|
|
physical.copy_(
|
|
torch.arange(
|
|
physical.numel(), dtype=physical.dtype, device=physical.device
|
|
).reshape_as(physical)
|
|
)
|
|
for step in range(speculative_num_draft_tokens):
|
|
torch.testing.assert_close(
|
|
view[:, :, step],
|
|
physical[:, :, step : step + window_size],
|
|
)
|
|
torch.testing.assert_close(view[:, :, :-1, 1:], view[:, :, 1:, :-1])
|
|
|
|
view[0, 0, 0, 1, 0] = -1
|
|
self.assertEqual(view[0, 0, 1, 0, 0].item(), -1)
|
|
|
|
def test_mamba_radix_cache_1(self):
|
|
tree, allocator, req_to_token_pool, make_dummy_req = (
|
|
self._setup_tree_and_allocator()
|
|
)
|
|
mamba_allocator = req_to_token_pool.mamba_allocator
|
|
mamba_pool = req_to_token_pool.mamba_pool
|
|
# test
|
|
print(
|
|
f"[Start] allocator mamba available size: {mamba_allocator.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}"
|
|
)
|
|
key = RadixKey(array("q", req1_token_ids))
|
|
result = tree.insert(
|
|
InsertParams(
|
|
key=key,
|
|
value=req1_kv_indices[: len(key)],
|
|
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_allocator.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}"
|
|
)
|
|
key = RadixKey(array("q", req2_token_ids))
|
|
result = tree.insert(
|
|
InsertParams(
|
|
key=key,
|
|
value=req2_kv_indices[: len(key)],
|
|
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_allocator.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}"
|
|
)
|
|
key = RadixKey(array("q", req3_token_ids))
|
|
result = tree.insert(
|
|
InsertParams(
|
|
key=key,
|
|
value=req3_kv_indices[: len(key)],
|
|
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_allocator.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}"
|
|
)
|
|
key = RadixKey(array("q", req4_token_ids))
|
|
result = tree.insert(
|
|
InsertParams(
|
|
key=key,
|
|
value=req4_kv_indices[: len(key)],
|
|
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_allocator.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(array("q", 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(array("q", 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(array("q", 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(array("q", 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(array("q", 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 test_mamba_radix_cache_kv_events(self):
|
|
tree, allocator, _, make_dummy_req = self._setup_tree_and_allocator(
|
|
enable_kv_cache_events=True
|
|
)
|
|
tree.take_events() # Clear the reset event.
|
|
|
|
stored_hashes = []
|
|
|
|
req1 = make_dummy_req()
|
|
key1 = RadixKey(array("q", [1, 2, 3]))
|
|
tree.insert(
|
|
InsertParams(
|
|
key=key1,
|
|
value=allocator.alloc(3)[: len(key1)],
|
|
mamba_value=req1.mamba_pool_idx.unsqueeze(0),
|
|
)
|
|
)
|
|
events = tree.take_events()
|
|
stored_events = [e for e in events if isinstance(e, BlockStored)]
|
|
self.assertEqual(len(stored_events), 3)
|
|
self.assertEqual([e.token_ids[0] for e in stored_events], [1, 2, 3])
|
|
stored_hashes.extend(e.block_hashes[0] for e in stored_events)
|
|
|
|
req2 = make_dummy_req()
|
|
key2 = RadixKey(array("q", [1, 2, 3, 4, 5]))
|
|
tree.insert(
|
|
InsertParams(
|
|
key=key2,
|
|
value=allocator.alloc(5)[: len(key2)],
|
|
mamba_value=req2.mamba_pool_idx.unsqueeze(0),
|
|
)
|
|
)
|
|
events = tree.take_events()
|
|
stored_events = [e for e in events if isinstance(e, BlockStored)]
|
|
self.assertEqual(len(stored_events), 2)
|
|
self.assertEqual([e.token_ids[0] for e in stored_events], [4, 5])
|
|
stored_hashes.extend(e.block_hashes[0] for e in stored_events)
|
|
|
|
# Evicting an internal mamba state creates a tombstone but does not
|
|
# remove full-attention KV blocks, so it must not emit BlockRemoved.
|
|
result = tree.evict(EvictParams(num_tokens=0, mamba_num=1))
|
|
self.assertEqual(result.num_tokens_evicted, 0)
|
|
self.assertEqual(result.mamba_num_evicted, 1)
|
|
events = tree.take_events()
|
|
self.assertEqual([e for e in events if isinstance(e, BlockRemoved)], [])
|
|
|
|
result = tree.evict(EvictParams(num_tokens=1))
|
|
self.assertGreaterEqual(result.num_tokens_evicted, 1)
|
|
events = tree.take_events()
|
|
removed_hashes = _event_hashes(
|
|
[e for e in events if isinstance(e, BlockRemoved)]
|
|
)
|
|
self.assertCountEqual(removed_hashes, stored_hashes)
|
|
|
|
def test_mamba_radix_cache_kv_events_split_hash(self):
|
|
tree, allocator, _, make_dummy_req = self._setup_tree_and_allocator(
|
|
enable_kv_cache_events=True
|
|
)
|
|
tree.take_events() # Clear the reset event.
|
|
|
|
req1 = make_dummy_req()
|
|
key1 = RadixKey(array("q", [1, 2, 3, 4]))
|
|
tree.insert(
|
|
InsertParams(
|
|
key=key1,
|
|
value=allocator.alloc(4)[: len(key1)],
|
|
mamba_value=req1.mamba_pool_idx.unsqueeze(0),
|
|
)
|
|
)
|
|
first_insert_events = [
|
|
e for e in tree.take_events() if isinstance(e, BlockStored)
|
|
]
|
|
self.assertEqual(len(first_insert_events), 4)
|
|
split_parent_hash = first_insert_events[1].block_hashes[0]
|
|
|
|
req2 = make_dummy_req()
|
|
key2 = RadixKey(array("q", [1, 2, 5, 6]))
|
|
tree.insert(
|
|
InsertParams(
|
|
key=key2,
|
|
value=allocator.alloc(4)[: len(key2)],
|
|
mamba_value=req2.mamba_pool_idx.unsqueeze(0),
|
|
)
|
|
)
|
|
second_insert_events = [
|
|
e for e in tree.take_events() if isinstance(e, BlockStored)
|
|
]
|
|
self.assertEqual(len(second_insert_events), 2)
|
|
self.assertEqual(list(second_insert_events[0].token_ids), [5])
|
|
self.assertEqual(second_insert_events[0].parent_block_hash, split_parent_hash)
|
|
|
|
def _setup_tree_and_allocator(self, enable_kv_cache_events=False):
|
|
"""Helper to create a MambaRadixCache with allocator for testing."""
|
|
server_args = ServerArgs(model_path="dummy", page_size=1)
|
|
# MambaRadixCache reads mamba_cache_chunk_size, whose property otherwise
|
|
# loads the HF config for self.model_path — impossible for the dummy model.
|
|
# Mirror the property's default for a dummy HF config: FLA_CHUNK_SIZE.
|
|
server_args._mamba_cache_chunk_size = FLA_CHUNK_SIZE
|
|
set_global_server_args_for_scheduler(server_args)
|
|
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,
|
|
mamba_layer_ids=mamba_layers,
|
|
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,
|
|
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,
|
|
enable_kv_cache_events=enable_kv_cache_events,
|
|
)
|
|
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=array("q"),
|
|
sampling_params=sampling_params,
|
|
)
|
|
req_to_token_pool.alloc([req])
|
|
return req
|
|
|
|
return tree, allocator, req_to_token_pool, make_dummy_req
|
|
|
|
def test_hi_mamba_tombstone_cleanup_respects_host_ref(self):
|
|
tree = object.__new__(HiMambaRadixCache)
|
|
root = TreeNode()
|
|
parent = TreeNode()
|
|
deleted = TreeNode()
|
|
|
|
root.key = RadixKey(array("q", []))
|
|
parent.key = RadixKey(array("q", [1]))
|
|
deleted.key = RadixKey(array("q", [2]))
|
|
parent.parent = root
|
|
deleted.parent = parent
|
|
parent.value = torch.tensor([1], dtype=torch.int64)
|
|
parent.protect_host()
|
|
root.children[parent.key.child_key(1)] = parent
|
|
|
|
class RecordingCacheController:
|
|
def __init__(self):
|
|
self.device_evictions = []
|
|
self.host_evictions = []
|
|
|
|
def evict_device(self, value):
|
|
self.device_evictions.append(value)
|
|
|
|
def evict_host(self, value):
|
|
self.host_evictions.append(value)
|
|
|
|
tree.root_node = root
|
|
tree.page_size = 1
|
|
tree.full_lru_list = LRUList(mamba=False)
|
|
tree.full_lru_list.insert_mru(parent)
|
|
tree.cache_controller = RecordingCacheController()
|
|
tree.full_evictable_size_ = len(parent.value)
|
|
tree.evictable_full_device_leaves = {parent}
|
|
tree.evictable_full_host_leaves = set()
|
|
|
|
result_node, full_evicted, mamba_evicted = (
|
|
tree._iteratively_delete_tombstone_leaf(deleted)
|
|
)
|
|
|
|
self.assertIs(result_node, deleted)
|
|
self.assertEqual(full_evicted, 0)
|
|
self.assertEqual(mamba_evicted, 0)
|
|
self.assertIs(root.children[parent.key.child_key(1)], parent)
|
|
self.assertTrue(tree.full_lru_list.in_list(parent))
|
|
self.assertEqual(tree.cache_controller.device_evictions, [])
|
|
self.assertEqual(tree.cache_controller.host_evictions, [])
|
|
|
|
def test_mamba_pool_cpu_offload(self):
|
|
"""MambaPool.get_cpu_copy / load_cpu_copy round-trips conv and temporal state."""
|
|
_, _, req_to_token_pool, _ = self._setup_tree_and_allocator()
|
|
mamba_pool = req_to_token_pool.mamba_pool
|
|
n = 3
|
|
indices = req_to_token_pool.mamba_allocator.alloc(n)
|
|
self.assertIsNotNone(indices)
|
|
|
|
# Write known sentinel values at the allocated slots.
|
|
for conv in mamba_pool.mamba_cache.conv:
|
|
conv[:, indices] = 1.0
|
|
mamba_pool.mamba_cache.temporal[:, indices] = 2.0
|
|
|
|
# Save to CPU.
|
|
conv_cpu, temporal_cpu = mamba_pool.get_cpu_copy(indices)
|
|
|
|
# Verify CPU tensors match what was written.
|
|
for i, conv in enumerate(mamba_pool.mamba_cache.conv):
|
|
expected = conv[:, indices].cpu()
|
|
self.assertTrue(
|
|
torch.allclose(conv_cpu[i].float(), expected.float()),
|
|
f"conv[{i}] CPU copy mismatch",
|
|
)
|
|
expected_t = mamba_pool.mamba_cache.temporal[:, indices].cpu()
|
|
self.assertTrue(
|
|
torch.allclose(temporal_cpu.float(), expected_t.float()),
|
|
"temporal CPU copy mismatch",
|
|
)
|
|
|
|
# Zero out GPU slots and restore from CPU copy.
|
|
for conv in mamba_pool.mamba_cache.conv:
|
|
conv[:, indices] = 0.0
|
|
mamba_pool.mamba_cache.temporal[:, indices] = 0.0
|
|
|
|
mamba_pool.load_cpu_copy((conv_cpu, temporal_cpu), indices)
|
|
|
|
# Verify restored values match the sentinels.
|
|
for conv in mamba_pool.mamba_cache.conv:
|
|
restored = conv[:, indices]
|
|
self.assertTrue(
|
|
torch.all(restored == 1.0),
|
|
"conv not restored after load_cpu_copy",
|
|
)
|
|
self.assertTrue(
|
|
torch.all(mamba_pool.mamba_cache.temporal[:, indices] == 2.0),
|
|
"temporal not restored after load_cpu_copy",
|
|
)
|
|
|
|
def test_hybrid_kv_pool_cpu_offload(self):
|
|
"""HybridLinearKVPool.get_cpu_copy / load_cpu_copy saves and restores both
|
|
the full-attention KV cache and Mamba state in a single round-trip."""
|
|
_, allocator, req_to_token_pool, _ = self._setup_tree_and_allocator()
|
|
mamba_pool = req_to_token_pool.mamba_pool
|
|
hybrid_pool = allocator._kvcache # HybridLinearKVPool
|
|
|
|
self.assertIsInstance(hybrid_pool, HybridLinearKVPool)
|
|
|
|
n_tokens = 4
|
|
kv_indices = allocator.alloc(n_tokens)
|
|
self.assertIsNotNone(kv_indices)
|
|
mamba_indices = req_to_token_pool.mamba_allocator.alloc(1)
|
|
self.assertIsNotNone(mamba_indices)
|
|
|
|
# Write sentinel values into KV buffers (all full-attention layers).
|
|
for layer_id in range(hybrid_pool.full_kv_pool.layer_num):
|
|
hybrid_pool.full_kv_pool.k_buffer[layer_id][kv_indices] = 3.0
|
|
hybrid_pool.full_kv_pool.v_buffer[layer_id][kv_indices] = 4.0
|
|
|
|
# Write sentinel values into Mamba state.
|
|
for conv in mamba_pool.mamba_cache.conv:
|
|
conv[:, mamba_indices] = 5.0
|
|
mamba_pool.mamba_cache.temporal[:, mamba_indices] = 6.0
|
|
|
|
# --- Round-trip with Mamba indices provided ---
|
|
cpu_copy = allocator.get_cpu_copy(kv_indices, mamba_indices=mamba_indices)
|
|
kv_cpu, mamba_cpu = cpu_copy
|
|
self.assertIsNotNone(
|
|
mamba_cpu, "mamba_cpu should be saved when mamba_indices given"
|
|
)
|
|
|
|
# Zero out GPU.
|
|
for layer_id in range(hybrid_pool.full_kv_pool.layer_num):
|
|
hybrid_pool.full_kv_pool.k_buffer[layer_id][kv_indices] = 0.0
|
|
hybrid_pool.full_kv_pool.v_buffer[layer_id][kv_indices] = 0.0
|
|
for conv in mamba_pool.mamba_cache.conv:
|
|
conv[:, mamba_indices] = 0.0
|
|
mamba_pool.mamba_cache.temporal[:, mamba_indices] = 0.0
|
|
|
|
allocator.load_cpu_copy(cpu_copy, kv_indices, mamba_indices=mamba_indices)
|
|
|
|
# Verify KV restored.
|
|
for layer_id in range(hybrid_pool.full_kv_pool.layer_num):
|
|
self.assertTrue(
|
|
torch.all(
|
|
hybrid_pool.full_kv_pool.k_buffer[layer_id][kv_indices] == 3.0
|
|
),
|
|
f"k_buffer layer {layer_id} not restored",
|
|
)
|
|
self.assertTrue(
|
|
torch.all(
|
|
hybrid_pool.full_kv_pool.v_buffer[layer_id][kv_indices] == 4.0
|
|
),
|
|
f"v_buffer layer {layer_id} not restored",
|
|
)
|
|
|
|
# Verify Mamba restored.
|
|
for conv in mamba_pool.mamba_cache.conv:
|
|
self.assertTrue(
|
|
torch.all(conv[:, mamba_indices] == 5.0),
|
|
"conv not restored after load_cpu_copy",
|
|
)
|
|
self.assertTrue(
|
|
torch.all(mamba_pool.mamba_cache.temporal[:, mamba_indices] == 6.0),
|
|
"temporal not restored after load_cpu_copy",
|
|
)
|
|
|
|
# --- Without mamba_indices: mamba_cpu must be None ---
|
|
cpu_copy_no_mamba = allocator.get_cpu_copy(kv_indices, mamba_indices=None)
|
|
_, mamba_cpu_none = cpu_copy_no_mamba
|
|
self.assertIsNone(
|
|
mamba_cpu_none, "mamba_cpu should be None when mamba_indices=None"
|
|
)
|
|
|
|
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()
|
|
key1 = RadixKey(array("q", [1, 2, 3]))
|
|
tree.insert(
|
|
InsertParams(
|
|
key=key1,
|
|
value=allocator.alloc(3)[: len(key1)],
|
|
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()
|
|
key2 = RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7]))
|
|
result = tree.insert(
|
|
InsertParams(
|
|
key=key2,
|
|
value=allocator.alloc(7)[: len(key2)],
|
|
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()
|
|
key3 = RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7, 8]))
|
|
result = tree.insert(
|
|
InsertParams(
|
|
key=key3,
|
|
value=allocator.alloc(8)[: len(key3)],
|
|
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()
|
|
key4 = RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7, 8, 9]))
|
|
result = tree.insert(
|
|
InsertParams(
|
|
key=key4,
|
|
value=allocator.alloc(9)[: len(key4)],
|
|
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()
|