prain

Serve a mixture-of-experts model with vLLM when its experts do not fit in VRAM. A TOML file places each expert of each layer in VRAM (pinned or in a cache slot), in RAM, or in a pack file on disk. Missing experts are copied to the GPU or computed on the CPU ([miss]). Eviction, heat, prefetch and miss policies are configurable, and prain tune picks them from a routing trace. Experts can be served as bf16, int8 or int4 (quant). Target: Gemma 4 26B-A4B (30 layers x 128 experts, 42.5 GiB of bf16 experts) on an 8 GB RX 6600 XT.

Build (GPU host, user space)

scripts/build.sh wrappers sdk fetch torch torchvision vllm prain

• sdk: venv with AMD's ROCm SDK wheels for PRAIN_GPU_ARCH (default gfx1032)
• torch, torchvision: from source with clang and -march=native (hours)
• vllm: v0.32.0 from source with gfx1032 added and patches/vllm-*.patch applied
• prain: editable install and the test suite

Check the stack: scripts/smoke.sh MODEL_DIR --attention-backend TRITON_ATTN

Tune vLLM's fused-MoE kernel for the GPU once (~45 min of Triton compiles the first time; scripts/env.sh points vLLM at the files). It keeps vLLM's default wherever no candidate beats it:

python scripts/tune_moe.py build/moe-configs

Serve

. scripts/env.sh
prain serve examples/gemma4-26b-a4b-rx6600xt.toml MODEL_DIR -- --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

The first start writes the bf16 pack (42.5 GiB, ~3 min); with quant = "int8" or "int4" it then quantizes that pack once on the GPU (~3 min). Once the int4 pack is written, the bf16 files (layerNNN.w13, .w2, .done) are not read and can be deleted; a later bf16 or int8 start writes them again. Chat with prain tui, or in the browser with prain web --config FILE: several users with drafts, and a monitor beside the chat (tokens/s, a live map of which experts are in VRAM and working, VRAM hit rate, context use, GPU load, temperature and power; the map needs stats_path in the placement file). --host 0.0.0.0 serves it to other machines; it has no login. Tool calls work through vLLM's gemma4 parser, so coding agents and MCP clients can use the OpenAI-compatible API.

Expert formats

quant	expert	all experts	weights	VRAM kernel	CPU kernel
bf16	11.34 MiB	42.5 GiB	as trained	vLLM Triton	torch bf16 matmul
int8	5.69 MiB	21.3 GiB	fp32 scale per row	vLLM Triton (w8a16)	torch._weight_int8pack_mm
int4	3.01 MiB	11.3 GiB	bf16 scale per 64, MSE clip	vLLM Triton (GPTQ/AWQ, w4a16); prain's on RDNA2	torch._weight_int4pack_mm_for_cpu

Measured on the RX 6600 XT behind Thunderbolt (PCIe 3.0 x4, 2.5 GB/s), i9-13950HX (16 CPU threads for experts), 31 GiB RAM, quiet host. Every format got 1.66 GiB of expert VRAM, its own placement from prain tune and tuned kernels; 4k context, 0.75 GiB KV:

	bf16	int8	int4
experts in VRAM (cache slots)	120	270	540
decode, tok/s (median of 3 x 48 tokens)	1.04	1.42	1.62
first token, short chat prompt (median)	22.9 s	5.0 s	4.9 s
3 greedy answers x 64 tokens	221 s	143 s	123 s
greedy answers identical to bf16	-	1 of 3	1 of 3

The smaller formats win twice: more experts fit in VRAM, and every miss moves fewer bytes over the link and out of the page cache (the int4 pack fits in free RAM, the bf16 one does not). On the GPU itself a decode round costs about the same in every format (2.3-2.8 ms). These decode numbers predate the dense-layer fix and the int4 kernels below, which make decode 3x faster.

Dense layers 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 just 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). With the agent setup below, decode went from 1.58 to 2.77 tok/s with unchanged answers (same quality score). Prefill keeps rocBLAS.

int4 experts on RDNA2

vLLM's int4 MoE kernel pads every token to a 16-row tile and multiplies in bf16, which RDNA2 lacks: 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 slots): decode is the median of 3 x 48 tokens, prefill a fresh 2.3k-token prompt, quality the greedy answers to the six quality.py prompts (identical to the reference answers, 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 (both 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).

Trading quality for speed

[miss] min_weight (default 1) is the share of each token's router weight that must be computed. Below 1, a token skips its lightest experts that are not in VRAM once it has that share, and its other routes are scaled up to its full weight: fewer copies, a slightly different model. Measured on the agent setup: int4, 16k context, 1 GiB KV and 450 cache slots (1.32 GiB) as in the example config, per-layer slot counts from prain tune. Quality is the mean negative log-likelihood of the min_weight = 1 answers to six chat prompts (scripts/quality.py; 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). This table predates the int4 kernels above.

Context length and coding agents

Gemma 4 has 25 sliding-window layers (window 1024) and 5 full-attention layers, so KV grows slowly with context: per request ~20 KiB per token plus ~0.6 GiB for the sliding layers, which also hold the tokens in flight (2 x --max-num-batched-tokens). kv_cache_gib = 1.0 holds 18,169 tokens, 1.11 requests of 16k (measured).

With the example config (int4, 16k, tools) a cold 4.7k-token prompt took 70 s and a 12k-token one 176 s (before the int4 kernels, which cut prefill by a quarter): every 1024-token prefill chunk needs nearly all experts, copied over the 2.5 GB/s link. 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 just the system prompt recomputes it once (176 s); later ones hit (4.5 s).

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 config even ~300-token prompts crashed the engine. scripts/quality.py score --greedy-only skips them.

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), so 0.3 GiB less for experts (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).

Tune for a machine

python scripts/measure.py PACK_DIR/layer029.w13 --quant int4 --threads 16 > cost.json  # server stopped
prain sim TRACE CONFIG MODEL_DIR/config.json --cost cost.json
prain tune TRACE MODEL_DIR/config.json --cost cost.json --quant int4 --expert-vram-gib 1.4 \
    --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, then
python scripts/quality.py http://127.0.0.1:8000 score ref.json  # on 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-predicts bf16 (0.6 vs 1.04 tok/s measured, the page cache helps) and over-predicted int4 (2.4 vs 1.62 tok/s).

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 decode step took 395 ms profiled, 307 ms without.

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