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 · 2026-10-28 22:07:25 +01:00
44 changed files with 3302 additions and 296 deletions
@@ -0,0 +1,359 @@
 # prain
 prain runs a mixture-of-experts model with vLLM on a GPU that is too small to hold it. A
 placement file (TOML) says where each expert of each layer lives: in VRAM, pinned or in a cache
 slot; in pinned RAM; or in a pack file on disk. When a token needs an expert that is not in VRAM,
 prain copies it in or computes it on the CPU. Experts can be stored as bf16, int8 or int4.
 It is tuned for one machine: Gemma 4 26B-A4B (30 layers x 128 experts, 42.5 GiB of bf16 experts)
 on an 8 GB Radeon RX 6600 XT (RDNA2, gfx1032) in a Thunderbolt enclosure. With int4 experts it
 answers at about 5 tokens/s, with 16k tokens of context and tool calls. Clients talk to it
 through vLLM's OpenAI-compatible API, a browser chat with a live monitor, or a terminal chat. Both
 chats start moving the experts a message needs into VRAM while it is still being typed.
 ## Quick start
 After the build below, start the model with the example placement (int4 experts):
    . scripts/env.sh
    M=models/gemma-4-26B-A4B-it-Uncensored-MAX
    FLAGS="--served-model-name gemma4 --language-model-only --max-model-len 16384
           --max-num-batched-tokens 1024 --enable-prefix-caching --attention-backend TRITON_ATTN
           --enable-auto-tool-choice --tool-call-parser gemma4"
    prain serve examples/gemma4-26b-a4b-rx6600xt.toml $M -- $FLAGS
 The model is ready when `curl -sf localhost:8000/health && echo ready` prints `ready`. Then, in a
 second terminal:
    . scripts/env.sh
    prain web --config examples/gemma4-26b-a4b-rx6600xt.toml --host 0.0.0.0 \
        --persona examples/personas/gentoo-admin.md
 and open http://HOST:8080. `--host 0.0.0.0` makes the page reachable from other machines; it has
 no login. Coding agents and other OpenAI clients use http://HOST:8000/v1 with the model name
 `gemma4`.
 ## Build (GPU host, user space)
    scripts/build.sh wrappers sdk fetch torch torchvision vllm prain
 - `sdk`: a venv with AMD's ROCm SDK wheels for `PRAIN_GPU_ARCH` (default gfx1032)
 - `torch`, `torchvision`: built from source with clang and `-march=native` (hours)
 - `vllm`: v0.32.0 from source, with gfx1032 added and `patches/vllm-*.patch` applied
 - `prain`: an editable install, then the test suite
 Check the stack with `scripts/smoke.sh MODEL_DIR --attention-backend TRITON_ATTN`. Then tune
 vLLM's fused-MoE kernel for the GPU once (about 45 minutes of Triton compiles the first time;
 `scripts/env.sh` points vLLM at the result). It keeps vLLM's default wherever no candidate beats
 it:
    python scripts/tune_moe.py build/moe-configs
 ## Serve a model
 `prain serve PLACEMENT MODEL_DIR -- VLLM_FLAGS` validates the placement file and starts
 `vllm serve` with prain's plugin. prain adds `--enforce-eager --moe-backend triton
 --scheduling-policy priority --chat-template-content-format string`, and it takes
 `--max-num-seqs` and `--kv-cache-memory-bytes` from the file's `[serve]` table (it refuses them
 after `--`). vLLM would otherwise pick the `openai` content format for Gemma 4, whose template
 then prints a system prompt as a Python list (`[{'type': 'text', 'text': '…\n…'}]`, newlines
 escaped); a flag after `--` still overrides it. The flags in `$FLAGS` above:
 | flag | why |
 |---|---|
 | `--served-model-name gemma4` | the model name clients ask for |
 | `--language-model-only` | keeps Gemma's vision and audio towers out of VRAM |
 | `--max-model-len 16384` | the context of one request; `kv_cache_gib` must hold it |
 | `--max-num-batched-tokens 1024` | the prefill chunk (see [Trade-offs](#trade-offs)) |
 | `--enable-prefix-caching` | reuses earlier turns, personas and warm-ups |
 | `--attention-backend TRITON_ATTN` | Triton attention, which runs on RDNA2 |
 | `--enable-auto-tool-choice --tool-call-parser gemma4` | tool calls for agents and MCP clients |
 `examples/gemma4-26b-a4b-rx6600xt.toml` is commented line by line. It keeps all experts in the
 pack on disk, served through the page cache, and gives each layer 15 cache slots in VRAM (1.32
 GiB), plus 1 GiB of KV cache for 16k tokens of context. `prain check PLACEMENT
 MODEL_DIR/config.json` prints what a placement costs in VRAM, RAM and disk before anything
 loads.
 The first start writes the bf16 pack (42.5 GiB, about 3 minutes); an int8 or int4 placement then
 converts it once on the GPU (about 3 minutes). Stop the server with Ctrl-C. Only one model fits on
 the GPU at a time, so stop it before starting another quantization.
 ## Choose a quantization
 `quant` in the placement file picks the expert format. Smaller experts win twice: more of them
 fit in VRAM, and every miss moves fewer bytes over the 2.5 GB/s Thunderbolt link and out of the
 page cache (on a 31 GiB host the int4 and int8 packs fit in free RAM, the bf16 one does not).
 They cost some accuracy:
 | quant | expert | pack on disk | slots per layer in 1.32 GiB | decode | first token, short prompt | first token, 2.3k-token prompt | answers equal to bf16 |
 |---|---|---|---|---|---|---|---|
 | bf16 | 11.34 MiB | 42.5 GiB | 4 | 0.90 tok/s | 14.7 s | 69.5 s | reference |
 | int8 | 5.69 MiB | 21.3 GiB | 8 | 2.19 tok/s | 5.9 s | 32.3 s | 4 of 6 (215 of 288 tokens) |
 | int4 | 3.01 MiB | 11.3 GiB | 15 | 5.19 tok/s | 3.3 s | 20.5 s | 0 of 6 (94 of 288 tokens) |
 Measured with the example placement on the RX 6600 XT behind Thunderbolt (PCIe 3.0 x4), an
 i9-13950HX and 31 GiB of RAM, changing only `quant` and `cache_slots`. Decode is the median of
 3 x 48 tokens. "Equal" counts how many of the six `scripts/quality.py` prompts got the same 48
 greedy tokens as bf16; the brackets count the tokens that matched before each answer's first
 difference.
 int4 is the default: it decodes 2.4 times as fast as int8 and reads a long prompt in two thirds
 of the time. Its answers are worded differently from bf16's from the first tokens on. Five of
 the six meant the same; the sixth, a Gentoo question, named a tool that does not exist
 ("oneway-way") where bf16 and int8 named `revdep-rebuild`. int8 stays closest to bf16 at less
 than half the speed of int4. bf16 is mostly the reference for quality checks: its 4 slots per
 layer are fewer than the 8 experts each token uses, so every decode step runs in several rounds
 (`prain check` warns about this).
 To launch int8 instead of int4, copy the example and change two lines. Relative paths in a
 placement file start at the file's directory, so keep the copy next to the example:
    cp examples/gemma4-26b-a4b-rx6600xt.toml examples/int8.toml
    sed -i -e 's/^quant = "int4"/quant = "int8"/' -e 's/^cache_slots = 15 /cache_slots = 8 /' \
        examples/int8.toml
    prain check examples/int8.toml $M/config.json     # VRAM 1.33 GiB, as before
    prain serve examples/int8.toml $M -- $FLAGS
 For bf16, set `cache_slots = 4`. Keep slots x expert size x 30 layers near 1.32 GiB. On the
 example placement vLLM reports "Model loading took 6.29 GiB" (the dense weights, 1.32 GiB of
 expert slots and prain's buffers), the KV cache takes 1 GiB (`kv_cache_gib`), and about 0.55 GiB
 of the card's 7.98 GiB stays free for activations and kernels. prain logs that free VRAM at
 startup ("MiB of VRAM free for kernels and activations", before the KV cache takes its share).
 More VRAM for experts means less for context, and the other way round.
 `prain tune` fits per-layer slot counts and policies to a routing trace, which the example records
 while it serves (`trace_path`). The tuned int4 placement of the live agent setup (450 slots,
 12-18 per layer) decodes at 5.2 tok/s. For another format:
    prain tune models/packs/gemma4-26b-a4b/trace.jsonl.gz $M/config.json --quant int8 \
        --expert-vram-gib 1.32 --ram-gib 0 --pack-dir models/packs/gemma4-26b-a4b > int8-tuned.toml
 Every pack stays on disk once written: `layerNNN.w13`, `.w2` and `.done` for bf16,
 `layerNNN.int8.*` and `layerNNN.int4.*` for the others. Delete the formats you do not use; the
 next start of that format writes its pack again, and int8 or int4 need the bf16 pack for that.
 How the formats are stored and computed:
 | quant | weights | GPU kernel | CPU kernel (`[miss] mode = "cpu"`) |
 |---|---|---|---|
 | bf16 | as trained | vLLM Triton | torch bf16 matmul |
 | int8 | fp32 scale per row | vLLM Triton (w8a16) | `torch._weight_int8pack_mm` |
 | int4 | bf16 scale per 64 weights, MSE-clipped | prain's on RDNA2, else vLLM Triton (w4a16) | `torch._weight_int4pack_mm_for_cpu` |
 ## Give the model a persona
 A persona is a short text file that `prain web` and `prain tui` send as the system message at
 the start of every conversation. It tells the model its task, how to handle questions, what its
 environment is and how to use it. One running model can serve several personas at once, one
 chat front end each:
    P=examples/gemma4-26b-a4b-rx6600xt.toml
    prain web --config $P --port 8080 --persona examples/personas/gentoo-admin.md --host 0.0.0.0
    prain web --config $P --port 8081 --persona examples/personas/pair-programmer.md --host 0.0.0.0
    prain tui --persona examples/personas/explainer.md
 The page shows the persona's name, its file name, next to the model's. A new persona needs no
 model restart: edit the file and restart `prain web` or `prain tui`. Without `--persona` the
 model gets no system message.
 A persona has four parts. `examples/personas/gentoo-admin.md`, shortened:
    You help one person run their Gentoo Linux desktop.
    ## Task
    Answer questions about this machine and fix problems together with its owner: installing and
    updating packages, USE flags, services, the kernel, drivers, disks and the Hyprland desktop.
    ## How to answer
    - Lead with the command or the fix, then explain it in a few sentences.
    - When the answer depends on this machine (USE flags, profile, kernel, installed versions),
      ask for the output of a specific command instead of guessing.
    ## Environment
    - Gentoo Linux with Portage and OpenRC. There is no systemd and no journalctl.
    - You cannot run anything yourself. The owner types your commands into a root shell and
      pastes the output back into the chat.
    ## How to use it
    - Preview before changing: `emerge --pretend --verbose PKG`, then `emerge --ask PKG`.
 The other examples are `pair-programmer.md` (reviews and debugs pasted Python and shell code)
 and `explainer.md` (plain answers for beginners, often read on a phone). When you write one:
 - Address the model as "you" and state what it cannot know: the system, the tools at hand, and
  what it can and cannot do itself.
 - Say how to handle missing information, for example by asking for a specific command's output
  instead of guessing.
 - Keep it short. A new persona text is prefilled once, before its first answer: the examples
  (241-467 tokens) added 2.0-3.2 s to that answer on the int4 agent setup, and at most 0.3 s to
  the next conversations, which find it in vLLM's prefix cache. The chats also send it with
  their warm-ups while you type.
 Agents and other API clients send their own system message. To give one a persona, put the
 file's text there:
    jq -n --rawfile p examples/personas/explainer.md \
        '{model: "gemma4", messages: [{role: "system", content: $p},
                                      {role: "user", content: "What is a page cache?"}]}' |
        curl -s localhost:8000/v1/chat/completions -H "Content-Type: application/json" -d @-
 ## Chat in the browser
    prain web --config PLACEMENT [--persona FILE] [--host 0.0.0.0] [--port 8080]
 The page renders Markdown answers with code blocks you can copy. It has a stop button and "New
 chat", and it keeps the conversation in the browser until you start a new one. Enter sends and
 Shift+Enter adds a line; on a phone, Return adds a line and the button sends.
 While you type, the page sends the unfinished message as a warm-up: vLLM prefills it and prain
 moves the experts it routes to into VRAM, so the answer starts sooner. "Loading experts for your
 message" under the box shows that one runs. `--warmup idle` (the default) warms up only while no
 answer streams, `always` also during one, `off` never.
 Several people can chat at once. `[serve] max_requests` answers run in parallel, and
 `[draft] full_requests` of them with all their experts. Later ones start at once as drafts from
 the experts already in VRAM and are regenerated with all experts when the GPU is free (`[draft]
 refine`); the page shows a draft in grey until then.
 The monitor beside the chat (behind the status pill on a phone) shows:
 - tokens per second, now or of the last answer, and a line for the last minute
 - the expert map: one row per layer, one column per expert. Violet experts are in VRAM, orange
  to yellow ones were used in the last seconds, cyan ones were just copied in. It needs
  `stats_path` in the placement file; `stats_every = 8` (forward passes) refreshes it about every
  1.5 s while the model writes.
 - VRAM hits (routed experts that were already in VRAM, last minute) and copies per second
 - context in use (the KV cache), prompt cache reuse, and requests running and waiting
 - the GPU's VRAM, load, temperature, power and fan (from amdgpu's sysfs files)
 ## Chat in the terminal
    prain tui [--persona FILE] [--url http://127.0.0.1:8000]
 The same warm-ups as the page. Enter sends, Esc quits.
 ## Coding agents and long context
 Gemma 4 has 25 sliding-window layers (window 1024) and 5 full-attention layers, so its KV cache
 grows slowly with context: ~20 KiB per token per request, plus ~0.6 GiB for the sliding layers,
 which also hold the tokens in flight (2 x `--max-num-batched-tokens`). `kv_cache_gib = 1.0`
 holds 18,169 tokens, 1.11 requests of 16k (measured).
 Every 1024-token prefill chunk needs nearly all experts, copied over the 2.5 GB/s link: a fresh
 2.3k-token prompt takes about 20 s. Before the int4 kernels, which cut that by a quarter, a cold
 4.7k-token prompt took 70 s and a 12k-token one 176 s. Agent turns reuse the prefix cache: turns
 2 and 3 of a conversation found 4,704 and 4,736 of ~4,750 prompt tokens cached and took ~20 s
 each (24-token answers). vLLM keeps sliding-window state only where earlier prompts ended, so the
 second conversation that shares only the system prompt recomputes it once (176 s); later ones hit
 (4.5 s).
 ## Trade-offs
 `[miss] min_weight` (default 1) is the share of each token's router weight that must be
 computed. Below 1, a token skips its lightest experts that are not in VRAM once it has that
 share, and its other routes are scaled up to its full weight: fewer copies, a slightly different
 model. Measured on the int4 agent setup before the int4 kernels; quality is the mean negative
 log-likelihood of the min_weight = 1 answers to the six `scripts/quality.py` prompts (0.11 is
 that config scoring its own answers):
 | min_weight | answer NLL | decode tok/s | first token, short prompt | first token, 2.5k-token prompt |
 |---|---|---|---|---|
 | 1.0 | 0.11 | 2.77 | 4.6 s | 28.5 s |
 | 0.6 | 0.19 | 3.15 | 2.7 s | 23.0 s |
 | 0.4 | 0.40 | 3.36 | 2.0 s | 19.9 s |
 At 0.6 the answers stay correct but take other words; at 0.4 they show slips ("Hereon any of the
 following one-liners"). Pinning only the hottest experts and never copying (`prain plan
 --cache-slots 1`, min_weight 0) broke the model (NLL 9) for 25% more decode speed (measured
 before the dense-layer fix).
 Bigger prefill chunks trade decode speed for prompt speed: at 16k, `--max-num-batched-tokens 2048`
 needs `kv_cache_gib = 1.3` (the sliding layers hold two chunks), which leaves 348 slots instead
 of 450. A 2.3k-token prompt then took 17.4 s instead of 20.6 s, and decode ran at 4.5 instead of
 5.2 tok/s. Long agent prompts gain more: half the chunks, half the expert copies.
 ## Why it is fast on RDNA2
 rocBLAS runs bf16 GEMVs at 13-16 GB/s on gfx1032 (RDNA2 has no bf16 dot instructions), and vLLM
 uses its own skinny-GEMM kernels only on gfx9/gfx11/gfx12. A decode step spent ~350 ms reading
 the dense weights (lm_head alone: 93 ms). On gfx10 prain sends every unquantized GEMM with at
 most 4 rows through a Triton GEMV (`prain/gemv.py`, ~210 GB/s; `PRAIN_GEMV=0` turns it off):
 decode went from 1.58 to 2.77 tok/s with unchanged answers. Prefill keeps rocBLAS.
 vLLM's int4 MoE kernel pads every token to a 16-row tile and multiplies in bf16: a decode round
 took 2.7 ms for one token. On gfx10 prain runs int4 rounds of up to 4 tokens through its own
 GEMV-style kernel (`prain/kernels.py`, 0.25 ms) and larger rounds through vLLM's kernel in fp16,
 twice as fast as bf16 there (a round that overflows fp16 is redone in bf16). Every int4 expert
 copy goes through a pinned staging buffer: one DMA and one unpack kernel instead of four pageable
 copies and a ten-kernel repack, asynchronous from every tier.
 A/B on the live agent setup (int4, 16k context, 1 GiB KV, 450 tuned slots); quality is the
 greedy answers to the six `quality.py` prompts (identical to the reference, mean logprob of their
 tokens):
 | | decode tok/s | first token, short | first token, 2.3k tokens | identical | mean logprob |
 |---|---|---|---|---|---|
 | before | 3.49 | 4.3 s | 27.5 s | 4 of 6 | -0.104 |
 | decode kernel | 4.73 | 4.2 s | 27.3 s | 4 of 6 | -0.093 |
 | staged copies | 3.67 | 3.7 s | 24.5 s | 4 of 6 | -0.098 |
 | fp16 rounds | 3.84 | 4.0 s | 23.3 s | 4 of 6 | -0.102 |
 | all three | 5.24 | 3.3 s | 20.6 s | 5 of 6 | -0.104 |
 Repeated runs of one config differ by up to 5% in decode speed and 1% in prefill. Measured and
 dropped: routing statistics on the CPU and a Triton GEMV for the fp32 router (within noise), and
 `[prefetch] mode = "pregate"`, which can now prefetch disk-tier experts too (4.9 vs 5.2 tok/s:
 the link is already busy with demand copies).
 ## Tune for a machine
    python scripts/measure.py PACK_DIR/layer029.w13 --quant int4 --threads 16 > cost.json  # server stopped
    prain sim TRACE PLACEMENT MODEL_DIR/config.json --cost cost.json
    prain tune TRACE MODEL_DIR/config.json --cost cost.json --quant int4 --expert-vram-gib 1.32 \
        --ram-gib 0 --pack-dir PACK_DIR --cpu-threads 16 > tuned.toml
    python scripts/bench.py http://127.0.0.1:8000 [--context FILES...]      # speed
    python scripts/quality.py http://127.0.0.1:8000 make ref.json           # reference answers
    python scripts/quality.py http://127.0.0.1:8000 score ref.json --greedy-only   # another config
 Judge quality on chat answers, not raw-text perplexity: this instruction-tuned model scores raw
 code at a negative log-likelihood of 5-9 per token on a correct engine. The simulator treats
 every disk-tier expert as a disk read, so it under-predicted bf16 (0.6 vs 1.04 tok/s measured;
 the page cache helps) and over-predicted int4 (2.4 vs 1.62 tok/s), both before the dense-layer
 fix and the int4 kernels.
 ## Profile
 Serve with `PRAIN_PROFILE=/abs/dir`. Writing a number N into `dir/start` records the next N
 engine steps: spans around the engine step, every module of the model and prain's MoE phases
 (route, plan, copy: memcpy/h2d/unpack for int4, round), each synced with the GPU. `python -m
 prain.profile dir` writes `decode.svg`/`prefill.svg` flame graphs and prints self time and calls
 per step. The syncs remove the CPU/GPU overlap: a profiled decode step took 395 ms where an
 unprofiled one takes 307 ms.
 ## Limits
 - Requests with `prompt_logprobs` (perplexity tools) materialize [prompt tokens x 262k vocab]
  fp32 logits that vLLM's memory plan does not cover: on the 16k setup even ~300-token prompts
  crashed the engine. `scripts/quality.py score --greedy-only` skips them.
 - Pinning host memory was very slow on the Thunderbolt host (minutes per GiB, cause unproven) and
  pinned copies were only 4% faster than pageable ones, so the example has no RAM tier and lets
  the page cache hold the pack.
 - `prain web` has no login. Bind it to 127.0.0.1 (the default) on untrusted networks.
 ## vLLM patches
 - `vllm-0.32-rocm-no-cuda-families`: gfx1032 reports capability 10.3, and vLLM's CUDA "family"
  checks took it for Blackwell (B200 attention tiles: 128 KiB of LDS on a 64 KiB GPU).
 - `vllm-0.32-gemma4-fit-8gb`: Gemma 4 shares `embed_tokens` as `lm_head` instead of allocating a
  second 1.4 GiB copy while loading, and sizes rope caches by `max_model_len` (saves ~0.4 GiB).
 ## Repository
 | path | what |
 |---|---|
 | `prain/` | the vLLM plugin (`plugin.py`, `tiered.py`, `experts.py`, `drafts.py`), RDNA2 kernels (`kernels.py`, `gemv.py`), placement files and slot policies (`config.py`, `policy.py`), quantization (`quant.py`), simulator and tuner (`sim.py`), the chats (`chat.py`, `web.py` + `web.html`, `tui.py`), the profiler (`profile.py`) and `cli.py` |
 | `examples/` | the RX 6600 XT placement file and `personas/` |
 | `scripts/` | build, environment, kernel tuning, benchmarks and quality checks |
 | `patches/` | the vLLM patches above |
 | `tests/` | `python -m unittest discover -s tests`; GPU tests skip without a GPU |
 | `models/`, `build/` | model weights and packs, build outputs and logs (not in git) |
@@ -1,6 +1,12 @@
-# prain placement: prithivMLmods/gemma-4-26B-A4B-it-Uncensored-MAX on an 8 GB RX 6600 XT
+# prain placement: prithivMLmods/gemma-4-26B-A4B-it-Uncensored-MAX on an 8 GB RX 6600 XT, set up
-# Model: 30 MoE layers x 128 experts, top-8 routing, 11.34 MiB per expert (BF16).
+# for a coding agent: int4 experts, 16k context, tool calls.
-# VRAM: 8 GiB - 4.45 GiB dense weights - 0.75 GiB KV cache - ~0.9 GiB runtime = ~1.7 GiB experts.
+# Model: 30 MoE layers x 128 experts, top-8 routing, 3.01 MiB per int4 expert (11.34 MiB bf16).
 # VRAM: 4.46 GiB dense weights + 30 x 15 x 3.01 MiB = 1.32 GiB of expert slots + 1.0 GiB KV cache.
 # Serve with --max-model-len 16384 --attention-backend TRITON_ATTN (ROCM_ATTN needs more LDS than
 # RDNA2 has) --enable-auto-tool-choice --tool-call-parser gemma4.
 # RAM tier: pinning host memory took minutes per GiB on the Thunderbolt eGPU host (cause unproven)
 # and pinned copies were only 4% faster than pageable ones, so this file keeps no RAM tier; the
 # kernel page cache holds the pack instead (the int4 pack, 11.3 GiB, fits in free RAM).
 #
 # Tiers:  vram = pinned in a VRAM slot for the whole run
 #         ram  = pinned host memory, copied into a cache slot when a token needs it
@@ -10,19 +16,23 @@
 # Layer selectors: "*", "12", "0-9,20".   Expert lists: "*", "0-15,32", [1, 2, 3].
 # Check what a file costs:   prain check <this file> <model config.json>
 # Compare variants offline:  prain sim <trace> <this file> <model config.json>
-#                            prain tune <trace> <model config.json> --expert-vram-gib 1.7 ...
+#                            prain tune <trace> <model config.json> --quant int4 --expert-vram-gib 1.4 ...
 # Relative paths start at this file's directory; "~" works too.
-pack_dir = "../models/packs/gemma4-26b-a4b"  # 42.5 GiB, written on the first run
+pack_dir = "../models/packs/gemma4-26b-a4b"  # bf16 pack (42.5 GiB) written on the first run,
 quant = "int4"                           # then quantized once into an int4 pack (11.3 GiB):
                                         # bf16 | int8 (5.69 MiB/expert) | int4 (3.01 MiB)
 default_tier = "disk"                    # runs on any RAM size; `prain tune` does better
-cache_slots = 5                          # 30 x 5 x 11.34 MiB = 1.66 GiB of VRAM
+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
-eviction = "lfu"      # lru | lfu (decayed frequency) | gdsf (frequency x refetch cost)
+eviction = "gdsf"     # lru | lfu (decayed frequency) | gdsf (frequency x refetch cost)
-heat = "weight"       # what makes an expert hot: count (tokens) | weight (router mass) |
+heat = "count"        # what makes an expert hot: count (tokens) | weight (router mass) |
                      # recency (router mass, later tokens weigh more: good for warm-ups)
 lfu_half_life = 64    # forwards until an LFU score halves
 disk_cost = 4.0       # gdsf: a disk expert costs this many RAM experts to refetch
@@ -39,16 +49,23 @@ min_prob = 0.05       # skip guesses below this router probability
 [miss]                # an expert that is not in VRAM
 mode = "copy"         # copy | cpu (always compute it on the CPU) | auto (CPU if it has
-cpu_max_tokens = 1    #   at most cpu_max_tokens tokens, else copy)
+cpu_max_tokens = 1    #   at most cpu_max_tokens tokens, else copy). Measured on the Thunderbolt
-cpu_threads = 4       # CPU threads of the vLLM worker (sets OMP_NUM_THREADS)
+                      #   eGPU host, per int4 expert: a copy ~1.3 ms + disk read, the CPU
-admit = 1             # cpu/auto: copy this many CPU-run experts into VRAM afterwards
+                      #   0.8-1.3 ms per token. `prain tune` picked copy for int4 (a copied
                      #   expert stays for later tokens) and auto <= 2 tokens for bf16 (copy 4.8 ms)
 cpu_threads = 16      # CPU threads of the vLLM worker (sets OMP_NUM_THREADS)
 admit = 0             # cpu/auto: copy this many CPU-run RAM-tier experts into VRAM afterwards
 min_weight = 1.0      # share of each token's router weight that must be computed. Below 1 a token
                      #   skips experts not in VRAM once it has that share: faster, dumber.
                      #   Measured here: 0.6 decodes 14% faster (answer NLL 0.11 -> 0.19), 0.4
                      #   21% (0.40, visible slips)
 [serve]               # parallel requests (`prain serve` passes these to vLLM)
 max_requests = 2      # requests vLLM runs at once (--max-num-seqs); warm-ups count too
-kv_cache_gib = 0.75   # KV cache VRAM shared by them (--kv-cache-memory-bytes). vLLM logs
+kv_cache_gib = 1.0    # KV cache VRAM shared by them (--kv-cache-memory-bytes). vLLM logs
-                      #   "Maximum concurrency for 4096 tokens per request: N.NNx" at start;
+                      #   "Maximum concurrency for 16384 tokens per request: 1.11x" with 1.0 GiB.
-                      #   below max_requests, long requests wait or get preempted. Every GiB
+                      #   Gemma 4 needs ~20 KiB per token of context plus ~0.6 GiB for its
-                      #   here is 90 expert cache slots less.
+                      #   sliding-window layers; every GiB here is 340 int4 expert slots less.
 [draft]               # answers that arrive while others use the GPU (`prain web --config`)
 full_requests = 1     # answers that run with all their experts; later ones run at once as
@@ -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.
@@ -0,0 +1,52 @@
 diff --git a/vllm/model_executor/models/gemma4.py b/vllm/model_executor/models/gemma4.py
 index d1498f1..745d1d3 100644
 --- a/vllm/model_executor/models/gemma4.py
 +++ b/vllm/model_executor/models/gemma4.py
@@ -26,7 +26,7 @@ import torch
 from torch import nn
 from vllm.compilation.decorators import support_torch_compile
 -from vllm.config import CacheConfig, VllmConfig
 +from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
 from vllm.config.utils import getattr_iter
 from vllm.distributed import (
     get_pp_group,
@@ -491,9 +491,12 @@ class Gemma4Attention(nn.Module):
                     f"{kv_shared_layer_index}.self_attn.attn"
                 )
 +        # Gemma 4's rope types (default, proportional) do not depend on max_position, so
 +        # size the cos/sin caches for max_model_len: 262144 positions take ~0.4 GiB.
 +        max_model_len = get_current_vllm_config().model_config.max_model_len
         self.rotary_emb = get_rope(
             self.head_dim,
 -            max_position=max_position_embeddings,
 +            max_position=min(max_position_embeddings, max_model_len),
             rope_parameters=rope_parameters,
             is_neox_style=True,
         )
@@ -1444,14 +1447,17 @@ class Gemma4ForCausalLM(
             prefix=maybe_prefix(prefix, "model"),
         )
 -        self.lm_head = ParallelLMHead(
 -            config.vocab_size,
 -            config.hidden_size,
 -            quant_config=quant_config,
 -            prefix=maybe_prefix(prefix, "lm_head"),
 -        )
         if config.tie_word_embeddings:
 -            self.lm_head = self.lm_head.tie_weights(self.model.embed_tokens)
 +            # Share embed_tokens: a ParallelLMHead tied after construction would hold the
 +            # vocab matrix (1.4 GiB) twice while loading, which overflows 8 GB GPUs.
 +            self.lm_head = self.model.embed_tokens
 +        else:
 +            self.lm_head = ParallelLMHead(
 +                config.vocab_size,
 +                config.hidden_size,
 +                quant_config=quant_config,
 +                prefix=maybe_prefix(prefix, "lm_head"),
 +            )
         self.logits_processor = LogitsProcessor(
             config.vocab_size,
@@ -0,0 +1,17 @@
 diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py
 index 807e571..bbc642a 100644
 --- a/vllm/platforms/rocm.py
 +++ b/vllm/platforms/rocm.py
@@ -983,6 +983,12 @@ class RocmPlatform(Platform):
         major, minor = torch.cuda.get_device_capability(device_id)
         return DeviceCapability(major=major, minor=minor)
 +    @classmethod
 +    def is_device_capability_family(cls, capability: int, device_id: int = 0) -> bool:
 +        # Families are a CUDA notion, and gfx majors collide with them (gfx1032 reads as
 +        # 10.3, i.e. Blackwell sm_103): no CUDA family check may match a ROCm GPU.
 +        return False
 +
     @classmethod
     @with_amdsmi_context
     def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
@@ -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,
@@ -7,19 +7,19 @@ import os
 import sys
 from dataclasses import asdict
-from prain.config import (GIB, Miss, Policy, Prefetch, expert_bytes, format_ids, load_config,
+from prain.config import (GIB, QUANTS, Miss, Policy, Prefetch, expert_bytes, format_ids, load_config,
                            moe_dims, plan, render_toml)
-def _model_dims(hf_config_path):
+def _model_dims(hf_config_path, quant="bf16"):
    with open(hf_config_path) as f:
        num_layers, num_experts, top_k, hidden, inter = moe_dims(json.load(f))
-    return num_layers, num_experts, top_k, expert_bytes(hidden, inter)
+    return num_layers, num_experts, top_k, expert_bytes(hidden, inter, quant)
 def check(config_path, hf_config_path):
    config = load_config(config_path)
-    num_layers, num_experts, top_k, size = _model_dims(hf_config_path)
+    num_layers, num_experts, top_k, size = _model_dims(hf_config_path, config.quant)
    print("layer  vram   ram  disk  cache_slots")
    vram_rows = ram_rows = few_slots = 0
    for layer in range(num_layers):
@@ -43,6 +43,8 @@ def _variant_text(variant):
    policy, prefetch, miss = variant
    fetch = "none" if prefetch.mode == "none" else f"pregate d{prefetch.distance} b{prefetch.budget}"
    cpu = miss.mode if miss.mode == "copy" else f"{miss.mode} admit{miss.admit}"
    if miss.mode == "auto":
        cpu += f" <={miss.cpu_max_tokens}tok"
    return (f"{policy.eviction:5} {policy.heat:8} keep{policy.scan_keep} "
            f"promote{policy.promote}  {fetch:15} {cpu}")
@@ -51,7 +53,7 @@ def sim(trace, config_path, hf_config_path, cost_path, expert_mib):
    from prain.sim import describe, load_cost, read_forwards, simulate
    config = load_config(config_path)
-    _, num_experts, _, size = _model_dims(hf_config_path)
+    _, num_experts, _, size = _model_dims(hf_config_path, config.quant)
    forwards = read_forwards(trace)
    cost = load_cost(cost_path, expert_mib or size / 2**20)
    layers = sorted({layer for fwd in forwards for layer in fwd})
@@ -64,13 +66,15 @@ def sim(trace, config_path, hf_config_path, cost_path, expert_mib):
 def tune_command(args):
    from prain.config import Miss, Policy, Prefetch
    from prain.sim import describe, load_cost, read_forwards, tune
-    _, num_experts, _, size = _model_dims(args.hf_config)
+    _, num_experts, _, size = _model_dims(args.hf_config, args.quant)
    cost = load_cost(args.cost, args.expert_mib or size / 2**20)
    forwards = read_forwards(args.trace)
    best, ram, per, final, bound, rows = tune(forwards, num_experts, args.expert_vram_gib,
-                                              args.ram_gib, cost)
+                                              args.ram_gib, cost,
                                              base=(Policy(), Prefetch(), Miss(cpu_threads=args.cpu_threads)))
    print(" rank  total_s  decode_ms  hit%  eviction heat keep promote  prefetch  miss",
          file=sys.stderr)
    for rank, (variant, r) in enumerate(rows[: args.top], 1):
@@ -84,7 +88,7 @@ def tune_command(args):
              for layer in per}
    print(render_toml(
        f"generated by `prain tune` from {os.path.basename(args.trace)}",
-        {"pack_dir": os.path.abspath(args.pack_dir), "default_tier": "disk",
+        {"pack_dir": os.path.abspath(args.pack_dir), "quant": None if args.quant == "bf16" else args.quant, "default_tier": "disk",
         "cache_slots": min(per.values()), "trace_path": os.path.abspath(args.trace)},
        {"policy": asdict(policy), "prefetch": asdict(prefetch), "miss": asdict(miss)},
        layers,
@@ -109,6 +113,9 @@ def serve(args):
                       "--scheduling-policy", "priority",
                       "--max-num-seqs", str(config.serve.max_requests),
                       "--kv-cache-memory-bytes", str(int(config.serve.kv_cache_gib * GIB)),
                       # vLLM detects "openai" for Gemma 4, whose template then prints a system
                       # prompt as a Python list; a flag in `extra` comes later and wins
                       "--chat-template-content-format", "string",
                       *extra])
@@ -126,6 +133,7 @@ def main(argv=None):
    p.add_argument("--ram-gib", type=float, required=True, help="pinned host memory for experts")
    p.add_argument("--cache-slots", type=int, default=16)
    p.add_argument("--pack-dir", required=True)
    p.add_argument("--quant", choices=QUANTS, default="bf16", help="expert weight format")
    for name, text in (("sim", "simulate a placement file on a routing trace"),
                       ("tune", "search policies and per-layer slots on a trace; prints a placement file")):
        s = sub.add_parser(name, help=text)
@@ -139,6 +147,9 @@ def main(argv=None):
            s.add_argument("--expert-vram-gib", type=float, required=True)
            s.add_argument("--ram-gib", type=float, required=True)
            s.add_argument("--pack-dir", required=True)
            s.add_argument("--quant", choices=QUANTS, default="bf16", help="expert weight format")
            s.add_argument("--cpu-threads", type=int, default=4,
                           help="[miss] cpu_threads to write: the --threads scripts/measure.py used")
            s.add_argument("--top", type=int, default=10, help="variants to list")
    y = sub.add_parser("synth", help="write a synthetic routing trace, to try sim and tune")
    y.add_argument("out")
@@ -157,20 +168,26 @@ def main(argv=None):
        ui.add_argument("--warmup-priority", type=int, default=1,
                        help="vLLM priority of warm-ups; 0 for a server without "
                             "--scheduling-policy priority")
        ui.add_argument("--persona", metavar="FILE",
                        help="text file sent as the system message of every chat: the model's "
                             "task, how to answer, its environment (see examples/personas)")
        if name == "web":
            ui.add_argument("--port", type=int, default=8080)
-            ui.add_argument("--config", help="placement file: [serve] max_requests and the "
+            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":
        check(args.config, args.hf_config)
    elif args.command == "plan":
-        _, num_experts, _, size = _model_dims(args.hf_config)
+        _, num_experts, _, size = _model_dims(args.hf_config, args.quant)
        with open(args.stats) as f:
            stats = json.load(f)
        print(plan(stats, num_experts, size, args.expert_vram_gib, args.ram_gib,
-                   args.cache_slots, os.path.abspath(args.pack_dir), os.path.abspath(args.stats)), end="")
+                   args.cache_slots, os.path.abspath(args.pack_dir), os.path.abspath(args.stats), args.quant), end="")
    elif args.command == "sim":
        sim(args.trace, args.config, args.hf_config, args.cost, args.expert_mib)
    elif args.command == "tune":
@@ -185,13 +202,23 @@ def main(argv=None):
    elif args.command == "tui":
        from prain.tui import main as tui  # curses only loads when needed
-        tui(args.url, args.warmup, args.warmup_priority)
+        tui(args.url, args.warmup, args.warmup_priority, system=_persona(args)[0])
    else:
        from prain.config import Config
        from prain.web import serve as web
        config = load_config(args.config) if args.config else Config(pack_dir="")  # defaults
-        web(args.url, args.port, args.warmup, args.warmup_priority, config.serve, config.draft)
+        system, persona = _persona(args)
        web(args.url, args.port, args.warmup, args.warmup_priority, config.serve, config.draft,
            host=args.host, stats_path=config.stats_path, system=system, persona=persona)
 def _persona(args):
    """(system message text, name) of --persona FILE, or (None, None)."""
    if not args.persona:
        return None, None
    with open(args.persona) as f:
        return f.read().strip(), os.path.splitext(os.path.basename(args.persona))[0]
 if __name__ == "__main__":
@@ -8,8 +8,10 @@ import tomllib
 from dataclasses import dataclass, field, fields
 TIERS = ("vram", "ram", "disk")
 QUANTS = ("bf16", "int8", "int4")  # expert weight formats, see prain.quant
 INT4_GROUP = 64  # int4 weights per scale
 GIB = 1024**3
-_TOP_KEYS = {"pack_dir", "default_tier", "cache_slots", "stats_path", "stats_every",
+_TOP_KEYS = {"pack_dir", "quant", "default_tier", "cache_slots", "stats_path", "stats_every",
             "trace_path", "policy", "prefetch", "miss", "serve", "draft", "layers"}
 _LAYER_KEYS = {"cache_slots", *TIERS}
@@ -104,6 +106,8 @@ class Miss:
    cpu_max_tokens: int = _opt(1, minimum=1)  # auto: CPU runs experts with <= this many tokens
    cpu_threads: int = _opt(4, minimum=1)  # OMP threads of the vLLM worker
    admit: int = _opt(1)  # CPU-run experts copied into VRAM afterwards, hottest first
    min_weight: float = _opt(1.0, maximum=1)  # router weight every token must compute; below 1
    #                                           it skips experts not in VRAM: faster, dumber
@dataclass(frozen=True)
@@ -143,6 +147,7 @@ def _section(cls, raw, where):
@dataclass(frozen=True)
 class Config:
    pack_dir: str
    quant: str = "bf16"
    default_tier: str = "ram"
    cache_slots: int = 16
    stats_path: str | None = None
@@ -221,6 +226,7 @@ def load_config(path):
        layers.append((selector, section))
    config = Config(
        pack_dir=_path(path, raw["pack_dir"]),
        quant=raw.get("quant", "bf16"),
        default_tier=raw.get("default_tier", "ram"),
        cache_slots=raw.get("cache_slots", 16),
        stats_path=_path(path, raw.get("stats_path")),
@@ -236,6 +242,8 @@ def load_config(path):
    if config.draft.full_requests > config.serve.max_requests:
        raise ValueError(f"{path}: [draft] full_requests must be <= [serve] max_requests "
                         f"({config.serve.max_requests})")
    if config.quant not in QUANTS:
        raise ValueError(f"{path}: quant must be one of {QUANTS}")
    if config.default_tier not in TIERS:
        raise ValueError(f"{path}: default_tier must be one of {TIERS}")
    _check_count(config.cache_slots, 0, f"{path}: cache_slots")
@@ -262,9 +270,14 @@ def moe_dims(hf_config):
    )
-def expert_bytes(hidden, moe_intermediate, dtype_bytes=2):
+def expert_bytes(hidden, moe_intermediate, quant="bf16"):
-    """One gated expert: gate+up [2I, H] plus down [H, I]."""
+    """One gated expert, gate+up [2I, H] plus down [H, I], stored as `quant`."""
-    return 3 * hidden * moe_intermediate * dtype_bytes
+    weights = 3 * hidden * moe_intermediate
    if quant == "bf16":
        return 2 * weights
    if quant == "int8":  # plus one fp32 scale per output row
        return weights + 4 * (2 * moe_intermediate + hidden)
    return weights // 2 + 2 * weights // INT4_GROUP  # plus one bf16 scale per group
 def render_toml(header, top, tables=None, layers=None):
@@ -278,7 +291,8 @@ def render_toml(header, top, tables=None, layers=None):
    return "\n".join(lines) + "\n"
-def plan(stats, num_experts, expert_size, expert_vram_gib, ram_gib, cache_slots, pack_dir, stats_path):
+def plan(stats, num_experts, expert_size, expert_vram_gib, ram_gib, cache_slots, pack_dir, stats_path,
         quant="bf16"):
    """Hottest experts (by steps that routed to them) pinned in VRAM, next in RAM, rest on disk."""
    uses = {int(layer): s["uses"] for layer, s in stats["layers"].items()}
    vram_rows = int(expert_vram_gib * GIB // expert_size) - cache_slots * len(uses)
@@ -298,7 +312,7 @@ def plan(stats, num_experts, expert_size, expert_vram_gib, ram_gib, cache_slots,
    }
    return render_toml(
        "generated by `prain plan`; experts not listed live on disk",
-        {"pack_dir": pack_dir, "default_tier": "disk", "cache_slots": cache_slots,
+        {"pack_dir": pack_dir, "quant": None if quant == "bf16" else quant, "default_tier": "disk", "cache_slots": cache_slots,
         "stats_path": stats_path},
        layers=layers,
    )
@@ -2,14 +2,17 @@
 vLLM builds this class instead of RoutedExperts once prain.plugin.register() ran.
 Every expert is written once to a per-layer pack file on disk while vLLM loads the
-checkpoint. On the first forward the RAM tier is copied from the pack into pinned host
+checkpoint. With quant = "int8" or "int4" the first forward quantizes that bf16 pack once into
-memory and the pinned VRAM slots are filled. After that each forward:
+a second pack (prain.quant), and the VRAM slots, copies and CPU path use the quantized
 experts. On the first forward the RAM tier is copied from the pack into pinned host memory
 and the pinned VRAM slots are filled. After that each forward:
  0. drops the routes of draft tokens to experts that are not available ([draft] min_weight;
     prain.drafts knows which rows are drafts),
  1. counts routes per expert (one host sync),
  2. lets the [miss] config send some missing experts to the CPU,
-  3. copies the other missing experts into slots chosen by the [policy] config and runs the
+  3. copies the other missing experts into slots chosen by the [policy] config and runs
-     stock fused-MoE kernel on the slot tensors, in rounds if needed,
+     vLLM's Triton fused-MoE kernel on the slot tensors, in rounds if needed (each round
     computes only its own routes),
  4. queues async copies on a side stream: [prefetch] guesses for a later layer, and the
     hottest CPU-run experts (admit-behind),
  5. adds the CPU result.
@@ -26,11 +29,14 @@ import os
 import torch
 from vllm.config import get_current_vllm_config
 from vllm.logger import init_logger
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
 from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
-from prain import drafts
+from prain import drafts, profile
 from prain.config import layer_index, load_config
 from prain.kernels import DECODE_ROWS, gfx10, moe_int4_decode, unpack_int4
 from prain.policy import SlotCache, write_stats
 from prain.quant import (GROUP, cpu_matmul, layout, moe_quant_config, quantize, to_cpu_int4)
 from prain.sim import write_trace
 from prain.tiered import draft_rows, tiered_forward
@@ -40,6 +46,68 @@ _COPY_STREAM = []  # one side stream shared by all layers: async copies stay in
 _TRACE = []  # trace records not written yet
 _DRAFT_MASK = [None, None]  # (spans, their bool row mask on the GPU) of the current forward
 _DRAFT_TOKENS = [0]  # draft token rows seen, for the stats log
 _STAGING = {}  # stream -> _Staging
 LOW_VRAM = 128 * 2**20  # driver-free bytes below which torch's cache goes back (_release_cache)
 def _release_cache():
    """Hand torch's cached free blocks (the freed one-row bf16 expert weights, ~11 MiB per layer,
    and load-time temporaries) back to the driver. HIP loads kernel code objects lazily into
    VRAM; with a full card and nothing returned, a kernel's first launch fails with
    hipErrorNoBinaryForGpu ("no kernel image")."""
    before = torch.cuda.mem_get_info()[0]
    torch.cuda.empty_cache()
    free, total = torch.cuda.mem_get_info()
    logger.info("prain: %.0f MiB of VRAM free for kernels and activations (%.0f MiB before "
                "releasing torch's cache, %.2f of %.2f GiB reserved)", free / 2**20,
                before / 2**20, torch.cuda.memory_reserved() / 2**30, total / 2**30)
 class _Staging:
    """A ring of pinned host buffers and GPU buffers for staged int4 copies: the CPU copies an
    expert's tensors from the page cache into one pinned buffer, one DMA moves it, and
    unpack_int4 writes the slot. The CPU fills the next buffer while the last DMA runs."""
    SLOTS = 4
    def __init__(self, nbytes, device):
        self.host = [torch.empty(nbytes, dtype=torch.uint8).pin_memory() for _ in range(self.SLOTS)]
        self.dev = [torch.empty(nbytes, dtype=torch.uint8, device=device)
                    for _ in range(self.SLOTS)]
        self.done = [None] * self.SLOTS  # event: the DMA that last read this pinned buffer
        self.i = 0
    def take(self):
        self.i = (self.i + 1) % self.SLOTS
        if self.done[self.i] is not None:
            self.done[self.i].synchronize()
        return self.host[self.i], self.dev[self.i], self.i
 def stage_int4(sources, slots, slot, ring, stream):
    """Copy one int4 expert: sources [[w13 q, w13 s], [w2 q, w2 s]] on the CPU (q in the CPU
    layout) -> slots[...][slot] on the GPU (q in the GPU layout), through `ring`."""
    host, dev, i = ring.take()
    parts, offset = [], 0
    with profile.span("memcpy"):
        for mats in sources:
            for src in mats:
                n = src.nbytes
                host[offset:offset + n].copy_(src.reshape(-1).view(torch.uint8))
                parts.append((offset, n))
                offset += n
    with profile.span("h2d"):
        dev[:offset].copy_(host[:offset], non_blocking=True)
        ring.done[i] = torch.cuda.Event()
        ring.done[i].record(stream)
    with profile.span("unpack"):
        parts = iter(parts)
        for dsts in slots:
            for t, dst in enumerate(dsts):
                o, n = next(parts)
                if t == 0:  # the q tensor: the pack keeps the CPU kernel's layout
                    unpack_int4(dev[o:o + n], dst[slot])
                else:
                    dst[slot].view(-1).view(torch.uint8).copy_(dev[o:o + n])
 def _install_draft_hook():
@@ -121,6 +189,8 @@ class TieredRoutedExperts(RoutedExperts):
            raise NotImplementedError("prain: one GPU only (no TP, DP or EP)")
        if vllm_config.use_v2_model_runner:
            _install_draft_hook()
            if profile.DIR:
                profile.install()
        elif not LAYERS:
            logger.warning("prain: drafts need vLLM's V2 model runner; with "
                           "VLLM_USE_V2_MODEL_RUNNER=0 every request runs with all its experts")
@@ -133,9 +203,16 @@ class TieredRoutedExperts(RoutedExperts):
            raise NotImplementedError("prain: prefetch.mode = pregate needs a router module "
                                      "passed as custom_routing_function (Gemma 4 does)")
        # vLLM sizes w13/w2 and the kernel by num_experts: make that the VRAM slot count.
-        moe_config.num_experts = moe_config.num_local_experts = placement.num_slots
+        # Quantized slots are prain's own tensors; vLLM's bf16 ones shrink to one row.
        quantized = config.quant != "bf16"
        moe_config.num_experts = moe_config.num_local_experts = 1 if quantized else placement.num_slots
        super().__init__(layer_name, params_dtype, moe_config, quant_config, *args, **kwargs)
        self.local_num_experts = placement.num_slots
        self.quant = config.quant
        self._shapes = (tuple(self.w13_weight.shape[1:]), tuple(self.w2_weight.shape[1:]))
        self._qslots = [[torch.zeros((placement.num_slots, *shape), dtype=dtype, device="cuda")
                         for shape, dtype in layout(config.quant, *dims)]
                        for dims in self._shapes] if quantized else None
        self.prain = config
        self.prain_layer = layer
        self.placement = placement
@@ -155,27 +232,69 @@ class TieredRoutedExperts(RoutedExperts):
        logger.info(
            "prain layer %d: %d vram, %d ram, %d disk, %d cache slots%s",
            layer, len(placement.vram), len(placement.ram), len(placement.disk),
-            placement.cache_slots, "" if self.pack_complete else ", writing pack",
+            placement.cache_slots, "" if self.pack_complete else
            ", writing pack" if self._writing else f", quantizing the pack to {self.quant}",
        )
    def _open_pack(self, pack_dir, layer, num_experts, model):
        """The loader writes the bf16 pack; a quantized pack is made from it once. Once the
        pack of config.quant is complete, the bf16 one is not read (it may be deleted)."""
        os.makedirs(pack_dir, exist_ok=True)
-        stem = os.path.join(pack_dir, f"layer{layer:03d}")
+        self._stem = os.path.join(pack_dir, f"layer{layer:03d}")
-        w13_shape = (num_experts, *self.w13_weight.shape[1:])
+        self._num_experts = num_experts
-        w2_shape = (num_experts, *self.w2_weight.shape[1:])
+        self._marker = {"model": model, "w13": [num_experts, *self._shapes[0]],
-        self._pack_files = (f"{stem}.w13", f"{stem}.w2")
+                        "w2": [num_experts, *self._shapes[1]], "dtype": str(self.w13_weight.dtype)}
-        self._pack_w13 = _file_tensor(self._pack_files[0], w13_shape, self.w13_weight.dtype)
+        self.pack_complete = self._complete(self.quant)
-        self._pack_w2 = _file_tensor(self._pack_files[1], w2_shape, self.w2_weight.dtype)
+        self._bf16 = None if self.pack_complete else self._pack_of("bf16")
-        self._marker_path = f"{stem}.done"
+        self._writing = self._bf16 is not None and not self._complete("bf16")
-        self._marker = {"model": model, "w13": list(w13_shape), "w2": list(w2_shape),
+
-                        "dtype": str(self.w13_weight.dtype)}
+    def _paths(self, quant):
-        self.pack_complete = os.path.exists(self._marker_path)
+        """[[w13 files], [w2 files]] of a pack format: weights, then scales."""
-        if self.pack_complete:
+        infix = "" if quant == "bf16" else f".{quant}"
-            with open(self._marker_path) as f:
+        return [[f"{self._stem}{infix}.{name}{'s' * i}" for i in range(len(layout(quant, *dims)))]
-                found = json.load(f)
+                for name, dims in zip(("w13", "w2"), self._shapes)]
-            if found != self._marker:
+
-                raise ValueError(f"prain: {self._marker_path} is for {found}, not "
+    def _pack_of(self, quant):
-                                 f"{self._marker}; delete {pack_dir} or use another pack_dir")
+        return [[_file_tensor(path, (self._num_experts, *shape), dtype)
                 for path, (shape, dtype) in zip(paths, layout(quant, *dims))]
                for paths, dims in zip(self._paths(quant), self._shapes)]
    def _marker_of(self, quant):
        if quant == "bf16":
            return f"{self._stem}.done", self._marker
        return f"{self._stem}.{quant}.done", {**self._marker, "quant": quant, "group": GROUP}
    def _complete(self, quant):
        path, marker = self._marker_of(quant)
        if not os.path.exists(path):
            return False
        with open(path) as f:
            found = json.load(f)
        if found != marker:
            raise ValueError(f"prain: {path} is for {found}, not {marker}; delete the pack "
                             "or use another pack_dir")
        return True
    def _seal(self, quant):
        for path in sum(self._paths(quant), []):  # mmap writes -> disk before we trust the pack
            fd = os.open(path, os.O_RDONLY)
            os.fsync(fd)
            os.close(fd)
        path, marker = self._marker_of(quant)
        with open(path, "w") as f:
            json.dump(marker, f)
    def _quantize_pack(self):
        """bf16 pack -> config.quant pack, one expert at a time on the GPU (int4's clip search
        takes ~20 min for the whole model on the CPU, under a minute on the GPU)."""
        out = self._pack_of(self.quant)
        for src, dst in zip(self._bf16, out):
            for e in range(self._num_experts):
                parts = quantize(src[0][e].cuda(), self.quant)
                if self.quant == "int4":
                    parts[0] = to_cpu_int4(parts[0])
                for tensor, part in zip(dst, parts):
                    tensor[e] = part.cpu()
    def get_expert_mapping(self, ckpt_gate_proj_name=None, ckpt_down_proj_name=None,
                           ckpt_up_proj_name=None, include_fused=False):
@@ -198,8 +317,8 @@ class TieredRoutedExperts(RoutedExperts):
    def weight_loader(self, param, loaded_weight, weight_name, shard_id, expert_id,
                      return_success=False):
        # Every expert goes to the pack; the VRAM and RAM tiers fill from it later.
-        if not self.pack_complete:
+        if self._writing:
-            pack = self._pack_w13 if param is self.w13_weight else self._pack_w2
+            pack = self._bf16[0][0] if param is self.w13_weight else self._bf16[1][0]
            super().weight_loader(pack, loaded_weight, weight_name, shard_id, expert_id,
                                  return_success=True)
        return True if return_success else None
@@ -208,15 +327,24 @@ class TieredRoutedExperts(RoutedExperts):
    def _materialize(self):
        if not self.pack_complete:
-            for path in self._pack_files:  # mmap writes -> disk before we trust the pack
+            if self._writing:
-                fd = os.open(path, os.O_RDONLY)
+                self._seal("bf16")
-                os.fsync(fd)
+            if self.quant != "bf16":
-                os.close(fd)
+                self._quantize_pack()
-            with open(self._marker_path, "w") as f:
+                self._seal(self.quant)
-                json.dump(self._marker, f)
+            self.pack_complete, self._writing, self._bf16 = True, False, None
-            self.pack_complete = True
+        self._pack = self._pack_of(self.quant)  # [[w13 tensors], [w2 tensors]] per expert row
        self._expert_bytes = sum(t[0].nbytes for mats in self._pack for t in mats)
        if self._qslots is None:
            self._slots, self._qconfig = [[self.w13_weight.data], [self.w2_weight.data]], None
        else:
            self._slots = self._qslots
            self._qconfig = moe_quant_config(self.quant, self._slots[0][1], self._slots[1][1])
            # vLLM's one-row bf16 w13/w2 are never read: free them (11 MiB per layer)
            self.w13_weight.data = self.w13_weight.data.new_empty(0)
            self.w2_weight.data = self.w2_weight.data.new_empty(0)
        if self.prain_layer == min(LAYERS):  # check the whole RAM tier once, before pinning
-            row = self._pack_w13[0].nbytes + self._pack_w2[0].nbytes
+            row = sum(t[0].nbytes for mats in self._pack for t in mats)
            need = row * sum(len(experts.placement.ram) for experts in LAYERS.values())
            free = _mem_available()
            if free is not None and need > 0.9 * free:
@@ -225,56 +353,72 @@ class TieredRoutedExperts(RoutedExperts):
                                   "from ram to disk")
        ram = list(self.placement.ram)
        self._ram_row = {e: row for row, e in enumerate(ram)}
-        index = torch.tensor(ram, dtype=torch.int64, device="cpu")
+        # ponytail: pinning is slow on the eGPU host (minutes for a 10 GiB tier, even per
-        self._ram_w13, self._ram_w2 = (
+        # expert); the page cache serves the pack about as fast, so keep the RAM tier small.
-            torch.index_select(pack, 0, index, out=torch.empty(
+        self._ram = [[[t[e].pin_memory() for t in mats] for mats in self._pack] for e in ram]
                (len(ram), *pack.shape[1:]), dtype=pack.dtype, device="cpu", pin_memory=True))
            for pack in (self._pack_w13, self._pack_w2)
        )
        for slot, expert in enumerate(self.placement.vram):
            self._load(expert, slot)
        if self._router and (self.prain.prefetch.mode == "pregate" or self.prain.trace_path):
            self._router[0].register_forward_pre_hook(self._predict)
    def _source(self, expert):
-        """CPU tensors of an expert: pinned rows for the RAM tier, else the mmap'd pack."""
+        """[[w13 tensors], [w2 tensors]] of an expert on the CPU: pinned rows for the RAM tier,
        else the mmap'd pack."""
        row = self._ram_row.get(expert)
        if row is None:
-            return self._pack_w13[expert], self._pack_w2[expert]
+            return [[t[expert] for t in mats] for mats in self._pack]
-        return self._ram_w13[row], self._ram_w2[row]
+        return self._ram[row]
    def _copy(self, expert, slot):
        if self.quant == "int4":  # the pack keeps the CPU kernel's layout: stage, then unpack
            return self._copy_staged(expert, slot)
        for dsts, srcs in zip(self._slots, self._source(expert)):
            for dst, src in zip(dsts, srcs):
                with profile.span("h2d" if src.device.type == "cpu" else "d2d"):
                    dst[slot].copy_(src, non_blocking=True)
    def _copy_staged(self, expert, slot):
        """int4 on the current stream: one pinned buffer, one DMA, unpack into the slot."""
        stream = torch.cuda.current_stream()
        ring = _STAGING.get(stream)
        if ring is None:
            ring = _STAGING[stream] = _Staging(self._expert_bytes, stream.device)
        stage_int4(self._source(expert), self._slots, slot, ring, stream)
    def _load(self, expert, slot):
        """Demand copy on the compute stream."""
        pending = self._ready_events.pop(slot, None)
        if pending is not None:  # an older async copy into this slot must land first
            torch.cuda.current_stream().wait_event(pending)
-        w13, w2 = self._source(expert)
+        self._copy(expert, slot)
        self.w13_weight.data[slot].copy_(w13, non_blocking=True)
        self.w2_weight.data[slot].copy_(w2, non_blocking=True)
    def _load_async(self, expert, slot):
-        """Copy on the side stream; only RAM-tier (pinned) experts are truly asynchronous."""
+        """Copy on the side stream; truly asynchronous for int4 (staged) and pinned RAM-tier
        experts."""
        stream = _copy_stream()
        last = self._last_use.get(slot)
        if last is not None:  # a queued kernel may still read the slot's old expert
            stream.wait_event(last)
        w13, w2 = self._source(expert)
        with torch.cuda.stream(stream):
-            self.w13_weight.data[slot].copy_(w13, non_blocking=True)
+            self._copy(expert, slot)
            self.w2_weight.data[slot].copy_(w2, non_blocking=True)
        event = torch.cuda.Event()
        event.record(stream)
        self._ready_events[slot] = event
-    def _run_slots(self, x, weights, slot_ids, used_slots):
+    def _run_slots(self, x, weights, ids, table, used_slots):
        """The fused-MoE kernel on the slot tensors (vLLM's, or prain's for int4 on RDNA2);
        expert_map=table skips the routes of experts outside this round."""
        current = torch.cuda.current_stream()
        for slot in used_slots:
            pending = self._ready_events.pop(slot, None)
            if pending is not None:
                current.wait_event(pending)
-        out = super().forward_modular(x, weights, slot_ids)
+        if self.quant == "int4" and gfx10(x.device.index or 0):
-        if not isinstance(out, torch.Tensor):
+            out = self._int4_gfx10(x, weights, ids, table)
-            raise NotImplementedError("prain: deferred MoE finalize is not supported")
+        else:
            out = fused_experts(x, self._slots[0][0], self._slots[1][0], weights, ids,
                                activation=self.activation, global_num_experts=table.numel(),
                                expert_map=table, quant_config=self._qconfig)
        if self._async:
            event = torch.cuda.Event()
            event.record(current)
@@ -282,6 +426,22 @@ class TieredRoutedExperts(RoutedExperts):
                self._last_use[slot] = event
        return out
    def _int4_gfx10(self, x, weights, ids, table):
        """RDNA2 has no bf16 dot instructions, which vLLM's int4 kernel uses: a GEMV-style
        kernel for a few tokens (10x faster at one), else vLLM's kernel in fp16 (2x faster),
        redone in bf16 if fp16 overflowed."""
        if (x.shape[0] <= DECODE_ROWS
                and getattr(self.activation, "value", self.activation) == "gelu_tanh"):
            return moe_int4_decode(x, weights, ids, table, *self._slots[0], *self._slots[1])
        run = lambda x, config: fused_experts(  # noqa: E731
            x, self._slots[0][0], self._slots[1][0], weights, ids, activation=self.activation,
            global_num_experts=table.numel(), expert_map=table, quant_config=config)
        out = run(x.half(), moe_quant_config("int4", self._slots[0][1].half(),
                                             self._slots[1][1].half()))
        if not torch.isfinite(out).all():  # host sync, once per round of a prefill
            out = run(x, self._qconfig)
        return out.to(x.dtype)
    def _predict(self, module, args):
        """Pre-hook of this layer's router: guess the experts of later layers from the same
        residual stream (their routers normalize their input, so the guess is cheap and
@@ -298,7 +458,10 @@ class TieredRoutedExperts(RoutedExperts):
            self._predictions[d] = list(zip(e.tolist(), p.tolist()))  # host sync
    def _cpu_weights(self, expert):
-        return self._source(expert)
+        mats = self._source(expert)
        if self.quant == "bf16":
            return mats[0][0], mats[1][0]
        return tuple(cpu_matmul(self.quant, *parts) for parts in mats)
    def _after_enqueue(self, cpu, stats):
        miss, prefetch = self.prain.miss, self.prain.prefetch
@@ -308,8 +471,9 @@ class TieredRoutedExperts(RoutedExperts):
                self._load_async(e, slot)
        target = LAYERS.get(self.prain_layer + prefetch.distance)
        if prefetch.mode == "pregate" and target is not None and target._ready:
            staged = target.quant == "int4"  # staged copies are async from every tier
            guesses = [(e, p) for e, p in self._predictions.get(prefetch.distance, ())
-                       if p >= prefetch.min_prob and e in target._ram_row]
+                       if p >= prefetch.min_prob and (staged or e in target._ram_row)]
            for e, slot in target.cache.prefetch(guesses, prefetch.budget):
                target._load_async(e, slot)
@@ -330,6 +494,8 @@ class TieredRoutedExperts(RoutedExperts):
        if not self._ready:
            self._materialize()
            self._ready = True
            if self.prain_layer == max(LAYERS):  # every layer has its slots now
                _release_cache()
        policy = self.prain.policy
        out = tiered_forward(
            x, topk_weights, topk_ids, self.cache, self._load, self._run_slots,
@@ -347,6 +513,8 @@ class TieredRoutedExperts(RoutedExperts):
        self._steps += 1
        if self._steps % self.prain.stats_every or self.prain_layer != min(LAYERS):
            return
        if torch.cuda.mem_get_info()[0] < LOW_VRAM:  # other processes or activations took it
            _release_cache()
        if self.prain.trace_path:
            _flush_trace()
        if not self.prain.stats_path:
@@ -0,0 +1,90 @@
 """Decode GEMMs (at most MAX_ROWS rows) for gfx10 GPUs as a Triton GEMV.
 vLLM sends unquantized linear layers on ROCm through rocBLAS, except on gfx9/gfx11/gfx12 where it
 has skinny-GEMM kernels. On an RX 6600 XT (gfx1032) rocBLAS runs a bf16 GEMV at 13-16 GB/s
 (RDNA2 has no bf16 dot instructions), so the dense part of a decode step reads its weights at
 6% of the card's bandwidth. This kernel reads them at ~210 GB/s (measured, lm_head 262144 x 2816:
 93 ms -> 6.7 ms). install() wraps vLLM's dispatch_unquantized_gemm; PRAIN_GEMV=0 turns it off.
 """
 import functools
 import os
 import torch
 import torch.nn.functional as F
 import triton
 import triton.language as tl
 from prain.kernels import gfx10
 MAX_ROWS = 4
 # ponytail: one config, the one autotuning picked for every Gemma 4 dense shape on gfx1032.
 # Triton's autotuner benchmarks with a 256 MiB buffer that torch keeps cached, and an 8 GB card
 # has no 256 MiB to spare (lazy kernel loads then fail). Re-measure if the GPU or model changes.
 BLOCK_N, BLOCK_K, NUM_WARPS = 4, 512, 8
@triton.jit
 def _gemv(w_ptr, x_ptr, y_ptr, M, N, K, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
          BLOCK_K: tl.constexpr):
    rows = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N)
    ms = tl.arange(0, BLOCK_M)
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k0 in range(0, K, BLOCK_K):
        cols = k0 + tl.arange(0, BLOCK_K)
        w = tl.load(w_ptr + rows[:, None] * K + cols[None, :],
                    mask=(rows[:, None] < N) & (cols[None, :] < K), other=0.0)
        x = tl.load(x_ptr + ms[:, None] * K + cols[None, :],
                    mask=(ms[:, None] < M) & (cols[None, :] < K), other=0.0)
        acc += tl.sum(x.to(tl.float32)[:, None, :] * w.to(tl.float32)[None, :, :], axis=2)
    tl.store(y_ptr + ms[:, None] * N + rows[None, :], acc.to(y_ptr.dtype.element_ty),
             mask=(ms[:, None] < M) & (rows[None, :] < N))
@functools.cache
 def fast_device(index):
    """gfx10 (RDNA2) GPUs, where vLLM has no skinny-GEMM kernels."""
    return gfx10(index) and os.environ.get("PRAIN_GEMV", "1") != "0"
 def _small(x, weight, bias):
    return (bias is None and x.is_cuda and x.dtype in (torch.bfloat16, torch.float16)
            and weight.dtype == x.dtype and weight.is_contiguous()
            and x.numel() <= MAX_ROWS * x.shape[-1] and fast_device(x.device.index or 0))
 def linear(x, weight, bias=None):
    """F.linear(x, weight, bias); the Triton GEMV for up to MAX_ROWS rows on gfx10."""
    if not _small(x, weight, bias):
        return F.linear(x, weight, bias)
    n, k = weight.shape
    x2 = x.reshape(-1, k).contiguous()
    m = x2.shape[0]
    y = torch.empty(m, n, dtype=x.dtype, device=x.device)
    _gemv[(triton.cdiv(n, BLOCK_N),)](weight, x2, y, m, n, k, BLOCK_M=triton.next_power_of_2(m),
                                      BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, num_warps=NUM_WARPS)
    return y.reshape(*x.shape[:-1], n)
 def wrap(dispatch):
    """vLLM's dispatch_unquantized_gemm(), with decode shapes sent to linear()."""
    def wrapped(*args, **kwargs):
        fallback = dispatch(*args, **kwargs)
        def gemm(layer, x, weight, bias=None):
            if _small(x, weight, bias):
                return linear(x, weight)
            return fallback(layer, x, weight, bias)
        return gemm
    return wrapped
 def install():
    """Before the model is built: every module that looks up dispatch_unquantized_gemm."""
    from vllm.model_executor.layers import linear as linear_layers, utils
    from vllm.model_executor.layers import vocab_parallel_embedding
    from vllm.model_executor.layers.fused_moe.runner import moe_runner
    wrapped = wrap(utils.dispatch_unquantized_gemm)
    for module in (linear_layers, vocab_parallel_embedding, moe_runner):
        module.dispatch_unquantized_gemm = wrapped
@@ -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])
@@ -10,7 +10,9 @@ def register():
    if "RoutedExperts" in op_registry_oot:  # vLLM may call register() more than once
        return
    from prain import gemv
    from prain.experts import TieredRoutedExperts
    # vLLM's documented out-of-tree hook: every RoutedExperts(...) now builds ours.
    PluggableLayer.register_oot(TieredRoutedExperts, name="RoutedExperts")
    gemv.install()  # decode GEMVs on gfx10
@@ -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}")
@@ -0,0 +1,89 @@
 """Weight-only int8 / int4 quantization of expert weights. Needs torch.
 int8: q int8 [..., N, K] and one fp32 scale per output row: w = q * scale. vLLM's Triton
      fused-MoE kernel and torch's _weight_int8pack_mm (CPU) both read this layout.
 int4: one bf16 scale per GROUP inputs of a row: w = (nibble - 8) * scale, with the scale's
      clip ratio chosen per group for the least squared error. Two layouts of the nibbles:
      gpu: [..., N, K/2], two values per byte along K, even k in the low nibble (vLLM's
           Triton GPTQ/AWQ MoE kernel);
      cpu: [N/32, K, 16] bytes, rows d and d+16 of each 32-row block share a byte (torch's AVX2
           _weight_int4pack_mm_for_cpu); packs store it, slots convert with to_gpu_int4().
 """
 import torch
 from prain.config import INT4_GROUP as GROUP  # divides Gemma 4 26B-A4B's 2816 and 704
 CLIP = (1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7)
 def layout(quant, rows, cols):
    """[(shape, dtype)] of the tensors that store one [rows, cols] weight matrix."""
    if quant == "bf16":
        return [((rows, cols), torch.bfloat16)]
    if quant == "int8":
        return [((rows, cols), torch.int8), ((rows,), torch.float32)]
    return [((rows, cols // 2), torch.uint8), ((rows, cols // GROUP), torch.bfloat16)]
 def quantize(w, quant):
    """Float weights [..., N, K] -> [q, scale], int4 in the gpu layout."""
    w = w.float()
    if quant == "int8":
        scale = w.abs().amax(-1).clamp_min(1e-12) / 127
        return [(w / scale[..., None]).round().to(torch.int8), scale]
    g = w.unflatten(-1, (-1, GROUP))
    amax = g.abs().amax(-1, keepdim=True).clamp_min(1e-12)
    best_err = best_scale = None
    for ratio in CLIP:
        scale = (amax * ratio / 7).bfloat16().float()  # the kernels see the bf16 scale
        err = ((g / scale).round().clamp(-8, 7) * scale - g).pow(2).sum(-1, keepdim=True)
        if best_err is None:
            best_err, best_scale = err, scale
        else:
            better = err < best_err
            best_err, best_scale = torch.where(better, err, best_err), torch.where(better, scale, best_scale)
    q = ((g / best_scale).round().clamp(-8, 7) + 8).flatten(-2).to(torch.uint8)
    return [q[..., 0::2] | q[..., 1::2] << 4, best_scale.squeeze(-1).bfloat16()]
 def dequantize(q, scale, quant, dtype=torch.bfloat16):
    """The weights [..., N, K] of quantize()'s output (int4 in the gpu layout)."""
    if quant == "int8":
        return q.to(dtype) * scale.to(dtype)[..., None]
    v = torch.stack((q & 15, q >> 4), -1).flatten(-2).to(dtype) - 8
    return (v.unflatten(-1, (-1, GROUP)) * scale.to(dtype)[..., None]).flatten(-2)
 def to_cpu_int4(q):
    """int4 nibbles [..., N, K/2]: gpu layout -> cpu layout (same shape)."""
    *lead, n, half = q.shape
    v = torch.stack((q & 15, q >> 4), -1).reshape(*lead, n // 32, 32, 2 * half).transpose(-1, -2)
    return (v[..., :16] | v[..., 16:] << 4).reshape(q.shape)
 def to_gpu_int4(c):
    """int4 nibbles [..., N, K/2]: cpu layout -> gpu layout (same shape; works on the GPU)."""
    *lead, n, half = c.shape
    b = c.reshape(*lead, n // 32, 2 * half, 16)
    v = torch.cat((b & 15, b >> 4), -1).transpose(-1, -2).reshape(*lead, n, 2 * half)
    return v[..., 0::2] | v[..., 1::2] << 4
 def cpu_matmul(quant, q, scale):
    """x -> x @ w.T on the CPU for stored int8 q / cpu-layout int4 q; x is bf16 [M, K]."""
    if quant == "int8":
        return lambda x: torch._weight_int8pack_mm(x, q, scale.to(x.dtype))
    zeros = torch.stack((scale.t(), torch.zeros_like(scale.t())), -1).contiguous()
    return lambda x: torch._weight_int4pack_mm_for_cpu(x, q, GROUP, zeros)
 def moe_quant_config(quant, w13_scale, w2_scale):
    """vLLM's fused-MoE quant config for int8 / gpu-layout int4 slot tensors."""
    from vllm.model_executor.layers.fused_moe import config
    if quant == "int8":
        return config.int8_w8a16_moe_quant_config(w1_scale=w13_scale, w2_scale=w2_scale,
                                                  w1_zp=None, w2_zp=None)
    return config.int4_w4a16_moe_quant_config(w1_scale=w13_scale, w2_scale=w2_scale, w1_zp=None,
                                              w2_zp=None, block_shape=[0, GROUP])
@@ -31,9 +31,13 @@ class Cost:
    sync_ms: float = 0.05  # host sync per layer
    cpu_ms: float = 0.5  # CPU work per (token, expert) route
    def read_ms(self):
        """Reading one expert from the pack (the CPU path pays this for disk experts too)."""
        return self.expert_mib * 1.048576 / self.disk_gbps
    def copy_ms(self, disk):
        mb = self.expert_mib * 1.048576
-        return mb / self.h2d_gbps + (mb / self.disk_gbps if disk else 0.0)
+        return mb / self.h2d_gbps + (self.read_ms() if disk else 0.0)
 def load_cost(path, expert_mib):
@@ -175,7 +179,8 @@ def simulate(forwards, placements, policy, prefetch, miss, cost, belady=False):
                    gpu = link
                land = max((ready.pop((layer, e), 0.0) for e in slot_of), default=0.0)
                gpu = max(gpu, land) + cost.round_ms + cost.route_ms * sum(counts[e] for e in slot_of)
-            cpu_end = t + cost.sync_ms + cost.cpu_ms * sum(counts[e] for e in cpu)
+            cpu_end = (t + cost.sync_ms + cost.cpu_ms * sum(counts[e] for e in cpu)
                       + cost.read_ms() * sum(e in disk[layer] for e in cpu))
            later = []
            if cpu and miss.admit:
                hot = sorted((e for e in cpu if e not in disk[layer]), key=lambda e: (-heat[e], e))
@@ -226,15 +231,19 @@ DIMENSIONS = (
    ("prefetch", "mode", ("none", "pregate")),
    ("prefetch", "distance", (1, 2)),
    ("prefetch", "budget", (1, 2, 4)),
-    ("miss", "mode", ("copy", "auto", "cpu")),
+    # One dimension: cpu_max_tokens only matters for auto, and coordinate descent would never
    # reach auto with a high threshold once auto <= 1 lost to cpu.
    ("miss", ("mode", "cpu_max_tokens"),
     (("copy", 1), ("cpu", 1)) + tuple(("auto", n) for n in (1, 2, 4, 8, 16))),
    ("miss", "admit", (0, 1, 2)),
 )
 _PART = {"policy": 0, "prefetch": 1, "miss": 2}
 def _with(variant, table, name, value):
    names, values = (name, value) if isinstance(name, tuple) else ((name,), (value,))
    parts = list(variant)
-    parts[_PART[table]] = replace(parts[_PART[table]], **{name: value})
+    parts[_PART[table]] = replace(parts[_PART[table]], **dict(zip(names, values)))
    return tuple(parts)
@@ -3,6 +3,8 @@
 import torch
 import torch.nn.functional as F
 from prain import profile
 HEAT_INDEX = {"count": 0, "weight": 1, "recency": 2}
@@ -49,7 +51,7 @@ def degrade(topk_weights, topk_ids, draft, available, min_weight):
    A draft token keeps its heaviest unavailable routes only while less than min_weight of its
    router weight would be computed; its kept routes are scaled up to its full router weight.
-    Dropped routes get id -1 and weight 0, which route_stats and remap ignore.
+    Dropped routes get id -1 and weight 0, which route_stats and slot_table ignore.
    draft: bool [tokens]; available: bool [num_experts]. Full tokens are left unchanged."""
    ids = topk_ids.long()
    valid = ids >= 0
@@ -68,37 +70,36 @@ def degrade(topk_weights, topk_ids, draft, available, min_weight):
    return weights, torch.where(drop, torch.full_like(topk_ids, -1), topk_ids)
-def remap(topk_weights, topk_ids, slot_of, num_experts):
+def slot_table(slot_of, num_experts, device):
-    """Point routes at VRAM slots for one round.
+    """int32 [num_experts + 1]: VRAM slot of each expert in one round, -1 for the others.
-
+    vLLM's fused-MoE kernel takes it as expert_map and skips routes that map to -1; the extra
-    Routes to experts outside this round (and padding routes with id -1) get slot 0
+    last row stands for the padding id (route id -1 arrives as num_experts)."""
-    and weight 0, so any fused-MoE kernel can run them unchanged."""
+    table = torch.full((num_experts + 1,), -1, dtype=torch.int32)
    table = torch.full((num_experts,), -1, dtype=torch.int64)
    if slot_of:
-        table[torch.tensor(list(slot_of), dtype=torch.int64)] = torch.tensor(
+        table[torch.tensor(list(slot_of))] = torch.tensor(list(slot_of.values()), dtype=torch.int32)
-            list(slot_of.values()), dtype=torch.int64
+    return table.to(device)
-        )
+
-    table = table.to(topk_ids.device)
+
-    ids = topk_ids.long()
+def _matmul(w):
-    slots = torch.where(ids >= 0, table[ids.clamp_min(0)], -1)
+    """x -> x @ w.T for a weight tensor; functions (quantized experts) pass through."""
-    weights = torch.where(slots >= 0, topk_weights, torch.zeros_like(topk_weights))
+    return w if callable(w) else (lambda v: v.to(w.dtype) @ w.t())
    return weights, slots.clamp_min(0).to(topk_ids.dtype)
 def cpu_experts(x, topk_weights, topk_ids, experts, weights_of):
    """The routes to `experts`, computed on the CPU like vLLM's fused kernel with activation
    "gelu_tanh": w13 = [gate; up], y = (gelu_tanh(x @ gate.T) * (x @ up.T)) @ w2.T times the
-    route weight. weights_of(e) -> (w13 [2I, H], w2 [H, I]). Returns [tokens, hidden] fp32."""
+    route weight. weights_of(e) -> (w13 [2I, H], w2 [H, I]), or two functions v -> v @ w.T
    (quantized experts). Returns [tokens, hidden] fp32."""
    out = torch.zeros(x.shape, dtype=torch.float32)
    for e in sorted(experts):
        rows, cols = (topk_ids == e).nonzero(as_tuple=True)
        if rows.numel() == 0:
            continue
-        w13, w2 = weights_of(e)
+        mm13, mm2 = map(_matmul, weights_of(e))
-        h = x[rows].to(w13.dtype) @ w13.t()
+        h = mm13(x[rows])
-        half = w13.shape[0] // 2
+        half = h.shape[1] // 2
        act = F.gelu(h[:, :half].float(), approximate="tanh") * h[:, half:].float()
-        y = (act.to(w2.dtype) @ w2.t()).float()
+        y = mm2(act.to(h.dtype)).float()
        out.index_add_(0, rows, y * topk_weights[rows, cols].float()[:, None])
    return out
@@ -108,21 +109,35 @@ def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight"
                   after_enqueue=None, draft=None, min_weight=0.5):
    """One MoE forward over VRAM slots.
-    cache: SlotCache. load(expert, slot) copies an expert into a slot. kernel(x, weights,
+    cache: SlotCache. load(expert, slot) copies an expert into a slot. kernel(x, weights, ids,
-    slot_ids, used_slots) runs the stock fused-MoE kernel on the slot tensors. With a
+    table, used_slots) runs the fused-MoE kernel on the slot tensors for the routes whose
    expert the slot_table() `table` maps to a slot. With a
    [miss] config other than "copy", the experts it picks run on the CPU (cpu_weights(e)
    gives their weights) while the GPU rounds run. on_stats(stats, tokens) sees the routing
    statistics; after_enqueue(cpu, stats) runs once all GPU work of this forward is queued.
-    draft: bool [tokens] for the rows of draft requests, or None; see degrade().
+    draft: bool [tokens] for the rows of draft requests, or None; see degrade(). Draft rows keep
    min_weight of their router weight, every row keeps [miss] min_weight (1 = everything).
    """
    raw = None
-    if draft is not None:
+    skip = miss is not None and miss.min_weight < 1
-        if on_stats is not None:  # the trace records the routing, not what drafts skipped
+    if draft is not None or skip:
-            raw = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)
+        with profile.span("degrade"):
-        available = available_experts(topk_ids, draft, set(cache.pinned) | set(cache.slots),
+            if on_stats is not None:  # the trace records the routing, not what was skipped
-                                      cache.num_experts)
+                raw = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)
-        topk_weights, topk_ids = degrade(topk_weights, topk_ids, draft, available, min_weight)
+            rows = draft if draft is not None else torch.zeros(
-    stats = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)  # host sync
+                topk_ids.shape[0], dtype=torch.bool, device=topk_ids.device)
            need = min_weight
            if skip:  # every row may skip; draft rows by the smaller share
                need = torch.where(rows, min(min_weight, miss.min_weight), miss.min_weight)[:, None]
                rows = torch.ones_like(rows)
            available = available_experts(topk_ids, rows, set(cache.pinned) | set(cache.slots),
                                          cache.num_experts)
            _, kept = degrade(topk_weights, topk_ids, rows, available, need)
            # experts some row copies anyway serve every row for free
            available |= torch.bincount(kept[kept >= 0].long(), minlength=cache.num_experts) > 0
            topk_weights, topk_ids = degrade(topk_weights, topk_ids, rows, available, need)
    with profile.span("route"):
        stats = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)  # sync
    tokens = int(topk_ids.shape[0])
    if on_stats is not None:
        on_stats(raw if raw is not None else stats, tokens)
@@ -131,22 +146,28 @@ def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight"
    if cpu:  # copy the inputs now: later the copy would wait for the queued GPU rounds
        host = x.cpu(), topk_weights.cpu(), topk_ids.cpu()
    index = HEAT_INDEX[heat]
-    outs = []
+    ids = torch.where(topk_ids >= 0, topk_ids, cache.num_experts)  # padding -> the table's last row
    for slot_of, loads in cache.plan({e: s[index] for e, s in stats.items()}, tokens, cpu):
        for expert, slot in loads:
            load(expert, slot)
        if slot_of:
            weights, slots = remap(topk_weights, topk_ids, slot_of, cache.num_experts)
            outs.append(kernel(x, weights, slots, list(slot_of.values())))
    if after_enqueue is not None:
        after_enqueue(cpu, stats)
    if cpu:  # runs while the GPU works through its rounds
        outs.append(cpu_experts(*host, cpu, cpu_weights).to(x.device))
    if not outs:
        return torch.zeros_like(x)
    if len(outs) == 1:
        return outs[0].to(x.dtype)
    # ponytail: every extra round re-runs all tokens with weight 0 for experts it does not
    # hold; only big batches need many rounds. Upgrade: compact tokens per round.
    acc = torch.promote_types(x.dtype, torch.float32)
-    return torch.stack([o.to(acc) for o in outs]).sum(0).to(x.dtype)
+    total = None  # sum of the round (and CPU) outputs: one buffer, not one per round
    def add(out):
        nonlocal total
        total = out if total is None else total.to(acc).add_(out)  # .to: a no-op once fp32
    with profile.span("plan"):
        rounds = cache.plan({e: s[index] for e, s in stats.items()}, tokens, cpu)
    for slot_of, loads in rounds:
        for expert, slot in loads:
            with profile.span("copy"):
                load(expert, slot)
        if slot_of:  # each round computes only its own routes
            with profile.span("round"):
                table = slot_table(slot_of, cache.num_experts, x.device)
                add(kernel(x, topk_weights, ids, table, list(slot_of.values())))
    if after_enqueue is not None:
        with profile.span("prefetch"):
            after_enqueue(cpu, stats)
    if cpu:  # runs while the GPU works through its rounds
        with profile.span("cpu experts"):
            add(cpu_experts(*host, cpu, cpu_weights).to(x.device))
    return torch.zeros_like(x) if total is None else total.to(x.dtype)
@@ -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"]
@@ -0,0 +1,62 @@
 """Time to first token and decode speed of an OpenAI-compatible server, as JSON lines.
 usage (prain installed): python3 scripts/bench.py http://127.0.0.1:8000 [--tokens 128] [--runs 3]
 Each run sends a fresh prompt (a random prefix defeats the prefix cache) and decodes exactly
 --tokens tokens (ignore_eos), so numbers compare across runs and configs. The last line holds
 the medians. --context FILE... puts the files' text (up to --context-chars) before every
 prompt and the random tag after it: run 1 prefills the whole context, later runs find it in
 the prefix cache, like an agent's system prompt and tools.
 """
 import argparse
 import json
 import statistics
 import time
 import uuid
 from prain.chat import chat_body, first_model, post, stream_text
 PROMPTS = [
    "Explain how a CPU cache works, step by step, for a beginner.",
    "Write a short story about a lighthouse keeper who finds a message in a bottle.",
    "Schreibe eine kurze Anleitung, wie man Brot backt.",
 ]
 def run(base_url, model, prompt, tokens, context=""):
    text = f"{context}[{uuid.uuid4().hex[:8]}] {prompt}"
    body = chat_body(model, [{"role": "user", "content": text}], max_tokens=tokens,
                     temperature=0, ignore_eos=True)
    start = time.monotonic()
    stamps = []
    with post(base_url, "/v1/chat/completions", body, timeout=3600) as response:
        for _ in stream_text(response):
            stamps.append(time.monotonic())
    return {"ttft_s": round(stamps[0] - start, 3),
            "decode_tok_s": round((tokens - 1) / max(stamps[-1] - stamps[0], 1e-9), 3)}
 def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("base_url")
    parser.add_argument("--tokens", type=int, default=128)
    parser.add_argument("--runs", type=int, default=3)
    parser.add_argument("--context", nargs="*", default=[], help="files to put before each prompt")
    parser.add_argument("--context-chars", type=int, default=48000)
    args = parser.parse_args()
    model = first_model(args.base_url)
    context = ""
    for path in args.context:
        with open(path) as f:
            context += f.read()
    context = context[: args.context_chars] + "\n\n" if context else ""
    results = []
    for i in range(args.runs):
        results.append(run(args.base_url, model, PROMPTS[i % len(PROMPTS)], args.tokens, context))
        print(json.dumps(results[-1]), flush=True)
    print(json.dumps({key: statistics.median(r[key] for r in results) for key in results[0]}
                     | {"median_of": len(results), "tokens": args.tokens}))
 if __name__ == "__main__":
    main()
@@ -4,14 +4,16 @@
 #   sdk       venv + AMD's ROCm SDK wheels for $PRAIN_GPU_ARCH
 #   fetch     clone PyTorch and vLLM into build/
 #   torch     PyTorch from source -> build/wheels, installed into the venv
 #   torchvision  torchvision from source -> build/wheels, installed into the venv
 #   vllm      vLLM from source (adds $PRAIN_GPU_ARCH to its HIP arch list), installed editable
 #   prain     this package, installed editable, and its test suite
-# usage: scripts/build.sh STEP...      Versions: PRAIN_ROCM_VERSION, PRAIN_TORCH_REF, PRAIN_VLLM_REF
+# usage: scripts/build.sh STEP...   Versions: PRAIN_ROCM_VERSION, PRAIN_TORCH_REF, PRAIN_VISION_REF, PRAIN_VLLM_REF
 set -euo pipefail
 cd "$(dirname "$0")/.."
 ROCM_INDEX=${PRAIN_ROCM_INDEX:-https://stable.repo.amd.com/rocm/whl-next/}
 ROCM_VERSION=${PRAIN_ROCM_VERSION:-10.1.0}
 TORCH_REF=${PRAIN_TORCH_REF:-v2.13.0}
 VISION_REF=${PRAIN_VISION_REF:-v0.28.0}
 VLLM_REF=${PRAIN_VLLM_REF:-v0.32.0}
 PYTHON=${PRAIN_PYTHON:-python3.13}
@@ -90,9 +92,23 @@ torch() {
    env -u LD_LIBRARY_PATH python -c 'import torch; x = torch.ones(4, device="cuda"); print(torch.__version__, torch.version.hip, torch.cuda.get_device_properties(0).gcnArchName, (x * 2).sum().item(), torch.backends.openmp.is_available(), torch.get_num_threads())'
 }
 # transformers' Gemma 4 processor imports torchvision, even for text-only serving.
 torchvision() {
    [ -d build/vision/.git ] || git clone --depth 1 --branch "$VISION_REF" \
        https://github.com/pytorch/vision build/vision
    (cd build/vision && env FORCE_CUDA=1 PYTORCH_ROCM_ARCH="$PRAIN_GPU_ARCH" MAX_JOBS="$PRAIN_JOBS" \
        nice -n 19 python -m pip wheel --no-build-isolation --no-deps -v -w "$PRAIN_ROOT/build/wheels" .)
    python -m pip install --no-deps --force-reinstall build/wheels/torchvision-*.whl
    env -u LD_LIBRARY_PATH python -c 'import torch, torchvision; print(torchvision.__version__, torch.ops.torchvision.nms(torch.rand(4, 4, device="cuda"), torch.rand(4, device="cuda"), 0.5).device)'
 }
 vllm() {
-    local src=build/vllm site
+    local src=build/vllm site p
    site=$(python -c 'import sysconfig; print(sysconfig.get_paths()["purelib"])')
    # prain's fixes to vLLM (patches/vllm-*.patch), each applied once.
    for p in patches/vllm-*.patch; do
        git -C $src apply --reverse --check "$PRAIN_ROOT/$p" 2>/dev/null || git -C $src apply "$PRAIN_ROOT/$p"
    done
    # vLLM lists gfx1030 but not its siblings gfx1031/gfx1032: add ours.
    grep -q "$PRAIN_GPU_ARCH" $src/CMakeLists.txt ||
        sed -i "s/;gfx1030;/;gfx1030;$PRAIN_GPU_ARCH;/" $src/CMakeLists.txt
@@ -23,4 +23,7 @@ export PATH
 # libc++/compiler-rt defaults: everything links libstdc++/libgcc like the ROCm libraries.
 export CC=$PRAIN_ROOT/build/bin/clang CXX=$PRAIN_ROOT/build/bin/clang++
 export CFLAGS="-march=native" CXXFLAGS="-march=native"
 # Fused-MoE kernel configs tuned for this GPU by scripts/tune_moe.py (vLLM's defaults when
 # the files are missing).
 export VLLM_TUNED_CONFIG_FOLDER=$PRAIN_ROOT/build/moe-configs
 unset _site _dev _core
@@ -1,7 +1,7 @@
 """Measure the prain cost model on this machine and print it as JSON for
 `prain sim --cost` and `prain tune --cost`. Needs the GPU, vLLM and prain.
-usage: python3 scripts/measure.py PACK_FILE [--threads N] [--dense-gib 4.45] [--layers 30] > cost.json
+usage: python3 scripts/measure.py PACK_FILE [--quant int4] [--threads N] [--dense-gib 4.45] [--layers 30] > cost.json
 PACK_FILE is any <pack_dir>/layerNNN.w13 file; it is read cold (dropped from the page cache).
 """
@@ -15,6 +15,8 @@ from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
 from prain.config import QUANTS, expert_bytes as size_of
 from prain.quant import cpu_matmul, moe_quant_config, quantize, to_cpu_int4
 from prain.tiered import cpu_experts
 HIDDEN, INTER, SLOTS = 2816, 704, 5  # Gemma 4 26B-A4B expert, 5 slots per layer
@@ -37,10 +39,11 @@ def main():
    parser.add_argument("--threads", type=int, default=4)
    parser.add_argument("--dense-gib", type=float, default=4.45, help="non-expert weights")
    parser.add_argument("--layers", type=int, default=30)
    parser.add_argument("--quant", choices=QUANTS, default="bf16", help="expert weight format")
    args = parser.parse_args()
    torch.set_num_threads(args.threads)
-    expert_bytes = 3 * HIDDEN * INTER * 2
+    expert_bytes = size_of(HIDDEN, INTER, args.quant)
-    host = torch.empty(expert_bytes // 2, dtype=torch.bfloat16).pin_memory()
+    host = torch.empty(expert_bytes, dtype=torch.uint8).pin_memory()
    device = torch.empty_like(host, device="cuda")
    h2d_gbps = expert_bytes / 1e6 / ms(lambda: device.copy_(host, non_blocking=True))
@@ -56,13 +59,17 @@ def main():
    w1 = torch.randn(SLOTS, 2 * INTER, HIDDEN, dtype=torch.bfloat16, device="cuda") * 0.02
    w2 = torch.randn(SLOTS, HIDDEN, INTER, dtype=torch.bfloat16, device="cuda") * 0.02
    quant_config = None
    if args.quant != "bf16":
        (w1, s1), (w2, s2) = quantize(w1, args.quant), quantize(w2, args.quant)
        quant_config = moe_quant_config(args.quant, s1, s2)
    def moe(tokens):
        x = torch.randn(tokens, HIDDEN, dtype=torch.bfloat16, device="cuda")
        ids = torch.arange(SLOTS, device="cuda", dtype=torch.int32).repeat(tokens, 1)
        weights = torch.full((tokens, SLOTS), 1 / SLOTS, device="cuda")
        return lambda: fused_experts(x, w1, w2, weights, ids, activation=MoEActivation.GELU_TANH,
-                                     global_num_experts=SLOTS)
+                                     global_num_experts=SLOTS, quant_config=quant_config)
    with set_current_vllm_config(VllmConfig()):
        round_ms = ms(moe(1))
@@ -73,15 +80,28 @@ def main():
    other_ms = args.dense_gib * 2**30 / 1e6 / device_gbps / args.layers
    flag = torch.zeros(1, device="cuda")
    sync_ms = ms(lambda: flag.add_(1).item())
-    w13_cpu = torch.randn(2 * INTER, HIDDEN, dtype=torch.bfloat16).pin_memory()
+    # Serving computes a different expert each time, read from RAM (not from the CPU caches).
-    w2_cpu = torch.randn(HIDDEN, INTER, dtype=torch.bfloat16).pin_memory()
+    experts = [(torch.randn(2 * INTER, HIDDEN, dtype=torch.bfloat16),
                torch.randn(HIDDEN, INTER, dtype=torch.bfloat16)) for _ in range(23)]
    if args.quant != "bf16":  # the CPU kernels on pack-layout weights, as the plugin runs them
        def stored(w):
            q, scale = quantize(w, args.quant)
            return cpu_matmul(args.quant, to_cpu_int4(q) if args.quant == "int4" else q, scale)
        experts = [(stored(a), stored(b)) for a, b in experts]
    x = torch.randn(1, HIDDEN, dtype=torch.bfloat16)
    one = torch.ones(1, 1)
    ids = torch.zeros(1, 1, dtype=torch.int32)
    def cpu_route(weights):
        cpu_experts(x, one, ids, {0}, lambda e: weights)
    for weights in experts[:3]:  # warm up like ms(): thread pool and BLAS setup are one-time costs
        cpu_route(weights)
    start = time.perf_counter()
-    for _ in range(20):
+    for weights in experts[3:]:
-        cpu_experts(x, one, ids, {0}, lambda e: (w13_cpu, w2_cpu))
+        cpu_route(weights)
-    cpu_ms = (time.perf_counter() - start) * 1000 / 20
+    cpu_ms = (time.perf_counter() - start) * 1000 / len(experts[3:])
    print(json.dumps({"h2d_gbps": round(h2d_gbps, 2), "disk_gbps": round(disk_gbps, 2),
                      "round_ms": round(round_ms, 4), "route_ms": round(route_ms, 6),
                      "other_ms": round(other_ms, 4), "sync_ms": round(sync_ms, 4),
@@ -0,0 +1,100 @@
 """Answer quality of a served config against reference answers from another config.
    python scripts/quality.py URL make REF.json    greedy reference answers (run on the exact config)
    python scripts/quality.py URL score REF.json   how likely this server finds those answers
 Both work on token ids: the prompt is the chat template with its generation prompt, and `score`
 re-scores the reference answer on exactly the ids it was generated from (prompt logprobs). A
 finished assistant message renders differently from the generation prompt, so re-scoring
 through the chat API skews the logprobs. `nll` is the mean negative log-likelihood per answer
 token (the reference config scores about its own decode logprobs); `same_tokens` is how many
 leading tokens this server's own greedy answer shares with the reference.
 """
 import argparse
 import json
 import urllib.request
 PROMPTS = [
    "Write a Python function that returns the n-th Fibonacci number iteratively.",
    "Explain what `git rebase -i HEAD~3` does, in three sentences.",
    "Write a bash one-liner that counts the lines of all .py files below the current directory.",
    "On Gentoo, how do I rebuild every package that links against a library I just updated?",
    "What is the difference between a process and a thread? Answer in two sentences.",
 ]
 CODE_PROMPT = "Here is some Python code:\n\n{code}\n\nWhat does this code do? Answer in two sentences."
 def post(url, path, body):
    request = urllib.request.Request(url + path, json.dumps(body).encode(),
                                     {"Content-Type": "application/json"})
    with urllib.request.urlopen(request, timeout=3600) as response:
        return json.load(response)
 def greedy(url, model, ids, tokens):
    """(text, token ids, their logprobs) of a greedy answer."""
    choice = post(url, "/v1/completions", {
        "model": model, "prompt": ids, "max_tokens": tokens, "temperature": 0, "logprobs": 1,
        "return_tokens_as_token_ids": True})["choices"][0]
    logprobs = choice["logprobs"]
    return (choice["text"], [int(t.split(":")[1]) for t in logprobs["tokens"]],
            logprobs["token_logprobs"])
 def nll(url, model, ids, answer):
    """Mean negative log-likelihood of the answer tokens after the prompt ids."""
    entries = post(url, "/v1/completions", {
        "model": model, "prompt": ids + answer, "max_tokens": 1, "temperature": 0,
        "prompt_logprobs": 0})["choices"][0]["prompt_logprobs"]
    return -sum(entries[len(ids) + i][str(t)]["logprob"] for i, t in enumerate(answer)) / len(answer)
 def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("url")
    parser.add_argument("action", choices=("make", "score"))
    parser.add_argument("ref")
    parser.add_argument("--model", default="gemma4")
    parser.add_argument("--tokens", type=int, default=48, help="answer tokens")
    parser.add_argument("--code", default="prain/policy.py", help="file for the code question")
    parser.add_argument("--code-chars", type=int, default=900)
    parser.add_argument("--greedy-only", action="store_true",
                        help="score: skip the nll (prompt logprobs need ~0.5 GiB of free VRAM "
                             "for the logits); compare greedy answers only")
    args = parser.parse_args()
    if args.action == "make":
        with open(args.code) as f:
            prompts = PROMPTS + [CODE_PROMPT.format(code=f.read()[: args.code_chars])]
        refs = []
        for prompt in prompts:
            ids = post(args.url, "/tokenize", {"model": args.model, "add_generation_prompt": True,
                                               "messages": [{"role": "user", "content": prompt}]})
            ids = ids["tokens"]
            text, answer, _ = greedy(args.url, args.model, ids, args.tokens)
            refs.append({"prompt": prompt, "ids": ids, "answer": answer, "text": text})
            print(json.dumps({"prompt": prompt[:60], "answer": text[:120]}), flush=True)
        with open(args.ref, "w") as f:
            json.dump(refs, f)
        return
    with open(args.ref) as f:
        refs = json.load(f)
    scores, same, own = [], [], []
    for ref in refs:
        if not args.greedy_only:
            scores.append(nll(args.url, args.model, ref["ids"], ref["answer"]))
        text, answer, logprobs = greedy(args.url, args.model, ref["ids"], len(ref["answer"]))
        own.append(sum(logprobs) / len(logprobs))
        same.append(next((i for i, (a, b) in enumerate(zip(answer, ref["answer"])) if a != b),
                         min(len(answer), len(ref["answer"]))))
        print(json.dumps({"prompt": ref["prompt"][:40],
                          "nll": round(scores[-1], 3) if scores else None,
                          "own_lp": round(own[-1], 3), "same_tokens": same[-1],
                          "answer": text[:100]}), flush=True)
    print(json.dumps({"mean_nll": round(sum(scores) / len(scores), 3) if scores else None,
                      "mean_own_lp": round(sum(own) / len(own), 3), "same_tokens": same,
                      "identical": sum(s == len(r["answer"]) for s, r in zip(same, refs))}))
 if __name__ == "__main__":
    main()
@@ -0,0 +1,29 @@
 #!/bin/bash
 # Go/no-go: does this vLLM build run Gemma 4 on this GPU? Serves a 2-layer copy of MODEL_DIR with
 # random weights (port 8000, log in build/smoke.log), then sends a greedy request, the warm-up
 # shape (continue_final_message, priority) and the draft fields (vllm_xargs, cache_salt).
 # usage: scripts/smoke.sh MODEL_DIR [vllm serve args...]
 # Expected: READY, a JSON list of 3 (gibberish) strings, two {"id":"chatcmpl-... lines, STOPPED.
 # RDNA2 needs --attention-backend TRITON_ATTN: ROCM_ATTN's prefill kernel wants 160 KiB of LDS.
 cd "$(dirname "$0")/.." && . scripts/env.sh
 model=$1
 shift
 setsid vllm serve "$model" --served-model-name gemma4 --language-model-only --max-model-len 4096 \
    --max-num-batched-tokens 1024 --max-num-seqs 2 --kv-cache-memory-bytes 805306368 --port 8000 \
    --enforce-eager --moe-backend triton --scheduling-policy priority --load-format dummy \
    --hf-overrides '{"text_config": {"num_hidden_layers": 2, "layer_types": ["sliding_attention", "full_attention"]}}' \
    "$@" > build/smoke.log 2>&1 < /dev/null &
 server=$!
 for _ in $(seq 120); do
    curl -sf localhost:8000/health > /dev/null && break
    kill -0 $server 2> /dev/null || break
    sleep 5
 done
 curl -sf localhost:8000/health > /dev/null && echo READY || { echo DIED; tail -40 build/smoke.log; exit 1; }
 python scripts/greedy.py http://127.0.0.1:8000 gemma4 | head -c 200; echo
 post() { curl -s localhost:8000/v1/chat/completions -H 'Content-Type: application/json' -d "$1" | head -c 60; echo; }
 post '{"model": "gemma4", "messages": [{"role": "user", "content": "Why is the sky"}], "max_tokens": 1, "add_generation_prompt": false, "continue_final_message": true, "priority": 1}'
 post '{"model": "gemma4", "messages": [{"role": "user", "content": "Hi"}], "max_tokens": 2, "vllm_xargs": {"prain": "draft"}, "cache_salt": "prain-draft"}'
 kill -INT -$server
 while curl -sf localhost:8000/health > /dev/null; do sleep 2; done
 echo STOPPED
@@ -0,0 +1,111 @@
 """Tune vLLM's Triton fused-MoE kernel for this GPU on prain's slot tensors.
 usage: python scripts/tune_moe.py OUT_DIR [FORMAT...]   (scripts/env.sh points
                                    VLLM_TUNED_CONFIG_FOLDER at build/moe-configs)
 Times candidate configs on one prain round (16 slots of 128 routed experts, expert_map;
 up to 32 tokens route into the slots like decode, more tokens spread like prefill),
 Gemma 4 expert shapes, bf16/int8/int4 slots, and writes the fastest per batch size as vLLM
 config files for every slot count E (the best config hardly depends on E). Every config in
 the files compiled, ran and matched the default config's output on this GPU.
 """
 import itertools
 import json
 import os
 import sys
 import time
 import torch
 from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.model_executor.layers.fused_moe import override_config
 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
 from vllm.model_executor.layers.fused_moe.config import (int4_w4a16_moe_quant_config,
                                                         int8_w8a16_moe_quant_config)
 from vllm.model_executor.layers.fused_moe.fused_moe import (fused_experts, get_config_file_name,
                                                            get_default_config)
 from prain.quant import GROUP, quantize
 H, I, TOPK, EXPERTS, SLOTS = 2816, 704, 8, 128, 16
 SIZES = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048)
 FORMATS = {"bf16": (None, None), "int8": ("int8_w8a16", None), "int4": ("int4_w4a16", [0, GROUP])}
 def candidates(tokens):
    # ponytail: ~8 s per Triton compile on this host, so a small space (~70 compiles per
    # format); widen it when a GPU compiles faster.
    if tokens <= 32:
        space = itertools.product((16,), (32, 64, 128), (64, 128, 256), (1,), (2, 4), (2,))
    else:
        space = itertools.product((32, 64), (64, 128), (32, 64), (8,), (4, 8), (2,))
    for bm, bn, bk, group, warps, stages in space:
        yield {"BLOCK_SIZE_M": bm, "BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk, "GROUP_SIZE_M": group,
               "SPLIT_K": 1, "num_warps": warps, "num_stages": stages}  # int4 needs SPLIT_K
 def timed(fn, repeat):
    fn()
    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(repeat):
        fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - start) * 1000 / repeat
 def main(out_dir, formats):
    os.makedirs(out_dir, exist_ok=True)
    torch.manual_seed(0)
    w13 = torch.randn(SLOTS, 2 * I, H, device="cuda").mul_(0.02).bfloat16()
    w2 = torch.randn(SLOTS, H, I, device="cuda").mul_(0.02).bfloat16()
    table = torch.full((EXPERTS + 1,), -1, dtype=torch.int32, device="cuda")
    table[:SLOTS] = torch.arange(SLOTS, dtype=torch.int32, device="cuda")
    for name in formats:
        dtype, block_shape = FORMATS[name]
        if name == "bf16":
            a, b, quant_config = w13, w2, None
        else:
            (a, s13), (b, s2) = quantize(w13, name), quantize(w2, name)
            make = int8_w8a16_moe_quant_config if name == "int8" else int4_w4a16_moe_quant_config
            quant_config = make(w1_scale=s13, w2_scale=s2, w1_zp=None, w2_zp=None,
                                block_shape=block_shape)
        best = {}
        for tokens in SIZES:
            x = torch.randn(tokens, H, device="cuda").mul_(0.5).bfloat16()
            pool = SLOTS if tokens <= 32 else EXPERTS
            ids = torch.stack([torch.randperm(pool, device="cuda")[:TOPK]
                               for _ in range(tokens)]).to(torch.int32)
            weights = torch.softmax(torch.randn(tokens, TOPK, device="cuda"), -1)
            def run():
                return fused_experts(x, a, b, weights, ids, activation=MoEActivation.GELU_TANH,
                                     global_num_experts=EXPERTS + 1, expert_map=table,
                                     quant_config=quant_config)
            default = get_default_config(tokens, SLOTS, I, a.shape[2], TOPK, dtype, block_shape)
            with override_config(default):  # vLLM's own choice when no config file exists
                default_ms = timed(run, 10 if tokens <= 32 else 3)
                reference = run()
            results = [(default_ms, json.dumps(default, sort_keys=True))]  # never worse than it
            for config in candidates(tokens):
                # catch inside: vLLM's override_config skips its restore when the body raises
                with override_config(config):
                    try:
                        ms = timed(run, 10 if tokens <= 32 else 3)
                        out = run()
                    except Exception:  # noqa: BLE001 - configs that do not compile or fit
                        continue
                if (out.float() - reference.float()).norm() <= 1e-2 * reference.float().norm() + 1e-6:
                    results.append((ms, json.dumps(config, sort_keys=True)))
            ms, config = min(results)
            best[str(tokens)] = json.loads(config)
            print(f"{name} tokens={tokens:4d}: default {default_ms:8.3f} ms, best {ms:8.3f} ms "
                  f"({len(results)} configs) {config}", flush=True)
        for e in range(1, EXPERTS + 1):  # any slot count a placement can give a layer
            with open(os.path.join(out_dir, get_config_file_name(e, I, dtype, block_shape)), "w") as f:
                json.dump(best, f, indent=1)
 if __name__ == "__main__":
    with set_current_vllm_config(VllmConfig()):
        main(sys.argv[1], sys.argv[2:] or list(FORMATS))
@@ -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")
@@ -76,3 +97,17 @@ class Cli(TmpCase):
        self.assertEqual((tuned.pack_dir, tuned.trace_path), ("/packs", trace))
        self.assertLessEqual(sum(tuned.placement(layer, 128).cache_slots for layer in range(3)),
                             int(0.2 * 1024 / 11.34) // 3 * 3)
    def test_quantized_experts_check_and_tune(self):
        hf = self.write("config.json", json.dumps(GEMMA4))
        out, _ = self.run_cli("check", self.write("q.toml", 'pack_dir = "/packs"\nquant = "int4"\n'), hf)
        self.assertIn("expert size 3.01 MiB\n", out)
        trace = os.path.join(self.dir, "trace.jsonl.gz")
        self.run_cli("synth", trace, "--forwards", "40", "--layers", "3")
        out, _ = self.run_cli("tune", trace, hf, "--quant", "int4", "--expert-vram-gib", "0.2",
                              "--ram-gib", "4", "--pack-dir", "/packs", "--top", "3")
        tuned = load_config(self.write("tuned.toml", out))
        self.assertEqual(tuned.quant, "int4")
        slots = sum(tuned.placement(layer, 128).cache_slots for layer in range(3))
        self.assertGreater(slots, int(0.2 * 1024 / 11.34))  # int4 experts are ~4x smaller
        self.assertLessEqual(slots, int(0.2 * 1024 / 3.01))
@@ -107,6 +107,23 @@ mode = "auto"
        self.assertEqual(config.stats_path, os.path.expanduser("~/s.json"))
        self.assertEqual(config.trace_path, "/abs/t.gz")
    def test_miss_min_weight(self):
        config = mc.load_config(self.write("m.toml", 'pack_dir = "/p"\n[miss]\nmin_weight = 0.7\n'))
        self.assertEqual(config.miss.min_weight, 0.7)
        self.assertEqual(mc.Miss().min_weight, 1.0)  # default: compute every routed expert
        with self.assertRaisesRegex(ValueError, "min_weight must be a number in"):
            mc.load_config(self.write("m2.toml", 'pack_dir = "/p"\n[miss]\nmin_weight = 1.5\n'))
    def test_quant_key_and_expert_sizes(self):
        self.assertEqual(mc.load_config(self.write("q.toml", 'pack_dir = "/p"\nquant = "int4"\n')).quant, "int4")
        self.assertEqual(mc.load_config(self.write("b.toml", 'pack_dir = "/p"\n')).quant, "bf16")
        with self.assertRaisesRegex(ValueError, "quant must be one of"):
            mc.load_config(self.write("bad.toml", 'pack_dir = "/p"\nquant = "int3"\n'))
        h, i = 2816, 704
        self.assertEqual(mc.expert_bytes(h, i), 3 * h * i * 2)
        self.assertEqual(mc.expert_bytes(h, i, "int8"), 3 * h * i + 4 * (2 * i + h))  # fp32 row scales
        self.assertEqual(mc.expert_bytes(h, i, "int4"), 3 * h * i // 2 + 2 * 3 * h * i // 64)
    def test_serve_and_draft_tables(self):
        config = mc.load_config(self.write("s.toml", 'pack_dir = "/p"\n[serve]\nmax_requests = 4\n'
                                                     'kv_cache_gib = 1.5\n[draft]\nfull_requests = 2\n'
@@ -0,0 +1,55 @@
 import unittest
 import torch
 import torch.nn.functional as F
 from prain import gemv
 GPU = torch.cuda.is_available() and gemv.fast_device(0)
 def rel(a, b):
    return ((a.float() - b.float()).norm() / b.float().norm()).item()
@unittest.skipUnless(GPU, "needs a gfx10 GPU")
 class TestGemv(unittest.TestCase):
    def test_up_to_four_rows_match_a_float32_linear(self):
        torch.manual_seed(0)
        for n, k in ((1000, 2816), (2816, 2112), (130, 70)):
            w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") * 0.02
            for rows in (1, 2, 3, 4):
                x = torch.randn(rows, k, dtype=torch.bfloat16, device="cuda")
                y = gemv.linear(x, w)
                self.assertEqual((y.shape, y.dtype), ((rows, n), torch.bfloat16))
                self.assertLess(rel(y, F.linear(x.float(), w.float())), 1e-2, (n, k, rows))
        x = torch.randn(2, 1, 70, dtype=torch.float16, device="cuda")  # leading dims, fp16
        w = torch.randn(130, 70, dtype=torch.float16, device="cuda")
        self.assertEqual(gemv.linear(x, w).shape, (2, 1, 130))
        self.assertLess(rel(gemv.linear(x, w), F.linear(x.float(), w.float())), 1e-2)
    def test_more_rows_bias_or_strided_weights_use_f_linear(self):
        w = torch.randn(64, 96, dtype=torch.bfloat16, device="cuda")
        bias = torch.randn(64, dtype=torch.bfloat16, device="cuda")
        x = torch.randn(5, 96, dtype=torch.bfloat16, device="cuda")
        self.assertTrue(torch.equal(gemv.linear(x, w), F.linear(x, w)))
        self.assertTrue(torch.equal(gemv.linear(x[:1], w, bias), F.linear(x[:1], w, bias)))
        strided = torch.randn(96, 64, dtype=torch.bfloat16, device="cuda").t()
        self.assertTrue(torch.equal(gemv.linear(x[:1], strided), F.linear(x[:1], strided)))
    def test_wrapped_dispatch_sends_only_decode_shapes_to_the_kernel(self):
        calls = []
        def fallback(layer, x, weight, bias=None):
            calls.append(x.shape[0])
            return F.linear(x, weight, bias)
        gemm = gemv.wrap(lambda *args, **kwargs: fallback)("auto")
        w = torch.randn(64, 96, dtype=torch.bfloat16, device="cuda")
        gemm(None, torch.randn(1, 96, dtype=torch.bfloat16, device="cuda"), w)
        gemm(None, torch.randn(9, 96, dtype=torch.bfloat16, device="cuda"), w)
        self.assertEqual(calls, [9])
 if __name__ == "__main__":
    unittest.main()
@@ -31,3 +31,64 @@ class KernelMatch(unittest.TestCase):
                                global_num_experts=experts).float().cpu()
        cpu = cpu_experts(x, topk_weights, topk_ids, set(range(experts)), lambda e: (w13[e], w2[e]))
        torch.testing.assert_close(cpu, gpu, rtol=3e-2, atol=3e-2)
    def test_a_round_computes_only_the_experts_its_slot_table_maps(self):
        from vllm.config import VllmConfig, set_current_vllm_config
        from vllm.model_executor.layers.fused_moe.activation import MoEActivation
        from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
        from prain.tiered import cpu_experts, slot_table
        torch.manual_seed(0)
        tokens, hidden, inter, experts = 64, 256, 128, 6
        x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
        w13 = torch.randn(experts, 2 * inter, hidden, dtype=torch.bfloat16) * 0.1
        w2 = torch.randn(experts, hidden, inter, dtype=torch.bfloat16) * 0.1
        topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int()
        topk_ids[0, 1] = experts  # a padding route
        topk_weights = torch.rand(tokens, 2)
        table = slot_table({1: 0, 4: 1}, experts, "cuda")
        with set_current_vllm_config(VllmConfig()):
            gpu = fused_experts(x.cuda(), w13[[1, 4]].cuda(), w2[[1, 4]].cuda(), topk_weights.cuda(),
                                topk_ids.cuda(), activation=MoEActivation.GELU_TANH,
                                global_num_experts=table.numel(), expert_map=table).float().cpu()
        cpu = cpu_experts(x, topk_weights, topk_ids, {1, 4}, lambda e: (w13[e], w2[e]))
        torch.testing.assert_close(cpu, gpu, rtol=3e-2, atol=3e-2)
    def test_quantized_slots_compute_what_the_cpu_path_computes(self):
        """Pack layout -> slot layout on the GPU -> vLLM's int8/int4 kernel for a round that
        holds experts 1 and 3, against the CPU kernels on the pack layout and against the
        dequantized weights."""
        from vllm.config import VllmConfig, set_current_vllm_config
        from vllm.model_executor.layers.fused_moe.activation import MoEActivation
        from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
        from prain.quant import (cpu_matmul, dequantize, moe_quant_config, quantize, to_cpu_int4,
                                 to_gpu_int4)
        from prain.tiered import cpu_experts, slot_table
        torch.manual_seed(0)
        tokens, hidden, inter, experts = 8, 256, 128, 4
        x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
        w13 = torch.randn(experts, 2 * inter, hidden) * 0.1
        w2 = torch.randn(experts, hidden, inter) * 0.1
        topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int()
        topk_weights = torch.rand(tokens, 2)
        table = slot_table({1: 0, 3: 1}, experts, "cuda")
        for quant in ("int8", "int4"):
            parts = [quantize(w, quant) for w in (w13, w2)]  # [q, scale], gpu layout
            pack = [[to_cpu_int4(q) if quant == "int4" else q, s] for q, s in parts]
            slots = [[to_gpu_int4(q[[1, 3]].cuda()) if quant == "int4" else q[[1, 3]].cuda(),
                      s[[1, 3]].cuda()] for q, s in pack]
            with set_current_vllm_config(VllmConfig()):
                gpu = fused_experts(x.cuda(), slots[0][0], slots[1][0], topk_weights.cuda(),
                                    topk_ids.cuda(), activation=MoEActivation.GELU_TANH,
                                    global_num_experts=table.numel(), expert_map=table,
                                    quant_config=moe_quant_config(quant, slots[0][1], slots[1][1]))
            cpu = cpu_experts(x, topk_weights, topk_ids, {1, 3}, lambda e: tuple(
                cpu_matmul(quant, q[e], s[e]) for q, s in pack))
            exact = cpu_experts(x, topk_weights, topk_ids, {1, 3}, lambda e: tuple(
                dequantize(q, s, quant)[e] for q, s in parts))
            self.assertGreater(exact.norm().item(), 1.0)
            torch.testing.assert_close(cpu, exact, rtol=3e-2, atol=3e-2, msg=quant)
            torch.testing.assert_close(gpu.float().cpu(), exact, rtol=3e-2, atol=3e-2, msg=quant)
@@ -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))
@@ -0,0 +1,49 @@
 import unittest
 try:
    import torch
 except ImportError:  # the planning host has no torch
    torch = None
 if torch is not None:
    from prain.quant import (GROUP, cpu_matmul, dequantize, layout, quantize, to_cpu_int4,
                             to_gpu_int4)
@unittest.skipIf(torch is None, "needs torch")
 class Quant(unittest.TestCase):
    def setUp(self):
        torch.manual_seed(0)
        self.w = torch.randn(64, 256) * 0.02
    def test_layouts_match_what_quantize_returns(self):
        for quant in ("int8", "int4"):
            got = [(tuple(t.shape), t.dtype) for t in quantize(self.w, quant)]
            self.assertEqual(got, layout(quant, 64, 256))
    def test_round_trip_error(self):
        for quant, bound in (("int8", 0.01), ("int4", 0.11)):
            err = (dequantize(*quantize(self.w, quant), quant, torch.float32) - self.w).norm()
            self.assertLess(err / self.w.norm(), bound)
    def test_int4_clipping_beats_plain_round_to_nearest(self):
        g = self.w.unflatten(-1, (-1, GROUP))
        scale = (g.abs().amax(-1, keepdim=True) / 7).bfloat16().float()
        plain = ((g / scale).round() * scale - g).norm()
        clipped = (dequantize(*quantize(self.w, "int4"), "int4", torch.float32) - self.w).norm()
        self.assertLess(clipped, plain)
    def test_cpu_layout_is_torch_avx2_int4pack_and_inverts(self):
        q, _ = quantize(self.w, "int4")
        nibbles = torch.stack((q & 15, q >> 4), -1).flatten(-2).to(torch.int32)
        self.assertTrue(torch.equal(to_cpu_int4(q), torch._convert_weight_to_int4pack_for_cpu(nibbles, 1)))
        self.assertTrue(torch.equal(to_gpu_int4(to_cpu_int4(q)), q))
    def test_cpu_matmul_matches_dequantized_weights(self):
        x = torch.randn(3, 256, dtype=torch.bfloat16)
        for quant in ("int8", "int4"):
            q, scale = quantize(self.w, quant)
            ref = x.float() @ dequantize(q, scale, quant, torch.float32).t()
            stored = to_cpu_int4(q) if quant == "int4" else q
            got = cpu_matmul(quant, stored, scale)(x).float()
            self.assertLess(((got - ref).norm() / ref.norm()).item(), 0.01, quant)
@@ -57,6 +57,14 @@ class Simulate(unittest.TestCase):
        copy = self.run_sim().ms_per_decode
        self.assertLess(self.run_sim(miss=Miss(mode="auto")).ms_per_decode, copy)
    def test_cpu_path_pays_the_disk_read_of_disk_experts(self):
        on_disk = {layer: LayerPlacement((), (), tuple(range(32)), 3) for layer in range(4)}
        cpu = Miss(mode="cpu", admit=0)
        ram = simulate(self.forwards, all_ram(3), Policy(), Prefetch(), cpu, Cost())
        disk = simulate(self.forwards, on_disk, Policy(), Prefetch(), cpu, Cost())
        self.assertAlmostEqual(disk.total_ms - ram.total_ms, ram.counters["cpu"] * Cost().read_ms(),
                               places=3)
    def test_prefetch_helps_when_predictions_are_right(self):
        none = self.run_sim().hit_rate
        good = self.run_sim(prefetch=Prefetch(mode="pregate", budget=4))
@@ -81,3 +89,13 @@ class Tune(unittest.TestCase):
        self.assertLessEqual(final.total_ms, default.total_ms)
        self.assertEqual(rows[0][1].total_ms, min(r.total_ms for _, r in rows))
        self.assertGreaterEqual(bound.hit_rate, 0.0)
    def test_tune_searches_the_auto_threshold_and_keeps_the_cpu_threads(self):
        forwards = forwards_of(synthetic_trace(**SMALL))
        cost = Cost(expert_mib=64.0, h2d_gbps=2.5, cpu_ms=0.5)  # eGPU: a copy is ~50 CPU routes
        base = (Policy(), Prefetch(), Miss(cpu_threads=16))
        best, ram, per, final, bound, rows = tune(forwards, 32, 1.0, 1.0, cost, base=base)
        self.assertEqual(best[2].cpu_threads, 16)
        self.assertNotEqual(best[2].mode, "copy")
        tried = {(v[2].mode, v[2].cpu_max_tokens) for v, _ in rows}
        self.assertLessEqual({("copy", 1), ("cpu", 1), ("auto", 1), ("auto", 8), ("auto", 16)}, tried)
@@ -9,7 +9,7 @@ from prain.config import Miss, Policy
 from prain.policy import SlotCache
 if torch is not None:
-    from prain.tiered import (available_experts, cpu_experts, degrade, remap, route_stats,
+    from prain.tiered import (available_experts, cpu_experts, degrade, route_stats, slot_table,
                                tiered_forward)
 HIDDEN, INTER, EXPERTS = 6, 4, 10
@@ -52,8 +52,9 @@ class TieredForward(unittest.TestCase):
            loads.append(expert)
            slots13[slot], slots2[slot] = self.w13[expert], self.w2[expert]
-        def kernel(x, w, ids, used):
+        def kernel(x, w, ids, table, used):
-            return naive_moe(x, slots13, slots2, w, ids)
+            self.assertTrue(bool((ids >= 0).all()))  # padding id -1 arrives as num_experts
            return naive_moe(x, slots13, slots2, w, table[ids.long()])
        out = tiered_forward(self.x, self.topk_weights, self.topk_ids, cache, load, kernel,
                             miss=miss, cpu_weights=lambda e: (self.w13[e], self.w2[e]),
@@ -80,6 +81,33 @@ class TieredForward(unittest.TestCase):
        self.assertEqual(cpu_experts(self.x, self.topk_weights, self.topk_ids, set(),
                                     None).abs().sum().item(), 0.0)
    def test_miss_min_weight_skips_missing_experts_for_every_token(self):
        """[miss] min_weight < 1: every token skips experts that are not in VRAM while it keeps
        that share of its router weight; experts some token copies anyway serve all tokens."""
        everyone = torch.ones(6, dtype=torch.bool)
        resident = torch.zeros(EXPERTS, dtype=torch.bool)
        resident[[0, 1, 2]] = True
        _, first = degrade(self.topk_weights, self.topk_ids, everyone, resident, 0.5)
        free = resident.clone()
        free[first[first >= 0].long()] = True
        w, i = degrade(self.topk_weights, self.topk_ids, everyone, free, 0.5)
        out, loads, _ = self.run_tiered(pinned=[0, 1, 2], cache_slots=3, miss=Miss(min_weight=0.5))
        torch.testing.assert_close(out, naive_moe(self.x, self.w13, self.w2, w, i))
        self.assertEqual(sorted(loads), sorted(set(i[i >= 0].tolist()) - {0, 1, 2}))
        self.assertLess(len(loads), 7)  # computing everything copies all 7 others
        self.assertEqual(self.seen_experts, list(range(EXPERTS)))  # the trace sees the routing
        out, loads, _ = self.run_tiered(pinned=[0, 1, 2], cache_slots=3, miss=Miss(min_weight=1.0))
        torch.testing.assert_close(out, self.expected)
        self.assertEqual(len(loads), 7)
    def test_cpu_experts_take_matmul_functions(self):
        """Quantized experts hand the CPU path x -> x @ w.T functions instead of weights."""
        def matmuls(e):
            return (lambda v: v @ self.w13[e].t()), (lambda v: v @ self.w2[e].t())
        out = cpu_experts(self.x, self.topk_weights, self.topk_ids, set(range(EXPERTS)), matmuls)
        torch.testing.assert_close(out, self.expected.float())
    def test_route_stats(self):
        stats = route_stats(torch.tensor([[0.5, 0.25], [1.0, 0.0]]),
                            torch.tensor([[2, 1], [2, -1]]), 4, recency_half_life=1.0)
@@ -89,10 +117,10 @@ class TieredForward(unittest.TestCase):
        self.assertAlmostEqual(stats[2][2], 0.5 * 0.5 + 1.0)  # the first token is one half-life old
        self.assertAlmostEqual(stats[1][2], 0.125)
-    def test_remap_zeroes_routes_outside_the_round(self):
+    def test_slot_table_maps_the_round_and_skips_the_rest(self):
-        weights, slots = remap(torch.ones(1, 3), torch.tensor([[2, 7, -1]]), {7: 4}, 10)
+        table = slot_table({7: 4, 2: 0}, 10, "cpu")
-        self.assertEqual(weights.tolist(), [[0.0, 1.0, 0.0]])
+        self.assertEqual(table.tolist(), [-1, -1, 0, -1, -1, -1, -1, 4, -1, -1, -1])
-        self.assertEqual(slots.tolist(), [[0, 4, 0]])
+        self.assertEqual(table.dtype, torch.int32)  # one extra row: padding id 10 -> -1
    def test_degrade_drops_unavailable_routes_of_draft_tokens_only(self):
        weights = torch.tensor([[0.5, 0.3, 0.2], [0.5, 0.3, 0.2]])
@@ -115,6 +143,8 @@ class TieredForward(unittest.TestCase):
        available[[0, 1, 2, 4, 5, 6, 7, 8]] = True
        self.assertTrue(torch.equal(
            available_experts(self.topk_ids, draft, {0}, EXPERTS), available))
        _, kept = degrade(self.topk_weights, self.topk_ids, draft, available, 0.5)
        available[kept[kept >= 0].long()] = True  # experts some row copies anyway serve all rows
        expected = naive_moe(self.x, self.w13, self.w2,
                             *degrade(self.topk_weights, self.topk_ids, draft, available, 0.5))
        torch.testing.assert_close(out, expected)
@@ -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
fabolous005	bb42ea3802	prain: README: expert formats, dense-layer fix, min_weight trade-off, 16k agent setup; example config int4, 16k, tools	2026-10-28 22:07:25 +01:00
fabolous005	216cd76a77	prain: Triton GEMV for decode GEMMs on gfx10 (rocBLAS ran bf16 GEMVs at 6% of the bandwidth; agent setup 1.58 -> 2.77 tok/s, same answers)	2026-10-28 22:07:25 +01:00
fabolous005	708e152549	prain: int8/int4 experts quantized once on the GPU, rounds through vLLM's expert_map, [miss] min_weight skips experts not in VRAM; scripts: tune_moe, quality (answer NLL on identical token ids), bench --context, measure --quant	2026-10-28 22:07:25 +01:00
fabolous005	3ae9bd73a7	prain: README; example config with the measured eGPU defaults (miss auto <=2 tokens, 16 CPU threads, no pinned RAM tier)	2026-10-28 01:10:17 +01:00
fabolous005	16c3bf81f7	prain: pin the RAM tier per expert (big pinned blocks take minutes on eGPU hosts); sim charges disk reads on the CPU path; measure CPU cost on cold experts	2026-10-28 00:33:34 +01:00
fabolous005	01dce2ffc6	prain: tune searches the auto CPU threshold and writes --cpu-threads; vLLM patch: fit Gemma 4 in 8 GB (tied lm_head without a copy, rope caches sized by max_model_len)	2026-10-27 23:52:53 +01:00
fabolous005	d69fd34492	prain: fix vLLM on ROCm (no CUDA capability families), smoke test and benchmark scripts	2026-10-27 16:44:40 +01:00
fabolous005	7be39c9287	prain: build torchvision too (transformers Gemma 4 processor imports it)	2026-10-27 16:04:22 +01:00