Compare commits
2 Commits
0f3911afb5
...
a14e3ade0b
| Author | SHA1 | Date | |
|---|---|---|---|
 fabolous005
|
a14e3ade0b | ||
|
fabolous005
|
d1ccc5c294 |
@@ -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.
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
+4
-3
@@ -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__":
|
||||
|
||||
+108
-15
@@ -34,9 +34,9 @@ from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
|
||||
|
||||
from prain import drafts, profile
|
||||
from prain.config import layer_index, load_config
|
||||
from prain.kernels import DECODE_ROWS, gfx10, moe_int4_decode, unpack_int4
|
||||
from prain.policy import SlotCache, write_stats
|
||||
from prain.quant import (GROUP, cpu_matmul, layout, moe_quant_config, quantize, to_cpu_int4,
|
||||
to_gpu_int4)
|
||||
from prain.quant import (GROUP, cpu_matmul, layout, moe_quant_config, quantize, to_cpu_int4)
|
||||
from prain.sim import write_trace
|
||||
from prain.tiered import draft_rows, tiered_forward
|
||||
|
||||
@@ -46,6 +46,68 @@ _COPY_STREAM = [] # one side stream shared by all layers: async copies stay in
|
||||
_TRACE = [] # trace records not written yet
|
||||
_DRAFT_MASK = [None, None] # (spans, their bool row mask on the GPU) of the current forward
|
||||
_DRAFT_TOKENS = [0] # draft token rows seen, for the stats log
|
||||
_STAGING = {} # stream -> _Staging
|
||||
LOW_VRAM = 128 * 2**20 # driver-free bytes below which torch's cache goes back (_release_cache)
|
||||
|
||||
|
||||
def _release_cache():
|
||||
"""Hand torch's cached free blocks (the freed one-row bf16 expert weights, ~11 MiB per layer,
|
||||
and load-time temporaries) back to the driver. HIP loads kernel code objects lazily into
|
||||
VRAM; with a full card and nothing returned, a kernel's first launch fails with
|
||||
hipErrorNoBinaryForGpu ("no kernel image")."""
|
||||
before = torch.cuda.mem_get_info()[0]
|
||||
torch.cuda.empty_cache()
|
||||
free, total = torch.cuda.mem_get_info()
|
||||
logger.info("prain: %.0f MiB of VRAM free for kernels and activations (%.0f MiB before "
|
||||
"releasing torch's cache, %.2f of %.2f GiB reserved)", free / 2**20,
|
||||
before / 2**20, torch.cuda.memory_reserved() / 2**30, total / 2**30)
|
||||
|
||||
|
||||
class _Staging:
|
||||
"""A ring of pinned host buffers and GPU buffers for staged int4 copies: the CPU copies an
|
||||
expert's tensors from the page cache into one pinned buffer, one DMA moves it, and
|
||||
unpack_int4 writes the slot. The CPU fills the next buffer while the last DMA runs."""
|
||||
SLOTS = 4
|
||||
|
||||
def __init__(self, nbytes, device):
|
||||
self.host = [torch.empty(nbytes, dtype=torch.uint8).pin_memory() for _ in range(self.SLOTS)]
|
||||
self.dev = [torch.empty(nbytes, dtype=torch.uint8, device=device)
|
||||
for _ in range(self.SLOTS)]
|
||||
self.done = [None] * self.SLOTS # event: the DMA that last read this pinned buffer
|
||||
self.i = 0
|
||||
|
||||
def take(self):
|
||||
self.i = (self.i + 1) % self.SLOTS
|
||||
if self.done[self.i] is not None:
|
||||
self.done[self.i].synchronize()
|
||||
return self.host[self.i], self.dev[self.i], self.i
|
||||
|
||||
|
||||
def stage_int4(sources, slots, slot, ring, stream):
|
||||
"""Copy one int4 expert: sources [[w13 q, w13 s], [w2 q, w2 s]] on the CPU (q in the CPU
|
||||
layout) -> slots[...][slot] on the GPU (q in the GPU layout), through `ring`."""
|
||||
host, dev, i = ring.take()
|
||||
parts, offset = [], 0
|
||||
with profile.span("memcpy"):
|
||||
for mats in sources:
|
||||
for src in mats:
|
||||
n = src.nbytes
|
||||
host[offset:offset + n].copy_(src.reshape(-1).view(torch.uint8))
|
||||
parts.append((offset, n))
|
||||
offset += n
|
||||
with profile.span("h2d"):
|
||||
dev[:offset].copy_(host[:offset], non_blocking=True)
|
||||
ring.done[i] = torch.cuda.Event()
|
||||
ring.done[i].record(stream)
|
||||
with profile.span("unpack"):
|
||||
parts = iter(parts)
|
||||
for dsts in slots:
|
||||
for t, dst in enumerate(dsts):
|
||||
o, n = next(parts)
|
||||
if t == 0: # the q tensor: the pack keeps the CPU kernel's layout
|
||||
unpack_int4(dev[o:o + n], dst[slot])
|
||||
else:
|
||||
dst[slot].view(-1).view(torch.uint8).copy_(dev[o:o + n])
|
||||
|
||||
|
||||
def _install_draft_hook():
|
||||
@@ -272,6 +334,7 @@ class TieredRoutedExperts(RoutedExperts):
|
||||
self._seal(self.quant)
|
||||
self.pack_complete, self._writing, self._bf16 = True, False, None
|
||||
self._pack = self._pack_of(self.quant) # [[w13 tensors], [w2 tensors]] per expert row
|
||||
self._expert_bytes = sum(t[0].nbytes for mats in self._pack for t in mats)
|
||||
if self._qslots is None:
|
||||
self._slots, self._qconfig = [[self.w13_weight.data], [self.w2_weight.data]], None
|
||||
else:
|
||||
@@ -307,16 +370,21 @@ class TieredRoutedExperts(RoutedExperts):
|
||||
return self._ram[row]
|
||||
|
||||
def _copy(self, expert, slot):
|
||||
if self.quant == "int4": # the pack keeps the CPU kernel's layout: stage, then unpack
|
||||
return self._copy_staged(expert, slot)
|
||||
for dsts, srcs in zip(self._slots, self._source(expert)):
|
||||
for i, (dst, src) in enumerate(zip(dsts, srcs)):
|
||||
if i == 0 and self.quant == "int4": # the pack keeps the CPU kernel's order
|
||||
with profile.span("h2d"):
|
||||
src = src.to(dst.device, non_blocking=True)
|
||||
with profile.span("repack"):
|
||||
src = to_gpu_int4(src)
|
||||
for dst, src in zip(dsts, srcs):
|
||||
with profile.span("h2d" if src.device.type == "cpu" else "d2d"):
|
||||
dst[slot].copy_(src, non_blocking=True)
|
||||
|
||||
def _copy_staged(self, expert, slot):
|
||||
"""int4 on the current stream: one pinned buffer, one DMA, unpack into the slot."""
|
||||
stream = torch.cuda.current_stream()
|
||||
ring = _STAGING.get(stream)
|
||||
if ring is None:
|
||||
ring = _STAGING[stream] = _Staging(self._expert_bytes, stream.device)
|
||||
stage_int4(self._source(expert), self._slots, slot, ring, stream)
|
||||
|
||||
def _load(self, expert, slot):
|
||||
"""Demand copy on the compute stream."""
|
||||
pending = self._ready_events.pop(slot, None)
|
||||
@@ -325,7 +393,8 @@ class TieredRoutedExperts(RoutedExperts):
|
||||
self._copy(expert, slot)
|
||||
|
||||
def _load_async(self, expert, slot):
|
||||
"""Copy on the side stream; only RAM-tier (pinned) experts are truly asynchronous."""
|
||||
"""Copy on the side stream; truly asynchronous for int4 (staged) and pinned RAM-tier
|
||||
experts."""
|
||||
stream = _copy_stream()
|
||||
last = self._last_use.get(slot)
|
||||
if last is not None: # a queued kernel may still read the slot's old expert
|
||||
@@ -337,16 +406,19 @@ class TieredRoutedExperts(RoutedExperts):
|
||||
self._ready_events[slot] = event
|
||||
|
||||
def _run_slots(self, x, weights, ids, table, used_slots):
|
||||
"""vLLM's Triton fused-MoE kernel on the slot tensors; expert_map=table skips the
|
||||
routes of experts outside this round."""
|
||||
"""The fused-MoE kernel on the slot tensors (vLLM's, or prain's for int4 on RDNA2);
|
||||
expert_map=table skips the routes of experts outside this round."""
|
||||
current = torch.cuda.current_stream()
|
||||
for slot in used_slots:
|
||||
pending = self._ready_events.pop(slot, None)
|
||||
if pending is not None:
|
||||
current.wait_event(pending)
|
||||
out = fused_experts(x, self._slots[0][0], self._slots[1][0], weights, ids,
|
||||
activation=self.activation, global_num_experts=table.numel(),
|
||||
expert_map=table, quant_config=self._qconfig)
|
||||
if self.quant == "int4" and gfx10(x.device.index or 0):
|
||||
out = self._int4_gfx10(x, weights, ids, table)
|
||||
else:
|
||||
out = fused_experts(x, self._slots[0][0], self._slots[1][0], weights, ids,
|
||||
activation=self.activation, global_num_experts=table.numel(),
|
||||
expert_map=table, quant_config=self._qconfig)
|
||||
if self._async:
|
||||
event = torch.cuda.Event()
|
||||
event.record(current)
|
||||
@@ -354,6 +426,22 @@ class TieredRoutedExperts(RoutedExperts):
|
||||
self._last_use[slot] = event
|
||||
return out
|
||||
|
||||
def _int4_gfx10(self, x, weights, ids, table):
|
||||
"""RDNA2 has no bf16 dot instructions, which vLLM's int4 kernel uses: a GEMV-style
|
||||
kernel for a few tokens (10x faster at one), else vLLM's kernel in fp16 (2x faster),
|
||||
redone in bf16 if fp16 overflowed."""
|
||||
if (x.shape[0] <= DECODE_ROWS
|
||||
and getattr(self.activation, "value", self.activation) == "gelu_tanh"):
|
||||
return moe_int4_decode(x, weights, ids, table, *self._slots[0], *self._slots[1])
|
||||
run = lambda x, config: fused_experts( # noqa: E731
|
||||
x, self._slots[0][0], self._slots[1][0], weights, ids, activation=self.activation,
|
||||
global_num_experts=table.numel(), expert_map=table, quant_config=config)
|
||||
out = run(x.half(), moe_quant_config("int4", self._slots[0][1].half(),
|
||||
self._slots[1][1].half()))
|
||||
if not torch.isfinite(out).all(): # host sync, once per round of a prefill
|
||||
out = run(x, self._qconfig)
|
||||
return out.to(x.dtype)
|
||||
|
||||
def _predict(self, module, args):
|
||||
"""Pre-hook of this layer's router: guess the experts of later layers from the same
|
||||
residual stream (their routers normalize their input, so the guess is cheap and
|
||||
@@ -383,8 +471,9 @@ class TieredRoutedExperts(RoutedExperts):
|
||||
self._load_async(e, slot)
|
||||
target = LAYERS.get(self.prain_layer + prefetch.distance)
|
||||
if prefetch.mode == "pregate" and target is not None and target._ready:
|
||||
staged = target.quant == "int4" # staged copies are async from every tier
|
||||
guesses = [(e, p) for e, p in self._predictions.get(prefetch.distance, ())
|
||||
if p >= prefetch.min_prob and e in target._ram_row]
|
||||
if p >= prefetch.min_prob and (staged or e in target._ram_row)]
|
||||
for e, slot in target.cache.prefetch(guesses, prefetch.budget):
|
||||
target._load_async(e, slot)
|
||||
|
||||
@@ -405,6 +494,8 @@ class TieredRoutedExperts(RoutedExperts):
|
||||
if not self._ready:
|
||||
self._materialize()
|
||||
self._ready = True
|
||||
if self.prain_layer == max(LAYERS): # every layer has its slots now
|
||||
_release_cache()
|
||||
policy = self.prain.policy
|
||||
out = tiered_forward(
|
||||
x, topk_weights, topk_ids, self.cache, self._load, self._run_slots,
|
||||
@@ -422,6 +513,8 @@ class TieredRoutedExperts(RoutedExperts):
|
||||
self._steps += 1
|
||||
if self._steps % self.prain.stats_every or self.prain_layer != min(LAYERS):
|
||||
return
|
||||
if torch.cuda.mem_get_info()[0] < LOW_VRAM: # other processes or activations took it
|
||||
_release_cache()
|
||||
if self.prain.trace_path:
|
||||
_flush_trace()
|
||||
if not self.prain.stats_path:
|
||||
|
||||
+9
-8
@@ -15,13 +15,15 @@ import torch.nn.functional as F
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from prain.kernels import gfx10
|
||||
|
||||
MAX_ROWS = 4
|
||||
# ponytail: one config, the one autotuning picked for every Gemma 4 dense shape on gfx1032.
|
||||
# Triton's autotuner benchmarks with a 256 MiB buffer that torch keeps cached, and an 8 GB card
|
||||
# has no 256 MiB to spare (lazy kernel loads then fail). Re-measure if the GPU or model changes.
|
||||
BLOCK_N, BLOCK_K, NUM_WARPS = 4, 512, 8
|
||||
|
||||
|
||||
@triton.autotune(
|
||||
configs=[triton.Config({"BLOCK_N": n, "BLOCK_K": k}, num_warps=w)
|
||||
for n, k, w in ((4, 512, 8), (8, 512, 8), (16, 256, 4), (32, 256, 2))],
|
||||
key=["M", "N", "K"])
|
||||
@triton.jit
|
||||
def _gemv(w_ptr, x_ptr, y_ptr, M, N, K, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr):
|
||||
@@ -42,8 +44,7 @@ def _gemv(w_ptr, x_ptr, y_ptr, M, N, K, BLOCK_M: tl.constexpr, BLOCK_N: tl.const
|
||||
@functools.cache
|
||||
def fast_device(index):
|
||||
"""gfx10 (RDNA2) GPUs, where vLLM has no skinny-GEMM kernels."""
|
||||
name = getattr(torch.cuda.get_device_properties(index), "gcnArchName", "")
|
||||
return name.startswith("gfx10") and os.environ.get("PRAIN_GEMV", "1") != "0"
|
||||
return gfx10(index) and os.environ.get("PRAIN_GEMV", "1") != "0"
|
||||
|
||||
|
||||
def _small(x, weight, bias):
|
||||
@@ -60,8 +61,8 @@ def linear(x, weight, bias=None):
|
||||
x2 = x.reshape(-1, k).contiguous()
|
||||
m = x2.shape[0]
|
||||
y = torch.empty(m, n, dtype=x.dtype, device=x.device)
|
||||
_gemv[lambda meta: (triton.cdiv(n, meta["BLOCK_N"]),)](
|
||||
weight, x2, y, m, n, k, BLOCK_M=triton.next_power_of_2(m))
|
||||
_gemv[(triton.cdiv(n, BLOCK_N),)](weight, x2, y, m, n, k, BLOCK_M=triton.next_power_of_2(m),
|
||||
BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, num_warps=NUM_WARPS)
|
||||
return y.reshape(*x.shape[:-1], n)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,133 @@
|
||||
"""Triton kernels for int4 experts on gfx10 (RDNA2), where vLLM's int4 fused-MoE kernel pads
|
||||
every token to a 16-row tile and dot-products in bf16, which RDNA2 lacks: 2.85 ms for one token.
|
||||
|
||||
GPU int4 layout (prain.quant): q uint8 [rows, K/2], even k in the low nibble, bf16 scales
|
||||
[rows, K/64], w = (nibble - 8) * scale. Experts use gelu_tanh(gate) * up, gate = the first half
|
||||
of the w13 rows (Gemma 4).
|
||||
"""
|
||||
|
||||
import functools
|
||||
|
||||
import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
GROUP = 64
|
||||
DECODE_ROWS = 4 # up to this many tokens per round use moe_int4_decode
|
||||
|
||||
|
||||
@functools.cache
|
||||
def gfx10(index):
|
||||
"""RDNA2 GPUs: no bf16 dot instructions, where vLLM's bf16 GEMV/int4 kernels are slow."""
|
||||
return getattr(torch.cuda.get_device_properties(index), "gcnArchName", "").startswith("gfx10")
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _gelu_tanh(x): # 0.5 x (1 + tanh(z)) = x sigmoid(2z), z = sqrt(2/pi) (x + 0.044715 x^3)
|
||||
return x * tl.sigmoid(1.5957691216057308 * (x + 0.044715 * x * x * x))
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _gate_up(x_ptr, ids_ptr, table_ptr, q_ptr, s_ptr, h_ptr, TOPK: tl.constexpr,
|
||||
H: tl.constexpr, N2: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
|
||||
"""h[pair] = W13[slot] @ x[token] for one (token, route) pair; routes off the table skip."""
|
||||
pair = tl.program_id(0)
|
||||
slot = tl.load(table_ptr + tl.load(ids_ptr + pair))
|
||||
if slot >= 0:
|
||||
m = pair // TOPK
|
||||
rows = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
|
||||
rmask = rows < N2
|
||||
half = tl.arange(0, BLOCK_K // 2)
|
||||
q_base = q_ptr + slot.to(tl.int64) * (N2 * (H // 2))
|
||||
s_base = s_ptr + slot.to(tl.int64) * (N2 * (H // 64))
|
||||
acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
|
||||
for k0 in range(0, H, BLOCK_K):
|
||||
j = k0 // 2 + half
|
||||
b = tl.load(q_base + rows[:, None] * (H // 2) + j[None, :], mask=rmask[:, None],
|
||||
other=0).to(tl.int32)
|
||||
s = tl.load(s_base + rows[:, None] * (H // 64) + (j // 32)[None, :],
|
||||
mask=rmask[:, None], other=0.0).to(tl.float32)
|
||||
xe = tl.load(x_ptr + m * H + 2 * j).to(tl.float32)
|
||||
xo = tl.load(x_ptr + m * H + 2 * j + 1).to(tl.float32)
|
||||
lo = (b & 15).to(tl.float32) - 8.0
|
||||
hi = (b >> 4).to(tl.float32) - 8.0
|
||||
acc += tl.sum((lo * xe[None, :] + hi * xo[None, :]) * s, axis=1)
|
||||
tl.store(h_ptr + pair * N2 + rows, acc, mask=rmask)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _down(h_ptr, ids_ptr, table_ptr, w_ptr, q_ptr, s_ptr, y_ptr, H: tl.constexpr,
|
||||
I: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
|
||||
"""y[pair] = weight * W2[slot] @ (gelu_tanh(gate) * up) for one (token, route) pair."""
|
||||
pair = tl.program_id(0)
|
||||
slot = tl.load(table_ptr + tl.load(ids_ptr + pair))
|
||||
if slot >= 0:
|
||||
rows = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
|
||||
rmask = rows < H
|
||||
half = tl.arange(0, BLOCK_K // 2)
|
||||
q_base = q_ptr + slot.to(tl.int64) * (H * (I // 2))
|
||||
s_base = s_ptr + slot.to(tl.int64) * (H * (I // 64))
|
||||
h_base = h_ptr + pair * (2 * I)
|
||||
acc = tl.zeros((BLOCK_N,), dtype=tl.float32)
|
||||
for k0 in range(0, I, BLOCK_K):
|
||||
j = k0 // 2 + half
|
||||
kmask = 2 * j < I
|
||||
ae = _gelu_tanh(tl.load(h_base + 2 * j, mask=kmask, other=0.0)) * \
|
||||
tl.load(h_base + I + 2 * j, mask=kmask, other=0.0)
|
||||
ao = _gelu_tanh(tl.load(h_base + 2 * j + 1, mask=kmask, other=0.0)) * \
|
||||
tl.load(h_base + I + 2 * j + 1, mask=kmask, other=0.0)
|
||||
mask = rmask[:, None] & kmask[None, :]
|
||||
b = tl.load(q_base + rows[:, None] * (I // 2) + j[None, :], mask=mask,
|
||||
other=0).to(tl.int32)
|
||||
s = tl.load(s_base + rows[:, None] * (I // 64) + (j // 32)[None, :], mask=mask,
|
||||
other=0.0).to(tl.float32)
|
||||
lo = (b & 15).to(tl.float32) - 8.0
|
||||
hi = (b >> 4).to(tl.float32) - 8.0
|
||||
acc += tl.sum((lo * ae[None, :] + hi * ao[None, :]) * s, axis=1)
|
||||
tl.store(y_ptr + pair * H + rows, acc * tl.load(w_ptr + pair), mask=rmask)
|
||||
|
||||
|
||||
def moe_int4_decode(x, weights, ids, table, w13, w13s, w2, w2s, blocks=(16, 256, 32, 128)):
|
||||
"""sum over routes r of weights[m, r] * expert(x[m]) for the routes whose expert `table`
|
||||
maps to a slot (others add nothing). x [M, H] bf16 with few rows; ids [M, TOPK], padding
|
||||
pointing at the table's last row (-1); slot tensors in the GPU int4 layout above."""
|
||||
m, topk = ids.shape
|
||||
hidden, n2 = x.shape[1], w13.shape[1]
|
||||
pairs = m * topk
|
||||
bn1, bk1, bn2, bk2 = blocks
|
||||
assert hidden % bk1 == 0, "gate/up blocks must tile the hidden size (no k mask there)"
|
||||
x, ids = x.contiguous(), ids.contiguous()
|
||||
h = torch.empty(pairs, n2, dtype=torch.float32, device=x.device)
|
||||
y = torch.zeros(pairs, hidden, dtype=torch.float32, device=x.device)
|
||||
_gate_up[(pairs, triton.cdiv(n2, bn1))](x, ids, table, w13, w13s, h, TOPK=topk, H=hidden,
|
||||
N2=n2, BLOCK_N=bn1, BLOCK_K=bk1, num_warps=4)
|
||||
_down[(pairs, triton.cdiv(hidden, bn2))](h, ids, table, weights.float().contiguous(), w2,
|
||||
w2s, y, H=hidden, I=n2 // 2, BLOCK_N=bn2,
|
||||
BLOCK_K=bk2, num_warps=4)
|
||||
return y.view(m, topk, hidden).sum(1).to(x.dtype)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _unpack(src_ptr, dst_ptr, N: tl.constexpr, K: tl.constexpr, BLOCK_R: tl.constexpr,
|
||||
BLOCK_J: tl.constexpr):
|
||||
"""CPU int4 layout (torch's AVX2 order: [N/32, K, 16] bytes, rows d and d+16 share a byte)
|
||||
-> GPU layout [N, K/2]; prain.quant.to_gpu_int4 in one pass."""
|
||||
rows = tl.program_id(0) * BLOCK_R + tl.arange(0, BLOCK_R)
|
||||
j = tl.program_id(1) * BLOCK_J + tl.arange(0, BLOCK_J)
|
||||
mask = (rows < N)[:, None] & (j < K // 2)[None, :]
|
||||
r = rows % 32
|
||||
base = (rows // 32) * (K * 16) + r % 16
|
||||
shift = ((r // 16) * 4)[:, None]
|
||||
even = tl.load(src_ptr + base[:, None] + (2 * j * 16)[None, :], mask=mask, other=0)
|
||||
odd = tl.load(src_ptr + base[:, None] + ((2 * j + 1) * 16)[None, :], mask=mask, other=0)
|
||||
lo = (even.to(tl.int32) >> shift) & 15
|
||||
hi = (odd.to(tl.int32) >> shift) & 15
|
||||
tl.store(dst_ptr + rows[:, None] * (K // 2) + j[None, :], (lo | (hi << 4)).to(tl.uint8),
|
||||
mask=mask)
|
||||
|
||||
|
||||
def unpack_int4(src, dst, block=(32, 64)):
|
||||
"""dst (uint8 [N, K/2], GPU layout) <- src (the same bytes in the CPU layout, flat)."""
|
||||
n, half = dst.shape
|
||||
_unpack[(triton.cdiv(n, block[0]), triton.cdiv(half, block[1]))](
|
||||
src, dst, N=n, K=2 * half, BLOCK_R=block[0], BLOCK_J=block[1])
|
||||
+4
-2
@@ -236,8 +236,10 @@ COUNTERS = ("hits", "misses", "cpu", "prefetches", "prefetch_hits", "prefetch_wa
|
||||
|
||||
|
||||
def write_stats(path, caches):
|
||||
"""Atomically write {"layers": {"<layer>": {"uses": [...], "hits": n, ...}}}."""
|
||||
data = {"layers": {str(layer): {"uses": c.uses, **{k: getattr(c, k) for k in COUNTERS}}
|
||||
"""Atomically write {"layers": {"<layer>": {"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:
|
||||
|
||||
+716
@@ -0,0 +1,716 @@
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1, viewport-fit=cover">
|
||||
<meta name="color-scheme" content="light dark">
|
||||
<title>prain</title>
|
||||
<style>
|
||||
:root {
|
||||
--bg: #f1f3f6; --surface: #fff; --ink: #161b22; --muted: #5b6673; --faint: #8a94a0;
|
||||
--line: #dde2e8; --bubble: #e5e9ef; --code: #f6f7f9; --focus: #7c4dcc; --live: #f26b1d;
|
||||
--ok: #23946a; --bad: #c2362f;
|
||||
--ui: system-ui, -apple-system, "Segoe UI", Roboto, "Noto Sans", Cantarell, sans-serif;
|
||||
--display: ui-rounded, "SF Pro Rounded", Nunito, system-ui, -apple-system, sans-serif;
|
||||
--mono: ui-monospace, "SF Mono", "JetBrains Mono", "Cascadia Code", "DejaVu Sans Mono", Menlo, Consolas, monospace;
|
||||
}
|
||||
@media (prefers-color-scheme: dark) {
|
||||
:root { --bg: #0f1216; --surface: #161a20; --ink: #e6eaf0; --muted: #9aa4b1; --faint: #6d7784;
|
||||
--line: #262d36; --bubble: #232a34; --code: #1b2027; --focus: #a58bea; }
|
||||
}
|
||||
* { box-sizing: border-box; }
|
||||
html, body { height: 100%; margin: 0; }
|
||||
body { background: var(--bg); color: var(--ink); font: 16px/1.6 var(--ui); -webkit-font-smoothing: antialiased; }
|
||||
button { font: inherit; color: inherit; }
|
||||
:focus-visible { outline: 2px solid var(--focus); outline-offset: 2px; }
|
||||
[hidden] { display: none !important; }
|
||||
|
||||
.app { display: grid; grid-template-columns: minmax(0, 1fr) 340px; height: 100dvh; }
|
||||
.chat { display: flex; flex-direction: column; min-width: 0; height: 100dvh; }
|
||||
|
||||
/* header */
|
||||
.bar { display: flex; align-items: center; justify-content: space-between; gap: 12px; padding: 12px 20px; }
|
||||
.brand { display: flex; align-items: center; gap: 10px; min-width: 0; }
|
||||
.mark { display: grid; grid-template-columns: repeat(4, 6px); gap: 2px; }
|
||||
.mark i { width: 6px; height: 6px; border-radius: 2px; background: var(--line); transition: background .3s; }
|
||||
.mark i.on { background: var(--live); }
|
||||
.word { font: 700 20px/1 var(--display); letter-spacing: -.01em; }
|
||||
.model { font: 13px var(--mono); color: var(--faint); white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
|
||||
.bar-actions { display: flex; gap: 8px; align-items: center; }
|
||||
.ghost { border: 1px solid var(--line); background: var(--surface); border-radius: 999px; padding: 6px 14px; font-size: 14px; cursor: pointer; }
|
||||
.ghost:hover { border-color: var(--faint); }
|
||||
.pill { display: none; align-items: center; gap: 8px; }
|
||||
.dot { width: 8px; height: 8px; border-radius: 50%; background: var(--faint); flex: none; }
|
||||
.dot.ok { background: var(--ok); }
|
||||
.dot.live { background: var(--live); animation: pulse 1.2s ease-in-out infinite; }
|
||||
.dot.bad { background: var(--bad); }
|
||||
@keyframes pulse { 50% { opacity: .35; } }
|
||||
|
||||
/* thread */
|
||||
.log { flex: 1; overflow-y: auto; padding: 8px 20px 24px; overscroll-behavior: contain; }
|
||||
.thread { max-width: 760px; margin: 0 auto; }
|
||||
.empty { max-width: 560px; margin: 14vh auto 0; }
|
||||
.empty h1 { font: 650 clamp(28px, 4.5vw, 38px)/1.15 var(--display); letter-spacing: -.015em; margin: 0 0 10px; }
|
||||
.empty p { color: var(--muted); margin: 0 0 6px; }
|
||||
.empty .live-line { font-size: 14px; color: var(--faint); margin-bottom: 22px; font-variant-numeric: tabular-nums; }
|
||||
.chips { display: flex; flex-wrap: wrap; gap: 8px; margin-top: 22px; }
|
||||
.chip { border: 1px solid var(--line); background: var(--surface); border-radius: 999px; padding: 8px 14px; font-size: 14px; cursor: pointer; text-align: left; }
|
||||
.chip:hover { border-color: var(--faint); }
|
||||
|
||||
.msg { margin: 22px 0; }
|
||||
.msg.user { display: flex; justify-content: flex-end; }
|
||||
.bubble { background: var(--bubble); padding: 10px 16px; border-radius: 20px 20px 6px 20px; max-width: min(85%, 600px); white-space: pre-wrap; overflow-wrap: anywhere; }
|
||||
.body { overflow-wrap: anywhere; }
|
||||
.body > :first-child { margin-top: 0; }
|
||||
.body > :last-child { margin-bottom: 0; }
|
||||
.body p { margin: 0 0 .8em; }
|
||||
.body h3, .body h4, .body h5 { font: 650 1.08em/1.3 var(--display); margin: 1.3em 0 .4em; }
|
||||
.body ul, .body ol { padding-left: 1.4em; margin: 0 0 .8em; }
|
||||
.body li { margin: .25em 0; }
|
||||
.body blockquote { border-left: 3px solid var(--line); margin: .8em 0; padding-left: 12px; color: var(--muted); }
|
||||
.body hr { border: 0; border-top: 1px solid var(--line); margin: 1.2em 0; }
|
||||
.body a { color: var(--focus); }
|
||||
.body code { font: .88em var(--mono); background: var(--code); border: 1px solid var(--line); padding: .05em .35em; border-radius: 6px; }
|
||||
.code { border: 1px solid var(--line); background: var(--code); border-radius: 12px; margin: .9em 0; overflow: hidden; }
|
||||
.code-head { display: flex; justify-content: space-between; align-items: center; padding: 4px 6px 4px 14px; border-bottom: 1px solid var(--line); font: 12px var(--mono); color: var(--faint); }
|
||||
.code pre { margin: 0; padding: 12px 14px; overflow-x: auto; font: 13.5px/1.55 var(--mono); tab-size: 4; }
|
||||
.code code { background: none; border: 0; padding: 0; font-size: inherit; }
|
||||
.table { overflow-x: auto; margin: .9em 0; }
|
||||
.body table { border-collapse: collapse; font-size: 14px; }
|
||||
.body th, .body td { border-bottom: 1px solid var(--line); padding: 6px 12px 6px 0; text-align: left; vertical-align: top; }
|
||||
.small { border: 0; background: none; border-radius: 6px; padding: 3px 8px; font-size: 12.5px; color: var(--faint); cursor: pointer; }
|
||||
.small:hover { color: var(--ink); background: var(--bubble); }
|
||||
.meta { display: flex; flex-wrap: wrap; gap: 4px 12px; align-items: center; margin-top: 6px; font-size: 12.5px; color: var(--faint); font-variant-numeric: tabular-nums; }
|
||||
.note { display: inline-block; font-size: 13px; color: var(--muted); background: var(--bubble); border-radius: 10px; padding: 5px 10px; margin-bottom: 10px; }
|
||||
.note.bad { color: var(--bad); }
|
||||
.is-draft .body, .rest { color: var(--faint); }
|
||||
.typing::after { content: ""; display: inline-block; width: .5em; height: 1.05em; margin-left: 3px; vertical-align: -.18em; border-radius: 2px; background: var(--live); animation: pulse 1s ease-in-out infinite; }
|
||||
|
||||
/* composer */
|
||||
.composer { padding: 0 20px calc(14px + env(safe-area-inset-bottom)); }
|
||||
.field { max-width: 760px; margin: 0 auto; display: flex; align-items: flex-end; gap: 8px; background: var(--surface); border: 1px solid var(--line); border-radius: 22px; padding: 7px 7px 7px 16px; box-shadow: 0 10px 30px -18px rgb(20 30 45 / .35); }
|
||||
.field:focus-within { border-color: color-mix(in srgb, var(--focus) 50%, var(--line)); }
|
||||
textarea { flex: 1; border: 0; outline: 0; resize: none; background: transparent; color: inherit; font: 16px/1.5 var(--ui); max-height: 40vh; padding: 6px 0; }
|
||||
.send { width: 36px; height: 36px; flex: none; border-radius: 50%; border: 0; background: var(--ink); color: var(--surface); display: grid; place-items: center; cursor: pointer; }
|
||||
.send:disabled { opacity: .3; cursor: default; }
|
||||
.send svg { width: 18px; height: 18px; }
|
||||
.hint { max-width: 760px; margin: 6px auto 0; min-height: 18px; padding: 0 16px; font-size: 12.5px; color: var(--faint); display: flex; align-items: center; gap: 8px; }
|
||||
.cells { display: inline-grid; grid-template-columns: repeat(8, 5px); gap: 2px; }
|
||||
.cells i { width: 5px; height: 5px; border-radius: 1.5px; background: var(--live); animation: pulse 1s ease-in-out infinite; }
|
||||
.cells i:nth-child(2n) { animation-delay: .25s; } .cells i:nth-child(3n) { animation-delay: .5s; }
|
||||
|
||||
/* monitor */
|
||||
.monitor { background: var(--surface); border-left: 1px solid var(--line); overflow-y: auto; padding: 18px 20px 28px; font-size: 14px; }
|
||||
.m-top { display: flex; justify-content: space-between; align-items: start; }
|
||||
.m-title { font: 650 17px/1.2 var(--display); }
|
||||
.m-sub { font: 12.5px var(--mono); color: var(--faint); margin-top: 2px; }
|
||||
.close { display: none; }
|
||||
.state { display: flex; align-items: center; gap: 8px; margin: 14px 0 6px; color: var(--muted); }
|
||||
.speed { display: flex; align-items: baseline; gap: 6px; }
|
||||
.speed b { font: 650 40px/1 var(--display); letter-spacing: -.02em; font-variant-numeric: tabular-nums; }
|
||||
.speed.idle b { color: var(--faint); }
|
||||
.speed span { color: var(--faint); }
|
||||
.spark { width: 100%; height: 28px; margin: 6px 0 16px; display: block; }
|
||||
.spark path { fill: none; stroke: var(--live); stroke-width: 1.5; vector-effect: non-scaling-stroke; }
|
||||
.scan { background: #16121d; color: #c9c1d6; border-radius: 14px; padding: 12px 12px 10px; }
|
||||
.scan-head { display: flex; justify-content: space-between; font-size: 12.5px; margin-bottom: 8px; font-variant-numeric: tabular-nums; }
|
||||
.scan-head b { font-weight: 600; color: #f2edf8; }
|
||||
.scan canvas { display: block; width: 100%; aspect-ratio: 128 / 64; image-rendering: pixelated; border-radius: 4px; background: #2c2538; }
|
||||
.axis { display: flex; justify-content: space-between; font: 11px var(--mono); color: #8f86a0; margin-top: 4px; }
|
||||
.legend { display: flex; flex-wrap: wrap; gap: 4px 12px; margin: 8px 0 0; padding: 0; list-style: none; font-size: 12px; }
|
||||
.legend i { display: inline-block; width: 9px; height: 9px; border-radius: 2px; margin-right: 5px; vertical-align: -1px; }
|
||||
.scan-empty { font-size: 12.5px; color: #8f86a0; padding: 18px 4px; }
|
||||
.kv { margin: 16px 0 0; }
|
||||
.kv div { display: flex; justify-content: space-between; align-items: baseline; gap: 12px; padding: 7px 0; border-bottom: 1px solid var(--line); }
|
||||
.kv dt { color: var(--muted); }
|
||||
.kv dd { margin: 0; font-variant-numeric: tabular-nums; text-align: right; }
|
||||
.kv small { color: var(--faint); }
|
||||
.meter { display: block; width: 84px; height: 4px; border-radius: 2px; background: var(--bubble); margin-top: 5px; margin-left: auto; overflow: hidden; }
|
||||
.meter i { display: block; height: 100%; width: 0; background: var(--ink); transition: width .4s; }
|
||||
.m-h { font: 650 13px/1 var(--display); color: var(--faint); text-transform: uppercase; letter-spacing: .06em; margin: 24px 0 0; }
|
||||
.scrim { display: none; }
|
||||
|
||||
@media (max-width: 999px) {
|
||||
.app { grid-template-columns: 1fr; }
|
||||
.pill, .close { display: inline-flex; }
|
||||
.monitor { position: fixed; top: 0; right: 0; bottom: 0; width: min(360px, 92vw); z-index: 20; transform: translateX(100%); transition: transform .25s ease; box-shadow: -20px 0 40px -20px rgb(0 0 0 / .35); padding-top: calc(18px + env(safe-area-inset-top)); }
|
||||
.monitor.open { transform: none; }
|
||||
.scrim.open { display: block; position: fixed; inset: 0; z-index: 10; background: rgb(10 14 20 / .35); }
|
||||
}
|
||||
@media (max-width: 560px) {
|
||||
.bar { padding: 10px 14px; } .log { padding: 6px 14px 18px; } .composer { padding-left: 10px; padding-right: 10px; }
|
||||
.model, .hint .keys { display: none; } .empty { margin-top: 8vh; }
|
||||
}
|
||||
@media (prefers-reduced-motion: reduce) {
|
||||
*, *::after { animation: none !important; transition: none !important; }
|
||||
}
|
||||
</style>
|
||||
|
||||
<div class="app">
|
||||
<main class="chat">
|
||||
<header class="bar">
|
||||
<div class="brand">
|
||||
<span class="mark" id="mark" aria-hidden="true"><i></i><i></i><i></i><i></i><i></i><i></i><i></i><i></i></span>
|
||||
<span class="word">prain</span>
|
||||
<span class="model" id="model"></span>
|
||||
</div>
|
||||
<div class="bar-actions">
|
||||
<button class="ghost pill" id="pill" aria-controls="monitor" aria-expanded="false">
|
||||
<span class="dot" id="pill-dot"></span><span id="pill-text">Monitor</span></button>
|
||||
<button class="ghost" id="new-chat">New chat</button>
|
||||
</div>
|
||||
</header>
|
||||
|
||||
<section class="log" id="log" role="log">
|
||||
<div class="thread" id="thread">
|
||||
<div class="empty" id="empty">
|
||||
<h1>Ready when you are</h1>
|
||||
<p>prain moves the experts your message needs into VRAM while you type, so the answer starts sooner.</p>
|
||||
<div class="live-line" id="empty-live"></div>
|
||||
<div class="chips">
|
||||
<button class="chip">Explain what git rebase -i HEAD~3 does</button>
|
||||
<button class="chip">Write a bash one-liner that counts lines of code per file type</button>
|
||||
<button class="chip">On Gentoo, how do I rebuild the packages affected by a USE flag change?</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</section>
|
||||
|
||||
<form class="composer" id="composer">
|
||||
<div class="field">
|
||||
<textarea id="draft" rows="1" placeholder="Message prain" aria-label="Message"></textarea>
|
||||
<button class="send" id="send" type="submit" aria-label="Send" disabled>
|
||||
<svg viewBox="0 0 24 24" aria-hidden="true"><path id="send-icon" fill="currentColor" d="M12 4l7 7-1.4 1.4-4.6-4.6V20h-2V7.8l-4.6 4.6L5 11z"/></svg>
|
||||
</button>
|
||||
</div>
|
||||
<div class="hint">
|
||||
<span id="warm" hidden><span class="cells" aria-hidden="true"><i></i><i></i><i></i><i></i><i></i><i></i><i></i><i></i></span> Loading experts for your message</span>
|
||||
<span class="keys" id="keys">Enter sends · Shift+Enter adds a line</span>
|
||||
</div>
|
||||
</form>
|
||||
</main>
|
||||
|
||||
<div class="scrim" id="scrim"></div>
|
||||
<aside class="monitor" id="monitor" aria-label="Model monitor">
|
||||
<div class="m-top">
|
||||
<div><div class="m-title">Monitor</div><div class="m-sub" id="m-model">connecting…</div></div>
|
||||
<button class="ghost close" id="close" aria-label="Close the monitor">Close</button>
|
||||
</div>
|
||||
<div class="state"><span class="dot" id="dot"></span><span id="state">Connecting…</span></div>
|
||||
<div class="speed idle" id="speed-box"><b id="speed">–</b><span id="speed-unit">tokens/s</span></div>
|
||||
<svg class="spark" viewBox="0 0 60 28" preserveAspectRatio="none" aria-hidden="true"><path id="spark"/></svg>
|
||||
|
||||
<section class="scan" aria-label="Experts by layer">
|
||||
<div class="scan-head"><span>Experts</span><span id="scan-count"></span></div>
|
||||
<div id="scan-body">
|
||||
<canvas id="scan" width="128" height="30" role="img" aria-label="Expert map"></canvas>
|
||||
<div class="axis"><span>layer 1 · expert 1</span><span id="axis-end"></span></div>
|
||||
<ul class="legend">
|
||||
<li><i style="background:#6a35a0"></i>In VRAM</li>
|
||||
<li><i style="background:#ff8a2a"></i>Used now</li>
|
||||
<li><i style="background:#45d3c7"></i>Copied in</li>
|
||||
</ul>
|
||||
</div>
|
||||
<div class="scan-empty" id="scan-empty" hidden>The map appears once prain writes expert stats: set stats_path in the placement file and pass it with --config.</div>
|
||||
</section>
|
||||
|
||||
<dl class="kv">
|
||||
<div><dt>VRAM hits</dt><dd id="hits">–</dd></div>
|
||||
<div><dt>Copies</dt><dd id="copies">–</dd></div>
|
||||
<div><dt>Context in use</dt><dd><span id="kv">–</span><span class="meter"><i id="kv-bar"></i></span></dd></div>
|
||||
<div><dt>Prompt cache reuse</dt><dd id="prefix">–</dd></div>
|
||||
<div><dt>Requests</dt><dd id="reqs">–</dd></div>
|
||||
</dl>
|
||||
<h3 class="m-h">GPU</h3>
|
||||
<dl class="kv">
|
||||
<div><dt>VRAM</dt><dd><span id="vram">–</span><span class="meter"><i id="vram-bar"></i></span></dd></div>
|
||||
<div><dt>Load</dt><dd id="load">–</dd></div>
|
||||
<div><dt>Temperature</dt><dd id="temp">–</dd></div>
|
||||
<div><dt>Power</dt><dd id="power">–</dd></div>
|
||||
<div><dt>Fan</dt><dd id="fan">–</dd></div>
|
||||
</dl>
|
||||
</aside>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
const $ = (id) => document.getElementById(id);
|
||||
const log = $("log"), thread = $("thread"), draft = $("draft"), send = $("send");
|
||||
const coarse = matchMedia("(pointer: coarse)").matches;
|
||||
const newId = () => Math.random().toString(36).slice(2) + Date.now().toString(36);
|
||||
let chat = load() || {session: newId(), messages: []};
|
||||
let streaming = false, controller = null;
|
||||
|
||||
function load() {
|
||||
try { return JSON.parse(localStorage.getItem("prain.chat")); } catch { return null; }
|
||||
}
|
||||
function save() {
|
||||
try { localStorage.setItem("prain.chat", JSON.stringify(chat)); } catch {}
|
||||
}
|
||||
const plain = () => chat.messages.map(({role, content}) => ({role, content}));
|
||||
|
||||
/* ---------- markdown: escape first, then a small safe subset ---------- */
|
||||
const esc = (s) => s.replace(/[&<>"']/g, (c) => ({"&": "&", "<": "<", ">": ">", '"': """, "'": "'"}[c]));
|
||||
function inline(text) {
|
||||
return text.split(/(`[^`\n]+`)/g).map((part, i) => i % 2
|
||||
? `<code>${esc(part.slice(1, -1))}</code>`
|
||||
: esc(part)
|
||||
.replace(/\[([^\]]+)\]\((https?:\/\/[^\s)]+)\)/g, '<a href="$2" target="_blank" rel="noopener noreferrer">$1</a>')
|
||||
.replace(/\*\*([^*]+)\*\*/g, "<strong>$1</strong>")
|
||||
.replace(/(^|[^*\w])\*([^*\s][^*]*)\*(?![*\w])/g, "$1<em>$2</em>")
|
||||
).join("");
|
||||
}
|
||||
const ITEM = /^\s*([-*+]|\d+[.)])\s+/;
|
||||
function md(src) {
|
||||
const lines = src.split("\n"), out = [];
|
||||
let i = 0, para = [];
|
||||
const flush = () => { if (para.length) out.push(`<p>${para.map(inline).join("<br>")}</p>`); para = []; };
|
||||
while (i < lines.length) {
|
||||
const line = lines[i];
|
||||
let m;
|
||||
if ((m = line.match(/^\s*```\s*([\w+#.-]*)/))) {
|
||||
flush();
|
||||
const code = [];
|
||||
for (i++; i < lines.length && !/^\s*```\s*$/.test(lines[i]); i++) code.push(lines[i]);
|
||||
i++;
|
||||
out.push(`<div class="code"><div class="code-head"><span>${esc(m[1] || "code")}</span>` +
|
||||
`<button class="small" data-copy>Copy</button></div><pre><code>${esc(code.join("\n"))}</code></pre></div>`);
|
||||
} else if ((m = line.match(/^(#{1,6})\s+(.*)/))) {
|
||||
flush();
|
||||
const level = Math.min(m[1].length + 2, 5);
|
||||
out.push(`<h${level}>${inline(m[2])}</h${level}>`);
|
||||
i++;
|
||||
} else if (/^\s*\|.*\|\s*$/.test(line) && /^\s*\|?\s*:?-{2,}/.test(lines[i + 1] || "")) {
|
||||
flush();
|
||||
const rows = [];
|
||||
while (i < lines.length && /^\s*\|.*\|\s*$/.test(lines[i])) rows.push(lines[i++]);
|
||||
const cells = (row, tag) => row.trim().replace(/^\||\|$/g, "").split("|")
|
||||
.map((c) => `<${tag}>${inline(c.trim())}</${tag}>`).join("");
|
||||
out.push(`<div class="table"><table><thead><tr>${cells(rows[0], "th")}</tr></thead><tbody>` +
|
||||
rows.slice(2).map((r) => `<tr>${cells(r, "td")}</tr>`).join("") + "</tbody></table></div>");
|
||||
} else if (ITEM.test(line)) {
|
||||
flush();
|
||||
const ordered = /^\s*\d/.test(line), items = [];
|
||||
const start = ordered ? parseInt(line.trim(), 10) : 1;
|
||||
while (i < lines.length && (ITEM.test(lines[i]) || (items.length && /^\s{2,}\S/.test(lines[i])))) {
|
||||
if (ITEM.test(lines[i])) items.push(lines[i].replace(ITEM, ""));
|
||||
else items[items.length - 1] += " " + lines[i].trim();
|
||||
i++;
|
||||
}
|
||||
const tag = ordered ? "ol" : "ul";
|
||||
out.push(`<${tag}${ordered && start !== 1 ? ` start="${start}"` : ""}>` +
|
||||
items.map((t) => `<li>${inline(t)}</li>`).join("") + `</${tag}>`);
|
||||
} else if (/^\s*>/.test(line)) {
|
||||
flush();
|
||||
const quote = [];
|
||||
while (i < lines.length && /^\s*>/.test(lines[i])) quote.push(lines[i++].replace(/^\s*>\s?/, ""));
|
||||
out.push(`<blockquote>${md(quote.join("\n"))}</blockquote>`);
|
||||
} else if (/^\s*([-*_])(\s*\1){2,}\s*$/.test(line)) {
|
||||
flush();
|
||||
out.push("<hr>");
|
||||
i++;
|
||||
} else if (!line.trim()) {
|
||||
flush();
|
||||
i++;
|
||||
} else {
|
||||
para.push(line);
|
||||
i++;
|
||||
}
|
||||
}
|
||||
flush();
|
||||
return out.join("");
|
||||
}
|
||||
|
||||
/* ---------- thread ---------- */
|
||||
function nearBottom() { return log.scrollHeight - log.scrollTop - log.clientHeight < 120; }
|
||||
function toBottom() { log.scrollTop = log.scrollHeight; }
|
||||
function el(tag, cls, text) {
|
||||
const node = document.createElement(tag);
|
||||
if (cls) node.className = cls;
|
||||
if (text !== undefined) node.textContent = text;
|
||||
return node;
|
||||
}
|
||||
function addUser(text) {
|
||||
const row = el("div", "msg user");
|
||||
row.appendChild(el("div", "bubble", text));
|
||||
thread.appendChild(row);
|
||||
}
|
||||
function addBot() {
|
||||
const row = el("div", "msg bot");
|
||||
const view = {row, note: el("div", "note"), body: el("div", "body"), meta: el("div", "meta")};
|
||||
view.note.hidden = true;
|
||||
view.meta.hidden = true;
|
||||
row.append(view.note, view.body, view.meta);
|
||||
thread.appendChild(row);
|
||||
return view;
|
||||
}
|
||||
function paint(view, text, rest, live) {
|
||||
view.body.innerHTML = md(text);
|
||||
if (rest) {
|
||||
const target = view.body.lastElementChild && view.body.lastElementChild.tagName === "P"
|
||||
? view.body.lastElementChild : view.body;
|
||||
target.appendChild(el("span", "rest", rest));
|
||||
}
|
||||
view.body.classList.toggle("typing", live);
|
||||
}
|
||||
function finish(view, message) {
|
||||
view.meta.hidden = false;
|
||||
view.meta.replaceChildren(el("span", "", message.meta || ""));
|
||||
const copy = el("button", "small", "Copy");
|
||||
copy.addEventListener("click", () => copied(copy, message.content));
|
||||
view.meta.appendChild(copy);
|
||||
}
|
||||
function render() {
|
||||
thread.querySelectorAll(".msg").forEach((n) => n.remove());
|
||||
$("empty").hidden = chat.messages.length > 0;
|
||||
for (const m of chat.messages) {
|
||||
if (m.role === "user") { addUser(m.content); continue; }
|
||||
const view = addBot();
|
||||
view.row.classList.toggle("is-draft", !!m.draft);
|
||||
paint(view, m.content, "", false);
|
||||
finish(view, m);
|
||||
}
|
||||
toBottom();
|
||||
}
|
||||
|
||||
async function copyText(text) {
|
||||
try { await navigator.clipboard.writeText(text); return true; } catch {}
|
||||
const area = el("textarea"); // plain http: no Clipboard API
|
||||
area.value = text;
|
||||
area.setAttribute("readonly", "");
|
||||
area.style.cssText = "position:fixed;top:0;opacity:0";
|
||||
document.body.appendChild(area);
|
||||
area.select();
|
||||
area.setSelectionRange(0, text.length);
|
||||
let ok = false;
|
||||
try { ok = document.execCommand("copy"); } catch {}
|
||||
area.remove();
|
||||
return ok;
|
||||
}
|
||||
async function copied(button, text) {
|
||||
const ok = await copyText(text);
|
||||
button.textContent = ok ? "Copied" : "Copy failed";
|
||||
setTimeout(() => (button.textContent = "Copy"), 1500);
|
||||
}
|
||||
thread.addEventListener("click", (event) => {
|
||||
const button = event.target.closest("[data-copy]");
|
||||
if (button) copied(button, button.closest(".code").querySelector("pre").textContent);
|
||||
const chip = event.target.closest(".chip");
|
||||
if (chip) { draft.value = chip.textContent; draft.focus(); draft.dispatchEvent(new Event("input")); }
|
||||
});
|
||||
|
||||
/* ---------- sending ---------- */
|
||||
const ICON_SEND = "M12 4l7 7-1.4 1.4-4.6-4.6V20h-2V7.8l-4.6 4.6L5 11z";
|
||||
const ICON_STOP = "M7 7h10v10H7z";
|
||||
function setStreaming(on) {
|
||||
streaming = on;
|
||||
log.setAttribute("aria-busy", on);
|
||||
$("send-icon").setAttribute("d", on ? ICON_STOP : ICON_SEND);
|
||||
send.setAttribute("aria-label", on ? "Stop" : "Send");
|
||||
send.disabled = !on && !draft.value.trim();
|
||||
if (!on) lightMark(0);
|
||||
}
|
||||
function lightMark(n) {
|
||||
const cells = [...$("mark").children];
|
||||
cells.forEach((c) => c.classList.remove("on"));
|
||||
for (let k = 0; k < n; k++) cells[Math.floor(Math.random() * cells.length)].classList.add("on");
|
||||
}
|
||||
|
||||
async function submit() {
|
||||
const text = draft.value.trim();
|
||||
if (streaming || !text) return;
|
||||
const conv = chat; // "New chat" mid-answer must not receive this answer
|
||||
$("empty").hidden = true;
|
||||
conv.messages.push({role: "user", content: text});
|
||||
addUser(text);
|
||||
save();
|
||||
draft.value = "";
|
||||
autosize();
|
||||
hideWarm();
|
||||
setStreaming(true);
|
||||
toBottom();
|
||||
const view = addBot();
|
||||
const st = {mode: "full", draft: "", final: "", refined: false};
|
||||
let start = performance.now(), first = 0, pieces = 0, stopped = false, failure = null;
|
||||
const show = () => {
|
||||
const follow = nearBottom();
|
||||
const text = st.mode === "draft" ? st.draft : st.final;
|
||||
const rest = st.mode === "refine" && st.draft.startsWith(st.final) ? st.draft.slice(st.final.length) : "";
|
||||
view.row.classList.toggle("is-draft", st.mode === "draft");
|
||||
paint(view, text, rest, true);
|
||||
if (follow) toBottom();
|
||||
};
|
||||
const note = (text, bad) => { view.note.hidden = !text; view.note.textContent = text || ""; view.note.classList.toggle("bad", !!bad); };
|
||||
controller = new AbortController();
|
||||
try {
|
||||
const response = await fetch("/chat", {method: "POST", signal: controller.signal,
|
||||
body: JSON.stringify({session: chat.session, messages: plain()})});
|
||||
if (!response.ok) throw new Error(`HTTP ${response.status}: ${(await response.text()).slice(0, 200)}`);
|
||||
const reader = response.body.getReader(), decoder = new TextDecoder();
|
||||
let buffer = "";
|
||||
show();
|
||||
for (;;) {
|
||||
const {done, value} = await reader.read();
|
||||
if (done) break;
|
||||
buffer += decoder.decode(value, {stream: true});
|
||||
const lines = buffer.split("\n");
|
||||
buffer = lines.pop();
|
||||
for (const line of lines) {
|
||||
if (!line.startsWith("data: ") || line === "data: [DONE]") continue;
|
||||
const data = JSON.parse(line.slice(6));
|
||||
if (data.prain === "draft") {
|
||||
st.mode = "draft";
|
||||
note("Another answer is using the GPU, so this is a draft from the experts already in VRAM.");
|
||||
} else if (data.prain === "refine") {
|
||||
st.mode = "refine";
|
||||
st.refined = true;
|
||||
start = performance.now(), first = 0, pieces = 0;
|
||||
note("Refining the draft with all experts…");
|
||||
} else if (data.prain === "error") {
|
||||
st.mode = "draft";
|
||||
st.refined = false;
|
||||
note(`Refining failed (HTTP ${data.status}); this is the draft.`, true);
|
||||
} else {
|
||||
const piece = data.choices?.[0]?.delta?.content; // usage chunks have no choices
|
||||
if (!piece) continue;
|
||||
if (pieces++ === 0) first = performance.now();
|
||||
if (st.mode === "draft") st.draft += piece; else st.final += piece;
|
||||
lightMark(2);
|
||||
}
|
||||
show();
|
||||
}
|
||||
}
|
||||
} catch (error) {
|
||||
if (error.name === "AbortError") stopped = true; else failure = error;
|
||||
}
|
||||
controller = null;
|
||||
const content = st.mode === "draft" ? st.draft : st.final;
|
||||
if (conv !== chat) { setStreaming(false); return; }
|
||||
if (failure && !content) {
|
||||
view.row.remove();
|
||||
conv.messages.pop();
|
||||
save();
|
||||
const box = addBot();
|
||||
box.note.hidden = false;
|
||||
box.note.classList.add("bad");
|
||||
box.note.textContent = `The model didn't answer (${failure.message}). Your message is back in the box: ` +
|
||||
"send it again once the monitor shows the model as ready.";
|
||||
draft.value = text;
|
||||
autosize();
|
||||
setStreaming(false);
|
||||
return;
|
||||
}
|
||||
const seconds = (performance.now() - first) / 1000;
|
||||
const parts = [];
|
||||
if (first) parts.push(`${((first - start) / 1000).toFixed(1)} s to first token`);
|
||||
if (pieces > 1) parts.push(`${((pieces - 1) / seconds).toFixed(1)} tokens/s`);
|
||||
if (st.refined && st.draft) {
|
||||
let same = 0;
|
||||
while (same < st.draft.length && st.draft[same] === st.final[same]) same++;
|
||||
parts.push(`refined, kept ${Math.round(100 * same / st.draft.length)}% of the draft`);
|
||||
} else if (st.mode === "draft") parts.push("draft only");
|
||||
if (stopped) parts.push("stopped");
|
||||
if (failure) parts.push("cut off: " + failure.message);
|
||||
const message = {role: "assistant", content, meta: parts.join(" · "), draft: st.mode === "draft"};
|
||||
conv.messages.push(message);
|
||||
save();
|
||||
if (st.refined || !view.note.classList.contains("bad")) note("");
|
||||
paint(view, content, "", false);
|
||||
finish(view, message);
|
||||
setStreaming(false);
|
||||
}
|
||||
|
||||
$("composer").addEventListener("submit", (event) => {
|
||||
event.preventDefault();
|
||||
if (streaming) controller && controller.abort(); else submit();
|
||||
});
|
||||
draft.addEventListener("keydown", (event) => {
|
||||
if (event.key === "Enter" && !event.shiftKey && !event.isComposing && !coarse) {
|
||||
event.preventDefault();
|
||||
submit();
|
||||
}
|
||||
});
|
||||
if (coarse) $("keys").textContent = "Return adds a line";
|
||||
|
||||
function autosize() {
|
||||
draft.style.height = "auto";
|
||||
draft.style.height = Math.min(draft.scrollHeight, innerHeight * 0.4) + "px";
|
||||
if (!streaming) send.disabled = !draft.value.trim();
|
||||
}
|
||||
let warmTimer = 0;
|
||||
function hideWarm() { clearTimeout(warmTimer); $("warm").hidden = true; $("keys").hidden = false; }
|
||||
draft.addEventListener("input", () => {
|
||||
autosize();
|
||||
if (streaming || !draft.value.trim()) return;
|
||||
fetch("/warm", {method: "POST", body: JSON.stringify({session: chat.session, messages: plain(), draft: draft.value})})
|
||||
.then((response) => {
|
||||
if (response.status !== 202) return;
|
||||
$("warm").hidden = false;
|
||||
$("keys").hidden = true;
|
||||
clearTimeout(warmTimer);
|
||||
warmTimer = setTimeout(hideWarm, 2500);
|
||||
}).catch(() => {});
|
||||
});
|
||||
|
||||
$("new-chat").addEventListener("click", () => {
|
||||
if (controller) controller.abort();
|
||||
chat = {session: newId(), messages: []};
|
||||
save();
|
||||
render();
|
||||
draft.focus();
|
||||
});
|
||||
|
||||
/* ---------- monitor ---------- */
|
||||
const monitor = $("monitor"), scrim = $("scrim"), pill = $("pill");
|
||||
function openMonitor(on) {
|
||||
monitor.classList.toggle("open", on);
|
||||
scrim.classList.toggle("open", on);
|
||||
pill.setAttribute("aria-expanded", on);
|
||||
}
|
||||
pill.addEventListener("click", () => openMonitor(!monitor.classList.contains("open")));
|
||||
$("close").addEventListener("click", () => openMonitor(false));
|
||||
scrim.addEventListener("click", () => openMonitor(false));
|
||||
addEventListener("keydown", (event) => { if (event.key === "Escape") openMonitor(false); });
|
||||
|
||||
const fmt = (n, digits = 0) => n.toLocaleString(undefined, {maximumFractionDigits: digits, minimumFractionDigits: digits});
|
||||
const GIB = 2 ** 30;
|
||||
const history = []; // polls of the last minute
|
||||
const speeds = [];
|
||||
let lastSpeed = 0, scanState = null;
|
||||
|
||||
function set(id, text) { $(id).textContent = text; }
|
||||
function meter(id, share) { $(id).style.width = `${Math.max(0, Math.min(1, share)) * 100}%`; }
|
||||
function setState(kind, text) {
|
||||
for (const id of ["dot", "pill-dot"]) $(id).className = "dot " + kind;
|
||||
set("state", text);
|
||||
}
|
||||
|
||||
const SCAN = {off: [44, 37, 56], vram: [106, 53, 160], warm: [255, 107, 26], hot: [255, 212, 71], copy: [69, 211, 199]};
|
||||
const mix = (a, b, t) => [0, 1, 2].map((k) => a[k] + (b[k] - a[k]) * t);
|
||||
function drawScan(experts) {
|
||||
const canvas = $("scan");
|
||||
const layers = Object.keys(experts.layers).map(Number).sort((a, b) => a - b).map((k) => experts.layers[k]);
|
||||
const rows = layers.length, cols = layers[0].uses.length, n = rows * cols;
|
||||
if (!scanState || scanState.heat.length !== n) {
|
||||
canvas.width = cols;
|
||||
canvas.height = rows;
|
||||
scanState = {heat: new Float32Array(n), flash: new Float32Array(n), uses: null, vram: null};
|
||||
}
|
||||
const uses = new Float64Array(n), vram = new Uint8Array(n);
|
||||
layers.forEach((layer, l) => {
|
||||
layer.uses.forEach((u, e) => (uses[l * cols + e] = u));
|
||||
(layer.vram || []).forEach((e) => (vram[l * cols + e] = 1));
|
||||
});
|
||||
const {heat, flash} = scanState, ctx = canvas.getContext("2d"), image = ctx.createImageData(cols, rows);
|
||||
let resident = 0, active = 0;
|
||||
for (let k = 0; k < n; k++) {
|
||||
const fired = scanState.uses ? uses[k] - scanState.uses[k] : 0;
|
||||
heat[k] = Math.min(1, heat[k] * 0.5 + (fired > 0 ? 0.55 + 0.1 * Math.min(fired, 4) : 0));
|
||||
flash[k] = scanState.vram && vram[k] && !scanState.vram[k] ? 1 : flash[k] * 0.4;
|
||||
resident += vram[k];
|
||||
active += fired > 0;
|
||||
let c = vram[k] ? SCAN.vram : SCAN.off;
|
||||
if (heat[k] > 0.02) c = mix(c, mix(SCAN.warm, SCAN.hot, heat[k]), Math.min(1, heat[k] * 1.5));
|
||||
if (flash[k] > 0.02) c = mix(c, SCAN.copy, flash[k] * 0.85);
|
||||
image.data.set([c[0], c[1], c[2], 255], k * 4);
|
||||
}
|
||||
ctx.putImageData(image, 0, 0);
|
||||
scanState.uses = uses;
|
||||
scanState.vram = vram;
|
||||
set("scan-count", `${fmt(resident)} of ${fmt(n)} in VRAM`);
|
||||
set("axis-end", `layer ${rows} · expert ${cols}`);
|
||||
canvas.setAttribute("aria-label", `${resident} of ${n} experts in VRAM; ${active} used since the last update`);
|
||||
set("empty-live", `${fmt(resident)} of ${fmt(n)} experts are in VRAM right now.`);
|
||||
}
|
||||
|
||||
function totals(experts) {
|
||||
const t = {hits: 0, misses: 0, cpu: 0};
|
||||
for (const layer of Object.values(experts.layers)) for (const k in t) t[k] += layer[k] || 0;
|
||||
return t;
|
||||
}
|
||||
|
||||
async function poll() {
|
||||
let s;
|
||||
try {
|
||||
const response = await fetch("/status", {cache: "no-store"});
|
||||
s = await response.json();
|
||||
} catch {
|
||||
setState("bad", "Can't reach prain web");
|
||||
set("pill-text", "Offline");
|
||||
return;
|
||||
}
|
||||
set("model", s.model);
|
||||
set("m-model", s.model);
|
||||
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);
|
||||
while (history.length > 2 && now - history[0].t > 60) history.shift();
|
||||
|
||||
if (!v) {
|
||||
setState("bad", "Model offline: starting or restarting");
|
||||
set("pill-text", "Offline");
|
||||
} else {
|
||||
const busy = v.running > 0;
|
||||
const prev = history[history.length - 2];
|
||||
const dgen = prev && prev.gen != null ? entry.gen - prev.gen : 0; // < 0 after a restart
|
||||
const rate = dgen > 0 ? dgen / Math.max(entry.t - prev.t, 0.1) : 0;
|
||||
if (busy && rate) lastSpeed = lastSpeed ? lastSpeed * 0.4 + rate * 0.6 : rate;
|
||||
speeds.push(busy ? rate : 0);
|
||||
if (speeds.length > 60) speeds.shift();
|
||||
$("speed-box").classList.toggle("idle", !busy || !rate);
|
||||
set("speed", lastSpeed ? fmt(lastSpeed, 1) : "–");
|
||||
set("speed-unit", busy && rate ? "tokens/s" : lastSpeed ? "tokens/s, last answer" : "tokens/s");
|
||||
const top = Math.max(...speeds, 1);
|
||||
$("spark").setAttribute("d", speeds.map((y, x) => `${x ? "L" : "M"}${x + 60 - speeds.length} ${27 - (y / top) * 25}`).join(""));
|
||||
const waiting = v.waiting ? `, ${fmt(v.waiting)} waiting` : "";
|
||||
if (busy) setState("live", `${rate ? "Generating" : "Reading the prompt"}${waiting}`);
|
||||
else setState("ok", "Ready");
|
||||
set("pill-text", !busy ? "Ready" : rate ? `${fmt(lastSpeed, 1)} tok/s` : "Working");
|
||||
set("reqs", `${fmt(v.running)} running${waiting}`);
|
||||
set("kv", `${fmt(v.kv_usage * 100)}%`);
|
||||
meter("kv-bar", v.kv_usage);
|
||||
set("prefix", v.prefix_queries ? `${fmt(100 * v.prefix_hits / v.prefix_queries)}%` : "–");
|
||||
}
|
||||
|
||||
if (s.experts && Object.keys(s.experts.layers).length) {
|
||||
$("scan-empty").hidden = true;
|
||||
$("scan-body").hidden = false;
|
||||
drawScan(s.experts);
|
||||
const old = history.find((h) => h.hits !== undefined);
|
||||
const d = old ? {hits: entry.hits - old.hits, misses: entry.misses - old.misses, cpu: entry.cpu - old.cpu} : null;
|
||||
const routed = d && d.hits >= 0 && d.misses >= 0 ? d.hits + d.misses + d.cpu : 0; // < 0 after a restart
|
||||
if (routed > 0) {
|
||||
set("hits", `${fmt(100 * d.hits / routed)}%`);
|
||||
set("copies", `${fmt(d.misses / Math.max(entry.t - old.t, 1), 0)}/s`);
|
||||
} else {
|
||||
const all = entry.hits + entry.misses + entry.cpu;
|
||||
set("hits", all ? `${fmt(100 * entry.hits / all)}% overall` : "–");
|
||||
set("copies", `${fmt(entry.misses)} so far`);
|
||||
}
|
||||
} else {
|
||||
$("scan-empty").hidden = false;
|
||||
$("scan-body").hidden = true;
|
||||
set("empty-live", "");
|
||||
}
|
||||
|
||||
const g = s.gpu;
|
||||
if (g) {
|
||||
set("vram", `${fmt(g.vram_used / GIB, 1)} / ${fmt(g.vram_total / GIB, 1)} GiB`);
|
||||
meter("vram-bar", g.vram_used / g.vram_total);
|
||||
set("load", g.busy == null ? "–" : `${g.busy}%`);
|
||||
set("temp", g.temp == null ? "–" : `${fmt(g.temp)} °C`);
|
||||
set("power", g.power == null ? "–" : `${fmt(g.power)} W`);
|
||||
set("fan", g.fan == null ? "–" : `${fmt(g.fan)} rpm`);
|
||||
}
|
||||
}
|
||||
|
||||
let timer = 0;
|
||||
function schedule() {
|
||||
clearTimeout(timer);
|
||||
if (document.visibilityState === "visible") timer = setTimeout(() => poll().finally(schedule), 1500);
|
||||
}
|
||||
document.addEventListener("visibilitychange", () => { if (document.visibilityState === "visible") poll().finally(schedule); });
|
||||
|
||||
render();
|
||||
autosize();
|
||||
poll().finally(schedule);
|
||||
if (!coarse) draft.focus();
|
||||
</script>
|
||||
+84
-113
@@ -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 = """<!doctype html>
|
||||
<meta charset="utf-8">
|
||||
<title>prain</title>
|
||||
<style>
|
||||
body { font: 15px/1.45 system-ui, sans-serif; max-width: 52rem; margin: 2rem auto; padding: 0 1rem; }
|
||||
.user { color: #666; margin-top: 1rem; }
|
||||
.assistant { white-space: pre-wrap; }
|
||||
.draft { color: #999; }
|
||||
textarea { width: 100%; height: 5rem; font: inherit; margin-top: 1rem; }
|
||||
#status { color: #888; font-size: 13px; }
|
||||
</style>
|
||||
<div id="log"></div>
|
||||
<textarea id="draft" autofocus
|
||||
placeholder="Experts load while you type. Enter sends, Shift+Enter adds a line."></textarea>
|
||||
<div id="status"></div>
|
||||
<script>
|
||||
const session = Math.random().toString(36).slice(2) + Date.now().toString(36);
|
||||
const messages = [];
|
||||
const log = document.getElementById("log");
|
||||
const draft = document.getElementById("draft");
|
||||
const status = document.getElementById("status");
|
||||
let streaming = false;
|
||||
PAGE = (Path(__file__).parent / "web.html").read_bytes()
|
||||
METRICS = {"vllm:num_requests_running": "running", "vllm:num_requests_waiting": "waiting",
|
||||
"vllm:kv_cache_usage_perc": "kv_usage", "vllm:prefix_cache_queries_total":
|
||||
"prefix_queries", "vllm:prefix_cache_hits_total": "prefix_hits",
|
||||
"vllm:prompt_tokens_total": "prompt_tokens",
|
||||
"vllm:generation_tokens_total": "generation_tokens"}
|
||||
|
||||
function show(role, text) {
|
||||
const div = document.createElement("div");
|
||||
div.className = role;
|
||||
div.textContent = text;
|
||||
log.appendChild(div);
|
||||
return div;
|
||||
}
|
||||
|
||||
draft.addEventListener("input", () => {
|
||||
if (!streaming && draft.value.trim())
|
||||
fetch("/warm", {method: "POST", body: JSON.stringify({session, messages, draft: draft.value})});
|
||||
});
|
||||
def vllm_metrics(base_url):
|
||||
"""The METRICS of vLLM's Prometheus text, summed over label sets; None when unreachable."""
|
||||
try:
|
||||
with urllib.request.urlopen(f"{base_url}/metrics", timeout=2) as response:
|
||||
text = response.read().decode()
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
out = dict.fromkeys(METRICS.values(), 0.0)
|
||||
for line in text.splitlines():
|
||||
key = METRICS.get(line.split("{", 1)[0].split(" ", 1)[0])
|
||||
if key and not line.startswith("#"):
|
||||
out[key] += float(line.rsplit(None, 1)[1])
|
||||
return out
|
||||
|
||||
|
||||
def _number(path):
|
||||
try:
|
||||
with open(path) as f:
|
||||
return int(f.read().strip())
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
|
||||
|
||||
def gpu_stats(sysfs="/sys/class/drm"):
|
||||
"""VRAM, load, temperature, power and fan of the GPU with the most VRAM (amdgpu sysfs);
|
||||
None without one."""
|
||||
cards = [(_number(os.path.join(d, "mem_info_vram_total")), d)
|
||||
for d in glob.glob(os.path.join(sysfs, "card*", "device"))]
|
||||
cards = sorted((total, d) for total, d in cards if total)
|
||||
if not cards:
|
||||
return None
|
||||
total, card = cards[-1]
|
||||
hwmon = sorted(glob.glob(os.path.join(card, "hwmon", "hwmon*")))
|
||||
|
||||
def sensor(name, scale):
|
||||
value = _number(os.path.join(hwmon[0], name)) if hwmon else None
|
||||
return None if value is None else value / scale
|
||||
|
||||
return {"vram_used": _number(os.path.join(card, "mem_info_vram_used")), "vram_total": total,
|
||||
"busy": _number(os.path.join(card, "gpu_busy_percent")),
|
||||
"temp": sensor("temp1_input", 1000), "power": sensor("power1_average", 1e6),
|
||||
"fan": sensor("fan1_input", 1)}
|
||||
|
||||
|
||||
def expert_stats(path):
|
||||
"""prain's stats file (prain.policy.write_stats), or None."""
|
||||
if not path:
|
||||
return None
|
||||
try:
|
||||
with open(path) as f:
|
||||
return json.load(f)
|
||||
except (OSError, ValueError):
|
||||
return None
|
||||
|
||||
draft.addEventListener("keydown", async (event) => {
|
||||
if (event.key !== "Enter" || event.shiftKey || streaming || !draft.value.trim()) return;
|
||||
event.preventDefault();
|
||||
const text = draft.value.trim();
|
||||
messages.push({role: "user", content: text});
|
||||
show("user", "you> " + text);
|
||||
draft.value = "";
|
||||
streaming = true;
|
||||
const out = show("assistant", "");
|
||||
const finalSpan = out.appendChild(document.createElement("span"));
|
||||
const draftSpan = out.appendChild(document.createElement("span"));
|
||||
draftSpan.className = "draft";
|
||||
let mode = "full", draftText = "", finalText = "", refined = false;
|
||||
const render = () => { // the full text, then the grey rest of the draft while they agree
|
||||
finalSpan.textContent = finalText;
|
||||
draftSpan.textContent = mode === "draft" ? draftText
|
||||
: mode === "refine" && draftText.startsWith(finalText) ? draftText.slice(finalText.length) : "";
|
||||
};
|
||||
let start = performance.now(), first = 0, pieces = 0, buffer = "";
|
||||
const response = await fetch("/chat", {method: "POST", body: JSON.stringify({session, messages})});
|
||||
if (!response.ok) {
|
||||
out.textContent = "error " + response.status + ": " + await response.text();
|
||||
messages.pop();
|
||||
streaming = false;
|
||||
return;
|
||||
}
|
||||
const reader = response.body.getReader();
|
||||
const decoder = new TextDecoder();
|
||||
for (;;) {
|
||||
const {done, value} = await reader.read();
|
||||
if (done) break;
|
||||
buffer += decoder.decode(value, {stream: true});
|
||||
const lines = buffer.split("\\n");
|
||||
buffer = lines.pop();
|
||||
for (const line of lines) {
|
||||
if (!line.startsWith("data: ") || line === "data: [DONE]") continue;
|
||||
const data = JSON.parse(line.slice(6));
|
||||
if (data.prain === "draft") {
|
||||
mode = "draft";
|
||||
status.textContent = "GPU busy: drafting from the experts in VRAM, refined when it is free";
|
||||
} else if (data.prain === "refine") {
|
||||
mode = "refine";
|
||||
refined = true;
|
||||
start = performance.now(), first = 0, pieces = 0;
|
||||
status.textContent = "refining with all experts";
|
||||
} else if (data.prain === "error") {
|
||||
mode = "draft";
|
||||
refined = false;
|
||||
status.textContent = "refining failed (HTTP " + data.status + "), keeping the draft";
|
||||
} else {
|
||||
const piece = data.choices[0].delta.content;
|
||||
if (!piece) continue;
|
||||
if (pieces++ === 0) first = performance.now();
|
||||
if (mode === "draft") draftText += piece; else finalText += piece;
|
||||
}
|
||||
render();
|
||||
}
|
||||
}
|
||||
messages.push({role: "assistant", content: mode === "draft" ? draftText : finalText});
|
||||
let same = 0;
|
||||
while (same < draftText.length && draftText[same] === finalText[same]) same++;
|
||||
const seconds = (performance.now() - first) / 1000;
|
||||
status.textContent = (refined ? `draft kept ${same} of ${draftText.length} chars; refined: `
|
||||
: mode === "draft" ? "draft, not refined: " : "") +
|
||||
`first token ${((first - start) / 1000).toFixed(1)} s, ` +
|
||||
`${pieces > 1 ? ((pieces - 1) / seconds).toFixed(1) : "0.0"} tok/s`;
|
||||
streaming = false;
|
||||
});
|
||||
</script>
|
||||
"""
|
||||
|
||||
|
||||
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()
|
||||
|
||||
@@ -17,3 +17,6 @@ build-backend = "setuptools.build_meta"
|
||||
|
||||
[tool.setuptools]
|
||||
packages = ["prain"]
|
||||
|
||||
[tool.setuptools.package-data]
|
||||
prain = ["web.html"]
|
||||
|
||||
+19
-8
@@ -33,10 +33,13 @@ def post(url, path, body):
|
||||
|
||||
|
||||
def greedy(url, model, ids, tokens):
|
||||
"""(text, token ids, their logprobs) of a greedy answer."""
|
||||
choice = post(url, "/v1/completions", {
|
||||
"model": model, "prompt": ids, "max_tokens": tokens, "temperature": 0, "logprobs": 1,
|
||||
"return_tokens_as_token_ids": True})["choices"][0]
|
||||
return choice["text"], [int(t.split(":")[1]) for t in choice["logprobs"]["tokens"]]
|
||||
logprobs = choice["logprobs"]
|
||||
return (choice["text"], [int(t.split(":")[1]) for t in logprobs["tokens"]],
|
||||
logprobs["token_logprobs"])
|
||||
|
||||
|
||||
def nll(url, model, ids, answer):
|
||||
@@ -56,6 +59,9 @@ def main():
|
||||
parser.add_argument("--tokens", type=int, default=48, help="answer tokens")
|
||||
parser.add_argument("--code", default="prain/policy.py", help="file for the code question")
|
||||
parser.add_argument("--code-chars", type=int, default=900)
|
||||
parser.add_argument("--greedy-only", action="store_true",
|
||||
help="score: skip the nll (prompt logprobs need ~0.5 GiB of free VRAM "
|
||||
"for the logits); compare greedy answers only")
|
||||
args = parser.parse_args()
|
||||
if args.action == "make":
|
||||
with open(args.code) as f:
|
||||
@@ -65,7 +71,7 @@ def main():
|
||||
ids = post(args.url, "/tokenize", {"model": args.model, "add_generation_prompt": True,
|
||||
"messages": [{"role": "user", "content": prompt}]})
|
||||
ids = ids["tokens"]
|
||||
text, answer = greedy(args.url, args.model, ids, args.tokens)
|
||||
text, answer, _ = greedy(args.url, args.model, ids, args.tokens)
|
||||
refs.append({"prompt": prompt, "ids": ids, "answer": answer, "text": text})
|
||||
print(json.dumps({"prompt": prompt[:60], "answer": text[:120]}), flush=True)
|
||||
with open(args.ref, "w") as f:
|
||||
@@ -73,15 +79,20 @@ def main():
|
||||
return
|
||||
with open(args.ref) as f:
|
||||
refs = json.load(f)
|
||||
scores, same = [], []
|
||||
scores, same, own = [], [], []
|
||||
for ref in refs:
|
||||
scores.append(nll(args.url, args.model, ref["ids"], ref["answer"]))
|
||||
text, answer = greedy(args.url, args.model, ref["ids"], len(ref["answer"]))
|
||||
if not args.greedy_only:
|
||||
scores.append(nll(args.url, args.model, ref["ids"], ref["answer"]))
|
||||
text, answer, logprobs = greedy(args.url, args.model, ref["ids"], len(ref["answer"]))
|
||||
own.append(sum(logprobs) / len(logprobs))
|
||||
same.append(next((i for i, (a, b) in enumerate(zip(answer, ref["answer"])) if a != b),
|
||||
min(len(answer), len(ref["answer"]))))
|
||||
print(json.dumps({"prompt": ref["prompt"][:40], "nll": round(scores[-1], 3),
|
||||
"same_tokens": same[-1], "answer": text[:100]}), flush=True)
|
||||
print(json.dumps({"mean_nll": round(sum(scores) / len(scores), 3), "same_tokens": same,
|
||||
print(json.dumps({"prompt": ref["prompt"][:40],
|
||||
"nll": round(scores[-1], 3) if scores else None,
|
||||
"own_lp": round(own[-1], 3), "same_tokens": same[-1],
|
||||
"answer": text[:100]}), flush=True)
|
||||
print(json.dumps({"mean_nll": round(sum(scores) / len(scores), 3) if scores else None,
|
||||
"mean_own_lp": round(sum(own) / len(own), 3), "same_tokens": same,
|
||||
"identical": sum(s == len(r["answer"]) for s, r in zip(same, refs))}))
|
||||
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
import importlib.util
|
||||
import unittest
|
||||
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
torch = None
|
||||
|
||||
GPU = torch is not None and torch.cuda.is_available() and importlib.util.find_spec("vllm") is not None
|
||||
|
||||
|
||||
@unittest.skipUnless(GPU, "needs vLLM and a GPU")
|
||||
class Kernels(unittest.TestCase):
|
||||
def setUp(self):
|
||||
torch.manual_seed(0)
|
||||
|
||||
def test_unpack_writes_what_to_gpu_int4_computes(self):
|
||||
from prain.kernels import unpack_int4
|
||||
from prain.quant import to_gpu_int4
|
||||
|
||||
for n, half in ((256, 128), (1408, 1408), (2816, 352)):
|
||||
cpu_layout = torch.randint(0, 256, (n, half), dtype=torch.uint8, device="cuda")
|
||||
dst = torch.zeros_like(cpu_layout)
|
||||
unpack_int4(cpu_layout.flatten(), dst)
|
||||
self.assertTrue(torch.equal(dst, to_gpu_int4(cpu_layout)), (n, half))
|
||||
|
||||
def test_decode_kernel_computes_what_fused_experts_computes(self):
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
|
||||
|
||||
from prain.kernels import moe_int4_decode
|
||||
from prain.quant import dequantize, moe_quant_config, quantize
|
||||
from prain.tiered import cpu_experts, slot_table
|
||||
|
||||
hidden, inter, experts, topk = 256, 128, 6, 4
|
||||
w13 = torch.randn(experts, 2 * inter, hidden) * 0.1
|
||||
w2 = torch.randn(experts, hidden, inter) * 0.1
|
||||
parts = [quantize(w, "int4") for w in (w13, w2)] # gpu layout
|
||||
held = [1, 3, 4] # the round's experts, in slots 0..2
|
||||
slots = [[q[held].cuda(), s[held].cuda()] for q, s in parts]
|
||||
table = slot_table({e: i for i, e in enumerate(held)}, experts, "cuda")
|
||||
for tokens in (1, 2, 3, 4):
|
||||
x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
|
||||
ids = torch.stack([torch.randperm(experts)[:topk] for _ in range(tokens)]).int()
|
||||
ids[0, -1] = experts # padding route: the table's last row
|
||||
w = torch.rand(tokens, topk)
|
||||
got = moe_int4_decode(x.cuda(), w.cuda(), ids.cuda(), table, *slots[0], *slots[1])
|
||||
with set_current_vllm_config(VllmConfig()):
|
||||
vllm = fused_experts(x.cuda(), slots[0][0], slots[1][0], w.cuda(), ids.cuda(),
|
||||
activation=MoEActivation.GELU_TANH,
|
||||
global_num_experts=table.numel(), expert_map=table,
|
||||
quant_config=moe_quant_config("int4", *[s[1] for s in slots]))
|
||||
exact = cpu_experts(x, w, torch.where(ids < experts, ids, -1), set(held), lambda e: tuple(
|
||||
dequantize(q, s, "int4")[e] for q, s in parts))
|
||||
self.assertEqual((got.shape, got.dtype), ((tokens, hidden), torch.bfloat16))
|
||||
torch.testing.assert_close(got.float().cpu(), exact, rtol=3e-2, atol=3e-2)
|
||||
torch.testing.assert_close(got.float(), vllm.float(), rtol=3e-2, atol=3e-2)
|
||||
|
||||
def test_staged_copy_fills_a_slot_like_the_repack_path(self):
|
||||
from prain.experts import _Staging, stage_int4
|
||||
from prain.quant import quantize, to_cpu_int4, to_gpu_int4
|
||||
|
||||
w13, w2 = torch.randn(2, 256, 256), torch.randn(2, 256, 128)
|
||||
pack = [[to_cpu_int4(q), s] for q, s in (quantize(w, "int4") for w in (w13, w2))]
|
||||
slots = [[torch.zeros((3, *t.shape[1:]), dtype=t.dtype, device="cuda") for t in mats]
|
||||
for mats in pack]
|
||||
ring = _Staging(sum(t[0].nbytes for mats in pack for t in mats), "cuda")
|
||||
stream = torch.cuda.current_stream()
|
||||
for expert, slot in ((1, 2), (0, 0), (1, 1), (0, 2), (1, 0)): # wraps the ring
|
||||
stage_int4([[t[expert] for t in mats] for mats in pack], slots, slot, ring, stream)
|
||||
torch.cuda.synchronize()
|
||||
for mats, dsts in zip(pack, slots):
|
||||
self.assertTrue(torch.equal(dsts[0][slot], to_gpu_int4(mats[0][expert].cuda())))
|
||||
self.assertTrue(torch.equal(dsts[1][slot], mats[1][expert].cuda()))
|
||||
|
||||
|
||||
def test_int4_rounds_on_rdna2_match_the_bf16_kernel_even_when_fp16_overflows(self):
|
||||
from types import SimpleNamespace
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
|
||||
|
||||
from prain.experts import TieredRoutedExperts
|
||||
from prain.quant import moe_quant_config, quantize
|
||||
from prain.tiered import slot_table
|
||||
|
||||
hidden, inter, experts = 256, 128, 4
|
||||
parts = [quantize(w, "int4") for w in (torch.randn(experts, 2 * inter, hidden) * 0.1,
|
||||
torch.randn(experts, hidden, inter) * 0.1)]
|
||||
slots = [[q.cuda(), s.cuda()] for q, s in parts]
|
||||
config = moe_quant_config("int4", slots[0][1], slots[1][1])
|
||||
layer = SimpleNamespace(_slots=slots, activation=MoEActivation.GELU_TANH, _qconfig=config)
|
||||
table = slot_table({e: e for e in range(experts)}, experts, "cuda")
|
||||
with set_current_vllm_config(VllmConfig()):
|
||||
for tokens, scale in ((2, 1.0), (64, 1.0), (64, 3000.0)): # decode kernel, fp16, overflow
|
||||
x = (torch.randn(tokens, hidden) * scale).bfloat16().cuda()
|
||||
ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int().cuda()
|
||||
w = torch.rand(tokens, 2).cuda()
|
||||
got = TieredRoutedExperts._int4_gfx10(layer, x, w, ids, table)
|
||||
want = fused_experts(x, slots[0][0], slots[1][0], w, ids, activation=layer.activation,
|
||||
global_num_experts=table.numel(), expert_map=table,
|
||||
quant_config=config)
|
||||
self.assertTrue(torch.isfinite(got).all(), (tokens, scale))
|
||||
torch.testing.assert_close(got.float(), want.float(), rtol=3e-2, atol=3e-2 * scale,
|
||||
msg=f"{tokens} tokens, x * {scale}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -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
|
||||
|
||||
+61
-5
@@ -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'<textarea id="draft"', response.read())
|
||||
page = response.read()
|
||||
self.assertIn(b'<textarea id="draft"', page)
|
||||
self.assertIn(b'"/status"', page) # the monitor polls it
|
||||
self.warm("s1", "Why is ")
|
||||
wait_for(lambda: self.stub.bodies)
|
||||
self.assertEqual(self.stub.bodies[0]["messages"], [{"role": "user", "content": "Why is"}])
|
||||
@@ -67,9 +72,9 @@ class Web(unittest.TestCase):
|
||||
|
||||
def test_no_warm_up_while_another_answer_streams(self):
|
||||
self.server.streams[0] = 1 # someone else's answer is streaming
|
||||
self.warm("bob", "typing along")
|
||||
self.assertEqual(self.warm("bob", "typing along"), 204) # 204: nothing warms up
|
||||
self.server.streams[0] = 0
|
||||
self.warm("bob", "now")
|
||||
self.assertEqual(self.warm("bob", "now"), 202) # 202: a warm-up starts
|
||||
wait_for(lambda: self.stub.bodies)
|
||||
self.assertEqual([b["messages"][-1]["content"] for b in self.stub.bodies], ["now"])
|
||||
|
||||
@@ -136,3 +141,54 @@ class Web(unittest.TestCase):
|
||||
self.assertEqual(events(r), ["Hel", "lo", "!", "<done>"])
|
||||
thread.join()
|
||||
self.assertEqual(sum("vllm_xargs" in b for b in self.stub.bodies), 1)
|
||||
|
||||
|
||||
class Status(TmpCase):
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.stub = StubVllm()
|
||||
self.addCleanup(self.stub.close)
|
||||
|
||||
def status(self, **options):
|
||||
server = make_server(self.stub.url, 0, **options)
|
||||
threading.Thread(target=server.serve_forever, daemon=True).start()
|
||||
self.addCleanup(server.server_close)
|
||||
self.addCleanup(server.shutdown)
|
||||
with urllib.request.urlopen(f"http://127.0.0.1:{server.server_port}/status") as r:
|
||||
return json.load(r)
|
||||
|
||||
def test_status_reports_vllm_the_gpu_and_the_experts(self):
|
||||
self.stub.metrics = (
|
||||
"# HELP vllm:num_requests_running Number of requests running.\n"
|
||||
'vllm:num_requests_running{engine="0",model_name="m"} 1.0\n'
|
||||
'vllm:num_requests_waiting{engine="0",model_name="m"} 0.0\n'
|
||||
'vllm:num_requests_waiting_by_reason{engine="0",reason="capacity"} 5.0\n'
|
||||
'vllm:kv_cache_usage_perc{engine="0",model_name="m"} 0.25\n'
|
||||
'vllm:prefix_cache_queries_total{engine="0",model_name="m"} 40.0\n'
|
||||
'vllm:prefix_cache_hits_total{engine="0",model_name="m"} 30.0\n'
|
||||
'vllm:prompt_tokens_total{engine="0",model_name="m"} 100.0\n'
|
||||
'vllm:generation_tokens_total{engine="0",model_name="m"} 7.0\n')
|
||||
os.makedirs(os.path.join(self.dir, "card0", "device")) # an iGPU: no VRAM files
|
||||
gpu = os.path.join(self.dir, "card1", "device")
|
||||
os.makedirs(os.path.join(gpu, "hwmon", "hwmon4"))
|
||||
for name, value in {"mem_info_vram_total": 8573157376, "mem_info_vram_used": 8000000000,
|
||||
"gpu_busy_percent": 96, "hwmon/hwmon4/temp1_input": 61000,
|
||||
"hwmon/hwmon4/power1_average": 120500000,
|
||||
"hwmon/hwmon4/fan1_input": 1800}.items():
|
||||
self.write(os.path.join("card1", "device", name), f"{value}\n")
|
||||
stats = {"layers": {"0": {"uses": [3, 0], "vram": [0], "hits": 2, "misses": 1}}}
|
||||
stats_path = self.write("stats.json", json.dumps(stats))
|
||||
got = self.status(stats_path=stats_path, sysfs=self.dir)
|
||||
self.assertEqual(got["model"], "stub-model")
|
||||
self.assertEqual(got["vllm"], {"running": 1.0, "waiting": 0.0, "kv_usage": 0.25,
|
||||
"prefix_queries": 40.0, "prefix_hits": 30.0,
|
||||
"prompt_tokens": 100.0, "generation_tokens": 7.0})
|
||||
self.assertEqual(got["gpu"], {"vram_used": 8000000000, "vram_total": 8573157376,
|
||||
"busy": 96, "temp": 61.0, "power": 120.5, "fan": 1800.0})
|
||||
self.assertEqual(got["experts"], stats)
|
||||
self.assertEqual(got["streams"], 0)
|
||||
|
||||
def test_status_without_metrics_gpu_or_stats(self):
|
||||
got = self.status(sysfs=self.dir) # the stub answers /metrics with JSON, not metrics
|
||||
self.assertEqual((got["gpu"], got["experts"]), (None, None))
|
||||
self.assertEqual(got["vllm"]["running"], 0.0)
|
||||
|
||||
@@ -0,0 +1,49 @@
|
||||
import json
|
||||
import shutil
|
||||
import subprocess
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
PAGE = (Path(__file__).parent.parent / "prain" / "web.html").read_text()
|
||||
|
||||
|
||||
@unittest.skipUnless(shutil.which("node"), "needs node")
|
||||
class Markdown(unittest.TestCase):
|
||||
"""The page's md() (model output -> innerHTML) run in node."""
|
||||
|
||||
def md(self, text):
|
||||
script = PAGE[PAGE.index("<script>") + 8:PAGE.index("/* ---------- thread")]
|
||||
script = script.split("/* ---------- markdown", 1)[1].split("*/", 1)[1]
|
||||
out = subprocess.run(["node", "-e", script + f"\nprocess.stdout.write(md({json.dumps(text)}))"],
|
||||
capture_output=True, text=True, check=True)
|
||||
return out.stdout
|
||||
|
||||
def test_html_in_model_output_is_text(self):
|
||||
html = self.md('<img src=x onerror="alert(1)"> **<b>hi</b>** `<i>` [x](javascript:alert(1))')
|
||||
self.assertNotIn("<img", html)
|
||||
self.assertNotIn("<b>", html)
|
||||
self.assertNotIn('href="javascript', html)
|
||||
self.assertIn("<img src=x onerror="alert(1)">", html)
|
||||
self.assertIn("<strong><b>hi</b></strong>", html)
|
||||
|
||||
def test_blocks(self):
|
||||
html = self.md("# Title\n\nSome *very* **bold** text with `code` and "
|
||||
"[a link](https://x.org/?a=1&b=2).\n\n```python\nif a < b:\n pass\n```\n"
|
||||
"3. three\n4. four\n\n- one\n more\n- two\n\n| a | b |\n|---|---|\n| 1 | 2 |\n"
|
||||
"\n> quoted\n\n---\nend")
|
||||
for part in ("<h3>Title</h3>", "<em>very</em>", "<strong>bold</strong>", "<code>code</code>",
|
||||
'<a href="https://x.org/?a=1&b=2" target="_blank" rel="noopener noreferrer">',
|
||||
'<span>python</span>', "<pre><code>if a < b:\n pass</code></pre>",
|
||||
'<ol start="3"><li>three</li><li>four</li></ol>',
|
||||
"<ul><li>one more</li><li>two</li></ul>",
|
||||
"<th>a</th><th>b</th></tr></thead><tbody><tr><td>1</td><td>2</td>",
|
||||
"<blockquote><p>quoted</p></blockquote>", "<hr>", "<p>end</p>"):
|
||||
self.assertIn(part, html)
|
||||
|
||||
def test_an_open_fence_while_streaming_is_code(self):
|
||||
self.assertIn("<pre><code>x = 1\n# not a heading</code></pre>",
|
||||
self.md("Here:\n```\nx = 1\n# not a heading"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user