prain: web UI with a model monitor
A clear chat page (markdown, code blocks with copy, stop, new chat, kept in the browser, drafts shown as such) and a monitor sidebar, a sheet on phones: tokens/s, a live map of the 30 x 128 experts (in VRAM, used now, copied in), VRAM hit rate, copies, context, prompt cache, GPU VRAM/load/temperature/power/fan. GET /status merges vLLM /metrics, amdgpu sysfs and the stats file, which now lists the experts in VRAM. /warm answers 202 when a warm-up starts. README: the int4 kernels, their A/B, 2048-token chunks as a trade, prompt logprobs limit.
This commit is contained in:
@@ -34,9 +34,12 @@ Tune vLLM's fused-MoE kernel for the GPU once (~45 min of Triton compiles the fi
|
||||
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 (`--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.
|
||||
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.
|
||||
|
||||
## Expert formats
|
||||
|
||||
@@ -44,7 +47,7 @@ vLLM's `gemma4` parser, so coding agents and MCP clients can use the OpenAI-comp
|
||||
|---|---|---|---|---|---|
|
||||
| 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` |
|
||||
| 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` |
|
||||
|
||||
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
|
||||
@@ -61,7 +64,7 @@ placement from `prain tune` and tuned kernels; 4k context, 0.75 GiB KV:
|
||||
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.
|
||||
numbers predate the dense-layer fix and the int4 kernels below, which make decode 3x faster.
|
||||
|
||||
## Dense layers on RDNA2
|
||||
|
||||
@@ -72,6 +75,32 @@ with at most 4 rows through a Triton GEMV (`prain/gemv.py`, ~210 GB/s; `PRAIN_GE
|
||||
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.
|
||||
|
||||
## int4 experts on RDNA2
|
||||
|
||||
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):
|
||||
|
||||
| | decode tok/s | first token, short | first token, 2.3k tokens | identical | mean logprob |
|
||||
|---|---|---|---|---|---|
|
||||
| before | 3.49 | 4.3 s | 27.5 s | 4 of 6 | -0.104 |
|
||||
| decode kernel | 4.73 | 4.2 s | 27.3 s | 4 of 6 | -0.093 |
|
||||
| staged copies | 3.67 | 3.7 s | 24.5 s | 4 of 6 | -0.098 |
|
||||
| fp16 rounds | 3.84 | 4.0 s | 23.3 s | 4 of 6 | -0.102 |
|
||||
| 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:
|
||||
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
|
||||
@@ -91,7 +120,7 @@ negative log-likelihood of the min_weight = 1 answers to six chat prompts (`scri
|
||||
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).
|
||||
the dense-layer fix). This table predates the int4 kernels above.
|
||||
|
||||
## Context length and coding agents
|
||||
|
||||
@@ -101,12 +130,22 @@ hold the tokens in flight (2 x `--max-num-batched-tokens`). `kv_cache_gib = 1.0`
|
||||
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.
|
||||
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
|
||||
@@ -127,7 +166,7 @@ tok/s measured, the page cache helps) and over-predicted int4 (2.4 vs 1.62 tok/s
|
||||
|
||||
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: h2d/repack/d2d, round), each synced with the GPU. `python -m prain.profile dir`
|
||||
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.
|
||||
|
||||
|
||||
Reference in New Issue
Block a user