# 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`, and it takes `--max-num-seqs` and `--kv-cache-memory-bytes` from
the file's `[serve]` table (it refuses them after `--`). 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) |