diff --git a/README.md b/README.md index 623c21c..50ad8d0 100644 --- a/README.md +++ b/README.md @@ -1,92 +1,293 @@ # 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`, and it takes `--max-num-seqs` and `--kv-cache-memory-bytes` from +the file's `[serve]` table (it refuses them after `--`). 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 +298,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 +343,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) | diff --git a/examples/gemma4-26b-a4b-rx6600xt.toml b/examples/gemma4-26b-a4b-rx6600xt.toml index aeacdb6..f8ba29b 100644 --- a/examples/gemma4-26b-a4b-rx6600xt.toml +++ b/examples/gemma4-26b-a4b-rx6600xt.toml @@ -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 diff --git a/examples/personas/explainer.md b/examples/personas/explainer.md new file mode 100644 index 0000000..ee7d2fd --- /dev/null +++ b/examples/personas/explainer.md @@ -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. diff --git a/examples/personas/gentoo-admin.md b/examples/personas/gentoo-admin.md new file mode 100644 index 0000000..82ceab2 --- /dev/null +++ b/examples/personas/gentoo-admin.md @@ -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. diff --git a/examples/personas/pair-programmer.md b/examples/personas/pair-programmer.md new file mode 100644 index 0000000..5bbb023 --- /dev/null +++ b/examples/personas/pair-programmer.md @@ -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. diff --git a/prain/chat.py b/prain/chat.py index f82a1d1..14dde53 100644 --- a/prain/chat.py +++ b/prain/chat.py @@ -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, diff --git a/prain/cli.py b/prain/cli.py index c0913cd..8a9b891 100644 --- a/prain/cli.py +++ b/prain/cli.py @@ -165,6 +165,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 +199,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__": diff --git a/prain/tui.py b/prain/tui.py index 1d4d002..f65f235 100644 --- a/prain/tui.py +++ b/prain/tui.py @@ -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) diff --git a/prain/web.html b/prain/web.html index d504ab4..b7f23e2 100644 --- a/prain/web.html +++ b/prain/web.html @@ -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); diff --git a/prain/web.py b/prain/web.py index ede3666..93b3851 100644 --- a/prain/web.py +++ b/prain/web.py @@ -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() diff --git a/tests/test_cli.py b/tests/test_cli.py index a286c4c..9285504 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -68,6 +68,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") diff --git a/tests/test_tui.py b/tests/test_tui.py index 51fe063..028a15f 100644 --- a/tests/test_tui.py +++ b/tests/test_tui.py @@ -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") diff --git a/tests/test_web.py b/tests/test_web.py index bf02588..4e02272 100644 --- a/tests/test_web.py +++ b/tests/test_web.py @@ -124,6 +124,19 @@ class Web(unittest.TestCase): ""]) 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