fabolous005/prain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

prain: README: expert formats, dense-layer fix, min_weight trade-off, 16k agent setup; example config int4, 16k, tools

Browse Source
This commit is contained in:
fabolous005
2026-10-28 22:07:25 +01:00
parent 216cd76a77
commit bb42ea3802
2 changed files with 126 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+97 -11
View File
@@ -3,8 +3,9 @@
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.
Target: Gemma 4 26B-A4B (30 layers x 128 experts, 42.5 GiB of experts) on an 8 GB RX 6600 XT.
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)
@@ -17,25 +18,110 @@ Target: Gemma 4 26B-A4B (30 layers x 128 experts, 42.5 GiB of experts) on an 8 G
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 4096 --max-num-batched-tokens 1024 \
--enable-prefix-caching --attention-backend TRITON_ATTN
--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 pack (~3 min). Chat with `prain tui`, or `prain web --config FILE` for
several users with drafts.
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 --threads 16 > cost.json # server stopped
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 --expert-vram-gib 1.67 --ram-gib 0 \
--pack-dir PACK_DIR --cpu-threads 16 > tuned.toml
python scripts/bench.py http://127.0.0.1:8000
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
The simulator treats every disk-tier expert as a disk read; the page cache makes real runs faster.
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
examples/gemma4-26b-a4b-rx6600xt.toml
+29 -18
View File
@@ -1,9 +1,12 @@
# prain placement: prithivMLmods/gemma-4-26B-A4B-it-Uncensored-MAX on an 8 GB RX 6600 XT
# Model: 30 MoE layers x 128 experts, top-8 routing, 11.34 MiB per expert (BF16).
# VRAM, measured: 6.22 GiB weights (4.46 dense + 5 slots x 30 layers) + 0.75 GiB KV cache: 7.9 of
# 8 GiB in use. Serve with --attention-backend TRITON_ATTN (ROCM_ATTN needs more LDS than RDNA2 has).
# RAM tier: pinning is very slow behind a strict IOMMU (Thunderbolt eGPU: minutes per GiB), so this
# file keeps no RAM tier; the kernel page cache holds the hot part of the pack instead.
# prain placement: prithivMLmods/gemma-4-26B-A4B-it-Uncensored-MAX on an 8 GB RX 6600 XT, set up
# for a coding agent: int4 experts, 16k context, tool calls.
# Model: 30 MoE layers x 128 experts, top-8 routing, 3.01 MiB per int4 expert (11.34 MiB bf16).
# VRAM: 4.46 GiB dense weights + 30 x 15 x 3.01 MiB = 1.32 GiB of expert slots + 1.0 GiB KV cache.
# Serve with --max-model-len 16384 --attention-backend TRITON_ATTN (ROCM_ATTN needs more LDS than
# RDNA2 has) --enable-auto-tool-choice --tool-call-parser gemma4.
# RAM tier: pinning host memory took minutes per GiB on the Thunderbolt eGPU host (cause unproven)
# and pinned copies were only 4% faster than pageable ones, so this file keeps no RAM tier; the
# kernel page cache holds the pack instead (the int4 pack, 11.3 GiB, fits in free RAM).
#
# Tiers: vram = pinned in a VRAM slot for the whole run
# ram = pinned host memory, copied into a cache slot when a token needs it
@@ -13,19 +16,21 @@
# Layer selectors: "*", "12", "0-9,20". Expert lists: "*", "0-15,32", [1, 2, 3].
# Check what a file costs: prain check <this file> <model config.json>
# Compare variants offline: prain sim <trace> <this file> <model config.json>
# prain tune <trace> <model config.json> --expert-vram-gib 1.7 ...
# prain tune <trace> <model config.json> --quant int4 --expert-vram-gib 1.4 ...
# Relative paths start at this file's directory; "~" works too.
pack_dir = "../models/packs/gemma4-26b-a4b" # 42.5 GiB, written on the first run
pack_dir = "../models/packs/gemma4-26b-a4b" # bf16 pack (42.5 GiB) written on the first run,
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 = 5 # 30 x 5 x 11.34 MiB = 1.66 GiB of VRAM
cache_slots = 15 # 30 x 15 x 3.01 MiB = 1.32 GiB of VRAM
stats_path = "../models/packs/gemma4-26b-a4b/stats.json"
stats_every = 64 # MoE forward passes between stats/trace writes
trace_path = "../models/packs/gemma4-26b-a4b/trace.jsonl.gz" # routing trace for sim/tune
[policy] # which expert a cache slot gives up
eviction = "lfu" # lru | lfu (decayed frequency) | gdsf (frequency x refetch cost)
heat = "weight" # what makes an expert hot: count (tokens) | weight (router mass) |
eviction = "gdsf" # lru | lfu (decayed frequency) | gdsf (frequency x refetch cost)
heat = "count" # what makes an expert hot: count (tokens) | weight (router mass) |
# recency (router mass, later tokens weigh more: good for warm-ups)
lfu_half_life = 64 # forwards until an LFU score halves
disk_cost = 4.0 # gdsf: a disk expert costs this many RAM experts to refetch
@@ -41,18 +46,24 @@ budget = 2 # experts per layer and forward
min_prob = 0.05 # skip guesses below this router probability
[miss] # an expert that is not in VRAM
mode = "auto" # copy | cpu (always compute it on the CPU) | auto (CPU if it has
cpu_max_tokens = 2 # at most cpu_max_tokens tokens, else copy). Measured on the Thunderbolt
# eGPU host: a copy costs ~4.6 ms + disk read, the CPU ~2.5 ms per token
mode = "copy" # copy | cpu (always compute it on the CPU) | auto (CPU if it has
cpu_max_tokens = 1 # at most cpu_max_tokens tokens, else copy). Measured on the Thunderbolt
# eGPU host, per int4 expert: a copy ~1.3 ms + disk read, the CPU
# 0.8-1.3 ms per token. `prain tune` picked copy for int4 (a copied
# expert stays for later tokens) and auto <= 2 tokens for bf16 (copy 4.8 ms)
cpu_threads = 16 # CPU threads of the vLLM worker (sets OMP_NUM_THREADS)
admit = 0 # cpu/auto: copy this many CPU-run RAM-tier experts into VRAM afterwards
min_weight = 1.0 # share of each token's router weight that must be computed. Below 1 a token
# skips experts not in VRAM once it has that share: faster, dumber.
# Measured here: 0.6 decodes 14% faster (answer NLL 0.11 -> 0.19), 0.4
# 21% (0.40, visible slips)
[serve] # parallel requests (`prain serve` passes these to vLLM)
max_requests = 2 # requests vLLM runs at once (--max-num-seqs); warm-ups count too
kv_cache_gib = 0.75 # KV cache VRAM shared by them (--kv-cache-memory-bytes). vLLM logs
# "Maximum concurrency for 4096 tokens per request: N.NNx" at start;
# below max_requests, long requests wait or get preempted. Every GiB
# here is 90 expert cache slots less.
kv_cache_gib = 1.0 # KV cache VRAM shared by them (--kv-cache-memory-bytes). vLLM logs
# "Maximum concurrency for 16384 tokens per request: 1.11x" with 1.0 GiB.
# Gemma 4 needs ~20 KiB per token of context plus ~0.6 GiB for its
# sliding-window layers; every GiB here is 340 int4 expert slots less.
[draft] # answers that arrive while others use the GPU (`prain web --config`)
full_requests = 1 # answers that run with all their experts; later ones run at once as