"""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")