Xiaoze Fan
250 lines
12 KiB
Python
250 lines
12 KiB
Python
"""Two coupled properties of the prefill-boundary commit.
|
|
|
|
1. It must hand ``inc_lock`` ONE window. ``cache_req(finished=False)`` inserts the forwarded
|
|
prompt as a single suffix node and locks it for the rest of decode (the decode driver floors
|
|
its own frees at the committed length, so the request cannot reclaim it either). inc_lock is
|
|
node-granular, so without a node boundary a window back it pins the entire last chunk's swa --
|
|
and PrefillAdder sizes the chunk to fill the pool, so a long prompt hands its whole pool to
|
|
the lock and the next decode step raises "SWA pool exhausted", which nothing handles. Whether
|
|
it fires depends on the final chunk's length, i.e. on prompt % chunk.
|
|
|
|
2. The pool floor must then cover ``max_running_req`` of those footprints: admission gates only
|
|
the incoming chunk and never reserves the decode growth of the requests already running.
|
|
|
|
CPU-only, no model: real PrefillManager/PrefillAdder chunking (so the cap is the shipped one)
|
|
in the scheduler's overlap order, then the decode driver.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from freetoken.core import Context, Req, SamplingParams, get_global_ctx, set_global_ctx
|
|
from freetoken.distributed import set_tp_info, try_get_tp_info
|
|
from freetoken.kvcache.hybrid_swa_pool import (
|
|
_swa_paged_num_tokens,
|
|
_swa_per_req_swa_floor,
|
|
_swa_pool_floor,
|
|
)
|
|
from freetoken.models.config import KVCacheGroupSpec
|
|
from freetoken.scheduler import cache as cache_mod
|
|
from freetoken.scheduler.cache import CacheManager
|
|
from freetoken.scheduler.decode import DecodeManager
|
|
from freetoken.scheduler.prefill import ChunkedReq, PrefillManager
|
|
from freetoken.scheduler.table import TableManager
|
|
from freetoken.scheduler.utils import PendingReq
|
|
|
|
DEVICE = torch.device("cpu")
|
|
MAX_RUNNING = 4
|
|
UID = 7
|
|
TOKEN_BUDGET = 8192 # max_extend_tokens: never the binding cap here, the swa pool is
|
|
GAP = cache_mod._SWA_RETAIN_GAP
|
|
|
|
if try_get_tp_info() is None:
|
|
set_tp_info(rank=0, size=1)
|
|
|
|
|
|
def _cfg(window: int, page_size: int = 1, max_running_req: int = 4, **kw):
|
|
"""The slice of EngineConfig the floor helpers read."""
|
|
groups = (
|
|
KVCacheGroupSpec(name="full", layer_ids=(1,), num_kv_heads=1, head_dim=8,
|
|
sliding_window=None),
|
|
KVCacheGroupSpec(name="swa", layer_ids=(0,), num_kv_heads=1, head_dim=8,
|
|
sliding_window=window),
|
|
)
|
|
return SimpleNamespace(
|
|
page_size=page_size, max_running_req=max_running_req,
|
|
model_config=SimpleNamespace(kv_cache_group_specs=lambda: groups),
|
|
swa_num_pages_override=kw.get("override"), swa_full_tokens_ratio=kw.get("ratio", 0.2),
|
|
)
|
|
|
|
|
|
def _managers(window: int, num_swa_tokens: int, ps: int = 1, num_pages: int = 4096, width=2048):
|
|
from freetoken.kvcache.hybrid_swa_pool import HybridSWAKVCache
|
|
|
|
try:
|
|
get_global_ctx()
|
|
except AssertionError:
|
|
set_global_ctx(Context(page_size=ps))
|
|
|
|
pool = HybridSWAKVCache(
|
|
groups=_cfg(window, ps).model_config.kv_cache_group_specs(), num_layers=2,
|
|
num_full_pages=num_pages, page_size=ps, dtype=torch.bfloat16, device=DEVICE,
|
|
num_swa_tokens=num_swa_tokens,
|
|
)
|
|
pt = torch.zeros((MAX_RUNNING + 1, width), dtype=torch.int32, device=DEVICE)
|
|
cm = CacheManager(num_pages=num_pages, page_size=ps, page_table=pt, type="swa_radix",
|
|
swa_pool=pool, sliding_window_size=window)
|
|
assert cm.swa_paged and cm.is_swa
|
|
tm = TableManager(max_running_reqs=MAX_RUNNING, page_table=pt)
|
|
return cm, tm, PrefillManager(cm, tm, DecodeManager(page_size=ps))
|
|
|
|
|
|
def _prefill(cm, pm, prompt_len: int, n_decode: int, base: int = 1):
|
|
"""Real chunked prefill (chunk sizes from PrefillAdder) in the scheduler's overlap order:
|
|
schedule+forward chunk N+1 before committing chunk N; only the final chunk commits.
|
|
``base`` shifts the token ids so concurrent requests share no prefix."""
|
|
pm.pending_list = [PendingReq(uid=UID,
|
|
input_ids=torch.arange(base, base + prompt_len,
|
|
dtype=torch.int32),
|
|
sampling_params=SamplingParams(max_tokens=n_decode))]
|
|
final = last_batch = None
|
|
while pm.runnable or last_batch is not None:
|
|
batch = pm.schedule_next_batch(TOKEN_BUDGET)
|
|
if batch is not None:
|
|
assert batch.reqs[0].extend_len > 0, "prefill stalled at a zero-length chunk"
|
|
cm.free_swa_out_of_window_extend(batch.reqs) # _prepare_batch
|
|
cm.allocate_paged(batch.reqs)
|
|
for r in batch.reqs:
|
|
r.complete_one()
|
|
if last_batch is not None:
|
|
for r in last_batch.reqs:
|
|
if not isinstance(r, ChunkedReq):
|
|
cm.cache_req(r, finished=False) # scheduler.py:328
|
|
final = r
|
|
last_batch = batch
|
|
return final
|
|
|
|
|
|
def _decode(cm, reqs, n: int):
|
|
"""Shared decode loop over the scheduler's global forward counter, which is what drives the
|
|
out-of-window eviction cadence."""
|
|
for i in range(n):
|
|
for r in reqs:
|
|
r.append_host(torch.tensor([9999], dtype=torch.int32))
|
|
r.decode_batch_idx = i + 1
|
|
cm.maybe_free_swa_out_of_window(reqs, forward_iter=i + 1)
|
|
cm.allocate_paged(reqs)
|
|
for r in reqs:
|
|
r.complete_one()
|
|
|
|
|
|
# --------------------------------------------------------------- what the commit locks
|
|
@pytest.mark.parametrize("ps", [1, 8, 128])
|
|
def test_commit_locks_one_window_not_the_whole_extend(ps, monkeypatch):
|
|
"""The lock covers the retained window (page-rounded); the head stays live and unlocked, so
|
|
ambient pressure can still reclaim it. ps == window == 128 is the DSV4 shape, where the
|
|
boundary must also stay page-aligned for the pool's whole-page free path."""
|
|
monkeypatch.setattr(cache_mod, "_SWA_EVICTION_INTERVAL", 1)
|
|
window, prompt = (8, 200) if ps == 1 else (ps, 24 * ps)
|
|
cm, _tm, pm = _managers(window, num_swa_tokens=64 * max(ps, 8) + 1, ps=ps,
|
|
num_pages=256 if ps > 1 else 4096, width=64 * max(ps, 32))
|
|
_prefill(cm, pm, prompt, n_decode=1)
|
|
|
|
retained = -(-(window + GAP) // ps) * ps
|
|
assert cm.prefix_cache.swa_protected == retained
|
|
assert cm.prefix_cache.swa_evictable == prompt - retained
|
|
cm.prefix_cache.check_integrity()
|
|
|
|
|
|
def test_short_final_chunk_locks_no_more_than_a_window(monkeypatch):
|
|
"""The boundary is the retain gap, not the window: the committed live region is
|
|
[L - c_last - window - 1, L), so keep_from (L - window - gap) only falls inside it once
|
|
c_last > gap. Below that the split lands in the tombstoned head and does nothing -- and does
|
|
not need to, the node is already short. Either way the lock stays <= window + gap."""
|
|
monkeypatch.setattr(cache_mod, "_SWA_EVICTION_INTERVAL", 1)
|
|
window, first_chunk, n_decode = 8, 100, 40
|
|
for c_last in (1, 14, 15, 16, 17, 60):
|
|
cm, tm, _pm = _managers(window, num_swa_tokens=4096)
|
|
total = first_chunk + c_last
|
|
req = Req(input_ids=torch.arange(1, total + 1, dtype=torch.int32), table_idx=0,
|
|
cached_len=0, output_len=n_decode, uid=UID,
|
|
sampling_params=SamplingParams(), cache_handle=None)
|
|
req.input_len = total
|
|
h = cm.match_req(req).cuda_handle
|
|
req.cache_handle = h
|
|
cm.lock(h)
|
|
for end in (first_chunk, total): # two explicit chunks -> c_last is exact
|
|
req.device_len = end
|
|
cm.free_swa_out_of_window_extend([req])
|
|
cm.allocate_paged([req])
|
|
req.cached_len = end
|
|
req.device_len = end + 1
|
|
live_node = total - req.swa_evicted_seqlen
|
|
cm.cache_req(req, finished=False)
|
|
|
|
assert cm.prefix_cache.swa_protected == min(live_node, window + GAP), f"c_last={c_last}"
|
|
_decode(cm, [req], n_decode)
|
|
cm.cache_req(req, finished=True)
|
|
tm.free(req.table_idx)
|
|
cm.check_integrity()
|
|
|
|
|
|
# --------------------------------------------------------------- what that buys at runtime
|
|
@pytest.mark.parametrize("interval", [1, 128])
|
|
def test_decode_survives_any_prompt_length(interval, monkeypatch):
|
|
"""The pool holds the window working set and nothing chunk-sized, at both an every-forward
|
|
cadence and the shipped 128-forward one (where it must also absorb a full interval of decode
|
|
growth before the first reclaim). The prompt length -- hence the final chunk's -- must not
|
|
decide whether decode completes."""
|
|
monkeypatch.setattr(cache_mod, "_SWA_EVICTION_INTERVAL", interval)
|
|
window, n_decode = 8, 300
|
|
pool = 2 * window + GAP + interval + 8
|
|
for prompt in (103, 129, 200, 400):
|
|
cm, tm, pm = _managers(window, num_swa_tokens=pool + 1)
|
|
req = _prefill(cm, pm, prompt, n_decode)
|
|
_decode(cm, [req], n_decode)
|
|
cm.cache_req(req, finished=True)
|
|
tm.free(req.table_idx)
|
|
cm.check_integrity()
|
|
|
|
|
|
@pytest.mark.parametrize("n_req", [1, 3])
|
|
def test_a_full_batch_decodes_on_a_pool_sized_to_the_declared_floor(n_req, monkeypatch):
|
|
"""Ties the sizing formula to the runtime: open the pool at exactly _swa_pool_floor and run
|
|
max_running_req requests through prefill into a shared decode. All are committed before any
|
|
decode, so every one holds its full non-evictable footprint at once."""
|
|
monkeypatch.setattr(cache_mod, "_SWA_EVICTION_INTERVAL", 1)
|
|
window, n_decode = 8, 200
|
|
floor = _swa_pool_floor(_cfg(window, max_running_req=n_req))
|
|
cm, tm, pm = _managers(window, num_swa_tokens=floor + 1)
|
|
reqs = [_prefill(cm, pm, 120 + 7 * i, n_decode, base=1 + 10_000 * i) for i in range(n_req)]
|
|
_decode(cm, reqs, n_decode)
|
|
for r in reqs:
|
|
cm.cache_req(r, finished=True)
|
|
tm.free(r.table_idx)
|
|
cm.check_integrity()
|
|
|
|
|
|
def test_a_per_request_sized_pool_cannot_hold_a_full_batch(monkeypatch):
|
|
"""Why the floor carries the max_running_req factor. Two manifestations, both fatal and
|
|
neither handled: alloc_swa raising once the pool is dry, or -- reached first here --
|
|
PrefillAdder's swa cap collapsing a continuation chunk to zero tokens (prefill.py:138),
|
|
which builds a Req with device_len == cached_len. Continuations take the chunked branch of
|
|
try_add_one, so the admission gate never sees them."""
|
|
monkeypatch.setattr(cache_mod, "_SWA_EVICTION_INTERVAL", 1)
|
|
window, n_decode = 8, 200
|
|
per_req = _swa_per_req_swa_floor(_cfg(window, max_running_req=3))
|
|
cm, _tm, pm = _managers(window, num_swa_tokens=per_req + 1)
|
|
with pytest.raises((RuntimeError, AssertionError)):
|
|
reqs = [_prefill(cm, pm, 120, n_decode, base=1 + 10_000 * i) for i in range(3)]
|
|
_decode(cm, reqs, n_decode)
|
|
|
|
|
|
# --------------------------------------------------------------- the floor formula
|
|
def test_floor_terms_and_page_rounding():
|
|
for window in (8, 128, 1024):
|
|
for ps in (1, 8, 64, 128):
|
|
locked = -(-(window + GAP) // ps) * ps # what the commit locks
|
|
tail = window + 2 * ps + cache_mod._SWA_EVICTION_INTERVAL # uncollected decode tail
|
|
assert _swa_per_req_swa_floor(_cfg(window, ps)) == locked + tail, (window, ps)
|
|
|
|
|
|
def test_pool_floor_scales_with_max_running_req():
|
|
for n in (1, 2, 4, 8, 32):
|
|
cfg = _cfg(128, max_running_req=n)
|
|
assert _swa_pool_floor(cfg) == n * _swa_per_req_swa_floor(cfg)
|
|
|
|
|
|
def test_ratio_and_override_never_dip_below_the_floor():
|
|
cfg = _cfg(128, max_running_req=4)
|
|
# Small full pool (ratio x full << floor) -> the floor, +1 slot-0 sentinel.
|
|
assert _swa_paged_num_tokens(cfg, num_full_pages=64) == _swa_pool_floor(cfg) + 1
|
|
# Generous full pool -> the ratio wins.
|
|
assert _swa_paged_num_tokens(cfg, num_full_pages=1 << 20) == int(0.2 * (1 << 20)) + 1
|
|
# A pinned window (the rebuild path only validates num_swa_pages > 0) is clamped up.
|
|
pinned = _cfg(128, max_running_req=4, override=1)
|
|
assert _swa_paged_num_tokens(pinned, num_full_pages=1024) == _swa_pool_floor(pinned) + 1
|