prain: RX 6600 XT example config

examples/gemma4-26b-a4b-rx6600xt.toml

@@ -0,0 +1,59 @@
+# prain placement: prithivMLmods/gemma-4-26B-A4B-it-Uncensored-MAX on an 8 GB RX 6600 XT
+# Model: 30 MoE layers x 128 experts, top-8 routing, 11.34 MiB per expert (BF16).
+# VRAM: 8 GiB - 4.45 GiB dense weights - 0.75 GiB KV cache - ~0.9 GiB runtime = ~1.7 GiB experts.
+#
+# Tiers:  vram = pinned in a VRAM slot for the whole run
+#         ram  = pinned host memory, copied into a cache slot when a token needs it
+#         disk = read from the pack file in pack_dir when a token needs it
+#                (the kernel page cache still caches it; nothing is pinned)
+# [layers."<selector>"] sections apply in file order; a later section wins per expert.
+# Layer selectors: "*", "12", "0-9,20".   Expert lists: "*", "0-15,32", [1, 2, 3].
+# Check what a file costs:   prain check <this file> <model config.json>
+# Compare variants offline:  prain sim <trace> <this file> <model config.json>
+#                            prain tune <trace> <model config.json> --expert-vram-gib 1.7 ...
+# Relative paths start at this file's directory; "~" works too.
+
+pack_dir = "../models/packs/gemma4-26b-a4b"  # 42.5 GiB, written on the first run
+default_tier = "disk"                    # runs on any RAM size; `prain tune` does better
+cache_slots = 5                          # 30 x 5 x 11.34 MiB = 1.66 GiB of VRAM
+stats_path = "../models/packs/gemma4-26b-a4b/stats.json"
+stats_every = 64                         # MoE forward passes between stats/trace writes
+trace_path = "../models/packs/gemma4-26b-a4b/trace.jsonl.gz"  # routing trace for sim/tune
+
+[policy]              # which expert a cache slot gives up
+eviction = "lfu"      # lru | lfu (decayed frequency) | gdsf (frequency x refetch cost)
+heat = "weight"       # what makes an expert hot: count (tokens) | weight (router mass) |
+                      # recency (router mass, later tokens weigh more: good for warm-ups)
+lfu_half_life = 64    # forwards until an LFU score halves
+disk_cost = 4.0       # gdsf: a disk expert costs this many RAM experts to refetch
+recency_half_life = 32  # heat = "recency": tokens until a token's weight halves
+scan_tokens = 64      # a forward with this many tokens (a prompt) is a scan:
+scan_keep = 0         #   it may not evict the scan_keep most valuable cached experts,
+promote = 2           #   and only its `promote` hottest experts get full priority
+
+[prefetch]            # copy experts before their layer runs, on a side stream
+mode = "none"         # none | pregate (run layer L+distance's router on layer L's input)
+distance = 1          # 1 or 2 layers ahead
+budget = 2            # experts per layer and forward
+min_prob = 0.05       # skip guesses below this router probability
+
+[miss]                # an expert that is not in VRAM
+mode = "copy"         # copy | cpu (always compute it on the CPU) | auto (CPU if it has
+cpu_max_tokens = 1    #   at most cpu_max_tokens tokens, else copy)
+cpu_threads = 4       # CPU threads of the vLLM worker (sets OMP_NUM_THREADS)
+admit = 1             # cpu/auto: copy this many CPU-run experts into VRAM afterwards
+
+[serve]               # parallel requests (`prain serve` passes these to vLLM)
+max_requests = 2      # requests vLLM runs at once (--max-num-seqs); warm-ups count too
+kv_cache_gib = 0.75   # KV cache VRAM shared by them (--kv-cache-memory-bytes). vLLM logs
+                      #   "Maximum concurrency for 4096 tokens per request: N.NNx" at start;
+                      #   below max_requests, long requests wait or get preempted. Every GiB
+                      #   here is 90 expert cache slots less.
+
+[draft]               # answers that arrive while others use the GPU (`prain web --config`)
+full_requests = 1     # answers that run with all their experts; later ones run at once as
+                      #   drafts on the experts already in VRAM (= max_requests: never draft)
+min_weight = 0.5      # a draft token still copies its heaviest missing experts until this
+                      #   share of its router weight is computed (0 = only what is in VRAM)
+refine = "regenerate" # regenerate: once a full slot frees, the draft is cut short and the
+                      #   answer generated again with all experts | none: keep the draft

prain/chat.py

@@ -0,0 +1,119 @@
+"""OpenAI-compatible client pieces shared by the TUI and the WebUI. Stdlib only."""
+
+import http.client
+import json
+import socket
+import threading
+import urllib.parse
+import urllib.request
+
+ABORTED = (OSError, ValueError, http.client.HTTPException)  # what a cancelled request raises
+
+
+def post(base_url, path, body, timeout=600):
+    request = urllib.request.Request(f"{base_url}{path}", json.dumps(body).encode(),
+                                     {"Content-Type": "application/json"})
+    return urllib.request.urlopen(request, timeout=timeout)
+
+
+def send_request(base_url, path, body, timeout=600):
+    """POST without waiting for the answer. Returns the connection: call getresponse() to
+    read it, or shutdown its socket from another thread to abort it (vLLM then cancels)."""
+    url = urllib.parse.urlsplit(base_url)
+    connection = http.client.HTTPConnection(url.hostname, url.port or 80, timeout=timeout)
+    connection.request("POST", path, json.dumps(body), {"Content-Type": "application/json"})
+    return connection
+
+
+def first_model(base_url):
+    with urllib.request.urlopen(f"{base_url}/v1/models", timeout=30) as response:
+        return json.load(response)["data"][0]["id"]
+
+
+def chat_body(model, messages, **extra):
+    return {"model": model, "messages": messages, "stream": True, **extra}
+
+
+def draft_body(model, messages):
+    """An answer computed from the experts already in VRAM: prain.drafts marks the request
+    by its vllm_xargs. Its own cache_salt keeps its KV blocks, computed with skipped experts,
+    out of the prefix-cache hits of full requests."""
+    return chat_body(model, messages, vllm_xargs={"prain": "draft"}, cache_salt="prain-draft")
+
+
+def warm_body(model, messages, draft, priority=0):
+    """Prefill the half-typed message and stop after one token. vLLM keeps its KV blocks
+    (prefix caching) and prain leaves the experts it routed to most in VRAM, so the real
+    request only prefills what was typed after this. Streamed, so aborting it is cheap.
+    priority > 0 queues it behind real requests (needs --scheduling-policy priority)."""
+    body = {"model": model, "messages": [*messages, {"role": "user", "content": draft.strip()}],
+            "max_tokens": 1, "temperature": 0, "stream": True,
+            "add_generation_prompt": False, "continue_final_message": True}
+    if priority:
+        body["priority"] = priority
+    return body
+
+
+def stream_text(response):
+    """Yield the content pieces of a streamed chat completion (server-sent events)."""
+    for raw in response:
+        line = raw.decode().strip()
+        if not line.startswith("data: ") or line == "data: [DONE]":
+            continue
+        piece = json.loads(line[6:])["choices"][0].get("delta", {}).get("content")
+        if piece:
+            yield piece
+
+
+class Warmer:
+    """Warm-up requests of one chat session. At most one is in flight; while it runs, newer
+    drafts replace each other and only the newest is sent next. cancel() drops the queued
+    draft and aborts the one in flight, so a sent message never waits behind a warm-up."""
+
+    def __init__(self, start):
+        self.start = start  # start(job) -> connection from send_request()
+        self.lock = threading.Lock()
+        self.pending = None
+        self.busy = False
+        self.connection = None
+        self.sent = 0
+
+    def update(self, job):
+        with self.lock:
+            self.pending = job
+            if self.busy:
+                return
+            self.busy = True
+        threading.Thread(target=self._run, daemon=True).start()
+
+    def cancel(self):
+        with self.lock:
+            self.pending = None
+            connection = self.connection
+        if connection is not None and connection.sock is not None:
+            try:
+                connection.sock.shutdown(socket.SHUT_RDWR)  # wakes the blocked reader
+            except OSError:
+                pass
+
+    def _run(self):
+        while True:
+            with self.lock:
+                job, self.pending = self.pending, None
+                if job is None:
+                    self.busy = False
+                    return
+            try:
+                connection = self.start(job)
+                with self.lock:
+                    self.connection = connection
+                    self.sent += 1
+                with connection.getresponse() as response:
+                    response.read()
+            except ABORTED:
+                pass  # best effort: the real request reports server errors
+            finally:
+                with self.lock:
+                    connection, self.connection = self.connection, None
+                if connection is not None:
+                    connection.close()

prain/cli.py

@@ -0,0 +1,198 @@
+"""prain command line: check, plan or tune a placement, simulate it, serve with vLLM, and
+chat in a TUI or the browser."""
+
+import argparse
+import json
+import os
+import sys
+from dataclasses import asdict
+
+from prain.config import (GIB, Miss, Policy, Prefetch, expert_bytes, format_ids, load_config,
+                            moe_dims, plan, render_toml)
+
+
+def _model_dims(hf_config_path):
+    with open(hf_config_path) as f:
+        num_layers, num_experts, top_k, hidden, inter = moe_dims(json.load(f))
+    return num_layers, num_experts, top_k, expert_bytes(hidden, inter)
+
+
+def check(config_path, hf_config_path):
+    config = load_config(config_path)
+    num_layers, num_experts, top_k, size = _model_dims(hf_config_path)
+    print("layer  vram   ram  disk  cache_slots")
+    vram_rows = ram_rows = few_slots = 0
+    for layer in range(num_layers):
+        p = config.placement(layer, num_experts)
+        vram_rows += p.num_slots
+        ram_rows += len(p.ram)
+        few_slots += bool((p.ram or p.disk) and p.cache_slots < top_k)
+        print(f"{layer:5}  {len(p.vram):4}  {len(p.ram):4}  {len(p.disk):4}  {p.cache_slots:11}")
+    if few_slots:
+        print(f"warning: {few_slots} layers have fewer cache slots than top_k={top_k}; "
+              "their decode steps run in several rounds", file=sys.stderr)
+    print(f"expert size {size / 2**20:.2f} MiB")
+    print(f"VRAM {vram_rows * size / GIB:.2f} GiB  (pinned experts + cache slots)")
+    print(f"KV   {config.serve.kv_cache_gib:.2f} GiB  (shared by up to "
+          f"{config.serve.max_requests} parallel requests)")
+    print(f"RAM  {ram_rows * size / GIB:.2f} GiB  (pinned host memory)")
+    print(f"disk {num_layers * num_experts * size / GIB:.2f} GiB  (pack in {config.pack_dir})")
+
+
+def _variant_text(variant):
+    policy, prefetch, miss = variant
+    fetch = "none" if prefetch.mode == "none" else f"pregate d{prefetch.distance} b{prefetch.budget}"
+    cpu = miss.mode if miss.mode == "copy" else f"{miss.mode} admit{miss.admit}"
+    return (f"{policy.eviction:5} {policy.heat:8} keep{policy.scan_keep} "
+            f"promote{policy.promote}  {fetch:15} {cpu}")
+
+
+def sim(trace, config_path, hf_config_path, cost_path, expert_mib):
+    from prain.sim import describe, load_cost, read_forwards, simulate
+
+    config = load_config(config_path)
+    _, num_experts, _, size = _model_dims(hf_config_path)
+    forwards = read_forwards(trace)
+    cost = load_cost(cost_path, expert_mib or size / 2**20)
+    layers = sorted({layer for fwd in forwards for layer in fwd})
+    placements = {layer: config.placement(layer, num_experts) for layer in layers}
+    variant = (config.policy, config.prefetch, config.miss)
+    result = simulate(forwards, placements, *variant, cost)
+    bound = simulate(forwards, placements, Policy(scan_tokens=config.policy.scan_tokens),
+                     Prefetch(), Miss(), cost, belady=True)
+    print("\n".join(describe(result, bound, cost)))
+
+
+def tune_command(args):
+    from prain.sim import describe, load_cost, read_forwards, tune
+
+    _, num_experts, _, size = _model_dims(args.hf_config)
+    cost = load_cost(args.cost, args.expert_mib or size / 2**20)
+    forwards = read_forwards(args.trace)
+    best, ram, per, final, bound, rows = tune(forwards, num_experts, args.expert_vram_gib,
+                                              args.ram_gib, cost)
+    print(" rank  total_s  decode_ms  hit%  eviction heat keep promote  prefetch  miss",
+          file=sys.stderr)
+    for rank, (variant, r) in enumerate(rows[: args.top], 1):
+        print(f"{rank:5}  {r.total_ms / 1000:7.2f}  {r.ms_per_decode:9.1f}  {100 * r.hit_rate:4.1f}"
+              f"  {_variant_text(variant)}", file=sys.stderr)
+    print("best, with per-layer cache slots:", file=sys.stderr)
+    print("\n".join("  " + line for line in describe(final, bound, cost)), file=sys.stderr)
+    policy, prefetch, miss = best
+    layers = {layer: {k: v for k, v in (("ram", format_ids(ram[layer])),
+                                         ("cache_slots", per[layer])) if v != ""}
+              for layer in per}
+    print(render_toml(
+        f"generated by `prain tune` from {os.path.basename(args.trace)}",
+        {"pack_dir": os.path.abspath(args.pack_dir), "default_tier": "disk",
+         "cache_slots": min(per.values()), "trace_path": os.path.abspath(args.trace)},
+        {"policy": asdict(policy), "prefetch": asdict(prefetch), "miss": asdict(miss)},
+        layers,
+    ), end="")
+
+
+SERVE_OWNED = ("--enforce-eager", "--moe-backend", "--scheduling-policy", "--max-num-seqs",
+               "--kv-cache-memory-bytes")  # set by `prain serve` from the placement file
+
+
+def serve(args):
+    config = load_config(args.config)  # fail now, not after vLLM spent minutes loading
+    extra = args.vllm_args[1:] if args.vllm_args[:1] == ["--"] else args.vllm_args
+    owned = [a for a in extra if a.split("=")[0] in SERVE_OWNED]
+    if owned:
+        raise SystemExit(f"prain serve: {' '.join(owned)} comes from the placement file "
+                         "([serve] max_requests, kv_cache_gib); remove it after --")
+    os.environ["PRAIN_CONFIG"] = os.path.abspath(args.config)
+    if config.miss.mode != "copy":  # vLLM pins torch to 1 CPU thread unless this is set
+        os.environ.setdefault("OMP_NUM_THREADS", str(config.miss.cpu_threads))
+    os.execvp("vllm", ["vllm", "serve", args.model, "--enforce-eager", "--moe-backend", "triton",
+                       "--scheduling-policy", "priority",
+                       "--max-num-seqs", str(config.serve.max_requests),
+                       "--kv-cache-memory-bytes", str(int(config.serve.kv_cache_gib * GIB)),
+                       *extra])
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(prog="prain")
+    sub = parser.add_subparsers(dest="command", required=True)
+    c = sub.add_parser("check", help="validate a placement file and print what it costs")
+    c.add_argument("config")
+    c.add_argument("hf_config", help="the model's config.json")
+    p = sub.add_parser("plan", help="print a placement file built from routing stats")
+    p.add_argument("stats", help="stats_path written by a prain run")
+    p.add_argument("hf_config", help="the model's config.json")
+    p.add_argument("--expert-vram-gib", type=float, required=True,
+                   help="VRAM for experts: pinned ones plus cache slots")
+    p.add_argument("--ram-gib", type=float, required=True, help="pinned host memory for experts")
+    p.add_argument("--cache-slots", type=int, default=16)
+    p.add_argument("--pack-dir", required=True)
+    for name, text in (("sim", "simulate a placement file on a routing trace"),
+                       ("tune", "search policies and per-layer slots on a trace; prints a placement file")):
+        s = sub.add_parser(name, help=text)
+        s.add_argument("trace", help="trace_path written by a prain run, or `prain synth`")
+        if name == "sim":
+            s.add_argument("config")
+        s.add_argument("hf_config", help="the model's config.json")
+        s.add_argument("--cost", help="JSON from scripts/measure.py (default: RX 6600 XT guesses)")
+        s.add_argument("--expert-mib", type=float, help="what-if: expert size, e.g. 2.9 for int4")
+        if name == "tune":
+            s.add_argument("--expert-vram-gib", type=float, required=True)
+            s.add_argument("--ram-gib", type=float, required=True)
+            s.add_argument("--pack-dir", required=True)
+            s.add_argument("--top", type=int, default=10, help="variants to list")
+    y = sub.add_parser("synth", help="write a synthetic routing trace, to try sim and tune")
+    y.add_argument("out")
+    for flag, default in (("forwards", 300), ("layers", 30), ("experts", 128), ("top-k", 8),
+                          ("seed", 0)):
+        y.add_argument(f"--{flag}", type=int, default=default)
+    s = sub.add_parser("serve", help="exec `vllm serve` with the placement active")
+    s.add_argument("config")
+    s.add_argument("model")
+    s.add_argument("vllm_args", nargs=argparse.REMAINDER, help="extra vllm serve args after --")
+    for name, what in (("tui", "the terminal"), ("web", "the browser")):
+        ui = sub.add_parser(name, help=f"chat in {what}; experts warm up while you type")
+        ui.add_argument("--url", default="http://127.0.0.1:8000", help="the vLLM server")
+        ui.add_argument("--warmup", choices=("idle", "always", "off"), default="idle",
+                        help="idle: not while an answer streams (default)")
+        ui.add_argument("--warmup-priority", type=int, default=1,
+                        help="vLLM priority of warm-ups; 0 for a server without "
+                             "--scheduling-policy priority")
+        if name == "web":
+            ui.add_argument("--port", type=int, default=8080)
+            ui.add_argument("--config", help="placement file: [serve] max_requests and the "
+                                             "[draft] table (default: their defaults)")
+    args = parser.parse_args(argv)
+
+    if args.command == "check":
+        check(args.config, args.hf_config)
+    elif args.command == "plan":
+        _, num_experts, _, size = _model_dims(args.hf_config)
+        with open(args.stats) as f:
+            stats = json.load(f)
+        print(plan(stats, num_experts, size, args.expert_vram_gib, args.ram_gib,
+                   args.cache_slots, os.path.abspath(args.pack_dir), os.path.abspath(args.stats)), end="")
+    elif args.command == "sim":
+        sim(args.trace, args.config, args.hf_config, args.cost, args.expert_mib)
+    elif args.command == "tune":
+        tune_command(args)
+    elif args.command == "synth":
+        from prain.sim import synthetic_trace, write_trace
+
+        write_trace(args.out, synthetic_trace(args.layers, args.experts, args.top_k,
+                                              args.forwards, seed=args.seed))
+    elif args.command == "serve":
+        serve(args)
+    elif args.command == "tui":
+        from prain.tui import main as tui  # curses only loads when needed
+
+        tui(args.url, args.warmup, args.warmup_priority)
+    else:
+        from prain.config import Config
+        from prain.web import serve as web
+
+        config = load_config(args.config) if args.config else Config(pack_dir="")  # defaults
+        web(args.url, args.port, args.warmup, args.warmup_priority, config.serve, config.draft)
+
+
+if __name__ == "__main__":
+    main()

prain/config.py

@@ -57,3 +57,248 @@ def layer_index(name):
     """'language_model.model.layers.12.experts' -> 12, or None unless exactly one number."""
     numbers = [int(part) for part in name.split(".") if part.isdigit()]
     return numbers[0] if len(numbers) == 1 else None
+
+
+@dataclass(frozen=True)
+class LayerPlacement:
+    vram: tuple  # pinned experts: slot i holds vram[i] for the whole run
+    ram: tuple  # pinned host memory, copied into a cache slot on demand
+    disk: tuple  # read from the pack file on demand
+    cache_slots: int  # VRAM slots for ram/disk experts, managed by [policy]
+
+    @property
+    def num_slots(self):
+        return len(self.vram) + self.cache_slots
+
+
+def _opt(default, choices=None, minimum=0, maximum=None):
+    return field(default=default, metadata={"choices": choices, "min": minimum, "max": maximum})
+
+
+@dataclass(frozen=True)
+class Policy:
+    """[policy]: how the VRAM cache slots of every layer are managed."""
+    eviction: str = _opt("lfu", ("lru", "lfu", "gdsf"))
+    heat: str = _opt("weight", ("count", "weight", "recency"))
+    lfu_half_life: float = _opt(64.0, minimum=1)  # forwards until an LFU score halves
+    disk_cost: float = _opt(4.0, minimum=1)  # gdsf: refetching a disk expert vs. a RAM one
+    recency_half_life: float = _opt(32.0, minimum=1)  # heat "recency": tokens per halving
+    scan_tokens: int = _opt(64, minimum=1)  # forwards with this many tokens are scans
+    scan_keep: int = _opt(0)  # cache slots a scan may not evict
+    promote: int = _opt(2)  # experts a scan admits at full priority
+
+
+@dataclass(frozen=True)
+class Prefetch:
+    """[prefetch]: copy experts before their layer runs."""
+    mode: str = _opt("none", ("none", "pregate"))
+    distance: int = _opt(1, minimum=1, maximum=2)  # layers ahead
+    budget: int = _opt(2)  # experts per layer and forward
+    min_prob: float = _opt(0.05, maximum=1)  # ignore predictions below this router probability
+
+
+@dataclass(frozen=True)
+class Miss:
+    """[miss]: what happens to an expert that is not in VRAM."""
+    mode: str = _opt("copy", ("copy", "cpu", "auto"))
+    cpu_max_tokens: int = _opt(1, minimum=1)  # auto: CPU runs experts with <= this many tokens
+    cpu_threads: int = _opt(4, minimum=1)  # OMP threads of the vLLM worker
+    admit: int = _opt(1)  # CPU-run experts copied into VRAM afterwards, hottest first
+
+
+@dataclass(frozen=True)
+class Serve:
+    """[serve]: the vLLM settings prain and its clients both need to know."""
+    max_requests: int = _opt(2, minimum=1)  # requests vLLM runs at once (--max-num-seqs)
+    kv_cache_gib: float = _opt(0.75, minimum=0.01)  # KV cache, shared by all requests
+
+
+@dataclass(frozen=True)
+class Draft:
+    """[draft]: answers that arrive while other answers already use the GPU."""
+    full_requests: int = _opt(1, minimum=1)  # answers that run with all their experts at once
+    min_weight: float = _opt(0.5, maximum=1)  # router weight a draft token must still compute
+    refine: str = _opt("regenerate", ("regenerate", "none"))  # what happens to a draft later
+
+
+def _section(cls, raw, where):
+    if not isinstance(raw, dict):
+        raise ValueError(f"{where} must be a table")
+    spec = {f.name: f for f in fields(cls)}
+    _check_keys(raw, set(spec), where)
+    for name, value in raw.items():
+        meta, default = spec[name].metadata, spec[name].default
+        if meta["choices"]:
+            if value not in meta["choices"]:
+                raise ValueError(f"{where} {name} must be one of {meta['choices']}, got {value!r}")
+            continue
+        number = type(value) is int or (type(default) is float and type(value) is float)
+        top = meta["max"]
+        if not number or value < meta["min"] or (top is not None and value > top):
+            raise ValueError(f"{where} {name} must be a number in "
+                             f"[{meta['min']}, {top if top is not None else 'inf'}], got {value!r}")
+    return cls(**raw)
+
+
+@dataclass(frozen=True)
+class Config:
+    pack_dir: str
+    default_tier: str = "ram"
+    cache_slots: int = 16
+    stats_path: str | None = None
+    stats_every: int = 256
+    trace_path: str | None = None
+    policy: Policy = Policy()
+    prefetch: Prefetch = Prefetch()
+    miss: Miss = Miss()
+    serve: Serve = Serve()
+    draft: Draft = Draft()
+    layers: tuple = ()  # ((selector, section), ...) in file order; later sections win
+
+    def placement(self, layer, num_experts):
+        tier = [self.default_tier] * num_experts
+        cache_slots = self.cache_slots
+        for selector, section in self.layers:
+            if selector != "*" and layer not in parse_ids(selector):
+                continue
+            seen = {}
+            for name in TIERS:
+                for expert in parse_ids(section.get(name, []), num_experts):
+                    if expert in seen:
+                        raise ValueError(
+                            f"[layers.\"{selector}\"] puts expert {expert} in both "
+                            f"{seen[expert]} and {name}"
+                        )
+                    seen[expert] = tier[expert] = name
+            cache_slots = section.get("cache_slots", cache_slots)
+        vram, ram, disk = (
+            tuple(e for e in range(num_experts) if tier[e] == name) for name in TIERS
+        )
+        if (ram or disk) and cache_slots < 1:
+            raise ValueError(
+                f"layer {layer}: cache_slots must be >= 1 while experts live in ram or disk"
+            )
+        if (ram or disk) and self.policy.scan_keep >= cache_slots:
+            raise ValueError(f"layer {layer}: policy.scan_keep must be < cache_slots ({cache_slots})")
+        return LayerPlacement(vram, ram, disk, cache_slots)
+
+
+def _check_keys(table, allowed, where):
+    unknown = set(table) - allowed
+    if unknown:
+        raise ValueError(f"{where}: unknown keys {sorted(unknown)}, allowed: {sorted(allowed)}")
+
+
+def _check_count(value, minimum, where):
+    if type(value) is not int or value < minimum:
+        raise ValueError(f"{where} must be an integer >= {minimum}, got {value!r}")
+
+
+def _path(config_path, value):
+    """A path from the file at config_path: '~' expands, relative paths start at that file's
+    directory, so a config can point at '../models/...' from anywhere."""
+    if value is None:
+        return None
+    return os.path.join(os.path.dirname(os.path.abspath(config_path)), os.path.expanduser(value))
+
+
+def load_config(path):
+    with open(path, "rb") as f:
+        raw = tomllib.load(f)
+    _check_keys(raw, _TOP_KEYS, path)
+    if "pack_dir" not in raw:
+        raise ValueError(f"{path}: pack_dir is required")
+    layers = []
+    for selector, section in raw.get("layers", {}).items():
+        where = f'{path} [layers."{selector}"]'
+        if selector != "*":
+            parse_ids(selector)
+        if not isinstance(section, dict):
+            raise ValueError(f"{where} must be a table")
+        _check_keys(section, _LAYER_KEYS, where)
+        if "cache_slots" in section:
+            _check_count(section["cache_slots"], 0, f"{where} cache_slots")
+        layers.append((selector, section))
+    config = Config(
+        pack_dir=_path(path, raw["pack_dir"]),
+        default_tier=raw.get("default_tier", "ram"),
+        cache_slots=raw.get("cache_slots", 16),
+        stats_path=_path(path, raw.get("stats_path")),
+        stats_every=raw.get("stats_every", 256),
+        trace_path=_path(path, raw.get("trace_path")),
+        policy=_section(Policy, raw.get("policy", {}), f"{path} [policy]"),
+        prefetch=_section(Prefetch, raw.get("prefetch", {}), f"{path} [prefetch]"),
+        miss=_section(Miss, raw.get("miss", {}), f"{path} [miss]"),
+        serve=_section(Serve, raw.get("serve", {}), f"{path} [serve]"),
+        draft=_section(Draft, raw.get("draft", {}), f"{path} [draft]"),
+        layers=tuple(layers),
+    )
+    if config.draft.full_requests > config.serve.max_requests:
+        raise ValueError(f"{path}: [draft] full_requests must be <= [serve] max_requests "
+                         f"({config.serve.max_requests})")
+    if config.default_tier not in TIERS:
+        raise ValueError(f"{path}: default_tier must be one of {TIERS}")
+    _check_count(config.cache_slots, 0, f"{path}: cache_slots")
+    _check_count(config.stats_every, 1, f"{path}: stats_every")
+    return config
+
+
+def moe_dims(hf_config):
+    """(num_layers, num_experts, top_k, hidden, moe_intermediate) of a HF config.json dict."""
+    c = hf_config.get("text_config") or hf_config
+
+    def first(*names):
+        for name in names:
+            if c.get(name) is not None:
+                return c[name]
+        raise KeyError(f"config.json has none of {names}")
+
+    return (
+        first("num_hidden_layers"),
+        first("num_experts", "num_local_experts", "n_routed_experts"),
+        first("top_k_experts", "num_experts_per_tok"),
+        first("hidden_size"),
+        first("moe_intermediate_size", "intermediate_size"),
+    )
+
+
+def expert_bytes(hidden, moe_intermediate, dtype_bytes=2):
+    """One gated expert: gate+up [2I, H] plus down [H, I]."""
+    return 3 * hidden * moe_intermediate * dtype_bytes
+
+
+def render_toml(header, top, tables=None, layers=None):
+    """A placement file: top-level keys, [tables], then [layers."N"] sections (empty ones skipped)."""
+    lines = [f"# {header}"] + [f"{k} = {json.dumps(v)}" for k, v in top.items() if v is not None]
+    for name, values in (tables or {}).items():
+        lines += ["", f"[{name}]"] + [f"{k} = {json.dumps(v)}" for k, v in values.items()]
+    for layer, values in sorted((layers or {}).items()):
+        if values:
+            lines += ["", f'[layers."{layer}"]'] + [f"{k} = {json.dumps(v)}" for k, v in values.items()]
+    return "\n".join(lines) + "\n"
+
+
+def plan(stats, num_experts, expert_size, expert_vram_gib, ram_gib, cache_slots, pack_dir, stats_path):
+    """Hottest experts (by steps that routed to them) pinned in VRAM, next in RAM, rest on disk."""
+    uses = {int(layer): s["uses"] for layer, s in stats["layers"].items()}
+    vram_rows = int(expert_vram_gib * GIB // expert_size) - cache_slots * len(uses)
+    if vram_rows < 0:
+        raise ValueError("expert VRAM budget is smaller than the cache slots alone")
+    ram_rows = int(ram_gib * GIB // expert_size)
+    ranked = sorted(
+        ((layer, e) for layer in uses for e in range(num_experts)),
+        key=lambda le: (-uses[le[0]][le[1]], le),
+    )
+    tier = {le: "vram" if i < vram_rows else "ram" if i < vram_rows + ram_rows else "disk"
+            for i, le in enumerate(ranked)}
+    layers = {
+        layer: {name: format_ids(ids) for name in ("vram", "ram")
+                if (ids := [e for e in range(num_experts) if tier[(layer, e)] == name])}
+        for layer in uses
+    }
+    return render_toml(
+        "generated by `prain plan`; experts not listed live on disk",
+        {"pack_dir": pack_dir, "default_tier": "disk", "cache_slots": cache_slots,
+         "stats_path": stats_path},
+        layers=layers,
+    )

prain/drafts.py

@@ -0,0 +1,34 @@
+"""Which token rows of the current forward belong to draft requests. Stdlib only.
+
+A request is a draft when its body carries "vllm_xargs": {"prain": "draft"}; vLLM hands
+that to the worker as sampling_params.extra_args. Every step, vLLM's model runner (V2, the
+default in 0.32) lays the tokens of its requests out back to back in batch order, and
+prain.experts wraps its prepare_inputs to call note_step() with that layout. The MoE
+layers of the forward that follows read spans().
+"""
+
+_DRAFTS = set()  # request ids of unfinished drafts
+_SPANS = [None]  # [(first_row, end_row), ...] of the current forward's draft tokens, or None
+
+
+def note_step(scheduler_output, req_ids, query_start_loc):
+    """scheduler_output: vLLM's SchedulerOutput of this step; req_ids: request id per batch
+    position; query_start_loc: first token row per batch position, plus the total at the end."""
+    for new in scheduler_output.scheduled_new_reqs:
+        args = (new.sampling_params.extra_args or {}) if new.sampling_params else {}
+        if args.get("prain") == "draft":
+            _DRAFTS.add(new.req_id)
+    _DRAFTS.difference_update(scheduler_output.finished_req_ids)
+    spans = [(int(query_start_loc[i]), int(query_start_loc[i + 1]))
+             for i, req_id in enumerate(req_ids) if req_id in _DRAFTS]
+    _SPANS[0] = spans or None
+
+
+def spans():
+    return _SPANS[0]
+
+
+def clear():
+    """Called after the last MoE layer: forwards without a note_step() (vLLM's warm-up and
+    profiling runs) must not see the previous step's layout."""
+    _SPANS[0] = None

prain/policy.py

@@ -0,0 +1,245 @@
+"""Per-layer VRAM slot cache: which slot holds which expert, what to evict, what to copy
+first, how scans and prefetched experts are admitted. Pure Python, no torch."""
+
+import json
+import math
+import os
+from bisect import bisect_right
+
+from prain.config import Policy
+
+
+class Lru:
+    """Evict the expert used longest ago."""
+
+    def __init__(self, num_experts, policy, cost):
+        self.last = [0.0] * num_experts
+
+    def score(self, e, now):
+        return self.last[e]
+
+    def hit(self, e, now, heat):
+        self.last[e] = now
+
+    def admit(self, e, now, heat, low):
+        self.last[e] = -math.inf if low else now
+
+    def evicted(self, e):
+        pass
+
+
+class Lfu:
+    """Evict the expert with the lowest exponentially decayed heat. Scores survive eviction,
+    so an expert that comes back keeps its history."""
+
+    def __init__(self, num_experts, policy, cost):
+        self.decay = 0.5 ** (1 / policy.lfu_half_life)
+        self.value = [0.0] * num_experts
+        self.when = [0] * num_experts
+
+    def score(self, e, now):
+        return self.value[e] * self.decay ** (now - self.when[e])
+
+    def hit(self, e, now, heat):
+        self.value[e] = self.score(e, now) + heat
+        self.when[e] = now
+
+    def admit(self, e, now, heat, low):
+        if not low:  # a low-priority admit keeps its history: never-seen experts score 0
+            self.hit(e, now, heat)
+
+    def evicted(self, e):
+        pass
+
+
+class Gdsf:
+    """GreedyDual-Size-Frequency: score = floor + frequency x refetch cost. The floor rises
+    to each victim's score, so experts that stopped being used age out. cost[e] is 1 for a
+    RAM expert and policy.disk_cost for a disk one, so slow-to-refetch experts stay longer."""
+
+    def __init__(self, num_experts, policy, cost):
+        self.cost = cost
+        self.floor = 0.0
+        self.freq = [0.0] * num_experts
+        self.value = [0.0] * num_experts
+
+    def score(self, e, now):
+        return self.value[e]
+
+    def hit(self, e, now, heat):
+        self.freq[e] += heat
+        self.value[e] = self.floor + self.freq[e] * self.cost[e]
+
+    def admit(self, e, now, heat, low):
+        self.freq[e] = 0.0 if low else heat
+        self.value[e] = self.floor + self.freq[e] * self.cost[e]
+
+    def evicted(self, e):
+        self.floor = max(self.floor, self.value[e])
+
+
+class Belady:
+    """The offline optimum, for the simulator only: evict the expert used furthest in the
+    future. future[e] lists the forward numbers that route to e, ascending."""
+
+    def __init__(self, future):
+        self.future = future
+
+    def score(self, e, now):
+        uses = self.future.get(e, ())
+        i = bisect_right(uses, now)
+        return -uses[i] if i < len(uses) else -math.inf
+
+    def hit(self, e, now, heat):
+        pass
+
+    def admit(self, e, now, heat, low):
+        pass
+
+    def evicted(self, e):
+        pass
+
+
+POLICIES = {"lru": Lru, "lfu": Lfu, "gdsf": Gdsf}
+
+
+class SlotCache:
+    """Slots [0, len(pinned)) hold the pinned experts for the whole run. The cache_slots
+    slots after them hold experts from RAM or disk, managed by an eviction policy."""
+
+    def __init__(self, pinned, num_experts, cache_slots, policy=Policy(), cost=None, evictor=None):
+        self.num_experts = num_experts
+        self.cache_slots = cache_slots
+        self.policy = policy
+        self.pinned = {e: slot for slot, e in enumerate(pinned)}
+        self.slots = {}  # cached expert -> slot
+        first = len(self.pinned)
+        self.free = list(range(first + cache_slots - 1, first - 1, -1))  # pop() -> lowest slot
+        cost = cost or [1.0] * num_experts
+        self.evictor = evictor or POLICIES[policy.eviction](num_experts, policy, cost)
+        self.now = 0  # forwards planned so far
+        self.recent = set()  # experts routed to in the last forward
+        self.prefetched = set()  # prefetched experts that were not used yet
+        self.uses = [0] * num_experts  # forwards that routed to each expert
+        self.hits = 0  # routed expert already in VRAM
+        self.misses = 0  # routed expert copied in
+        self.cpu = 0  # routed expert computed on the CPU instead
+        self.prefetches = 0  # experts copied in ahead of time (prefetch or admit-behind)
+        self.prefetch_hits = 0  # ... that were used before being evicted
+        self.prefetch_wasted = 0  # ... that were evicted unused
+
+    def choose_cpu(self, counts, miss):
+        """Experts the CPU should compute: counts maps expert -> routed tokens."""
+        if miss.mode == "copy":
+            return set()
+        away = (e for e in counts if e not in self.pinned and e not in self.slots)
+        if miss.mode == "cpu":
+            return set(away)
+        return {e for e in away if counts[e] <= miss.cpu_max_tokens}
+
+    def plan(self, heat, tokens=1, cpu=()):
+        """Split one forward pass into rounds.
+
+        heat: {expert: heat > 0} for every expert routed to in this pass.
+        tokens: tokens in the pass; a pass with >= policy.scan_tokens tokens is a scan.
+        cpu: routed experts the CPU computes; they get no slot.
+        Returns [(slot_of, loads), ...]: slot_of maps each expert computed in that round to
+        its slot; loads lists the (expert, slot) copies to do before the round. Experts in
+        VRAM run in round 0. The others follow coldest first, so the hottest are copied last
+        and are the ones left in VRAM afterwards.
+        """
+        self.now += 1
+        now, scan = self.now, tokens >= self.policy.scan_tokens
+        norm = {e: h / tokens for e, h in heat.items()}  # per token, so scans don't dominate
+        for e in heat:
+            self.uses[e] += 1
+        self.recent = set(heat)
+        pinned = [e for e in heat if e in self.pinned]
+        resident = [e for e in heat if e in self.slots]
+        missing = sorted((e for e in heat if e not in self.pinned and e not in self.slots
+                          and e not in cpu), key=lambda e: (heat[e], e))
+        self.hits += len(pinned) + len(resident)
+        self.misses += len(missing)
+        self.cpu += len(cpu)
+        for e in resident:
+            self.evictor.hit(e, now, norm[e])
+            if e in self.prefetched:
+                self.prefetched.discard(e)
+                self.prefetch_hits += 1
+        if missing and not self.cache_slots:
+            raise RuntimeError(f"experts {missing} are not in VRAM and there are no cache slots")
+        protected = set()
+        if scan and self.policy.scan_keep:
+            ranked = sorted(self.slots, key=lambda e: (self.evictor.score(e, now), e))
+            protected = set(ranked[-self.policy.scan_keep:])
+        busy = set(resident)  # experts computed in the round being built
+        slot_of = {e: self.pinned[e] for e in pinned} | {e: self.slots[e] for e in resident}
+        rounds, i = [], 0
+        while True:
+            room = len(self.free) + sum(1 for e in self.slots if e not in busy and e not in protected)
+            group = missing[i:i + room]
+            if i < len(missing) and not group and not busy:
+                raise RuntimeError("no evictable cache slot: lower policy.scan_keep")
+            i += len(group)
+            loads = []
+            for e in group:
+                slot = self._take_slot(now, busy | protected)
+                self.slots[e] = slot
+                self.evictor.admit(e, now, norm[e], low=scan)
+                busy.add(e)
+                loads.append((e, slot))
+            slot_of.update(loads)
+            rounds.append((slot_of, loads))
+            if i >= len(missing):
+                break
+            slot_of, busy = {}, set()
+        if scan and self.policy.promote:
+            for e in missing[-self.policy.promote:]:
+                if e in self.slots:
+                    self.evictor.hit(e, now, norm[e])
+        return rounds
+
+    def prefetch(self, predicted, budget):
+        """Copy predicted experts in ahead of time, at low priority: they are evicted first
+        unless a forward uses them. Never evicts an expert routed to in the last forward.
+        predicted: [(expert, probability)], most likely first. Returns the (expert, slot) loads."""
+        loads, keep = [], set(self.recent)
+        for e, _ in predicted:
+            if len(loads) >= budget:
+                break
+            if e in self.pinned or e in self.slots:
+                continue
+            if not self.free and all(r in keep for r in self.slots):
+                break
+            slot = self._take_slot(self.now, keep)
+            self.slots[e] = slot
+            self.evictor.admit(e, self.now, 0.0, low=True)
+            self.prefetched.add(e)
+            keep.add(e)
+            loads.append((e, slot))
+        self.prefetches += len(loads)
+        return loads
+
+    def _take_slot(self, now, keep):
+        if self.free:
+            return self.free.pop()
+        victim = min((e for e in self.slots if e not in keep),
+                     key=lambda e: (self.evictor.score(e, now), e))
+        self.evictor.evicted(victim)
+        if victim in self.prefetched:
+            self.prefetched.discard(victim)
+            self.prefetch_wasted += 1
+        return self.slots.pop(victim)
+
+
+COUNTERS = ("hits", "misses", "cpu", "prefetches", "prefetch_hits", "prefetch_wasted")
+
+
+def write_stats(path, caches):
+    """Atomically write {"layers": {"<layer>": {"uses": [...], "hits": n, ...}}}."""
+    data = {"layers": {str(layer): {"uses": c.uses, **{k: getattr(c, k) for k in COUNTERS}}
+                       for layer, c in sorted(caches.items())}}
+    tmp = f"{path}.tmp"
+    with open(tmp, "w") as f:
+        json.dump(data, f)
+    os.replace(tmp, path)

prain/sim.py

@@ -0,0 +1,311 @@
+"""Routing traces, a synthetic trace generator, a simulator with a simple cost model, and a
+tuner that searches policy variants and per-layer slot counts. Pure Python, no torch.
+
+A trace is gzip JSON lines, one per MoE layer per forward, in execution order:
+{"l": layer, "t": tokens, "e": [experts], "c": [tokens per expert], "w": [router weight
+mass], "r": [recency-weighted mass], "p": {"1": [[expert, prob], ...], "2": [...]}}.
+"p" holds predictions for layers l+1 and l+2 (pregate) and is optional. Routing does not
+depend on placement, so one trace can evaluate any placement and policy.
+"""
+
+import gzip
+import json
+import random
+from dataclasses import dataclass, field, replace
+
+from prain.config import GIB, LayerPlacement, Miss, Policy, Prefetch
+from prain.policy import COUNTERS, Belady, SlotCache
+
+HEAT_FIELD = {"count": "c", "weight": "w", "recency": "r"}
+
+
+@dataclass(frozen=True)
+class Cost:
+    """Rough RX 6600 XT numbers. scripts/measure.py writes measured ones as JSON."""
+    expert_mib: float = 11.34  # one expert
+    h2d_gbps: float = 12.0  # pinned RAM -> VRAM, PCIe 4.0 x8
+    disk_gbps: float = 3.0  # pack file -> RAM
+    round_ms: float = 0.12  # one fused-MoE kernel launch
+    route_ms: float = 0.0012  # GPU work per (token, expert) route
+    other_ms: float = 0.6  # attention + dense MLP of one layer per forward
+    sync_ms: float = 0.05  # host sync per layer
+    cpu_ms: float = 0.5  # CPU work per (token, expert) route
+
+    def copy_ms(self, disk):
+        mb = self.expert_mib * 1.048576
+        return mb / self.h2d_gbps + (mb / self.disk_gbps if disk else 0.0)
+
+
+def load_cost(path, expert_mib):
+    values = {"expert_mib": expert_mib}
+    if path:
+        with open(path) as f:
+            values |= json.load(f)
+        values["expert_mib"] = expert_mib
+    return Cost(**values)
+
+
+def write_trace(path, records, mode="wt"):
+    with gzip.open(path, mode) as f:
+        for r in records:
+            f.write(json.dumps(r, separators=(",", ":")) + "\n")
+
+
+def read_forwards(path):
+    """[{layer: record}] per forward, in execution order."""
+    forwards, last = [], None
+    with gzip.open(path, "rt") as f:
+        for line in f:
+            if not line.strip():
+                continue
+            r = json.loads(line)
+            if last is None or r["l"] <= last:
+                forwards.append({})
+            forwards[-1][r["l"]] = r
+            last = r["l"]
+    return forwards
+
+
+def synthetic_trace(layers=30, experts=128, top_k=8, forwards=300, prompt_every=100,
+                    prompt_tokens=128, zipf=1.1, stay=0.5, accuracy=0.7, seed=0):
+    """Zipf-popular experts per layer, decode tokens that keep some of the previous token's
+    experts, a prompt every prompt_every forwards, and pregate predictions that are right
+    with probability accuracy (accuracy**2 two layers ahead)."""
+    rng = random.Random(seed)
+    rank = [rng.sample(range(experts), experts) for _ in range(layers)]
+    weights = [1 / (i + 1) ** zipf for i in range(experts)]
+    previous = [[] for _ in range(layers)]
+    records = []
+    for f in range(forwards):
+        tokens = prompt_tokens if f % prompt_every == 0 else 1
+        routed = []
+        for layer in range(layers):
+            count, recency = {}, {}
+            for t in range(tokens):
+                chosen = {e for e in previous[layer] if rng.random() < stay} if tokens == 1 else set()
+                while len(chosen) < top_k:
+                    chosen.add(rank[layer][rng.choices(range(experts), weights)[0]])
+                decay = 0.5 ** ((tokens - 1 - t) / 32)
+                for e in chosen:
+                    count[e] = count.get(e, 0) + 1
+                    recency[e] = recency.get(e, 0.0) + decay / top_k
+                if tokens == 1:
+                    previous[layer] = sorted(chosen)
+            routed.append((count, recency))
+        for layer, (count, recency) in enumerate(routed):
+            ids = sorted(count)
+            record = {"l": layer, "t": tokens, "e": ids, "c": [count[e] for e in ids],
+                      "w": [round(count[e] / top_k, 4) for e in ids],
+                      "r": [round(recency[e], 4) for e in ids]}
+            predictions = {}
+            for d in (1, 2):
+                if layer + d < layers:
+                    nxt = routed[layer + d][0]
+                    guess = []
+                    for e in sorted(nxt, key=lambda e: (-nxt[e], e))[:top_k]:
+                        g = e if rng.random() < accuracy ** d else rng.randrange(experts)
+                        if g not in guess:
+                            guess.append(g)
+                    predictions[str(d)] = [[g, round(1 / top_k, 4)] for g in guess]
+            if predictions:
+                record["p"] = predictions
+            records.append(record)
+    return records
+
+
+@dataclass
+class Result:
+    total_ms: float = 0.0
+    decode_ms: float = 0.0  # summed over forwards with fewer than scan_tokens tokens
+    decode_forwards: int = 0
+    prefill_ms: float = 0.0
+    prefill_tokens: int = 0
+    counters: dict = field(default_factory=dict)
+    layer_ms: dict = field(default_factory=dict)
+
+    @property
+    def hit_rate(self):
+        c = self.counters
+        return c["hits"] / max(1, c["hits"] + c["misses"] + c["cpu"])
+
+    @property
+    def ms_per_decode(self):
+        return self.decode_ms / max(1, self.decode_forwards)
+
+
+def _future(forwards, layer):
+    future, n = {}, 0
+    for fwd in forwards:
+        if layer in fwd:
+            n += 1
+            for e in fwd[layer]["e"]:
+                future.setdefault(e, []).append(n)
+    return future
+
+
+def simulate(forwards, placements, policy, prefetch, miss, cost, belady=False):
+    """Replay a trace. Model: one PCIe link serves copies in issue order; demand copies of
+    a round wait for the previous round's kernel; async copies (prefetch, admit-behind) are
+    issued after the layer's demand copies; the CPU path runs alongside the GPU."""
+    if prefetch.mode == "pregate" and not any("p" in r for fwd in forwards for r in fwd.values()):
+        raise ValueError("prefetch.mode = pregate needs a trace with predictions ('p')")
+    caches, disk = {}, {}
+    for layer, p in placements.items():
+        n = len(p.vram) + len(p.ram) + len(p.disk)
+        disk[layer] = set(p.disk)
+        cost_of = [policy.disk_cost if e in disk[layer] else 1.0 for e in range(n)]
+        evictor = Belady(_future(forwards, layer)) if belady else None
+        caches[layer] = SlotCache(p.vram, n, p.cache_slots, policy, cost_of, evictor)
+    result = Result(layer_ms={layer: 0.0 for layer in caches})
+    key, t, link, ready = HEAT_FIELD[policy.heat], 0.0, 0.0, {}
+    for fwd in forwards:
+        start, tokens = t, 0
+        for layer in sorted(fwd):
+            if layer not in caches:
+                continue
+            rec, cache = fwd[layer], caches[layer]
+            tokens = rec["t"]
+            counts = dict(zip(rec["e"], rec["c"]))
+            heat = dict(zip(rec["e"], rec[key]))
+            cpu = cache.choose_cpu(counts, miss)
+            gpu = t + cost.sync_ms
+            for slot_of, loads in cache.plan(heat, tokens, cpu):
+                if loads:
+                    link = max(link, gpu) + sum(cost.copy_ms(e in disk[layer]) for e, _ in loads)
+                    gpu = link
+                land = max((ready.pop((layer, e), 0.0) for e in slot_of), default=0.0)
+                gpu = max(gpu, land) + cost.round_ms + cost.route_ms * sum(counts[e] for e in slot_of)
+            cpu_end = t + cost.sync_ms + cost.cpu_ms * sum(counts[e] for e in cpu)
+            later = []
+            if cpu and miss.admit:
+                hot = sorted((e for e in cpu if e not in disk[layer]), key=lambda e: (-heat[e], e))
+                later += [(layer, e) for e, _ in cache.prefetch([(e, 1.0) for e in hot], miss.admit)]
+            target = layer + prefetch.distance
+            if prefetch.mode == "pregate" and target in caches:
+                guesses = [(e, p) for e, p in rec.get("p", {}).get(str(prefetch.distance), [])
+                           if p >= prefetch.min_prob and e not in disk[target]]
+                later += [(target, e) for e, _ in caches[target].prefetch(guesses, prefetch.budget)]
+            for where in later:
+                link = max(link, t + cost.sync_ms) + cost.copy_ms(False)
+                ready[where] = link
+            end = max(gpu, cpu_end) + cost.other_ms
+            result.layer_ms[layer] += end - t
+            t = end
+        if tokens < policy.scan_tokens:
+            result.decode_ms += t - start
+            result.decode_forwards += 1
+        else:
+            result.prefill_ms += t - start
+            result.prefill_tokens += tokens
+    result.total_ms = t
+    result.counters = {k: sum(getattr(c, k) for c in caches.values()) for k in COUNTERS}
+    return result
+
+
+def describe(result, bound, cost):
+    c = result.counters
+    copied = (c["misses"] + c["prefetches"]) * cost.expert_mib / 1024
+    rate = 1000 / result.ms_per_decode if result.decode_forwards else 0.0
+    return [
+        f"forwards {result.decode_forwards} decode + prefill of {result.prefill_tokens} tokens",
+        f"VRAM hit rate {100 * result.hit_rate:.1f}%  "
+        f"(eviction alone at best: {100 * bound.hit_rate:.1f}%)",
+        f"copies {c['misses']} + prefetched {c['prefetches']} (used {c['prefetch_hits']}, "
+        f"wasted {c['prefetch_wasted']}) = {copied:.1f} GiB, CPU-run {c['cpu']}",
+        f"decode {result.ms_per_decode:.1f} ms/forward ({rate:.1f} tok/s), "
+        f"prefill {result.prefill_ms / max(1, result.prefill_tokens):.2f} ms/token",
+        f"simulated total {result.total_ms / 1000:.2f} s",
+    ]
+
+
+DIMENSIONS = (
+    ("policy", "eviction", ("lru", "lfu", "gdsf")),
+    ("policy", "heat", ("count", "weight", "recency")),
+    ("policy", "scan_keep", (0, 1, 2)),
+    ("policy", "promote", (0, 2, 4)),
+    ("prefetch", "mode", ("none", "pregate")),
+    ("prefetch", "distance", (1, 2)),
+    ("prefetch", "budget", (1, 2, 4)),
+    ("miss", "mode", ("copy", "auto", "cpu")),
+    ("miss", "admit", (0, 1, 2)),
+)
+_PART = {"policy": 0, "prefetch": 1, "miss": 2}
+
+
+def _with(variant, table, name, value):
+    parts = list(variant)
+    parts[_PART[table]] = replace(parts[_PART[table]], **{name: value})
+    return tuple(parts)
+
+
+def ram_tier(forwards, layers, num_experts, rows):
+    """The `rows` most-used (layer, expert) pairs of the trace -> {layer: set of experts}."""
+    uses = {(layer, e): 0 for layer in layers for e in range(num_experts)}
+    for fwd in forwards:
+        for layer, rec in fwd.items():
+            for e in rec["e"]:
+                uses[(layer, e)] += 1
+    ranked = sorted(uses, key=lambda le: (-uses[le], le))[:rows]
+    ram = {layer: set() for layer in layers}
+    for layer, e in ranked:
+        ram[layer].add(e)
+    return ram
+
+
+def tune(forwards, num_experts, expert_vram_gib, ram_gib, cost, base=(Policy(), Prefetch(), Miss()),
+         passes=2, window=3):
+    """Coordinate descent over DIMENSIONS with uniform slots, then per-layer slot counts by
+    greedy marginal gain. Returns (variant, ram, slots_per_layer, result, bound, rows)."""
+    layers = sorted({layer for fwd in forwards for layer in fwd})
+    expert_size = cost.expert_mib * 2**20
+    slots = int(expert_vram_gib * GIB // expert_size) // len(layers)
+    if slots < 1:
+        raise ValueError("expert VRAM budget is less than one slot per layer")
+    ram = ram_tier(forwards, layers, num_experts, int(ram_gib * GIB // expert_size))
+
+    def placements(per_layer):
+        return {layer: LayerPlacement((), tuple(sorted(ram[layer])),
+                                      tuple(e for e in range(num_experts) if e not in ram[layer]),
+                                      per_layer[layer]) for layer in layers}
+
+    uniform = placements({layer: slots for layer in layers})
+    predicted = any("p" in r for fwd in forwards for r in fwd.values())
+    results = {}
+
+    def run(variant):
+        if variant not in results:
+            ok = variant[0].scan_keep < slots
+            results[variant] = simulate(forwards, uniform, *variant, cost) if ok else None
+        return results[variant]
+
+    best = base
+    if best[1].mode == "pregate" and not predicted:
+        best = _with(best, "prefetch", "mode", "none")
+    for _ in range(passes):
+        for table, name, values in DIMENSIONS:
+            if table == "prefetch" and (not predicted or (name != "mode" and best[1].mode == "none")):
+                continue
+            if (table, name) == ("miss", "admit") and best[2].mode == "copy":
+                continue
+            scored = [(r.total_ms, i, v) for i, v in enumerate(_with(best, table, name, x) for x in values)
+                      if (r := run(v)) is not None]
+            best = min(scored)[2]
+    lo, hi = max(1, slots - window, best[0].scan_keep + 1), slots + window
+    curves = {}
+    for s in range(lo, hi + 1):
+        r = simulate(forwards, placements({layer: s for layer in layers}), *best, cost)
+        curves |= {(layer, s): ms for layer, ms in r.layer_ms.items()}
+    per = {layer: lo for layer in layers}
+    for _ in range((slots - lo) * len(layers)):
+        gains = [(curves[(layer, per[layer])] - curves[(layer, per[layer] + 1)], -layer)
+                 for layer in layers if per[layer] < hi]
+        if not gains:
+            break
+        per[-max(gains)[1]] += 1
+    final = simulate(forwards, placements(per), *best, cost)
+    if final.total_ms > results[best].total_ms:  # greedy on approximate curves can lose
+        per, final = {layer: slots for layer in layers}, results[best]
+    bound = simulate(forwards, placements(per), Policy(scan_tokens=best[0].scan_tokens),
+                     Prefetch(), Miss(), cost, belady=True)
+    rows = sorted((r.total_ms, str(v), v, r) for v, r in results.items() if r is not None)
+    return best, ram, per, final, bound, [(v, r) for _, _, v, r in rows]

prain/tui.py

@@ -0,0 +1,115 @@
+"""prain tui: a terminal chat that warms experts up while you type. Stdlib curses."""
+
+import curses
+import locale
+import os
+import queue
+import textwrap
+import threading
+import time
+
+from prain.chat import (ABORTED, Warmer, chat_body, first_model, post, send_request,
+                          stream_text, warm_body)
+
+WARMUP_MODES = ("idle", "always", "off")
+
+
+class Chat:
+    """Conversation state without curses, so tests can drive it.
+    warmup: "idle" warms up only while no answer streams, "always" also during one, "off" never."""
+
+    def __init__(self, base_url, model, warmup="idle", priority=1):
+        self.base_url, self.model, self.warmup = base_url, model, warmup
+        self.messages = []  # finished turns
+        self.draft = ""
+        self.reply = None  # the answer text while it streams, else None
+        self.status = "type a message; Enter sends, Esc quits"
+        self.events = queue.Queue()
+        self.warmer = Warmer(lambda draft: send_request(
+            base_url, "/v1/chat/completions", warm_body(model, self.messages, draft, priority)))
+
+    def type(self, text):
+        self.draft = text
+        idle = self.reply is None
+        if text.strip() and (self.warmup == "always" or (self.warmup == "idle" and idle)):
+            self.warmer.update(text)
+
+    def send(self):
+        if self.reply is not None or not self.draft.strip():
+            return
+        self.warmer.cancel()  # a queued draft is stale now; the in-flight one only delays us
+        self.messages.append({"role": "user", "content": self.draft.strip()})
+        self.draft, self.reply = "", ""
+        threading.Thread(target=self._stream, args=(list(self.messages),), daemon=True).start()
+
+    def _stream(self, messages):
+        start, first, pieces = time.monotonic(), None, 0
+        try:
+            with post(self.base_url, "/v1/chat/completions", chat_body(self.model, messages)) as r:
+                for piece in stream_text(r):
+                    first = first or time.monotonic()
+                    pieces += 1
+                    self.events.put(("piece", piece))
+            rate = (pieces - 1) / (time.monotonic() - first) if pieces > 1 else 0.0
+            ttft = (first or time.monotonic()) - start
+            self.events.put(("done", f"first token {ttft:.1f} s, {rate:.1f} tok/s"))
+        except ABORTED as e:
+            self.events.put(("done", f"error: {e}"))
+
+    def poll(self):
+        """Apply streamed pieces; call from the UI thread."""
+        while True:
+            try:
+                kind, value = self.events.get_nowait()
+            except queue.Empty:
+                return
+            if kind == "piece":
+                self.reply += value
+            else:
+                self.messages.append({"role": "assistant", "content": self.reply})
+                self.reply, self.status = None, value
+
+
+def draw(screen, chat):
+    screen.erase()
+    height, width = screen.getmaxyx()
+    turns = chat.messages + ([{"role": "assistant", "content": chat.reply}]
+                             if chat.reply is not None else [])
+    lines = []
+    for turn in turns:
+        prefix = "you> " if turn["role"] == "user" else "ai>  "
+        for paragraph in (prefix + turn["content"]).split("\n"):
+            lines += textwrap.wrap(paragraph, width - 1) or [""]
+    for y, line in enumerate(lines[-(height - 2):]):
+        screen.addstr(y, 0, line)
+    warming = "warming experts... " if chat.warmer.busy else ""
+    screen.addstr(height - 2, 0, f"[{chat.model}] {warming}{chat.status}"[: width - 1],
+                  curses.A_REVERSE)
+    screen.addstr(height - 1, 0, ("> " + chat.draft)[-(width - 1):])
+    screen.refresh()
+
+
+def run(screen, chat):
+    screen.timeout(50)  # redraw at least every 50 ms while an answer streams
+    while True:
+        chat.poll()
+        draw(screen, chat)
+        try:
+            key = screen.get_wch()
+        except curses.error:
+            continue  # nothing typed
+        if key in ("\x1b", "\x04"):  # Esc or Ctrl-D
+            return
+        if key in ("\n", "\r", curses.KEY_ENTER):
+            chat.send()
+        elif key in (curses.KEY_BACKSPACE, "\x7f", "\b"):
+            chat.type(chat.draft[:-1])
+        elif isinstance(key, str) and key.isprintable():
+            chat.type(chat.draft + key)
+
+
+def main(base_url, warmup="idle", priority=1):
+    locale.setlocale(locale.LC_ALL, "")
+    os.environ.setdefault("ESCDELAY", "25")  # Esc quits at once, not after 1 s
+    chat = Chat(base_url, first_model(base_url), warmup, priority)
+    curses.wrapper(run, chat)

prain/web.py

@@ -0,0 +1,274 @@
+"""prain web: a one-page browser chat that warms experts up while you type.
+
+A small stdlib proxy in front of vLLM. GET / serves the page. POST /warm takes
+{"session", "messages", "draft"} and returns at once. POST /chat takes {"session",
+"messages"} and streams server-sent events: vLLM's chunks, plus {"prain": "draft" |
+"refine" | "error"} events of its own. Binds to 127.0.0.1; use an SSH tunnel from other
+machines.
+
+Several browsers at once ([serve] and [draft] tables of the placement file):
+- At most draft.full_requests answers run with all their experts. An answer that arrives
+  while they are taken runs as a draft right away, on the experts already in VRAM (see
+  prain.drafts), instead of waiting or evicting the other answers' experts.
+- With draft.refine = "regenerate", the oldest waiting draft takes the next free full slot:
+  its draft request is cancelled and the answer is generated again with all experts. The
+  page shows the draft in grey and the full answer over it.
+- Every session has its own Warmer. warmup="idle": no warm-ups while any answer streams;
+  "always": warm-ups while vLLM has a free request slot (serve.max_requests). Sending cancels
+  that session's warm-up. A browser that goes away closes the upstream request, and vLLM
+  cancels it.
+"""
+
+import json
+import threading
+import urllib.error
+from collections import OrderedDict, deque
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+from prain.chat import Warmer, chat_body, draft_body, first_model, post, send_request, warm_body
+from prain.config import Draft, Serve
+
+MAX_SESSIONS = 64
+
+PAGE = """<!doctype html>
+<meta charset="utf-8">
+<title>prain</title>
+<style>
+  body { font: 15px/1.45 system-ui, sans-serif; max-width: 52rem; margin: 2rem auto; padding: 0 1rem; }
+  .user { color: #666; margin-top: 1rem; }
+  .assistant { white-space: pre-wrap; }
+  .draft { color: #999; }
+  textarea { width: 100%; height: 5rem; font: inherit; margin-top: 1rem; }
+  #status { color: #888; font-size: 13px; }
+</style>
+<div id="log"></div>
+<textarea id="draft" autofocus
+  placeholder="Experts load while you type. Enter sends, Shift+Enter adds a line."></textarea>
+<div id="status"></div>
+<script>
+const session = Math.random().toString(36).slice(2) + Date.now().toString(36);
+const messages = [];
+const log = document.getElementById("log");
+const draft = document.getElementById("draft");
+const status = document.getElementById("status");
+let streaming = false;
+
+function show(role, text) {
+  const div = document.createElement("div");
+  div.className = role;
+  div.textContent = text;
+  log.appendChild(div);
+  return div;
+}
+
+draft.addEventListener("input", () => {
+  if (!streaming && draft.value.trim())
+    fetch("/warm", {method: "POST", body: JSON.stringify({session, messages, draft: draft.value})});
+});
+
+draft.addEventListener("keydown", async (event) => {
+  if (event.key !== "Enter" || event.shiftKey || streaming || !draft.value.trim()) return;
+  event.preventDefault();
+  const text = draft.value.trim();
+  messages.push({role: "user", content: text});
+  show("user", "you> " + text);
+  draft.value = "";
+  streaming = true;
+  const out = show("assistant", "");
+  const finalSpan = out.appendChild(document.createElement("span"));
+  const draftSpan = out.appendChild(document.createElement("span"));
+  draftSpan.className = "draft";
+  let mode = "full", draftText = "", finalText = "", refined = false;
+  const render = () => {  // the full text, then the grey rest of the draft while they agree
+    finalSpan.textContent = finalText;
+    draftSpan.textContent = mode === "draft" ? draftText
+      : mode === "refine" && draftText.startsWith(finalText) ? draftText.slice(finalText.length) : "";
+  };
+  let start = performance.now(), first = 0, pieces = 0, buffer = "";
+  const response = await fetch("/chat", {method: "POST", body: JSON.stringify({session, messages})});
+  if (!response.ok) {
+    out.textContent = "error " + response.status + ": " + await response.text();
+    messages.pop();
+    streaming = false;
+    return;
+  }
+  const reader = response.body.getReader();
+  const decoder = new TextDecoder();
+  for (;;) {
+    const {done, value} = await reader.read();
+    if (done) break;
+    buffer += decoder.decode(value, {stream: true});
+    const lines = buffer.split("\\n");
+    buffer = lines.pop();
+    for (const line of lines) {
+      if (!line.startsWith("data: ") || line === "data: [DONE]") continue;
+      const data = JSON.parse(line.slice(6));
+      if (data.prain === "draft") {
+        mode = "draft";
+        status.textContent = "GPU busy: drafting from the experts in VRAM, refined when it is free";
+      } else if (data.prain === "refine") {
+        mode = "refine";
+        refined = true;
+        start = performance.now(), first = 0, pieces = 0;
+        status.textContent = "refining with all experts";
+      } else if (data.prain === "error") {
+        mode = "draft";
+        refined = false;
+        status.textContent = "refining failed (HTTP " + data.status + "), keeping the draft";
+      } else {
+        const piece = data.choices[0].delta.content;
+        if (!piece) continue;
+        if (pieces++ === 0) first = performance.now();
+        if (mode === "draft") draftText += piece; else finalText += piece;
+      }
+      render();
+    }
+  }
+  messages.push({role: "assistant", content: mode === "draft" ? draftText : finalText});
+  let same = 0;
+  while (same < draftText.length && draftText[same] === finalText[same]) same++;
+  const seconds = (performance.now() - first) / 1000;
+  status.textContent = (refined ? `draft kept ${same} of ${draftText.length} chars; refined: `
+                        : mode === "draft" ? "draft, not refined: " : "") +
+    `first token ${((first - start) / 1000).toFixed(1)} s, ` +
+    `${pieces > 1 ? ((pieces - 1) / seconds).toFixed(1) : "0.0"} tok/s`;
+  streaming = false;
+});
+</script>
+"""
+
+
+def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, serve=Serve(),
+                draft=Draft()):
+    model = first_model(base_url)
+    lock = threading.Condition()
+    sessions = OrderedDict()  # session -> Warmer, least recently used first
+    streams = [0]  # answers streaming right now, all sessions, full or draft
+    full = [0]  # answers holding one of the draft.full_requests full slots
+    waiting = deque()  # drafts waiting for a full slot to refine, oldest first
+    refine = draft.refine == "regenerate"
+
+    def warmer(session):
+        with lock:
+            if session not in sessions:
+                sessions[session] = Warmer(lambda job: send_request(
+                    base_url, "/v1/chat/completions", warm_body(model, *job, priority)))
+                if len(sessions) > MAX_SESSIONS:
+                    sessions.popitem(last=False)[1].cancel()
+            sessions.move_to_end(session)
+            return sessions[session]
+
+    def my_turn(ticket):  # call with the lock held
+        return full[0] < draft.full_requests and bool(waiting) and waiting[0] is ticket
+
+    def turn_now(ticket):
+        with lock:
+            return my_turn(ticket)
+
+    class Handler(BaseHTTPRequestHandler):
+        def do_GET(self):
+            self.reply(200, "text/html; charset=utf-8", PAGE.encode())
+
+        def do_POST(self):
+            body = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
+            session = str(body.get("session", ""))
+            if self.path == "/warm":
+                with lock:
+                    busy = streams[0]
+                if body["draft"].strip() and ((warmup == "idle" and not busy) or
+                                              (warmup == "always" and busy < serve.max_requests)):
+                    warmer(session).update((body["messages"], body["draft"]))
+                return self.reply(204, "text/plain", b"")
+            if self.path != "/chat":
+                return self.reply(404, "text/plain", b"not found")
+            warmer(session).cancel()
+            messages, ticket = body["messages"], object()
+            with lock:
+                streams[0] += 1
+                holds_full = full[0] < draft.full_requests and not waiting
+                if holds_full:
+                    full[0] += 1
+                elif refine:
+                    waiting.append(ticket)
+            try:
+                if holds_full:
+                    if not self.relay(chat_body(model, messages), opens=True):
+                        return
+                else:
+                    if not self.relay(draft_body(model, messages), opens=True, event="draft",
+                                      stop=(lambda: turn_now(ticket)) if refine else None):
+                        return
+                    if refine:
+                        with lock:
+                            lock.wait_for(lambda: my_turn(ticket))
+                            waiting.popleft()
+                            full[0] += 1
+                            holds_full = True
+                        self.relay(chat_body(model, messages), opens=False, event="refine")
+                self.send_event("[DONE]")
+            except (BrokenPipeError, ConnectionResetError):
+                pass  # the browser left; closing upstream makes vLLM cancel the request
+            finally:
+                with lock:
+                    streams[0] -= 1
+                    if holds_full:
+                        full[0] -= 1
+                    if ticket in waiting:
+                        waiting.remove(ticket)
+                    lock.notify_all()
+
+        def relay(self, body, opens, event=None, stop=None):
+            """Stream one vLLM answer to the browser, without its [DONE]. opens: this starts
+            the response, so a vLLM error becomes the HTTP status (returns False). stop():
+            checked after every chunk; True closes the upstream request early."""
+            try:
+                upstream = post(base_url, "/v1/chat/completions", body)
+            except urllib.error.HTTPError as e:  # vLLM rejected it: pass its answer on
+                with e:
+                    data = e.read()
+                if opens:
+                    self.reply(e.code, "application/json", data)
+                    return False
+                self.send_event({"prain": "error", "status": e.code})
+                return True
+            with upstream:
+                if opens:
+                    self.send_response(200)
+                    self.send_header("Content-Type", "text/event-stream")
+                    self.end_headers()
+                if event:
+                    self.send_event({"prain": event})
+                for line in upstream:
+                    if line.strip() in (b"", b"data: [DONE]"):
+                        continue
+                    self.wfile.write(line.rstrip(b"\n") + b"\n\n")
+                    self.wfile.flush()
+                    if stop is not None and stop():
+                        break
+            return True
+
+        def send_event(self, data):
+            text = data if isinstance(data, str) else json.dumps(data)
+            self.wfile.write(f"data: {text}\n\n".encode())
+            self.wfile.flush()
+
+        def reply(self, code, kind, data):
+            self.send_response(code)
+            self.send_header("Content-Type", kind)
+            self.send_header("Content-Length", str(len(data)))
+            self.end_headers()
+            self.wfile.write(data)
+
+        def log_message(self, *args):
+            pass
+
+    server = ThreadingHTTPServer((host, port), Handler)
+    server.sessions, server.streams, server.full = sessions, streams, full  # for tests
+    return server
+
+
+def serve(base_url, port, warmup="idle", priority=1, serve=Serve(), draft=Draft()):
+    server = make_server(base_url, port, warmup=warmup, priority=priority, serve=serve, draft=draft)
+    print(f"prain web: http://127.0.0.1:{server.server_port}  (vLLM at {base_url}; "
+          f"{draft.full_requests} full of {serve.max_requests} parallel answers)")
+    server.serve_forever()

tests/test_chat.py

@@ -0,0 +1,60 @@
+import threading
+import time
+import unittest
+
+import prain.chat as chat
+from tests.helpers import StubVllm, wait_for
+
+
+class Chat(unittest.TestCase):
+    def setUp(self):
+        self.stub = StubVllm()
+        self.addCleanup(self.stub.close)
+
+    def warmer(self):
+        return chat.Warmer(lambda draft: chat.send_request(
+            self.stub.url, "/v1/chat/completions", chat.warm_body("m", [], draft, priority=1)))
+
+    def test_warm_body_continues_the_half_typed_message(self):
+        body = chat.warm_body("m", [{"role": "assistant", "content": "hi"}], "  Explain mix ", 1)
+        self.assertEqual(body["messages"][-1], {"role": "user", "content": "Explain mix"})
+        self.assertEqual(
+            (body["max_tokens"], body["add_generation_prompt"], body["continue_final_message"],
+             body["stream"], body["priority"]), (1, False, True, True, 1))
+        self.assertNotIn("priority", chat.warm_body("m", [], "x"))
+
+    def test_draft_body_marks_the_request_and_salts_its_cache(self):
+        body = chat.draft_body("m", [{"role": "user", "content": "hi"}])
+        self.assertEqual(body, {"model": "m", "messages": [{"role": "user", "content": "hi"}],
+                                "stream": True, "vllm_xargs": {"prain": "draft"},
+                                "cache_salt": "prain-draft"})
+
+    def test_first_model_and_streamed_text(self):
+        self.assertEqual(chat.first_model(self.stub.url), "stub-model")
+        with chat.post(self.stub.url, "/v1/chat/completions",
+                       chat.chat_body("stub-model", [], max_tokens=5)) as response:
+            self.assertEqual(list(chat.stream_text(response)), ["Hel", "lo", "!"])
+        self.assertEqual((self.stub.bodies[0]["stream"], self.stub.bodies[0]["max_tokens"]),
+                         (True, 5))
+
+    def test_warmer_keeps_one_request_in_flight_and_sends_the_newest_draft(self):
+        self.stub.delay = 0.3
+        warmer = self.warmer()
+        warmer.update("a")
+        wait_for(lambda: len(self.stub.bodies) == 1)
+        warmer.update("ab")  # "a" is still in flight: these two queue up,
+        warmer.update("abc")  # and only the newest is sent
+        wait_for(lambda: not warmer.busy)
+        self.assertEqual([b["messages"][-1]["content"] for b in self.stub.bodies], ["a", "abc"])
+
+    def test_cancel_drops_the_queue_and_aborts_the_request_in_flight(self):
+        self.stub.delay = 3.0
+        warmer = self.warmer()
+        warmer.update("slow")
+        wait_for(lambda: self.stub.bodies)
+        warmer.update("queued")
+        start = time.monotonic()
+        warmer.cancel()
+        wait_for(lambda: not warmer.busy, timeout=2.0)
+        self.assertLess(time.monotonic() - start, 1.0)  # did not wait out the 3 s answer
+        self.assertEqual(len(self.stub.bodies), 1)  # "queued" was never sent

tests/test_cli.py

@@ -0,0 +1,78 @@
+import contextlib
+import io
+import json
+import os
+import unittest
+from unittest import mock
+
+from prain.cli import main
+from prain.config import load_config
+from tests.helpers import TmpCase
+
+GEMMA4 = {"text_config": {"num_hidden_layers": 30, "num_experts": 128, "top_k_experts": 8,
+                          "hidden_size": 2816, "moe_intermediate_size": 704}}
+
+
+class Cli(TmpCase):
+    def run_cli(self, *argv):
+        out, err = io.StringIO(), io.StringIO()
+        with contextlib.redirect_stdout(out), contextlib.redirect_stderr(err):
+            main(list(argv))
+        return out.getvalue(), err.getvalue()
+
+    def test_check_prints_tier_sizes(self):
+        hf = self.write("config.json", json.dumps(GEMMA4))
+        config = self.write("p.toml", 'pack_dir = "/packs"\ncache_slots = 4\n')
+        out, err = self.run_cli("check", config, hf)
+        self.assertIn("    0     0   128     0            4", out)
+        self.assertTrue(out.endswith(
+            "expert size 11.34 MiB\n"
+            "VRAM 1.33 GiB  (pinned experts + cache slots)\n"
+            "KV   0.75 GiB  (shared by up to 2 parallel requests)\n"
+            "RAM  42.54 GiB  (pinned host memory)\n"
+            "disk 42.54 GiB  (pack in /packs)\n"), out)
+        self.assertIn("warning: 30 layers have fewer cache slots than top_k=8", err)
+
+    def test_plan_output_is_a_valid_config(self):
+        hf = self.write("config.json", json.dumps(GEMMA4))
+        stats = {"layers": {str(layer): {"uses": list(range(128))} for layer in range(30)}}
+        stats_path = self.write("stats.json", json.dumps(stats))
+        out, _ = self.run_cli("plan", stats_path, hf, "--expert-vram-gib", "8",
+                              "--ram-gib", "16", "--pack-dir", "/packs")
+        config = self.write("planned.toml", out)
+        check_out, _ = self.run_cli("check", config, hf)
+        self.assertIn("VRAM 8.00 GiB", check_out)  # budgets are filled, never exceeded
+        self.assertIn("RAM  16.00 GiB", check_out)
+
+    def test_serve_execs_vllm_with_the_required_flags(self):
+        config = self.write("p.toml", 'pack_dir = "/packs"\n[miss]\nmode = "auto"\ncpu_threads = 6\n'
+                                      '[serve]\nmax_requests = 3\nkv_cache_gib = 1.0\n')
+        with mock.patch("os.execvp") as execvp, mock.patch.dict(os.environ, clear=False):
+            os.environ.pop("OMP_NUM_THREADS", None)
+            main(["serve", config, "/models/m", "--", "--port", "8001"])
+            self.assertEqual(os.environ["PRAIN_CONFIG"], config)
+            self.assertEqual(os.environ["OMP_NUM_THREADS"], "6")
+        execvp.assert_called_once_with("vllm", [
+            "vllm", "serve", "/models/m", "--enforce-eager", "--moe-backend", "triton",
+            "--scheduling-policy", "priority", "--max-num-seqs", "3",
+            "--kv-cache-memory-bytes", "1073741824", "--port", "8001"])
+        with mock.patch("os.execvp") as execvp, self.assertRaisesRegex(
+                SystemExit, "--max-num-seqs=4 comes from the placement file"):
+            main(["serve", config, "/models/m", "--", "--max-num-seqs=4"])
+        execvp.assert_not_called()
+
+    def test_synth_sim_and_tune(self):
+        hf = self.write("config.json", json.dumps(GEMMA4))
+        trace = os.path.join(self.dir, "trace.jsonl.gz")
+        self.run_cli("synth", trace, "--forwards", "40", "--layers", "3")
+        config = self.write("p.toml", 'pack_dir = "/packs"\ncache_slots = 5\n')
+        out, _ = self.run_cli("sim", trace, config, hf)
+        self.assertRegex(out, r"VRAM hit rate \d+\.\d%  \(eviction alone at best: \d+\.\d%\)")
+        self.assertIn("decode ", out)
+        out, err = self.run_cli("tune", trace, hf, "--expert-vram-gib", "0.2",
+                                "--ram-gib", "4", "--pack-dir", "/packs", "--top", "3")
+        self.assertIn("    1  ", err)
+        tuned = load_config(self.write("tuned.toml", out))
+        self.assertEqual((tuned.pack_dir, tuned.trace_path), ("/packs", trace))
+        self.assertLessEqual(sum(tuned.placement(layer, 128).cache_slots for layer in range(3)),
+                             int(0.2 * 1024 / 11.34) // 3 * 3)

tests/test_config.py

@@ -23,3 +23,137 @@ class ParseIds(unittest.TestCase):
     def test_layer_index(self):
         self.assertEqual(mc.layer_index("language_model.model.layers.12.experts"), 12)
         self.assertIsNone(mc.layer_index("model.experts"))
+
+
+class Placement(TmpCase):
+    def test_sections_apply_in_file_order(self):
+        config = mc.load_config(self.write("p.toml", """
+pack_dir = "/packs"
+default_tier = "ram"
+cache_slots = 4
+
+[layers."*"]
+vram = "0-1"
+
+[layers."2-3"]
+disk = "6-7"
+vram = "1-2"
+cache_slots = 2
+"""))
+        p0 = config.placement(0, 8)
+        self.assertEqual((p0.vram, p0.ram, p0.disk, p0.cache_slots),
+                         ((0, 1), (2, 3, 4, 5, 6, 7), (), 4))
+        p3 = config.placement(3, 8)
+        self.assertEqual((p3.vram, p3.ram, p3.disk, p3.cache_slots),
+                         ((0, 1, 2), (3, 4, 5), (6, 7), 2))
+        self.assertEqual(p3.num_slots, 5)
+
+    def test_rejects_typos_conflicts_and_missing_slots(self):
+        with self.assertRaisesRegex(ValueError, "unknown keys"):
+            mc.load_config(self.write("a.toml", 'pack_dir = "/p"\ncache_slot = 3\n'))
+        with self.assertRaisesRegex(ValueError, "default_tier"):
+            mc.load_config(self.write("b.toml", 'pack_dir = "/p"\ndefault_tier = "gpu"\n'))
+        both = mc.load_config(
+            self.write("c.toml", 'pack_dir = "/p"\n[layers."*"]\nvram = "1"\ndisk = "1"\n'))
+        with self.assertRaisesRegex(ValueError, "both vram and disk"):
+            both.placement(0, 4)
+        no_slots = mc.load_config(self.write("d.toml", 'pack_dir = "/p"\ncache_slots = 0\n'))
+        with self.assertRaisesRegex(ValueError, "cache_slots must be >= 1"):
+            no_slots.placement(0, 4)
+
+    def test_all_in_vram_needs_no_cache_slots(self):
+        config = mc.load_config(
+            self.write("v.toml", 'pack_dir = "/p"\ndefault_tier = "vram"\ncache_slots = 0\n'))
+        self.assertEqual(config.placement(5, 4).num_slots, 4)
+
+    def test_policy_prefetch_and_miss_tables(self):
+        config = mc.load_config(self.write("t.toml", """
+pack_dir = "/p"
+cache_slots = 3
+[policy]
+eviction = "gdsf"
+lfu_half_life = 16
+scan_keep = 1
+[prefetch]
+mode = "pregate"
+min_prob = 0.1
+[miss]
+mode = "auto"
+"""))
+        self.assertEqual((config.policy.eviction, config.policy.lfu_half_life, config.policy.heat),
+                         ("gdsf", 16, "weight"))
+        self.assertEqual((config.prefetch.mode, config.prefetch.distance, config.miss.mode),
+                         ("pregate", 1, "auto"))
+        self.assertEqual(mc.load_config(self.write("u.toml", 'pack_dir = "/p"\n')).policy, mc.Policy())
+
+    def test_rejects_bad_table_values(self):
+        for body, error in (('[policy]\neviction = "fifo"', "eviction must be one of"),
+                            ('[policy]\nscan_keep = -1', "scan_keep must be a number"),
+                            ('[prefetch]\ndistance = 3', "distance must be a number in \\[1, 2\\]"),
+                            ('[prefetch]\nbudget = 1.5', "budget must be a number"),
+                            ('[miss]\nmode = "gpu"', "mode must be one of"),
+                            ('[miss]\nthreads = 4', "unknown keys")):
+            with self.assertRaisesRegex(ValueError, error, msg=body):
+                mc.load_config(self.write("bad.toml", f'pack_dir = "/p"\n{body}\n'))
+        keep_all = mc.load_config(self.write("k.toml", 'pack_dir = "/p"\ncache_slots = 2\n'
+                                                       '[policy]\nscan_keep = 2\n'))
+        with self.assertRaisesRegex(ValueError, "scan_keep must be < cache_slots"):
+            keep_all.placement(0, 4)
+
+    def test_paths_are_relative_to_the_config_file(self):
+        config = mc.load_config(self.write("r.toml", 'pack_dir = "packs/x"\nstats_path = "~/s.json"\n'
+                                                     'trace_path = "/abs/t.gz"\n'))
+        self.assertEqual(config.pack_dir, os.path.join(self.dir, "packs/x"))
+        self.assertEqual(config.stats_path, os.path.expanduser("~/s.json"))
+        self.assertEqual(config.trace_path, "/abs/t.gz")
+
+    def test_serve_and_draft_tables(self):
+        config = mc.load_config(self.write("s.toml", 'pack_dir = "/p"\n[serve]\nmax_requests = 4\n'
+                                                     'kv_cache_gib = 1.5\n[draft]\nfull_requests = 2\n'
+                                                     'refine = "none"\n'))
+        self.assertEqual(config.serve, mc.Serve(max_requests=4, kv_cache_gib=1.5))
+        self.assertEqual(config.draft, mc.Draft(full_requests=2, min_weight=0.5, refine="none"))
+        self.assertEqual(mc.load_config(self.write("d.toml", 'pack_dir = "/p"\n')).serve.max_requests, 2)
+        for body, error in (('[serve]\nmax_requests = 0', "max_requests must be a number"),
+                            ('[draft]\nmin_weight = 1.5', "min_weight must be a number in \\[0, 1\\]"),
+                            ('[draft]\nfull_requests = 3', "full_requests must be <= \\[serve\\] max_requests \\(2\\)")):
+            with self.assertRaisesRegex(ValueError, error, msg=body):
+                mc.load_config(self.write("bad.toml", f'pack_dir = "/p"\n{body}\n'))
+
+
+class Plan(TmpCase):
+    def test_hottest_to_vram_then_ram_then_disk(self):
+        stats = {"layers": {"0": {"uses": [5, 0, 9, 1]}, "1": {"uses": [2, 7, 0, 0]}}}
+        # 1 GiB experts: 4 GiB VRAM = 2 cache slots + 2 pinned, 3 GiB RAM = 3 experts
+        text = mc.plan(stats, 4, mc.GIB, 4, 3, 1, "/packs", "/s.json")
+        self.assertEqual(text, """\
+# generated by `prain plan`; experts not listed live on disk
+pack_dir = "/packs"
+default_tier = "disk"
+cache_slots = 1
+stats_path = "/s.json"
+
+[layers."0"]
+vram = "2"
+ram = "0,3"
+
+[layers."1"]
+vram = "1"
+ram = "0"
+""")
+        p0 = mc.load_config(self.write("plan.toml", text)).placement(0, 4)
+        self.assertEqual((p0.vram, p0.ram, p0.disk), ((2,), (0, 3), (1,)))
+
+    def test_budget_smaller_than_cache_slots(self):
+        with self.assertRaisesRegex(ValueError, "smaller than the cache slots"):
+            mc.plan({"layers": {"0": {"uses": [1]}}}, 1, mc.GIB, 1, 0, 2, "/p", "/s")
+
+    def test_render_toml_round_trips_tables(self):
+        text = mc.render_toml("h", {"pack_dir": "/p", "stats_path": None},
+                              {"policy": {"eviction": "lru", "lfu_half_life": 8.0}},
+                              {3: {"cache_slots": 2}, 4: {}})
+        config = mc.load_config(self.write("r.toml", text))
+        self.assertEqual((config.policy.eviction, config.policy.lfu_half_life), ("lru", 8.0))
+        self.assertEqual(config.placement(3, 4).cache_slots, 2)
+        self.assertNotIn("stats_path", text)
+        self.assertNotIn('"4"', text)

tests/test_drafts.py

@@ -0,0 +1,42 @@
+import unittest
+from types import SimpleNamespace
+
+from prain import drafts
+from tests.helpers import TmpCase  # noqa: F401  (every test module imports the helpers)
+
+
+def new_request(req_id, xargs):
+    return SimpleNamespace(req_id=req_id, sampling_params=SimpleNamespace(extra_args=xargs))
+
+
+def step(new=(), finished=()):
+    return SimpleNamespace(scheduled_new_reqs=list(new), finished_req_ids=set(finished))
+
+
+class Drafts(unittest.TestCase):
+    def setUp(self):
+        drafts._DRAFTS.clear()
+        drafts.clear()
+
+    def test_draft_rows_follow_the_batch_layout(self):
+        # a: 3 prompt tokens (full); b: 2 prompt tokens (draft)
+        drafts.note_step(step([new_request("a", None), new_request("b", {"prain": "draft"})]),
+                         ["a", "b"], [0, 3, 5])
+        self.assertEqual(drafts.spans(), [(3, 5)])
+        # next step: both decode one token each, in a new order
+        drafts.note_step(step(), ["b", "a"], [0, 1, 2])
+        self.assertEqual(drafts.spans(), [(0, 1)])
+        drafts.clear()
+        self.assertIsNone(drafts.spans())
+
+    def test_finished_drafts_are_forgotten(self):
+        drafts.note_step(step([new_request("b", {"prain": "draft"})]), ["b"], [0, 4])
+        drafts.note_step(step(finished=["b"]), [], [0])
+        self.assertIsNone(drafts.spans())
+        self.assertEqual(drafts._DRAFTS, set())
+
+    def test_requests_without_the_flag_are_full(self):
+        drafts.note_step(step([new_request("a", {"other": 1}),
+                               SimpleNamespace(req_id="c", sampling_params=None)]),
+                         ["a", "c"], [0, 1, 2])
+        self.assertIsNone(drafts.spans())

tests/test_policy.py

@@ -0,0 +1,121 @@
+import json
+import math
+import os
+import random
+import unittest
+
+from prain.config import Miss, Policy
+from prain.policy import SlotCache, write_stats
+from tests.helpers import TmpCase
+
+
+def lru(**kw):
+    return Policy(eviction="lru", **kw)
+
+
+class Eviction(unittest.TestCase):
+    def test_pinned_experts_stay_and_lru_evicts_the_oldest(self):
+        cache = SlotCache(pinned=[3, 5], num_experts=8, cache_slots=2, policy=lru())
+        self.assertEqual(cache.plan({0: 1, 3: 1}), [({0: 2, 3: 0}, [(0, 2)])])
+        self.assertEqual(cache.plan({0: 1, 1: 1}), [({0: 2, 1: 3}, [(1, 3)])])
+        self.assertEqual(cache.plan({1: 1, 2: 1}), [({1: 3, 2: 2}, [(2, 2)])])  # 0 was oldest
+        self.assertEqual((cache.hits, cache.misses), (3, 3))
+        self.assertEqual(cache.uses, [2, 2, 1, 1, 0, 0, 0, 0])
+
+    def test_lfu_keeps_the_frequent_expert_lru_the_recent_one(self):
+        for policy, survivor in ((lru(), 1), (Policy(eviction="lfu"), 0)):
+            cache = SlotCache([], 3, 2, policy)
+            for e in (0, 0, 0, 1, 2):
+                cache.plan({e: 1})
+            self.assertIn(survivor, cache.slots, policy.eviction)
+
+    def test_gdsf_keeps_the_expensive_disk_expert(self):
+        cost = [1.0, 4.0, 1.0]  # expert 1 lives on disk
+        for eviction, survivor in (("lfu", 0), ("gdsf", 1)):
+            cache = SlotCache([], 3, 2, Policy(eviction=eviction), cost)
+            for e in (0, 0, 1, 2):
+                cache.plan({e: 1})
+            self.assertIn(survivor, cache.slots, eviction)
+
+
+class Rounds(unittest.TestCase):
+    def test_hottest_experts_are_copied_last_and_stay(self):
+        cache = SlotCache([], 10, 2, lru())
+        rounds = cache.plan({0: 1, 1: 9, 2: 3, 3: 8, 4: 2, 5: 1})
+        self.assertEqual([sorted(slot_of) for slot_of, _ in rounds], [[0, 5], [2, 4], [1, 3]])
+        self.assertEqual(set(cache.slots), {1, 3})
+
+    def test_scan_keeps_the_protected_expert_and_promotes_the_hottest(self):
+        cache = SlotCache([], 6, 3, lru(scan_tokens=4, scan_keep=1, promote=1))
+        cache.plan({0: 1})
+        cache.plan({1: 1})
+        rounds = cache.plan({2: 1, 3: 5, 4: 2, 5: 3}, tokens=8)
+        self.assertEqual(rounds, [({2: 2, 4: 0}, [(2, 2), (4, 0)]),
+                                  ({5: 2, 3: 0}, [(5, 2), (3, 0)])])
+        self.assertEqual(set(cache.slots), {1, 3, 5})  # 1 was protected
+        self.assertEqual(cache.evictor.last[3], cache.now)  # promoted
+        self.assertEqual(cache.evictor.last[5], -math.inf)  # admitted at low priority
+
+    def test_cpu_experts_get_no_slot(self):
+        cache = SlotCache([0], 6, 2, lru())
+        cache.plan({1: 1})
+        counts = {0: 1, 1: 1, 2: 1, 3: 4}
+        self.assertEqual(cache.choose_cpu(counts, Miss()), set())
+        self.assertEqual(cache.choose_cpu(counts, Miss(mode="cpu")), {2, 3})
+        cpu = cache.choose_cpu(counts, Miss(mode="auto", cpu_max_tokens=2))
+        self.assertEqual(cpu, {2})
+        self.assertEqual(cache.plan(counts, 4, cpu), [({0: 0, 1: 1, 3: 2}, [(3, 2)])])
+        self.assertEqual((cache.cpu, cache.misses), (1, 2))
+
+    def test_random_plans_always_compute_every_expert_from_the_right_slot(self):
+        rng = random.Random(1)
+        for _ in range(300):
+            n = rng.randint(2, 12)
+            pinned = rng.sample(range(n), rng.randint(0, n // 2))
+            slots = rng.randint(1, 4)
+            policy = Policy(eviction=rng.choice(["lru", "lfu", "gdsf"]),
+                            scan_tokens=rng.randint(1, 6), scan_keep=rng.randint(0, slots - 1),
+                            promote=rng.randint(0, 3))
+            cache = SlotCache(pinned, n, slots, policy, [rng.choice([1.0, 4.0]) for _ in range(n)])
+            content = {slot: e for e, slot in cache.pinned.items()}  # what each slot really holds
+            for _ in range(12):
+                if rng.random() < 0.3:
+                    guesses = [(e, 0.5) for e in rng.sample(range(n), rng.randint(1, n))]
+                    recent = set(cache.recent)
+                    for e, slot in cache.prefetch(guesses, rng.randint(0, 3)):
+                        self.assertNotIn(content.get(slot), recent - {None})
+                        content[slot] = e
+                routed = rng.sample(range(n), rng.randint(1, n))
+                heat = {e: rng.randint(1, 9) for e in routed}
+                cpu = cache.choose_cpu(heat, Miss(mode=rng.choice(["copy", "auto"]), cpu_max_tokens=3))
+                computed = []
+                for slot_of, loads in cache.plan(heat, rng.randint(1, 8), cpu):
+                    self.assertEqual(len(set(slot_of.values())), len(slot_of))
+                    for e, slot in loads:
+                        self.assertEqual(slot_of[e], slot)
+                        content[slot] = e
+                    for e, slot in slot_of.items():
+                        self.assertEqual(content[slot], e)
+                    computed += slot_of
+                self.assertEqual(sorted(computed + sorted(cpu)), sorted(routed))
+                held = list(cache.slots.values())
+                self.assertEqual(len(set(held)), len(held))
+                self.assertEqual(len(held) + len(cache.free), slots)
+
+
+class Prefetch(TmpCase):
+    def test_prefetch_is_low_priority_and_never_evicts_the_last_forward(self):
+        cache = SlotCache([], 4, 2, Policy(eviction="lfu"))
+        cache.plan({0: 1})
+        self.assertEqual(cache.prefetch([(1, 0.5), (2, 0.4)], budget=2), [(1, 1)])  # 0 is recent
+        cache.plan({1: 1})  # the prefetch paid off
+        self.assertEqual(cache.prefetch([(3, 0.9)], budget=1), [(3, 0)])  # evicts 0, not 1
+        cache.plan({2: 1})  # 3 was never used: it goes first
+        self.assertEqual((cache.prefetches, cache.prefetch_hits, cache.prefetch_wasted), (2, 1, 1))
+        self.assertEqual(set(cache.slots), {1, 2})
+        path = os.path.join(self.dir, "stats.json")
+        write_stats(path, {7: cache})
+        with open(path) as f:
+            stats = json.load(f)["layers"]["7"]
+        self.assertEqual((stats["uses"], stats["prefetch_wasted"], stats["misses"]),
+                         ([1, 1, 1, 0], 1, 2))

tests/test_sim.py

@@ -0,0 +1,83 @@
+import os
+import unittest
+
+from prain.config import LayerPlacement, Miss, Policy, Prefetch
+from prain.sim import (Cost, read_forwards, simulate, synthetic_trace, tune, write_trace)
+from tests.helpers import TmpCase
+
+SMALL = dict(layers=4, experts=32, top_k=4, forwards=120, prompt_every=60, prompt_tokens=32)
+
+
+def forwards_of(records):
+    forwards, last = [], None
+    for r in records:
+        if last is None or r["l"] <= last:
+            forwards.append({})
+        forwards[-1][r["l"]] = r
+        last = r["l"]
+    return forwards
+
+
+def all_ram(slots, layers=4, experts=32):
+    return {layer: LayerPlacement((), tuple(range(experts)), (), slots) for layer in range(layers)}
+
+
+class Traces(TmpCase):
+    def test_round_trip_and_determinism(self):
+        records = synthetic_trace(**SMALL)
+        self.assertEqual(records, synthetic_trace(**SMALL))
+        path = os.path.join(self.dir, "t.jsonl.gz")
+        write_trace(path, records[:8])
+        write_trace(path, records[8:], mode="at")  # runs append to the same trace
+        forwards = read_forwards(path)
+        self.assertEqual(len(forwards), 120)
+        self.assertEqual(forwards[0][0], records[0])
+        self.assertEqual(forwards[0][0]["t"], 32)
+        self.assertEqual(sorted(forwards[1][1]["p"]), ["1", "2"])
+        self.assertNotIn("p", forwards[1][3])  # nothing after the last layer
+
+    def test_copy_cost(self):
+        self.assertAlmostEqual(Cost().copy_ms(False), 0.991, places=3)
+        self.assertAlmostEqual(Cost().copy_ms(True), 4.955, places=3)
+
+
+class Simulate(unittest.TestCase):
+    def setUp(self):
+        self.forwards = forwards_of(synthetic_trace(**SMALL))
+
+    def run_sim(self, policy=Policy(), prefetch=Prefetch(), miss=Miss(), slots=3, **kw):
+        return simulate(self.forwards, all_ram(slots), policy, prefetch, miss, Cost(), **kw)
+
+    def test_belady_bounds_every_policy(self):
+        bound = self.run_sim(belady=True).hit_rate
+        for eviction in ("lru", "lfu", "gdsf"):
+            self.assertLessEqual(self.run_sim(Policy(eviction=eviction)).hit_rate, bound)
+
+    def test_cpu_path_beats_copies_when_the_cpu_is_faster_than_pcie(self):
+        copy = self.run_sim().ms_per_decode
+        self.assertLess(self.run_sim(miss=Miss(mode="auto")).ms_per_decode, copy)
+
+    def test_prefetch_helps_when_predictions_are_right(self):
+        none = self.run_sim().hit_rate
+        good = self.run_sim(prefetch=Prefetch(mode="pregate", budget=4))
+        self.assertGreater(good.hit_rate, none)
+        self.assertGreater(good.counters["prefetch_hits"], good.counters["prefetch_wasted"])
+
+    def test_pregate_needs_predictions(self):
+        stripped = [{layer: {k: v for k, v in r.items() if k != "p"} for layer, r in f.items()}
+                    for f in self.forwards]
+        with self.assertRaisesRegex(ValueError, "needs a trace with predictions"):
+            simulate(stripped, all_ram(3), Policy(), Prefetch(mode="pregate"), Miss(), Cost())
+
+
+class Tune(unittest.TestCase):
+    def test_tune_beats_the_defaults_within_the_budget(self):
+        forwards = forwards_of(synthetic_trace(**SMALL))
+        cost = Cost(expert_mib=64.0)  # 1 GiB = 16 slots = 4 per layer
+        best, ram, per, final, bound, rows = tune(forwards, 32, 1.0, 1.0, cost)
+        self.assertLessEqual(sum(per.values()), 16)
+        self.assertEqual(sum(len(r) for r in ram.values()), 16)  # 1 GiB of RAM tier
+        default = next(r for v, r in rows if v == (Policy(), Prefetch(), Miss()))
+        self.assertLessEqual(final.total_ms, default.total_ms)
+        self.assertEqual(rows[0][1].total_ms, min(r.total_ms for _, r in rows))
+        self.assertGreaterEqual(bound.hit_rate, 0.0)

tests/test_tui.py

@@ -0,0 +1,70 @@
+import fcntl
+import os
+import pty
+import struct
+import subprocess
+import sys
+import termios
+import threading
+import unittest
+
+from prain.tui import Chat
+from tests.helpers import StubVllm, wait_for
+
+
+class Tui(unittest.TestCase):
+    def setUp(self):
+        self.stub = StubVllm()
+        self.addCleanup(self.stub.close)
+
+    def test_typing_warms_up_and_enter_streams_the_answer(self):
+        chat = Chat(self.stub.url, "stub-model")
+        chat.type("Why is the sky")
+        wait_for(lambda: self.stub.bodies and not chat.warmer.busy)
+        self.assertEqual(self.stub.bodies[0]["priority"], 1)
+        chat.send()
+        chat.type("next question")  # an answer is streaming: no warm-up in "idle" mode
+        wait_for(lambda: chat.poll() or chat.reply is None)
+        self.assertEqual(chat.messages, [{"role": "user", "content": "Why is the sky"},
+                                         {"role": "assistant", "content": "Hello!"}])
+        self.assertEqual(len(self.stub.bodies), 2)
+        self.assertIn("tok/s", chat.status)
+
+    def test_warmup_off_and_cancel_on_send(self):
+        off = Chat(self.stub.url, "stub-model", warmup="off")
+        off.type("hello")
+        self.assertFalse(off.warmer.busy)
+        self.stub.delay = 3.0
+        chat = Chat(self.stub.url, "stub-model")
+        chat.type("a slow warm-up")
+        wait_for(lambda: self.stub.bodies)
+        chat.send()  # aborts the warm-up instead of queueing behind it
+        wait_for(lambda: not chat.warmer.busy, timeout=2.0)
+
+    def test_curses_screen_runs_in_a_real_terminal(self):
+        master, slave = pty.openpty()
+        fcntl.ioctl(slave, termios.TIOCSWINSZ, struct.pack("HHHH", 24, 80, 0, 0))
+        tui = subprocess.Popen([sys.executable, "-m", "prain.cli", "tui", "--url", self.stub.url],
+                               stdin=slave, stdout=slave, stderr=slave,
+                               env={**os.environ, "TERM": "xterm"}, start_new_session=True)
+        os.close(slave)
+        screen = bytearray()
+        threading.Thread(target=self._drain, args=(master, screen), daemon=True).start()
+        wait_for(lambda: b"stub-model" in screen)  # status line drawn
+        os.write(master, b"hi\r")
+        wait_for(lambda: any(body.get("stream") and "priority" not in body
+                             for body in self.stub.bodies))
+        wait_for(lambda: b"tok/s" in screen)
+        os.write(master, b"\x1b")
+        self.assertEqual(tui.wait(timeout=5), 0)
+
+    @staticmethod
+    def _drain(fd, screen):
+        while True:
+            try:
+                data = os.read(fd, 4096)
+            except OSError:
+                return
+            if not data:
+                return
+            screen.extend(data)

tests/test_web.py

@@ -0,0 +1,138 @@
+import json
+import threading
+import unittest
+import urllib.error
+import urllib.request
+
+import prain.chat as chat
+from prain.config import Draft, Serve
+from prain.web import make_server
+from tests.helpers import StubVllm, wait_for
+
+
+def events(response):
+    """The proxy's stream as a list: content pieces, and "<draft>"-style markers."""
+    out = []
+    for raw in response:
+        line = raw.decode().strip()
+        if line == "data: [DONE]":
+            out.append("<done>")
+        elif line.startswith("data: "):
+            data = json.loads(line[6:])
+            out.append(f"<{data['prain']}>" if "prain" in data
+                       else data["choices"][0]["delta"]["content"])
+    return out
+
+
+class Web(unittest.TestCase):
+    def setUp(self):
+        self.stub = StubVllm()
+        self.addCleanup(self.stub.close)
+        self.start()
+
+    def start(self, **options):
+        self.server = make_server(self.stub.url, 0, **options)
+        threading.Thread(target=self.server.serve_forever, daemon=True).start()
+        self.addCleanup(self.server.server_close)
+        self.addCleanup(self.server.shutdown)
+        self.url = f"http://127.0.0.1:{self.server.server_port}"
+
+    def warm(self, session, draft):
+        chat.post(self.url, "/warm", {"session": session, "messages": [], "draft": draft}).read()
+
+    def ask(self, session, text):
+        return chat.post(self.url, "/chat", {"session": session,
+                                             "messages": [{"role": "user", "content": text}]})
+
+    def test_page_warm_up_and_chat(self):
+        with urllib.request.urlopen(self.url) as response:
+            self.assertIn(b'<textarea id="draft"', response.read())
+        self.warm("s1", "Why is ")
+        wait_for(lambda: self.stub.bodies)
+        self.assertEqual(self.stub.bodies[0]["messages"], [{"role": "user", "content": "Why is"}])
+        self.assertEqual(self.stub.bodies[0]["priority"], 1)
+        with self.ask("s1", "Why is") as r:
+            self.assertEqual(events(r), ["Hel", "lo", "!", "<done>"])
+        self.assertEqual(self.stub.bodies[1]["model"], "stub-model")
+        self.assertNotIn("vllm_xargs", self.stub.bodies[1])
+
+    def test_sessions_warm_up_independently(self):
+        self.stub.delay = 0.3
+        self.warm("alice", "first")
+        self.warm("bob", "second")  # alice's warm-up is in flight; bob's is not dropped
+        wait_for(lambda: len(self.stub.bodies) == 2)
+        self.assertEqual({b["messages"][-1]["content"] for b in self.stub.bodies},
+                         {"first", "second"})
+        self.assertEqual(set(self.server.sessions), {"alice", "bob"})
+
+    def test_no_warm_up_while_another_answer_streams(self):
+        self.server.streams[0] = 1  # someone else's answer is streaming
+        self.warm("bob", "typing along")
+        self.server.streams[0] = 0
+        self.warm("bob", "now")
+        wait_for(lambda: self.stub.bodies)
+        self.assertEqual([b["messages"][-1]["content"] for b in self.stub.bodies], ["now"])
+
+    def test_always_mode_warms_up_while_a_request_slot_is_free(self):
+        self.start(warmup="always", serve=Serve(max_requests=3))
+        self.server.streams[0] = 3  # every vLLM request slot is busy
+        self.warm("bob", "no slot")
+        self.server.streams[0] = 2
+        self.warm("bob", "one slot")
+        wait_for(lambda: self.stub.bodies)
+        self.assertEqual([b["messages"][-1]["content"] for b in self.stub.bodies], ["one slot"])
+
+    def test_vllm_errors_reach_the_browser(self):
+        self.stub.fail = 400
+        with self.assertRaises(urllib.error.HTTPError) as caught:
+            chat.post(self.url, "/chat", {"session": "s", "messages": []})
+        self.assertEqual(caught.exception.code, 400)
+        with caught.exception as error:
+            self.assertIn(b"rejected", error.read())
+        wait_for(lambda: self.server.streams[0] == 0)
+
+    def test_second_answer_drafts_now_and_refines_when_the_first_one_ends(self):
+        self.stub.delay = lambda body: 0.0 if "vllm_xargs" in body else 0.5  # full answers are slow
+        alice = []
+        thread = threading.Thread(target=lambda: alice.extend(events(self.ask("alice", "a"))))
+        thread.start()
+        wait_for(lambda: self.server.full[0] == 1)
+        with self.ask("bob", "b") as r:
+            bob = events(r)
+        thread.join()
+        self.assertEqual(alice, ["Hel", "lo", "!", "<done>"])
+        self.assertEqual(bob, ["<draft>", "Hel", "lo", "!", "<refine>", "Hel", "lo", "!", "<done>"])
+        self.assertEqual([(b["messages"][0]["content"], b.get("vllm_xargs"), b.get("cache_salt"))
+                          for b in self.stub.bodies],
+                         [("a", None, None), ("b", {"prain": "draft"}, "prain-draft"),
+                          ("b", None, None)])
+        wait_for(lambda: self.server.full[0] == 0 and self.server.streams[0] == 0)
+
+    def test_a_running_draft_is_cut_short_once_a_full_slot_frees(self):
+        self.stub.delay = lambda body: 0.0 if "vllm_xargs" in body else 0.3
+        self.stub.pause = lambda body: 0.6 if "vllm_xargs" in body else 0.0  # drafts trickle
+        thread = threading.Thread(target=lambda: self.ask("alice", "a").read())
+        thread.start()
+        wait_for(lambda: self.server.full[0] == 1)
+        with self.ask("bob", "b") as r:  # alice ends at 0.3 s, between bob's "Hel" and "lo"
+            self.assertEqual(events(r), ["<draft>", "Hel", "lo", "<refine>", "Hel", "lo", "!",
+                                         "<done>"])
+        thread.join()
+
+    def test_refine_none_keeps_the_draft_and_more_full_slots_mean_no_draft(self):
+        self.start(draft=Draft(refine="none"))
+        self.stub.delay = lambda body: 0.0 if "vllm_xargs" in body else 0.5
+        thread = threading.Thread(target=lambda: self.ask("alice", "a").read())
+        thread.start()
+        wait_for(lambda: self.server.full[0] == 1)
+        with self.ask("bob", "b") as r:
+            self.assertEqual(events(r), ["<draft>", "Hel", "lo", "!", "<done>"])
+        thread.join()
+        self.start(serve=Serve(max_requests=2), draft=Draft(full_requests=2))
+        thread = threading.Thread(target=lambda: self.ask("alice", "a").read())
+        thread.start()
+        wait_for(lambda: self.server.full[0] == 1)
+        with self.ask("bob", "b") as r:
+            self.assertEqual(events(r), ["Hel", "lo", "!", "<done>"])
+        thread.join()
+        self.assertEqual(sum("vllm_xargs" in b for b in self.stub.bodies), 1)