From bb42ea38025be211b241728386b95c9ef1f60be2 Mon Sep 17 00:00:00 2001 From: fabolous005 Date: Wed, 28 Oct 2026 22:07:25 +0100 Subject: [PATCH] prain: README: expert formats, dense-layer fix, min_weight trade-off, 16k agent setup; example config int4, 16k, tools --- README.md | 108 +++++++++++++++++++++++--- examples/gemma4-26b-a4b-rx6600xt.toml | 47 ++++++----- 2 files changed, 126 insertions(+), 29 deletions(-) diff --git a/README.md b/README.md index ed858ca..7c284bf 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/examples/gemma4-26b-a4b-rx6600xt.toml b/examples/gemma4-26b-a4b-rx6600xt.toml index ee099e1..71c579a 100644 --- a/examples/gemma4-26b-a4b-rx6600xt.toml +++ b/examples/gemma4-26b-a4b-rx6600xt.toml @@ -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 # Compare variants offline: prain sim -# prain tune --expert-vram-gib 1.7 ... +# prain tune --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