Xiaoze Fan
143 lines
6.0 KiB
Python
143 lines
6.0 KiB
Python
"""An unfinished commit that dedups against a prefix another request already published frees
|
|
the request's OWN pages for the shared span. Its page-table row must then name the tree's pages
|
|
instead: the attention backends read that row every decode step, and the freed pages go to the
|
|
next allocation. CPU, real CacheManager + real trees, no 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
|
|
|
|
PROMPT = [1, 2, 3, 4, 5, 6, 7, 8]
|
|
|
|
|
|
def _pend(ids):
|
|
t = torch.tensor(ids, dtype=torch.int32)
|
|
return SimpleNamespace(input_ids=t, input_len=len(ids), mm_embeds=None)
|
|
|
|
|
|
def _admit(cm, page_table, table_idx, ids, handle):
|
|
req = Req(input_ids=torch.tensor(ids, dtype=torch.int32), table_idx=table_idx,
|
|
cached_len=0, output_len=0, uid=table_idx, sampling_params=SamplingParams(),
|
|
cache_handle=handle)
|
|
req.device_len = len(ids)
|
|
cm.lock(handle)
|
|
cm.allocate_paged([req])
|
|
req.cached_len = len(ids)
|
|
return req
|
|
|
|
|
|
def _live_row(page_table, req):
|
|
return set(page_table[req.table_idx, : req.cached_len].tolist())
|
|
|
|
|
|
def test_radix_unfinished_commit_repoints_the_row_off_the_freed_pages():
|
|
page_table = torch.zeros(4, 32, dtype=torch.int32)
|
|
cm = CacheManager(32, 1, page_table, "radix")
|
|
|
|
a = _admit(cm, page_table, 0, PROMPT, cm.match_req(_pend(PROMPT)).cuda_handle)
|
|
b = _admit(cm, page_table, 1, PROMPT, cm.match_req(_pend(PROMPT)).cuda_handle)
|
|
assert _live_row(page_table, a).isdisjoint(_live_row(page_table, b))
|
|
|
|
with cm.lazy_free_region(): # the scheduler drains commits inside this region
|
|
cm.cache_req(a, finished=False)
|
|
cm.cache_req(b, finished=False)
|
|
|
|
free = set(cm.free_slots.tolist())
|
|
# b's own duplicate pages went back to the pool ...
|
|
assert free, "the later committer's duplicate pages should have been freed"
|
|
# ... and no page b still reads is on the free list.
|
|
assert _live_row(page_table, b).isdisjoint(free)
|
|
assert _live_row(page_table, b) == set(b.cache_handle.get_matched_indices().tolist())
|
|
|
|
|
|
def test_hybrid_unfinished_commit_repoints_the_row_off_the_freed_pages():
|
|
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,
|
|
)
|
|
pool = LinearStatePool(group=g, num_slots=16, dtype=torch.bfloat16,
|
|
device=torch.device("cpu"), tp_size=1)
|
|
page_table = torch.zeros(4, 32, dtype=torch.int32)
|
|
cm = CacheManager(32, 1, page_table, "hybrid_radix", linear_state_pool=pool)
|
|
|
|
reqs = []
|
|
for idx in (0, 1):
|
|
r = _admit(cm, page_table, idx, PROMPT, cm.match_req(_pend(PROMPT)).cuda_handle)
|
|
r.linear_slot_idx = pool.alloc(1)[0]
|
|
r.mamba_ping_pong = tuple(pool.alloc(2))
|
|
r.mamba_next_track_idx = 1
|
|
r.mamba_last_track_seqlen = len(PROMPT)
|
|
reqs.append(r)
|
|
|
|
with cm.lazy_free_region():
|
|
for r in reqs:
|
|
cm.cache_req(r, finished=False)
|
|
|
|
free = set(cm.free_slots.tolist())
|
|
assert free, "the later committer's duplicate pages should have been freed"
|
|
assert _live_row(page_table, reqs[1]).isdisjoint(free)
|
|
|
|
|
|
def test_lazy_free_snapshots_the_rows_it_was_handed():
|
|
"""The deferred free list must not be rewritten by a later re-point of the same row."""
|
|
page_table = torch.zeros(2, 8, dtype=torch.int32)
|
|
cm = CacheManager(8, 1, page_table, "radix")
|
|
page_table[0, :4] = torch.tensor([4, 5, 6, 7], dtype=torch.int32)
|
|
before = cm.free_slots.clone()
|
|
|
|
with cm.lazy_free_region():
|
|
cm._free(page_table[0, :4])
|
|
page_table[0, :4] = torch.tensor([0, 1, 2, 3], dtype=torch.int32) # a re-point
|
|
|
|
appended = cm.free_slots[len(before):].tolist()
|
|
assert appended == [4, 5, 6, 7]
|
|
|
|
|
|
def test_radix_subspan_commit_repoints_only_the_deduped_slice():
|
|
"""Same defect with old_handle.cached_len > 0: the committer admitted on top of a
|
|
published prefix, so the dedup free and the re-point cover only the sub-span
|
|
[old_cached, new_cached) -- the slice arithmetic the zero-prefix tests never touch."""
|
|
LONG = list(range(1, 17))
|
|
SHORT = LONG[:8]
|
|
page_table = torch.zeros(4, 32, dtype=torch.int32)
|
|
cm = CacheManager(32, 1, page_table, "radix")
|
|
|
|
seed = _admit(cm, page_table, 0, SHORT, cm.match_req(_pend(SHORT)).cuda_handle)
|
|
with cm.lazy_free_region():
|
|
cm.cache_req(seed, finished=False)
|
|
|
|
def _admit_on_prefix(table_idx):
|
|
m = cm.match_req(_pend(LONG))
|
|
matched = m.cuda_handle.cached_len
|
|
assert matched > 0, "the seeded prefix should match"
|
|
req = Req(input_ids=torch.tensor(LONG, dtype=torch.int32), table_idx=table_idx,
|
|
cached_len=matched, output_len=0, uid=table_idx,
|
|
sampling_params=SamplingParams(), cache_handle=m.cuda_handle)
|
|
req.device_len = len(LONG)
|
|
cm.lock(m.cuda_handle)
|
|
page_table[table_idx, :matched] = m.cuda_handle.get_matched_indices()[:matched]
|
|
cm.allocate_paged([req]) # only [matched, 16) -- the row prefix is canonical
|
|
req.cached_len = len(LONG)
|
|
return req, matched
|
|
|
|
b, matched = _admit_on_prefix(1)
|
|
d, _ = _admit_on_prefix(2)
|
|
own_suffix_d = set(page_table[2, matched:].tolist())
|
|
|
|
with cm.lazy_free_region():
|
|
cm.cache_req(b, finished=False) # b publishes [matched, 15)
|
|
cm.cache_req(d, finished=False) # d dedups against b: frees its own sub-span
|
|
|
|
free = set(cm.free_slots.tolist())
|
|
assert free & own_suffix_d, "d's duplicate sub-span pages should have been freed"
|
|
assert _live_row(page_table, d).isdisjoint(free)
|
|
canonical = d.cache_handle.get_matched_indices()
|
|
assert page_table[2, : d.cache_handle.cached_len].tolist() == canonical[: d.cache_handle.cached_len].tolist()
|
|
cm.check_integrity()
|