Xiaoze Fan
95 lines
3.6 KiB
Python
95 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
|
|
|
|
def test_kv_usage_pages_excludes_evictable_prefix_cache():
|
|
# page_usage now lives on the CacheManagerLike interface (polymorphic vs DSV4); test the
|
|
# generic formula there. The unbound method works on a duck-typed namespace.
|
|
from freetoken.scheduler.cache import CacheManager
|
|
|
|
cache_manager = SimpleNamespace(
|
|
num_pages=100,
|
|
page_size=4,
|
|
free_slots=[0] * 20,
|
|
is_hybrid=False,
|
|
is_swa=False,
|
|
prefix_cache=SimpleNamespace(
|
|
size_info=SimpleNamespace(evictable_size=120, protected_size=40)
|
|
),
|
|
)
|
|
|
|
used_pages, total_pages = CacheManager.page_usage(cache_manager)
|
|
|
|
assert used_pages == 50
|
|
assert total_pages == 100
|
|
|
|
|
|
def test_swa_token_usage_counts_window_pool():
|
|
from freetoken.scheduler.scheduler import Scheduler
|
|
|
|
# swa_paged model: 129-token pool -> 128 allocatable (sentinel excluded from total);
|
|
# 64 available (free + evictable tree) -> 64 used.
|
|
cache_manager = SimpleNamespace(
|
|
swa_paged=True,
|
|
swa_pool=SimpleNamespace(swa_num_tokens=129),
|
|
swa_available_size=64,
|
|
)
|
|
assert Scheduler._swa_token_usage(SimpleNamespace(cache_manager=cache_manager)) == (64, 128)
|
|
|
|
# non-SWA models report None (no swa field in logs/stats).
|
|
plain = SimpleNamespace(cache_manager=SimpleNamespace(swa_paged=False))
|
|
assert Scheduler._swa_token_usage(plain) is None
|
|
|
|
|
|
def _fake_engine(swa=True, moe=True, mamba=True):
|
|
import torch
|
|
|
|
# Superset fake engine driving the REAL compute_cache_pools/compute_cache_unit_bytes:
|
|
# KV 64 pages x 16 tokens x 1 MiB/token = 1.00 GiB; swa 512 tokens x 2 MiB = 1.00 GiB;
|
|
# mamba 64 slots x 16 MiB = 1.00 GiB; MoE 24 slots of 8x16 experts, 4 KiB/slot.
|
|
return SimpleNamespace(
|
|
num_pages=64,
|
|
config=SimpleNamespace(
|
|
page_size=16,
|
|
cache_type="swa_radix",
|
|
model_config=SimpleNamespace(dsv4_args=None, has_swa_attention=swa),
|
|
),
|
|
kv_cache=SimpleNamespace(
|
|
swa_num_tokens=513, unit_bytes=lambda: (1 << 20, 1 << 21)
|
|
),
|
|
moe_offload_cache=SimpleNamespace(
|
|
cache_size=24, num_layers=8, num_experts=16,
|
|
bank_caches={"w": torch.zeros((24, 1024), dtype=torch.float32)},
|
|
) if moe else None,
|
|
linear_state_pool=SimpleNamespace(num_slots=65, bytes_per_slot=lambda: 1 << 24)
|
|
if mamba else None,
|
|
)
|
|
|
|
|
|
def test_log_cache_geometry_reports_all_pools(monkeypatch):
|
|
import freetoken.scheduler.scheduler as sched_mod
|
|
from freetoken.scheduler.scheduler import Scheduler
|
|
|
|
lines: list[str] = []
|
|
monkeypatch.setattr(sched_mod.logger, "info_rank0", lines.append)
|
|
Scheduler._log_cache_geometry(SimpleNamespace(engine=_fake_engine()), "Cache rebuilt")
|
|
line = lines[-1]
|
|
assert "Cache rebuilt: KV 64 pages (1024 tokens, 1.00 GiB)" in line
|
|
assert "swa 512 pages (512 tokens, 1.00 GiB)" in line
|
|
assert "mamba 64 slots (1.00 GiB)" in line
|
|
assert "MoE cache 24/128 (0.00 GiB)" in line
|
|
|
|
|
|
def test_log_cache_geometry_plain_model_kv_only(monkeypatch):
|
|
import freetoken.scheduler.scheduler as sched_mod
|
|
from freetoken.scheduler.scheduler import Scheduler
|
|
|
|
lines: list[str] = []
|
|
monkeypatch.setattr(sched_mod.logger, "info_rank0", lines.append)
|
|
engine = _fake_engine(swa=False, moe=False, mamba=False)
|
|
Scheduler._log_cache_geometry(SimpleNamespace(engine=engine), "Cache rebuilt")
|
|
line = lines[-1]
|
|
assert "Cache rebuilt: KV 64 pages (1024 tokens, 1.00 GiB)" in line
|
|
assert "swa" not in line and "mamba" not in line and "MoE" not in line
|