Compare commits

...

2 Commits

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
13 changed files with 472 additions and 151 deletions
+307 -126
View File
@@ -1,92 +1,296 @@
# 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.
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`: 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
- `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: `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:
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
## Serve a model
. 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
`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:
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.
| 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 |
## Expert formats
`examples/gemma4-26b-a4b-rx6600xt.toml` is commented line by line. It keeps all experts in the
pack on disk, served through the page cache, and gives each layer 15 cache slots in VRAM (1.32
GiB), plus 1 GiB of KV cache for 16k tokens of context. `prain check PLACEMENT
MODEL_DIR/config.json` prints what a placement costs in VRAM, RAM and disk before anything
loads.
| quant | expert | all experts | weights | VRAM kernel | CPU kernel |
|---|---|---|---|---|---|
| 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` |
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.
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:
## Choose a quantization
| | bf16 | int8 | int4 |
`quant` in the placement file picks the expert format. Smaller experts win twice: more of them
fit in VRAM, and every miss moves fewer bytes over the 2.5 GB/s Thunderbolt link and out of the
page cache (on a 31 GiB host the int4 and int8 packs fit in free RAM, the bf16 one does not).
They cost some accuracy:
| quant | expert | pack on disk | slots per layer in 1.32 GiB | decode | first token, short prompt | first token, 2.3k-token prompt | answers equal to bf16 |
|---|---|---|---|---|---|---|---|
| bf16 | 11.34 MiB | 42.5 GiB | 4 | 0.90 tok/s | 14.7 s | 69.5 s | reference |
| int8 | 5.69 MiB | 21.3 GiB | 8 | 2.19 tok/s | 5.9 s | 32.3 s | 4 of 6 (215 of 288 tokens) |
| int4 | 3.01 MiB | 11.3 GiB | 15 | 5.19 tok/s | 3.3 s | 20.5 s | 0 of 6 (94 of 288 tokens) |
Measured with the example placement on the RX 6600 XT behind Thunderbolt (PCIe 3.0 x4), an
i9-13950HX and 31 GiB of RAM, changing only `quant` and `cache_slots`. Decode is the median of
3 x 48 tokens. "Equal" counts how many of the six `scripts/quality.py` prompts got the same 48
greedy tokens as bf16; the brackets count the tokens that matched before each answer's first
difference.
int4 is the default: it decodes 2.4 times as fast as int8 and reads a long prompt in two thirds
of the time. Its answers are worded differently from bf16's from the first tokens on. Five of
the six meant the same; the sixth, a Gentoo question, named a tool that does not exist
("oneway-way") where bf16 and int8 named `revdep-rebuild`. int8 stays closest to bf16 at less
than half the speed of int4. bf16 is mostly the reference for quality checks: its 4 slots per
layer are fewer than the 8 experts each token uses, so every decode step runs in several rounds
(`prain check` warns about this).
To launch int8 instead of int4, copy the example and change two lines. Relative paths in a
placement file start at the file's directory, so keep the copy next to the example:
cp examples/gemma4-26b-a4b-rx6600xt.toml examples/int8.toml
sed -i -e 's/^quant = "int4"/quant = "int8"/' -e 's/^cache_slots = 15 /cache_slots = 8 /' \
examples/int8.toml
prain check examples/int8.toml $M/config.json # VRAM 1.33 GiB, as before
prain serve examples/int8.toml $M -- $FLAGS
For bf16, set `cache_slots = 4`. Keep slots x expert size x 30 layers near 1.32 GiB. On the
example placement vLLM reports "Model loading took 6.29 GiB" (the dense weights, 1.32 GiB of
expert slots and prain's buffers), the KV cache takes 1 GiB (`kv_cache_gib`), and about 0.55 GiB
of the card's 7.98 GiB stays free for activations and kernels. prain logs that free VRAM at
startup ("MiB of VRAM free for kernels and activations", before the KV cache takes its share).
More VRAM for experts means less for context, and the other way round.
`prain tune` fits per-layer slot counts and policies to a routing trace, which the example records
while it serves (`trace_path`). The tuned int4 placement of the live agent setup (450 slots,
12-18 per layer) decodes at 5.2 tok/s. For another format:
prain tune models/packs/gemma4-26b-a4b/trace.jsonl.gz $M/config.json --quant int8 \
--expert-vram-gib 1.32 --ram-gib 0 --pack-dir models/packs/gemma4-26b-a4b > int8-tuned.toml
Every pack stays on disk once written: `layerNNN.w13`, `.w2` and `.done` for bf16,
`layerNNN.int8.*` and `layerNNN.int4.*` for the others. Delete the formats you do not use; the
next start of that format writes its pack again, and int8 or int4 need the bf16 pack for that.
How the formats are stored and computed:
| quant | weights | GPU kernel | CPU kernel (`[miss] mode = "cpu"`) |
|---|---|---|---|
| experts in VRAM (cache slots) | 120 | 270 | 540 |
| 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 |
| 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` |
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.
## Give the model a persona
## Dense layers on RDNA2
A persona is a short text file that `prain web` and `prain tui` send as the system message at
the start of every conversation. It tells the model its task, how to handle questions, what its
environment is and how to use it. One running model can serve several personas at once, one
chat front end each:
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 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.
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.
## int4 experts on RDNA2
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.
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):
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 |
|---|---|---|---|---|---|
@@ -97,78 +301,44 @@ A/B on the live agent setup (int4, 16k context, 1 GiB KV, 450 slots): decode is
| 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:
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).
## 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 \
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, then
python scripts/quality.py http://127.0.0.1:8000 score ref.json # on another config
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-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).
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 decode step took 395 ms profiled, 307 ms without.
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
@@ -176,3 +346,14 @@ syncs remove the CPU/GPU overlap: a decode step took 395 ms profiled, 307 ms wit
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) |
+3 -1
View File
@@ -23,7 +23,9 @@ pack_dir = "../models/packs/gemma4-26b-a4b" # bf16 pack (42.5 GiB) written on t
quant = "int4" # then quantized once into an int4 pack (11.3 GiB):
# bf16 | int8 (5.69 MiB/expert) | int4 (3.01 MiB)
default_tier = "disk" # runs on any RAM size; `prain tune` does better
cache_slots = 15 # 30 x 15 x 3.01 MiB = 1.32 GiB of VRAM
cache_slots = 15 # 30 x 15 x 3.01 MiB = 1.32 GiB of VRAM. The same VRAM
# holds 8 int8 or 4 bf16 experts per layer: change
# 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
+20
View File
@@ -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.
+28
View File
@@ -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.
+21
View File
@@ -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.
+5
View File
@@ -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,
+17 -2
View File
@@ -113,6 +113,9 @@ def serve(args):
"--scheduling-policy", "priority",
"--max-num-seqs", str(config.serve.max_requests),
"--kv-cache-memory-bytes", str(int(config.serve.kv_cache_gib * GIB)),
# vLLM detects "openai" for Gemma 4, whose template then prints a system
# prompt as a Python list; a flag in `extra` comes later and wins
"--chat-template-content-format", "string",
*extra])
@@ -165,6 +168,9 @@ 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("--host", default="127.0.0.1",
@@ -196,14 +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
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)
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__":
+11 -8
View File
@@ -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):
self.base_url, self.model, self.warmup = base_url, model, warmup
def __init__(self, base_url, model, warmup="idle", priority=1, system=None):
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)
+3 -2
View File
@@ -637,8 +637,9 @@ async function poll() {
set("pill-text", "Offline");
return;
}
set("model", s.model);
set("m-model", s.model);
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);
+18 -11
View File
@@ -5,8 +5,9 @@ what the monitor shows: vLLM's counters (/metrics), the GPU (amdgpu sysfs) and p
stats (the placement file's stats_path, with the experts in VRAM). POST /warm takes
{"session", "messages", "draft"} and returns at once: 202 when a warm-up starts, 204 when not.
POST /chat takes {"session", "messages"} and streams server-sent events: vLLM's chunks, plus
{"prain": "draft" | "refine" | "error"} events of its own. Binds to 127.0.0.1 unless --host
says otherwise (0.0.0.0 serves every interface, without a login).
{"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
@@ -32,7 +33,8 @@ 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
@@ -102,7 +104,10 @@ def expert_stats(path):
def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, serve=Serve(),
draft=Draft(), stats_path=None, sysfs="/sys/class/drm"):
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
@@ -133,8 +138,9 @@ def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, ser
if self.path == "/status":
with lock:
busy = streams[0]
data = {"model": model, "vllm": vllm_metrics(base_url), "gpu": gpu_stats(sysfs),
"experts": expert_stats(stats_path), "streams": busy, "time": time.time()}
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)
@@ -146,13 +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
@@ -238,9 +244,10 @@ def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, ser
def serve(base_url, port, warmup="idle", priority=1, serve=Serve(), draft=Draft(), host="127.0.0.1",
stats_path=None):
stats_path=None, system=None, persona=None):
server = make_server(base_url, port, host=host, warmup=warmup, priority=priority, serve=serve,
draft=draft, stats_path=stats_path)
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"{draft.full_requests} full of {serve.max_requests} parallel answers"
f"{f'; persona {persona}' if persona else ''})")
server.serve_forever()
+15 -1
View File
@@ -55,7 +55,8 @@ 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"])
@@ -68,6 +69,19 @@ class Cli(TmpCase):
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")
+11
View File
@@ -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")
+13
View File
@@ -124,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