prain/README.md

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