Xiaoze Fan
134 lines
6.4 KiB
Python
134 lines
6.4 KiB
Python
"""P2b integration: CacheManager hybrid path (match_req -> cache_req donate -> prefix hit).
|
||
CPU, real LinearStatePool + page_table, hand-built Reqs. Exercises the two-currency wiring
|
||
without the full scheduler/engine."""
|
||
from __future__ import annotations
|
||
|
||
from types import SimpleNamespace
|
||
|
||
import torch
|
||
|
||
from freetoken.core import Req, SamplingParams
|
||
from freetoken.kvcache.linear_state_pool import LinearStatePool
|
||
from freetoken.models.config import LinearGatedDeltaGroupConfig
|
||
from freetoken.scheduler.cache import CacheManager
|
||
|
||
|
||
def _pool(num_slots=16):
|
||
g = LinearGatedDeltaGroupConfig(
|
||
name="linear", layer_ids=(0,), num_key_heads=2, num_value_heads=4,
|
||
key_head_dim=16, value_head_dim=16, conv_kernel_dim=4, output_gate=True,
|
||
)
|
||
return LinearStatePool(group=g, num_slots=num_slots, dtype=torch.bfloat16,
|
||
device=torch.device("cpu"), tp_size=1)
|
||
|
||
|
||
def _pend(ids):
|
||
# int32 to match production Req.input_ids dtype (fast_compare_key needs consistent dtype)
|
||
t = torch.tensor(ids, dtype=torch.int32)
|
||
return SimpleNamespace(input_ids=t, input_len=len(ids), mm_embeds=None)
|
||
|
||
|
||
def test_hybrid_cache_manager_donate_then_hit():
|
||
pool = _pool()
|
||
page_table = torch.zeros(4, 64, dtype=torch.int32)
|
||
cm = CacheManager(64, 1, page_table, "hybrid_radix", linear_state_pool=pool)
|
||
assert cm.is_hybrid
|
||
|
||
# cold match on an empty tree
|
||
mr = cm.match_req(_pend([1, 2, 3, 4, 5]))
|
||
assert mr.cuda_handle.cached_len == 0 and mr.mamba_value is None
|
||
|
||
# admit req A: allocate live + ping-pong, stage KV pages, mark a ×N snapshot at boundary 4
|
||
live, pp = pool.alloc(1)[0], tuple(pool.alloc(2))
|
||
page_table[0, :4] = torch.tensor([100, 101, 102, 103], dtype=torch.int32)
|
||
reqA = Req(input_ids=torch.tensor([1, 2, 3, 4, 5], dtype=torch.int32), table_idx=0,
|
||
cached_len=4, output_len=1, uid=0, sampling_params=SamplingParams(),
|
||
cache_handle=mr.cuda_handle)
|
||
reqA.linear_slot_idx, reqA.mamba_ping_pong = live, pp
|
||
reqA.mamba_next_track_idx = 1 # flipped from 0 in build_fla_metadata; frozen = pp[0]
|
||
reqA.mamba_last_track_seqlen = 4
|
||
cm.lock(mr.cuda_handle)
|
||
|
||
free_before = pool.num_free_slots
|
||
cm.cache_req(reqA, finished=False) # donate pp[0] at boundary 4; replace it in the pair
|
||
# pp[0] donated to the tree; a fresh replacement was alloc'd -> net free-slot count unchanged
|
||
assert pool.num_free_slots == free_before - 1 # one replacement alloc'd (donated slot now tree-owned)
|
||
assert reqA.mamba_ping_pong[0] != pp[0] # slot 0 replaced; pp[0] now lives in the tree
|
||
|
||
# req B shares the [1,2,3,4] prefix -> HIT: returns the donated snapshot + reused KV
|
||
mrB = cm.match_req(_pend([1, 2, 3, 4, 9]))
|
||
assert mrB.cuda_handle.cached_len == 4
|
||
assert mrB.mamba_value == pp[0]
|
||
assert mrB.cuda_handle.get_matched_indices().tolist() == [100, 101, 102, 103]
|
||
|
||
|
||
def test_hybrid_finish_donates_live_slot():
|
||
pool = _pool()
|
||
page_table = torch.zeros(4, 64, dtype=torch.int32)
|
||
cm = CacheManager(64, 1, page_table, "hybrid_radix", linear_state_pool=pool)
|
||
|
||
mr = cm.match_req(_pend([7, 8, 9, 10]))
|
||
live, pp = pool.alloc(1)[0], tuple(pool.alloc(2))
|
||
page_table[1, :3] = torch.tensor([200, 201, 202], dtype=torch.int32)
|
||
req = Req(input_ids=torch.tensor([7, 8, 9, 10], dtype=torch.int32), table_idx=1,
|
||
cached_len=3, output_len=1, uid=1, sampling_params=SamplingParams(),
|
||
cache_handle=mr.cuda_handle)
|
||
req.linear_slot_idx, req.mamba_ping_pong = live, pp
|
||
cm.lock(mr.cuda_handle)
|
||
|
||
cm.cache_req(req, finished=True) # donate the live slot directly (final state)
|
||
# ping-pong pair freed; live slot kept (now owned by the tree)
|
||
mr2 = cm.match_req(_pend([7, 8, 9, 10]))
|
||
assert mr2.cuda_handle.cached_len == 3 and mr2.mamba_value == live
|
||
|
||
|
||
def test_free_req_slots_idempotent():
|
||
"""C2: a finish/abort double-free of the same request must NOT push its GDN slots twice."""
|
||
pool = _pool()
|
||
pt = torch.zeros(4, 64, dtype=torch.int32)
|
||
cm = CacheManager(64, 1, pt, "hybrid_radix", linear_state_pool=pool)
|
||
live, pp = pool.alloc(1)[0], tuple(pool.alloc(2))
|
||
req = Req(input_ids=torch.tensor([1, 2, 3], dtype=torch.int32), table_idx=0, cached_len=2,
|
||
output_len=1, uid=0, sampling_params=SamplingParams(), cache_handle=None)
|
||
req.linear_slot_idx, req.mamba_ping_pong = live, pp
|
||
base = pool.num_free_slots
|
||
cm._free_req_slots(req)
|
||
assert pool.num_free_slots == base + 3 # live + 2 ping-pong returned once
|
||
cm._free_req_slots(req) # second free (abort/finish race)
|
||
assert pool.num_free_slots == base + 3 # idempotent: nothing pushed twice
|
||
|
||
|
||
def test_rebuild_reclaims_donated_gdn_slots():
|
||
"""C5: a runtime cache rebuild must return the discarded tree's GDN snapshot slots (idle)."""
|
||
pool = _pool(num_slots=16)
|
||
pt = torch.zeros(4, 64, dtype=torch.int32)
|
||
cm = CacheManager(64, 1, pt, "hybrid_radix", linear_state_pool=pool)
|
||
mr = cm.match_req(_pend([7, 8, 9, 10]))
|
||
live, pp = pool.alloc(1)[0], tuple(pool.alloc(2))
|
||
pt[1, :3] = torch.tensor([200, 201, 202], dtype=torch.int32)
|
||
req = Req(input_ids=torch.tensor([7, 8, 9, 10], dtype=torch.int32), table_idx=1, cached_len=3,
|
||
output_len=1, uid=1, sampling_params=SamplingParams(), cache_handle=mr.cuda_handle)
|
||
req.linear_slot_idx, req.mamba_ping_pong = live, pp
|
||
cm.lock(mr.cuda_handle)
|
||
cm.cache_req(req, finished=True) # donates `live` to the tree, frees ping-pong
|
||
assert pool.num_free_slots < pool.num_slots - 1 # a slot is now tree-owned
|
||
cm.rebuild(64, pt) # idle rebuild discards the tree
|
||
assert pool.num_free_slots == pool.num_slots - 1 # all GDN slots reclaimed (no leak)
|
||
|
||
|
||
def test_pool_sizing_covers_4mr_floor():
|
||
"""C6: pool must reserve the 4-slot-per-request non-evictable floor even at a tiny ratio."""
|
||
from types import SimpleNamespace
|
||
from freetoken.kvcache.linear_state_pool import _linear_pool_num_slots
|
||
for mr in (1, 8, 64):
|
||
c = SimpleNamespace(max_running_req=mr, cache_type="hybrid_radix",
|
||
linear_state_cache_ratio=0.1)
|
||
assert _linear_pool_num_slots(c) >= 4 * mr + 1, (mr, _linear_pool_num_slots(c))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
for name, fn in list(globals().items()):
|
||
if name.startswith("test_") and callable(fn):
|
||
fn()
|
||
print(f"{name}: PASS")
|