prain: serve with the string chat template content format

prain web/tui --persona FILE sends the file as the system message ahead of every warm-up and
answer; the page shows its name. Three examples in examples/personas (task, how to answer,
environment, how to use it). A persona costs 2.0-3.2 s once (241-467 tokens), then at most
0.3 s per conversation from the prefix cache.

README rewritten around launching: quick start, the vLLM flags, a quantization table measured
today at equal expert VRAM (bf16 0.90 / int8 2.19 / int4 5.19 tok/s; 2.3k-token prompt 69.5 /
32.3 / 20.5 s; 4 of 6 int8 and 0 of 6 int4 answers identical to bf16, one int4 answer wrong),
how to switch formats, the VRAM budget, writing personas, the chats and the monitor.

README.md

@@ -0,0 +1,359 @@
+# prain
+
+prain runs a mixture-of-experts model with vLLM on a GPU that is too small to hold it. A
+placement file (TOML) says where each expert of each layer lives: in VRAM, pinned or in a cache
+slot; in pinned RAM; or in a pack file on disk. When a token needs an expert that is not in VRAM,
+prain copies it in or computes it on the CPU. Experts can be stored as bf16, int8 or int4.
+
+It is tuned for one machine: Gemma 4 26B-A4B (30 layers x 128 experts, 42.5 GiB of bf16 experts)
+on an 8 GB Radeon RX 6600 XT (RDNA2, gfx1032) in a Thunderbolt enclosure. With int4 experts it
+answers at about 5 tokens/s, with 16k tokens of context and tool calls. Clients talk to it
+through vLLM's OpenAI-compatible API, a browser chat with a live monitor, or a terminal chat. Both
+chats start moving the experts a message needs into VRAM while it is still being typed.
+
+## Quick start
+
+After the build below, start the model with the example placement (int4 experts):
+
+    . scripts/env.sh
+    M=models/gemma-4-26B-A4B-it-Uncensored-MAX
+    FLAGS="--served-model-name gemma4 --language-model-only --max-model-len 16384
+           --max-num-batched-tokens 1024 --enable-prefix-caching --attention-backend TRITON_ATTN
+           --enable-auto-tool-choice --tool-call-parser gemma4"
+    prain serve examples/gemma4-26b-a4b-rx6600xt.toml $M -- $FLAGS
+
+The model is ready when `curl -sf localhost:8000/health && echo ready` prints `ready`. Then, in a
+second terminal:
+
+    . scripts/env.sh
+    prain web --config examples/gemma4-26b-a4b-rx6600xt.toml --host 0.0.0.0 \
+        --persona examples/personas/gentoo-admin.md
+
+and open http://HOST:8080. `--host 0.0.0.0` makes the page reachable from other machines; it has
+no login. Coding agents and other OpenAI clients use http://HOST:8000/v1 with the model name
+`gemma4`.
+
+## Build (GPU host, user space)
+
+    scripts/build.sh wrappers sdk fetch torch torchvision vllm prain
+
+- `sdk`: a venv with AMD's ROCm SDK wheels for `PRAIN_GPU_ARCH` (default gfx1032)
+- `torch`, `torchvision`: built from source with clang and `-march=native` (hours)
+- `vllm`: v0.32.0 from source, with gfx1032 added and `patches/vllm-*.patch` applied
+- `prain`: an editable install, then the test suite
+
+Check the stack with `scripts/smoke.sh MODEL_DIR --attention-backend TRITON_ATTN`. Then tune
+vLLM's fused-MoE kernel for the GPU once (about 45 minutes of Triton compiles the first time;
+`scripts/env.sh` points vLLM at the result). It keeps vLLM's default wherever no candidate beats
+it:
+
+    python scripts/tune_moe.py build/moe-configs
+
+## Serve a model
+
+`prain serve PLACEMENT MODEL_DIR -- VLLM_FLAGS` validates the placement file and starts
+`vllm serve` with prain's plugin. prain adds `--enforce-eager --moe-backend triton
+--scheduling-policy priority --chat-template-content-format string`, and it takes
+`--max-num-seqs` and `--kv-cache-memory-bytes` from the file's `[serve]` table (it refuses them
+after `--`). vLLM would otherwise pick the `openai` content format for Gemma 4, whose template
+then prints a system prompt as a Python list (`[{'type': 'text', 'text': '…\n…'}]`, newlines
+escaped); a flag after `--` still overrides it. The flags in `$FLAGS` above:
+
+| flag | why |
+|---|---|
+| `--served-model-name gemma4` | the model name clients ask for |
+| `--language-model-only` | keeps Gemma's vision and audio towers out of VRAM |
+| `--max-model-len 16384` | the context of one request; `kv_cache_gib` must hold it |
+| `--max-num-batched-tokens 1024` | the prefill chunk (see [Trade-offs](#trade-offs)) |
+| `--enable-prefix-caching` | reuses earlier turns, personas and warm-ups |
+| `--attention-backend TRITON_ATTN` | Triton attention, which runs on RDNA2 |
+| `--enable-auto-tool-choice --tool-call-parser gemma4` | tool calls for agents and MCP clients |
+
+`examples/gemma4-26b-a4b-rx6600xt.toml` is commented line by line. It keeps all experts in the
+pack on disk, served through the page cache, and gives each layer 15 cache slots in VRAM (1.32
+GiB), plus 1 GiB of KV cache for 16k tokens of context. `prain check PLACEMENT
+MODEL_DIR/config.json` prints what a placement costs in VRAM, RAM and disk before anything
+loads.
+
+The first start writes the bf16 pack (42.5 GiB, about 3 minutes); an int8 or int4 placement then
+converts it once on the GPU (about 3 minutes). Stop the server with Ctrl-C. Only one model fits on
+the GPU at a time, so stop it before starting another quantization.
+
+## Choose a quantization
+
+`quant` in the placement file picks the expert format. Smaller experts win twice: more of them
+fit in VRAM, and every miss moves fewer bytes over the 2.5 GB/s Thunderbolt link and out of the
+page cache (on a 31 GiB host the int4 and int8 packs fit in free RAM, the bf16 one does not).
+They cost some accuracy:
+
+| quant | expert | pack on disk | slots per layer in 1.32 GiB | decode | first token, short prompt | first token, 2.3k-token prompt | answers equal to bf16 |
+|---|---|---|---|---|---|---|---|
+| bf16 | 11.34 MiB | 42.5 GiB | 4 | 0.90 tok/s | 14.7 s | 69.5 s | reference |
+| int8 | 5.69 MiB | 21.3 GiB | 8 | 2.19 tok/s | 5.9 s | 32.3 s | 4 of 6 (215 of 288 tokens) |
+| int4 | 3.01 MiB | 11.3 GiB | 15 | 5.19 tok/s | 3.3 s | 20.5 s | 0 of 6 (94 of 288 tokens) |
+
+Measured with the example placement on the RX 6600 XT behind Thunderbolt (PCIe 3.0 x4), an
+i9-13950HX and 31 GiB of RAM, changing only `quant` and `cache_slots`. Decode is the median of
+3 x 48 tokens. "Equal" counts how many of the six `scripts/quality.py` prompts got the same 48
+greedy tokens as bf16; the brackets count the tokens that matched before each answer's first
+difference.
+
+int4 is the default: it decodes 2.4 times as fast as int8 and reads a long prompt in two thirds
+of the time. Its answers are worded differently from bf16's from the first tokens on. Five of
+the six meant the same; the sixth, a Gentoo question, named a tool that does not exist
+("oneway-way") where bf16 and int8 named `revdep-rebuild`. int8 stays closest to bf16 at less
+than half the speed of int4. bf16 is mostly the reference for quality checks: its 4 slots per
+layer are fewer than the 8 experts each token uses, so every decode step runs in several rounds
+(`prain check` warns about this).
+
+To launch int8 instead of int4, copy the example and change two lines. Relative paths in a
+placement file start at the file's directory, so keep the copy next to the example:
+
+    cp examples/gemma4-26b-a4b-rx6600xt.toml examples/int8.toml
+    sed -i -e 's/^quant = "int4"/quant = "int8"/' -e 's/^cache_slots = 15 /cache_slots = 8 /' \
+        examples/int8.toml
+    prain check examples/int8.toml $M/config.json     # VRAM 1.33 GiB, as before
+    prain serve examples/int8.toml $M -- $FLAGS
+
+For bf16, set `cache_slots = 4`. Keep slots x expert size x 30 layers near 1.32 GiB. On the
+example placement vLLM reports "Model loading took 6.29 GiB" (the dense weights, 1.32 GiB of
+expert slots and prain's buffers), the KV cache takes 1 GiB (`kv_cache_gib`), and about 0.55 GiB
+of the card's 7.98 GiB stays free for activations and kernels. prain logs that free VRAM at
+startup ("MiB of VRAM free for kernels and activations", before the KV cache takes its share).
+More VRAM for experts means less for context, and the other way round.
+
+`prain tune` fits per-layer slot counts and policies to a routing trace, which the example records
+while it serves (`trace_path`). The tuned int4 placement of the live agent setup (450 slots,
+12-18 per layer) decodes at 5.2 tok/s. For another format:
+
+    prain tune models/packs/gemma4-26b-a4b/trace.jsonl.gz $M/config.json --quant int8 \
+        --expert-vram-gib 1.32 --ram-gib 0 --pack-dir models/packs/gemma4-26b-a4b > int8-tuned.toml
+
+Every pack stays on disk once written: `layerNNN.w13`, `.w2` and `.done` for bf16,
+`layerNNN.int8.*` and `layerNNN.int4.*` for the others. Delete the formats you do not use; the
+next start of that format writes its pack again, and int8 or int4 need the bf16 pack for that.
+
+How the formats are stored and computed:
+
+| quant | weights | GPU kernel | CPU kernel (`[miss] mode = "cpu"`) |
+|---|---|---|---|
+| bf16 | as trained | vLLM Triton | torch bf16 matmul |
+| int8 | fp32 scale per row | vLLM Triton (w8a16) | `torch._weight_int8pack_mm` |
+| int4 | bf16 scale per 64 weights, MSE-clipped | prain's on RDNA2, else vLLM Triton (w4a16) | `torch._weight_int4pack_mm_for_cpu` |
+
+## Give the model a persona
+
+A persona is a short text file that `prain web` and `prain tui` send as the system message at
+the start of every conversation. It tells the model its task, how to handle questions, what its
+environment is and how to use it. One running model can serve several personas at once, one
+chat front end each:
+
+    P=examples/gemma4-26b-a4b-rx6600xt.toml
+    prain web --config $P --port 8080 --persona examples/personas/gentoo-admin.md --host 0.0.0.0
+    prain web --config $P --port 8081 --persona examples/personas/pair-programmer.md --host 0.0.0.0
+    prain tui --persona examples/personas/explainer.md
+
+The page shows the persona's name, its file name, next to the model's. A new persona needs no
+model restart: edit the file and restart `prain web` or `prain tui`. Without `--persona` the
+model gets no system message.
+
+A persona has four parts. `examples/personas/gentoo-admin.md`, shortened:
+
+    You help one person run their Gentoo Linux desktop.
+
+    ## Task
+    Answer questions about this machine and fix problems together with its owner: installing and
+    updating packages, USE flags, services, the kernel, drivers, disks and the Hyprland desktop.
+
+    ## How to answer
+    - Lead with the command or the fix, then explain it in a few sentences.
+    - When the answer depends on this machine (USE flags, profile, kernel, installed versions),
+      ask for the output of a specific command instead of guessing.
+
+    ## Environment
+    - Gentoo Linux with Portage and OpenRC. There is no systemd and no journalctl.
+    - You cannot run anything yourself. The owner types your commands into a root shell and
+      pastes the output back into the chat.
+
+    ## How to use it
+    - Preview before changing: `emerge --pretend --verbose PKG`, then `emerge --ask PKG`.
+
+The other examples are `pair-programmer.md` (reviews and debugs pasted Python and shell code)
+and `explainer.md` (plain answers for beginners, often read on a phone). When you write one:
+
+- Address the model as "you" and state what it cannot know: the system, the tools at hand, and
+  what it can and cannot do itself.
+- Say how to handle missing information, for example by asking for a specific command's output
+  instead of guessing.
+- Keep it short. A new persona text is prefilled once, before its first answer: the examples
+  (241-467 tokens) added 2.0-3.2 s to that answer on the int4 agent setup, and at most 0.3 s to
+  the next conversations, which find it in vLLM's prefix cache. The chats also send it with
+  their warm-ups while you type.
+
+Agents and other API clients send their own system message. To give one a persona, put the
+file's text there:
+
+    jq -n --rawfile p examples/personas/explainer.md \
+        '{model: "gemma4", messages: [{role: "system", content: $p},
+                                      {role: "user", content: "What is a page cache?"}]}' |
+        curl -s localhost:8000/v1/chat/completions -H "Content-Type: application/json" -d @-
+
+## Chat in the browser
+
+    prain web --config PLACEMENT [--persona FILE] [--host 0.0.0.0] [--port 8080]
+
+The page renders Markdown answers with code blocks you can copy. It has a stop button and "New
+chat", and it keeps the conversation in the browser until you start a new one. Enter sends and
+Shift+Enter adds a line; on a phone, Return adds a line and the button sends.
+
+While you type, the page sends the unfinished message as a warm-up: vLLM prefills it and prain
+moves the experts it routes to into VRAM, so the answer starts sooner. "Loading experts for your
+message" under the box shows that one runs. `--warmup idle` (the default) warms up only while no
+answer streams, `always` also during one, `off` never.
+
+Several people can chat at once. `[serve] max_requests` answers run in parallel, and
+`[draft] full_requests` of them with all their experts. Later ones start at once as drafts from
+the experts already in VRAM and are regenerated with all experts when the GPU is free (`[draft]
+refine`); the page shows a draft in grey until then.
+
+The monitor beside the chat (behind the status pill on a phone) shows:
+
+- tokens per second, now or of the last answer, and a line for the last minute
+- the expert map: one row per layer, one column per expert. Violet experts are in VRAM, orange
+  to yellow ones were used in the last seconds, cyan ones were just copied in. It needs
+  `stats_path` in the placement file; `stats_every = 8` (forward passes) refreshes it about every
+  1.5 s while the model writes.
+- VRAM hits (routed experts that were already in VRAM, last minute) and copies per second
+- context in use (the KV cache), prompt cache reuse, and requests running and waiting
+- the GPU's VRAM, load, temperature, power and fan (from amdgpu's sysfs files)
+
+## Chat in the terminal
+
+    prain tui [--persona FILE] [--url http://127.0.0.1:8000]
+
+The same warm-ups as the page. Enter sends, Esc quits.
+
+## Coding agents and long context
+
+Gemma 4 has 25 sliding-window layers (window 1024) and 5 full-attention layers, so its KV cache
+grows slowly with context: ~20 KiB per token per request, plus ~0.6 GiB for the sliding layers,
+which also hold the tokens in flight (2 x `--max-num-batched-tokens`). `kv_cache_gib = 1.0`
+holds 18,169 tokens, 1.11 requests of 16k (measured).
+
+Every 1024-token prefill chunk needs nearly all experts, copied over the 2.5 GB/s link: a fresh
+2.3k-token prompt takes about 20 s. Before the int4 kernels, which cut that by a quarter, a cold
+4.7k-token prompt took 70 s and a 12k-token one 176 s. Agent turns reuse the prefix cache: turns
+2 and 3 of a conversation found 4,704 and 4,736 of ~4,750 prompt tokens cached and took ~20 s
+each (24-token answers). vLLM keeps sliding-window state only where earlier prompts ended, so the
+second conversation that shares only the system prompt recomputes it once (176 s); later ones hit
+(4.5 s).
+
+## Trade-offs
+
+`[miss] min_weight` (default 1) is the share of each token's router weight that must be
+computed. Below 1, a token skips its lightest experts that are not in VRAM once it has that
+share, and its other routes are scaled up to its full weight: fewer copies, a slightly different
+model. Measured on the int4 agent setup before the int4 kernels; quality is the mean negative
+log-likelihood of the min_weight = 1 answers to the six `scripts/quality.py` prompts (0.11 is
+that config scoring its own answers):
+
+| min_weight | answer NLL | decode tok/s | first token, short prompt | first token, 2.5k-token prompt |
+|---|---|---|---|---|
+| 1.0 | 0.11 | 2.77 | 4.6 s | 28.5 s |
+| 0.6 | 0.19 | 3.15 | 2.7 s | 23.0 s |
+| 0.4 | 0.40 | 3.36 | 2.0 s | 19.9 s |
+
+At 0.6 the answers stay correct but take other words; at 0.4 they show slips ("Hereon any of the
+following one-liners"). Pinning only the hottest experts and never copying (`prain plan
+--cache-slots 1`, min_weight 0) broke the model (NLL 9) for 25% more decode speed (measured
+before the dense-layer fix).
+
+Bigger prefill chunks trade decode speed for prompt speed: at 16k, `--max-num-batched-tokens 2048`
+needs `kv_cache_gib = 1.3` (the sliding layers hold two chunks), which leaves 348 slots instead
+of 450. A 2.3k-token prompt then took 17.4 s instead of 20.6 s, and decode ran at 4.5 instead of
+5.2 tok/s. Long agent prompts gain more: half the chunks, half the expert copies.
+
+## Why it is fast on RDNA2
+
+rocBLAS runs bf16 GEMVs at 13-16 GB/s on gfx1032 (RDNA2 has no bf16 dot instructions), and vLLM
+uses its own skinny-GEMM kernels only on gfx9/gfx11/gfx12. A decode step spent ~350 ms reading
+the dense weights (lm_head alone: 93 ms). On gfx10 prain sends every unquantized GEMM with at
+most 4 rows through a Triton GEMV (`prain/gemv.py`, ~210 GB/s; `PRAIN_GEMV=0` turns it off):
+decode went from 1.58 to 2.77 tok/s with unchanged answers. Prefill keeps rocBLAS.
+
+vLLM's int4 MoE kernel pads every token to a 16-row tile and multiplies in bf16: a decode round
+took 2.7 ms for one token. On gfx10 prain runs int4 rounds of up to 4 tokens through its own
+GEMV-style kernel (`prain/kernels.py`, 0.25 ms) and larger rounds through vLLM's kernel in fp16,
+twice as fast as bf16 there (a round that overflows fp16 is redone in bf16). Every int4 expert
+copy goes through a pinned staging buffer: one DMA and one unpack kernel instead of four pageable
+copies and a ten-kernel repack, asynchronous from every tier.
+
+A/B on the live agent setup (int4, 16k context, 1 GiB KV, 450 tuned slots); quality is the
+greedy answers to the six `quality.py` prompts (identical to the reference, mean logprob of their
+tokens):
+
+| | decode tok/s | first token, short | first token, 2.3k tokens | identical | mean logprob |
+|---|---|---|---|---|---|
+| before | 3.49 | 4.3 s | 27.5 s | 4 of 6 | -0.104 |
+| decode kernel | 4.73 | 4.2 s | 27.3 s | 4 of 6 | -0.093 |
+| staged copies | 3.67 | 3.7 s | 24.5 s | 4 of 6 | -0.098 |
+| fp16 rounds | 3.84 | 4.0 s | 23.3 s | 4 of 6 | -0.102 |
+| all three | 5.24 | 3.3 s | 20.6 s | 5 of 6 | -0.104 |
+
+Repeated runs of one config differ by up to 5% in decode speed and 1% in prefill. Measured and
+dropped: routing statistics on the CPU and a Triton GEMV for the fp32 router (within noise), and
+`[prefetch] mode = "pregate"`, which can now prefetch disk-tier experts too (4.9 vs 5.2 tok/s:
+the link is already busy with demand copies).
+
+## Tune for a machine
+
+    python scripts/measure.py PACK_DIR/layer029.w13 --quant int4 --threads 16 > cost.json  # server stopped
+    prain sim TRACE PLACEMENT MODEL_DIR/config.json --cost cost.json
+    prain tune TRACE MODEL_DIR/config.json --cost cost.json --quant int4 --expert-vram-gib 1.32 \
+        --ram-gib 0 --pack-dir PACK_DIR --cpu-threads 16 > tuned.toml
+    python scripts/bench.py http://127.0.0.1:8000 [--context FILES...]      # speed
+    python scripts/quality.py http://127.0.0.1:8000 make ref.json           # reference answers
+    python scripts/quality.py http://127.0.0.1:8000 score ref.json --greedy-only   # another config
+
+Judge quality on chat answers, not raw-text perplexity: this instruction-tuned model scores raw
+code at a negative log-likelihood of 5-9 per token on a correct engine. The simulator treats
+every disk-tier expert as a disk read, so it under-predicted bf16 (0.6 vs 1.04 tok/s measured;
+the page cache helps) and over-predicted int4 (2.4 vs 1.62 tok/s), both before the dense-layer
+fix and the int4 kernels.
+
+## Profile
+
+Serve with `PRAIN_PROFILE=/abs/dir`. Writing a number N into `dir/start` records the next N
+engine steps: spans around the engine step, every module of the model and prain's MoE phases
+(route, plan, copy: memcpy/h2d/unpack for int4, round), each synced with the GPU. `python -m
+prain.profile dir` writes `decode.svg`/`prefill.svg` flame graphs and prints self time and calls
+per step. The syncs remove the CPU/GPU overlap: a profiled decode step took 395 ms where an
+unprofiled one takes 307 ms.
+
+## Limits
+
+- Requests with `prompt_logprobs` (perplexity tools) materialize [prompt tokens x 262k vocab]
+  fp32 logits that vLLM's memory plan does not cover: on the 16k setup even ~300-token prompts
+  crashed the engine. `scripts/quality.py score --greedy-only` skips them.
+- Pinning host memory was very slow on the Thunderbolt host (minutes per GiB, cause unproven) and
+  pinned copies were only 4% faster than pageable ones, so the example has no RAM tier and lets
+  the page cache hold the pack.
+- `prain web` has no login. Bind it to 127.0.0.1 (the default) on untrusted networks.
+
+## vLLM patches
+
+- `vllm-0.32-rocm-no-cuda-families`: gfx1032 reports capability 10.3, and vLLM's CUDA "family"
+  checks took it for Blackwell (B200 attention tiles: 128 KiB of LDS on a 64 KiB GPU).
+- `vllm-0.32-gemma4-fit-8gb`: Gemma 4 shares `embed_tokens` as `lm_head` instead of allocating a
+  second 1.4 GiB copy while loading, and sizes rope caches by `max_model_len` (saves ~0.4 GiB).
+
+## Repository
+
+| path | what |
+|---|---|
+| `prain/` | the vLLM plugin (`plugin.py`, `tiered.py`, `experts.py`, `drafts.py`), RDNA2 kernels (`kernels.py`, `gemv.py`), placement files and slot policies (`config.py`, `policy.py`), quantization (`quant.py`), simulator and tuner (`sim.py`), the chats (`chat.py`, `web.py` + `web.html`, `tui.py`), the profiler (`profile.py`) and `cli.py` |
+| `examples/` | the RX 6600 XT placement file and `personas/` |
+| `scripts/` | build, environment, kernel tuning, benchmarks and quality checks |
+| `patches/` | the vLLM patches above |
+| `tests/` | `python -m unittest discover -s tests`; GPU tests skip without a GPU |
+| `models/`, `build/` | model weights and packs, build outputs and logs (not in git) |

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

@@ -0,0 +1,76 @@
+# prain placement: prithivMLmods/gemma-4-26B-A4B-it-Uncensored-MAX on an 8 GB RX 6600 XT, set up
+# for a coding agent: int4 experts, 16k context, tool calls.
+# Model: 30 MoE layers x 128 experts, top-8 routing, 3.01 MiB per int4 expert (11.34 MiB bf16).
+# VRAM: 4.46 GiB dense weights + 30 x 15 x 3.01 MiB = 1.32 GiB of expert slots + 1.0 GiB KV cache.
+# Serve with --max-model-len 16384 --attention-backend TRITON_ATTN (ROCM_ATTN needs more LDS than
+# RDNA2 has) --enable-auto-tool-choice --tool-call-parser gemma4.
+# RAM tier: pinning host memory took minutes per GiB on the Thunderbolt eGPU host (cause unproven)
+# and pinned copies were only 4% faster than pageable ones, so this file keeps no RAM tier; the
+# kernel page cache holds the pack instead (the int4 pack, 11.3 GiB, fits in free RAM).
+#
+# 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> --quant int4 --expert-vram-gib 1.4 ...
+# Relative paths start at this file's directory; "~" works too.
+
+pack_dir = "../models/packs/gemma4-26b-a4b"  # bf16 pack (42.5 GiB) written on the first run,
+quant = "int4"                           # then quantized once into an int4 pack (11.3 GiB):
+                                         # bf16 | int8 (5.69 MiB/expert) | int4 (3.01 MiB)
+default_tier = "disk"                    # runs on any RAM size; `prain tune` does better
+cache_slots = 15                         # 30 x 15 x 3.01 MiB = 1.32 GiB of VRAM. The same VRAM
+                                         # holds 8 int8 or 4 bf16 experts per layer: change
+                                         # quant and cache_slots together
+stats_path = "../models/packs/gemma4-26b-a4b/stats.json"  # prain plan and the web monitor read it
+stats_every = 8                          # MoE forward passes between stats/trace writes (~3 ms each)
+trace_path = "../models/packs/gemma4-26b-a4b/trace.jsonl.gz"  # routing trace for sim/tune
+
+[policy]              # which expert a cache slot gives up
+eviction = "gdsf"     # lru | lfu (decayed frequency) | gdsf (frequency x refetch cost)
+heat = "count"        # 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). Measured on the Thunderbolt
+                      #   eGPU host, per int4 expert: a copy ~1.3 ms + disk read, the CPU
+                      #   0.8-1.3 ms per token. `prain tune` picked copy for int4 (a copied
+                      #   expert stays for later tokens) and auto <= 2 tokens for bf16 (copy 4.8 ms)
+cpu_threads = 16      # CPU threads of the vLLM worker (sets OMP_NUM_THREADS)
+admit = 0             # cpu/auto: copy this many CPU-run RAM-tier experts into VRAM afterwards
+min_weight = 1.0      # share of each token's router weight that must be computed. Below 1 a token
+                      #   skips experts not in VRAM once it has that share: faster, dumber.
+                      #   Measured here: 0.6 decodes 14% faster (answer NLL 0.11 -> 0.19), 0.4
+                      #   21% (0.40, visible slips)
+
+[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 = 1.0    # KV cache VRAM shared by them (--kv-cache-memory-bytes). vLLM logs
+                      #   "Maximum concurrency for 16384 tokens per request: 1.11x" with 1.0 GiB.
+                      #   Gemma 4 needs ~20 KiB per token of context plus ~0.6 GiB for its
+                      #   sliding-window layers; every GiB here is 340 int4 expert 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

examples/personas/explainer.md

@@ -0,0 +1,20 @@
+You explain computers, Linux and programming to a curious beginner.
+
+## Task
+Answer questions so that someone without the background understands the answer and can ask a
+better follow-up question.
+
+## How to answer
+- Start with the answer in plain words, in at most three short paragraphs. Offer more detail
+  instead of giving all of it at once.
+- Use one concrete example or comparison per answer, and say where it stops fitting.
+- Explain every abbreviation the first time it appears.
+- If a question can be read two ways, answer the likelier one and mention the other in one
+  sentence.
+
+## Environment
+- A chat window in a browser, often on a phone. The reader has no terminal open.
+
+## How to use it
+- Markdown is rendered: use a short list or a code block only when it helps.
+- End with one suggestion for what to ask next, phrased as a question the reader could send.

examples/personas/gentoo-admin.md

@@ -0,0 +1,28 @@
+You help one person run their Gentoo Linux desktop.
+
+## Task
+Answer questions about this machine and fix problems together with its owner: installing and
+updating packages, USE flags, services, the kernel, drivers, disks and the Hyprland desktop.
+
+## How to answer
+- Lead with the command or the fix, then explain it in a few sentences.
+- Put commands in code blocks, one per line and ready to paste, without `$` or `#` prompts.
+- Before a command that deletes data, touches the bootloader or rebuilds many packages, say
+  what it changes and how to undo it.
+- When the answer depends on this machine (USE flags, profile, kernel, installed versions), ask
+  for the output of a specific command instead of guessing.
+- Say when you are not sure. Never invent package names, USE flags or options.
+
+## Environment
+- Gentoo Linux with Portage and OpenRC. There is no systemd and no journalctl.
+- cron is dcron; the system log is metalog, which writes to /var/log/everything/.
+- Hyprland on Wayland. Graphics: an Intel iGPU and an AMD RX 6600 XT (ROCm) in a Thunderbolt
+  enclosure.
+- You cannot run anything yourself. The owner types your commands into a root shell and pastes
+  the output back into the chat.
+
+## How to use it
+- Preview before changing: `emerge --pretend --verbose PKG`, then `emerge --ask PKG`.
+- Look things up with `eix PKG`, `equery uses PKG`, `emerge --info` and `rc-status`.
+- Services: `rc-service NAME status|start|stop|restart`, `rc-update add NAME default`.
+- Logs: `dmesg`, `tail -n 50 /var/log/everything/current` and the service's own files in /var/log.

examples/personas/pair-programmer.md

@@ -0,0 +1,21 @@
+You are a pair programmer for Python and shell code.
+
+## Task
+Help write, review and debug code that the user pastes into the chat.
+
+## How to answer
+- Read the code first and point at the lines you mean.
+- Prefer the smallest change that fixes the problem. Show it as a unified diff or a complete
+  function, never as fragments with "..." in them.
+- Use the standard library before adding a dependency, and say why when you add one.
+- When a bug report is vague, ask for the exact error message and the command that failed.
+- Write tests as plain `unittest` cases that run with `python -m unittest`.
+
+## Environment
+- Python 3 and Bash on Linux; every project lives in a git repository.
+- You see only what the user pastes: no files, no terminal, no internet.
+
+## How to use it
+- Ask for `git diff`, the full traceback or the failing test's output instead of guessing.
+- Ask for the whole file when a fragment is not enough to answer.
+- Keep an answer within one screen unless the user asks for more.

patches/vllm-0.32-gemma4-fit-8gb.patch

@@ -0,0 +1,52 @@
+diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py
+index d1498f1..745d1d3 100644
+--- a/vllm/model_executor/models/gemma4.py
++++ b/vllm/model_executor/models/gemma4.py
+@@ -26,7 +26,7 @@ import torch
+ from torch import nn
+ 
+ from vllm.compilation.decorators import support_torch_compile
+-from vllm.config import CacheConfig, VllmConfig
++from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
+ from vllm.config.utils import getattr_iter
+ from vllm.distributed import (
+     get_pp_group,
+@@ -491,9 +491,12 @@ class Gemma4Attention(nn.Module):
+                     f"{kv_shared_layer_index}.self_attn.attn"
+                 )
+ 
++        # Gemma 4's rope types (default, proportional) do not depend on max_position, so
++        # size the cos/sin caches for max_model_len: 262144 positions take ~0.4 GiB.
++        max_model_len = get_current_vllm_config().model_config.max_model_len
+         self.rotary_emb = get_rope(
+             self.head_dim,
+-            max_position=max_position_embeddings,
++            max_position=min(max_position_embeddings, max_model_len),
+             rope_parameters=rope_parameters,
+             is_neox_style=True,
+         )
+@@ -1444,14 +1447,17 @@ class Gemma4ForCausalLM(
+             prefix=maybe_prefix(prefix, "model"),
+         )
+ 
+-        self.lm_head = ParallelLMHead(
+-            config.vocab_size,
+-            config.hidden_size,
+-            quant_config=quant_config,
+-            prefix=maybe_prefix(prefix, "lm_head"),
+-        )
+         if config.tie_word_embeddings:
+-            self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
++            # Share embed_tokens: a ParallelLMHead tied after construction would hold the
++            # vocab matrix (1.4 GiB) twice while loading, which overflows 8 GB GPUs.
++            self.lm_head = self.model.embed_tokens
++        else:
++            self.lm_head = ParallelLMHead(
++                config.vocab_size,
++                config.hidden_size,
++                quant_config=quant_config,
++                prefix=maybe_prefix(prefix, "lm_head"),
++            )
+ 
+         self.logits_processor = LogitsProcessor(
+             config.vocab_size,

patches/vllm-0.32-rocm-no-cuda-families.patch

@@ -0,0 +1,17 @@
+diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py
+index 807e571..bbc642a 100644
+--- a/vllm/platforms/rocm.py
++++ b/vllm/platforms/rocm.py
+@@ -983,6 +983,12 @@ class RocmPlatform(Platform):
+         major, minor = torch.cuda.get_device_capability(device_id)
+         return DeviceCapability(major=major, minor=minor)
+ 
++    @classmethod
++    def is_device_capability_family(cls, capability: int, device_id: int = 0) -> bool:
++        # Families are a CUDA notion, and gfx majors collide with them (gfx1032 reads as
++        # 10.3, i.e. Blackwell sm_103): no CUDA family check may match a ROCm GPU.
++        return False
++
+     @classmethod
+     @with_amdsmi_context
+     def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:

prain/chat.py

@@ -0,0 +1,124 @@
+"""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 with_system(system, messages):
+    """messages behind the persona's system message; unchanged without one."""
+    return [{"role": "system", "content": system}, *messages] if system else list(messages)
+
+
+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,225 @@
+"""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, QUANTS, Miss, Policy, Prefetch, expert_bytes, format_ids, load_config,
+                            moe_dims, plan, render_toml)
+
+
+def _model_dims(hf_config_path, quant="bf16"):
+    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, quant)
+
+
+def check(config_path, hf_config_path):
+    config = load_config(config_path)
+    num_layers, num_experts, top_k, size = _model_dims(hf_config_path, config.quant)
+    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}"
+    if miss.mode == "auto":
+        cpu += f" <={miss.cpu_max_tokens}tok"
+    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, config.quant)
+    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.config import Miss, Policy, Prefetch
+    from prain.sim import describe, load_cost, read_forwards, tune
+
+    _, num_experts, _, size = _model_dims(args.hf_config, args.quant)
+    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,
+                                              base=(Policy(), Prefetch(), Miss(cpu_threads=args.cpu_threads)))
+    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), "quant": None if args.quant == "bf16" else args.quant, "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)),
+                       # vLLM detects "openai" for Gemma 4, whose template then prints a system
+                       # prompt as a Python list; a flag in `extra` comes later and wins
+                       "--chat-template-content-format", "string",
+                       *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)
+    p.add_argument("--quant", choices=QUANTS, default="bf16", help="expert weight format")
+    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("--quant", choices=QUANTS, default="bf16", help="expert weight format")
+            s.add_argument("--cpu-threads", type=int, default=4,
+                           help="[miss] cpu_threads to write: the --threads scripts/measure.py used")
+            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")
+        ui.add_argument("--persona", metavar="FILE",
+                        help="text file sent as the system message of every chat: the model's "
+                             "task, how to answer, its environment (see examples/personas)")
+        if name == "web":
+            ui.add_argument("--port", type=int, default=8080)
+            ui.add_argument("--host", default="127.0.0.1",
+                            help="address to bind; 0.0.0.0 for every interface (no login!)")
+            ui.add_argument("--config", help="placement file: [serve] max_requests, the [draft] "
+                                             "table and stats_path for the expert monitor "
+                                             "(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, args.quant)
+        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), args.quant), 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, system=_persona(args)[0])
+    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
+        system, persona = _persona(args)
+        web(args.url, args.port, args.warmup, args.warmup_priority, config.serve, config.draft,
+            host=args.host, stats_path=config.stats_path, system=system, persona=persona)
+
+
+def _persona(args):
+    """(system message text, name) of --persona FILE, or (None, None)."""
+    if not args.persona:
+        return None, None
+    with open(args.persona) as f:
+        return f.read().strip(), os.path.splitext(os.path.basename(args.persona))[0]
+
+
+if __name__ == "__main__":
+    main()

prain/config.py

@@ -8,8 +8,10 @@ import tomllib
 from dataclasses import dataclass, field, fields
 
 TIERS = ("vram", "ram", "disk")
+QUANTS = ("bf16", "int8", "int4")  # expert weight formats, see prain.quant
+INT4_GROUP = 64  # int4 weights per scale
 GIB = 1024**3
-_TOP_KEYS = {"pack_dir", "default_tier", "cache_slots", "stats_path", "stats_every",
+_TOP_KEYS = {"pack_dir", "quant", "default_tier", "cache_slots", "stats_path", "stats_every",
              "trace_path", "policy", "prefetch", "miss", "serve", "draft", "layers"}
 _LAYER_KEYS = {"cache_slots", *TIERS}
 
@@ -57,3 +59,260 @@ 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
+    min_weight: float = _opt(1.0, maximum=1)  # router weight every token must compute; below 1
+    #                                           it skips experts not in VRAM: faster, dumber
+
+
+@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
+    quant: str = "bf16"
+    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"]),
+        quant=raw.get("quant", "bf16"),
+        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.quant not in QUANTS:
+        raise ValueError(f"{path}: quant must be one of {QUANTS}")
+    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, quant="bf16"):
+    """One gated expert, gate+up [2I, H] plus down [H, I], stored as `quant`."""
+    weights = 3 * hidden * moe_intermediate
+    if quant == "bf16":
+        return 2 * weights
+    if quant == "int8":  # plus one fp32 scale per output row
+        return weights + 4 * (2 * moe_intermediate + hidden)
+    return weights // 2 + 2 * weights // INT4_GROUP  # plus one bf16 scale per group
+
+
+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,
+         quant="bf16"):
+    """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, "quant": None if quant == "bf16" else quant, "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/experts.py

@@ -0,0 +1,529 @@
+"""TieredRoutedExperts: vLLM's RoutedExperts with only some experts in VRAM.
+
+vLLM builds this class instead of RoutedExperts once prain.plugin.register() ran.
+Every expert is written once to a per-layer pack file on disk while vLLM loads the
+checkpoint. With quant = "int8" or "int4" the first forward quantizes that bf16 pack once into
+a second pack (prain.quant), and the VRAM slots, copies and CPU path use the quantized
+experts. On the first forward the RAM tier is copied from the pack into pinned host memory
+and the pinned VRAM slots are filled. After that each forward:
+  0. drops the routes of draft tokens to experts that are not available ([draft] min_weight;
+     prain.drafts knows which rows are drafts),
+  1. counts routes per expert (one host sync),
+  2. lets the [miss] config send some missing experts to the CPU,
+  3. copies the other missing experts into slots chosen by the [policy] config and runs
+     vLLM's Triton fused-MoE kernel on the slot tensors, in rounds if needed (each round
+     computes only its own routes),
+  4. queues async copies on a side stream: [prefetch] guesses for a later layer, and the
+     hottest CPU-run experts (admit-behind),
+  5. adds the CPU result.
+Slot safety with async copies: a copy into a slot waits for the last kernel that read it
+(`_last_use`), and a kernel or demand copy waits for a pending async copy into its slot
+(`_ready_events`).
+"""
+
+import atexit
+import json
+import math
+import os
+
+import torch
+from vllm.config import get_current_vllm_config
+from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
+from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
+
+from prain import drafts, profile
+from prain.config import layer_index, load_config
+from prain.kernels import DECODE_ROWS, gfx10, moe_int4_decode, unpack_int4
+from prain.policy import SlotCache, write_stats
+from prain.quant import (GROUP, cpu_matmul, layout, moe_quant_config, quantize, to_cpu_int4)
+from prain.sim import write_trace
+from prain.tiered import draft_rows, tiered_forward
+
+logger = init_logger("vllm.prain")  # under "vllm." so vLLM's log handler prints it
+LAYERS = {}  # layer index -> TieredRoutedExperts
+_COPY_STREAM = []  # one side stream shared by all layers: async copies stay in issue order
+_TRACE = []  # trace records not written yet
+_DRAFT_MASK = [None, None]  # (spans, their bool row mask on the GPU) of the current forward
+_DRAFT_TOKENS = [0]  # draft token rows seen, for the stats log
+_STAGING = {}  # stream -> _Staging
+LOW_VRAM = 128 * 2**20  # driver-free bytes below which torch's cache goes back (_release_cache)
+
+
+def _release_cache():
+    """Hand torch's cached free blocks (the freed one-row bf16 expert weights, ~11 MiB per layer,
+    and load-time temporaries) back to the driver. HIP loads kernel code objects lazily into
+    VRAM; with a full card and nothing returned, a kernel's first launch fails with
+    hipErrorNoBinaryForGpu ("no kernel image")."""
+    before = torch.cuda.mem_get_info()[0]
+    torch.cuda.empty_cache()
+    free, total = torch.cuda.mem_get_info()
+    logger.info("prain: %.0f MiB of VRAM free for kernels and activations (%.0f MiB before "
+                "releasing torch's cache, %.2f of %.2f GiB reserved)", free / 2**20,
+                before / 2**20, torch.cuda.memory_reserved() / 2**30, total / 2**30)
+
+
+class _Staging:
+    """A ring of pinned host buffers and GPU buffers for staged int4 copies: the CPU copies an
+    expert's tensors from the page cache into one pinned buffer, one DMA moves it, and
+    unpack_int4 writes the slot. The CPU fills the next buffer while the last DMA runs."""
+    SLOTS = 4
+
+    def __init__(self, nbytes, device):
+        self.host = [torch.empty(nbytes, dtype=torch.uint8).pin_memory() for _ in range(self.SLOTS)]
+        self.dev = [torch.empty(nbytes, dtype=torch.uint8, device=device)
+                    for _ in range(self.SLOTS)]
+        self.done = [None] * self.SLOTS  # event: the DMA that last read this pinned buffer
+        self.i = 0
+
+    def take(self):
+        self.i = (self.i + 1) % self.SLOTS
+        if self.done[self.i] is not None:
+            self.done[self.i].synchronize()
+        return self.host[self.i], self.dev[self.i], self.i
+
+
+def stage_int4(sources, slots, slot, ring, stream):
+    """Copy one int4 expert: sources [[w13 q, w13 s], [w2 q, w2 s]] on the CPU (q in the CPU
+    layout) -> slots[...][slot] on the GPU (q in the GPU layout), through `ring`."""
+    host, dev, i = ring.take()
+    parts, offset = [], 0
+    with profile.span("memcpy"):
+        for mats in sources:
+            for src in mats:
+                n = src.nbytes
+                host[offset:offset + n].copy_(src.reshape(-1).view(torch.uint8))
+                parts.append((offset, n))
+                offset += n
+    with profile.span("h2d"):
+        dev[:offset].copy_(host[:offset], non_blocking=True)
+        ring.done[i] = torch.cuda.Event()
+        ring.done[i].record(stream)
+    with profile.span("unpack"):
+        parts = iter(parts)
+        for dsts in slots:
+            for t, dst in enumerate(dsts):
+                o, n = next(parts)
+                if t == 0:  # the q tensor: the pack keeps the CPU kernel's layout
+                    unpack_int4(dev[o:o + n], dst[slot])
+                else:
+                    dst[slot].view(-1).view(torch.uint8).copy_(dev[o:o + n])
+
+
+def _install_draft_hook():
+    """Wrap the V2 model runner's prepare_inputs so prain.drafts sees each step's batch
+    layout: which request ids sit at which batch positions, and their token rows."""
+    from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
+    original = GPUModelRunner.prepare_inputs
+    if getattr(original, "prain", False):
+        return
+
+    def prepare_inputs(self, scheduler_output, *args, **kwargs):
+        batch = original(self, scheduler_output, *args, **kwargs)
+        if hasattr(scheduler_output, "scheduled_new_reqs"):
+            drafts.note_step(scheduler_output, batch.req_ids, batch.query_start_loc_np)
+        return batch
+
+    prepare_inputs.prain = True
+    GPUModelRunner.prepare_inputs = prepare_inputs
+
+
+def _draft_mask(x):
+    """Bool row mask of draft tokens for this forward, built once and shared by all layers."""
+    spans = drafts.spans()
+    if spans is None:
+        return None
+    if max(end for _, end in spans) > x.shape[0]:  # not this forward's layout: run full
+        return None
+    if _DRAFT_MASK[0] is not spans:
+        _DRAFT_MASK[:] = [spans, draft_rows(spans, x.shape[0], x.device)]
+        _DRAFT_TOKENS[0] += sum(end - first for first, end in spans)
+    return _DRAFT_MASK[1]
+
+
+def _copy_stream():
+    if not _COPY_STREAM:
+        _COPY_STREAM.append(torch.cuda.Stream())
+    return _COPY_STREAM[0]
+
+
+def _file_tensor(path, shape, dtype):
+    """A CPU tensor backed by `path` (created if missing); writes land in the file."""
+    return torch.from_file(
+        path, shared=True, size=math.prod(shape), dtype=dtype, device="cpu"
+    ).view(shape)
+
+
+def _mem_available():
+    with open("/proc/meminfo") as f:
+        for line in f:
+            if line.startswith("MemAvailable:"):
+                return int(line.split()[1]) * 1024
+    return None
+
+
+def _flush_trace():
+    if _TRACE and LAYERS:
+        write_trace(next(iter(LAYERS.values())).prain.trace_path, _TRACE, mode="at")
+        _TRACE.clear()
+
+
+atexit.register(_flush_trace)
+
+
+class TieredRoutedExperts(RoutedExperts):
+    def __init__(self, layer_name, params_dtype, moe_config, quant_config, *args, **kwargs):
+        vllm_config = get_current_vllm_config()
+        parallel = moe_config.moe_parallel_config
+        layer = layer_index(layer_name)
+        if layer is None:
+            raise ValueError(f"prain: cannot find the layer number in {layer_name!r}")
+        if quant_config is not None:
+            raise NotImplementedError("prain: quantized MoE weights are not supported")
+        if not vllm_config.model_config.enforce_eager:
+            raise ValueError("prain: pass --enforce-eager (expert copies wait on the host)")
+        if vllm_config.kernel_config.moe_backend != "triton":
+            raise ValueError("prain: pass --moe-backend triton")
+        if parallel.tp_size > 1 or parallel.dp_size > 1 or parallel.use_ep:
+            raise NotImplementedError("prain: one GPU only (no TP, DP or EP)")
+        if vllm_config.use_v2_model_runner:
+            _install_draft_hook()
+            if profile.DIR:
+                profile.install()
+        elif not LAYERS:
+            logger.warning("prain: drafts need vLLM's V2 model runner; with "
+                           "VLLM_USE_V2_MODEL_RUNNER=0 every request runs with all its experts")
+        config = load_config(os.environ["PRAIN_CONFIG"])
+        num_experts = moe_config.num_experts
+        placement = config.placement(layer, num_experts)
+        # Gemma 4 passes its router module's routing_function: keep the module for pregate.
+        router = getattr(kwargs.get("custom_routing_function"), "__self__", None)
+        if config.prefetch.mode == "pregate" and not isinstance(router, torch.nn.Module):
+            raise NotImplementedError("prain: prefetch.mode = pregate needs a router module "
+                                      "passed as custom_routing_function (Gemma 4 does)")
+        # vLLM sizes w13/w2 and the kernel by num_experts: make that the VRAM slot count.
+        # Quantized slots are prain's own tensors; vLLM's bf16 ones shrink to one row.
+        quantized = config.quant != "bf16"
+        moe_config.num_experts = moe_config.num_local_experts = 1 if quantized else placement.num_slots
+        super().__init__(layer_name, params_dtype, moe_config, quant_config, *args, **kwargs)
+        self.local_num_experts = placement.num_slots
+        self.quant = config.quant
+        self._shapes = (tuple(self.w13_weight.shape[1:]), tuple(self.w2_weight.shape[1:]))
+        self._qslots = [[torch.zeros((placement.num_slots, *shape), dtype=dtype, device="cuda")
+                         for shape, dtype in layout(config.quant, *dims)]
+                        for dims in self._shapes] if quantized else None
+        self.prain = config
+        self.prain_layer = layer
+        self.placement = placement
+        disk = set(placement.disk)
+        cost = [config.policy.disk_cost if e in disk else 1.0 for e in range(num_experts)]
+        self.cache = SlotCache(placement.vram, num_experts, placement.cache_slots, config.policy, cost)
+        self._router = [router] if isinstance(router, torch.nn.Module) else []  # not a submodule
+        self._async = config.prefetch.mode != "none" or (config.miss.mode != "copy" and config.miss.admit)
+        self._ready_events = {}  # slot -> event: an async copy into the slot lands
+        self._last_use = {}  # slot -> event: the last kernel reading the slot finished
+        self._predictions = {}  # distance -> [(expert, probability)] from this layer's router
+        self._ram_row = None
+        self._open_pack(config.pack_dir, layer, num_experts, vllm_config.model_config.model)
+        self._ready = False
+        self._steps = 0
+        LAYERS[layer] = self
+        logger.info(
+            "prain layer %d: %d vram, %d ram, %d disk, %d cache slots%s",
+            layer, len(placement.vram), len(placement.ram), len(placement.disk),
+            placement.cache_slots, "" if self.pack_complete else
+            ", writing pack" if self._writing else f", quantizing the pack to {self.quant}",
+        )
+
+    def _open_pack(self, pack_dir, layer, num_experts, model):
+        """The loader writes the bf16 pack; a quantized pack is made from it once. Once the
+        pack of config.quant is complete, the bf16 one is not read (it may be deleted)."""
+        os.makedirs(pack_dir, exist_ok=True)
+        self._stem = os.path.join(pack_dir, f"layer{layer:03d}")
+        self._num_experts = num_experts
+        self._marker = {"model": model, "w13": [num_experts, *self._shapes[0]],
+                        "w2": [num_experts, *self._shapes[1]], "dtype": str(self.w13_weight.dtype)}
+        self.pack_complete = self._complete(self.quant)
+        self._bf16 = None if self.pack_complete else self._pack_of("bf16")
+        self._writing = self._bf16 is not None and not self._complete("bf16")
+
+    def _paths(self, quant):
+        """[[w13 files], [w2 files]] of a pack format: weights, then scales."""
+        infix = "" if quant == "bf16" else f".{quant}"
+        return [[f"{self._stem}{infix}.{name}{'s' * i}" for i in range(len(layout(quant, *dims)))]
+                for name, dims in zip(("w13", "w2"), self._shapes)]
+
+    def _pack_of(self, quant):
+        return [[_file_tensor(path, (self._num_experts, *shape), dtype)
+                 for path, (shape, dtype) in zip(paths, layout(quant, *dims))]
+                for paths, dims in zip(self._paths(quant), self._shapes)]
+
+    def _marker_of(self, quant):
+        if quant == "bf16":
+            return f"{self._stem}.done", self._marker
+        return f"{self._stem}.{quant}.done", {**self._marker, "quant": quant, "group": GROUP}
+
+    def _complete(self, quant):
+        path, marker = self._marker_of(quant)
+        if not os.path.exists(path):
+            return False
+        with open(path) as f:
+            found = json.load(f)
+        if found != marker:
+            raise ValueError(f"prain: {path} is for {found}, not {marker}; delete the pack "
+                             "or use another pack_dir")
+        return True
+
+    def _seal(self, quant):
+        for path in sum(self._paths(quant), []):  # mmap writes -> disk before we trust the pack
+            fd = os.open(path, os.O_RDONLY)
+            os.fsync(fd)
+            os.close(fd)
+        path, marker = self._marker_of(quant)
+        with open(path, "w") as f:
+            json.dump(marker, f)
+
+    def _quantize_pack(self):
+        """bf16 pack -> config.quant pack, one expert at a time on the GPU (int4's clip search
+        takes ~20 min for the whole model on the CPU, under a minute on the GPU)."""
+        out = self._pack_of(self.quant)
+        for src, dst in zip(self._bf16, out):
+            for e in range(self._num_experts):
+                parts = quantize(src[0][e].cuda(), self.quant)
+                if self.quant == "int4":
+                    parts[0] = to_cpu_int4(parts[0])
+                for tensor, part in zip(dst, parts):
+                    tensor[e] = part.cpu()
+
+    def get_expert_mapping(self, ckpt_gate_proj_name=None, ckpt_down_proj_name=None,
+                           ckpt_up_proj_name=None, include_fused=False):
+        # The base class derives "redundant experts" as num_experts - num_logical_experts,
+        # which is negative here because num_experts counts VRAM slots.
+        return self.build_expert_params_mapping(
+            ckpt_gate_proj_name or self.ckpt_gate_proj_name,
+            ckpt_down_proj_name or self.ckpt_down_proj_name,
+            ckpt_up_proj_name or self.ckpt_up_proj_name,
+            num_experts=self.moe_config.num_logical_experts,
+            routed_experts_prefix="",
+            lora_base_layer_prefix=self.lora_base_layer_prefix,
+            include_fused=include_fused,
+            is_gated=self.moe_config.is_act_and_mul,
+        )
+
+    def _map_global_expert_id_to_local_expert_id(self, expert_id):
+        return expert_id  # the loader only ever writes into the [num_experts, ...] pack
+
+    def weight_loader(self, param, loaded_weight, weight_name, shard_id, expert_id,
+                      return_success=False):
+        # Every expert goes to the pack; the VRAM and RAM tiers fill from it later.
+        if self._writing:
+            pack = self._bf16[0][0] if param is self.w13_weight else self._bf16[1][0]
+            super().weight_loader(pack, loaded_weight, weight_name, shard_id, expert_id,
+                                  return_success=True)
+        return True if return_success else None
+
+    weight_loader.supports_moe_loading = True
+
+    def _materialize(self):
+        if not self.pack_complete:
+            if self._writing:
+                self._seal("bf16")
+            if self.quant != "bf16":
+                self._quantize_pack()
+                self._seal(self.quant)
+            self.pack_complete, self._writing, self._bf16 = True, False, None
+        self._pack = self._pack_of(self.quant)  # [[w13 tensors], [w2 tensors]] per expert row
+        self._expert_bytes = sum(t[0].nbytes for mats in self._pack for t in mats)
+        if self._qslots is None:
+            self._slots, self._qconfig = [[self.w13_weight.data], [self.w2_weight.data]], None
+        else:
+            self._slots = self._qslots
+            self._qconfig = moe_quant_config(self.quant, self._slots[0][1], self._slots[1][1])
+            # vLLM's one-row bf16 w13/w2 are never read: free them (11 MiB per layer)
+            self.w13_weight.data = self.w13_weight.data.new_empty(0)
+            self.w2_weight.data = self.w2_weight.data.new_empty(0)
+        if self.prain_layer == min(LAYERS):  # check the whole RAM tier once, before pinning
+            row = sum(t[0].nbytes for mats in self._pack for t in mats)
+            need = row * sum(len(experts.placement.ram) for experts in LAYERS.values())
+            free = _mem_available()
+            if free is not None and need > 0.9 * free:
+                raise RuntimeError(f"prain: the RAM tier needs {need / 2**30:.1f} GiB pinned, "
+                                   f"but only {free / 2**30:.1f} GiB are available; move experts "
+                                   "from ram to disk")
+        ram = list(self.placement.ram)
+        self._ram_row = {e: row for row, e in enumerate(ram)}
+        # ponytail: pinning is slow on the eGPU host (minutes for a 10 GiB tier, even per
+        # expert); the page cache serves the pack about as fast, so keep the RAM tier small.
+        self._ram = [[[t[e].pin_memory() for t in mats] for mats in self._pack] for e in ram]
+        for slot, expert in enumerate(self.placement.vram):
+            self._load(expert, slot)
+        if self._router and (self.prain.prefetch.mode == "pregate" or self.prain.trace_path):
+            self._router[0].register_forward_pre_hook(self._predict)
+
+    def _source(self, expert):
+        """[[w13 tensors], [w2 tensors]] of an expert on the CPU: pinned rows for the RAM tier,
+        else the mmap'd pack."""
+        row = self._ram_row.get(expert)
+        if row is None:
+            return [[t[expert] for t in mats] for mats in self._pack]
+        return self._ram[row]
+
+    def _copy(self, expert, slot):
+        if self.quant == "int4":  # the pack keeps the CPU kernel's layout: stage, then unpack
+            return self._copy_staged(expert, slot)
+        for dsts, srcs in zip(self._slots, self._source(expert)):
+            for dst, src in zip(dsts, srcs):
+                with profile.span("h2d" if src.device.type == "cpu" else "d2d"):
+                    dst[slot].copy_(src, non_blocking=True)
+
+    def _copy_staged(self, expert, slot):
+        """int4 on the current stream: one pinned buffer, one DMA, unpack into the slot."""
+        stream = torch.cuda.current_stream()
+        ring = _STAGING.get(stream)
+        if ring is None:
+            ring = _STAGING[stream] = _Staging(self._expert_bytes, stream.device)
+        stage_int4(self._source(expert), self._slots, slot, ring, stream)
+
+    def _load(self, expert, slot):
+        """Demand copy on the compute stream."""
+        pending = self._ready_events.pop(slot, None)
+        if pending is not None:  # an older async copy into this slot must land first
+            torch.cuda.current_stream().wait_event(pending)
+        self._copy(expert, slot)
+
+    def _load_async(self, expert, slot):
+        """Copy on the side stream; truly asynchronous for int4 (staged) and pinned RAM-tier
+        experts."""
+        stream = _copy_stream()
+        last = self._last_use.get(slot)
+        if last is not None:  # a queued kernel may still read the slot's old expert
+            stream.wait_event(last)
+        with torch.cuda.stream(stream):
+            self._copy(expert, slot)
+        event = torch.cuda.Event()
+        event.record(stream)
+        self._ready_events[slot] = event
+
+    def _run_slots(self, x, weights, ids, table, used_slots):
+        """The fused-MoE kernel on the slot tensors (vLLM's, or prain's for int4 on RDNA2);
+        expert_map=table skips the routes of experts outside this round."""
+        current = torch.cuda.current_stream()
+        for slot in used_slots:
+            pending = self._ready_events.pop(slot, None)
+            if pending is not None:
+                current.wait_event(pending)
+        if self.quant == "int4" and gfx10(x.device.index or 0):
+            out = self._int4_gfx10(x, weights, ids, table)
+        else:
+            out = fused_experts(x, self._slots[0][0], self._slots[1][0], weights, ids,
+                                activation=self.activation, global_num_experts=table.numel(),
+                                expert_map=table, quant_config=self._qconfig)
+        if self._async:
+            event = torch.cuda.Event()
+            event.record(current)
+            for slot in used_slots:
+                self._last_use[slot] = event
+        return out
+
+    def _int4_gfx10(self, x, weights, ids, table):
+        """RDNA2 has no bf16 dot instructions, which vLLM's int4 kernel uses: a GEMV-style
+        kernel for a few tokens (10x faster at one), else vLLM's kernel in fp16 (2x faster),
+        redone in bf16 if fp16 overflowed."""
+        if (x.shape[0] <= DECODE_ROWS
+                and getattr(self.activation, "value", self.activation) == "gelu_tanh"):
+            return moe_int4_decode(x, weights, ids, table, *self._slots[0], *self._slots[1])
+        run = lambda x, config: fused_experts(  # noqa: E731
+            x, self._slots[0][0], self._slots[1][0], weights, ids, activation=self.activation,
+            global_num_experts=table.numel(), expert_map=table, quant_config=config)
+        out = run(x.half(), moe_quant_config("int4", self._slots[0][1].half(),
+                                             self._slots[1][1].half()))
+        if not torch.isfinite(out).all():  # host sync, once per round of a prefill
+            out = run(x, self._qconfig)
+        return out.to(x.dtype)
+
+    def _predict(self, module, args):
+        """Pre-hook of this layer's router: guess the experts of later layers from the same
+        residual stream (their routers normalize their input, so the guess is cheap and
+        usually close)."""
+        distances = (1, 2) if self.prain.trace_path else (self.prain.prefetch.distance,)
+        self._predictions = {}
+        for d in distances:
+            target = LAYERS.get(self.prain_layer + d)
+            if target is None or not target._router:
+                continue
+            logits = target._router[0].forward(args[0])  # .forward: skip the target's own hook
+            probs = torch.softmax(logits.float(), dim=-1).mean(0)
+            p, e = probs.topk(min(2 * self.top_k, probs.numel()))
+            self._predictions[d] = list(zip(e.tolist(), p.tolist()))  # host sync
+
+    def _cpu_weights(self, expert):
+        mats = self._source(expert)
+        if self.quant == "bf16":
+            return mats[0][0], mats[1][0]
+        return tuple(cpu_matmul(self.quant, *parts) for parts in mats)
+
+    def _after_enqueue(self, cpu, stats):
+        miss, prefetch = self.prain.miss, self.prain.prefetch
+        if cpu and miss.admit:  # admit-behind: the hottest CPU-run RAM experts go to VRAM
+            hot = sorted((e for e in cpu if e in self._ram_row), key=lambda e: (-stats[e][1], e))
+            for e, slot in self.cache.prefetch([(e, 1.0) for e in hot], miss.admit):
+                self._load_async(e, slot)
+        target = LAYERS.get(self.prain_layer + prefetch.distance)
+        if prefetch.mode == "pregate" and target is not None and target._ready:
+            staged = target.quant == "int4"  # staged copies are async from every tier
+            guesses = [(e, p) for e, p in self._predictions.get(prefetch.distance, ())
+                       if p >= prefetch.min_prob and (staged or e in target._ram_row)]
+            for e, slot in target.cache.prefetch(guesses, prefetch.budget):
+                target._load_async(e, slot)
+
+    def _record(self, stats, tokens):
+        ids = sorted(stats)
+        record = {"l": self.prain_layer, "t": tokens, "e": ids,
+                  "c": [stats[e][0] for e in ids], "w": [round(stats[e][1], 5) for e in ids],
+                  "r": [round(stats[e][2], 5) for e in ids]}
+        if self._predictions:
+            record["p"] = {str(d): [[e, round(p, 5)] for e, p in guesses]
+                           for d, guesses in self._predictions.items()}
+        _TRACE.append(record)
+
+    def forward_modular(self, x, topk_weights, topk_ids, shared_experts=None,
+                        shared_experts_input=None):
+        if shared_experts is not None:
+            raise NotImplementedError("prain: fused shared experts are not supported")
+        if not self._ready:
+            self._materialize()
+            self._ready = True
+            if self.prain_layer == max(LAYERS):  # every layer has its slots now
+                _release_cache()
+        policy = self.prain.policy
+        out = tiered_forward(
+            x, topk_weights, topk_ids, self.cache, self._load, self._run_slots,
+            policy.heat, policy.recency_half_life, self.prain.miss, self._cpu_weights,
+            on_stats=self._record if self.prain.trace_path else None,
+            after_enqueue=self._after_enqueue,
+            draft=_draft_mask(x), min_weight=self.prain.draft.min_weight,
+        )
+        if self.prain_layer == max(LAYERS):
+            drafts.clear()  # the next forward may be a warm-up run without a batch layout
+        self._maybe_write_stats()
+        return out
+
+    def _maybe_write_stats(self):
+        self._steps += 1
+        if self._steps % self.prain.stats_every or self.prain_layer != min(LAYERS):
+            return
+        if torch.cuda.mem_get_info()[0] < LOW_VRAM:  # other processes or activations took it
+            _release_cache()
+        if self.prain.trace_path:
+            _flush_trace()
+        if not self.prain.stats_path:
+            return
+        write_stats(self.prain.stats_path, {layer: e.cache for layer, e in LAYERS.items()})
+        c = {k: sum(getattr(e.cache, k) for e in LAYERS.values())
+             for k in ("hits", "misses", "cpu", "prefetches", "prefetch_hits")}
+        logger.info("prain: VRAM hit rate %.1f%% (%d hits, %d copies, %d on CPU, "
+                    "%d prefetched, %d of them used, %d draft tokens), stats in %s",
+                    100 * c["hits"] / max(1, c["hits"] + c["misses"] + c["cpu"]), c["hits"],
+                    c["misses"], c["cpu"], c["prefetches"], c["prefetch_hits"],
+                    _DRAFT_TOKENS[0], self.prain.stats_path)

prain/gemv.py

@@ -0,0 +1,90 @@
+"""Decode GEMMs (at most MAX_ROWS rows) for gfx10 GPUs as a Triton GEMV.
+
+vLLM sends unquantized linear layers on ROCm through rocBLAS, except on gfx9/gfx11/gfx12 where it
+has skinny-GEMM kernels. On an RX 6600 XT (gfx1032) rocBLAS runs a bf16 GEMV at 13-16 GB/s
+(RDNA2 has no bf16 dot instructions), so the dense part of a decode step reads its weights at
+6% of the card's bandwidth. This kernel reads them at ~210 GB/s (measured, lm_head 262144 x 2816:
+93 ms -> 6.7 ms). install() wraps vLLM's dispatch_unquantized_gemm; PRAIN_GEMV=0 turns it off.
+"""
+
+import functools
+import os
+
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+from prain.kernels import gfx10
+
+MAX_ROWS = 4
+# ponytail: one config, the one autotuning picked for every Gemma 4 dense shape on gfx1032.
+# Triton's autotuner benchmarks with a 256 MiB buffer that torch keeps cached, and an 8 GB card
+# has no 256 MiB to spare (lazy kernel loads then fail). Re-measure if the GPU or model changes.
+BLOCK_N, BLOCK_K, NUM_WARPS = 4, 512, 8
+
+
+@triton.jit
+def _gemv(w_ptr, x_ptr, y_ptr, M, N, K, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
+          BLOCK_K: tl.constexpr):
+    rows = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N)
+    ms = tl.arange(0, BLOCK_M)
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k0 in range(0, K, BLOCK_K):
+        cols = k0 + tl.arange(0, BLOCK_K)
+        w = tl.load(w_ptr + rows[:, None] * K + cols[None, :],
+                    mask=(rows[:, None] < N) & (cols[None, :] < K), other=0.0)
+        x = tl.load(x_ptr + ms[:, None] * K + cols[None, :],
+                    mask=(ms[:, None] < M) & (cols[None, :] < K), other=0.0)
+        acc += tl.sum(x.to(tl.float32)[:, None, :] * w.to(tl.float32)[None, :, :], axis=2)
+    tl.store(y_ptr + ms[:, None] * N + rows[None, :], acc.to(y_ptr.dtype.element_ty),
+             mask=(ms[:, None] < M) & (rows[None, :] < N))
+
+
+@functools.cache
+def fast_device(index):
+    """gfx10 (RDNA2) GPUs, where vLLM has no skinny-GEMM kernels."""
+    return gfx10(index) and os.environ.get("PRAIN_GEMV", "1") != "0"
+
+
+def _small(x, weight, bias):
+    return (bias is None and x.is_cuda and x.dtype in (torch.bfloat16, torch.float16)
+            and weight.dtype == x.dtype and weight.is_contiguous()
+            and x.numel() <= MAX_ROWS * x.shape[-1] and fast_device(x.device.index or 0))
+
+
+def linear(x, weight, bias=None):
+    """F.linear(x, weight, bias); the Triton GEMV for up to MAX_ROWS rows on gfx10."""
+    if not _small(x, weight, bias):
+        return F.linear(x, weight, bias)
+    n, k = weight.shape
+    x2 = x.reshape(-1, k).contiguous()
+    m = x2.shape[0]
+    y = torch.empty(m, n, dtype=x.dtype, device=x.device)
+    _gemv[(triton.cdiv(n, BLOCK_N),)](weight, x2, y, m, n, k, BLOCK_M=triton.next_power_of_2(m),
+                                      BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, num_warps=NUM_WARPS)
+    return y.reshape(*x.shape[:-1], n)
+
+
+def wrap(dispatch):
+    """vLLM's dispatch_unquantized_gemm(), with decode shapes sent to linear()."""
+    def wrapped(*args, **kwargs):
+        fallback = dispatch(*args, **kwargs)
+
+        def gemm(layer, x, weight, bias=None):
+            if _small(x, weight, bias):
+                return linear(x, weight)
+            return fallback(layer, x, weight, bias)
+        return gemm
+    return wrapped
+
+
+def install():
+    """Before the model is built: every module that looks up dispatch_unquantized_gemm."""
+    from vllm.model_executor.layers import linear as linear_layers, utils
+    from vllm.model_executor.layers import vocab_parallel_embedding
+    from vllm.model_executor.layers.fused_moe.runner import moe_runner
+
+    wrapped = wrap(utils.dispatch_unquantized_gemm)
+    for module in (linear_layers, vocab_parallel_embedding, moe_runner):
+        module.dispatch_unquantized_gemm = wrapped

prain/kernels.py

@@ -0,0 +1,133 @@
+"""Triton kernels for int4 experts on gfx10 (RDNA2), where vLLM's int4 fused-MoE kernel pads
+every token to a 16-row tile and dot-products in bf16, which RDNA2 lacks: 2.85 ms for one token.
+
+GPU int4 layout (prain.quant): q uint8 [rows, K/2], even k in the low nibble, bf16 scales
+[rows, K/64], w = (nibble - 8) * scale. Experts use gelu_tanh(gate) * up, gate = the first half
+of the w13 rows (Gemma 4).
+"""
+
+import functools
+
+import torch
+import triton
+import triton.language as tl
+
+GROUP = 64
+DECODE_ROWS = 4  # up to this many tokens per round use moe_int4_decode
+
+
+@functools.cache
+def gfx10(index):
+    """RDNA2 GPUs: no bf16 dot instructions, where vLLM's bf16 GEMV/int4 kernels are slow."""
+    return getattr(torch.cuda.get_device_properties(index), "gcnArchName", "").startswith("gfx10")
+
+
+@triton.jit
+def _gelu_tanh(x):  # 0.5 x (1 + tanh(z)) = x sigmoid(2z), z = sqrt(2/pi) (x + 0.044715 x^3)
+    return x * tl.sigmoid(1.5957691216057308 * (x + 0.044715 * x * x * x))
+
+
+@triton.jit
+def _gate_up(x_ptr, ids_ptr, table_ptr, q_ptr, s_ptr, h_ptr, TOPK: tl.constexpr,
+             H: tl.constexpr, N2: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    """h[pair] = W13[slot] @ x[token] for one (token, route) pair; routes off the table skip."""
+    pair = tl.program_id(0)
+    slot = tl.load(table_ptr + tl.load(ids_ptr + pair))
+    if slot >= 0:
+        m = pair // TOPK
+        rows = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
+        rmask = rows < N2
+        half = tl.arange(0, BLOCK_K // 2)
+        q_base = q_ptr + slot.to(tl.int64) * (N2 * (H // 2))
+        s_base = s_ptr + slot.to(tl.int64) * (N2 * (H // 64))
+        acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+        for k0 in range(0, H, BLOCK_K):
+            j = k0 // 2 + half
+            b = tl.load(q_base + rows[:, None] * (H // 2) + j[None, :], mask=rmask[:, None],
+                        other=0).to(tl.int32)
+            s = tl.load(s_base + rows[:, None] * (H // 64) + (j // 32)[None, :],
+                        mask=rmask[:, None], other=0.0).to(tl.float32)
+            xe = tl.load(x_ptr + m * H + 2 * j).to(tl.float32)
+            xo = tl.load(x_ptr + m * H + 2 * j + 1).to(tl.float32)
+            lo = (b & 15).to(tl.float32) - 8.0
+            hi = (b >> 4).to(tl.float32) - 8.0
+            acc += tl.sum((lo * xe[None, :] + hi * xo[None, :]) * s, axis=1)
+        tl.store(h_ptr + pair * N2 + rows, acc, mask=rmask)
+
+
+@triton.jit
+def _down(h_ptr, ids_ptr, table_ptr, w_ptr, q_ptr, s_ptr, y_ptr, H: tl.constexpr,
+          I: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
+    """y[pair] = weight * W2[slot] @ (gelu_tanh(gate) * up) for one (token, route) pair."""
+    pair = tl.program_id(0)
+    slot = tl.load(table_ptr + tl.load(ids_ptr + pair))
+    if slot >= 0:
+        rows = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
+        rmask = rows < H
+        half = tl.arange(0, BLOCK_K // 2)
+        q_base = q_ptr + slot.to(tl.int64) * (H * (I // 2))
+        s_base = s_ptr + slot.to(tl.int64) * (H * (I // 64))
+        h_base = h_ptr + pair * (2 * I)
+        acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
+        for k0 in range(0, I, BLOCK_K):
+            j = k0 // 2 + half
+            kmask = 2 * j < I
+            ae = _gelu_tanh(tl.load(h_base + 2 * j, mask=kmask, other=0.0)) * \
+                tl.load(h_base + I + 2 * j, mask=kmask, other=0.0)
+            ao = _gelu_tanh(tl.load(h_base + 2 * j + 1, mask=kmask, other=0.0)) * \
+                tl.load(h_base + I + 2 * j + 1, mask=kmask, other=0.0)
+            mask = rmask[:, None] & kmask[None, :]
+            b = tl.load(q_base + rows[:, None] * (I // 2) + j[None, :], mask=mask,
+                        other=0).to(tl.int32)
+            s = tl.load(s_base + rows[:, None] * (I // 64) + (j // 32)[None, :], mask=mask,
+                        other=0.0).to(tl.float32)
+            lo = (b & 15).to(tl.float32) - 8.0
+            hi = (b >> 4).to(tl.float32) - 8.0
+            acc += tl.sum((lo * ae[None, :] + hi * ao[None, :]) * s, axis=1)
+        tl.store(y_ptr + pair * H + rows, acc * tl.load(w_ptr + pair), mask=rmask)
+
+
+def moe_int4_decode(x, weights, ids, table, w13, w13s, w2, w2s, blocks=(16, 256, 32, 128)):
+    """sum over routes r of weights[m, r] * expert(x[m]) for the routes whose expert `table`
+    maps to a slot (others add nothing). x [M, H] bf16 with few rows; ids [M, TOPK], padding
+    pointing at the table's last row (-1); slot tensors in the GPU int4 layout above."""
+    m, topk = ids.shape
+    hidden, n2 = x.shape[1], w13.shape[1]
+    pairs = m * topk
+    bn1, bk1, bn2, bk2 = blocks
+    assert hidden % bk1 == 0, "gate/up blocks must tile the hidden size (no k mask there)"
+    x, ids = x.contiguous(), ids.contiguous()
+    h = torch.empty(pairs, n2, dtype=torch.float32, device=x.device)
+    y = torch.zeros(pairs, hidden, dtype=torch.float32, device=x.device)
+    _gate_up[(pairs, triton.cdiv(n2, bn1))](x, ids, table, w13, w13s, h, TOPK=topk, H=hidden,
+                                            N2=n2, BLOCK_N=bn1, BLOCK_K=bk1, num_warps=4)
+    _down[(pairs, triton.cdiv(hidden, bn2))](h, ids, table, weights.float().contiguous(), w2,
+                                             w2s, y, H=hidden, I=n2 // 2, BLOCK_N=bn2,
+                                             BLOCK_K=bk2, num_warps=4)
+    return y.view(m, topk, hidden).sum(1).to(x.dtype)
+
+
+@triton.jit
+def _unpack(src_ptr, dst_ptr, N: tl.constexpr, K: tl.constexpr, BLOCK_R: tl.constexpr,
+            BLOCK_J: tl.constexpr):
+    """CPU int4 layout (torch's AVX2 order: [N/32, K, 16] bytes, rows d and d+16 share a byte)
+    -> GPU layout [N, K/2]; prain.quant.to_gpu_int4 in one pass."""
+    rows = tl.program_id(0) * BLOCK_R + tl.arange(0, BLOCK_R)
+    j = tl.program_id(1) * BLOCK_J + tl.arange(0, BLOCK_J)
+    mask = (rows < N)[:, None] & (j < K // 2)[None, :]
+    r = rows % 32
+    base = (rows // 32) * (K * 16) + r % 16
+    shift = ((r // 16) * 4)[:, None]
+    even = tl.load(src_ptr + base[:, None] + (2 * j * 16)[None, :], mask=mask, other=0)
+    odd = tl.load(src_ptr + base[:, None] + ((2 * j + 1) * 16)[None, :], mask=mask, other=0)
+    lo = (even.to(tl.int32) >> shift) & 15
+    hi = (odd.to(tl.int32) >> shift) & 15
+    tl.store(dst_ptr + rows[:, None] * (K // 2) + j[None, :], (lo | (hi << 4)).to(tl.uint8),
+             mask=mask)
+
+
+def unpack_int4(src, dst, block=(32, 64)):
+    """dst (uint8 [N, K/2], GPU layout) <- src (the same bytes in the CPU layout, flat)."""
+    n, half = dst.shape
+    _unpack[(triton.cdiv(n, block[0]), triton.cdiv(half, block[1]))](
+        src, dst, N=n, K=2 * half, BLOCK_R=block[0], BLOCK_J=block[1])

prain/plugin.py

@@ -0,0 +1,18 @@
+"""vLLM general plugin entry point (see pyproject.toml). Inactive unless PRAIN_CONFIG is set."""
+
+import os
+
+
+def register():
+    if not os.environ.get("PRAIN_CONFIG"):
+        return
+    from vllm.model_executor.custom_op import PluggableLayer, op_registry_oot
+
+    if "RoutedExperts" in op_registry_oot:  # vLLM may call register() more than once
+        return
+    from prain import gemv
+    from prain.experts import TieredRoutedExperts
+
+    # vLLM's documented out-of-tree hook: every RoutedExperts(...) now builds ours.
+    PluggableLayer.register_oot(TieredRoutedExperts, name="RoutedExperts")
+    gemv.install()  # decode GEMVs on gfx10

prain/policy.py

@@ -0,0 +1,247 @@
+"""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": [...], "vram": [experts in VRAM now],
+    "hits": n, ...}}}."""
+    data = {"layers": {str(layer): {"uses": c.uses, "vram": sorted({*c.pinned, *c.slots}),
+                                    **{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/profile.py

@@ -0,0 +1,277 @@
+"""Span profiler for the vLLM engine process: flame-graph data per engine step.
+
+Active only with PRAIN_PROFILE=DIR. Writing a number N into DIR/start records the next N engine
+steps: the step, the runner's execute_model and sample_tokens, every module forward of the model
+and prain's MoE phases become nested spans. Each span syncs the GPU at both ends, so GPU work is
+charged to the span that queued it (the overlap of CPU and GPU work is lost; a sync costs ~10 us).
+Then it writes DIR/decode.folded and DIR/prefill.folded ("frame;frame;frame microseconds" lines:
+time spent in the last frame itself, summed over the steps of that kind; a step with more than
+DECODE_TOKENS tokens is a prefill step), DIR/steps.json, and removes DIR/start.
+scripts/flamegraph.py draws them.
+"""
+
+import contextlib
+import json
+import os
+import threading
+import time
+
+import torch
+
+DIR = os.environ.get("PRAIN_PROFILE")
+SYNC = [torch.cuda.synchronize]  # tests replace it
+DECODE_TOKENS = 4
+
+
+class _Recorder:
+    def __init__(self):
+        self.active, self.left, self.thread = False, 0, None
+        self.stack = []  # [frame, start, time in children]
+        self.totals, self.steps, self.hooks = {}, {}, []
+        self.kind = "decode"
+
+
+_R = _Recorder()
+_NULL = contextlib.nullcontext()
+
+
+def _enter(frame):
+    SYNC[0]()
+    _R.stack.append([frame, time.perf_counter(), 0.0])
+
+
+def _exit():
+    SYNC[0]()
+    frame, start, children = _R.stack.pop()
+    took = time.perf_counter() - start
+    stack = (*(f[0] for f in _R.stack), frame)
+    _R.step[stack] = _R.step.get(stack, 0.0) + took - children
+    _R.step_calls[stack] = _R.step_calls.get(stack, 0) + 1
+    if _R.stack:
+        _R.stack[-1][2] += took
+
+
+def _recording():
+    return _R.active and threading.get_ident() == _R.thread
+
+
+def _pre_hook(frame):
+    def hook(module, args):  # must return None: anything else replaces the module's inputs
+        if _recording():
+            _enter(frame)
+    return hook
+
+
+def _post_hook(module, args, out):  # must return None, or it replaces the output
+    if _recording():
+        _exit()
+
+
+@contextlib.contextmanager
+def _span(name):
+    _enter(name)
+    try:
+        yield
+    finally:
+        _exit()
+
+
+def span(name):
+    """with span("copy"): ... -- a no-op unless a recording runs."""
+    return _span(name) if _recording() else _NULL
+
+
+def step_begin():
+    if not _R.active and DIR and os.path.exists(os.path.join(DIR, "start")):
+        with open(os.path.join(DIR, "start")) as f:
+            text = f.read().strip()
+        _R.left = int(text) if text.isdigit() else 1
+        _R.active, _R.thread, _R.totals, _R.steps = True, threading.get_ident(), {}, {}
+        _R.calls = {}
+    if _recording():
+        _R.step, _R.step_calls, _R.kind = {}, {}, "decode"
+        _enter("step")
+
+
+def note_step(model, tokens):
+    """In execute_model: the step's token count, and the module hooks on first use."""
+    if not _recording():
+        return
+    if tokens > DECODE_TOKENS:
+        _R.kind = "prefill"
+    if not _R.hooks:
+        for name, module in model.named_modules():
+            attr, cls = name.rsplit(".", 1)[-1], type(module).__name__
+            frame = cls if not name or attr.isdigit() else f"{attr}:{cls}"
+            _R.hooks.append(module.register_forward_pre_hook(_pre_hook(frame)))
+            _R.hooks.append(module.register_forward_hook(_post_hook))
+
+
+def step_end():
+    if not _recording():
+        return
+    while _R.stack:  # also closes spans an exception left open
+        _exit()
+    for stack, seconds in _R.step.items():
+        key = (_R.kind, *stack)
+        _R.totals[key] = _R.totals.get(key, 0.0) + seconds
+        _R.calls[key] = _R.calls.get(key, 0) + _R.step_calls[stack]
+    _R.steps[_R.kind] = _R.steps.get(_R.kind, 0) + 1
+    _R.left -= 1
+    if _R.left <= 0:
+        _finish()
+
+
+def _finish():
+    _R.active = False
+    for hook in _R.hooks:
+        hook.remove()
+    _R.hooks = []
+    for kind in _R.steps:
+        for suffix, values, fmt in (("folded", _R.totals, lambda s: round(s * 1e6)),
+                                    ("calls", _R.calls, int)):
+            with open(os.path.join(DIR, f"{kind}.{suffix}"), "w") as f:
+                for key, value in sorted(values.items()):
+                    if key[0] == kind:
+                        f.write(f"{';'.join(key[1:])} {fmt(value)}\n")
+    with open(os.path.join(DIR, "steps.json"), "w") as f:
+        json.dump({"steps": _R.steps, "decode_tokens": DECODE_TOKENS}, f)
+    with contextlib.suppress(OSError):
+        os.remove(os.path.join(DIR, "start"))
+
+
+def install():
+    """Spans around vLLM's engine step and the model runner (the vLLM process with the GPU)."""
+    from vllm.v1.engine.core import EngineCoreProc
+    from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
+    def wrap(cls, name, around):
+        original = getattr(cls, name)
+        if not getattr(original, "prain_profile", False):
+            wrapped = around(original)
+            wrapped.prain_profile = True
+            setattr(cls, name, wrapped)
+
+    def engine_step(original):
+        def step(self, *args, **kwargs):
+            step_begin()
+            try:
+                return original(self, *args, **kwargs)
+            finally:
+                step_end()
+        return step
+
+    def execute_model(original):
+        def execute(self, scheduler_output, *args, **kwargs):
+            with span("execute_model"):
+                note_step(self.model, getattr(scheduler_output, "total_num_scheduled_tokens", 0))
+                return original(self, scheduler_output, *args, **kwargs)
+        return execute
+
+    def sample_tokens(original):
+        def sample(self, *args, **kwargs):
+            with span("sample_tokens"):
+                return original(self, *args, **kwargs)
+        return sample
+
+    wrap(EngineCoreProc, "_process_engine_step", engine_step)
+    wrap(GPUModelRunner, "execute_model", execute_model)
+    wrap(GPUModelRunner, "sample_tokens", sample_tokens)
+
+
+_COLORS = (("copy", "#e8603c"), ("round", "#f0a030"), ("expert", "#f0a030"), ("moe", "#f0a030"),
+           ("route", "#e8c440"), ("plan", "#e8c440"), ("degrade", "#e8c440"), ("cpu", "#e8c440"),
+           ("attn", "#5b9bd5"), ("rotary", "#5b9bd5"), ("mlp", "#6abf69"), ("gate_up", "#6abf69"),
+           ("down_proj", "#6abf69"), ("norm", "#a0a0a0"), ("logits", "#a070c0"),
+           ("lm_head", "#a070c0"), ("sample", "#a070c0"), ("embed", "#a070c0"))
+
+
+def _tree(lines):
+    root = {"name": "all", "self": 0.0, "kids": {}}
+    for line in lines:
+        stack, _, us = line.strip().rpartition(" ")
+        node = root
+        for frame in stack.split(";"):
+            node = node["kids"].setdefault(frame, {"name": frame, "self": 0.0, "kids": {}})
+        node["self"] += float(us)
+
+    def total(node):
+        node["total"] = node["self"] + sum(total(k) for k in node["kids"].values())
+        return node["total"]
+    total(root)
+    return root
+
+
+def top(lines, steps=1, n=25, calls=None):
+    """[(frame, ms per step, % of all, calls per step or None)] by time in the frame itself;
+    calls: the .calls lines of the same recording."""
+    own, count = {}, {}
+    for line in lines:
+        stack, _, us = line.strip().rpartition(" ")
+        frame = stack.rsplit(";", 1)[-1]
+        own[frame] = own.get(frame, 0.0) + float(us)
+    for line in calls or ():
+        stack, _, n_calls = line.strip().rpartition(" ")
+        frame = stack.rsplit(";", 1)[-1]
+        count[frame] = count.get(frame, 0) + int(n_calls)
+    whole = sum(own.values()) or 1.0
+    return [(frame, round(us / 1e3 / steps, 1), round(100 * us / whole, 1),
+             round(count[frame] / steps, 1) if frame in count else None)
+            for frame, us in sorted(own.items(), key=lambda kv: -kv[1])[:n]]
+
+
+def svg(lines, steps=1, title="", width=1200, row=18):
+    """A flame graph of folded stack lines (root on top, widest child first); each box's
+    tooltip gives its time per step, its share of all time and its own (self) time."""
+    from html import escape
+
+    root = _tree(lines)
+    whole = root["total"] or 1.0
+    boxes = []
+
+    def place(node, x, depth):
+        boxes.append((node, x, depth, node["total"] / whole * width))
+        for kid in sorted(node["kids"].values(), key=lambda k: -k["total"]):
+            place(kid, x, depth + 1)
+            x += kid["total"] / whole * width
+    place(root, 0.0, 0)
+    height = (max(d for _, _, d, _ in boxes) + 1) * row + 30
+    out = [f'<svg xmlns="http://www.w3.org/2000/svg" width="{width}" height="{height}" '
+           f'font-family="sans-serif" font-size="11">',
+           f'<text x="4" y="16" font-size="14">{escape(title)}: '
+           f'{whole / 1e3 / steps:.1f} ms per step over {steps} steps (hover for times)</text>']
+    for node, x, depth, w in boxes:
+        if w < 0.5:
+            continue
+        name, y = node["name"], 26 + depth * row
+        color = next((c for key, c in _COLORS if key in name.lower()), "#d8b890")
+        ms = node["total"] / 1e3 / steps
+        text = f"{name} {ms:.1f} ms" if ms >= 0.1 else name  # readable without hovering
+        label = text if len(text) * 6.5 < w - 4 else text[:max(int((w - 4) / 6.5) - 1, 0)] + "…"
+        out.append(
+            f'<g><title>{escape(name)}: {ms:.1f} ms per step '
+            f'({100 * node["total"] / whole:.1f}%), self {node["self"] / 1e3 / steps:.1f} ms</title>'
+            f'<rect x="{x:.1f}" y="{y}" width="{w:.1f}" height="{row - 1}" fill="{color}" rx="2"/>'
+            + (f'<text x="{x + 3:.1f}" y="{y + 12}">{escape(label)}</text>' if w > 14 else "")
+            + "</g>")
+    return "\n".join(out + ["</svg>"])
+
+
+if __name__ == "__main__":  # python -m prain.profile DIR: DIR/KIND.svg and a summary per kind
+    import sys
+
+    directory = sys.argv[1]
+    with open(os.path.join(directory, "steps.json")) as f:
+        counts = json.load(f)["steps"]
+    for kind, steps in counts.items():
+        with open(os.path.join(directory, f"{kind}.folded")) as f:
+            lines = f.read().splitlines()
+        calls_path = os.path.join(directory, f"{kind}.calls")
+        calls = open(calls_path).read().splitlines() if os.path.exists(calls_path) else None
+        with open(os.path.join(directory, f"{kind}.svg"), "w") as f:
+            f.write(svg(lines, steps, title=kind))
+        print(f"== {kind}: {steps} steps, {_tree(lines)['total'] / 1e3 / steps:.1f} ms per step"
+              f"  (self time per step, share, calls per step, ms per call)")
+        for frame, ms, share, n in top(lines, steps, calls=calls):
+            per_call = f"{n:8.1f} x {ms / n:7.3f} ms" if n else ""
+            print(f"{ms:9.1f} ms {share:5.1f}%  {per_call}  {frame}")

prain/quant.py

@@ -0,0 +1,89 @@
+"""Weight-only int8 / int4 quantization of expert weights. Needs torch.
+
+int8: q int8 [..., N, K] and one fp32 scale per output row: w = q * scale. vLLM's Triton
+      fused-MoE kernel and torch's _weight_int8pack_mm (CPU) both read this layout.
+int4: one bf16 scale per GROUP inputs of a row: w = (nibble - 8) * scale, with the scale's
+      clip ratio chosen per group for the least squared error. Two layouts of the nibbles:
+      gpu: [..., N, K/2], two values per byte along K, even k in the low nibble (vLLM's
+           Triton GPTQ/AWQ MoE kernel);
+      cpu: [N/32, K, 16] bytes, rows d and d+16 of each 32-row block share a byte (torch's AVX2
+           _weight_int4pack_mm_for_cpu); packs store it, slots convert with to_gpu_int4().
+"""
+
+import torch
+
+from prain.config import INT4_GROUP as GROUP  # divides Gemma 4 26B-A4B's 2816 and 704
+
+CLIP = (1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7)
+
+
+def layout(quant, rows, cols):
+    """[(shape, dtype)] of the tensors that store one [rows, cols] weight matrix."""
+    if quant == "bf16":
+        return [((rows, cols), torch.bfloat16)]
+    if quant == "int8":
+        return [((rows, cols), torch.int8), ((rows,), torch.float32)]
+    return [((rows, cols // 2), torch.uint8), ((rows, cols // GROUP), torch.bfloat16)]
+
+
+def quantize(w, quant):
+    """Float weights [..., N, K] -> [q, scale], int4 in the gpu layout."""
+    w = w.float()
+    if quant == "int8":
+        scale = w.abs().amax(-1).clamp_min(1e-12) / 127
+        return [(w / scale[..., None]).round().to(torch.int8), scale]
+    g = w.unflatten(-1, (-1, GROUP))
+    amax = g.abs().amax(-1, keepdim=True).clamp_min(1e-12)
+    best_err = best_scale = None
+    for ratio in CLIP:
+        scale = (amax * ratio / 7).bfloat16().float()  # the kernels see the bf16 scale
+        err = ((g / scale).round().clamp(-8, 7) * scale - g).pow(2).sum(-1, keepdim=True)
+        if best_err is None:
+            best_err, best_scale = err, scale
+        else:
+            better = err < best_err
+            best_err, best_scale = torch.where(better, err, best_err), torch.where(better, scale, best_scale)
+    q = ((g / best_scale).round().clamp(-8, 7) + 8).flatten(-2).to(torch.uint8)
+    return [q[..., 0::2] | q[..., 1::2] << 4, best_scale.squeeze(-1).bfloat16()]
+
+
+def dequantize(q, scale, quant, dtype=torch.bfloat16):
+    """The weights [..., N, K] of quantize()'s output (int4 in the gpu layout)."""
+    if quant == "int8":
+        return q.to(dtype) * scale.to(dtype)[..., None]
+    v = torch.stack((q & 15, q >> 4), -1).flatten(-2).to(dtype) - 8
+    return (v.unflatten(-1, (-1, GROUP)) * scale.to(dtype)[..., None]).flatten(-2)
+
+
+def to_cpu_int4(q):
+    """int4 nibbles [..., N, K/2]: gpu layout -> cpu layout (same shape)."""
+    *lead, n, half = q.shape
+    v = torch.stack((q & 15, q >> 4), -1).reshape(*lead, n // 32, 32, 2 * half).transpose(-1, -2)
+    return (v[..., :16] | v[..., 16:] << 4).reshape(q.shape)
+
+
+def to_gpu_int4(c):
+    """int4 nibbles [..., N, K/2]: cpu layout -> gpu layout (same shape; works on the GPU)."""
+    *lead, n, half = c.shape
+    b = c.reshape(*lead, n // 32, 2 * half, 16)
+    v = torch.cat((b & 15, b >> 4), -1).transpose(-1, -2).reshape(*lead, n, 2 * half)
+    return v[..., 0::2] | v[..., 1::2] << 4
+
+
+def cpu_matmul(quant, q, scale):
+    """x -> x @ w.T on the CPU for stored int8 q / cpu-layout int4 q; x is bf16 [M, K]."""
+    if quant == "int8":
+        return lambda x: torch._weight_int8pack_mm(x, q, scale.to(x.dtype))
+    zeros = torch.stack((scale.t(), torch.zeros_like(scale.t())), -1).contiguous()
+    return lambda x: torch._weight_int4pack_mm_for_cpu(x, q, GROUP, zeros)
+
+
+def moe_quant_config(quant, w13_scale, w2_scale):
+    """vLLM's fused-MoE quant config for int8 / gpu-layout int4 slot tensors."""
+    from vllm.model_executor.layers.fused_moe import config
+
+    if quant == "int8":
+        return config.int8_w8a16_moe_quant_config(w1_scale=w13_scale, w2_scale=w2_scale,
+                                                  w1_zp=None, w2_zp=None)
+    return config.int4_w4a16_moe_quant_config(w1_scale=w13_scale, w2_scale=w2_scale, w1_zp=None,
+                                              w2_zp=None, block_shape=[0, GROUP])

prain/sim.py

@@ -0,0 +1,320 @@
+"""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 read_ms(self):
+        """Reading one expert from the pack (the CPU path pays this for disk experts too)."""
+        return self.expert_mib * 1.048576 / self.disk_gbps
+
+    def copy_ms(self, disk):
+        mb = self.expert_mib * 1.048576
+        return mb / self.h2d_gbps + (self.read_ms() 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)
+                       + cost.read_ms() * sum(e in disk[layer] 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)),
+    # One dimension: cpu_max_tokens only matters for auto, and coordinate descent would never
+    # reach auto with a high threshold once auto <= 1 lost to cpu.
+    ("miss", ("mode", "cpu_max_tokens"),
+     (("copy", 1), ("cpu", 1)) + tuple(("auto", n) for n in (1, 2, 4, 8, 16))),
+    ("miss", "admit", (0, 1, 2)),
+)
+_PART = {"policy": 0, "prefetch": 1, "miss": 2}
+
+
+def _with(variant, table, name, value):
+    names, values = (name, value) if isinstance(name, tuple) else ((name,), (value,))
+    parts = list(variant)
+    parts[_PART[table]] = replace(parts[_PART[table]], **dict(zip(names, values)))
+    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/tiered.py

@@ -0,0 +1,173 @@
+"""Run one MoE layer when only some experts fit in VRAM. Needs torch, not vLLM."""
+
+import torch
+import torch.nn.functional as F
+
+from prain import profile
+
+HEAT_INDEX = {"count": 0, "weight": 1, "recency": 2}
+
+
+def route_stats(topk_weights, topk_ids, num_experts, recency_half_life):
+    """{expert: (tokens, router weight mass, recency-weighted mass)} for every routed expert,
+    with one host sync. Later tokens weigh more for recency. Route id -1 counts for nothing."""
+    tokens = topk_ids.shape[0]
+    valid = (topk_ids >= 0).float()
+    ids = topk_ids.clamp_min(0).long().flatten()
+    position = torch.arange(tokens, device=topk_ids.device, dtype=torch.float32)
+    decay = torch.pow(0.5, (tokens - 1 - position) / recency_half_life)[:, None]
+    weight = topk_weights.float() * valid
+    stats = torch.stack([
+        torch.bincount(ids, weights=valid.flatten(), minlength=num_experts),
+        torch.bincount(ids, weights=weight.flatten(), minlength=num_experts),
+        torch.bincount(ids, weights=(weight * decay).flatten(), minlength=num_experts),
+    ]).tolist()
+    return {e: (round(stats[0][e]), stats[1][e], stats[2][e])
+            for e in range(num_experts) if stats[0][e] > 0}
+
+
+def draft_rows(spans, rows, device):
+    """Bool [rows]: True for the token rows of draft requests (spans from prain.drafts)."""
+    mask = torch.zeros(rows, dtype=torch.bool)
+    for first, end in spans:
+        mask[first:min(end, rows)] = True
+    return mask.to(device)
+
+
+def available_experts(topk_ids, draft, resident, num_experts):
+    """Bool [num_experts]: experts in VRAM, plus every expert a full (non-draft) token routes
+    to, because those are copied in this forward anyway."""
+    ids = topk_ids.long()
+    full_routes = ((ids >= 0) & ~draft[:, None]).flatten().float()
+    available = torch.bincount(ids.clamp_min(0).flatten(), weights=full_routes,
+                               minlength=num_experts) > 0
+    if resident:
+        available[torch.tensor(sorted(resident), device=ids.device)] = True
+    return available
+
+
+def degrade(topk_weights, topk_ids, draft, available, min_weight):
+    """Let draft tokens skip experts that are not available.
+
+    A draft token keeps its heaviest unavailable routes only while less than min_weight of its
+    router weight would be computed; its kept routes are scaled up to its full router weight.
+    Dropped routes get id -1 and weight 0, which route_stats and slot_table ignore.
+    draft: bool [tokens]; available: bool [num_experts]. Full tokens are left unchanged."""
+    ids = topk_ids.long()
+    valid = ids >= 0
+    w = topk_weights.to(torch.promote_types(topk_weights.dtype, torch.float32)) * valid
+    missing = valid & draft[:, None] & ~available[ids.clamp_min(0)]
+    total = w.sum(1, keepdim=True)
+    short = (min_weight * total - (w * ~missing).sum(1, keepdim=True)).clamp_min(0)
+    order = (w * missing).argsort(1, descending=True)  # heaviest missing route first
+    heavy = torch.gather(w * missing, 1, order)
+    keep_sorted = (heavy > 0) & (heavy.cumsum(1) - heavy < short)  # still short before it
+    keep = torch.zeros_like(keep_sorted).scatter(1, order, keep_sorted)
+    drop = missing & ~keep
+    kept = (w * ~drop).sum(1, keepdim=True)
+    scale = torch.where(kept > 0, total / kept, torch.ones_like(kept))
+    weights = torch.where(drop, torch.zeros_like(w), w * scale).to(topk_weights.dtype)
+    return weights, torch.where(drop, torch.full_like(topk_ids, -1), topk_ids)
+
+
+def slot_table(slot_of, num_experts, device):
+    """int32 [num_experts + 1]: VRAM slot of each expert in one round, -1 for the others.
+    vLLM's fused-MoE kernel takes it as expert_map and skips routes that map to -1; the extra
+    last row stands for the padding id (route id -1 arrives as num_experts)."""
+    table = torch.full((num_experts + 1,), -1, dtype=torch.int32)
+    if slot_of:
+        table[torch.tensor(list(slot_of))] = torch.tensor(list(slot_of.values()), dtype=torch.int32)
+    return table.to(device)
+
+
+def _matmul(w):
+    """x -> x @ w.T for a weight tensor; functions (quantized experts) pass through."""
+    return w if callable(w) else (lambda v: v.to(w.dtype) @ w.t())
+
+
+def cpu_experts(x, topk_weights, topk_ids, experts, weights_of):
+    """The routes to `experts`, computed on the CPU like vLLM's fused kernel with activation
+    "gelu_tanh": w13 = [gate; up], y = (gelu_tanh(x @ gate.T) * (x @ up.T)) @ w2.T times the
+    route weight. weights_of(e) -> (w13 [2I, H], w2 [H, I]), or two functions v -> v @ w.T
+    (quantized experts). Returns [tokens, hidden] fp32."""
+    out = torch.zeros(x.shape, dtype=torch.float32)
+    for e in sorted(experts):
+        rows, cols = (topk_ids == e).nonzero(as_tuple=True)
+        if rows.numel() == 0:
+            continue
+        mm13, mm2 = map(_matmul, weights_of(e))
+        h = mm13(x[rows])
+        half = h.shape[1] // 2
+        act = F.gelu(h[:, :half].float(), approximate="tanh") * h[:, half:].float()
+        y = mm2(act.to(h.dtype)).float()
+        out.index_add_(0, rows, y * topk_weights[rows, cols].float()[:, None])
+    return out
+
+
+def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight",
+                   recency_half_life=32.0, miss=None, cpu_weights=None, on_stats=None,
+                   after_enqueue=None, draft=None, min_weight=0.5):
+    """One MoE forward over VRAM slots.
+
+    cache: SlotCache. load(expert, slot) copies an expert into a slot. kernel(x, weights, ids,
+    table, used_slots) runs the fused-MoE kernel on the slot tensors for the routes whose
+    expert the slot_table() `table` maps to a slot. With a
+    [miss] config other than "copy", the experts it picks run on the CPU (cpu_weights(e)
+    gives their weights) while the GPU rounds run. on_stats(stats, tokens) sees the routing
+    statistics; after_enqueue(cpu, stats) runs once all GPU work of this forward is queued.
+    draft: bool [tokens] for the rows of draft requests, or None; see degrade(). Draft rows keep
+    min_weight of their router weight, every row keeps [miss] min_weight (1 = everything).
+    """
+    raw = None
+    skip = miss is not None and miss.min_weight < 1
+    if draft is not None or skip:
+        with profile.span("degrade"):
+            if on_stats is not None:  # the trace records the routing, not what was skipped
+                raw = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)
+            rows = draft if draft is not None else torch.zeros(
+                topk_ids.shape[0], dtype=torch.bool, device=topk_ids.device)
+            need = min_weight
+            if skip:  # every row may skip; draft rows by the smaller share
+                need = torch.where(rows, min(min_weight, miss.min_weight), miss.min_weight)[:, None]
+                rows = torch.ones_like(rows)
+            available = available_experts(topk_ids, rows, set(cache.pinned) | set(cache.slots),
+                                          cache.num_experts)
+            _, kept = degrade(topk_weights, topk_ids, rows, available, need)
+            # experts some row copies anyway serve every row for free
+            available |= torch.bincount(kept[kept >= 0].long(), minlength=cache.num_experts) > 0
+            topk_weights, topk_ids = degrade(topk_weights, topk_ids, rows, available, need)
+    with profile.span("route"):
+        stats = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)  # sync
+    tokens = int(topk_ids.shape[0])
+    if on_stats is not None:
+        on_stats(raw if raw is not None else stats, tokens)
+    counts = {e: s[0] for e, s in stats.items()}
+    cpu = cache.choose_cpu(counts, miss) if miss is not None else set()
+    if cpu:  # copy the inputs now: later the copy would wait for the queued GPU rounds
+        host = x.cpu(), topk_weights.cpu(), topk_ids.cpu()
+    index = HEAT_INDEX[heat]
+    ids = torch.where(topk_ids >= 0, topk_ids, cache.num_experts)  # padding -> the table's last row
+    acc = torch.promote_types(x.dtype, torch.float32)
+    total = None  # sum of the round (and CPU) outputs: one buffer, not one per round
+
+    def add(out):
+        nonlocal total
+        total = out if total is None else total.to(acc).add_(out)  # .to: a no-op once fp32
+
+    with profile.span("plan"):
+        rounds = cache.plan({e: s[index] for e, s in stats.items()}, tokens, cpu)
+    for slot_of, loads in rounds:
+        for expert, slot in loads:
+            with profile.span("copy"):
+                load(expert, slot)
+        if slot_of:  # each round computes only its own routes
+            with profile.span("round"):
+                table = slot_table(slot_of, cache.num_experts, x.device)
+                add(kernel(x, topk_weights, ids, table, list(slot_of.values())))
+    if after_enqueue is not None:
+        with profile.span("prefetch"):
+            after_enqueue(cpu, stats)
+    if cpu:  # runs while the GPU works through its rounds
+        with profile.span("cpu experts"):
+            add(cpu_experts(*host, cpu, cpu_weights).to(x.device))
+    return torch.zeros_like(x) if total is None else total.to(x.dtype)

prain/tui.py

@@ -0,0 +1,118 @@
+"""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, with_system)
+
+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.
+    system: the persona, sent ahead of every request but not kept in `messages`."""
+
+    def __init__(self, base_url, model, warmup="idle", priority=1, system=None):
+        self.base_url, self.model, self.warmup, self.system = base_url, model, warmup, system
+        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, with_system(system, 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=(with_system(self.system, 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, system=None):
+    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, system)
+    curses.wrapper(run, chat)

prain/web.html

@@ -0,0 +1,717 @@
+<!doctype html>
+<html lang="en">
+<meta charset="utf-8">
+<meta name="viewport" content="width=device-width, initial-scale=1, viewport-fit=cover">
+<meta name="color-scheme" content="light dark">
+<title>prain</title>
+<style>
+:root {
+  --bg: #f1f3f6; --surface: #fff; --ink: #161b22; --muted: #5b6673; --faint: #8a94a0;
+  --line: #dde2e8; --bubble: #e5e9ef; --code: #f6f7f9; --focus: #7c4dcc; --live: #f26b1d;
+  --ok: #23946a; --bad: #c2362f;
+  --ui: system-ui, -apple-system, "Segoe UI", Roboto, "Noto Sans", Cantarell, sans-serif;
+  --display: ui-rounded, "SF Pro Rounded", Nunito, system-ui, -apple-system, sans-serif;
+  --mono: ui-monospace, "SF Mono", "JetBrains Mono", "Cascadia Code", "DejaVu Sans Mono", Menlo, Consolas, monospace;
+}
+@media (prefers-color-scheme: dark) {
+  :root { --bg: #0f1216; --surface: #161a20; --ink: #e6eaf0; --muted: #9aa4b1; --faint: #6d7784;
+          --line: #262d36; --bubble: #232a34; --code: #1b2027; --focus: #a58bea; }
+}
+* { box-sizing: border-box; }
+html, body { height: 100%; margin: 0; }
+body { background: var(--bg); color: var(--ink); font: 16px/1.6 var(--ui); -webkit-font-smoothing: antialiased; }
+button { font: inherit; color: inherit; }
+:focus-visible { outline: 2px solid var(--focus); outline-offset: 2px; }
+[hidden] { display: none !important; }
+
+.app { display: grid; grid-template-columns: minmax(0, 1fr) 340px; height: 100dvh; }
+.chat { display: flex; flex-direction: column; min-width: 0; height: 100dvh; }
+
+/* header */
+.bar { display: flex; align-items: center; justify-content: space-between; gap: 12px; padding: 12px 20px; }
+.brand { display: flex; align-items: center; gap: 10px; min-width: 0; }
+.mark { display: grid; grid-template-columns: repeat(4, 6px); gap: 2px; }
+.mark i { width: 6px; height: 6px; border-radius: 2px; background: var(--line); transition: background .3s; }
+.mark i.on { background: var(--live); }
+.word { font: 700 20px/1 var(--display); letter-spacing: -.01em; }
+.model { font: 13px var(--mono); color: var(--faint); white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
+.bar-actions { display: flex; gap: 8px; align-items: center; }
+.ghost { border: 1px solid var(--line); background: var(--surface); border-radius: 999px; padding: 6px 14px; font-size: 14px; cursor: pointer; }
+.ghost:hover { border-color: var(--faint); }
+.pill { display: none; align-items: center; gap: 8px; }
+.dot { width: 8px; height: 8px; border-radius: 50%; background: var(--faint); flex: none; }
+.dot.ok { background: var(--ok); }
+.dot.live { background: var(--live); animation: pulse 1.2s ease-in-out infinite; }
+.dot.bad { background: var(--bad); }
+@keyframes pulse { 50% { opacity: .35; } }
+
+/* thread */
+.log { flex: 1; overflow-y: auto; padding: 8px 20px 24px; overscroll-behavior: contain; }
+.thread { max-width: 760px; margin: 0 auto; }
+.empty { max-width: 560px; margin: 14vh auto 0; }
+.empty h1 { font: 650 clamp(28px, 4.5vw, 38px)/1.15 var(--display); letter-spacing: -.015em; margin: 0 0 10px; }
+.empty p { color: var(--muted); margin: 0 0 6px; }
+.empty .live-line { font-size: 14px; color: var(--faint); margin-bottom: 22px; font-variant-numeric: tabular-nums; }
+.chips { display: flex; flex-wrap: wrap; gap: 8px; margin-top: 22px; }
+.chip { border: 1px solid var(--line); background: var(--surface); border-radius: 999px; padding: 8px 14px; font-size: 14px; cursor: pointer; text-align: left; }
+.chip:hover { border-color: var(--faint); }
+
+.msg { margin: 22px 0; }
+.msg.user { display: flex; justify-content: flex-end; }
+.bubble { background: var(--bubble); padding: 10px 16px; border-radius: 20px 20px 6px 20px; max-width: min(85%, 600px); white-space: pre-wrap; overflow-wrap: anywhere; }
+.body { overflow-wrap: anywhere; }
+.body > :first-child { margin-top: 0; }
+.body > :last-child { margin-bottom: 0; }
+.body p { margin: 0 0 .8em; }
+.body h3, .body h4, .body h5 { font: 650 1.08em/1.3 var(--display); margin: 1.3em 0 .4em; }
+.body ul, .body ol { padding-left: 1.4em; margin: 0 0 .8em; }
+.body li { margin: .25em 0; }
+.body blockquote { border-left: 3px solid var(--line); margin: .8em 0; padding-left: 12px; color: var(--muted); }
+.body hr { border: 0; border-top: 1px solid var(--line); margin: 1.2em 0; }
+.body a { color: var(--focus); }
+.body code { font: .88em var(--mono); background: var(--code); border: 1px solid var(--line); padding: .05em .35em; border-radius: 6px; }
+.code { border: 1px solid var(--line); background: var(--code); border-radius: 12px; margin: .9em 0; overflow: hidden; }
+.code-head { display: flex; justify-content: space-between; align-items: center; padding: 4px 6px 4px 14px; border-bottom: 1px solid var(--line); font: 12px var(--mono); color: var(--faint); }
+.code pre { margin: 0; padding: 12px 14px; overflow-x: auto; font: 13.5px/1.55 var(--mono); tab-size: 4; }
+.code code { background: none; border: 0; padding: 0; font-size: inherit; }
+.table { overflow-x: auto; margin: .9em 0; }
+.body table { border-collapse: collapse; font-size: 14px; }
+.body th, .body td { border-bottom: 1px solid var(--line); padding: 6px 12px 6px 0; text-align: left; vertical-align: top; }
+.small { border: 0; background: none; border-radius: 6px; padding: 3px 8px; font-size: 12.5px; color: var(--faint); cursor: pointer; }
+.small:hover { color: var(--ink); background: var(--bubble); }
+.meta { display: flex; flex-wrap: wrap; gap: 4px 12px; align-items: center; margin-top: 6px; font-size: 12.5px; color: var(--faint); font-variant-numeric: tabular-nums; }
+.note { display: inline-block; font-size: 13px; color: var(--muted); background: var(--bubble); border-radius: 10px; padding: 5px 10px; margin-bottom: 10px; }
+.note.bad { color: var(--bad); }
+.is-draft .body, .rest { color: var(--faint); }
+.typing::after { content: ""; display: inline-block; width: .5em; height: 1.05em; margin-left: 3px; vertical-align: -.18em; border-radius: 2px; background: var(--live); animation: pulse 1s ease-in-out infinite; }
+
+/* composer */
+.composer { padding: 0 20px calc(14px + env(safe-area-inset-bottom)); }
+.field { max-width: 760px; margin: 0 auto; display: flex; align-items: flex-end; gap: 8px; background: var(--surface); border: 1px solid var(--line); border-radius: 22px; padding: 7px 7px 7px 16px; box-shadow: 0 10px 30px -18px rgb(20 30 45 / .35); }
+.field:focus-within { border-color: color-mix(in srgb, var(--focus) 50%, var(--line)); }
+textarea { flex: 1; border: 0; outline: 0; resize: none; background: transparent; color: inherit; font: 16px/1.5 var(--ui); max-height: 40vh; padding: 6px 0; }
+.send { width: 36px; height: 36px; flex: none; border-radius: 50%; border: 0; background: var(--ink); color: var(--surface); display: grid; place-items: center; cursor: pointer; }
+.send:disabled { opacity: .3; cursor: default; }
+.send svg { width: 18px; height: 18px; }
+.hint { max-width: 760px; margin: 6px auto 0; min-height: 18px; padding: 0 16px; font-size: 12.5px; color: var(--faint); display: flex; align-items: center; gap: 8px; }
+.cells { display: inline-grid; grid-template-columns: repeat(8, 5px); gap: 2px; }
+.cells i { width: 5px; height: 5px; border-radius: 1.5px; background: var(--live); animation: pulse 1s ease-in-out infinite; }
+.cells i:nth-child(2n) { animation-delay: .25s; } .cells i:nth-child(3n) { animation-delay: .5s; }
+
+/* monitor */
+.monitor { background: var(--surface); border-left: 1px solid var(--line); overflow-y: auto; padding: 18px 20px 28px; font-size: 14px; }
+.m-top { display: flex; justify-content: space-between; align-items: start; }
+.m-title { font: 650 17px/1.2 var(--display); }
+.m-sub { font: 12.5px var(--mono); color: var(--faint); margin-top: 2px; }
+.close { display: none; }
+.state { display: flex; align-items: center; gap: 8px; margin: 14px 0 6px; color: var(--muted); }
+.speed { display: flex; align-items: baseline; gap: 6px; }
+.speed b { font: 650 40px/1 var(--display); letter-spacing: -.02em; font-variant-numeric: tabular-nums; }
+.speed.idle b { color: var(--faint); }
+.speed span { color: var(--faint); }
+.spark { width: 100%; height: 28px; margin: 6px 0 16px; display: block; }
+.spark path { fill: none; stroke: var(--live); stroke-width: 1.5; vector-effect: non-scaling-stroke; }
+.scan { background: #16121d; color: #c9c1d6; border-radius: 14px; padding: 12px 12px 10px; }
+.scan-head { display: flex; justify-content: space-between; font-size: 12.5px; margin-bottom: 8px; font-variant-numeric: tabular-nums; }
+.scan-head b { font-weight: 600; color: #f2edf8; }
+.scan canvas { display: block; width: 100%; aspect-ratio: 128 / 64; image-rendering: pixelated; border-radius: 4px; background: #2c2538; }
+.axis { display: flex; justify-content: space-between; font: 11px var(--mono); color: #8f86a0; margin-top: 4px; }
+.legend { display: flex; flex-wrap: wrap; gap: 4px 12px; margin: 8px 0 0; padding: 0; list-style: none; font-size: 12px; }
+.legend i { display: inline-block; width: 9px; height: 9px; border-radius: 2px; margin-right: 5px; vertical-align: -1px; }
+.scan-empty { font-size: 12.5px; color: #8f86a0; padding: 18px 4px; }
+.kv { margin: 16px 0 0; }
+.kv div { display: flex; justify-content: space-between; align-items: baseline; gap: 12px; padding: 7px 0; border-bottom: 1px solid var(--line); }
+.kv dt { color: var(--muted); }
+.kv dd { margin: 0; font-variant-numeric: tabular-nums; text-align: right; }
+.kv small { color: var(--faint); }
+.meter { display: block; width: 84px; height: 4px; border-radius: 2px; background: var(--bubble); margin-top: 5px; margin-left: auto; overflow: hidden; }
+.meter i { display: block; height: 100%; width: 0; background: var(--ink); transition: width .4s; }
+.m-h { font: 650 13px/1 var(--display); color: var(--faint); text-transform: uppercase; letter-spacing: .06em; margin: 24px 0 0; }
+.scrim { display: none; }
+
+@media (max-width: 999px) {
+  .app { grid-template-columns: 1fr; }
+  .pill, .close { display: inline-flex; }
+  .monitor { position: fixed; top: 0; right: 0; bottom: 0; width: min(360px, 92vw); z-index: 20; transform: translateX(100%); transition: transform .25s ease; box-shadow: -20px 0 40px -20px rgb(0 0 0 / .35); padding-top: calc(18px + env(safe-area-inset-top)); }
+  .monitor.open { transform: none; }
+  .scrim.open { display: block; position: fixed; inset: 0; z-index: 10; background: rgb(10 14 20 / .35); }
+}
+@media (max-width: 560px) {
+  .bar { padding: 10px 14px; } .log { padding: 6px 14px 18px; } .composer { padding-left: 10px; padding-right: 10px; }
+  .model, .hint .keys { display: none; } .empty { margin-top: 8vh; }
+}
+@media (prefers-reduced-motion: reduce) {
+  *, *::after { animation: none !important; transition: none !important; }
+}
+</style>
+
+<div class="app">
+  <main class="chat">
+    <header class="bar">
+      <div class="brand">
+        <span class="mark" id="mark" aria-hidden="true"><i></i><i></i><i></i><i></i><i></i><i></i><i></i><i></i></span>
+        <span class="word">prain</span>
+        <span class="model" id="model"></span>
+      </div>
+      <div class="bar-actions">
+        <button class="ghost pill" id="pill" aria-controls="monitor" aria-expanded="false">
+          <span class="dot" id="pill-dot"></span><span id="pill-text">Monitor</span></button>
+        <button class="ghost" id="new-chat">New chat</button>
+      </div>
+    </header>
+
+    <section class="log" id="log" role="log">
+      <div class="thread" id="thread">
+        <div class="empty" id="empty">
+          <h1>Ready when you are</h1>
+          <p>prain moves the experts your message needs into VRAM while you type, so the answer starts sooner.</p>
+          <div class="live-line" id="empty-live"></div>
+          <div class="chips">
+            <button class="chip">Explain what git rebase -i HEAD~3 does</button>
+            <button class="chip">Write a bash one-liner that counts lines of code per file type</button>
+            <button class="chip">On Gentoo, how do I rebuild the packages affected by a USE flag change?</button>
+          </div>
+        </div>
+      </div>
+    </section>
+
+    <form class="composer" id="composer">
+      <div class="field">
+        <textarea id="draft" rows="1" placeholder="Message prain" aria-label="Message"></textarea>
+        <button class="send" id="send" type="submit" aria-label="Send" disabled>
+          <svg viewBox="0 0 24 24" aria-hidden="true"><path id="send-icon" fill="currentColor" d="M12 4l7 7-1.4 1.4-4.6-4.6V20h-2V7.8l-4.6 4.6L5 11z"/></svg>
+        </button>
+      </div>
+      <div class="hint">
+        <span id="warm" hidden><span class="cells" aria-hidden="true"><i></i><i></i><i></i><i></i><i></i><i></i><i></i><i></i></span> Loading experts for your message</span>
+        <span class="keys" id="keys">Enter sends · Shift+Enter adds a line</span>
+      </div>
+    </form>
+  </main>
+
+  <div class="scrim" id="scrim"></div>
+  <aside class="monitor" id="monitor" aria-label="Model monitor">
+    <div class="m-top">
+      <div><div class="m-title">Monitor</div><div class="m-sub" id="m-model">connecting…</div></div>
+      <button class="ghost close" id="close" aria-label="Close the monitor">Close</button>
+    </div>
+    <div class="state"><span class="dot" id="dot"></span><span id="state">Connecting…</span></div>
+    <div class="speed idle" id="speed-box"><b id="speed">–</b><span id="speed-unit">tokens/s</span></div>
+    <svg class="spark" viewBox="0 0 60 28" preserveAspectRatio="none" aria-hidden="true"><path id="spark"/></svg>
+
+    <section class="scan" aria-label="Experts by layer">
+      <div class="scan-head"><span>Experts</span><span id="scan-count"></span></div>
+      <div id="scan-body">
+        <canvas id="scan" width="128" height="30" role="img" aria-label="Expert map"></canvas>
+        <div class="axis"><span>layer 1 · expert 1</span><span id="axis-end"></span></div>
+        <ul class="legend">
+          <li><i style="background:#6a35a0"></i>In VRAM</li>
+          <li><i style="background:#ff8a2a"></i>Used now</li>
+          <li><i style="background:#45d3c7"></i>Copied in</li>
+        </ul>
+      </div>
+      <div class="scan-empty" id="scan-empty" hidden>The map appears once prain writes expert stats: set stats_path in the placement file and pass it with --config.</div>
+    </section>
+
+    <dl class="kv">
+      <div><dt>VRAM hits</dt><dd id="hits">–</dd></div>
+      <div><dt>Copies</dt><dd id="copies">–</dd></div>
+      <div><dt>Context in use</dt><dd><span id="kv">–</span><span class="meter"><i id="kv-bar"></i></span></dd></div>
+      <div><dt>Prompt cache reuse</dt><dd id="prefix">–</dd></div>
+      <div><dt>Requests</dt><dd id="reqs">–</dd></div>
+    </dl>
+    <h3 class="m-h">GPU</h3>
+    <dl class="kv">
+      <div><dt>VRAM</dt><dd><span id="vram">–</span><span class="meter"><i id="vram-bar"></i></span></dd></div>
+      <div><dt>Load</dt><dd id="load">–</dd></div>
+      <div><dt>Temperature</dt><dd id="temp">–</dd></div>
+      <div><dt>Power</dt><dd id="power">–</dd></div>
+      <div><dt>Fan</dt><dd id="fan">–</dd></div>
+    </dl>
+  </aside>
+</div>
+
+<script>
+const $ = (id) => document.getElementById(id);
+const log = $("log"), thread = $("thread"), draft = $("draft"), send = $("send");
+const coarse = matchMedia("(pointer: coarse)").matches;
+const newId = () => Math.random().toString(36).slice(2) + Date.now().toString(36);
+let chat = load() || {session: newId(), messages: []};
+let streaming = false, controller = null;
+
+function load() {
+  try { return JSON.parse(localStorage.getItem("prain.chat")); } catch { return null; }
+}
+function save() {
+  try { localStorage.setItem("prain.chat", JSON.stringify(chat)); } catch {}
+}
+const plain = () => chat.messages.map(({role, content}) => ({role, content}));
+
+/* ---------- markdown: escape first, then a small safe subset ---------- */
+const esc = (s) => s.replace(/[&<>"']/g, (c) => ({"&": "&amp;", "<": "&lt;", ">": "&gt;", '"': "&quot;", "'": "&#39;"}[c]));
+function inline(text) {
+  return text.split(/(`[^`\n]+`)/g).map((part, i) => i % 2
+    ? `<code>${esc(part.slice(1, -1))}</code>`
+    : esc(part)
+        .replace(/\[([^\]]+)\]\((https?:\/\/[^\s)]+)\)/g, '<a href="$2" target="_blank" rel="noopener noreferrer">$1</a>')
+        .replace(/\*\*([^*]+)\*\*/g, "<strong>$1</strong>")
+        .replace(/(^|[^*\w])\*([^*\s][^*]*)\*(?![*\w])/g, "$1<em>$2</em>")
+  ).join("");
+}
+const ITEM = /^\s*([-*+]|\d+[.)])\s+/;
+function md(src) {
+  const lines = src.split("\n"), out = [];
+  let i = 0, para = [];
+  const flush = () => { if (para.length) out.push(`<p>${para.map(inline).join("<br>")}</p>`); para = []; };
+  while (i < lines.length) {
+    const line = lines[i];
+    let m;
+    if ((m = line.match(/^\s*```\s*([\w+#.-]*)/))) {
+      flush();
+      const code = [];
+      for (i++; i < lines.length && !/^\s*```\s*$/.test(lines[i]); i++) code.push(lines[i]);
+      i++;
+      out.push(`<div class="code"><div class="code-head"><span>${esc(m[1] || "code")}</span>` +
+               `<button class="small" data-copy>Copy</button></div><pre><code>${esc(code.join("\n"))}</code></pre></div>`);
+    } else if ((m = line.match(/^(#{1,6})\s+(.*)/))) {
+      flush();
+      const level = Math.min(m[1].length + 2, 5);
+      out.push(`<h${level}>${inline(m[2])}</h${level}>`);
+      i++;
+    } else if (/^\s*\|.*\|\s*$/.test(line) && /^\s*\|?\s*:?-{2,}/.test(lines[i + 1] || "")) {
+      flush();
+      const rows = [];
+      while (i < lines.length && /^\s*\|.*\|\s*$/.test(lines[i])) rows.push(lines[i++]);
+      const cells = (row, tag) => row.trim().replace(/^\||\|$/g, "").split("|")
+        .map((c) => `<${tag}>${inline(c.trim())}</${tag}>`).join("");
+      out.push(`<div class="table"><table><thead><tr>${cells(rows[0], "th")}</tr></thead><tbody>` +
+               rows.slice(2).map((r) => `<tr>${cells(r, "td")}</tr>`).join("") + "</tbody></table></div>");
+    } else if (ITEM.test(line)) {
+      flush();
+      const ordered = /^\s*\d/.test(line), items = [];
+      const start = ordered ? parseInt(line.trim(), 10) : 1;
+      while (i < lines.length && (ITEM.test(lines[i]) || (items.length && /^\s{2,}\S/.test(lines[i])))) {
+        if (ITEM.test(lines[i])) items.push(lines[i].replace(ITEM, ""));
+        else items[items.length - 1] += " " + lines[i].trim();
+        i++;
+      }
+      const tag = ordered ? "ol" : "ul";
+      out.push(`<${tag}${ordered && start !== 1 ? ` start="${start}"` : ""}>` +
+               items.map((t) => `<li>${inline(t)}</li>`).join("") + `</${tag}>`);
+    } else if (/^\s*>/.test(line)) {
+      flush();
+      const quote = [];
+      while (i < lines.length && /^\s*>/.test(lines[i])) quote.push(lines[i++].replace(/^\s*>\s?/, ""));
+      out.push(`<blockquote>${md(quote.join("\n"))}</blockquote>`);
+    } else if (/^\s*([-*_])(\s*\1){2,}\s*$/.test(line)) {
+      flush();
+      out.push("<hr>");
+      i++;
+    } else if (!line.trim()) {
+      flush();
+      i++;
+    } else {
+      para.push(line);
+      i++;
+    }
+  }
+  flush();
+  return out.join("");
+}
+
+/* ---------- thread ---------- */
+function nearBottom() { return log.scrollHeight - log.scrollTop - log.clientHeight < 120; }
+function toBottom() { log.scrollTop = log.scrollHeight; }
+function el(tag, cls, text) {
+  const node = document.createElement(tag);
+  if (cls) node.className = cls;
+  if (text !== undefined) node.textContent = text;
+  return node;
+}
+function addUser(text) {
+  const row = el("div", "msg user");
+  row.appendChild(el("div", "bubble", text));
+  thread.appendChild(row);
+}
+function addBot() {
+  const row = el("div", "msg bot");
+  const view = {row, note: el("div", "note"), body: el("div", "body"), meta: el("div", "meta")};
+  view.note.hidden = true;
+  view.meta.hidden = true;
+  row.append(view.note, view.body, view.meta);
+  thread.appendChild(row);
+  return view;
+}
+function paint(view, text, rest, live) {
+  view.body.innerHTML = md(text);
+  if (rest) {
+    const target = view.body.lastElementChild && view.body.lastElementChild.tagName === "P"
+      ? view.body.lastElementChild : view.body;
+    target.appendChild(el("span", "rest", rest));
+  }
+  view.body.classList.toggle("typing", live);
+}
+function finish(view, message) {
+  view.meta.hidden = false;
+  view.meta.replaceChildren(el("span", "", message.meta || ""));
+  const copy = el("button", "small", "Copy");
+  copy.addEventListener("click", () => copied(copy, message.content));
+  view.meta.appendChild(copy);
+}
+function render() {
+  thread.querySelectorAll(".msg").forEach((n) => n.remove());
+  $("empty").hidden = chat.messages.length > 0;
+  for (const m of chat.messages) {
+    if (m.role === "user") { addUser(m.content); continue; }
+    const view = addBot();
+    view.row.classList.toggle("is-draft", !!m.draft);
+    paint(view, m.content, "", false);
+    finish(view, m);
+  }
+  toBottom();
+}
+
+async function copyText(text) {
+  try { await navigator.clipboard.writeText(text); return true; } catch {}
+  const area = el("textarea");  // plain http: no Clipboard API
+  area.value = text;
+  area.setAttribute("readonly", "");
+  area.style.cssText = "position:fixed;top:0;opacity:0";
+  document.body.appendChild(area);
+  area.select();
+  area.setSelectionRange(0, text.length);
+  let ok = false;
+  try { ok = document.execCommand("copy"); } catch {}
+  area.remove();
+  return ok;
+}
+async function copied(button, text) {
+  const ok = await copyText(text);
+  button.textContent = ok ? "Copied" : "Copy failed";
+  setTimeout(() => (button.textContent = "Copy"), 1500);
+}
+thread.addEventListener("click", (event) => {
+  const button = event.target.closest("[data-copy]");
+  if (button) copied(button, button.closest(".code").querySelector("pre").textContent);
+  const chip = event.target.closest(".chip");
+  if (chip) { draft.value = chip.textContent; draft.focus(); draft.dispatchEvent(new Event("input")); }
+});
+
+/* ---------- sending ---------- */
+const ICON_SEND = "M12 4l7 7-1.4 1.4-4.6-4.6V20h-2V7.8l-4.6 4.6L5 11z";
+const ICON_STOP = "M7 7h10v10H7z";
+function setStreaming(on) {
+  streaming = on;
+  log.setAttribute("aria-busy", on);
+  $("send-icon").setAttribute("d", on ? ICON_STOP : ICON_SEND);
+  send.setAttribute("aria-label", on ? "Stop" : "Send");
+  send.disabled = !on && !draft.value.trim();
+  if (!on) lightMark(0);
+}
+function lightMark(n) {
+  const cells = [...$("mark").children];
+  cells.forEach((c) => c.classList.remove("on"));
+  for (let k = 0; k < n; k++) cells[Math.floor(Math.random() * cells.length)].classList.add("on");
+}
+
+async function submit() {
+  const text = draft.value.trim();
+  if (streaming || !text) return;
+  const conv = chat;  // "New chat" mid-answer must not receive this answer
+  $("empty").hidden = true;
+  conv.messages.push({role: "user", content: text});
+  addUser(text);
+  save();
+  draft.value = "";
+  autosize();
+  hideWarm();
+  setStreaming(true);
+  toBottom();
+  const view = addBot();
+  const st = {mode: "full", draft: "", final: "", refined: false};
+  let start = performance.now(), first = 0, pieces = 0, stopped = false, failure = null;
+  const show = () => {
+    const follow = nearBottom();
+    const text = st.mode === "draft" ? st.draft : st.final;
+    const rest = st.mode === "refine" && st.draft.startsWith(st.final) ? st.draft.slice(st.final.length) : "";
+    view.row.classList.toggle("is-draft", st.mode === "draft");
+    paint(view, text, rest, true);
+    if (follow) toBottom();
+  };
+  const note = (text, bad) => { view.note.hidden = !text; view.note.textContent = text || ""; view.note.classList.toggle("bad", !!bad); };
+  controller = new AbortController();
+  try {
+    const response = await fetch("/chat", {method: "POST", signal: controller.signal,
+                                           body: JSON.stringify({session: chat.session, messages: plain()})});
+    if (!response.ok) throw new Error(`HTTP ${response.status}: ${(await response.text()).slice(0, 200)}`);
+    const reader = response.body.getReader(), decoder = new TextDecoder();
+    let buffer = "";
+    show();
+    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") {
+          st.mode = "draft";
+          note("Another answer is using the GPU, so this is a draft from the experts already in VRAM.");
+        } else if (data.prain === "refine") {
+          st.mode = "refine";
+          st.refined = true;
+          start = performance.now(), first = 0, pieces = 0;
+          note("Refining the draft with all experts…");
+        } else if (data.prain === "error") {
+          st.mode = "draft";
+          st.refined = false;
+          note(`Refining failed (HTTP ${data.status}); this is the draft.`, true);
+        } else {
+          const piece = data.choices?.[0]?.delta?.content;  // usage chunks have no choices
+          if (!piece) continue;
+          if (pieces++ === 0) first = performance.now();
+          if (st.mode === "draft") st.draft += piece; else st.final += piece;
+          lightMark(2);
+        }
+        show();
+      }
+    }
+  } catch (error) {
+    if (error.name === "AbortError") stopped = true; else failure = error;
+  }
+  controller = null;
+  const content = st.mode === "draft" ? st.draft : st.final;
+  if (conv !== chat) { setStreaming(false); return; }
+  if (failure && !content) {
+    view.row.remove();
+    conv.messages.pop();
+    save();
+    const box = addBot();
+    box.note.hidden = false;
+    box.note.classList.add("bad");
+    box.note.textContent = `The model didn't answer (${failure.message}). Your message is back in the box: ` +
+                           "send it again once the monitor shows the model as ready.";
+    draft.value = text;
+    autosize();
+    setStreaming(false);
+    return;
+  }
+  const seconds = (performance.now() - first) / 1000;
+  const parts = [];
+  if (first) parts.push(`${((first - start) / 1000).toFixed(1)} s to first token`);
+  if (pieces > 1) parts.push(`${((pieces - 1) / seconds).toFixed(1)} tokens/s`);
+  if (st.refined && st.draft) {
+    let same = 0;
+    while (same < st.draft.length && st.draft[same] === st.final[same]) same++;
+    parts.push(`refined, kept ${Math.round(100 * same / st.draft.length)}% of the draft`);
+  } else if (st.mode === "draft") parts.push("draft only");
+  if (stopped) parts.push("stopped");
+  if (failure) parts.push("cut off: " + failure.message);
+  const message = {role: "assistant", content, meta: parts.join(" · "), draft: st.mode === "draft"};
+  conv.messages.push(message);
+  save();
+  if (st.refined || !view.note.classList.contains("bad")) note("");
+  paint(view, content, "", false);
+  finish(view, message);
+  setStreaming(false);
+}
+
+$("composer").addEventListener("submit", (event) => {
+  event.preventDefault();
+  if (streaming) controller && controller.abort(); else submit();
+});
+draft.addEventListener("keydown", (event) => {
+  if (event.key === "Enter" && !event.shiftKey && !event.isComposing && !coarse) {
+    event.preventDefault();
+    submit();
+  }
+});
+if (coarse) $("keys").textContent = "Return adds a line";
+
+function autosize() {
+  draft.style.height = "auto";
+  draft.style.height = Math.min(draft.scrollHeight, innerHeight * 0.4) + "px";
+  if (!streaming) send.disabled = !draft.value.trim();
+}
+let warmTimer = 0;
+function hideWarm() { clearTimeout(warmTimer); $("warm").hidden = true; $("keys").hidden = false; }
+draft.addEventListener("input", () => {
+  autosize();
+  if (streaming || !draft.value.trim()) return;
+  fetch("/warm", {method: "POST", body: JSON.stringify({session: chat.session, messages: plain(), draft: draft.value})})
+    .then((response) => {
+      if (response.status !== 202) return;
+      $("warm").hidden = false;
+      $("keys").hidden = true;
+      clearTimeout(warmTimer);
+      warmTimer = setTimeout(hideWarm, 2500);
+    }).catch(() => {});
+});
+
+$("new-chat").addEventListener("click", () => {
+  if (controller) controller.abort();
+  chat = {session: newId(), messages: []};
+  save();
+  render();
+  draft.focus();
+});
+
+/* ---------- monitor ---------- */
+const monitor = $("monitor"), scrim = $("scrim"), pill = $("pill");
+function openMonitor(on) {
+  monitor.classList.toggle("open", on);
+  scrim.classList.toggle("open", on);
+  pill.setAttribute("aria-expanded", on);
+}
+pill.addEventListener("click", () => openMonitor(!monitor.classList.contains("open")));
+$("close").addEventListener("click", () => openMonitor(false));
+scrim.addEventListener("click", () => openMonitor(false));
+addEventListener("keydown", (event) => { if (event.key === "Escape") openMonitor(false); });
+
+const fmt = (n, digits = 0) => n.toLocaleString(undefined, {maximumFractionDigits: digits, minimumFractionDigits: digits});
+const GIB = 2 ** 30;
+const history = [];  // polls of the last minute
+const speeds = [];
+let lastSpeed = 0, scanState = null;
+
+function set(id, text) { $(id).textContent = text; }
+function meter(id, share) { $(id).style.width = `${Math.max(0, Math.min(1, share)) * 100}%`; }
+function setState(kind, text) {
+  for (const id of ["dot", "pill-dot"]) $(id).className = "dot " + kind;
+  set("state", text);
+}
+
+const SCAN = {off: [44, 37, 56], vram: [106, 53, 160], warm: [255, 107, 26], hot: [255, 212, 71], copy: [69, 211, 199]};
+const mix = (a, b, t) => [0, 1, 2].map((k) => a[k] + (b[k] - a[k]) * t);
+function drawScan(experts) {
+  const canvas = $("scan");
+  const layers = Object.keys(experts.layers).map(Number).sort((a, b) => a - b).map((k) => experts.layers[k]);
+  const rows = layers.length, cols = layers[0].uses.length, n = rows * cols;
+  if (!scanState || scanState.heat.length !== n) {
+    canvas.width = cols;
+    canvas.height = rows;
+    scanState = {heat: new Float32Array(n), flash: new Float32Array(n), uses: null, vram: null};
+  }
+  const uses = new Float64Array(n), vram = new Uint8Array(n);
+  layers.forEach((layer, l) => {
+    layer.uses.forEach((u, e) => (uses[l * cols + e] = u));
+    (layer.vram || []).forEach((e) => (vram[l * cols + e] = 1));
+  });
+  const {heat, flash} = scanState, ctx = canvas.getContext("2d"), image = ctx.createImageData(cols, rows);
+  let resident = 0, active = 0;
+  for (let k = 0; k < n; k++) {
+    const fired = scanState.uses ? uses[k] - scanState.uses[k] : 0;
+    heat[k] = Math.min(1, heat[k] * 0.5 + (fired > 0 ? 0.55 + 0.1 * Math.min(fired, 4) : 0));
+    flash[k] = scanState.vram && vram[k] && !scanState.vram[k] ? 1 : flash[k] * 0.4;
+    resident += vram[k];
+    active += fired > 0;
+    let c = vram[k] ? SCAN.vram : SCAN.off;
+    if (heat[k] > 0.02) c = mix(c, mix(SCAN.warm, SCAN.hot, heat[k]), Math.min(1, heat[k] * 1.5));
+    if (flash[k] > 0.02) c = mix(c, SCAN.copy, flash[k] * 0.85);
+    image.data.set([c[0], c[1], c[2], 255], k * 4);
+  }
+  ctx.putImageData(image, 0, 0);
+  scanState.uses = uses;
+  scanState.vram = vram;
+  set("scan-count", `${fmt(resident)} of ${fmt(n)} in VRAM`);
+  set("axis-end", `layer ${rows} · expert ${cols}`);
+  canvas.setAttribute("aria-label", `${resident} of ${n} experts in VRAM; ${active} used since the last update`);
+  set("empty-live", `${fmt(resident)} of ${fmt(n)} experts are in VRAM right now.`);
+}
+
+function totals(experts) {
+  const t = {hits: 0, misses: 0, cpu: 0};
+  for (const layer of Object.values(experts.layers)) for (const k in t) t[k] += layer[k] || 0;
+  return t;
+}
+
+async function poll() {
+  let s;
+  try {
+    const response = await fetch("/status", {cache: "no-store"});
+    s = await response.json();
+  } catch {
+    setState("bad", "Can't reach prain web");
+    set("pill-text", "Offline");
+    return;
+  }
+  const name = s.persona ? `${s.model} · ${s.persona}` : s.model;
+  set("model", name);
+  set("m-model", name);
+  const v = s.vllm, now = s.time;
+  const entry = {t: now, gen: v ? v.generation_tokens : null, ...(s.experts ? totals(s.experts) : {})};
+  history.push(entry);
+  while (history.length > 2 && now - history[0].t > 60) history.shift();
+
+  if (!v) {
+    setState("bad", "Model offline: starting or restarting");
+    set("pill-text", "Offline");
+  } else {
+    const busy = v.running > 0;
+    const prev = history[history.length - 2];
+    const dgen = prev && prev.gen != null ? entry.gen - prev.gen : 0;  // < 0 after a restart
+    const rate = dgen > 0 ? dgen / Math.max(entry.t - prev.t, 0.1) : 0;
+    if (busy && rate) lastSpeed = lastSpeed ? lastSpeed * 0.4 + rate * 0.6 : rate;
+    speeds.push(busy ? rate : 0);
+    if (speeds.length > 60) speeds.shift();
+    $("speed-box").classList.toggle("idle", !busy || !rate);
+    set("speed", lastSpeed ? fmt(lastSpeed, 1) : "–");
+    set("speed-unit", busy && rate ? "tokens/s" : lastSpeed ? "tokens/s, last answer" : "tokens/s");
+    const top = Math.max(...speeds, 1);
+    $("spark").setAttribute("d", speeds.map((y, x) => `${x ? "L" : "M"}${x + 60 - speeds.length} ${27 - (y / top) * 25}`).join(""));
+    const waiting = v.waiting ? `, ${fmt(v.waiting)} waiting` : "";
+    if (busy) setState("live", `${rate ? "Generating" : "Reading the prompt"}${waiting}`);
+    else setState("ok", "Ready");
+    set("pill-text", !busy ? "Ready" : rate ? `${fmt(lastSpeed, 1)} tok/s` : "Working");
+    set("reqs", `${fmt(v.running)} running${waiting}`);
+    set("kv", `${fmt(v.kv_usage * 100)}%`);
+    meter("kv-bar", v.kv_usage);
+    set("prefix", v.prefix_queries ? `${fmt(100 * v.prefix_hits / v.prefix_queries)}%` : "–");
+  }
+
+  if (s.experts && Object.keys(s.experts.layers).length) {
+    $("scan-empty").hidden = true;
+    $("scan-body").hidden = false;
+    drawScan(s.experts);
+    const old = history.find((h) => h.hits !== undefined);
+    const d = old ? {hits: entry.hits - old.hits, misses: entry.misses - old.misses, cpu: entry.cpu - old.cpu} : null;
+    const routed = d && d.hits >= 0 && d.misses >= 0 ? d.hits + d.misses + d.cpu : 0;  // < 0 after a restart
+    if (routed > 0) {
+      set("hits", `${fmt(100 * d.hits / routed)}%`);
+      set("copies", `${fmt(d.misses / Math.max(entry.t - old.t, 1), 0)}/s`);
+    } else {
+      const all = entry.hits + entry.misses + entry.cpu;
+      set("hits", all ? `${fmt(100 * entry.hits / all)}% overall` : "–");
+      set("copies", `${fmt(entry.misses)} so far`);
+    }
+  } else {
+    $("scan-empty").hidden = false;
+    $("scan-body").hidden = true;
+    set("empty-live", "");
+  }
+
+  const g = s.gpu;
+  if (g) {
+    set("vram", `${fmt(g.vram_used / GIB, 1)} / ${fmt(g.vram_total / GIB, 1)} GiB`);
+    meter("vram-bar", g.vram_used / g.vram_total);
+    set("load", g.busy == null ? "–" : `${g.busy}%`);
+    set("temp", g.temp == null ? "–" : `${fmt(g.temp)} °C`);
+    set("power", g.power == null ? "–" : `${fmt(g.power)} W`);
+    set("fan", g.fan == null ? "–" : `${fmt(g.fan)} rpm`);
+  }
+}
+
+let timer = 0;
+function schedule() {
+  clearTimeout(timer);
+  if (document.visibilityState === "visible") timer = setTimeout(() => poll().finally(schedule), 1500);
+}
+document.addEventListener("visibilitychange", () => { if (document.visibilityState === "visible") poll().finally(schedule); });
+
+render();
+autosize();
+poll().finally(schedule);
+if (!coarse) draft.focus();
+</script>

prain/web.py

@@ -0,0 +1,253 @@
+"""prain web: a browser chat that warms experts up while you type, with a model monitor.
+
+A small stdlib proxy in front of vLLM. GET / serves the page (web.html). GET /status returns
+what the monitor shows: vLLM's counters (/metrics), the GPU (amdgpu sysfs) and prain's expert
+stats (the placement file's stats_path, with the experts in VRAM). POST /warm takes
+{"session", "messages", "draft"} and returns at once: 202 when a warm-up starts, 204 when not.
+POST /chat takes {"session", "messages"} and streams server-sent events: vLLM's chunks, plus
+{"prain": "draft" | "refine" | "error"} events of its own. With --persona, the proxy puts the
+persona's text in front of both as the system message, so the page never sees it. Binds to
+127.0.0.1 unless --host says otherwise (0.0.0.0 serves every interface, without a login).
+
+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 glob
+import json
+import os
+import threading
+import time
+import urllib.error
+import urllib.request
+from collections import OrderedDict, deque
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from pathlib import Path
+
+from prain.chat import (Warmer, chat_body, draft_body, first_model, post, send_request,
+                        warm_body, with_system)
+from prain.config import Draft, Serve
+
+MAX_SESSIONS = 64
+
+PAGE = (Path(__file__).parent / "web.html").read_bytes()
+METRICS = {"vllm:num_requests_running": "running", "vllm:num_requests_waiting": "waiting",
+           "vllm:kv_cache_usage_perc": "kv_usage", "vllm:prefix_cache_queries_total":
+           "prefix_queries", "vllm:prefix_cache_hits_total": "prefix_hits",
+           "vllm:prompt_tokens_total": "prompt_tokens",
+           "vllm:generation_tokens_total": "generation_tokens"}
+
+
+def vllm_metrics(base_url):
+    """The METRICS of vLLM's Prometheus text, summed over label sets; None when unreachable."""
+    try:
+        with urllib.request.urlopen(f"{base_url}/metrics", timeout=2) as response:
+            text = response.read().decode()
+    except (OSError, ValueError):
+        return None
+    out = dict.fromkeys(METRICS.values(), 0.0)
+    for line in text.splitlines():
+        key = METRICS.get(line.split("{", 1)[0].split(" ", 1)[0])
+        if key and not line.startswith("#"):
+            out[key] += float(line.rsplit(None, 1)[1])
+    return out
+
+
+def _number(path):
+    try:
+        with open(path) as f:
+            return int(f.read().strip())
+    except (OSError, ValueError):
+        return None
+
+
+def gpu_stats(sysfs="/sys/class/drm"):
+    """VRAM, load, temperature, power and fan of the GPU with the most VRAM (amdgpu sysfs);
+    None without one."""
+    cards = [(_number(os.path.join(d, "mem_info_vram_total")), d)
+             for d in glob.glob(os.path.join(sysfs, "card*", "device"))]
+    cards = sorted((total, d) for total, d in cards if total)
+    if not cards:
+        return None
+    total, card = cards[-1]
+    hwmon = sorted(glob.glob(os.path.join(card, "hwmon", "hwmon*")))
+
+    def sensor(name, scale):
+        value = _number(os.path.join(hwmon[0], name)) if hwmon else None
+        return None if value is None else value / scale
+
+    return {"vram_used": _number(os.path.join(card, "mem_info_vram_used")), "vram_total": total,
+            "busy": _number(os.path.join(card, "gpu_busy_percent")),
+            "temp": sensor("temp1_input", 1000), "power": sensor("power1_average", 1e6),
+            "fan": sensor("fan1_input", 1)}
+
+
+def expert_stats(path):
+    """prain's stats file (prain.policy.write_stats), or None."""
+    if not path:
+        return None
+    try:
+        with open(path) as f:
+            return json.load(f)
+    except (OSError, ValueError):
+        return None
+
+
+
+def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, serve=Serve(),
+                draft=Draft(), stats_path=None, sysfs="/sys/class/drm", system=None,
+                persona=None):
+    """system: the persona text, sent as the system message ahead of every warm-up and answer;
+    persona: its name for the page."""
+    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):
+            if self.path == "/status":
+                with lock:
+                    busy = streams[0]
+                data = {"model": model, "persona": persona, "vllm": vllm_metrics(base_url),
+                        "gpu": gpu_stats(sysfs), "experts": expert_stats(stats_path),
+                        "streams": busy, "time": time.time()}
+                return self.reply(200, "application/json", json.dumps(data).encode())
+            self.reply(200, "text/html; charset=utf-8", PAGE)
+
+        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((with_system(system, body["messages"]), body["draft"]))
+                    return self.reply(202, "text/plain", b"")
+                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 = with_system(system, 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(), host="127.0.0.1",
+          stats_path=None, system=None, persona=None):
+    server = make_server(base_url, port, host=host, warmup=warmup, priority=priority, serve=serve,
+                         draft=draft, stats_path=stats_path, system=system, persona=persona)
+    print(f"prain web: http://{host}:{server.server_port}  (vLLM at {base_url}; "
+          f"{draft.full_requests} full of {serve.max_requests} parallel answers"
+          f"{f'; persona {persona}' if persona else ''})")
+    server.serve_forever()

pyproject.toml

@@ -17,3 +17,6 @@ build-backend = "setuptools.build_meta"
 
 [tool.setuptools]
 packages = ["prain"]
+
+[tool.setuptools.package-data]
+prain = ["web.html"]

scripts/bench.py

@@ -0,0 +1,62 @@
+"""Time to first token and decode speed of an OpenAI-compatible server, as JSON lines.
+
+usage (prain installed): python3 scripts/bench.py http://127.0.0.1:8000 [--tokens 128] [--runs 3]
+Each run sends a fresh prompt (a random prefix defeats the prefix cache) and decodes exactly
+--tokens tokens (ignore_eos), so numbers compare across runs and configs. The last line holds
+the medians. --context FILE... puts the files' text (up to --context-chars) before every
+prompt and the random tag after it: run 1 prefills the whole context, later runs find it in
+the prefix cache, like an agent's system prompt and tools.
+"""
+
+import argparse
+import json
+import statistics
+import time
+import uuid
+
+from prain.chat import chat_body, first_model, post, stream_text
+
+PROMPTS = [
+    "Explain how a CPU cache works, step by step, for a beginner.",
+    "Write a short story about a lighthouse keeper who finds a message in a bottle.",
+    "Schreibe eine kurze Anleitung, wie man Brot backt.",
+]
+
+
+def run(base_url, model, prompt, tokens, context=""):
+    text = f"{context}[{uuid.uuid4().hex[:8]}] {prompt}"
+    body = chat_body(model, [{"role": "user", "content": text}], max_tokens=tokens,
+                     temperature=0, ignore_eos=True)
+    start = time.monotonic()
+    stamps = []
+    with post(base_url, "/v1/chat/completions", body, timeout=3600) as response:
+        for _ in stream_text(response):
+            stamps.append(time.monotonic())
+    return {"ttft_s": round(stamps[0] - start, 3),
+            "decode_tok_s": round((tokens - 1) / max(stamps[-1] - stamps[0], 1e-9), 3)}
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("base_url")
+    parser.add_argument("--tokens", type=int, default=128)
+    parser.add_argument("--runs", type=int, default=3)
+    parser.add_argument("--context", nargs="*", default=[], help="files to put before each prompt")
+    parser.add_argument("--context-chars", type=int, default=48000)
+    args = parser.parse_args()
+    model = first_model(args.base_url)
+    context = ""
+    for path in args.context:
+        with open(path) as f:
+            context += f.read()
+    context = context[: args.context_chars] + "\n\n" if context else ""
+    results = []
+    for i in range(args.runs):
+        results.append(run(args.base_url, model, PROMPTS[i % len(PROMPTS)], args.tokens, context))
+        print(json.dumps(results[-1]), flush=True)
+    print(json.dumps({key: statistics.median(r[key] for r in results) for key in results[0]}
+                     | {"median_of": len(results), "tokens": args.tokens}))
+
+
+if __name__ == "__main__":
+    main()

scripts/build.sh

@@ -0,0 +1,137 @@
+#!/bin/bash
+# Build prain's stack on the GPU host, with clang, tuned for this machine:
+#   wrappers  clang/clang++ wrappers in build/bin (libstdc++ + libgcc, lld)
+#   sdk       venv + AMD's ROCm SDK wheels for $PRAIN_GPU_ARCH
+#   fetch     clone PyTorch and vLLM into build/
+#   torch     PyTorch from source -> build/wheels, installed into the venv
+#   torchvision  torchvision from source -> build/wheels, installed into the venv
+#   vllm      vLLM from source (adds $PRAIN_GPU_ARCH to its HIP arch list), installed editable
+#   prain     this package, installed editable, and its test suite
+# usage: scripts/build.sh STEP...   Versions: PRAIN_ROCM_VERSION, PRAIN_TORCH_REF, PRAIN_VISION_REF, PRAIN_VLLM_REF
+set -euo pipefail
+cd "$(dirname "$0")/.."
+ROCM_INDEX=${PRAIN_ROCM_INDEX:-https://stable.repo.amd.com/rocm/whl-next/}
+ROCM_VERSION=${PRAIN_ROCM_VERSION:-10.1.0}
+TORCH_REF=${PRAIN_TORCH_REF:-v2.13.0}
+VISION_REF=${PRAIN_VISION_REF:-v0.28.0}
+VLLM_REF=${PRAIN_VLLM_REF:-v0.32.0}
+PYTHON=${PRAIN_PYTHON:-python3.13}
+
+wrappers() {
+    mkdir -p build/bin
+    local gccdir
+    gccdir=$(gcc -print-search-dirs | sed -n 's/^install: //p')
+    for c in clang clang++; do
+        cat > build/bin/$c <<EOF
+#!/bin/sh
+case "\$*" in
+    # clang 23 spins for hours on PyTorch's baseline (non-AVX) linear algebra kernel with
+    # -march=native. This CPU runs the AVX2 variant, so the baseline one is built with -O1.
+    *LinearAlgebraKernel.cpp.DEFAULT.cpp*) set -- "\$@" -O1 ;;
+esac
+exec $(command -v $c) --no-default-config --gcc-install-dir=${gccdir%/} -stdlib=libstdc++ -rtlib=libgcc -unwindlib=libgcc -fuse-ld=lld -Wno-unused-command-line-argument "\$@"
+EOF
+        chmod +x build/bin/$c
+    done
+    printf '#include <iostream>\nint main() { std::cout << "ok" << std::endl; }\n' > build/bin/probe.cpp
+    build/bin/clang++ -march=native -O2 build/bin/probe.cpp -o build/bin/probe
+    build/bin/probe
+    ldd build/bin/probe | grep -E "c\+\+|gcc|unwind"
+}
+
+sdk() {
+    [ -x .venv/bin/python ] || $PYTHON -m venv .venv
+    .venv/bin/python -m pip install -U pip wheel 'setuptools<80' 'cmake<4' ninja
+    .venv/bin/python -m pip install --index-url "$ROCM_INDEX" \
+        "rocm[libraries,devel]==$ROCM_VERSION" "rocm-sdk-device-$PRAIN_GPU_ARCH==$ROCM_VERSION"
+    .venv/bin/rocm-sdk init
+    .venv/bin/rocm-sdk version
+    # libstdc++ 16's <format> has [[__gnu__::__noinline__]], and HIP's host_defines.h empties
+    # __noinline__ for host compilers: parse <format> first, like Gentoo's hip ebuild does
+    # (https://github.com/ROCm/rocm-systems/issues/1874).
+    python - "$(echo .venv/lib/python3*/site-packages)/_rocm_sdk_devel/include/hip/amd_detail/host_defines.h" <<'EOF'
+import sys
+path = sys.argv[1]
+text = open(path).read()
+fix = ("#if defined(__cplusplus) && defined(__has_include)\n#  if __has_include(<__config>)\n"
+       "#    include <__config>\n#  else\n#    include <format>\n#  endif\n#endif\n")
+if fix not in text:
+    assert text.count("\n#define __noinline__\n") == 1, path
+    open(path, "w").write(text.replace("\n#define __noinline__\n", "\n" + fix + "#define __noinline__\n"))
+print("host_defines.h: <format> before __noinline__")
+EOF
+}
+
+fetch() {
+    [ -d build/pytorch/.git ] || git clone --depth 1 --branch "$TORCH_REF" --recurse-submodules \
+        --shallow-submodules -j 8 https://github.com/pytorch/pytorch build/pytorch
+    [ -d build/vllm/.git ] || git clone --depth 1 --branch "$VLLM_REF" \
+        https://github.com/vllm-project/vllm build/vllm
+}
+
+# PyTorch preloads these SDK libraries on import (torch/_rocm_init.py), so no LD_LIBRARY_PATH.
+TORCH_PRELOADS='"amd_comgr", "amdhip64", "hiprtc", "roctx64", "hipblas", "hipblaslt", "hipfft",
+    "hiprand", "hipsparse", "hipsparselt", "hipsolver", "miopen"'
+
+torch() {
+    local src=build/pytorch
+    python -m pip install -r $src/requirements-build.txt filelock 'fsspec>=0.8.5' jinja2 \
+        'networkx>=2.5.1' 'optree>=0.13.0' psutil 'sympy>=1.13.3'
+    (cd $src && python tools/amd_build/build_amd.py > ../logs/hipify.log)
+    printf 'def initialize():\n    import rocm_sdk\n    rocm_sdk.initialize_process(\n        preload_shortnames=[%s],\n        check_version="%s")\n' \
+        "$TORCH_PRELOADS" "$ROCM_VERSION" > $src/torch/_rocm_init.py
+    # Only what prain needs on one RDNA2 GPU: no CK/aotriton attention (not built for gfx103x),
+    # no NCCL (vLLM falls back to gloo on one GPU), no profiler, no mobile/quantized CPU backends.
+    (cd $src && env USE_ROCM=1 USE_CUDA=0 USE_XPU=0 PYTORCH_ROCM_ARCH="$PRAIN_GPU_ARCH" \
+        USE_ROCM_CK_GEMM=0 USE_ROCM_CK_SDPA=0 USE_FLASH_ATTENTION=0 USE_MEM_EFF_ATTENTION=0 \
+        USE_MSLK=0 USE_NCCL=0 USE_MPI=0 USE_KINETO=0 USE_FBGEMM=0 USE_NNPACK=0 \
+        USE_PYTORCH_QNNPACK=0 USE_XNNPACK=0 USE_ITT=0 USE_MAGMA=0 BUILD_TEST=0 BLAS=OpenBLAS \
+        MAX_JOBS="$PRAIN_JOBS" nice -n 19 python -m pip wheel --no-build-isolation --no-deps -v \
+        -w "$PRAIN_ROOT/build/wheels" .)
+    python -m pip install --no-deps --force-reinstall build/wheels/torch-*.whl
+    env -u LD_LIBRARY_PATH python -c 'import torch; x = torch.ones(4, device="cuda"); print(torch.__version__, torch.version.hip, torch.cuda.get_device_properties(0).gcnArchName, (x * 2).sum().item(), torch.backends.openmp.is_available(), torch.get_num_threads())'
+}
+
+# transformers' Gemma 4 processor imports torchvision, even for text-only serving.
+torchvision() {
+    [ -d build/vision/.git ] || git clone --depth 1 --branch "$VISION_REF" \
+        https://github.com/pytorch/vision build/vision
+    (cd build/vision && env FORCE_CUDA=1 PYTORCH_ROCM_ARCH="$PRAIN_GPU_ARCH" MAX_JOBS="$PRAIN_JOBS" \
+        nice -n 19 python -m pip wheel --no-build-isolation --no-deps -v -w "$PRAIN_ROOT/build/wheels" .)
+    python -m pip install --no-deps --force-reinstall build/wheels/torchvision-*.whl
+    env -u LD_LIBRARY_PATH python -c 'import torch, torchvision; print(torchvision.__version__, torch.ops.torchvision.nms(torch.rand(4, 4, device="cuda"), torch.rand(4, device="cuda"), 0.5).device)'
+}
+
+vllm() {
+    local src=build/vllm site p
+    site=$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')
+    # prain's fixes to vLLM (patches/vllm-*.patch), each applied once.
+    for p in patches/vllm-*.patch; do
+        git -C $src apply --reverse --check "$PRAIN_ROOT/$p" 2>/dev/null || git -C $src apply "$PRAIN_ROOT/$p"
+    done
+    # vLLM lists gfx1030 but not its siblings gfx1031/gfx1032: add ours.
+    grep -q "$PRAIN_GPU_ARCH" $src/CMakeLists.txt ||
+        sed -i "s/;gfx1030;/;gfx1030;$PRAIN_GPU_ARCH;/" $src/CMakeLists.txt
+    # Keep our torch and AMD's Triton for this ROCm release: pip fails rather than swap them.
+    printf 'torch==%s\ntriton==3.8.0+git669b31ac.rocm%s\n' \
+        "$(python -c 'import torch; print(torch.__version__)')" "$ROCM_VERSION" > build/constraints.txt
+    python -m pip install -c build/constraints.txt --extra-index-url "$ROCM_INDEX" triton \
+        -r $src/requirements/common.txt 'setuptools-scm>=9.2.0' 'setuptools-rust>=1.9.0' 'jinja2>=3.1.6'
+    # The SDK's own amdsmi package (vLLM's ROCm platform imports it), not PyPI's.
+    echo "$site/_rocm_sdk_core/share/amd_smi" > "$site/rocm-sdk-amdsmi.pth"
+    (cd $src && env VLLM_TARGET_DEVICE=rocm PYTORCH_ROCM_ARCH="$PRAIN_GPU_ARCH" \
+        MAX_JOBS="$PRAIN_JOBS" nice -n 19 python -m pip install --no-build-isolation --no-deps -v -e .)
+    env -u LD_LIBRARY_PATH python -c 'import vllm, triton, amdsmi; print(vllm.__version__, triton.__version__)'
+}
+
+prain() {
+    python -m pip install --no-deps -e .
+    python -m unittest discover -s tests -t .
+}
+
+for step in "$@"; do
+    . scripts/env.sh
+    mkdir -p build/logs
+    echo "== $step $(date -Is)"
+    "$step" 2>&1 | tee "build/logs/$step.log"
+done

scripts/draft_effect.py

@@ -0,0 +1,54 @@
+"""A second answer that arrives while a first one streams, through `prain web`: when does
+its first token come, and how much of its draft survives the refine?
+
+Run it against a proxy with drafts and one without ([draft] full_requests = max_requests)
+and compare. usage: python3 scripts/draft_effect.py http://127.0.0.1:8080
+"""
+
+import json
+import os
+import sys
+import threading
+import time
+
+from prain.chat import post
+
+LONG = [{"role": "user", "content": "Write a detailed 400-word history of the printing press."}]
+SHORT = [{"role": "user", "content": "Name three uses of copper, one line each."}]
+
+
+def answer(url, session, messages, out):
+    start, mode = time.monotonic(), "full"
+    out.update(first=None, draft="", final="")
+    with post(url, "/chat", {"session": session, "messages": messages}) as response:
+        for raw in response:
+            line = raw.decode().strip()
+            if not line.startswith("data: ") or line == "data: [DONE]":
+                continue
+            data = json.loads(line[6:])
+            if "prain" in data:
+                mode = data["prain"]
+                continue
+            piece = data["choices"][0]["delta"].get("content")
+            if piece:
+                out["first"] = out["first"] or time.monotonic() - start
+                out["draft" if mode == "draft" else "final"] += piece
+    out["total"] = time.monotonic() - start
+
+
+def main():
+    url = sys.argv[1]
+    first, second = {}, {}
+    thread = threading.Thread(target=answer, args=(url, "first", LONG, first))
+    thread.start()
+    time.sleep(5)  # the first answer is decoding now
+    answer(url, "second", SHORT, second)
+    thread.join()
+    kept = len(os.path.commonprefix([second["draft"], second["final"]]))
+    print(f"second answer: first token {second['first']:.1f} s, done after {second['total']:.1f} s; "
+          + (f"draft {len(second['draft'])} chars, {kept} kept by the refine"
+             if second["draft"] else "no draft"))
+
+
+if __name__ == "__main__":
+    main()

scripts/env.sh

@@ -0,0 +1,29 @@
+# prain environment: the project venv with AMD's ROCm SDK wheels (TheRock) and a clang host
+# toolchain tuned for this CPU. Source it from bash:   . scripts/env.sh
+# Overrides: PRAIN_GPU_ARCH (default gfx1032), PRAIN_JOBS (parallel build jobs, default 16).
+PRAIN_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)
+export PRAIN_ROOT
+export PRAIN_GPU_ARCH=${PRAIN_GPU_ARCH:-gfx1032}
+export PRAIN_JOBS=${PRAIN_JOBS:-16}
+export VIRTUAL_ENV=$PRAIN_ROOT/.venv
+case ":$PATH:" in *":$VIRTUAL_ENV/bin:"*) ;; *) PATH=$VIRTUAL_ENV/bin:$PATH ;; esac
+_site=$(echo "$VIRTUAL_ENV"/lib/python3*/site-packages)
+if [ -d "$_site/_rocm_sdk_devel" ]; then
+    _dev=$_site/_rocm_sdk_devel
+    _core=$_site/_rocm_sdk_core
+    export ROCM_PATH=$_dev ROCM_HOME=$_dev HIP_PATH=$_dev CMAKE_PREFIX_PATH=$_dev
+    export HIP_CLANG_PATH=$_dev/lib/llvm/bin HIP_DEVICE_LIB_PATH=$_core/lib/llvm/amdgcn/bitcode
+    case ":${LD_LIBRARY_PATH:-}:" in *":$_dev/lib:"*) ;; *)
+        export LD_LIBRARY_PATH=$_dev/lib:$_dev/lib/rocm_sysdeps/lib:$_core/lib:$_core/lib/rocm_sysdeps/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH} ;;
+    esac
+    case ":$PATH:" in *":$_dev/bin:"*) ;; *) PATH=$_dev/bin:$PATH ;; esac
+fi
+export PATH
+# Host C/C++: the system clang through wrappers (scripts/build.sh wrappers) that drop Gentoo's
+# libc++/compiler-rt defaults: everything links libstdc++/libgcc like the ROCm libraries.
+export CC=$PRAIN_ROOT/build/bin/clang CXX=$PRAIN_ROOT/build/bin/clang++
+export CFLAGS="-march=native" CXXFLAGS="-march=native"
+# Fused-MoE kernel configs tuned for this GPU by scripts/tune_moe.py (vLLM's defaults when
+# the files are missing).
+export VLLM_TUNED_CONFIG_FOLDER=$PRAIN_ROOT/build/moe-configs
+unset _site _dev _core

scripts/measure.py

@@ -0,0 +1,112 @@
+"""Measure the prain cost model on this machine and print it as JSON for
+`prain sim --cost` and `prain tune --cost`. Needs the GPU, vLLM and prain.
+
+usage: python3 scripts/measure.py PACK_FILE [--quant int4] [--threads N] [--dense-gib 4.45] [--layers 30] > cost.json
+PACK_FILE is any <pack_dir>/layerNNN.w13 file; it is read cold (dropped from the page cache).
+"""
+
+import argparse
+import json
+import os
+import time
+
+import torch
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
+
+from prain.config import QUANTS, expert_bytes as size_of
+from prain.quant import cpu_matmul, moe_quant_config, quantize, to_cpu_int4
+from prain.tiered import cpu_experts
+
+HIDDEN, INTER, SLOTS = 2816, 704, 5  # Gemma 4 26B-A4B expert, 5 slots per layer
+
+
+def ms(fn, repeat=20):
+    for _ in range(3):
+        fn()
+    torch.cuda.synchronize()
+    start = time.perf_counter()
+    for _ in range(repeat):
+        fn()
+    torch.cuda.synchronize()
+    return (time.perf_counter() - start) * 1000 / repeat
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("pack_file")
+    parser.add_argument("--threads", type=int, default=4)
+    parser.add_argument("--dense-gib", type=float, default=4.45, help="non-expert weights")
+    parser.add_argument("--layers", type=int, default=30)
+    parser.add_argument("--quant", choices=QUANTS, default="bf16", help="expert weight format")
+    args = parser.parse_args()
+    torch.set_num_threads(args.threads)
+    expert_bytes = size_of(HIDDEN, INTER, args.quant)
+    host = torch.empty(expert_bytes, dtype=torch.uint8).pin_memory()
+    device = torch.empty_like(host, device="cuda")
+    h2d_gbps = expert_bytes / 1e6 / ms(lambda: device.copy_(host, non_blocking=True))
+
+    fd = os.open(args.pack_file, os.O_RDONLY)
+    size = os.fstat(fd).st_size
+    os.posix_fadvise(fd, 0, size, os.POSIX_FADV_DONTNEED)  # read it cold
+    chunk = 2 * HIDDEN * INTER * 2  # one w13 row
+    start, reads = time.perf_counter(), 0
+    for offset in range(0, min(size, 32 * chunk), chunk):
+        reads += len(os.pread(fd, chunk, offset))
+    disk_gbps = reads / 1e6 / ((time.perf_counter() - start) * 1000)
+    os.close(fd)
+
+    w1 = torch.randn(SLOTS, 2 * INTER, HIDDEN, dtype=torch.bfloat16, device="cuda") * 0.02
+    w2 = torch.randn(SLOTS, HIDDEN, INTER, dtype=torch.bfloat16, device="cuda") * 0.02
+    quant_config = None
+    if args.quant != "bf16":
+        (w1, s1), (w2, s2) = quantize(w1, args.quant), quantize(w2, args.quant)
+        quant_config = moe_quant_config(args.quant, s1, s2)
+
+    def moe(tokens):
+        x = torch.randn(tokens, HIDDEN, dtype=torch.bfloat16, device="cuda")
+        ids = torch.arange(SLOTS, device="cuda", dtype=torch.int32).repeat(tokens, 1)
+        weights = torch.full((tokens, SLOTS), 1 / SLOTS, device="cuda")
+        return lambda: fused_experts(x, w1, w2, weights, ids, activation=MoEActivation.GELU_TANH,
+                                     global_num_experts=SLOTS, quant_config=quant_config)
+
+    with set_current_vllm_config(VllmConfig()):
+        round_ms = ms(moe(1))
+        route_ms = max(0.0, (ms(moe(256)) - round_ms) / (256 * SLOTS))
+    big = torch.empty(256 * 2**20, dtype=torch.uint8, device="cuda")
+    copy = torch.empty_like(big)
+    device_gbps = 2 * big.numel() / 1e6 / ms(lambda: copy.copy_(big))  # read + write
+    other_ms = args.dense_gib * 2**30 / 1e6 / device_gbps / args.layers
+    flag = torch.zeros(1, device="cuda")
+    sync_ms = ms(lambda: flag.add_(1).item())
+    # Serving computes a different expert each time, read from RAM (not from the CPU caches).
+    experts = [(torch.randn(2 * INTER, HIDDEN, dtype=torch.bfloat16),
+                torch.randn(HIDDEN, INTER, dtype=torch.bfloat16)) for _ in range(23)]
+    if args.quant != "bf16":  # the CPU kernels on pack-layout weights, as the plugin runs them
+        def stored(w):
+            q, scale = quantize(w, args.quant)
+            return cpu_matmul(args.quant, to_cpu_int4(q) if args.quant == "int4" else q, scale)
+
+        experts = [(stored(a), stored(b)) for a, b in experts]
+    x = torch.randn(1, HIDDEN, dtype=torch.bfloat16)
+    one = torch.ones(1, 1)
+    ids = torch.zeros(1, 1, dtype=torch.int32)
+
+    def cpu_route(weights):
+        cpu_experts(x, one, ids, {0}, lambda e: weights)
+
+    for weights in experts[:3]:  # warm up like ms(): thread pool and BLAS setup are one-time costs
+        cpu_route(weights)
+    start = time.perf_counter()
+    for weights in experts[3:]:
+        cpu_route(weights)
+    cpu_ms = (time.perf_counter() - start) * 1000 / len(experts[3:])
+    print(json.dumps({"h2d_gbps": round(h2d_gbps, 2), "disk_gbps": round(disk_gbps, 2),
+                      "round_ms": round(round_ms, 4), "route_ms": round(route_ms, 6),
+                      "other_ms": round(other_ms, 4), "sync_ms": round(sync_ms, 4),
+                      "cpu_ms": round(cpu_ms, 4)}, indent=1))
+
+
+if __name__ == "__main__":
+    main()

scripts/quality.py

@@ -0,0 +1,100 @@
+"""Answer quality of a served config against reference answers from another config.
+
+    python scripts/quality.py URL make REF.json    greedy reference answers (run on the exact config)
+    python scripts/quality.py URL score REF.json   how likely this server finds those answers
+
+Both work on token ids: the prompt is the chat template with its generation prompt, and `score`
+re-scores the reference answer on exactly the ids it was generated from (prompt logprobs). A
+finished assistant message renders differently from the generation prompt, so re-scoring
+through the chat API skews the logprobs. `nll` is the mean negative log-likelihood per answer
+token (the reference config scores about its own decode logprobs); `same_tokens` is how many
+leading tokens this server's own greedy answer shares with the reference.
+"""
+
+import argparse
+import json
+import urllib.request
+
+PROMPTS = [
+    "Write a Python function that returns the n-th Fibonacci number iteratively.",
+    "Explain what `git rebase -i HEAD~3` does, in three sentences.",
+    "Write a bash one-liner that counts the lines of all .py files below the current directory.",
+    "On Gentoo, how do I rebuild every package that links against a library I just updated?",
+    "What is the difference between a process and a thread? Answer in two sentences.",
+]
+CODE_PROMPT = "Here is some Python code:\n\n{code}\n\nWhat does this code do? Answer in two sentences."
+
+
+def post(url, path, body):
+    request = urllib.request.Request(url + path, json.dumps(body).encode(),
+                                     {"Content-Type": "application/json"})
+    with urllib.request.urlopen(request, timeout=3600) as response:
+        return json.load(response)
+
+
+def greedy(url, model, ids, tokens):
+    """(text, token ids, their logprobs) of a greedy answer."""
+    choice = post(url, "/v1/completions", {
+        "model": model, "prompt": ids, "max_tokens": tokens, "temperature": 0, "logprobs": 1,
+        "return_tokens_as_token_ids": True})["choices"][0]
+    logprobs = choice["logprobs"]
+    return (choice["text"], [int(t.split(":")[1]) for t in logprobs["tokens"]],
+            logprobs["token_logprobs"])
+
+
+def nll(url, model, ids, answer):
+    """Mean negative log-likelihood of the answer tokens after the prompt ids."""
+    entries = post(url, "/v1/completions", {
+        "model": model, "prompt": ids + answer, "max_tokens": 1, "temperature": 0,
+        "prompt_logprobs": 0})["choices"][0]["prompt_logprobs"]
+    return -sum(entries[len(ids) + i][str(t)]["logprob"] for i, t in enumerate(answer)) / len(answer)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("url")
+    parser.add_argument("action", choices=("make", "score"))
+    parser.add_argument("ref")
+    parser.add_argument("--model", default="gemma4")
+    parser.add_argument("--tokens", type=int, default=48, help="answer tokens")
+    parser.add_argument("--code", default="prain/policy.py", help="file for the code question")
+    parser.add_argument("--code-chars", type=int, default=900)
+    parser.add_argument("--greedy-only", action="store_true",
+                        help="score: skip the nll (prompt logprobs need ~0.5 GiB of free VRAM "
+                             "for the logits); compare greedy answers only")
+    args = parser.parse_args()
+    if args.action == "make":
+        with open(args.code) as f:
+            prompts = PROMPTS + [CODE_PROMPT.format(code=f.read()[: args.code_chars])]
+        refs = []
+        for prompt in prompts:
+            ids = post(args.url, "/tokenize", {"model": args.model, "add_generation_prompt": True,
+                                               "messages": [{"role": "user", "content": prompt}]})
+            ids = ids["tokens"]
+            text, answer, _ = greedy(args.url, args.model, ids, args.tokens)
+            refs.append({"prompt": prompt, "ids": ids, "answer": answer, "text": text})
+            print(json.dumps({"prompt": prompt[:60], "answer": text[:120]}), flush=True)
+        with open(args.ref, "w") as f:
+            json.dump(refs, f)
+        return
+    with open(args.ref) as f:
+        refs = json.load(f)
+    scores, same, own = [], [], []
+    for ref in refs:
+        if not args.greedy_only:
+            scores.append(nll(args.url, args.model, ref["ids"], ref["answer"]))
+        text, answer, logprobs = greedy(args.url, args.model, ref["ids"], len(ref["answer"]))
+        own.append(sum(logprobs) / len(logprobs))
+        same.append(next((i for i, (a, b) in enumerate(zip(answer, ref["answer"])) if a != b),
+                         min(len(answer), len(ref["answer"]))))
+        print(json.dumps({"prompt": ref["prompt"][:40],
+                          "nll": round(scores[-1], 3) if scores else None,
+                          "own_lp": round(own[-1], 3), "same_tokens": same[-1],
+                          "answer": text[:100]}), flush=True)
+    print(json.dumps({"mean_nll": round(sum(scores) / len(scores), 3) if scores else None,
+                      "mean_own_lp": round(sum(own) / len(own), 3), "same_tokens": same,
+                      "identical": sum(s == len(r["answer"]) for s, r in zip(same, refs))}))
+
+
+if __name__ == "__main__":
+    main()

scripts/smoke.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+# Go/no-go: does this vLLM build run Gemma 4 on this GPU? Serves a 2-layer copy of MODEL_DIR with
+# random weights (port 8000, log in build/smoke.log), then sends a greedy request, the warm-up
+# shape (continue_final_message, priority) and the draft fields (vllm_xargs, cache_salt).
+# usage: scripts/smoke.sh MODEL_DIR [vllm serve args...]
+# Expected: READY, a JSON list of 3 (gibberish) strings, two {"id":"chatcmpl-... lines, STOPPED.
+# RDNA2 needs --attention-backend TRITON_ATTN: ROCM_ATTN's prefill kernel wants 160 KiB of LDS.
+cd "$(dirname "$0")/.." && . scripts/env.sh
+model=$1
+shift
+setsid vllm serve "$model" --served-model-name gemma4 --language-model-only --max-model-len 4096 \
+    --max-num-batched-tokens 1024 --max-num-seqs 2 --kv-cache-memory-bytes 805306368 --port 8000 \
+    --enforce-eager --moe-backend triton --scheduling-policy priority --load-format dummy \
+    --hf-overrides '{"text_config": {"num_hidden_layers": 2, "layer_types": ["sliding_attention", "full_attention"]}}' \
+    "$@" > build/smoke.log 2>&1 < /dev/null &
+server=$!
+for _ in $(seq 120); do
+    curl -sf localhost:8000/health > /dev/null && break
+    kill -0 $server 2> /dev/null || break
+    sleep 5
+done
+curl -sf localhost:8000/health > /dev/null && echo READY || { echo DIED; tail -40 build/smoke.log; exit 1; }
+python scripts/greedy.py http://127.0.0.1:8000 gemma4 | head -c 200; echo
+post() { curl -s localhost:8000/v1/chat/completions -H 'Content-Type: application/json' -d "$1" | head -c 60; echo; }
+post '{"model": "gemma4", "messages": [{"role": "user", "content": "Why is the sky"}], "max_tokens": 1, "add_generation_prompt": false, "continue_final_message": true, "priority": 1}'
+post '{"model": "gemma4", "messages": [{"role": "user", "content": "Hi"}], "max_tokens": 2, "vllm_xargs": {"prain": "draft"}, "cache_salt": "prain-draft"}'
+kill -INT -$server
+while curl -sf localhost:8000/health > /dev/null; do sleep 2; done
+echo STOPPED

scripts/tune_moe.py

@@ -0,0 +1,111 @@
+"""Tune vLLM's Triton fused-MoE kernel for this GPU on prain's slot tensors.
+
+usage: python scripts/tune_moe.py OUT_DIR [FORMAT...]   (scripts/env.sh points
+                                    VLLM_TUNED_CONFIG_FOLDER at build/moe-configs)
+Times candidate configs on one prain round (16 slots of 128 routed experts, expert_map;
+up to 32 tokens route into the slots like decode, more tokens spread like prefill),
+Gemma 4 expert shapes, bf16/int8/int4 slots, and writes the fastest per batch size as vLLM
+config files for every slot count E (the best config hardly depends on E). Every config in
+the files compiled, ran and matched the default config's output on this GPU.
+"""
+
+import itertools
+import json
+import os
+import sys
+import time
+
+import torch
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.fused_moe import override_config
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (int4_w4a16_moe_quant_config,
+                                                         int8_w8a16_moe_quant_config)
+from vllm.model_executor.layers.fused_moe.fused_moe import (fused_experts, get_config_file_name,
+                                                            get_default_config)
+
+from prain.quant import GROUP, quantize
+
+H, I, TOPK, EXPERTS, SLOTS = 2816, 704, 8, 128, 16
+SIZES = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048)
+FORMATS = {"bf16": (None, None), "int8": ("int8_w8a16", None), "int4": ("int4_w4a16", [0, GROUP])}
+
+
+def candidates(tokens):
+    # ponytail: ~8 s per Triton compile on this host, so a small space (~70 compiles per
+    # format); widen it when a GPU compiles faster.
+    if tokens <= 32:
+        space = itertools.product((16,), (32, 64, 128), (64, 128, 256), (1,), (2, 4), (2,))
+    else:
+        space = itertools.product((32, 64), (64, 128), (32, 64), (8,), (4, 8), (2,))
+    for bm, bn, bk, group, warps, stages in space:
+        yield {"BLOCK_SIZE_M": bm, "BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk, "GROUP_SIZE_M": group,
+               "SPLIT_K": 1, "num_warps": warps, "num_stages": stages}  # int4 needs SPLIT_K
+
+
+def timed(fn, repeat):
+    fn()
+    torch.cuda.synchronize()
+    start = time.perf_counter()
+    for _ in range(repeat):
+        fn()
+    torch.cuda.synchronize()
+    return (time.perf_counter() - start) * 1000 / repeat
+
+
+def main(out_dir, formats):
+    os.makedirs(out_dir, exist_ok=True)
+    torch.manual_seed(0)
+    w13 = torch.randn(SLOTS, 2 * I, H, device="cuda").mul_(0.02).bfloat16()
+    w2 = torch.randn(SLOTS, H, I, device="cuda").mul_(0.02).bfloat16()
+    table = torch.full((EXPERTS + 1,), -1, dtype=torch.int32, device="cuda")
+    table[:SLOTS] = torch.arange(SLOTS, dtype=torch.int32, device="cuda")
+    for name in formats:
+        dtype, block_shape = FORMATS[name]
+        if name == "bf16":
+            a, b, quant_config = w13, w2, None
+        else:
+            (a, s13), (b, s2) = quantize(w13, name), quantize(w2, name)
+            make = int8_w8a16_moe_quant_config if name == "int8" else int4_w4a16_moe_quant_config
+            quant_config = make(w1_scale=s13, w2_scale=s2, w1_zp=None, w2_zp=None,
+                                block_shape=block_shape)
+        best = {}
+        for tokens in SIZES:
+            x = torch.randn(tokens, H, device="cuda").mul_(0.5).bfloat16()
+            pool = SLOTS if tokens <= 32 else EXPERTS
+            ids = torch.stack([torch.randperm(pool, device="cuda")[:TOPK]
+                               for _ in range(tokens)]).to(torch.int32)
+            weights = torch.softmax(torch.randn(tokens, TOPK, device="cuda"), -1)
+
+            def run():
+                return fused_experts(x, a, b, weights, ids, activation=MoEActivation.GELU_TANH,
+                                     global_num_experts=EXPERTS + 1, expert_map=table,
+                                     quant_config=quant_config)
+
+            default = get_default_config(tokens, SLOTS, I, a.shape[2], TOPK, dtype, block_shape)
+            with override_config(default):  # vLLM's own choice when no config file exists
+                default_ms = timed(run, 10 if tokens <= 32 else 3)
+                reference = run()
+            results = [(default_ms, json.dumps(default, sort_keys=True))]  # never worse than it
+            for config in candidates(tokens):
+                # catch inside: vLLM's override_config skips its restore when the body raises
+                with override_config(config):
+                    try:
+                        ms = timed(run, 10 if tokens <= 32 else 3)
+                        out = run()
+                    except Exception:  # noqa: BLE001 - configs that do not compile or fit
+                        continue
+                if (out.float() - reference.float()).norm() <= 1e-2 * reference.float().norm() + 1e-6:
+                    results.append((ms, json.dumps(config, sort_keys=True)))
+            ms, config = min(results)
+            best[str(tokens)] = json.loads(config)
+            print(f"{name} tokens={tokens:4d}: default {default_ms:8.3f} ms, best {ms:8.3f} ms "
+                  f"({len(results)} configs) {config}", flush=True)
+        for e in range(1, EXPERTS + 1):  # any slot count a placement can give a layer
+            with open(os.path.join(out_dir, get_config_file_name(e, I, dtype, block_shape)), "w") as f:
+                json.dump(best, f, indent=1)
+
+
+if __name__ == "__main__":
+    with set_current_vllm_config(VllmConfig()):
+        main(sys.argv[1], sys.argv[2:] or list(FORMATS))

scripts/warm_effect.py

@@ -0,0 +1,38 @@
+"""Time to first token for a fresh prompt: cold, and warmed up as if it had been typed.
+
+usage (prain installed): python3 scripts/warm_effect.py http://127.0.0.1:8000
+"""
+
+import sys
+import time
+import uuid
+
+from prain.chat import chat_body, first_model, post, stream_text, warm_body
+
+PROMPT = ("You are reviewing a design for a home server that runs a mixture-of-experts language "
+          "model on a single consumer graphics card with eight gigabytes of memory. The model has "
+          "thirty layers with one hundred and twenty-eight experts each, and only a few experts per "
+          "layer fit on the card at once. The rest wait in system memory or on an NVMe drive and "
+          "are copied over PCI Express when a token needs them. List the three biggest risks of "
+          "this design and, for each, one concrete way to measure it.")
+
+
+def first_token_seconds(base_url, model, text):
+    start = time.monotonic()
+    body = chat_body(model, [{"role": "user", "content": text}], max_tokens=1)
+    with post(base_url, "/v1/chat/completions", body) as response:
+        next(stream_text(response), None)
+    return time.monotonic() - start
+
+
+def main(base_url):
+    model = first_model(base_url)
+    cold = first_token_seconds(base_url, model, f"[{uuid.uuid4().hex[:8]}] {PROMPT}")
+    text = f"[{uuid.uuid4().hex[:8]}] {PROMPT}"  # a new tag: no prefix-cache hit from above
+    post(base_url, "/v1/chat/completions", warm_body(model, [], text[:-40])).read()
+    warm = first_token_seconds(base_url, model, text)
+    print(f"first token: cold {cold:.2f} s, warmed {warm:.2f} s")
+
+
+if __name__ == "__main__":
+    main(sys.argv[1])

tests/helpers.py

@@ -41,10 +41,13 @@ class StubVllm:
         self.delay = 0.0
         self.pause = 0.0
         self.fail = None
+        self.metrics = None  # Prometheus text served at GET /metrics
         stub = self
 
         class Handler(BaseHTTPRequestHandler):
             def do_GET(self):
+                if self.path == "/metrics" and stub.metrics is not None:
+                    return self.send(200, "text/plain", stub.metrics)
                 self.send(200, "application/json", {"data": [{"id": "stub-model"}]})
 
             def do_POST(self):

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,113 @@
+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", "--chat-template-content-format", "string",
+            "--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_web_binds_the_host_it_is_given(self):
+        with mock.patch("prain.web.serve") as web:
+            main(["web", "--port", "8081"])
+            self.assertEqual(web.call_args.kwargs["host"], "127.0.0.1")
+            main(["web", "--host", "0.0.0.0"])
+            self.assertEqual(web.call_args.kwargs["host"], "0.0.0.0")
+
+    def test_a_persona_file_reaches_web_and_tui(self):
+        path = self.write("gentoo-admin.md", "You look after a Gentoo box.\n\n")
+        with mock.patch("prain.web.serve") as web:
+            main(["web", "--persona", path])
+            self.assertEqual((web.call_args.kwargs["system"], web.call_args.kwargs["persona"]),
+                             ("You look after a Gentoo box.", "gentoo-admin"))
+            main(["web"])
+            self.assertEqual((web.call_args.kwargs["system"], web.call_args.kwargs["persona"]),
+                             (None, None))
+        with mock.patch("prain.tui.main") as tui:
+            main(["tui", "--persona", path])
+            self.assertEqual(tui.call_args.kwargs["system"], "You look after a Gentoo box.")
+
+    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)
+
+    def test_quantized_experts_check_and_tune(self):
+        hf = self.write("config.json", json.dumps(GEMMA4))
+        out, _ = self.run_cli("check", self.write("q.toml", 'pack_dir = "/packs"\nquant = "int4"\n'), hf)
+        self.assertIn("expert size 3.01 MiB\n", out)
+        trace = os.path.join(self.dir, "trace.jsonl.gz")
+        self.run_cli("synth", trace, "--forwards", "40", "--layers", "3")
+        out, _ = self.run_cli("tune", trace, hf, "--quant", "int4", "--expert-vram-gib", "0.2",
+                              "--ram-gib", "4", "--pack-dir", "/packs", "--top", "3")
+        tuned = load_config(self.write("tuned.toml", out))
+        self.assertEqual(tuned.quant, "int4")
+        slots = sum(tuned.placement(layer, 128).cache_slots for layer in range(3))
+        self.assertGreater(slots, int(0.2 * 1024 / 11.34))  # int4 experts are ~4x smaller
+        self.assertLessEqual(slots, int(0.2 * 1024 / 3.01))

tests/test_config.py

@@ -23,3 +23,154 @@ 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_miss_min_weight(self):
+        config = mc.load_config(self.write("m.toml", 'pack_dir = "/p"\n[miss]\nmin_weight = 0.7\n'))
+        self.assertEqual(config.miss.min_weight, 0.7)
+        self.assertEqual(mc.Miss().min_weight, 1.0)  # default: compute every routed expert
+        with self.assertRaisesRegex(ValueError, "min_weight must be a number in"):
+            mc.load_config(self.write("m2.toml", 'pack_dir = "/p"\n[miss]\nmin_weight = 1.5\n'))
+
+    def test_quant_key_and_expert_sizes(self):
+        self.assertEqual(mc.load_config(self.write("q.toml", 'pack_dir = "/p"\nquant = "int4"\n')).quant, "int4")
+        self.assertEqual(mc.load_config(self.write("b.toml", 'pack_dir = "/p"\n')).quant, "bf16")
+        with self.assertRaisesRegex(ValueError, "quant must be one of"):
+            mc.load_config(self.write("bad.toml", 'pack_dir = "/p"\nquant = "int3"\n'))
+        h, i = 2816, 704
+        self.assertEqual(mc.expert_bytes(h, i), 3 * h * i * 2)
+        self.assertEqual(mc.expert_bytes(h, i, "int8"), 3 * h * i + 4 * (2 * i + h))  # fp32 row scales
+        self.assertEqual(mc.expert_bytes(h, i, "int4"), 3 * h * i // 2 + 2 * 3 * h * i // 64)
+
+    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_gemv.py

@@ -0,0 +1,55 @@
+import unittest
+
+import torch
+import torch.nn.functional as F
+
+from prain import gemv
+
+GPU = torch.cuda.is_available() and gemv.fast_device(0)
+
+
+def rel(a, b):
+    return ((a.float() - b.float()).norm() / b.float().norm()).item()
+
+
+@unittest.skipUnless(GPU, "needs a gfx10 GPU")
+class TestGemv(unittest.TestCase):
+    def test_up_to_four_rows_match_a_float32_linear(self):
+        torch.manual_seed(0)
+        for n, k in ((1000, 2816), (2816, 2112), (130, 70)):
+            w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") * 0.02
+            for rows in (1, 2, 3, 4):
+                x = torch.randn(rows, k, dtype=torch.bfloat16, device="cuda")
+                y = gemv.linear(x, w)
+                self.assertEqual((y.shape, y.dtype), ((rows, n), torch.bfloat16))
+                self.assertLess(rel(y, F.linear(x.float(), w.float())), 1e-2, (n, k, rows))
+        x = torch.randn(2, 1, 70, dtype=torch.float16, device="cuda")  # leading dims, fp16
+        w = torch.randn(130, 70, dtype=torch.float16, device="cuda")
+        self.assertEqual(gemv.linear(x, w).shape, (2, 1, 130))
+        self.assertLess(rel(gemv.linear(x, w), F.linear(x.float(), w.float())), 1e-2)
+
+    def test_more_rows_bias_or_strided_weights_use_f_linear(self):
+        w = torch.randn(64, 96, dtype=torch.bfloat16, device="cuda")
+        bias = torch.randn(64, dtype=torch.bfloat16, device="cuda")
+        x = torch.randn(5, 96, dtype=torch.bfloat16, device="cuda")
+        self.assertTrue(torch.equal(gemv.linear(x, w), F.linear(x, w)))
+        self.assertTrue(torch.equal(gemv.linear(x[:1], w, bias), F.linear(x[:1], w, bias)))
+        strided = torch.randn(96, 64, dtype=torch.bfloat16, device="cuda").t()
+        self.assertTrue(torch.equal(gemv.linear(x[:1], strided), F.linear(x[:1], strided)))
+
+    def test_wrapped_dispatch_sends_only_decode_shapes_to_the_kernel(self):
+        calls = []
+
+        def fallback(layer, x, weight, bias=None):
+            calls.append(x.shape[0])
+            return F.linear(x, weight, bias)
+
+        gemm = gemv.wrap(lambda *args, **kwargs: fallback)("auto")
+        w = torch.randn(64, 96, dtype=torch.bfloat16, device="cuda")
+        gemm(None, torch.randn(1, 96, dtype=torch.bfloat16, device="cuda"), w)
+        gemm(None, torch.randn(9, 96, dtype=torch.bfloat16, device="cuda"), w)
+        self.assertEqual(calls, [9])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_kernel_match.py

@@ -0,0 +1,94 @@
+import importlib.util
+import unittest
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+GPU = torch is not None and torch.cuda.is_available() and importlib.util.find_spec("vllm") is not None
+
+
+@unittest.skipUnless(GPU, "needs vLLM and a GPU")
+class KernelMatch(unittest.TestCase):
+    def test_cpu_experts_compute_what_the_vllm_kernel_computes(self):
+        from vllm.config import VllmConfig, set_current_vllm_config
+        from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+        from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
+
+        from prain.tiered import cpu_experts
+
+        torch.manual_seed(0)
+        tokens, hidden, inter, experts = 8, 256, 128, 4
+        x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
+        w13 = torch.randn(experts, 2 * inter, hidden, dtype=torch.bfloat16) * 0.1
+        w2 = torch.randn(experts, hidden, inter, dtype=torch.bfloat16) * 0.1
+        topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int()
+        topk_weights = torch.rand(tokens, 2)
+        with set_current_vllm_config(VllmConfig()):
+            gpu = fused_experts(x.cuda(), w13.cuda(), w2.cuda(), topk_weights.cuda(),
+                                topk_ids.cuda(), activation=MoEActivation.GELU_TANH,
+                                global_num_experts=experts).float().cpu()
+        cpu = cpu_experts(x, topk_weights, topk_ids, set(range(experts)), lambda e: (w13[e], w2[e]))
+        torch.testing.assert_close(cpu, gpu, rtol=3e-2, atol=3e-2)
+
+    def test_a_round_computes_only_the_experts_its_slot_table_maps(self):
+        from vllm.config import VllmConfig, set_current_vllm_config
+        from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+        from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
+
+        from prain.tiered import cpu_experts, slot_table
+
+        torch.manual_seed(0)
+        tokens, hidden, inter, experts = 64, 256, 128, 6
+        x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
+        w13 = torch.randn(experts, 2 * inter, hidden, dtype=torch.bfloat16) * 0.1
+        w2 = torch.randn(experts, hidden, inter, dtype=torch.bfloat16) * 0.1
+        topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int()
+        topk_ids[0, 1] = experts  # a padding route
+        topk_weights = torch.rand(tokens, 2)
+        table = slot_table({1: 0, 4: 1}, experts, "cuda")
+        with set_current_vllm_config(VllmConfig()):
+            gpu = fused_experts(x.cuda(), w13[[1, 4]].cuda(), w2[[1, 4]].cuda(), topk_weights.cuda(),
+                                topk_ids.cuda(), activation=MoEActivation.GELU_TANH,
+                                global_num_experts=table.numel(), expert_map=table).float().cpu()
+        cpu = cpu_experts(x, topk_weights, topk_ids, {1, 4}, lambda e: (w13[e], w2[e]))
+        torch.testing.assert_close(cpu, gpu, rtol=3e-2, atol=3e-2)
+
+    def test_quantized_slots_compute_what_the_cpu_path_computes(self):
+        """Pack layout -> slot layout on the GPU -> vLLM's int8/int4 kernel for a round that
+        holds experts 1 and 3, against the CPU kernels on the pack layout and against the
+        dequantized weights."""
+        from vllm.config import VllmConfig, set_current_vllm_config
+        from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+        from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
+
+        from prain.quant import (cpu_matmul, dequantize, moe_quant_config, quantize, to_cpu_int4,
+                                 to_gpu_int4)
+        from prain.tiered import cpu_experts, slot_table
+
+        torch.manual_seed(0)
+        tokens, hidden, inter, experts = 8, 256, 128, 4
+        x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
+        w13 = torch.randn(experts, 2 * inter, hidden) * 0.1
+        w2 = torch.randn(experts, hidden, inter) * 0.1
+        topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int()
+        topk_weights = torch.rand(tokens, 2)
+        table = slot_table({1: 0, 3: 1}, experts, "cuda")
+        for quant in ("int8", "int4"):
+            parts = [quantize(w, quant) for w in (w13, w2)]  # [q, scale], gpu layout
+            pack = [[to_cpu_int4(q) if quant == "int4" else q, s] for q, s in parts]
+            slots = [[to_gpu_int4(q[[1, 3]].cuda()) if quant == "int4" else q[[1, 3]].cuda(),
+                      s[[1, 3]].cuda()] for q, s in pack]
+            with set_current_vllm_config(VllmConfig()):
+                gpu = fused_experts(x.cuda(), slots[0][0], slots[1][0], topk_weights.cuda(),
+                                    topk_ids.cuda(), activation=MoEActivation.GELU_TANH,
+                                    global_num_experts=table.numel(), expert_map=table,
+                                    quant_config=moe_quant_config(quant, slots[0][1], slots[1][1]))
+            cpu = cpu_experts(x, topk_weights, topk_ids, {1, 3}, lambda e: tuple(
+                cpu_matmul(quant, q[e], s[e]) for q, s in pack))
+            exact = cpu_experts(x, topk_weights, topk_ids, {1, 3}, lambda e: tuple(
+                dequantize(q, s, quant)[e] for q, s in parts))
+            self.assertGreater(exact.norm().item(), 1.0)
+            torch.testing.assert_close(cpu, exact, rtol=3e-2, atol=3e-2, msg=quant)
+            torch.testing.assert_close(gpu.float().cpu(), exact, rtol=3e-2, atol=3e-2, msg=quant)

tests/test_kernels.py

@@ -0,0 +1,111 @@
+import importlib.util
+import unittest
+
+try:
+    import torch
+except ImportError:
+    torch = None
+
+GPU = torch is not None and torch.cuda.is_available() and importlib.util.find_spec("vllm") is not None
+
+
+@unittest.skipUnless(GPU, "needs vLLM and a GPU")
+class Kernels(unittest.TestCase):
+    def setUp(self):
+        torch.manual_seed(0)
+
+    def test_unpack_writes_what_to_gpu_int4_computes(self):
+        from prain.kernels import unpack_int4
+        from prain.quant import to_gpu_int4
+
+        for n, half in ((256, 128), (1408, 1408), (2816, 352)):
+            cpu_layout = torch.randint(0, 256, (n, half), dtype=torch.uint8, device="cuda")
+            dst = torch.zeros_like(cpu_layout)
+            unpack_int4(cpu_layout.flatten(), dst)
+            self.assertTrue(torch.equal(dst, to_gpu_int4(cpu_layout)), (n, half))
+
+    def test_decode_kernel_computes_what_fused_experts_computes(self):
+        from vllm.config import VllmConfig, set_current_vllm_config
+        from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+        from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
+
+        from prain.kernels import moe_int4_decode
+        from prain.quant import dequantize, moe_quant_config, quantize
+        from prain.tiered import cpu_experts, slot_table
+
+        hidden, inter, experts, topk = 256, 128, 6, 4
+        w13 = torch.randn(experts, 2 * inter, hidden) * 0.1
+        w2 = torch.randn(experts, hidden, inter) * 0.1
+        parts = [quantize(w, "int4") for w in (w13, w2)]  # gpu layout
+        held = [1, 3, 4]  # the round's experts, in slots 0..2
+        slots = [[q[held].cuda(), s[held].cuda()] for q, s in parts]
+        table = slot_table({e: i for i, e in enumerate(held)}, experts, "cuda")
+        for tokens in (1, 2, 3, 4):
+            x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
+            ids = torch.stack([torch.randperm(experts)[:topk] for _ in range(tokens)]).int()
+            ids[0, -1] = experts  # padding route: the table's last row
+            w = torch.rand(tokens, topk)
+            got = moe_int4_decode(x.cuda(), w.cuda(), ids.cuda(), table, *slots[0], *slots[1])
+            with set_current_vllm_config(VllmConfig()):
+                vllm = fused_experts(x.cuda(), slots[0][0], slots[1][0], w.cuda(), ids.cuda(),
+                                     activation=MoEActivation.GELU_TANH,
+                                     global_num_experts=table.numel(), expert_map=table,
+                                     quant_config=moe_quant_config("int4", *[s[1] for s in slots]))
+            exact = cpu_experts(x, w, torch.where(ids < experts, ids, -1), set(held), lambda e: tuple(
+                dequantize(q, s, "int4")[e] for q, s in parts))
+            self.assertEqual((got.shape, got.dtype), ((tokens, hidden), torch.bfloat16))
+            torch.testing.assert_close(got.float().cpu(), exact, rtol=3e-2, atol=3e-2)
+            torch.testing.assert_close(got.float(), vllm.float(), rtol=3e-2, atol=3e-2)
+
+    def test_staged_copy_fills_a_slot_like_the_repack_path(self):
+        from prain.experts import _Staging, stage_int4
+        from prain.quant import quantize, to_cpu_int4, to_gpu_int4
+
+        w13, w2 = torch.randn(2, 256, 256), torch.randn(2, 256, 128)
+        pack = [[to_cpu_int4(q), s] for q, s in (quantize(w, "int4") for w in (w13, w2))]
+        slots = [[torch.zeros((3, *t.shape[1:]), dtype=t.dtype, device="cuda") for t in mats]
+                 for mats in pack]
+        ring = _Staging(sum(t[0].nbytes for mats in pack for t in mats), "cuda")
+        stream = torch.cuda.current_stream()
+        for expert, slot in ((1, 2), (0, 0), (1, 1), (0, 2), (1, 0)):  # wraps the ring
+            stage_int4([[t[expert] for t in mats] for mats in pack], slots, slot, ring, stream)
+            torch.cuda.synchronize()
+            for mats, dsts in zip(pack, slots):
+                self.assertTrue(torch.equal(dsts[0][slot], to_gpu_int4(mats[0][expert].cuda())))
+                self.assertTrue(torch.equal(dsts[1][slot], mats[1][expert].cuda()))
+
+
+    def test_int4_rounds_on_rdna2_match_the_bf16_kernel_even_when_fp16_overflows(self):
+        from types import SimpleNamespace
+
+        from vllm.config import VllmConfig, set_current_vllm_config
+        from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+        from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
+
+        from prain.experts import TieredRoutedExperts
+        from prain.quant import moe_quant_config, quantize
+        from prain.tiered import slot_table
+
+        hidden, inter, experts = 256, 128, 4
+        parts = [quantize(w, "int4") for w in (torch.randn(experts, 2 * inter, hidden) * 0.1,
+                                                 torch.randn(experts, hidden, inter) * 0.1)]
+        slots = [[q.cuda(), s.cuda()] for q, s in parts]
+        config = moe_quant_config("int4", slots[0][1], slots[1][1])
+        layer = SimpleNamespace(_slots=slots, activation=MoEActivation.GELU_TANH, _qconfig=config)
+        table = slot_table({e: e for e in range(experts)}, experts, "cuda")
+        with set_current_vllm_config(VllmConfig()):
+            for tokens, scale in ((2, 1.0), (64, 1.0), (64, 3000.0)):  # decode kernel, fp16, overflow
+                x = (torch.randn(tokens, hidden) * scale).bfloat16().cuda()
+                ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int().cuda()
+                w = torch.rand(tokens, 2).cuda()
+                got = TieredRoutedExperts._int4_gfx10(layer, x, w, ids, table)
+                want = fused_experts(x, slots[0][0], slots[1][0], w, ids, activation=layer.activation,
+                                     global_num_experts=table.numel(), expert_map=table,
+                                     quant_config=config)
+                self.assertTrue(torch.isfinite(got).all(), (tokens, scale))
+                torch.testing.assert_close(got.float(), want.float(), rtol=3e-2, atol=3e-2 * scale,
+                                           msg=f"{tokens} tokens, x * {scale}")
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_plugin.py

@@ -0,0 +1,40 @@
+import importlib.metadata
+import importlib.util
+import os
+import unittest
+from unittest import mock
+
+from tests.helpers import TmpCase
+
+
+@unittest.skipIf(importlib.util.find_spec("vllm") is None, "needs vllm")
+class Plugin(TmpCase):
+    def test_entry_point_is_installed(self):
+        plugins = importlib.metadata.entry_points(group="vllm.general_plugins")
+        self.assertIn("prain", [p.name for p in plugins])
+
+    def test_register_swaps_in_tiered_experts(self):
+        config = self.write("p.toml", f'pack_dir = "{self.dir}/pack"\n')
+        with mock.patch.dict(os.environ, {"PRAIN_CONFIG": config}):
+            from prain.plugin import register
+
+            register()
+            register()  # vLLM may call it again; must not raise
+        from vllm.model_executor.custom_op import op_registry_oot
+
+        from prain.experts import TieredRoutedExperts
+
+        self.assertIs(op_registry_oot["RoutedExperts"], TieredRoutedExperts)
+
+    def test_draft_hook_wraps_the_v2_model_runner_once(self):
+        from vllm.v1.worker.gpu.input_batch import InputBatch
+        from vllm.v1.worker.gpu.model_runner import GPUModelRunner
+
+        from prain.experts import _install_draft_hook
+
+        _install_draft_hook()
+        wrapped = GPUModelRunner.prepare_inputs
+        _install_draft_hook()
+        self.assertIs(GPUModelRunner.prepare_inputs, wrapped)
+        self.assertTrue(wrapped.prain)
+        self.assertLessEqual({"req_ids", "query_start_loc_np"}, set(InputBatch.__dataclass_fields__))

tests/test_policy.py

@@ -0,0 +1,125 @@
+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))
+        self.assertEqual(stats["vram"], [1, 2])  # the experts in VRAM now, for the web monitor
+        write_stats(path, {7: SlotCache([5, 3], 8, 2, Policy(eviction="lfu"))})
+        with open(path) as f:
+            self.assertEqual(json.load(f)["layers"]["7"]["vram"], [3, 5])  # pinned ones too

tests/test_profile.py

@@ -0,0 +1,100 @@
+import json
+import os
+import threading
+import time
+
+import torch
+
+from prain import profile
+from tests.helpers import TmpCase
+
+
+class Inner(torch.nn.Module):
+    def forward(self, x):
+        time.sleep(0.002)
+        return x
+
+
+class Outer(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.layers = torch.nn.ModuleList([Inner(), Inner()])
+        self.head = Inner()
+
+    def forward(self, x):
+        for layer in self.layers:
+            x = layer(x)
+        with profile.span("copy"):
+            time.sleep(0.001)
+        return self.head(x)
+
+
+class TestProfile(TmpCase):
+    def setUp(self):
+        super().setUp()
+        self.syncs = []
+        saved = profile.DIR, profile.SYNC[0]
+        profile.DIR, profile.SYNC[0] = self.dir, lambda: self.syncs.append(1)
+        self.addCleanup(lambda: (setattr(profile, "DIR", saved[0]),
+                                 profile.SYNC.__setitem__(0, saved[1])))
+
+    def steps(self, model, tokens):
+        for n in tokens:
+            profile.step_begin()
+            with profile.span("execute_model"):
+                profile.note_step(model, n)
+                model(torch.zeros(1))
+            profile.step_end()
+
+    def folded(self, kind):
+        with open(os.path.join(self.dir, f"{kind}.folded")) as f:
+            return {stack: int(us) for stack, us in (l.rsplit(" ", 1) for l in f.read().splitlines())}
+
+    def test_records_the_requested_steps_as_folded_stacks(self):
+        model = Outer()
+        self.steps(model, [1])  # no start file: nothing recorded, nothing synced
+        self.assertEqual(self.syncs, [])
+        self.write("start", "3")
+        self.steps(model, [1, 1, 300, 1])  # records 2 decode steps and 1 prefill, not the 4th
+        self.assertFalse(os.path.exists(os.path.join(self.dir, "start")))
+        with open(os.path.join(self.dir, "steps.json")) as f:
+            self.assertEqual(json.load(f)["steps"], {"decode": 2, "prefill": 1})
+        decode = self.folded("decode")
+        self.assertEqual(set(decode), {"step", "step;execute_model", "step;execute_model;Outer",
+                                       "step;execute_model;Outer;Inner",  # both layers, merged
+                                       "step;execute_model;Outer;copy",
+                                       "step;execute_model;Outer;head:Inner"})
+        for stack, ms in (("Outer;Inner", 8), ("Outer;head:Inner", 4), ("Outer;copy", 2)):
+            self.assertGreaterEqual(decode[f"step;execute_model;{stack}"], ms * 1000)
+            self.assertLess(decode[f"step;execute_model;{stack}"], ms * 1000 * 3)
+        self.assertGreater(self.folded("prefill")["step;execute_model;Outer;Inner"], 4000)
+        with open(os.path.join(self.dir, "decode.calls")) as f:
+            calls = dict(line.rsplit(" ", 1) for line in f.read().splitlines())
+        self.assertEqual(calls["step;execute_model;Outer;Inner"], "4")  # 2 layers x 2 steps
+        self.assertEqual(calls["step"], "2")
+        recorded = len(self.syncs)
+        self.steps(model, [1])  # finished: the module hooks are gone
+        self.assertEqual(len(self.syncs), recorded)
+        self.assertFalse(model.head._forward_hooks or model.head._forward_pre_hooks)
+
+    def test_hooks_leave_other_threads_inputs_and_outputs_alone(self):
+        model = Outer()
+        self.write("start", "1")
+        profile.step_begin()
+        profile.note_step(model, 1)  # installs the hooks
+        out = []
+        thread = threading.Thread(target=lambda: out.append(model(torch.ones(1))))
+        thread.start()
+        thread.join()
+        self.assertTrue(torch.equal(out[0], torch.ones(1)))
+        profile.step_end()
+
+    def test_svg_shows_time_per_step(self):
+        lines = ["step 500", "step;attn 3000", "step;attn;copy 1000"]
+        svg = profile.svg(lines, steps=2, title="decode")
+        self.assertTrue(svg.startswith("<svg"))
+        self.assertIn("attn: 2.0 ms per step (88.9%)", svg)  # 4000 us of 4500, over 2 steps
+        self.assertIn("copy: 0.5 ms per step", svg)
+        self.assertEqual(profile.top(lines, steps=2)[0], ("attn", 1.5, 66.7, None))  # self time
+        calls = ["step;attn 10", "step;attn;copy 4", "step 2"]
+        self.assertEqual(profile.top(lines, steps=2, calls=calls)[0], ("attn", 1.5, 66.7, 5.0))

tests/test_quant.py

@@ -0,0 +1,49 @@
+import unittest
+
+try:
+    import torch
+except ImportError:  # the planning host has no torch
+    torch = None
+
+if torch is not None:
+    from prain.quant import (GROUP, cpu_matmul, dequantize, layout, quantize, to_cpu_int4,
+                             to_gpu_int4)
+
+
+@unittest.skipIf(torch is None, "needs torch")
+class Quant(unittest.TestCase):
+    def setUp(self):
+        torch.manual_seed(0)
+        self.w = torch.randn(64, 256) * 0.02
+
+    def test_layouts_match_what_quantize_returns(self):
+        for quant in ("int8", "int4"):
+            got = [(tuple(t.shape), t.dtype) for t in quantize(self.w, quant)]
+            self.assertEqual(got, layout(quant, 64, 256))
+
+    def test_round_trip_error(self):
+        for quant, bound in (("int8", 0.01), ("int4", 0.11)):
+            err = (dequantize(*quantize(self.w, quant), quant, torch.float32) - self.w).norm()
+            self.assertLess(err / self.w.norm(), bound)
+
+    def test_int4_clipping_beats_plain_round_to_nearest(self):
+        g = self.w.unflatten(-1, (-1, GROUP))
+        scale = (g.abs().amax(-1, keepdim=True) / 7).bfloat16().float()
+        plain = ((g / scale).round() * scale - g).norm()
+        clipped = (dequantize(*quantize(self.w, "int4"), "int4", torch.float32) - self.w).norm()
+        self.assertLess(clipped, plain)
+
+    def test_cpu_layout_is_torch_avx2_int4pack_and_inverts(self):
+        q, _ = quantize(self.w, "int4")
+        nibbles = torch.stack((q & 15, q >> 4), -1).flatten(-2).to(torch.int32)
+        self.assertTrue(torch.equal(to_cpu_int4(q), torch._convert_weight_to_int4pack_for_cpu(nibbles, 1)))
+        self.assertTrue(torch.equal(to_gpu_int4(to_cpu_int4(q)), q))
+
+    def test_cpu_matmul_matches_dequantized_weights(self):
+        x = torch.randn(3, 256, dtype=torch.bfloat16)
+        for quant in ("int8", "int4"):
+            q, scale = quantize(self.w, quant)
+            ref = x.float() @ dequantize(q, scale, quant, torch.float32).t()
+            stored = to_cpu_int4(q) if quant == "int4" else q
+            got = cpu_matmul(quant, stored, scale)(x).float()
+            self.assertLess(((got - ref).norm() / ref.norm()).item(), 0.01, quant)

tests/test_sim.py

@@ -0,0 +1,101 @@
+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_cpu_path_pays_the_disk_read_of_disk_experts(self):
+        on_disk = {layer: LayerPlacement((), (), tuple(range(32)), 3) for layer in range(4)}
+        cpu = Miss(mode="cpu", admit=0)
+        ram = simulate(self.forwards, all_ram(3), Policy(), Prefetch(), cpu, Cost())
+        disk = simulate(self.forwards, on_disk, Policy(), Prefetch(), cpu, Cost())
+        self.assertAlmostEqual(disk.total_ms - ram.total_ms, ram.counters["cpu"] * Cost().read_ms(),
+                               places=3)
+
+    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)
+
+    def test_tune_searches_the_auto_threshold_and_keeps_the_cpu_threads(self):
+        forwards = forwards_of(synthetic_trace(**SMALL))
+        cost = Cost(expert_mib=64.0, h2d_gbps=2.5, cpu_ms=0.5)  # eGPU: a copy is ~50 CPU routes
+        base = (Policy(), Prefetch(), Miss(cpu_threads=16))
+        best, ram, per, final, bound, rows = tune(forwards, 32, 1.0, 1.0, cost, base=base)
+        self.assertEqual(best[2].cpu_threads, 16)
+        self.assertNotEqual(best[2].mode, "copy")
+        tried = {(v[2].mode, v[2].cpu_max_tokens) for v, _ in rows}
+        self.assertLessEqual({("copy", 1), ("cpu", 1), ("auto", 1), ("auto", 8), ("auto", 16)}, tried)

tests/test_tiered.py

@@ -0,0 +1,158 @@
+import unittest
+
+try:
+    import torch
+except ImportError:  # the planning host has no torch; the vLLM container does
+    torch = None
+
+from prain.config import Miss, Policy
+from prain.policy import SlotCache
+
+if torch is not None:
+    from prain.tiered import (available_experts, cpu_experts, degrade, route_stats, slot_table,
+                                tiered_forward)
+
+HIDDEN, INTER, EXPERTS = 6, 4, 10
+
+
+def naive_moe(x, w13, w2, topk_weights, topk_ids):
+    """Gated-GELU experts, one route at a time; id -1 = no route."""
+    out = torch.zeros(x.shape, dtype=torch.float64)
+    for t in range(x.shape[0]):
+        for j in range(topk_ids.shape[1]):
+            e = int(topk_ids[t, j])
+            if e >= 0:
+                h = w13[e] @ x[t]
+                act = torch.nn.functional.gelu(h[:INTER], approximate="tanh") * h[INTER:]
+                out[t] += topk_weights[t, j] * (w2[e] @ act)
+    return out
+
+
+@unittest.skipIf(torch is None, "needs torch")
+class TieredForward(unittest.TestCase):
+    def setUp(self):
+        torch.manual_seed(0)
+        self.w13 = torch.randn(EXPERTS, 2 * INTER, HIDDEN, dtype=torch.float64)
+        self.w2 = torch.randn(EXPERTS, HIDDEN, INTER, dtype=torch.float64)
+        self.x = torch.randn(6, HIDDEN, dtype=torch.float64)
+        self.topk_ids = torch.tensor([[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 0, 1],
+                                      [2, 3, -1], [4, 5, 6]], dtype=torch.int32)
+        self.topk_weights = torch.rand(6, 3, dtype=torch.float64)
+        self.expected = naive_moe(self.x, self.w13, self.w2, self.topk_weights, self.topk_ids)
+
+    def run_tiered(self, pinned, cache_slots, miss=None, policy=Policy(), draft=None):
+        cache = SlotCache(pinned, EXPERTS, cache_slots, policy)
+        slots13 = torch.zeros(len(pinned) + cache_slots, 2 * INTER, HIDDEN, dtype=torch.float64)
+        slots2 = torch.zeros(len(pinned) + cache_slots, HIDDEN, INTER, dtype=torch.float64)
+        for slot, e in enumerate(pinned):
+            slots13[slot], slots2[slot] = self.w13[e], self.w2[e]
+        loads, seen = [], []
+
+        def load(expert, slot):
+            loads.append(expert)
+            slots13[slot], slots2[slot] = self.w13[expert], self.w2[expert]
+
+        def kernel(x, w, ids, table, used):
+            self.assertTrue(bool((ids >= 0).all()))  # padding id -1 arrives as num_experts
+            return naive_moe(x, slots13, slots2, w, table[ids.long()])
+
+        out = tiered_forward(self.x, self.topk_weights, self.topk_ids, cache, load, kernel,
+                             miss=miss, cpu_weights=lambda e: (self.w13[e], self.w2[e]),
+                             on_stats=lambda stats, tokens: seen.append((tokens, sorted(stats))),
+                             draft=draft, min_weight=0.5)
+        self.assertEqual(seen[0][0], 6)
+        self.seen_experts = seen[0][1]
+        return out, loads, cache
+
+    def test_matches_dense_moe_when_everything_needs_rounds(self):
+        out, loads, _ = self.run_tiered(pinned=[0, 1], cache_slots=3)  # 8 others -> 3 rounds
+        torch.testing.assert_close(out, self.expected)
+        self.assertEqual(sorted(loads), [2, 3, 4, 5, 6, 7, 8, 9])
+
+    def test_matches_dense_moe_in_one_round(self):
+        out, _, _ = self.run_tiered(pinned=[], cache_slots=10)
+        torch.testing.assert_close(out, self.expected)
+
+    def test_cpu_path_matches_dense_moe(self):
+        for mode in ("cpu", "auto"):
+            out, loads, cache = self.run_tiered([0], 3, Miss(mode=mode, cpu_max_tokens=1))
+            torch.testing.assert_close(out, self.expected, rtol=1e-5, atol=1e-5)
+            self.assertEqual(cache.cpu + cache.misses + cache.hits, 10)
+        self.assertEqual(cpu_experts(self.x, self.topk_weights, self.topk_ids, set(),
+                                     None).abs().sum().item(), 0.0)
+
+    def test_miss_min_weight_skips_missing_experts_for_every_token(self):
+        """[miss] min_weight < 1: every token skips experts that are not in VRAM while it keeps
+        that share of its router weight; experts some token copies anyway serve all tokens."""
+        everyone = torch.ones(6, dtype=torch.bool)
+        resident = torch.zeros(EXPERTS, dtype=torch.bool)
+        resident[[0, 1, 2]] = True
+        _, first = degrade(self.topk_weights, self.topk_ids, everyone, resident, 0.5)
+        free = resident.clone()
+        free[first[first >= 0].long()] = True
+        w, i = degrade(self.topk_weights, self.topk_ids, everyone, free, 0.5)
+        out, loads, _ = self.run_tiered(pinned=[0, 1, 2], cache_slots=3, miss=Miss(min_weight=0.5))
+        torch.testing.assert_close(out, naive_moe(self.x, self.w13, self.w2, w, i))
+        self.assertEqual(sorted(loads), sorted(set(i[i >= 0].tolist()) - {0, 1, 2}))
+        self.assertLess(len(loads), 7)  # computing everything copies all 7 others
+        self.assertEqual(self.seen_experts, list(range(EXPERTS)))  # the trace sees the routing
+        out, loads, _ = self.run_tiered(pinned=[0, 1, 2], cache_slots=3, miss=Miss(min_weight=1.0))
+        torch.testing.assert_close(out, self.expected)
+        self.assertEqual(len(loads), 7)
+
+    def test_cpu_experts_take_matmul_functions(self):
+        """Quantized experts hand the CPU path x -> x @ w.T functions instead of weights."""
+        def matmuls(e):
+            return (lambda v: v @ self.w13[e].t()), (lambda v: v @ self.w2[e].t())
+
+        out = cpu_experts(self.x, self.topk_weights, self.topk_ids, set(range(EXPERTS)), matmuls)
+        torch.testing.assert_close(out, self.expected.float())
+
+    def test_route_stats(self):
+        stats = route_stats(torch.tensor([[0.5, 0.25], [1.0, 0.0]]),
+                            torch.tensor([[2, 1], [2, -1]]), 4, recency_half_life=1.0)
+        self.assertEqual(sorted(stats), [1, 2])
+        self.assertEqual(stats[2][0], 2)
+        self.assertAlmostEqual(stats[2][1], 1.5)
+        self.assertAlmostEqual(stats[2][2], 0.5 * 0.5 + 1.0)  # the first token is one half-life old
+        self.assertAlmostEqual(stats[1][2], 0.125)
+
+    def test_slot_table_maps_the_round_and_skips_the_rest(self):
+        table = slot_table({7: 4, 2: 0}, 10, "cpu")
+        self.assertEqual(table.tolist(), [-1, -1, 0, -1, -1, -1, -1, 4, -1, -1, -1])
+        self.assertEqual(table.dtype, torch.int32)  # one extra row: padding id 10 -> -1
+
+    def test_degrade_drops_unavailable_routes_of_draft_tokens_only(self):
+        weights = torch.tensor([[0.5, 0.3, 0.2], [0.5, 0.3, 0.2]])
+        ids = torch.tensor([[0, 1, 2], [0, 1, 2]], dtype=torch.int32)
+        draft = torch.tensor([False, True])
+        available = torch.tensor([False, True, False, False])
+        w, i = degrade(weights, ids, draft, available, min_weight=0.5)
+        # the draft token computes 0.3 of 1.0: it keeps expert 0 (heaviest) to reach 0.5
+        self.assertEqual(i.tolist(), [[0, 1, 2], [0, 1, -1]])
+        torch.testing.assert_close(w, torch.tensor([[0.5, 0.3, 0.2], [0.625, 0.375, 0.0]]))
+        w, i = degrade(weights, ids, draft, available, min_weight=0.0)
+        self.assertEqual(i[1].tolist(), [-1, 1, -1])
+        torch.testing.assert_close(w[1], torch.tensor([0.0, 1.0, 0.0]))
+
+    def test_draft_tokens_use_vram_and_what_full_tokens_load(self):
+        draft = torch.tensor([False, True, False, True, True, False])
+        out, loads, _ = self.run_tiered(pinned=[0], cache_slots=3, draft=draft)
+        # available: pinned expert 0 plus every expert of the full rows 0, 2 and 5
+        available = torch.zeros(EXPERTS, dtype=torch.bool)
+        available[[0, 1, 2, 4, 5, 6, 7, 8]] = True
+        self.assertTrue(torch.equal(
+            available_experts(self.topk_ids, draft, {0}, EXPERTS), available))
+        _, kept = degrade(self.topk_weights, self.topk_ids, draft, available, 0.5)
+        available[kept[kept >= 0].long()] = True  # experts some row copies anyway serve all rows
+        expected = naive_moe(self.x, self.w13, self.w2,
+                             *degrade(self.topk_weights, self.topk_ids, draft, available, 0.5))
+        torch.testing.assert_close(out, expected)
+        torch.testing.assert_close(out[[0, 2, 5]], self.expected[[0, 2, 5]])  # full rows exact
+        self.assertEqual(self.seen_experts, list(range(EXPERTS)))  # trace sees the routing
+
+    def test_drafts_are_exact_when_every_expert_is_in_vram(self):
+        out, loads, _ = self.run_tiered(pinned=list(range(EXPERTS)), cache_slots=0,
+                                        draft=torch.ones(6, dtype=torch.bool))
+        torch.testing.assert_close(out, self.expected)
+        self.assertEqual(loads, [])

tests/test_tui.py

@@ -0,0 +1,81 @@
+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_a_persona_leads_every_request_but_stays_off_screen(self):
+        chat = Chat(self.stub.url, "stub-model", system="Be brief.")
+        chat.type("Why")
+        wait_for(lambda: self.stub.bodies and not chat.warmer.busy)
+        chat.send()
+        wait_for(lambda: chat.poll() or chat.reply is None)
+        system = {"role": "system", "content": "Be brief."}
+        self.assertEqual([b["messages"][0] for b in self.stub.bodies], [system, system])
+        self.assertEqual(chat.messages, [{"role": "user", "content": "Why"},
+                                         {"role": "assistant", "content": "Hello!"}])
+
+    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,207 @@
+import json
+import os
+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, TmpCase, 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):
+        with chat.post(self.url, "/warm", {"session": session, "messages": [],
+                                           "draft": draft}) as response:
+            return response.status
+
+    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:
+            page = response.read()
+        self.assertIn(b'<textarea id="draft"', page)
+        self.assertIn(b'"/status"', page)  # the monitor polls it
+        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.assertEqual(self.warm("bob", "typing along"), 204)  # 204: nothing warms up
+        self.server.streams[0] = 0
+        self.assertEqual(self.warm("bob", "now"), 202)  # 202: a warm-up starts
+        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_a_persona_leads_every_warm_up_and_answer(self):
+        self.start(system="You answer in German.", persona="german")
+        self.warm("s1", "Wie spät")
+        wait_for(lambda: self.stub.bodies)
+        with self.ask("s1", "Wie spät ist es?") as r:
+            r.read()
+        system = {"role": "system", "content": "You answer in German."}
+        self.assertEqual([b["messages"][0] for b in self.stub.bodies], [system, system])
+        self.assertEqual(self.stub.bodies[1]["messages"][1:],
+                         [{"role": "user", "content": "Wie spät ist es?"}])
+        with urllib.request.urlopen(self.url + "/status") as r:
+            self.assertEqual(json.load(r)["persona"], "german")
+
+    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)
+
+
+class Status(TmpCase):
+    def setUp(self):
+        super().setUp()
+        self.stub = StubVllm()
+        self.addCleanup(self.stub.close)
+
+    def status(self, **options):
+        server = make_server(self.stub.url, 0, **options)
+        threading.Thread(target=server.serve_forever, daemon=True).start()
+        self.addCleanup(server.server_close)
+        self.addCleanup(server.shutdown)
+        with urllib.request.urlopen(f"http://127.0.0.1:{server.server_port}/status") as r:
+            return json.load(r)
+
+    def test_status_reports_vllm_the_gpu_and_the_experts(self):
+        self.stub.metrics = (
+            "# HELP vllm:num_requests_running Number of requests running.\n"
+            'vllm:num_requests_running{engine="0",model_name="m"} 1.0\n'
+            'vllm:num_requests_waiting{engine="0",model_name="m"} 0.0\n'
+            'vllm:num_requests_waiting_by_reason{engine="0",reason="capacity"} 5.0\n'
+            'vllm:kv_cache_usage_perc{engine="0",model_name="m"} 0.25\n'
+            'vllm:prefix_cache_queries_total{engine="0",model_name="m"} 40.0\n'
+            'vllm:prefix_cache_hits_total{engine="0",model_name="m"} 30.0\n'
+            'vllm:prompt_tokens_total{engine="0",model_name="m"} 100.0\n'
+            'vllm:generation_tokens_total{engine="0",model_name="m"} 7.0\n')
+        os.makedirs(os.path.join(self.dir, "card0", "device"))  # an iGPU: no VRAM files
+        gpu = os.path.join(self.dir, "card1", "device")
+        os.makedirs(os.path.join(gpu, "hwmon", "hwmon4"))
+        for name, value in {"mem_info_vram_total": 8573157376, "mem_info_vram_used": 8000000000,
+                            "gpu_busy_percent": 96, "hwmon/hwmon4/temp1_input": 61000,
+                            "hwmon/hwmon4/power1_average": 120500000,
+                            "hwmon/hwmon4/fan1_input": 1800}.items():
+            self.write(os.path.join("card1", "device", name), f"{value}\n")
+        stats = {"layers": {"0": {"uses": [3, 0], "vram": [0], "hits": 2, "misses": 1}}}
+        stats_path = self.write("stats.json", json.dumps(stats))
+        got = self.status(stats_path=stats_path, sysfs=self.dir)
+        self.assertEqual(got["model"], "stub-model")
+        self.assertEqual(got["vllm"], {"running": 1.0, "waiting": 0.0, "kv_usage": 0.25,
+                                       "prefix_queries": 40.0, "prefix_hits": 30.0,
+                                       "prompt_tokens": 100.0, "generation_tokens": 7.0})
+        self.assertEqual(got["gpu"], {"vram_used": 8000000000, "vram_total": 8573157376,
+                                      "busy": 96, "temp": 61.0, "power": 120.5, "fan": 1800.0})
+        self.assertEqual(got["experts"], stats)
+        self.assertEqual(got["streams"], 0)
+
+    def test_status_without_metrics_gpu_or_stats(self):
+        got = self.status(sysfs=self.dir)  # the stub answers /metrics with JSON, not metrics
+        self.assertEqual((got["gpu"], got["experts"]), (None, None))
+        self.assertEqual(got["vllm"]["running"], 0.0)

tests/test_web_page.py

@@ -0,0 +1,49 @@
+import json
+import shutil
+import subprocess
+import unittest
+from pathlib import Path
+
+PAGE = (Path(__file__).parent.parent / "prain" / "web.html").read_text()
+
+
+@unittest.skipUnless(shutil.which("node"), "needs node")
+class Markdown(unittest.TestCase):
+    """The page's md() (model output -> innerHTML) run in node."""
+
+    def md(self, text):
+        script = PAGE[PAGE.index("<script>") + 8:PAGE.index("/* ---------- thread")]
+        script = script.split("/* ---------- markdown", 1)[1].split("*/", 1)[1]
+        out = subprocess.run(["node", "-e", script + f"\nprocess.stdout.write(md({json.dumps(text)}))"],
+                             capture_output=True, text=True, check=True)
+        return out.stdout
+
+    def test_html_in_model_output_is_text(self):
+        html = self.md('<img src=x onerror="alert(1)"> **<b>hi</b>** `<i>` [x](javascript:alert(1))')
+        self.assertNotIn("<img", html)
+        self.assertNotIn("<b>", html)
+        self.assertNotIn('href="javascript', html)
+        self.assertIn("&lt;img src=x onerror=&quot;alert(1)&quot;&gt;", html)
+        self.assertIn("<strong>&lt;b&gt;hi&lt;/b&gt;</strong>", html)
+
+    def test_blocks(self):
+        html = self.md("# Title\n\nSome *very* **bold** text with `code` and "
+                       "[a link](https://x.org/?a=1&b=2).\n\n```python\nif a < b:\n    pass\n```\n"
+                       "3. three\n4. four\n\n- one\n  more\n- two\n\n| a | b |\n|---|---|\n| 1 | 2 |\n"
+                       "\n> quoted\n\n---\nend")
+        for part in ("<h3>Title</h3>", "<em>very</em>", "<strong>bold</strong>", "<code>code</code>",
+                     '<a href="https://x.org/?a=1&amp;b=2" target="_blank" rel="noopener noreferrer">',
+                     '<span>python</span>', "<pre><code>if a &lt; b:\n    pass</code></pre>",
+                     '<ol start="3"><li>three</li><li>four</li></ol>',
+                     "<ul><li>one more</li><li>two</li></ul>",
+                     "<th>a</th><th>b</th></tr></thead><tbody><tr><td>1</td><td>2</td>",
+                     "<blockquote><p>quoted</p></blockquote>", "<hr>", "<p>end</p>"):
+            self.assertIn(part, html)
+
+    def test_an_open_fence_while_streaming_is_code(self):
+        self.assertIn("<pre><code>x = 1\n# not a heading</code></pre>",
+                      self.md("Here:\n```\nx = 1\n# not a heading"))
+
+
+if __name__ == "__main__":
+    unittest.main()