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

@@ -1,6 +1,12 @@
-# prain placement: prithivMLmods/gemma-4-26B-A4B-it-Uncensored-MAX on an 8 GB RX 6600 XT
-# Model: 30 MoE layers x 128 experts, top-8 routing, 11.34 MiB per expert (BF16).
-# VRAM: 8 GiB - 4.45 GiB dense weights - 0.75 GiB KV cache - ~0.9 GiB runtime = ~1.7 GiB experts.
+# 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
@@ -10,19 +16,23 @@
 # Layer selectors: "*", "12", "0-9,20".   Expert lists: "*", "0-15,32", [1, 2, 3].
 # Check what a file costs:   prain check <this file> <model config.json>
 # Compare variants offline:  prain sim <trace> <this file> <model config.json>
-#                            prain tune <trace> <model config.json> --expert-vram-gib 1.7 ...
+#                            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"  # 42.5 GiB, written on the first run
+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 = 5                          # 30 x 5 x 11.34 MiB = 1.66 GiB of VRAM
-stats_path = "../models/packs/gemma4-26b-a4b/stats.json"
-stats_every = 64                         # MoE forward passes between stats/trace writes
+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 = "lfu"      # lru | lfu (decayed frequency) | gdsf (frequency x refetch cost)
-heat = "weight"       # what makes an expert hot: count (tokens) | weight (router mass) |
+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
@@ -39,16 +49,23 @@ min_prob = 0.05       # skip guesses below this router probability
 
 [miss]                # an expert that is not in VRAM
 mode = "copy"         # copy | cpu (always compute it on the CPU) | auto (CPU if it has
-cpu_max_tokens = 1    #   at most cpu_max_tokens tokens, else copy)
-cpu_threads = 4       # CPU threads of the vLLM worker (sets OMP_NUM_THREADS)
-admit = 1             # cpu/auto: copy this many CPU-run experts into VRAM afterwards
+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 = 0.75   # KV cache VRAM shared by them (--kv-cache-memory-bytes). vLLM logs
-                      #   "Maximum concurrency for 4096 tokens per request: N.NNx" at start;
-                      #   below max_requests, long requests wait or get preempted. Every GiB
-                      #   here is 90 expert cache slots less.
+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

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

@@ -34,6 +34,11 @@ 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,

prain/cli.py

@@ -7,19 +7,19 @@ import os
 import sys
 from dataclasses import asdict
 
-from prain.config import (GIB, Miss, Policy, Prefetch, expert_bytes, format_ids, load_config,
+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):
+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)
+    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)
+    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):
@@ -43,6 +43,8 @@ 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}")
 
@@ -51,7 +53,7 @@ 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)
+    _, 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})
@@ -64,13 +66,15 @@ def sim(trace, config_path, hf_config_path, cost_path, expert_mib):
 
 
 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)
+    _, 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)
+                                              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):
@@ -84,7 +88,7 @@ def tune_command(args):
               for layer in per}
     print(render_toml(
         f"generated by `prain tune` from {os.path.basename(args.trace)}",
-        {"pack_dir": os.path.abspath(args.pack_dir), "default_tier": "disk",
+        {"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,
@@ -109,6 +113,9 @@ def serve(args):
                        "--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])
 
 
@@ -126,6 +133,7 @@ def main(argv=None):
     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)
@@ -139,6 +147,9 @@ def main(argv=None):
             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")
@@ -157,20 +168,26 @@ def main(argv=None):
         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("--config", help="placement file: [serve] max_requests and the "
-                                             "[draft] table (default: their defaults)")
+            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)
+        _, 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)), end="")
+                   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":
@@ -185,13 +202,23 @@ def main(argv=None):
     elif args.command == "tui":
         from prain.tui import main as tui  # curses only loads when needed
 
-        tui(args.url, args.warmup, args.warmup_priority)
+        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
-        web(args.url, args.port, args.warmup, args.warmup_priority, config.serve, config.draft)
+        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__":

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}
 
@@ -104,6 +106,8 @@ class Miss:
     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)
@@ -143,6 +147,7 @@ def _section(cls, raw, where):
 @dataclass(frozen=True)
 class Config:
     pack_dir: str
+    quant: str = "bf16"
     default_tier: str = "ram"
     cache_slots: int = 16
     stats_path: str | None = None
@@ -221,6 +226,7 @@ def load_config(path):
         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")),
@@ -236,6 +242,8 @@ def load_config(path):
     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")
@@ -262,9 +270,14 @@ def moe_dims(hf_config):
     )
 
 
-def expert_bytes(hidden, moe_intermediate, dtype_bytes=2):
-    """One gated expert: gate+up [2I, H] plus down [H, I]."""
-    return 3 * hidden * moe_intermediate * dtype_bytes
+def 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):
@@ -278,7 +291,8 @@ def render_toml(header, top, tables=None, layers=None):
     return "\n".join(lines) + "\n"
 
 
-def plan(stats, num_experts, expert_size, expert_vram_gib, ram_gib, cache_slots, pack_dir, stats_path):
+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)
@@ -298,7 +312,7 @@ def plan(stats, num_experts, expert_size, expert_vram_gib, ram_gib, cache_slots,
     }
     return render_toml(
         "generated by `prain plan`; experts not listed live on disk",
-        {"pack_dir": pack_dir, "default_tier": "disk", "cache_slots": cache_slots,
+        {"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/experts.py

@@ -2,14 +2,17 @@
 
 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. 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:
+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 the
-     stock fused-MoE kernel on the slot tensors, in rounds if needed,
+  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.
@@ -26,11 +29,14 @@ 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
+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
 
@@ -40,6 +46,68 @@ _COPY_STREAM = []  # one side stream shared by all layers: async copies stay in
 _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():
@@ -121,6 +189,8 @@ class TieredRoutedExperts(RoutedExperts):
             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")
@@ -133,9 +203,16 @@ class TieredRoutedExperts(RoutedExperts):
             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.
-        moe_config.num_experts = moe_config.num_local_experts = placement.num_slots
+        # 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
@@ -155,27 +232,69 @@ class TieredRoutedExperts(RoutedExperts):
         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",
+            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)
-        stem = os.path.join(pack_dir, f"layer{layer:03d}")
-        w13_shape = (num_experts, *self.w13_weight.shape[1:])
-        w2_shape = (num_experts, *self.w2_weight.shape[1:])
-        self._pack_files = (f"{stem}.w13", f"{stem}.w2")
-        self._pack_w13 = _file_tensor(self._pack_files[0], w13_shape, self.w13_weight.dtype)
-        self._pack_w2 = _file_tensor(self._pack_files[1], w2_shape, self.w2_weight.dtype)
-        self._marker_path = f"{stem}.done"
-        self._marker = {"model": model, "w13": list(w13_shape), "w2": list(w2_shape),
-                        "dtype": str(self.w13_weight.dtype)}
-        self.pack_complete = os.path.exists(self._marker_path)
-        if self.pack_complete:
-            with open(self._marker_path) as f:
+        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 != self._marker:
-                raise ValueError(f"prain: {self._marker_path} is for {found}, not "
-                                 f"{self._marker}; delete {pack_dir} or use another pack_dir")
+        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):
@@ -198,8 +317,8 @@ class TieredRoutedExperts(RoutedExperts):
     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 not self.pack_complete:
-            pack = self._pack_w13 if param is self.w13_weight else self._pack_w2
+        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
@@ -208,15 +327,24 @@ class TieredRoutedExperts(RoutedExperts):
 
     def _materialize(self):
         if not self.pack_complete:
-            for path in self._pack_files:  # mmap writes -> disk before we trust the pack
-                fd = os.open(path, os.O_RDONLY)
-                os.fsync(fd)
-                os.close(fd)
-            with open(self._marker_path, "w") as f:
-                json.dump(self._marker, f)
-            self.pack_complete = True
+            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 = self._pack_w13[0].nbytes + self._pack_w2[0].nbytes
+            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:
@@ -225,56 +353,72 @@ class TieredRoutedExperts(RoutedExperts):
                                    "from ram to disk")
         ram = list(self.placement.ram)
         self._ram_row = {e: row for row, e in enumerate(ram)}
-        index = torch.tensor(ram, dtype=torch.int64, device="cpu")
-        self._ram_w13, self._ram_w2 = (
-            torch.index_select(pack, 0, index, out=torch.empty(
-                (len(ram), *pack.shape[1:]), dtype=pack.dtype, device="cpu", pin_memory=True))
-            for pack in (self._pack_w13, self._pack_w2)
-        )
+        # 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):
-        """CPU tensors of an expert: pinned rows for the RAM tier, else the mmap'd pack."""
+        """[[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 self._pack_w13[expert], self._pack_w2[expert]
-        return self._ram_w13[row], self._ram_w2[row]
+            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)
-        w13, w2 = self._source(expert)
-        self.w13_weight.data[slot].copy_(w13, non_blocking=True)
-        self.w2_weight.data[slot].copy_(w2, non_blocking=True)
+        self._copy(expert, slot)
 
     def _load_async(self, expert, slot):
-        """Copy on the side stream; only RAM-tier (pinned) experts are truly asynchronous."""
+        """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)
-        w13, w2 = self._source(expert)
         with torch.cuda.stream(stream):
-            self.w13_weight.data[slot].copy_(w13, non_blocking=True)
-            self.w2_weight.data[slot].copy_(w2, non_blocking=True)
+            self._copy(expert, slot)
         event = torch.cuda.Event()
         event.record(stream)
         self._ready_events[slot] = event
 
-    def _run_slots(self, x, weights, slot_ids, used_slots):
+    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)
-        out = super().forward_modular(x, weights, slot_ids)
-        if not isinstance(out, torch.Tensor):
-            raise NotImplementedError("prain: deferred MoE finalize is not supported")
+        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)
@@ -282,6 +426,22 @@ class TieredRoutedExperts(RoutedExperts):
                 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
@@ -298,7 +458,10 @@ class TieredRoutedExperts(RoutedExperts):
             self._predictions[d] = list(zip(e.tolist(), p.tolist()))  # host sync
 
     def _cpu_weights(self, expert):
-        return self._source(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
@@ -308,8 +471,9 @@ class TieredRoutedExperts(RoutedExperts):
                 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 e in target._ram_row]
+                       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)
 
@@ -330,6 +494,8 @@ class TieredRoutedExperts(RoutedExperts):
         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,
@@ -347,6 +513,8 @@ class TieredRoutedExperts(RoutedExperts):
         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:

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

@@ -10,7 +10,9 @@ def register():
 
     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

@@ -236,8 +236,10 @@ COUNTERS = ("hits", "misses", "cpu", "prefetches", "prefetch_hits", "prefetch_wa
 
 
 def write_stats(path, caches):
-    """Atomically write {"layers": {"<layer>": {"uses": [...], "hits": n, ...}}}."""
-    data = {"layers": {str(layer): {"uses": c.uses, **{k: getattr(c, k) for k in COUNTERS}}
+    """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:

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

@@ -31,9 +31,13 @@ class Cost:
     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 + (mb / self.disk_gbps if disk else 0.0)
+        return mb / self.h2d_gbps + (self.read_ms() if disk else 0.0)
 
 
 def load_cost(path, expert_mib):
@@ -175,7 +179,8 @@ def simulate(forwards, placements, policy, prefetch, miss, cost, belady=False):
                     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)
+            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))
@@ -226,15 +231,19 @@ DIMENSIONS = (
     ("prefetch", "mode", ("none", "pregate")),
     ("prefetch", "distance", (1, 2)),
     ("prefetch", "budget", (1, 2, 4)),
-    ("miss", "mode", ("copy", "auto", "cpu")),
+    # 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]], **{name: value})
+    parts[_PART[table]] = replace(parts[_PART[table]], **dict(zip(names, values)))
     return tuple(parts)
 
 

prain/tiered.py

@@ -3,6 +3,8 @@
 import torch
 import torch.nn.functional as F
 
+from prain import profile
+
 HEAT_INDEX = {"count": 0, "weight": 1, "recency": 2}
 
 
@@ -49,7 +51,7 @@ def degrade(topk_weights, topk_ids, draft, available, min_weight):
 
     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 remap ignore.
+    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
@@ -68,37 +70,36 @@ def degrade(topk_weights, topk_ids, draft, available, min_weight):
     return weights, torch.where(drop, torch.full_like(topk_ids, -1), topk_ids)
 
 
-def remap(topk_weights, topk_ids, slot_of, num_experts):
-    """Point routes at VRAM slots for one round.
-
-    Routes to experts outside this round (and padding routes with id -1) get slot 0
-    and weight 0, so any fused-MoE kernel can run them unchanged."""
-    table = torch.full((num_experts,), -1, dtype=torch.int64)
+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), dtype=torch.int64)] = torch.tensor(
-            list(slot_of.values()), dtype=torch.int64
-        )
-    table = table.to(topk_ids.device)
-    ids = topk_ids.long()
-    slots = torch.where(ids >= 0, table[ids.clamp_min(0)], -1)
-    weights = torch.where(slots >= 0, topk_weights, torch.zeros_like(topk_weights))
-    return weights, slots.clamp_min(0).to(topk_ids.dtype)
+        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]). Returns [tokens, hidden] fp32."""
+    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
-        w13, w2 = weights_of(e)
-        h = x[rows].to(w13.dtype) @ w13.t()
-        half = w13.shape[0] // 2
+        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 = (act.to(w2.dtype) @ w2.t()).float()
+        y = mm2(act.to(h.dtype)).float()
         out.index_add_(0, rows, y * topk_weights[rows, cols].float()[:, None])
     return out
 
@@ -108,21 +109,35 @@ def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight"
                    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,
-    slot_ids, used_slots) runs the stock fused-MoE kernel on the slot tensors. With a
+    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: 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
-    if draft is not None:
-        if on_stats is not None:  # the trace records the routing, not what drafts skipped
+    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)
-        available = available_experts(topk_ids, draft, set(cache.pinned) | set(cache.slots),
+            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)
-        topk_weights, topk_ids = degrade(topk_weights, topk_ids, draft, available, min_weight)
-    stats = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)  # host sync
+            _, 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)
@@ -131,22 +146,28 @@ def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight"
     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]
-    outs = []
-    for slot_of, loads in cache.plan({e: s[index] for e, s in stats.items()}, tokens, cpu):
+    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:
-            weights, slots = remap(topk_weights, topk_ids, slot_of, cache.num_experts)
-            outs.append(kernel(x, weights, slots, list(slot_of.values())))
+        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
-        outs.append(cpu_experts(*host, cpu, cpu_weights).to(x.device))
-    if not outs:
-        return torch.zeros_like(x)
-    if len(outs) == 1:
-        return outs[0].to(x.dtype)
-    # ponytail: every extra round re-runs all tokens with weight 0 for experts it does not
-    # hold; only big batches need many rounds. Upgrade: compact tokens per round.
-    acc = torch.promote_types(x.dtype, torch.float32)
-    return torch.stack([o.to(acc) for o in outs]).sum(0).to(x.dtype)
+        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

@@ -9,24 +9,26 @@ import threading
 import time
 
 from prain.chat import (ABORTED, Warmer, chat_body, first_model, post, send_request,
-                          stream_text, warm_body)
+                          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."""
+    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):
-        self.base_url, self.model, self.warmup = base_url, model, warmup
+    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, self.messages, draft, priority)))
+            base_url, "/v1/chat/completions",
+            warm_body(model, with_system(system, self.messages), draft, priority)))
 
     def type(self, text):
         self.draft = text
@@ -40,7 +42,8 @@ class Chat:
         self.warmer.cancel()  # a queued draft is stale now; the in-flight one only delays us
         self.messages.append({"role": "user", "content": self.draft.strip()})
         self.draft, self.reply = "", ""
-        threading.Thread(target=self._stream, args=(list(self.messages),), daemon=True).start()
+        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
@@ -108,8 +111,8 @@ def run(screen, chat):
             chat.type(chat.draft + key)
 
 
-def main(base_url, warmup="idle", priority=1):
+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)
+    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

@@ -1,10 +1,13 @@
-"""prain web: a one-page browser chat that warms experts up while you type.
+"""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. POST /warm takes
-{"session", "messages", "draft"} and returns at once. POST /chat takes {"session",
-"messages"} and streams server-sent events: vLLM's chunks, plus {"prain": "draft" |
-"refine" | "error"} events of its own. Binds to 127.0.0.1; use an SSH tunnel from other
-machines.
+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
@@ -19,127 +22,92 @@ Several browsers at once ([serve] and [draft] tables of the placement file):
   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
+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 = """<!doctype html>
-<meta charset="utf-8">
-<title>prain</title>
-<style>
-  body { font: 15px/1.45 system-ui, sans-serif; max-width: 52rem; margin: 2rem auto; padding: 0 1rem; }
-  .user { color: #666; margin-top: 1rem; }
-  .assistant { white-space: pre-wrap; }
-  .draft { color: #999; }
-  textarea { width: 100%; height: 5rem; font: inherit; margin-top: 1rem; }
-  #status { color: #888; font-size: 13px; }
-</style>
-<div id="log"></div>
-<textarea id="draft" autofocus
-  placeholder="Experts load while you type. Enter sends, Shift+Enter adds a line."></textarea>
-<div id="status"></div>
-<script>
-const session = Math.random().toString(36).slice(2) + Date.now().toString(36);
-const messages = [];
-const log = document.getElementById("log");
-const draft = document.getElementById("draft");
-const status = document.getElementById("status");
-let streaming = false;
+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"}
 
-function show(role, text) {
-  const div = document.createElement("div");
-  div.className = role;
-  div.textContent = text;
-  log.appendChild(div);
-  return div;
-}
 
-draft.addEventListener("input", () => {
-  if (!streaming && draft.value.trim())
-    fetch("/warm", {method: "POST", body: JSON.stringify({session, messages, draft: draft.value})});
-});
+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
 
-draft.addEventListener("keydown", async (event) => {
-  if (event.key !== "Enter" || event.shiftKey || streaming || !draft.value.trim()) return;
-  event.preventDefault();
-  const text = draft.value.trim();
-  messages.push({role: "user", content: text});
-  show("user", "you> " + text);
-  draft.value = "";
-  streaming = true;
-  const out = show("assistant", "");
-  const finalSpan = out.appendChild(document.createElement("span"));
-  const draftSpan = out.appendChild(document.createElement("span"));
-  draftSpan.className = "draft";
-  let mode = "full", draftText = "", finalText = "", refined = false;
-  const render = () => {  // the full text, then the grey rest of the draft while they agree
-    finalSpan.textContent = finalText;
-    draftSpan.textContent = mode === "draft" ? draftText
-      : mode === "refine" && draftText.startsWith(finalText) ? draftText.slice(finalText.length) : "";
-  };
-  let start = performance.now(), first = 0, pieces = 0, buffer = "";
-  const response = await fetch("/chat", {method: "POST", body: JSON.stringify({session, messages})});
-  if (!response.ok) {
-    out.textContent = "error " + response.status + ": " + await response.text();
-    messages.pop();
-    streaming = false;
-    return;
-  }
-  const reader = response.body.getReader();
-  const decoder = new TextDecoder();
-  for (;;) {
-    const {done, value} = await reader.read();
-    if (done) break;
-    buffer += decoder.decode(value, {stream: true});
-    const lines = buffer.split("\\n");
-    buffer = lines.pop();
-    for (const line of lines) {
-      if (!line.startsWith("data: ") || line === "data: [DONE]") continue;
-      const data = JSON.parse(line.slice(6));
-      if (data.prain === "draft") {
-        mode = "draft";
-        status.textContent = "GPU busy: drafting from the experts in VRAM, refined when it is free";
-      } else if (data.prain === "refine") {
-        mode = "refine";
-        refined = true;
-        start = performance.now(), first = 0, pieces = 0;
-        status.textContent = "refining with all experts";
-      } else if (data.prain === "error") {
-        mode = "draft";
-        refined = false;
-        status.textContent = "refining failed (HTTP " + data.status + "), keeping the draft";
-      } else {
-        const piece = data.choices[0].delta.content;
-        if (!piece) continue;
-        if (pieces++ === 0) first = performance.now();
-        if (mode === "draft") draftText += piece; else finalText += piece;
-      }
-      render();
-    }
-  }
-  messages.push({role: "assistant", content: mode === "draft" ? draftText : finalText});
-  let same = 0;
-  while (same < draftText.length && draftText[same] === finalText[same]) same++;
-  const seconds = (performance.now() - first) / 1000;
-  status.textContent = (refined ? `draft kept ${same} of ${draftText.length} chars; refined: `
-                        : mode === "draft" ? "draft, not refined: " : "") +
-    `first token ${((first - start) / 1000).toFixed(1)} s, ` +
-    `${pieces > 1 ? ((pieces - 1) / seconds).toFixed(1) : "0.0"} tok/s`;
-  streaming = false;
-});
-</script>
-"""
 
 
 def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, serve=Serve(),
-                draft=Draft()):
+                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
@@ -167,7 +135,14 @@ def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, ser
 
     class Handler(BaseHTTPRequestHandler):
         def do_GET(self):
-            self.reply(200, "text/html; charset=utf-8", PAGE.encode())
+            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"])))
@@ -177,12 +152,13 @@ def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, ser
                     busy = streams[0]
                 if body["draft"].strip() and ((warmup == "idle" and not busy) or
                                               (warmup == "always" and busy < serve.max_requests)):
-                    warmer(session).update((body["messages"], body["draft"]))
+                    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 = body["messages"], object()
+            messages, ticket = with_system(system, body["messages"]), object()
             with lock:
                 streams[0] += 1
                 holds_full = full[0] < draft.full_requests and not waiting
@@ -267,8 +243,11 @@ def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, ser
     return server
 
 
-def serve(base_url, port, warmup="idle", priority=1, serve=Serve(), draft=Draft()):
-    server = make_server(base_url, port, warmup=warmup, priority=priority, serve=serve, draft=draft)
-    print(f"prain web: http://127.0.0.1:{server.server_port}  (vLLM at {base_url}; "
-          f"{draft.full_requests} full of {serve.max_requests} parallel answers)")
+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

@@ -4,14 +4,16 @@
 #   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_VLLM_REF
+# 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}
 
@@ -90,9 +92,23 @@ torch() {
     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
+    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

scripts/env.sh

@@ -23,4 +23,7 @@ export PATH
 # 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

@@ -1,7 +1,7 @@
 """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 [--threads N] [--dense-gib 4.45] [--layers 30] > cost.json
+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).
 """
 
@@ -15,6 +15,8 @@ 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
@@ -37,10 +39,11 @@ def main():
     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 = 3 * HIDDEN * INTER * 2
-    host = torch.empty(expert_bytes // 2, dtype=torch.bfloat16).pin_memory()
+    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))
 
@@ -56,13 +59,17 @@ def main():
 
     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)
+                                     global_num_experts=SLOTS, quant_config=quant_config)
 
     with set_current_vllm_config(VllmConfig()):
         round_ms = ms(moe(1))
@@ -73,15 +80,28 @@ def main():
     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())
-    w13_cpu = torch.randn(2 * INTER, HIDDEN, dtype=torch.bfloat16).pin_memory()
-    w2_cpu = torch.randn(HIDDEN, INTER, dtype=torch.bfloat16).pin_memory()
+    # 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 _ in range(20):
-        cpu_experts(x, one, ids, {0}, lambda e: (w13_cpu, w2_cpu))
-    cpu_ms = (time.perf_counter() - start) * 1000 / 20
+    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),

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))

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_cli.py

@@ -55,12 +55,33 @@ class Cli(TmpCase):
         execvp.assert_called_once_with("vllm", [
             "vllm", "serve", "/models/m", "--enforce-eager", "--moe-backend", "triton",
             "--scheduling-policy", "priority", "--max-num-seqs", "3",
-            "--kv-cache-memory-bytes", "1073741824", "--port", "8001"])
+            "--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")
@@ -76,3 +97,17 @@ class Cli(TmpCase):
         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

@@ -107,6 +107,23 @@ mode = "auto"
         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'

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

@@ -31,3 +31,64 @@ class KernelMatch(unittest.TestCase):
                                 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_policy.py

@@ -119,3 +119,7 @@ class Prefetch(TmpCase):
             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

@@ -57,6 +57,14 @@ class Simulate(unittest.TestCase):
         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))
@@ -81,3 +89,13 @@ class Tune(unittest.TestCase):
         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

@@ -9,7 +9,7 @@ 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, remap, route_stats,
+    from prain.tiered import (available_experts, cpu_experts, degrade, route_stats, slot_table,
                                 tiered_forward)
 
 HIDDEN, INTER, EXPERTS = 6, 4, 10
@@ -52,8 +52,9 @@ class TieredForward(unittest.TestCase):
             loads.append(expert)
             slots13[slot], slots2[slot] = self.w13[expert], self.w2[expert]
 
-        def kernel(x, w, ids, used):
-            return naive_moe(x, slots13, slots2, w, ids)
+        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]),
@@ -80,6 +81,33 @@ class TieredForward(unittest.TestCase):
         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)
@@ -89,10 +117,10 @@ class TieredForward(unittest.TestCase):
         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_remap_zeroes_routes_outside_the_round(self):
-        weights, slots = remap(torch.ones(1, 3), torch.tensor([[2, 7, -1]]), {7: 4}, 10)
-        self.assertEqual(weights.tolist(), [[0.0, 1.0, 0.0]])
-        self.assertEqual(slots.tolist(), [[0, 4, 0]])
+    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]])
@@ -115,6 +143,8 @@ class TieredForward(unittest.TestCase):
         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)

tests/test_tui.py

@@ -30,6 +30,17 @@ class Tui(unittest.TestCase):
         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")

tests/test_web.py

@@ -1,4 +1,5 @@
 import json
+import os
 import threading
 import unittest
 import urllib.error
@@ -7,7 +8,7 @@ import urllib.request
 import prain.chat as chat
 from prain.config import Draft, Serve
 from prain.web import make_server
-from tests.helpers import StubVllm, wait_for
+from tests.helpers import StubVllm, TmpCase, wait_for
 
 
 def events(response):
@@ -38,7 +39,9 @@ class Web(unittest.TestCase):
         self.url = f"http://127.0.0.1:{self.server.server_port}"
 
     def warm(self, session, draft):
-        chat.post(self.url, "/warm", {"session": session, "messages": [], "draft": draft}).read()
+        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,
@@ -46,7 +49,9 @@ class Web(unittest.TestCase):
 
     def test_page_warm_up_and_chat(self):
         with urllib.request.urlopen(self.url) as response:
-            self.assertIn(b'<textarea id="draft"', response.read())
+            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"}])
@@ -67,9 +72,9 @@ class Web(unittest.TestCase):
 
     def test_no_warm_up_while_another_answer_streams(self):
         self.server.streams[0] = 1  # someone else's answer is streaming
-        self.warm("bob", "typing along")
+        self.assertEqual(self.warm("bob", "typing along"), 204)  # 204: nothing warms up
         self.server.streams[0] = 0
-        self.warm("bob", "now")
+        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"])
 
@@ -119,6 +124,19 @@ class Web(unittest.TestCase):
                                          "<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
@@ -136,3 +154,54 @@ class Web(unittest.TestCase):
             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()