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