feat: initial open-source release

Co-authored-by: Shuo Yang <andy_yang@berkeley.edu>

tests/scheduler/test_scheduler_kv_usage.py

@@ -0,0 +1,94 @@
+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