132 lines
7.3 KiB
Markdown
132 lines
7.3 KiB
Markdown
# 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 `prain web --config FILE` for several users with
|
|
drafts. 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) | `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 below, which makes every format 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.
|
|
|
|
## 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).
|
|
|
|
## 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: 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).
|
|
|
|
## 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).
|
|
|
|
## 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).
|