diff --git a/README.md b/README.md index f768d35..623c21c 100644 --- a/README.md +++ b/README.md @@ -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. diff --git a/examples/gemma4-26b-a4b-rx6600xt.toml b/examples/gemma4-26b-a4b-rx6600xt.toml index 71c579a..aeacdb6 100644 --- a/examples/gemma4-26b-a4b-rx6600xt.toml +++ b/examples/gemma4-26b-a4b-rx6600xt.toml @@ -24,8 +24,8 @@ quant = "int4" # then quantized once into an int4 pack # 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 -stats_path = "../models/packs/gemma4-26b-a4b/stats.json" -stats_every = 64 # MoE forward passes between stats/trace writes +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 [policy] # which expert a cache slot gives up diff --git a/prain/cli.py b/prain/cli.py index e44b023..c0913cd 100644 --- a/prain/cli.py +++ b/prain/cli.py @@ -169,8 +169,9 @@ def main(argv=None): ui.add_argument("--port", type=int, default=8080) ui.add_argument("--host", default="127.0.0.1", help="address to bind; 0.0.0.0 for every interface (no login!)") - ui.add_argument("--config", help="placement file: [serve] max_requests and the " - "[draft] table (default: their defaults)") + ui.add_argument("--config", help="placement file: [serve] max_requests, the [draft] " + "table and stats_path for the expert monitor " + "(default: their defaults)") args = parser.parse_args(argv) if args.command == "check": @@ -202,7 +203,7 @@ def main(argv=None): config = load_config(args.config) if args.config else Config(pack_dir="") # defaults web(args.url, args.port, args.warmup, args.warmup_priority, config.serve, config.draft, - host=args.host) + host=args.host, stats_path=config.stats_path) if __name__ == "__main__": diff --git a/prain/policy.py b/prain/policy.py index c823156..311ee58 100644 --- a/prain/policy.py +++ b/prain/policy.py @@ -236,8 +236,10 @@ COUNTERS = ("hits", "misses", "cpu", "prefetches", "prefetch_hits", "prefetch_wa def write_stats(path, caches): - """Atomically write {"layers": {"": {"uses": [...], "hits": n, ...}}}.""" - data = {"layers": {str(layer): {"uses": c.uses, **{k: getattr(c, k) for k in COUNTERS}} + """Atomically write {"layers": {"": {"uses": [...], "vram": [experts in VRAM now], + "hits": n, ...}}}.""" + data = {"layers": {str(layer): {"uses": c.uses, "vram": sorted({*c.pinned, *c.slots}), + **{k: getattr(c, k) for k in COUNTERS}} for layer, c in sorted(caches.items())}} tmp = f"{path}.tmp" with open(tmp, "w") as f: diff --git a/prain/web.html b/prain/web.html new file mode 100644 index 0000000..d504ab4 --- /dev/null +++ b/prain/web.html @@ -0,0 +1,716 @@ + + + + + +prain + + +
+
+
+
+ + prain + +
+
+ + +
+
+ +
+
+
+

Ready when you are

+

prain moves the experts your message needs into VRAM while you type, so the answer starts sooner.

+
+
+ + + +
+
+
+
+ +
+
+ + +
+
+ + Enter sends · Shift+Enter adds a line +
+
+
+ +
+ +
+ + diff --git a/prain/web.py b/prain/web.py index 0ea9a9b..ede3666 100644 --- a/prain/web.py +++ b/prain/web.py @@ -1,10 +1,12 @@ -"""prain web: a one-page browser chat that warms experts up while you type. +"""prain web: a browser chat that warms experts up while you type, with a model monitor. -A small stdlib proxy in front of vLLM. GET / serves the page. POST /warm takes -{"session", "messages", "draft"} and returns at once. 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). +A small stdlib proxy in front of vLLM. GET / serves the page (web.html). GET /status returns +what the monitor shows: vLLM's counters (/metrics), the GPU (amdgpu sysfs) and prain's expert +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). 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 @@ -19,127 +21,88 @@ Several browsers at once ([serve] and [draft] tables of the placement file): cancels it. """ +import glob import json +import os import threading +import time import urllib.error +import urllib.request 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.config import Draft, Serve MAX_SESSIONS = 64 -PAGE = """ - -prain - -
- -
- -""" def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, serve=Serve(), - draft=Draft()): + draft=Draft(), stats_path=None, sysfs="/sys/class/drm"): model = first_model(base_url) lock = threading.Condition() sessions = OrderedDict() # session -> Warmer, least recently used first @@ -167,7 +130,13 @@ def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, ser class Handler(BaseHTTPRequestHandler): def do_GET(self): - self.reply(200, "text/html; charset=utf-8", PAGE.encode()) + 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()} + return self.reply(200, "application/json", json.dumps(data).encode()) + self.reply(200, "text/html; charset=utf-8", PAGE) def do_POST(self): body = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) @@ -178,6 +147,7 @@ def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, ser 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"])) + 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") @@ -267,9 +237,10 @@ def make_server(base_url, port, host="127.0.0.1", warmup="idle", priority=1, ser return server -def serve(base_url, port, warmup="idle", priority=1, serve=Serve(), draft=Draft(), host="127.0.0.1"): +def serve(base_url, port, warmup="idle", priority=1, serve=Serve(), draft=Draft(), host="127.0.0.1", + stats_path=None): server = make_server(base_url, port, host=host, warmup=warmup, priority=priority, serve=serve, - draft=draft) + draft=draft, stats_path=stats_path) print(f"prain web: http://{host}:{server.server_port} (vLLM at {base_url}; " f"{draft.full_requests} full of {serve.max_requests} parallel answers)") server.serve_forever() diff --git a/pyproject.toml b/pyproject.toml index 3beb54b..7365345 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,3 +17,6 @@ build-backend = "setuptools.build_meta" [tool.setuptools] packages = ["prain"] + +[tool.setuptools.package-data] +prain = ["web.html"] diff --git a/tests/helpers.py b/tests/helpers.py index ae3e3b5..a5213e2 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -41,10 +41,13 @@ class StubVllm: self.delay = 0.0 self.pause = 0.0 self.fail = None + self.metrics = None # Prometheus text served at GET /metrics stub = self class Handler(BaseHTTPRequestHandler): def do_GET(self): + if self.path == "/metrics" and stub.metrics is not None: + return self.send(200, "text/plain", stub.metrics) self.send(200, "application/json", {"data": [{"id": "stub-model"}]}) def do_POST(self): diff --git a/tests/test_policy.py b/tests/test_policy.py index 414ff4b..8ecc02f 100644 --- a/tests/test_policy.py +++ b/tests/test_policy.py @@ -119,3 +119,7 @@ class Prefetch(TmpCase): stats = json.load(f)["layers"]["7"] self.assertEqual((stats["uses"], stats["prefetch_wasted"], stats["misses"]), ([1, 1, 1, 0], 1, 2)) + self.assertEqual(stats["vram"], [1, 2]) # the experts in VRAM now, for the web monitor + write_stats(path, {7: SlotCache([5, 3], 8, 2, Policy(eviction="lfu"))}) + with open(path) as f: + self.assertEqual(json.load(f)["layers"]["7"]["vram"], [3, 5]) # pinned ones too diff --git a/tests/test_web.py b/tests/test_web.py index 24aa835..bf02588 100644 --- a/tests/test_web.py +++ b/tests/test_web.py @@ -1,4 +1,5 @@ import json +import os import threading import unittest import urllib.error @@ -7,7 +8,7 @@ import urllib.request import prain.chat as chat from prain.config import Draft, Serve from prain.web import make_server -from tests.helpers import StubVllm, wait_for +from tests.helpers import StubVllm, TmpCase, wait_for def events(response): @@ -38,7 +39,9 @@ class Web(unittest.TestCase): self.url = f"http://127.0.0.1:{self.server.server_port}" def warm(self, session, draft): - chat.post(self.url, "/warm", {"session": session, "messages": [], "draft": draft}).read() + with chat.post(self.url, "/warm", {"session": session, "messages": [], + "draft": draft}) as response: + return response.status def ask(self, session, text): return chat.post(self.url, "/chat", {"session": session, @@ -46,7 +49,9 @@ class Web(unittest.TestCase): def test_page_warm_up_and_chat(self): with urllib.request.urlopen(self.url) as response: - self.assertIn(b'