prain: serve with the string chat template content format

prain: personas for the chats; README on quantization levels and personas
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.
2026-11-04 21:50:46 +01:00 · 2026-10-29 23:55:36 +01:00 · 2026-10-29 14:20:23 +01:00 · 2026-10-29 14:20:23 +01:00 · 2026-10-29 08:39:28 +01:00
26 changed files with 2196 additions and 270 deletions
@@ -1,86 +1,261 @@
 # prain
-Serve a mixture-of-experts model with vLLM when its experts do not fit in VRAM. A TOML file places
+prain runs a mixture-of-experts model with vLLM on a GPU that is too small to hold it. A
-each expert of each layer in VRAM (pinned or in a cache slot), in RAM, or in a pack file on disk.
+placement file (TOML) says where each expert of each layer lives: in VRAM, pinned or in a cache
-Missing experts are copied to the GPU or computed on the CPU (`[miss]`). Eviction, heat, prefetch
+slot; in pinned RAM; or in a pack file on disk. When a token needs an expert that is not in VRAM,
-and miss policies are configurable, and `prain tune` picks them from a routing trace. Experts can be
+prain copies it in or computes it on the CPU. Experts can be stored as bf16, int8 or int4.
-served as bf16, int8 or int4 (`quant`).
+
-Target: Gemma 4 26B-A4B (30 layers x 128 experts, 42.5 GiB of bf16 experts) on an 8 GB RX 6600 XT.
+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`: venv with AMD's ROCm SDK wheels for `PRAIN_GPU_ARCH` (default gfx1032)
+- `sdk`: a venv with AMD's ROCm SDK wheels for `PRAIN_GPU_ARCH` (default gfx1032)
- `torch`, `torchvision`: from source with clang and `-march=native` (hours)
+- `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
+- `vllm`: v0.32.0 from source, with gfx1032 added and `patches/vllm-*.patch` applied
- `prain`: editable install and the test suite
+- `prain`: an editable install, then the test suite
-Check the stack: `scripts/smoke.sh MODEL_DIR --attention-backend TRITON_ATTN`
+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;
-Tune vLLM's fused-MoE kernel for the GPU once (~45 min of Triton compiles the first time;
+`scripts/env.sh` points vLLM at the result). It keeps vLLM's default wherever no candidate beats
-`scripts/env.sh` points vLLM at the files). It keeps vLLM's default wherever no candidate beats it:
+it:
    python scripts/tune_moe.py build/moe-configs
-## Serve
+## Serve a model
-    . scripts/env.sh
+`prain serve PLACEMENT MODEL_DIR -- VLLM_FLAGS` validates the placement file and starts
-    prain serve examples/gemma4-26b-a4b-rx6600xt.toml MODEL_DIR -- --served-model-name gemma4 \
+`vllm serve` with prain's plugin. prain adds `--enforce-eager --moe-backend triton
-        --language-model-only --max-model-len 16384 --max-num-batched-tokens 1024 \
+--scheduling-policy priority --chat-template-content-format string`, and it takes
-        --enable-prefix-caching --attention-backend TRITON_ATTN \
+`--max-num-seqs` and `--kv-cache-memory-bytes` from the file's `[serve]` table (it refuses them
-        --enable-auto-tool-choice --tool-call-parser gemma4
+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:
-The first start writes the bf16 pack (42.5 GiB, ~3 min); with `quant = "int8"` or `"int4"` it
+| flag | why |
-then quantizes that pack once on the GPU (~3 min). Once the int4 pack is written, the bf16 files
+|---|---|
-(`layerNNN.w13`, `.w2`, `.done`) are not read and can be deleted; a later bf16 or int8 start
+| `--served-model-name gemma4` | the model name clients ask for |
-writes them again. Chat with `prain tui`, or `prain web --config FILE` for several users with
+| `--language-model-only` | keeps Gemma's vision and audio towers out of VRAM |
-drafts. Tool calls work through vLLM's `gemma4` parser, so coding agents and MCP clients can use
+| `--max-model-len 16384` | the context of one request; `kv_cache_gib` must hold it |
-the OpenAI-compatible API.
+| `--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 |
-## Expert formats
+`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.
-| quant | expert | all experts | weights | VRAM kernel | CPU kernel |
+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
-| bf16 | 11.34 MiB | 42.5 GiB | as trained | vLLM Triton | torch bf16 matmul |
+the GPU at a time, so stop it before starting another quantization.
 | int8 | 5.69 MiB | 21.3 GiB | fp32 scale per row | vLLM Triton (w8a16) | `torch._weight_int8pack_mm` |
 | int4 | 3.01 MiB | 11.3 GiB | bf16 scale per 64, MSE clip | vLLM Triton (GPTQ/AWQ, w4a16) | `torch._weight_int4pack_mm_for_cpu` |
-Measured on the RX 6600 XT behind Thunderbolt (PCIe 3.0 x4, 2.5 GB/s), i9-13950HX (16 CPU
+## Choose a quantization
 threads for experts), 31 GiB RAM, quiet host. Every format got 1.66 GiB of expert VRAM, its own
 placement from `prain tune` and tuned kernels; 4k context, 0.75 GiB KV:
-| | bf16 | int8 | int4 |
+`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"`) |
 |---|---|---|---|
-| experts in VRAM (cache slots) | 120 | 270 | 540 |
+| bf16 | as trained | vLLM Triton | torch bf16 matmul |
-| decode, tok/s (median of 3 x 48 tokens) | 1.04 | 1.42 | 1.62 |
+| int8 | fp32 scale per row | vLLM Triton (w8a16) | `torch._weight_int8pack_mm` |
-| first token, short chat prompt (median) | 22.9 s | 5.0 s | 4.9 s |
+| int4 | bf16 scale per 64 weights, MSE-clipped | prain's on RDNA2, else vLLM Triton (w4a16) | `torch._weight_int4pack_mm_for_cpu` |
 | 3 greedy answers x 64 tokens | 221 s | 143 s | 123 s |
 | greedy answers identical to bf16 | - | 1 of 3 | 1 of 3 |
-The smaller formats win twice: more experts fit in VRAM, and every miss moves fewer bytes over
+## Give the model a persona
 the link and out of the page cache (the int4 pack fits in free RAM, the bf16 one does not). On
 the GPU itself a decode round costs about the same in every format (2.3-2.8 ms). These decode
 numbers predate the dense-layer fix below, which makes every format faster.
-## Dense layers on RDNA2
+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:
-rocBLAS runs bf16 GEMVs at 13-16 GB/s on gfx1032 (RDNA2 has no bf16 dot instructions), and vLLM
+    P=examples/gemma4-26b-a4b-rx6600xt.toml
-uses its own skinny-GEMM kernels only on gfx9/gfx11/gfx12. A decode step spent ~350 ms just
+    prain web --config $P --port 8080 --persona examples/personas/gentoo-admin.md --host 0.0.0.0
-reading the dense weights (lm_head alone: 93 ms). On gfx10 prain sends every unquantized GEMM
+    prain web --config $P --port 8081 --persona examples/personas/pair-programmer.md --host 0.0.0.0
-with at most 4 rows through a Triton GEMV (`prain/gemv.py`, ~210 GB/s; `PRAIN_GEMV=0` turns it
+    prain tui --persona examples/personas/explainer.md
 off). With the agent setup below, decode went from 1.58 to 2.77 tok/s with unchanged answers
 (same quality score). Prefill keeps rocBLAS.
-## Trading quality for speed
+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 agent setup: int4, 16k context, 1 GiB KV and 450 cache slots (1.32 GiB) as
+model. Measured on the int4 agent setup before the int4 kernels; quality is the mean negative
-in the example config, per-layer slot counts from `prain tune`. Quality is the mean
+log-likelihood of the min_weight = 1 answers to the six `scripts/quality.py` prompts (0.11 is
-negative log-likelihood of the min_weight = 1 answers to six chat prompts (`scripts/quality.py`;
+that config scoring its own answers):
 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 |
 |---|---|---|---|---|
@@ -90,38 +265,80 @@ negative log-likelihood of the min_weight = 1 answers to six chat prompts (`scri
 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
+--cache-slots 1`, min_weight 0) broke the model (NLL 9) for 25% more decode speed (measured
-the dense-layer fix).
+before the dense-layer fix).
-## Context length and coding agents
+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.
-Gemma 4 has 25 sliding-window layers (window 1024) and 5 full-attention layers, so KV grows slowly
+## Why it is fast on RDNA2
 with context: per request ~20 KiB per token 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).
-With the example config (int4, 16k, tools) a cold 4.7k-token prompt took 70 s and a 12k-token
+rocBLAS runs bf16 GEMVs at 13-16 GB/s on gfx1032 (RDNA2 has no bf16 dot instructions), and vLLM
-one 176 s: every 1024-token prefill chunk needs nearly all experts, copied over the 2.5 GB/s link.
+uses its own skinny-GEMM kernels only on gfx9/gfx11/gfx12. A decode step spent ~350 ms reading
-Agent turns reuse the prefix cache: turns 2 and 3 of a conversation found 4,704 and 4,736 of
+the dense weights (lm_head alone: 93 ms). On gfx10 prain sends every unquantized GEMM with at
-~4,750 prompt tokens cached and took ~20 s each (24-token answers). vLLM keeps sliding-window
+most 4 rows through a Triton GEMV (`prain/gemv.py`, ~210 GB/s; `PRAIN_GEMV=0` turns it off):
-state only where earlier prompts ended, so the second conversation that shares just the system
+decode went from 1.58 to 2.77 tok/s with unchanged answers. Prefill keeps rocBLAS.
-prompt recomputes it once (176 s); later ones hit (4.5 s).
+
 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 CONFIG MODEL_DIR/config.json --cost cost.json
+    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.4 \
+    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/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, then
+    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  # on another config
+    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.
+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.
-The simulator treats every disk-tier expert as a disk read, so it under-predicts bf16 (0.6 vs 1.04
+## Profile
-tok/s measured, the page cache helps) and over-predicted int4 (2.4 vs 1.62 tok/s).
+
 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
@@ -129,3 +346,14 @@ tok/s measured, the page cache helps) and over-predicted int4 (2.4 vs 1.62 tok/s
  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) |
@@ -23,9 +23,11 @@ pack_dir = "../models/packs/gemma4-26b-a4b"  # bf16 pack (42.5 GiB) written on t
 quant = "int4"                           # then quantized once into an int4 pack (11.3 GiB):
                                         # bf16 | int8 (5.69 MiB/expert) | int4 (3.01 MiB)
 default_tier = "disk"                    # runs on any RAM size; `prain tune` does better
-cache_slots = 15                         # 30 x 15 x 3.01 MiB = 1.32 GiB of VRAM
+cache_slots = 15                         # 30 x 15 x 3.01 MiB = 1.32 GiB of VRAM. The same VRAM
-stats_path = "../models/packs/gemma4-26b-a4b/stats.json"
+                                         # holds 8 int8 or 4 bf16 experts per layer: change
-stats_every = 64                         # MoE forward passes between stats/trace writes
+                                         # 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
@@ -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.
@@ -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.
@@ -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.
@@ -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,
@@ -113,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])
@@ -165,10 +168,16 @@ 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 "
+            ui.add_argument("--host", default="127.0.0.1",
-                                             "[draft] table (default: their defaults)")
+                            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":
@@ -193,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__":
@@ -32,11 +32,11 @@ 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.quant import (GROUP, cpu_matmul, layout, moe_quant_config, quantize, to_cpu_int4)
                         to_gpu_int4)
 from prain.sim import write_trace
 from prain.tiered import draft_rows, tiered_forward
@@ -46,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():
@@ -127,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")
@@ -270,6 +334,7 @@ class TieredRoutedExperts(RoutedExperts):
                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:
@@ -305,11 +370,20 @@ class TieredRoutedExperts(RoutedExperts):
        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 i, (dst, src) in enumerate(zip(dsts, srcs)):
+            for dst, src in zip(dsts, srcs):
-                if i == 0 and self.quant == "int4":  # the pack keeps the CPU kernel's order
+                with profile.span("h2d" if src.device.type == "cpu" else "d2d"):
-                    src = to_gpu_int4(src.to(dst.device, non_blocking=True))
+                    dst[slot].copy_(src, non_blocking=True)
-                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."""
@@ -319,7 +393,8 @@ class TieredRoutedExperts(RoutedExperts):
        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
@@ -331,16 +406,19 @@ class TieredRoutedExperts(RoutedExperts):
        self._ready_events[slot] = event
    def _run_slots(self, x, weights, ids, table, used_slots):
-        """vLLM's Triton fused-MoE kernel on the slot tensors; expert_map=table skips the
+        """The fused-MoE kernel on the slot tensors (vLLM's, or prain's for int4 on RDNA2);
-        routes of experts outside this round."""
+        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 = fused_experts(x, self._slots[0][0], self._slots[1][0], weights, ids,
+        if self.quant == "int4" and gfx10(x.device.index or 0):
-                            activation=self.activation, global_num_experts=table.numel(),
+            out = self._int4_gfx10(x, weights, ids, table)
-                            expert_map=table, quant_config=self._qconfig)
+        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)
@@ -348,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
@@ -377,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)
@@ -399,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,
@@ -416,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:
@@ -15,13 +15,15 @@ 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.autotune(
    configs=[triton.Config({"BLOCK_N": n, "BLOCK_K": k}, num_warps=w)
             for n, k, w in ((4, 512, 8), (8, 512, 8), (16, 256, 4), (32, 256, 2))],
    key=["M", "N", "K"])
@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):
@@ -42,8 +44,7 @@ def _gemv(w_ptr, x_ptr, y_ptr, M, N, K, BLOCK_M: tl.constexpr, BLOCK_N: tl.const
@functools.cache
 def fast_device(index):
    """gfx10 (RDNA2) GPUs, where vLLM has no skinny-GEMM kernels."""
-    name = getattr(torch.cuda.get_device_properties(index), "gcnArchName", "")
+    return gfx10(index) and os.environ.get("PRAIN_GEMV", "1") != "0"
    return name.startswith("gfx10") and os.environ.get("PRAIN_GEMV", "1") != "0"
 def _small(x, weight, bias):
@@ -60,8 +61,8 @@ def linear(x, weight, bias=None):
    x2 = x.reshape(-1, k).contiguous()
    m = x2.shape[0]
    y = torch.empty(m, n, dtype=x.dtype, device=x.device)
-    _gemv[lambda meta: (triton.cdiv(n, meta["BLOCK_N"]),)](
+    _gemv[(triton.cdiv(n, BLOCK_N),)](weight, x2, y, m, n, k, BLOCK_M=triton.next_power_of_2(m),
-        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)
@@ -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])
@@ -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, ...}}}."""
+    """Atomically write {"layers": {"<layer>": {"uses": [...], "vram": [experts in VRAM now],
-    data = {"layers": {str(layer): {"uses": c.uses, **{k: getattr(c, k) for k in COUNTERS}}
+    "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:
@@ -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}")
@@ -3,6 +3,8 @@
 import torch
 import torch.nn.functional as F
 from prain import profile
 HEAT_INDEX = {"count": 0, "weight": 1, "recency": 2}
@@ -119,21 +121,23 @@ def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight"
    raw = None
    skip = miss is not None and miss.min_weight < 1
    if draft is not None or skip:
-        if on_stats is not None:  # the trace records the routing, not what was skipped
+        with profile.span("degrade"):
-            raw = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)
+            if on_stats is not None:  # the trace records the routing, not what was skipped
-        rows = draft if draft is not None else torch.zeros(topk_ids.shape[0], dtype=torch.bool,
+                raw = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)
-                                                           device=topk_ids.device)
+            rows = draft if draft is not None else torch.zeros(
-        need = min_weight
+                topk_ids.shape[0], dtype=torch.bool, device=topk_ids.device)
-        if skip:  # every row may skip; draft rows by the smaller share
+            need = min_weight
-            need = torch.where(rows, min(min_weight, miss.min_weight), miss.min_weight)[:, None]
+            if skip:  # every row may skip; draft rows by the smaller share
-            rows = torch.ones_like(rows)
+                need = torch.where(rows, min(min_weight, miss.min_weight), miss.min_weight)[:, None]
-        available = available_experts(topk_ids, rows, set(cache.pinned) | set(cache.slots),
+                rows = torch.ones_like(rows)
-                                      cache.num_experts)
+            available = available_experts(topk_ids, rows, set(cache.pinned) | set(cache.slots),
-        _, kept = degrade(topk_weights, topk_ids, rows, available, need)
+                                          cache.num_experts)
-        # experts some row copies anyway serve every row for free
+            _, kept = degrade(topk_weights, topk_ids, rows, available, need)
-        available |= torch.bincount(kept[kept >= 0].long(), minlength=cache.num_experts) > 0
+            # experts some row copies anyway serve every row for free
-        topk_weights, topk_ids = degrade(topk_weights, topk_ids, rows, available, need)
+            available |= torch.bincount(kept[kept >= 0].long(), minlength=cache.num_experts) > 0
-    stats = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)  # host sync
+            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)
@@ -150,14 +154,20 @@ def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight"
        nonlocal total
        total = out if total is None else total.to(acc).add_(out)  # .to: a no-op once fp32
-    for slot_of, loads in cache.plan({e: s[index] for e, s in stats.items()}, tokens, cpu):
+    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:
-            load(expert, slot)
+            with profile.span("copy"):
                load(expert, slot)
        if slot_of:  # each round computes only its own routes
-            table = slot_table(slot_of, cache.num_experts, x.device)
+            with profile.span("round"):
-            add(kernel(x, topk_weights, ids, table, list(slot_of.values())))
+                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:
-        after_enqueue(cpu, stats)
+        with profile.span("prefetch"):
            after_enqueue(cpu, stats)
    if cpu:  # runs while the GPU works through its rounds
-        add(cpu_experts(*host, cpu, cpu_weights).to(x.device))
+        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)
@@ -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):
+    def __init__(self, base_url, model, warmup="idle", priority=1, system=None):
-        self.base_url, self.model, self.warmup = base_url, model, warmup
+        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)
@@ -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>
@@ -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
+A small stdlib proxy in front of vLLM. GET / serves the page (web.html). GET /status returns
-{"session", "messages", "draft"} and returns at once. POST /chat takes {"session",
+what the monitor shows: vLLM's counters (/metrics), the GPU (amdgpu sysfs) and prain's expert
-"messages"} and streams server-sent events: vLLM's chunks, plus {"prain": "draft" |
+stats (the placement file's stats_path, with the experts in VRAM). POST /warm takes
-"refine" | "error"} events of its own. Binds to 127.0.0.1; use an SSH tunnel from other
+{"session", "messages", "draft"} and returns at once: 202 when a warm-up starts, 204 when not.
-machines.
+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>
+PAGE = (Path(__file__).parent / "web.html").read_bytes()
-<meta charset="utf-8">
+METRICS = {"vllm:num_requests_running": "running", "vllm:num_requests_waiting": "waiting",
-<title>prain</title>
+           "vllm:kv_cache_usage_perc": "kv_usage", "vllm:prefix_cache_queries_total":
-<style>
+           "prefix_queries", "vllm:prefix_cache_hits_total": "prefix_hits",
-  body { font: 15px/1.45 system-ui, sans-serif; max-width: 52rem; margin: 2rem auto; padding: 0 1rem; }
+           "vllm:prompt_tokens_total": "prompt_tokens",
-  .user { color: #666; margin-top: 1rem; }
+           "vllm:generation_tokens_total": "generation_tokens"}
  .assistant { white-space: pre-wrap; }
  .draft { color: #999; }
  textarea { width: 100%; height: 5rem; font: inherit; margin-top: 1rem; }
  #status { color: #888; font-size: 13px; }
 </style>
 <div id="log"></div>
 <textarea id="draft" autofocus
  placeholder="Experts load while you type. Enter sends, Shift+Enter adds a line."></textarea>
 <div id="status"></div>
 <script>
 const session = Math.random().toString(36).slice(2) + Date.now().toString(36);
 const messages = [];
 const log = document.getElementById("log");
 const draft = document.getElementById("draft");
 const status = document.getElementById("status");
 let streaming = false;
 function show(role, text) {
  const div = document.createElement("div");
  div.className = role;
  div.textContent = text;
  log.appendChild(div);
  return div;
 }
-draft.addEventListener("input", () => {
+def vllm_metrics(base_url):
-  if (!streaming && draft.value.trim())
+    """The METRICS of vLLM's Prometheus text, summed over label sets; None when unreachable."""
-    fetch("/warm", {method: "POST", body: JSON.stringify({session, messages, draft: draft.value})});
+    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()):
+def serve(base_url, port, warmup="idle", priority=1, serve=Serve(), draft=Draft(), host="127.0.0.1",
-    server = make_server(base_url, port, warmup=warmup, priority=priority, serve=serve, draft=draft)
+          stats_path=None, system=None, persona=None):
-    print(f"prain web: http://127.0.0.1:{server.server_port}  (vLLM at {base_url}; "
+    server = make_server(base_url, port, host=host, warmup=warmup, priority=priority, serve=serve,
-          f"{draft.full_requests} full of {serve.max_requests} parallel answers)")
+                         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()
@@ -17,3 +17,6 @@ build-backend = "setuptools.build_meta"
 [tool.setuptools]
 packages = ["prain"]
 [tool.setuptools.package-data]
 prain = ["web.html"]
@@ -33,10 +33,13 @@ def post(url, path, body):
 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]
-    return choice["text"], [int(t.split(":")[1]) for t in choice["logprobs"]["tokens"]]
+    logprobs = choice["logprobs"]
    return (choice["text"], [int(t.split(":")[1]) for t in logprobs["tokens"]],
            logprobs["token_logprobs"])
 def nll(url, model, ids, answer):
@@ -56,6 +59,9 @@ def main():
    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:
@@ -65,7 +71,7 @@ def main():
            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)
+            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:
@@ -73,15 +79,20 @@ def main():
        return
    with open(args.ref) as f:
        refs = json.load(f)
-    scores, same = [], []
+    scores, same, own = [], [], []
    for ref in refs:
-        scores.append(nll(args.url, args.model, ref["ids"], ref["answer"]))
+        if not args.greedy_only:
-        text, answer = greedy(args.url, args.model, ref["ids"], len(ref["answer"]))
+            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),
+        print(json.dumps({"prompt": ref["prompt"][:40],
-                          "same_tokens": same[-1], "answer": text[:100]}), flush=True)
+                          "nll": round(scores[-1], 3) if scores else None,
-    print(json.dumps({"mean_nll": round(sum(scores) / len(scores), 3), "same_tokens": same,
+                          "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))}))
@@ -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):
@@ -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")
@@ -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()
@@ -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
@@ -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))
@@ -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")
@@ -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)
@@ -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()
Author	SHA1	Message	Date
fabolous005	1e3a412563	prain: serve with the string chat template content format	2026-11-04 21:50:46 +01:00
fabolous005	a67238f66b	prain: personas for the chats; README on quantization levels and personas 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.	2026-10-29 23:55:36 +01:00
fabolous005	a14e3ade0b	prain: web UI with a model monitor A clear chat page (markdown, code blocks with copy, stop, new chat, kept in the browser, drafts shown as such) and a monitor sidebar, a sheet on phones: tokens/s, a live map of the 30 x 128 experts (in VRAM, used now, copied in), VRAM hit rate, copies, context, prompt cache, GPU VRAM/load/temperature/power/fan. GET /status merges vLLM /metrics, amdgpu sysfs and the stats file, which now lists the experts in VRAM. /warm answers 202 when a warm-up starts. README: the int4 kernels, their A/B, 2048-token chunks as a trade, prompt logprobs limit.	2026-10-29 14:20:23 +01:00
fabolous005	d1ccc5c294	prain: int4 MoE on RDNA2: GEMV-style decode kernel, fp16 rounds, staged expert copies Measured on the live 16k agent setup (A/B with runtime flags, greedy answers vs reference): decode 3.49 -> 5.24 tok/s, first token 4.3 -> 3.3 s, 2.3k-token prompt 27.5 -> 20.6 s, same answers. Dropped after measuring: route stats on the CPU, a Triton router GEMV, pregate prefetch from disk (all within noise or slower). The GEMV keeps one config instead of autotuning: Triton benchmarks with a 256 MiB buffer that torch keeps cached, which left 8 MiB of VRAM and made lazy kernel loads fail (hipErrorNoBinaryForGpu). prain also returns torch cached blocks after materializing and whenever the driver has < 128 MiB free. quality.py score --greedy-only skips prompt logprobs, which OOM the 16k config.	2026-10-29 14:20:23 +01:00
fabolous005	0f3911afb5	prain: span profiler with flame graphs (PRAIN_PROFILE: engine step, every model module, MoE route/plan/copy/round, call counts; python -m prain.profile); prain web --host	2026-10-29 08:39:28 +01:00