prain: README: expert formats, dense-layer fix, min_weight trade-off, 16k agent setup; example config int4, 16k, tools

README.md

@@ -3,8 +3,9 @@
 Serve a mixture-of-experts model with vLLM when its experts do not fit in VRAM. A TOML file places
 each expert of each layer in VRAM (pinned or in a cache slot), in RAM, or in a pack file on disk.
 Missing experts are copied to the GPU or computed on the CPU (`[miss]`). Eviction, heat, prefetch
-and miss policies are configurable, and `prain tune` picks them from a routing trace.
+and miss policies are configurable, and `prain tune` picks them from a routing trace. Experts can be
-Target: Gemma 4 26B-A4B (30 layers x 128 experts, 42.5 GiB of experts) on an 8 GB RX 6600 XT.
+served as bf16, int8 or int4 (`quant`).
+Target: Gemma 4 26B-A4B (30 layers x 128 experts, 42.5 GiB of bf16 experts) on an 8 GB RX 6600 XT.
 
 ## Build (GPU host, user space)
 
@@ -17,25 +18,110 @@ Target: Gemma 4 26B-A4B (30 layers x 128 experts, 42.5 GiB of experts) on an 8 G
 
 Check the stack: `scripts/smoke.sh MODEL_DIR --attention-backend TRITON_ATTN`
 
+Tune vLLM's fused-MoE kernel for the GPU once (~45 min of Triton compiles the first time;
+`scripts/env.sh` points vLLM at the files). It keeps vLLM's default wherever no candidate beats it:
+
+    python scripts/tune_moe.py build/moe-configs
+
 ## Serve
 
     . scripts/env.sh
     prain serve examples/gemma4-26b-a4b-rx6600xt.toml MODEL_DIR -- --served-model-name gemma4 \
-        --language-model-only --max-model-len 4096 --max-num-batched-tokens 1024 \
+        --language-model-only --max-model-len 16384 --max-num-batched-tokens 1024 \
-        --enable-prefix-caching --attention-backend TRITON_ATTN
+        --enable-prefix-caching --attention-backend TRITON_ATTN \
+        --enable-auto-tool-choice --tool-call-parser gemma4
 
-The first start writes the pack (~3 min). Chat with `prain tui`, or `prain web --config FILE` for
+The first start writes the bf16 pack (42.5 GiB, ~3 min); with `quant = "int8"` or `"int4"` it
-several users with drafts.
+then quantizes that pack once on the GPU (~3 min). Once the int4 pack is written, the bf16 files
+(`layerNNN.w13`, `.w2`, `.done`) are not read and can be deleted; a later bf16 or int8 start
+writes them again. Chat with `prain tui`, or `prain web --config FILE` for several users with
+drafts. Tool calls work through vLLM's `gemma4` parser, so coding agents and MCP clients can use
+the OpenAI-compatible API.
+
+## Expert formats
+
+| quant | expert | all experts | weights | VRAM kernel | CPU kernel |
+|---|---|---|---|---|---|
+| bf16 | 11.34 MiB | 42.5 GiB | as trained | vLLM Triton | torch bf16 matmul |
+| int8 | 5.69 MiB | 21.3 GiB | fp32 scale per row | vLLM Triton (w8a16) | `torch._weight_int8pack_mm` |
+| int4 | 3.01 MiB | 11.3 GiB | bf16 scale per 64, MSE clip | vLLM Triton (GPTQ/AWQ, w4a16) | `torch._weight_int4pack_mm_for_cpu` |
+
+Measured on the RX 6600 XT behind Thunderbolt (PCIe 3.0 x4, 2.5 GB/s), i9-13950HX (16 CPU
+threads for experts), 31 GiB RAM, quiet host. Every format got 1.66 GiB of expert VRAM, its own
+placement from `prain tune` and tuned kernels; 4k context, 0.75 GiB KV:
+
+| | bf16 | int8 | int4 |
+|---|---|---|---|
+| experts in VRAM (cache slots) | 120 | 270 | 540 |
+| decode, tok/s (median of 3 x 48 tokens) | 1.04 | 1.42 | 1.62 |
+| first token, short chat prompt (median) | 22.9 s | 5.0 s | 4.9 s |
+| 3 greedy answers x 64 tokens | 221 s | 143 s | 123 s |
+| greedy answers identical to bf16 | - | 1 of 3 | 1 of 3 |
+
+The smaller formats win twice: more experts fit in VRAM, and every miss moves fewer bytes over
+the link and out of the page cache (the int4 pack fits in free RAM, the bf16 one does not). On
+the GPU itself a decode round costs about the same in every format (2.3-2.8 ms). These decode
+numbers predate the dense-layer fix below, which makes every format faster.
+
+## Dense layers on RDNA2
+
+rocBLAS runs bf16 GEMVs at 13-16 GB/s on gfx1032 (RDNA2 has no bf16 dot instructions), and vLLM
+uses its own skinny-GEMM kernels only on gfx9/gfx11/gfx12. A decode step spent ~350 ms just
+reading the dense weights (lm_head alone: 93 ms). On gfx10 prain sends every unquantized GEMM
+with at most 4 rows through a Triton GEMV (`prain/gemv.py`, ~210 GB/s; `PRAIN_GEMV=0` turns it
+off). With the agent setup below, decode went from 1.58 to 2.77 tok/s with unchanged answers
+(same quality score). Prefill keeps rocBLAS.
+
+## Trading quality for speed
+
+`[miss] min_weight` (default 1) is the share of each token's router weight that must be
+computed. Below 1, a token skips its lightest experts that are not in VRAM once it has that
+share, and its other routes are scaled up to its full weight: fewer copies, a slightly different
+model. Measured on the agent setup: int4, 16k context, 1 GiB KV and 450 cache slots (1.32 GiB) as
+in the example config, per-layer slot counts from `prain tune`. Quality is the mean
+negative log-likelihood of the min_weight = 1 answers to six chat prompts (`scripts/quality.py`;
+0.11 is that config scoring its own answers):
+
+| min_weight | answer NLL | decode tok/s | first token, short prompt | first token, 2.5k-token prompt |
+|---|---|---|---|---|
+| 1.0 | 0.11 | 2.77 | 4.6 s | 28.5 s |
+| 0.6 | 0.19 | 3.15 | 2.7 s | 23.0 s |
+| 0.4 | 0.40 | 3.36 | 2.0 s | 19.9 s |
+
+At 0.6 the answers stay correct but take other words; at 0.4 they show slips ("Hereon any of the
+following one-liners"). Pinning only the hottest experts and never copying (`prain plan
+--cache-slots 1`, min_weight 0) broke the model (NLL 9) for 25% more decode speed (measured before
+the dense-layer fix).
+
+## Context length and coding agents
+
+Gemma 4 has 25 sliding-window layers (window 1024) and 5 full-attention layers, so KV grows slowly
+with context: per request ~20 KiB per token plus ~0.6 GiB for the sliding layers, which also
+hold the tokens in flight (2 x `--max-num-batched-tokens`). `kv_cache_gib = 1.0` holds 18,169
+tokens, 1.11 requests of 16k (measured).
+
+With the example config (int4, 16k, tools) a cold 4.7k-token prompt took 70 s and a 12k-token
+one 176 s: every 1024-token prefill chunk needs nearly all experts, copied over the 2.5 GB/s link.
+Agent turns reuse the prefix cache: turns 2 and 3 of a conversation found 4,704 and 4,736 of
+~4,750 prompt tokens cached and took ~20 s each (24-token answers). vLLM keeps sliding-window
+state only where earlier prompts ended, so the second conversation that shares just the system
+prompt recomputes it once (176 s); later ones hit (4.5 s).
 
 ## Tune for a machine
 
-    python scripts/measure.py PACK_DIR/layer029.w13 --threads 16 > cost.json   # server stopped
+    python scripts/measure.py PACK_DIR/layer029.w13 --quant int4 --threads 16 > cost.json  # server stopped
     prain sim TRACE CONFIG MODEL_DIR/config.json --cost cost.json
-    prain tune TRACE MODEL_DIR/config.json --cost cost.json --expert-vram-gib 1.67 --ram-gib 0 \
+    prain tune TRACE MODEL_DIR/config.json --cost cost.json --quant int4 --expert-vram-gib 1.4 \
-        --pack-dir PACK_DIR --cpu-threads 16 > tuned.toml
+        --ram-gib 0 --pack-dir PACK_DIR --cpu-threads 16 > tuned.toml
-    python scripts/bench.py http://127.0.0.1:8000
+    python scripts/bench.py http://127.0.0.1:8000 [--context FILES...]          # speed
+    python scripts/quality.py http://127.0.0.1:8000 make ref.json   # reference answers, then
+    python scripts/quality.py http://127.0.0.1:8000 score ref.json  # on another config
 
-The simulator treats every disk-tier expert as a disk read; the page cache makes real runs faster.
+Judge quality on chat answers, not raw-text perplexity: this instruction-tuned model scores raw
+code at a negative log-likelihood of 5-9 per token on a correct engine.
+
+The simulator treats every disk-tier expert as a disk read, so it under-predicts bf16 (0.6 vs 1.04
+tok/s measured, the page cache helps) and over-predicted int4 (2.4 vs 1.62 tok/s).
 
 ## vLLM patches
 

examples/gemma4-26b-a4b-rx6600xt.toml

@@ -1,9 +1,12 @@
-# prain placement: prithivMLmods/gemma-4-26B-A4B-it-Uncensored-MAX on an 8 GB RX 6600 XT
+# prain placement: prithivMLmods/gemma-4-26B-A4B-it-Uncensored-MAX on an 8 GB RX 6600 XT, set up
-# Model: 30 MoE layers x 128 experts, top-8 routing, 11.34 MiB per expert (BF16).
+# for a coding agent: int4 experts, 16k context, tool calls.
-# VRAM, measured: 6.22 GiB weights (4.46 dense + 5 slots x 30 layers) + 0.75 GiB KV cache: 7.9 of
+# Model: 30 MoE layers x 128 experts, top-8 routing, 3.01 MiB per int4 expert (11.34 MiB bf16).
-# 8 GiB in use. Serve with --attention-backend TRITON_ATTN (ROCM_ATTN needs more LDS than RDNA2 has).
+# VRAM: 4.46 GiB dense weights + 30 x 15 x 3.01 MiB = 1.32 GiB of expert slots + 1.0 GiB KV cache.
-# RAM tier: pinning is very slow behind a strict IOMMU (Thunderbolt eGPU: minutes per GiB), so this
+# Serve with --max-model-len 16384 --attention-backend TRITON_ATTN (ROCM_ATTN needs more LDS than
-# file keeps no RAM tier; the kernel page cache holds the hot part of the pack instead.
+# RDNA2 has) --enable-auto-tool-choice --tool-call-parser gemma4.
+# RAM tier: pinning host memory took minutes per GiB on the Thunderbolt eGPU host (cause unproven)
+# and pinned copies were only 4% faster than pageable ones, so this file keeps no RAM tier; the
+# kernel page cache holds the pack instead (the int4 pack, 11.3 GiB, fits in free RAM).
 #
 # Tiers:  vram = pinned in a VRAM slot for the whole run
 #         ram  = pinned host memory, copied into a cache slot when a token needs it
@@ -13,19 +16,21 @@
 # Layer selectors: "*", "12", "0-9,20".   Expert lists: "*", "0-15,32", [1, 2, 3].
 # Check what a file costs:   prain check <this file> <model config.json>
 # Compare variants offline:  prain sim <trace> <this file> <model config.json>
-#                            prain tune <trace> <model config.json> --expert-vram-gib 1.7 ...
+#                            prain tune <trace> <model config.json> --quant int4 --expert-vram-gib 1.4 ...
 # Relative paths start at this file's directory; "~" works too.
 
-pack_dir = "../models/packs/gemma4-26b-a4b"  # 42.5 GiB, written on the first run
+pack_dir = "../models/packs/gemma4-26b-a4b"  # bf16 pack (42.5 GiB) written on the first run,
+quant = "int4"                           # then quantized once into an int4 pack (11.3 GiB):
+                                         # bf16 | int8 (5.69 MiB/expert) | int4 (3.01 MiB)
 default_tier = "disk"                    # runs on any RAM size; `prain tune` does better
-cache_slots = 5                          # 30 x 5 x 11.34 MiB = 1.66 GiB of VRAM
+cache_slots = 15                         # 30 x 15 x 3.01 MiB = 1.32 GiB of VRAM
 stats_path = "../models/packs/gemma4-26b-a4b/stats.json"
 stats_every = 64                         # MoE forward passes between stats/trace writes
 trace_path = "../models/packs/gemma4-26b-a4b/trace.jsonl.gz"  # routing trace for sim/tune
 
 [policy]              # which expert a cache slot gives up
-eviction = "lfu"      # lru | lfu (decayed frequency) | gdsf (frequency x refetch cost)
+eviction = "gdsf"     # lru | lfu (decayed frequency) | gdsf (frequency x refetch cost)
-heat = "weight"       # what makes an expert hot: count (tokens) | weight (router mass) |
+heat = "count"        # what makes an expert hot: count (tokens) | weight (router mass) |
                       # recency (router mass, later tokens weigh more: good for warm-ups)
 lfu_half_life = 64    # forwards until an LFU score halves
 disk_cost = 4.0       # gdsf: a disk expert costs this many RAM experts to refetch
@@ -41,18 +46,24 @@ budget = 2            # experts per layer and forward
 min_prob = 0.05       # skip guesses below this router probability
 
 [miss]                # an expert that is not in VRAM
-mode = "auto"         # copy | cpu (always compute it on the CPU) | auto (CPU if it has
+mode = "copy"         # copy | cpu (always compute it on the CPU) | auto (CPU if it has
-cpu_max_tokens = 2    #   at most cpu_max_tokens tokens, else copy). Measured on the Thunderbolt
+cpu_max_tokens = 1    #   at most cpu_max_tokens tokens, else copy). Measured on the Thunderbolt
-                      #   eGPU host: a copy costs ~4.6 ms + disk read, the CPU ~2.5 ms per token
+                      #   eGPU host, per int4 expert: a copy ~1.3 ms + disk read, the CPU
+                      #   0.8-1.3 ms per token. `prain tune` picked copy for int4 (a copied
+                      #   expert stays for later tokens) and auto <= 2 tokens for bf16 (copy 4.8 ms)
 cpu_threads = 16      # CPU threads of the vLLM worker (sets OMP_NUM_THREADS)
 admit = 0             # cpu/auto: copy this many CPU-run RAM-tier experts into VRAM afterwards
+min_weight = 1.0      # share of each token's router weight that must be computed. Below 1 a token
+                      #   skips experts not in VRAM once it has that share: faster, dumber.
+                      #   Measured here: 0.6 decodes 14% faster (answer NLL 0.11 -> 0.19), 0.4
+                      #   21% (0.40, visible slips)
 
 [serve]               # parallel requests (`prain serve` passes these to vLLM)
 max_requests = 2      # requests vLLM runs at once (--max-num-seqs); warm-ups count too
-kv_cache_gib = 0.75   # KV cache VRAM shared by them (--kv-cache-memory-bytes). vLLM logs
+kv_cache_gib = 1.0    # KV cache VRAM shared by them (--kv-cache-memory-bytes). vLLM logs
-                      #   "Maximum concurrency for 4096 tokens per request: N.NNx" at start;
+                      #   "Maximum concurrency for 16384 tokens per request: 1.11x" with 1.0 GiB.
-                      #   below max_requests, long requests wait or get preempted. Every GiB
+                      #   Gemma 4 needs ~20 KiB per token of context plus ~0.6 GiB for its
-                      #   here is 90 expert cache slots less.
+                      #   sliding-window layers; every GiB here is 340 int4 expert slots less.
 
 [draft]               # answers that arrive while others use the GPU (`prain web --config`)
 full_requests = 1     # answers that run with all their experts; later ones run at once as

prain/cli.py

@@ -7,19 +7,19 @@ import os
 import sys
 from dataclasses import asdict
 
-from prain.config import (GIB, Miss, Policy, Prefetch, expert_bytes, format_ids, load_config,
+from prain.config import (GIB, QUANTS, Miss, Policy, Prefetch, expert_bytes, format_ids, load_config,
                             moe_dims, plan, render_toml)
 
 
-def _model_dims(hf_config_path):
+def _model_dims(hf_config_path, quant="bf16"):
     with open(hf_config_path) as f:
         num_layers, num_experts, top_k, hidden, inter = moe_dims(json.load(f))
-    return num_layers, num_experts, top_k, expert_bytes(hidden, inter)
+    return num_layers, num_experts, top_k, expert_bytes(hidden, inter, quant)
 
 
 def check(config_path, hf_config_path):
     config = load_config(config_path)
-    num_layers, num_experts, top_k, size = _model_dims(hf_config_path)
+    num_layers, num_experts, top_k, size = _model_dims(hf_config_path, config.quant)
     print("layer  vram   ram  disk  cache_slots")
     vram_rows = ram_rows = few_slots = 0
     for layer in range(num_layers):
@@ -53,7 +53,7 @@ def sim(trace, config_path, hf_config_path, cost_path, expert_mib):
     from prain.sim import describe, load_cost, read_forwards, simulate
 
     config = load_config(config_path)
-    _, num_experts, _, size = _model_dims(hf_config_path)
+    _, num_experts, _, size = _model_dims(hf_config_path, config.quant)
     forwards = read_forwards(trace)
     cost = load_cost(cost_path, expert_mib or size / 2**20)
     layers = sorted({layer for fwd in forwards for layer in fwd})
@@ -69,7 +69,7 @@ def tune_command(args):
     from prain.config import Miss, Policy, Prefetch
     from prain.sim import describe, load_cost, read_forwards, tune
 
-    _, num_experts, _, size = _model_dims(args.hf_config)
+    _, num_experts, _, size = _model_dims(args.hf_config, args.quant)
     cost = load_cost(args.cost, args.expert_mib or size / 2**20)
     forwards = read_forwards(args.trace)
     best, ram, per, final, bound, rows = tune(forwards, num_experts, args.expert_vram_gib,
@@ -88,7 +88,7 @@ def tune_command(args):
               for layer in per}
     print(render_toml(
         f"generated by `prain tune` from {os.path.basename(args.trace)}",
-        {"pack_dir": os.path.abspath(args.pack_dir), "default_tier": "disk",
+        {"pack_dir": os.path.abspath(args.pack_dir), "quant": None if args.quant == "bf16" else args.quant, "default_tier": "disk",
          "cache_slots": min(per.values()), "trace_path": os.path.abspath(args.trace)},
         {"policy": asdict(policy), "prefetch": asdict(prefetch), "miss": asdict(miss)},
         layers,
@@ -130,6 +130,7 @@ def main(argv=None):
     p.add_argument("--ram-gib", type=float, required=True, help="pinned host memory for experts")
     p.add_argument("--cache-slots", type=int, default=16)
     p.add_argument("--pack-dir", required=True)
+    p.add_argument("--quant", choices=QUANTS, default="bf16", help="expert weight format")
     for name, text in (("sim", "simulate a placement file on a routing trace"),
                        ("tune", "search policies and per-layer slots on a trace; prints a placement file")):
         s = sub.add_parser(name, help=text)
@@ -143,6 +144,7 @@ def main(argv=None):
             s.add_argument("--expert-vram-gib", type=float, required=True)
             s.add_argument("--ram-gib", type=float, required=True)
             s.add_argument("--pack-dir", required=True)
+            s.add_argument("--quant", choices=QUANTS, default="bf16", help="expert weight format")
             s.add_argument("--cpu-threads", type=int, default=4,
                            help="[miss] cpu_threads to write: the --threads scripts/measure.py used")
             s.add_argument("--top", type=int, default=10, help="variants to list")
@@ -172,11 +174,11 @@ def main(argv=None):
     if args.command == "check":
         check(args.config, args.hf_config)
     elif args.command == "plan":
-        _, num_experts, _, size = _model_dims(args.hf_config)
+        _, num_experts, _, size = _model_dims(args.hf_config, args.quant)
         with open(args.stats) as f:
             stats = json.load(f)
         print(plan(stats, num_experts, size, args.expert_vram_gib, args.ram_gib,
-                   args.cache_slots, os.path.abspath(args.pack_dir), os.path.abspath(args.stats)), end="")
+                   args.cache_slots, os.path.abspath(args.pack_dir), os.path.abspath(args.stats), args.quant), end="")
     elif args.command == "sim":
         sim(args.trace, args.config, args.hf_config, args.cost, args.expert_mib)
     elif args.command == "tune":

prain/config.py

@@ -8,8 +8,10 @@ import tomllib
 from dataclasses import dataclass, field, fields
 
 TIERS = ("vram", "ram", "disk")
+QUANTS = ("bf16", "int8", "int4")  # expert weight formats, see prain.quant
+INT4_GROUP = 64  # int4 weights per scale
 GIB = 1024**3
-_TOP_KEYS = {"pack_dir", "default_tier", "cache_slots", "stats_path", "stats_every",
+_TOP_KEYS = {"pack_dir", "quant", "default_tier", "cache_slots", "stats_path", "stats_every",
              "trace_path", "policy", "prefetch", "miss", "serve", "draft", "layers"}
 _LAYER_KEYS = {"cache_slots", *TIERS}
 
@@ -104,6 +106,8 @@ class Miss:
     cpu_max_tokens: int = _opt(1, minimum=1)  # auto: CPU runs experts with <= this many tokens
     cpu_threads: int = _opt(4, minimum=1)  # OMP threads of the vLLM worker
     admit: int = _opt(1)  # CPU-run experts copied into VRAM afterwards, hottest first
+    min_weight: float = _opt(1.0, maximum=1)  # router weight every token must compute; below 1
+    #                                           it skips experts not in VRAM: faster, dumber
 
 
 @dataclass(frozen=True)
@@ -143,6 +147,7 @@ def _section(cls, raw, where):
 @dataclass(frozen=True)
 class Config:
     pack_dir: str
+    quant: str = "bf16"
     default_tier: str = "ram"
     cache_slots: int = 16
     stats_path: str | None = None
@@ -221,6 +226,7 @@ def load_config(path):
         layers.append((selector, section))
     config = Config(
         pack_dir=_path(path, raw["pack_dir"]),
+        quant=raw.get("quant", "bf16"),
         default_tier=raw.get("default_tier", "ram"),
         cache_slots=raw.get("cache_slots", 16),
         stats_path=_path(path, raw.get("stats_path")),
@@ -236,6 +242,8 @@ def load_config(path):
     if config.draft.full_requests > config.serve.max_requests:
         raise ValueError(f"{path}: [draft] full_requests must be <= [serve] max_requests "
                          f"({config.serve.max_requests})")
+    if config.quant not in QUANTS:
+        raise ValueError(f"{path}: quant must be one of {QUANTS}")
     if config.default_tier not in TIERS:
         raise ValueError(f"{path}: default_tier must be one of {TIERS}")
     _check_count(config.cache_slots, 0, f"{path}: cache_slots")
@@ -262,9 +270,14 @@ def moe_dims(hf_config):
     )
 
 
-def expert_bytes(hidden, moe_intermediate, dtype_bytes=2):
+def expert_bytes(hidden, moe_intermediate, quant="bf16"):
-    """One gated expert: gate+up [2I, H] plus down [H, I]."""
+    """One gated expert, gate+up [2I, H] plus down [H, I], stored as `quant`."""
-    return 3 * hidden * moe_intermediate * dtype_bytes
+    weights = 3 * hidden * moe_intermediate
+    if quant == "bf16":
+        return 2 * weights
+    if quant == "int8":  # plus one fp32 scale per output row
+        return weights + 4 * (2 * moe_intermediate + hidden)
+    return weights // 2 + 2 * weights // INT4_GROUP  # plus one bf16 scale per group
 
 
 def render_toml(header, top, tables=None, layers=None):
@@ -278,7 +291,8 @@ def render_toml(header, top, tables=None, layers=None):
     return "\n".join(lines) + "\n"
 
 
-def plan(stats, num_experts, expert_size, expert_vram_gib, ram_gib, cache_slots, pack_dir, stats_path):
+def plan(stats, num_experts, expert_size, expert_vram_gib, ram_gib, cache_slots, pack_dir, stats_path,
+         quant="bf16"):
     """Hottest experts (by steps that routed to them) pinned in VRAM, next in RAM, rest on disk."""
     uses = {int(layer): s["uses"] for layer, s in stats["layers"].items()}
     vram_rows = int(expert_vram_gib * GIB // expert_size) - cache_slots * len(uses)
@@ -298,7 +312,7 @@ def plan(stats, num_experts, expert_size, expert_vram_gib, ram_gib, cache_slots,
     }
     return render_toml(
         "generated by `prain plan`; experts not listed live on disk",
-        {"pack_dir": pack_dir, "default_tier": "disk", "cache_slots": cache_slots,
+        {"pack_dir": pack_dir, "quant": None if quant == "bf16" else quant, "default_tier": "disk", "cache_slots": cache_slots,
          "stats_path": stats_path},
         layers=layers,
     )

prain/experts.py

@@ -2,14 +2,17 @@
 
 vLLM builds this class instead of RoutedExperts once prain.plugin.register() ran.
 Every expert is written once to a per-layer pack file on disk while vLLM loads the
-checkpoint. On the first forward the RAM tier is copied from the pack into pinned host
+checkpoint. With quant = "int8" or "int4" the first forward quantizes that bf16 pack once into
-memory and the pinned VRAM slots are filled. After that each forward:
+a second pack (prain.quant), and the VRAM slots, copies and CPU path use the quantized
+experts. On the first forward the RAM tier is copied from the pack into pinned host memory
+and the pinned VRAM slots are filled. After that each forward:
   0. drops the routes of draft tokens to experts that are not available ([draft] min_weight;
      prain.drafts knows which rows are drafts),
   1. counts routes per expert (one host sync),
   2. lets the [miss] config send some missing experts to the CPU,
-  3. copies the other missing experts into slots chosen by the [policy] config and runs the
+  3. copies the other missing experts into slots chosen by the [policy] config and runs
-     stock fused-MoE kernel on the slot tensors, in rounds if needed,
+     vLLM's Triton fused-MoE kernel on the slot tensors, in rounds if needed (each round
+     computes only its own routes),
   4. queues async copies on a side stream: [prefetch] guesses for a later layer, and the
      hottest CPU-run experts (admit-behind),
   5. adds the CPU result.
@@ -26,11 +29,14 @@ import os
 import torch
 from vllm.config import get_current_vllm_config
 from vllm.logger import init_logger
+from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts
 from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
 
 from prain import drafts
 from prain.config import layer_index, load_config
 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.sim import write_trace
 from prain.tiered import draft_rows, tiered_forward
 
@@ -133,9 +139,16 @@ class TieredRoutedExperts(RoutedExperts):
             raise NotImplementedError("prain: prefetch.mode = pregate needs a router module "
                                       "passed as custom_routing_function (Gemma 4 does)")
         # vLLM sizes w13/w2 and the kernel by num_experts: make that the VRAM slot count.
-        moe_config.num_experts = moe_config.num_local_experts = placement.num_slots
+        # Quantized slots are prain's own tensors; vLLM's bf16 ones shrink to one row.
+        quantized = config.quant != "bf16"
+        moe_config.num_experts = moe_config.num_local_experts = 1 if quantized else placement.num_slots
         super().__init__(layer_name, params_dtype, moe_config, quant_config, *args, **kwargs)
         self.local_num_experts = placement.num_slots
+        self.quant = config.quant
+        self._shapes = (tuple(self.w13_weight.shape[1:]), tuple(self.w2_weight.shape[1:]))
+        self._qslots = [[torch.zeros((placement.num_slots, *shape), dtype=dtype, device="cuda")
+                         for shape, dtype in layout(config.quant, *dims)]
+                        for dims in self._shapes] if quantized else None
         self.prain = config
         self.prain_layer = layer
         self.placement = placement
@@ -155,27 +168,69 @@ class TieredRoutedExperts(RoutedExperts):
         logger.info(
             "prain layer %d: %d vram, %d ram, %d disk, %d cache slots%s",
             layer, len(placement.vram), len(placement.ram), len(placement.disk),
-            placement.cache_slots, "" if self.pack_complete else ", writing pack",
+            placement.cache_slots, "" if self.pack_complete else
+            ", writing pack" if self._writing else f", quantizing the pack to {self.quant}",
         )
 
     def _open_pack(self, pack_dir, layer, num_experts, model):
+        """The loader writes the bf16 pack; a quantized pack is made from it once. Once the
+        pack of config.quant is complete, the bf16 one is not read (it may be deleted)."""
         os.makedirs(pack_dir, exist_ok=True)
-        stem = os.path.join(pack_dir, f"layer{layer:03d}")
+        self._stem = os.path.join(pack_dir, f"layer{layer:03d}")
-        w13_shape = (num_experts, *self.w13_weight.shape[1:])
+        self._num_experts = num_experts
-        w2_shape = (num_experts, *self.w2_weight.shape[1:])
+        self._marker = {"model": model, "w13": [num_experts, *self._shapes[0]],
-        self._pack_files = (f"{stem}.w13", f"{stem}.w2")
+                        "w2": [num_experts, *self._shapes[1]], "dtype": str(self.w13_weight.dtype)}
-        self._pack_w13 = _file_tensor(self._pack_files[0], w13_shape, self.w13_weight.dtype)
+        self.pack_complete = self._complete(self.quant)
-        self._pack_w2 = _file_tensor(self._pack_files[1], w2_shape, self.w2_weight.dtype)
+        self._bf16 = None if self.pack_complete else self._pack_of("bf16")
-        self._marker_path = f"{stem}.done"
+        self._writing = self._bf16 is not None and not self._complete("bf16")
-        self._marker = {"model": model, "w13": list(w13_shape), "w2": list(w2_shape),
+
-                        "dtype": str(self.w13_weight.dtype)}
+    def _paths(self, quant):
-        self.pack_complete = os.path.exists(self._marker_path)
+        """[[w13 files], [w2 files]] of a pack format: weights, then scales."""
-        if self.pack_complete:
+        infix = "" if quant == "bf16" else f".{quant}"
-            with open(self._marker_path) as f:
+        return [[f"{self._stem}{infix}.{name}{'s' * i}" for i in range(len(layout(quant, *dims)))]
-                found = json.load(f)
+                for name, dims in zip(("w13", "w2"), self._shapes)]
-            if found != self._marker:
+
-                raise ValueError(f"prain: {self._marker_path} is for {found}, not "
+    def _pack_of(self, quant):
-                                 f"{self._marker}; delete {pack_dir} or use another pack_dir")
+        return [[_file_tensor(path, (self._num_experts, *shape), dtype)
+                 for path, (shape, dtype) in zip(paths, layout(quant, *dims))]
+                for paths, dims in zip(self._paths(quant), self._shapes)]
+
+    def _marker_of(self, quant):
+        if quant == "bf16":
+            return f"{self._stem}.done", self._marker
+        return f"{self._stem}.{quant}.done", {**self._marker, "quant": quant, "group": GROUP}
+
+    def _complete(self, quant):
+        path, marker = self._marker_of(quant)
+        if not os.path.exists(path):
+            return False
+        with open(path) as f:
+            found = json.load(f)
+        if found != marker:
+            raise ValueError(f"prain: {path} is for {found}, not {marker}; delete the pack "
+                             "or use another pack_dir")
+        return True
+
+    def _seal(self, quant):
+        for path in sum(self._paths(quant), []):  # mmap writes -> disk before we trust the pack
+            fd = os.open(path, os.O_RDONLY)
+            os.fsync(fd)
+            os.close(fd)
+        path, marker = self._marker_of(quant)
+        with open(path, "w") as f:
+            json.dump(marker, f)
+
+    def _quantize_pack(self):
+        """bf16 pack -> config.quant pack, one expert at a time on the GPU (int4's clip search
+        takes ~20 min for the whole model on the CPU, under a minute on the GPU)."""
+        out = self._pack_of(self.quant)
+        for src, dst in zip(self._bf16, out):
+            for e in range(self._num_experts):
+                parts = quantize(src[0][e].cuda(), self.quant)
+                if self.quant == "int4":
+                    parts[0] = to_cpu_int4(parts[0])
+                for tensor, part in zip(dst, parts):
+                    tensor[e] = part.cpu()
 
     def get_expert_mapping(self, ckpt_gate_proj_name=None, ckpt_down_proj_name=None,
                            ckpt_up_proj_name=None, include_fused=False):
@@ -198,8 +253,8 @@ class TieredRoutedExperts(RoutedExperts):
     def weight_loader(self, param, loaded_weight, weight_name, shard_id, expert_id,
                       return_success=False):
         # Every expert goes to the pack; the VRAM and RAM tiers fill from it later.
-        if not self.pack_complete:
+        if self._writing:
-            pack = self._pack_w13 if param is self.w13_weight else self._pack_w2
+            pack = self._bf16[0][0] if param is self.w13_weight else self._bf16[1][0]
             super().weight_loader(pack, loaded_weight, weight_name, shard_id, expert_id,
                                   return_success=True)
         return True if return_success else None
@@ -208,15 +263,23 @@ class TieredRoutedExperts(RoutedExperts):
 
     def _materialize(self):
         if not self.pack_complete:
-            for path in self._pack_files:  # mmap writes -> disk before we trust the pack
+            if self._writing:
-                fd = os.open(path, os.O_RDONLY)
+                self._seal("bf16")
-                os.fsync(fd)
+            if self.quant != "bf16":
-                os.close(fd)
+                self._quantize_pack()
-            with open(self._marker_path, "w") as f:
+                self._seal(self.quant)
-                json.dump(self._marker, f)
+            self.pack_complete, self._writing, self._bf16 = True, False, None
-            self.pack_complete = True
+        self._pack = self._pack_of(self.quant)  # [[w13 tensors], [w2 tensors]] per expert row
+        if self._qslots is None:
+            self._slots, self._qconfig = [[self.w13_weight.data], [self.w2_weight.data]], None
+        else:
+            self._slots = self._qslots
+            self._qconfig = moe_quant_config(self.quant, self._slots[0][1], self._slots[1][1])
+            # vLLM's one-row bf16 w13/w2 are never read: free them (11 MiB per layer)
+            self.w13_weight.data = self.w13_weight.data.new_empty(0)
+            self.w2_weight.data = self.w2_weight.data.new_empty(0)
         if self.prain_layer == min(LAYERS):  # check the whole RAM tier once, before pinning
-            row = self._pack_w13[0].nbytes + self._pack_w2[0].nbytes
+            row = sum(t[0].nbytes for mats in self._pack for t in mats)
             need = row * sum(len(experts.placement.ram) for experts in LAYERS.values())
             free = _mem_available()
             if free is not None and need > 0.9 * free:
@@ -225,30 +288,35 @@ class TieredRoutedExperts(RoutedExperts):
                                    "from ram to disk")
         ram = list(self.placement.ram)
         self._ram_row = {e: row for row, e in enumerate(ram)}
-        # One pinned tensor per expert: pinning a single 340 MiB block took 6-23 s on an
+        # ponytail: pinning is slow on the eGPU host (minutes for a 10 GiB tier, even per
-        # eGPU host (hipHostMalloc), small blocks pin at the normal rate.
+        # expert); the page cache serves the pack about as fast, so keep the RAM tier small.
-        self._ram_w13 = [self._pack_w13[e].pin_memory() for e in ram]
+        self._ram = [[[t[e].pin_memory() for t in mats] for mats in self._pack] for e in ram]
-        self._ram_w2 = [self._pack_w2[e].pin_memory() for e in ram]
         for slot, expert in enumerate(self.placement.vram):
             self._load(expert, slot)
         if self._router and (self.prain.prefetch.mode == "pregate" or self.prain.trace_path):
             self._router[0].register_forward_pre_hook(self._predict)
 
     def _source(self, expert):
-        """CPU tensors of an expert: pinned rows for the RAM tier, else the mmap'd pack."""
+        """[[w13 tensors], [w2 tensors]] of an expert on the CPU: pinned rows for the RAM tier,
+        else the mmap'd pack."""
         row = self._ram_row.get(expert)
         if row is None:
-            return self._pack_w13[expert], self._pack_w2[expert]
+            return [[t[expert] for t in mats] for mats in self._pack]
-        return self._ram_w13[row], self._ram_w2[row]
+        return self._ram[row]
+
+    def _copy(self, 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
+                    src = to_gpu_int4(src.to(dst.device, non_blocking=True))
+                dst[slot].copy_(src, non_blocking=True)
 
     def _load(self, expert, slot):
         """Demand copy on the compute stream."""
         pending = self._ready_events.pop(slot, None)
         if pending is not None:  # an older async copy into this slot must land first
             torch.cuda.current_stream().wait_event(pending)
-        w13, w2 = self._source(expert)
+        self._copy(expert, slot)
-        self.w13_weight.data[slot].copy_(w13, non_blocking=True)
-        self.w2_weight.data[slot].copy_(w2, non_blocking=True)
 
     def _load_async(self, expert, slot):
         """Copy on the side stream; only RAM-tier (pinned) experts are truly asynchronous."""
@@ -256,23 +324,23 @@ class TieredRoutedExperts(RoutedExperts):
         last = self._last_use.get(slot)
         if last is not None:  # a queued kernel may still read the slot's old expert
             stream.wait_event(last)
-        w13, w2 = self._source(expert)
         with torch.cuda.stream(stream):
-            self.w13_weight.data[slot].copy_(w13, non_blocking=True)
+            self._copy(expert, slot)
-            self.w2_weight.data[slot].copy_(w2, non_blocking=True)
         event = torch.cuda.Event()
         event.record(stream)
         self._ready_events[slot] = event
 
-    def _run_slots(self, x, weights, slot_ids, used_slots):
+    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."""
         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 = super().forward_modular(x, weights, slot_ids)
+        out = fused_experts(x, self._slots[0][0], self._slots[1][0], weights, ids,
-        if not isinstance(out, torch.Tensor):
+                            activation=self.activation, global_num_experts=table.numel(),
-            raise NotImplementedError("prain: deferred MoE finalize is not supported")
+                            expert_map=table, quant_config=self._qconfig)
         if self._async:
             event = torch.cuda.Event()
             event.record(current)
@@ -296,7 +364,10 @@ class TieredRoutedExperts(RoutedExperts):
             self._predictions[d] = list(zip(e.tolist(), p.tolist()))  # host sync
 
     def _cpu_weights(self, expert):
-        return self._source(expert)
+        mats = self._source(expert)
+        if self.quant == "bf16":
+            return mats[0][0], mats[1][0]
+        return tuple(cpu_matmul(self.quant, *parts) for parts in mats)
 
     def _after_enqueue(self, cpu, stats):
         miss, prefetch = self.prain.miss, self.prain.prefetch

prain/gemv.py

@@ -0,0 +1,89 @@
+"""Decode GEMMs (at most MAX_ROWS rows) for gfx10 GPUs as a Triton GEMV.
+
+vLLM sends unquantized linear layers on ROCm through rocBLAS, except on gfx9/gfx11/gfx12 where it
+has skinny-GEMM kernels. On an RX 6600 XT (gfx1032) rocBLAS runs a bf16 GEMV at 13-16 GB/s
+(RDNA2 has no bf16 dot instructions), so the dense part of a decode step reads its weights at
+6% of the card's bandwidth. This kernel reads them at ~210 GB/s (measured, lm_head 262144 x 2816:
+93 ms -> 6.7 ms). install() wraps vLLM's dispatch_unquantized_gemm; PRAIN_GEMV=0 turns it off.
+"""
+
+import functools
+import os
+
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+
+MAX_ROWS = 4
+
+
+@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):
+    rows = tl.program_id(0) * BLOCK_N + tl.arange(0, BLOCK_N)
+    ms = tl.arange(0, BLOCK_M)
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for k0 in range(0, K, BLOCK_K):
+        cols = k0 + tl.arange(0, BLOCK_K)
+        w = tl.load(w_ptr + rows[:, None] * K + cols[None, :],
+                    mask=(rows[:, None] < N) & (cols[None, :] < K), other=0.0)
+        x = tl.load(x_ptr + ms[:, None] * K + cols[None, :],
+                    mask=(ms[:, None] < M) & (cols[None, :] < K), other=0.0)
+        acc += tl.sum(x.to(tl.float32)[:, None, :] * w.to(tl.float32)[None, :, :], axis=2)
+    tl.store(y_ptr + ms[:, None] * N + rows[None, :], acc.to(y_ptr.dtype.element_ty),
+             mask=(ms[:, None] < M) & (rows[None, :] < N))
+
+
+@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"
+
+
+def _small(x, weight, bias):
+    return (bias is None and x.is_cuda and x.dtype in (torch.bfloat16, torch.float16)
+            and weight.dtype == x.dtype and weight.is_contiguous()
+            and x.numel() <= MAX_ROWS * x.shape[-1] and fast_device(x.device.index or 0))
+
+
+def linear(x, weight, bias=None):
+    """F.linear(x, weight, bias); the Triton GEMV for up to MAX_ROWS rows on gfx10."""
+    if not _small(x, weight, bias):
+        return F.linear(x, weight, bias)
+    n, k = weight.shape
+    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))
+    return y.reshape(*x.shape[:-1], n)
+
+
+def wrap(dispatch):
+    """vLLM's dispatch_unquantized_gemm(), with decode shapes sent to linear()."""
+    def wrapped(*args, **kwargs):
+        fallback = dispatch(*args, **kwargs)
+
+        def gemm(layer, x, weight, bias=None):
+            if _small(x, weight, bias):
+                return linear(x, weight)
+            return fallback(layer, x, weight, bias)
+        return gemm
+    return wrapped
+
+
+def install():
+    """Before the model is built: every module that looks up dispatch_unquantized_gemm."""
+    from vllm.model_executor.layers import linear as linear_layers, utils
+    from vllm.model_executor.layers import vocab_parallel_embedding
+    from vllm.model_executor.layers.fused_moe.runner import moe_runner
+
+    wrapped = wrap(utils.dispatch_unquantized_gemm)
+    for module in (linear_layers, vocab_parallel_embedding, moe_runner):
+        module.dispatch_unquantized_gemm = wrapped

prain/plugin.py

@@ -10,7 +10,9 @@ def register():
 
     if "RoutedExperts" in op_registry_oot:  # vLLM may call register() more than once
         return
+    from prain import gemv
     from prain.experts import TieredRoutedExperts
 
     # vLLM's documented out-of-tree hook: every RoutedExperts(...) now builds ours.
     PluggableLayer.register_oot(TieredRoutedExperts, name="RoutedExperts")
+    gemv.install()  # decode GEMVs on gfx10

prain/quant.py

@@ -0,0 +1,89 @@
+"""Weight-only int8 / int4 quantization of expert weights. Needs torch.
+
+int8: q int8 [..., N, K] and one fp32 scale per output row: w = q * scale. vLLM's Triton
+      fused-MoE kernel and torch's _weight_int8pack_mm (CPU) both read this layout.
+int4: one bf16 scale per GROUP inputs of a row: w = (nibble - 8) * scale, with the scale's
+      clip ratio chosen per group for the least squared error. Two layouts of the nibbles:
+      gpu: [..., N, K/2], two values per byte along K, even k in the low nibble (vLLM's
+           Triton GPTQ/AWQ MoE kernel);
+      cpu: [N/32, K, 16] bytes, rows d and d+16 of each 32-row block share a byte (torch's AVX2
+           _weight_int4pack_mm_for_cpu); packs store it, slots convert with to_gpu_int4().
+"""
+
+import torch
+
+from prain.config import INT4_GROUP as GROUP  # divides Gemma 4 26B-A4B's 2816 and 704
+
+CLIP = (1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7)
+
+
+def layout(quant, rows, cols):
+    """[(shape, dtype)] of the tensors that store one [rows, cols] weight matrix."""
+    if quant == "bf16":
+        return [((rows, cols), torch.bfloat16)]
+    if quant == "int8":
+        return [((rows, cols), torch.int8), ((rows,), torch.float32)]
+    return [((rows, cols // 2), torch.uint8), ((rows, cols // GROUP), torch.bfloat16)]
+
+
+def quantize(w, quant):
+    """Float weights [..., N, K] -> [q, scale], int4 in the gpu layout."""
+    w = w.float()
+    if quant == "int8":
+        scale = w.abs().amax(-1).clamp_min(1e-12) / 127
+        return [(w / scale[..., None]).round().to(torch.int8), scale]
+    g = w.unflatten(-1, (-1, GROUP))
+    amax = g.abs().amax(-1, keepdim=True).clamp_min(1e-12)
+    best_err = best_scale = None
+    for ratio in CLIP:
+        scale = (amax * ratio / 7).bfloat16().float()  # the kernels see the bf16 scale
+        err = ((g / scale).round().clamp(-8, 7) * scale - g).pow(2).sum(-1, keepdim=True)
+        if best_err is None:
+            best_err, best_scale = err, scale
+        else:
+            better = err < best_err
+            best_err, best_scale = torch.where(better, err, best_err), torch.where(better, scale, best_scale)
+    q = ((g / best_scale).round().clamp(-8, 7) + 8).flatten(-2).to(torch.uint8)
+    return [q[..., 0::2] | q[..., 1::2] << 4, best_scale.squeeze(-1).bfloat16()]
+
+
+def dequantize(q, scale, quant, dtype=torch.bfloat16):
+    """The weights [..., N, K] of quantize()'s output (int4 in the gpu layout)."""
+    if quant == "int8":
+        return q.to(dtype) * scale.to(dtype)[..., None]
+    v = torch.stack((q & 15, q >> 4), -1).flatten(-2).to(dtype) - 8
+    return (v.unflatten(-1, (-1, GROUP)) * scale.to(dtype)[..., None]).flatten(-2)
+
+
+def to_cpu_int4(q):
+    """int4 nibbles [..., N, K/2]: gpu layout -> cpu layout (same shape)."""
+    *lead, n, half = q.shape
+    v = torch.stack((q & 15, q >> 4), -1).reshape(*lead, n // 32, 32, 2 * half).transpose(-1, -2)
+    return (v[..., :16] | v[..., 16:] << 4).reshape(q.shape)
+
+
+def to_gpu_int4(c):
+    """int4 nibbles [..., N, K/2]: cpu layout -> gpu layout (same shape; works on the GPU)."""
+    *lead, n, half = c.shape
+    b = c.reshape(*lead, n // 32, 2 * half, 16)
+    v = torch.cat((b & 15, b >> 4), -1).transpose(-1, -2).reshape(*lead, n, 2 * half)
+    return v[..., 0::2] | v[..., 1::2] << 4
+
+
+def cpu_matmul(quant, q, scale):
+    """x -> x @ w.T on the CPU for stored int8 q / cpu-layout int4 q; x is bf16 [M, K]."""
+    if quant == "int8":
+        return lambda x: torch._weight_int8pack_mm(x, q, scale.to(x.dtype))
+    zeros = torch.stack((scale.t(), torch.zeros_like(scale.t())), -1).contiguous()
+    return lambda x: torch._weight_int4pack_mm_for_cpu(x, q, GROUP, zeros)
+
+
+def moe_quant_config(quant, w13_scale, w2_scale):
+    """vLLM's fused-MoE quant config for int8 / gpu-layout int4 slot tensors."""
+    from vllm.model_executor.layers.fused_moe import config
+
+    if quant == "int8":
+        return config.int8_w8a16_moe_quant_config(w1_scale=w13_scale, w2_scale=w2_scale,
+                                                  w1_zp=None, w2_zp=None)
+    return config.int4_w4a16_moe_quant_config(w1_scale=w13_scale, w2_scale=w2_scale, w1_zp=None,
+                                              w2_zp=None, block_shape=[0, GROUP])

prain/tiered.py

@@ -49,7 +49,7 @@ def degrade(topk_weights, topk_ids, draft, available, min_weight):
 
     A draft token keeps its heaviest unavailable routes only while less than min_weight of its
     router weight would be computed; its kept routes are scaled up to its full router weight.
-    Dropped routes get id -1 and weight 0, which route_stats and remap ignore.
+    Dropped routes get id -1 and weight 0, which route_stats and slot_table ignore.
     draft: bool [tokens]; available: bool [num_experts]. Full tokens are left unchanged."""
     ids = topk_ids.long()
     valid = ids >= 0
@@ -68,37 +68,36 @@ def degrade(topk_weights, topk_ids, draft, available, min_weight):
     return weights, torch.where(drop, torch.full_like(topk_ids, -1), topk_ids)
 
 
-def remap(topk_weights, topk_ids, slot_of, num_experts):
+def slot_table(slot_of, num_experts, device):
-    """Point routes at VRAM slots for one round.
+    """int32 [num_experts + 1]: VRAM slot of each expert in one round, -1 for the others.
-
+    vLLM's fused-MoE kernel takes it as expert_map and skips routes that map to -1; the extra
-    Routes to experts outside this round (and padding routes with id -1) get slot 0
+    last row stands for the padding id (route id -1 arrives as num_experts)."""
-    and weight 0, so any fused-MoE kernel can run them unchanged."""
+    table = torch.full((num_experts + 1,), -1, dtype=torch.int32)
-    table = torch.full((num_experts,), -1, dtype=torch.int64)
     if slot_of:
-        table[torch.tensor(list(slot_of), dtype=torch.int64)] = torch.tensor(
+        table[torch.tensor(list(slot_of))] = torch.tensor(list(slot_of.values()), dtype=torch.int32)
-            list(slot_of.values()), dtype=torch.int64
+    return table.to(device)
-        )
+
-    table = table.to(topk_ids.device)
+
-    ids = topk_ids.long()
+def _matmul(w):
-    slots = torch.where(ids >= 0, table[ids.clamp_min(0)], -1)
+    """x -> x @ w.T for a weight tensor; functions (quantized experts) pass through."""
-    weights = torch.where(slots >= 0, topk_weights, torch.zeros_like(topk_weights))
+    return w if callable(w) else (lambda v: v.to(w.dtype) @ w.t())
-    return weights, slots.clamp_min(0).to(topk_ids.dtype)
 
 
 def cpu_experts(x, topk_weights, topk_ids, experts, weights_of):
     """The routes to `experts`, computed on the CPU like vLLM's fused kernel with activation
     "gelu_tanh": w13 = [gate; up], y = (gelu_tanh(x @ gate.T) * (x @ up.T)) @ w2.T times the
-    route weight. weights_of(e) -> (w13 [2I, H], w2 [H, I]). Returns [tokens, hidden] fp32."""
+    route weight. weights_of(e) -> (w13 [2I, H], w2 [H, I]), or two functions v -> v @ w.T
+    (quantized experts). Returns [tokens, hidden] fp32."""
     out = torch.zeros(x.shape, dtype=torch.float32)
     for e in sorted(experts):
         rows, cols = (topk_ids == e).nonzero(as_tuple=True)
         if rows.numel() == 0:
             continue
-        w13, w2 = weights_of(e)
+        mm13, mm2 = map(_matmul, weights_of(e))
-        h = x[rows].to(w13.dtype) @ w13.t()
+        h = mm13(x[rows])
-        half = w13.shape[0] // 2
+        half = h.shape[1] // 2
         act = F.gelu(h[:, :half].float(), approximate="tanh") * h[:, half:].float()
-        y = (act.to(w2.dtype) @ w2.t()).float()
+        y = mm2(act.to(h.dtype)).float()
         out.index_add_(0, rows, y * topk_weights[rows, cols].float()[:, None])
     return out
 
@@ -108,20 +107,32 @@ def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight"
                    after_enqueue=None, draft=None, min_weight=0.5):
     """One MoE forward over VRAM slots.
 
-    cache: SlotCache. load(expert, slot) copies an expert into a slot. kernel(x, weights,
+    cache: SlotCache. load(expert, slot) copies an expert into a slot. kernel(x, weights, ids,
-    slot_ids, used_slots) runs the stock fused-MoE kernel on the slot tensors. With a
+    table, used_slots) runs the fused-MoE kernel on the slot tensors for the routes whose
+    expert the slot_table() `table` maps to a slot. With a
     [miss] config other than "copy", the experts it picks run on the CPU (cpu_weights(e)
     gives their weights) while the GPU rounds run. on_stats(stats, tokens) sees the routing
     statistics; after_enqueue(cpu, stats) runs once all GPU work of this forward is queued.
-    draft: bool [tokens] for the rows of draft requests, or None; see degrade().
+    draft: bool [tokens] for the rows of draft requests, or None; see degrade(). Draft rows keep
+    min_weight of their router weight, every row keeps [miss] min_weight (1 = everything).
     """
     raw = None
-    if draft is not None:
+    skip = miss is not None and miss.min_weight < 1
-        if on_stats is not None:  # the trace records the routing, not what drafts skipped
+    if draft is not None or skip:
+        if on_stats is not None:  # the trace records the routing, not what was skipped
             raw = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)
-        available = available_experts(topk_ids, draft, set(cache.pinned) | set(cache.slots),
+        rows = draft if draft is not None else torch.zeros(topk_ids.shape[0], dtype=torch.bool,
+                                                           device=topk_ids.device)
+        need = min_weight
+        if skip:  # every row may skip; draft rows by the smaller share
+            need = torch.where(rows, min(min_weight, miss.min_weight), miss.min_weight)[:, None]
+            rows = torch.ones_like(rows)
+        available = available_experts(topk_ids, rows, set(cache.pinned) | set(cache.slots),
                                       cache.num_experts)
-        topk_weights, topk_ids = degrade(topk_weights, topk_ids, draft, available, min_weight)
+        _, kept = degrade(topk_weights, topk_ids, rows, available, need)
+        # experts some row copies anyway serve every row for free
+        available |= torch.bincount(kept[kept >= 0].long(), minlength=cache.num_experts) > 0
+        topk_weights, topk_ids = degrade(topk_weights, topk_ids, rows, available, need)
     stats = route_stats(topk_weights, topk_ids, cache.num_experts, recency_half_life)  # host sync
     tokens = int(topk_ids.shape[0])
     if on_stats is not None:
@@ -131,22 +142,22 @@ def tiered_forward(x, topk_weights, topk_ids, cache, load, kernel, heat="weight"
     if cpu:  # copy the inputs now: later the copy would wait for the queued GPU rounds
         host = x.cpu(), topk_weights.cpu(), topk_ids.cpu()
     index = HEAT_INDEX[heat]
-    outs = []
+    ids = torch.where(topk_ids >= 0, topk_ids, cache.num_experts)  # padding -> the table's last row
+    acc = torch.promote_types(x.dtype, torch.float32)
+    total = None  # sum of the round (and CPU) outputs: one buffer, not one per round
+
+    def add(out):
+        nonlocal total
+        total = out if total is None else total.to(acc).add_(out)  # .to: a no-op once fp32
+
     for slot_of, loads in cache.plan({e: s[index] for e, s in stats.items()}, tokens, cpu):
         for expert, slot in loads:
             load(expert, slot)
-        if slot_of:
+        if slot_of:  # each round computes only its own routes
-            weights, slots = remap(topk_weights, topk_ids, slot_of, cache.num_experts)
+            table = slot_table(slot_of, cache.num_experts, x.device)
-            outs.append(kernel(x, weights, slots, list(slot_of.values())))
+            add(kernel(x, topk_weights, ids, table, list(slot_of.values())))
     if after_enqueue is not None:
         after_enqueue(cpu, stats)
     if cpu:  # runs while the GPU works through its rounds
-        outs.append(cpu_experts(*host, cpu, cpu_weights).to(x.device))
+        add(cpu_experts(*host, cpu, cpu_weights).to(x.device))
-    if not outs:
+    return torch.zeros_like(x) if total is None else total.to(x.dtype)
-        return torch.zeros_like(x)
-    if len(outs) == 1:
-        return outs[0].to(x.dtype)
-    # ponytail: every extra round re-runs all tokens with weight 0 for experts it does not
-    # hold; only big batches need many rounds. Upgrade: compact tokens per round.
-    acc = torch.promote_types(x.dtype, torch.float32)
-    return torch.stack([o.to(acc) for o in outs]).sum(0).to(x.dtype)

scripts/bench.py

@@ -3,7 +3,9 @@
 usage (prain installed): python3 scripts/bench.py http://127.0.0.1:8000 [--tokens 128] [--runs 3]
 Each run sends a fresh prompt (a random prefix defeats the prefix cache) and decodes exactly
 --tokens tokens (ignore_eos), so numbers compare across runs and configs. The last line holds
-the medians.
+the medians. --context FILE... puts the files' text (up to --context-chars) before every
+prompt and the random tag after it: run 1 prefills the whole context, later runs find it in
+the prefix cache, like an agent's system prompt and tools.
 """
 
 import argparse
@@ -21,8 +23,8 @@ PROMPTS = [
 ]
 
 
-def run(base_url, model, prompt, tokens):
+def run(base_url, model, prompt, tokens, context=""):
-    text = f"[{uuid.uuid4().hex[:8]}] {prompt}"
+    text = f"{context}[{uuid.uuid4().hex[:8]}] {prompt}"
     body = chat_body(model, [{"role": "user", "content": text}], max_tokens=tokens,
                      temperature=0, ignore_eos=True)
     start = time.monotonic()
@@ -39,11 +41,18 @@ def main():
     parser.add_argument("base_url")
     parser.add_argument("--tokens", type=int, default=128)
     parser.add_argument("--runs", type=int, default=3)
+    parser.add_argument("--context", nargs="*", default=[], help="files to put before each prompt")
+    parser.add_argument("--context-chars", type=int, default=48000)
     args = parser.parse_args()
     model = first_model(args.base_url)
+    context = ""
+    for path in args.context:
+        with open(path) as f:
+            context += f.read()
+    context = context[: args.context_chars] + "\n\n" if context else ""
     results = []
     for i in range(args.runs):
-        results.append(run(args.base_url, model, PROMPTS[i % len(PROMPTS)], args.tokens))
+        results.append(run(args.base_url, model, PROMPTS[i % len(PROMPTS)], args.tokens, context))
         print(json.dumps(results[-1]), flush=True)
     print(json.dumps({key: statistics.median(r[key] for r in results) for key in results[0]}
                      | {"median_of": len(results), "tokens": args.tokens}))

scripts/env.sh

@@ -23,4 +23,7 @@ export PATH
 # libc++/compiler-rt defaults: everything links libstdc++/libgcc like the ROCm libraries.
 export CC=$PRAIN_ROOT/build/bin/clang CXX=$PRAIN_ROOT/build/bin/clang++
 export CFLAGS="-march=native" CXXFLAGS="-march=native"
+# Fused-MoE kernel configs tuned for this GPU by scripts/tune_moe.py (vLLM's defaults when
+# the files are missing).
+export VLLM_TUNED_CONFIG_FOLDER=$PRAIN_ROOT/build/moe-configs
 unset _site _dev _core

scripts/measure.py

@@ -1,7 +1,7 @@
 """Measure the prain cost model on this machine and print it as JSON for
 `prain sim --cost` and `prain tune --cost`. Needs the GPU, vLLM and prain.
 
-usage: python3 scripts/measure.py PACK_FILE [--threads N] [--dense-gib 4.45] [--layers 30] > cost.json
+usage: python3 scripts/measure.py PACK_FILE [--quant int4] [--threads N] [--dense-gib 4.45] [--layers 30] > cost.json
 PACK_FILE is any <pack_dir>/layerNNN.w13 file; it is read cold (dropped from the page cache).
 """
 
@@ -15,6 +15,8 @@ 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.config import QUANTS, expert_bytes as size_of
+from prain.quant import cpu_matmul, moe_quant_config, quantize, to_cpu_int4
 from prain.tiered import cpu_experts
 
 HIDDEN, INTER, SLOTS = 2816, 704, 5  # Gemma 4 26B-A4B expert, 5 slots per layer
@@ -37,10 +39,11 @@ def main():
     parser.add_argument("--threads", type=int, default=4)
     parser.add_argument("--dense-gib", type=float, default=4.45, help="non-expert weights")
     parser.add_argument("--layers", type=int, default=30)
+    parser.add_argument("--quant", choices=QUANTS, default="bf16", help="expert weight format")
     args = parser.parse_args()
     torch.set_num_threads(args.threads)
-    expert_bytes = 3 * HIDDEN * INTER * 2
+    expert_bytes = size_of(HIDDEN, INTER, args.quant)
-    host = torch.empty(expert_bytes // 2, dtype=torch.bfloat16).pin_memory()
+    host = torch.empty(expert_bytes, dtype=torch.uint8).pin_memory()
     device = torch.empty_like(host, device="cuda")
     h2d_gbps = expert_bytes / 1e6 / ms(lambda: device.copy_(host, non_blocking=True))
 
@@ -56,13 +59,17 @@ def main():
 
     w1 = torch.randn(SLOTS, 2 * INTER, HIDDEN, dtype=torch.bfloat16, device="cuda") * 0.02
     w2 = torch.randn(SLOTS, HIDDEN, INTER, dtype=torch.bfloat16, device="cuda") * 0.02
+    quant_config = None
+    if args.quant != "bf16":
+        (w1, s1), (w2, s2) = quantize(w1, args.quant), quantize(w2, args.quant)
+        quant_config = moe_quant_config(args.quant, s1, s2)
 
     def moe(tokens):
         x = torch.randn(tokens, HIDDEN, dtype=torch.bfloat16, device="cuda")
         ids = torch.arange(SLOTS, device="cuda", dtype=torch.int32).repeat(tokens, 1)
         weights = torch.full((tokens, SLOTS), 1 / SLOTS, device="cuda")
         return lambda: fused_experts(x, w1, w2, weights, ids, activation=MoEActivation.GELU_TANH,
-                                     global_num_experts=SLOTS)
+                                     global_num_experts=SLOTS, quant_config=quant_config)
 
     with set_current_vllm_config(VllmConfig()):
         round_ms = ms(moe(1))
@@ -76,6 +83,12 @@ def main():
     # Serving computes a different expert each time, read from RAM (not from the CPU caches).
     experts = [(torch.randn(2 * INTER, HIDDEN, dtype=torch.bfloat16),
                 torch.randn(HIDDEN, INTER, dtype=torch.bfloat16)) for _ in range(23)]
+    if args.quant != "bf16":  # the CPU kernels on pack-layout weights, as the plugin runs them
+        def stored(w):
+            q, scale = quantize(w, args.quant)
+            return cpu_matmul(args.quant, to_cpu_int4(q) if args.quant == "int4" else q, scale)
+
+        experts = [(stored(a), stored(b)) for a, b in experts]
     x = torch.randn(1, HIDDEN, dtype=torch.bfloat16)
     one = torch.ones(1, 1)
     ids = torch.zeros(1, 1, dtype=torch.int32)

scripts/quality.py

@@ -0,0 +1,89 @@
+"""Answer quality of a served config against reference answers from another config.
+
+    python scripts/quality.py URL make REF.json    greedy reference answers (run on the exact config)
+    python scripts/quality.py URL score REF.json   how likely this server finds those answers
+
+Both work on token ids: the prompt is the chat template with its generation prompt, and `score`
+re-scores the reference answer on exactly the ids it was generated from (prompt logprobs). A
+finished assistant message renders differently from the generation prompt, so re-scoring
+through the chat API skews the logprobs. `nll` is the mean negative log-likelihood per answer
+token (the reference config scores about its own decode logprobs); `same_tokens` is how many
+leading tokens this server's own greedy answer shares with the reference.
+"""
+
+import argparse
+import json
+import urllib.request
+
+PROMPTS = [
+    "Write a Python function that returns the n-th Fibonacci number iteratively.",
+    "Explain what `git rebase -i HEAD~3` does, in three sentences.",
+    "Write a bash one-liner that counts the lines of all .py files below the current directory.",
+    "On Gentoo, how do I rebuild every package that links against a library I just updated?",
+    "What is the difference between a process and a thread? Answer in two sentences.",
+]
+CODE_PROMPT = "Here is some Python code:\n\n{code}\n\nWhat does this code do? Answer in two sentences."
+
+
+def post(url, path, body):
+    request = urllib.request.Request(url + path, json.dumps(body).encode(),
+                                     {"Content-Type": "application/json"})
+    with urllib.request.urlopen(request, timeout=3600) as response:
+        return json.load(response)
+
+
+def greedy(url, model, ids, tokens):
+    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"]]
+
+
+def nll(url, model, ids, answer):
+    """Mean negative log-likelihood of the answer tokens after the prompt ids."""
+    entries = post(url, "/v1/completions", {
+        "model": model, "prompt": ids + answer, "max_tokens": 1, "temperature": 0,
+        "prompt_logprobs": 0})["choices"][0]["prompt_logprobs"]
+    return -sum(entries[len(ids) + i][str(t)]["logprob"] for i, t in enumerate(answer)) / len(answer)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("url")
+    parser.add_argument("action", choices=("make", "score"))
+    parser.add_argument("ref")
+    parser.add_argument("--model", default="gemma4")
+    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)
+    args = parser.parse_args()
+    if args.action == "make":
+        with open(args.code) as f:
+            prompts = PROMPTS + [CODE_PROMPT.format(code=f.read()[: args.code_chars])]
+        refs = []
+        for prompt in prompts:
+            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)
+            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:
+            json.dump(refs, f)
+        return
+    with open(args.ref) as f:
+        refs = json.load(f)
+    scores, same = [], []
+    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"]))
+        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,
+                      "identical": sum(s == len(r["answer"]) for s, r in zip(same, refs))}))
+
+
+if __name__ == "__main__":
+    main()

scripts/tune_moe.py

@@ -0,0 +1,111 @@
+"""Tune vLLM's Triton fused-MoE kernel for this GPU on prain's slot tensors.
+
+usage: python scripts/tune_moe.py OUT_DIR [FORMAT...]   (scripts/env.sh points
+                                    VLLM_TUNED_CONFIG_FOLDER at build/moe-configs)
+Times candidate configs on one prain round (16 slots of 128 routed experts, expert_map;
+up to 32 tokens route into the slots like decode, more tokens spread like prefill),
+Gemma 4 expert shapes, bf16/int8/int4 slots, and writes the fastest per batch size as vLLM
+config files for every slot count E (the best config hardly depends on E). Every config in
+the files compiled, ran and matched the default config's output on this GPU.
+"""
+
+import itertools
+import json
+import os
+import sys
+import time
+
+import torch
+from vllm.config import VllmConfig, set_current_vllm_config
+from vllm.model_executor.layers.fused_moe import override_config
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
+from vllm.model_executor.layers.fused_moe.config import (int4_w4a16_moe_quant_config,
+                                                         int8_w8a16_moe_quant_config)
+from vllm.model_executor.layers.fused_moe.fused_moe import (fused_experts, get_config_file_name,
+                                                            get_default_config)
+
+from prain.quant import GROUP, quantize
+
+H, I, TOPK, EXPERTS, SLOTS = 2816, 704, 8, 128, 16
+SIZES = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048)
+FORMATS = {"bf16": (None, None), "int8": ("int8_w8a16", None), "int4": ("int4_w4a16", [0, GROUP])}
+
+
+def candidates(tokens):
+    # ponytail: ~8 s per Triton compile on this host, so a small space (~70 compiles per
+    # format); widen it when a GPU compiles faster.
+    if tokens <= 32:
+        space = itertools.product((16,), (32, 64, 128), (64, 128, 256), (1,), (2, 4), (2,))
+    else:
+        space = itertools.product((32, 64), (64, 128), (32, 64), (8,), (4, 8), (2,))
+    for bm, bn, bk, group, warps, stages in space:
+        yield {"BLOCK_SIZE_M": bm, "BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk, "GROUP_SIZE_M": group,
+               "SPLIT_K": 1, "num_warps": warps, "num_stages": stages}  # int4 needs SPLIT_K
+
+
+def timed(fn, repeat):
+    fn()
+    torch.cuda.synchronize()
+    start = time.perf_counter()
+    for _ in range(repeat):
+        fn()
+    torch.cuda.synchronize()
+    return (time.perf_counter() - start) * 1000 / repeat
+
+
+def main(out_dir, formats):
+    os.makedirs(out_dir, exist_ok=True)
+    torch.manual_seed(0)
+    w13 = torch.randn(SLOTS, 2 * I, H, device="cuda").mul_(0.02).bfloat16()
+    w2 = torch.randn(SLOTS, H, I, device="cuda").mul_(0.02).bfloat16()
+    table = torch.full((EXPERTS + 1,), -1, dtype=torch.int32, device="cuda")
+    table[:SLOTS] = torch.arange(SLOTS, dtype=torch.int32, device="cuda")
+    for name in formats:
+        dtype, block_shape = FORMATS[name]
+        if name == "bf16":
+            a, b, quant_config = w13, w2, None
+        else:
+            (a, s13), (b, s2) = quantize(w13, name), quantize(w2, name)
+            make = int8_w8a16_moe_quant_config if name == "int8" else int4_w4a16_moe_quant_config
+            quant_config = make(w1_scale=s13, w2_scale=s2, w1_zp=None, w2_zp=None,
+                                block_shape=block_shape)
+        best = {}
+        for tokens in SIZES:
+            x = torch.randn(tokens, H, device="cuda").mul_(0.5).bfloat16()
+            pool = SLOTS if tokens <= 32 else EXPERTS
+            ids = torch.stack([torch.randperm(pool, device="cuda")[:TOPK]
+                               for _ in range(tokens)]).to(torch.int32)
+            weights = torch.softmax(torch.randn(tokens, TOPK, device="cuda"), -1)
+
+            def run():
+                return fused_experts(x, a, b, weights, ids, activation=MoEActivation.GELU_TANH,
+                                     global_num_experts=EXPERTS + 1, expert_map=table,
+                                     quant_config=quant_config)
+
+            default = get_default_config(tokens, SLOTS, I, a.shape[2], TOPK, dtype, block_shape)
+            with override_config(default):  # vLLM's own choice when no config file exists
+                default_ms = timed(run, 10 if tokens <= 32 else 3)
+                reference = run()
+            results = [(default_ms, json.dumps(default, sort_keys=True))]  # never worse than it
+            for config in candidates(tokens):
+                # catch inside: vLLM's override_config skips its restore when the body raises
+                with override_config(config):
+                    try:
+                        ms = timed(run, 10 if tokens <= 32 else 3)
+                        out = run()
+                    except Exception:  # noqa: BLE001 - configs that do not compile or fit
+                        continue
+                if (out.float() - reference.float()).norm() <= 1e-2 * reference.float().norm() + 1e-6:
+                    results.append((ms, json.dumps(config, sort_keys=True)))
+            ms, config = min(results)
+            best[str(tokens)] = json.loads(config)
+            print(f"{name} tokens={tokens:4d}: default {default_ms:8.3f} ms, best {ms:8.3f} ms "
+                  f"({len(results)} configs) {config}", flush=True)
+        for e in range(1, EXPERTS + 1):  # any slot count a placement can give a layer
+            with open(os.path.join(out_dir, get_config_file_name(e, I, dtype, block_shape)), "w") as f:
+                json.dump(best, f, indent=1)
+
+
+if __name__ == "__main__":
+    with set_current_vllm_config(VllmConfig()):
+        main(sys.argv[1], sys.argv[2:] or list(FORMATS))

tests/test_cli.py

@@ -76,3 +76,17 @@ class Cli(TmpCase):
         self.assertEqual((tuned.pack_dir, tuned.trace_path), ("/packs", trace))
         self.assertLessEqual(sum(tuned.placement(layer, 128).cache_slots for layer in range(3)),
                              int(0.2 * 1024 / 11.34) // 3 * 3)
+
+    def test_quantized_experts_check_and_tune(self):
+        hf = self.write("config.json", json.dumps(GEMMA4))
+        out, _ = self.run_cli("check", self.write("q.toml", 'pack_dir = "/packs"\nquant = "int4"\n'), hf)
+        self.assertIn("expert size 3.01 MiB\n", out)
+        trace = os.path.join(self.dir, "trace.jsonl.gz")
+        self.run_cli("synth", trace, "--forwards", "40", "--layers", "3")
+        out, _ = self.run_cli("tune", trace, hf, "--quant", "int4", "--expert-vram-gib", "0.2",
+                              "--ram-gib", "4", "--pack-dir", "/packs", "--top", "3")
+        tuned = load_config(self.write("tuned.toml", out))
+        self.assertEqual(tuned.quant, "int4")
+        slots = sum(tuned.placement(layer, 128).cache_slots for layer in range(3))
+        self.assertGreater(slots, int(0.2 * 1024 / 11.34))  # int4 experts are ~4x smaller
+        self.assertLessEqual(slots, int(0.2 * 1024 / 3.01))

tests/test_config.py

@@ -107,6 +107,23 @@ mode = "auto"
         self.assertEqual(config.stats_path, os.path.expanduser("~/s.json"))
         self.assertEqual(config.trace_path, "/abs/t.gz")
 
+    def test_miss_min_weight(self):
+        config = mc.load_config(self.write("m.toml", 'pack_dir = "/p"\n[miss]\nmin_weight = 0.7\n'))
+        self.assertEqual(config.miss.min_weight, 0.7)
+        self.assertEqual(mc.Miss().min_weight, 1.0)  # default: compute every routed expert
+        with self.assertRaisesRegex(ValueError, "min_weight must be a number in"):
+            mc.load_config(self.write("m2.toml", 'pack_dir = "/p"\n[miss]\nmin_weight = 1.5\n'))
+
+    def test_quant_key_and_expert_sizes(self):
+        self.assertEqual(mc.load_config(self.write("q.toml", 'pack_dir = "/p"\nquant = "int4"\n')).quant, "int4")
+        self.assertEqual(mc.load_config(self.write("b.toml", 'pack_dir = "/p"\n')).quant, "bf16")
+        with self.assertRaisesRegex(ValueError, "quant must be one of"):
+            mc.load_config(self.write("bad.toml", 'pack_dir = "/p"\nquant = "int3"\n'))
+        h, i = 2816, 704
+        self.assertEqual(mc.expert_bytes(h, i), 3 * h * i * 2)
+        self.assertEqual(mc.expert_bytes(h, i, "int8"), 3 * h * i + 4 * (2 * i + h))  # fp32 row scales
+        self.assertEqual(mc.expert_bytes(h, i, "int4"), 3 * h * i // 2 + 2 * 3 * h * i // 64)
+
     def test_serve_and_draft_tables(self):
         config = mc.load_config(self.write("s.toml", 'pack_dir = "/p"\n[serve]\nmax_requests = 4\n'
                                                      'kv_cache_gib = 1.5\n[draft]\nfull_requests = 2\n'

tests/test_gemv.py

@@ -0,0 +1,55 @@
+import unittest
+
+import torch
+import torch.nn.functional as F
+
+from prain import gemv
+
+GPU = torch.cuda.is_available() and gemv.fast_device(0)
+
+
+def rel(a, b):
+    return ((a.float() - b.float()).norm() / b.float().norm()).item()
+
+
+@unittest.skipUnless(GPU, "needs a gfx10 GPU")
+class TestGemv(unittest.TestCase):
+    def test_up_to_four_rows_match_a_float32_linear(self):
+        torch.manual_seed(0)
+        for n, k in ((1000, 2816), (2816, 2112), (130, 70)):
+            w = torch.randn(n, k, dtype=torch.bfloat16, device="cuda") * 0.02
+            for rows in (1, 2, 3, 4):
+                x = torch.randn(rows, k, dtype=torch.bfloat16, device="cuda")
+                y = gemv.linear(x, w)
+                self.assertEqual((y.shape, y.dtype), ((rows, n), torch.bfloat16))
+                self.assertLess(rel(y, F.linear(x.float(), w.float())), 1e-2, (n, k, rows))
+        x = torch.randn(2, 1, 70, dtype=torch.float16, device="cuda")  # leading dims, fp16
+        w = torch.randn(130, 70, dtype=torch.float16, device="cuda")
+        self.assertEqual(gemv.linear(x, w).shape, (2, 1, 130))
+        self.assertLess(rel(gemv.linear(x, w), F.linear(x.float(), w.float())), 1e-2)
+
+    def test_more_rows_bias_or_strided_weights_use_f_linear(self):
+        w = torch.randn(64, 96, dtype=torch.bfloat16, device="cuda")
+        bias = torch.randn(64, dtype=torch.bfloat16, device="cuda")
+        x = torch.randn(5, 96, dtype=torch.bfloat16, device="cuda")
+        self.assertTrue(torch.equal(gemv.linear(x, w), F.linear(x, w)))
+        self.assertTrue(torch.equal(gemv.linear(x[:1], w, bias), F.linear(x[:1], w, bias)))
+        strided = torch.randn(96, 64, dtype=torch.bfloat16, device="cuda").t()
+        self.assertTrue(torch.equal(gemv.linear(x[:1], strided), F.linear(x[:1], strided)))
+
+    def test_wrapped_dispatch_sends_only_decode_shapes_to_the_kernel(self):
+        calls = []
+
+        def fallback(layer, x, weight, bias=None):
+            calls.append(x.shape[0])
+            return F.linear(x, weight, bias)
+
+        gemm = gemv.wrap(lambda *args, **kwargs: fallback)("auto")
+        w = torch.randn(64, 96, dtype=torch.bfloat16, device="cuda")
+        gemm(None, torch.randn(1, 96, dtype=torch.bfloat16, device="cuda"), w)
+        gemm(None, torch.randn(9, 96, dtype=torch.bfloat16, device="cuda"), w)
+        self.assertEqual(calls, [9])
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_kernel_match.py

@@ -31,3 +31,64 @@ class KernelMatch(unittest.TestCase):
                                 global_num_experts=experts).float().cpu()
         cpu = cpu_experts(x, topk_weights, topk_ids, set(range(experts)), lambda e: (w13[e], w2[e]))
         torch.testing.assert_close(cpu, gpu, rtol=3e-2, atol=3e-2)
+
+    def test_a_round_computes_only_the_experts_its_slot_table_maps(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.tiered import cpu_experts, slot_table
+
+        torch.manual_seed(0)
+        tokens, hidden, inter, experts = 64, 256, 128, 6
+        x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
+        w13 = torch.randn(experts, 2 * inter, hidden, dtype=torch.bfloat16) * 0.1
+        w2 = torch.randn(experts, hidden, inter, dtype=torch.bfloat16) * 0.1
+        topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int()
+        topk_ids[0, 1] = experts  # a padding route
+        topk_weights = torch.rand(tokens, 2)
+        table = slot_table({1: 0, 4: 1}, experts, "cuda")
+        with set_current_vllm_config(VllmConfig()):
+            gpu = fused_experts(x.cuda(), w13[[1, 4]].cuda(), w2[[1, 4]].cuda(), topk_weights.cuda(),
+                                topk_ids.cuda(), activation=MoEActivation.GELU_TANH,
+                                global_num_experts=table.numel(), expert_map=table).float().cpu()
+        cpu = cpu_experts(x, topk_weights, topk_ids, {1, 4}, lambda e: (w13[e], w2[e]))
+        torch.testing.assert_close(cpu, gpu, rtol=3e-2, atol=3e-2)
+
+    def test_quantized_slots_compute_what_the_cpu_path_computes(self):
+        """Pack layout -> slot layout on the GPU -> vLLM's int8/int4 kernel for a round that
+        holds experts 1 and 3, against the CPU kernels on the pack layout and against the
+        dequantized weights."""
+        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.quant import (cpu_matmul, dequantize, moe_quant_config, quantize, to_cpu_int4,
+                                 to_gpu_int4)
+        from prain.tiered import cpu_experts, slot_table
+
+        torch.manual_seed(0)
+        tokens, hidden, inter, experts = 8, 256, 128, 4
+        x = torch.randn(tokens, hidden, dtype=torch.bfloat16)
+        w13 = torch.randn(experts, 2 * inter, hidden) * 0.1
+        w2 = torch.randn(experts, hidden, inter) * 0.1
+        topk_ids = torch.stack([torch.randperm(experts)[:2] for _ in range(tokens)]).int()
+        topk_weights = torch.rand(tokens, 2)
+        table = slot_table({1: 0, 3: 1}, experts, "cuda")
+        for quant in ("int8", "int4"):
+            parts = [quantize(w, quant) for w in (w13, w2)]  # [q, scale], gpu layout
+            pack = [[to_cpu_int4(q) if quant == "int4" else q, s] for q, s in parts]
+            slots = [[to_gpu_int4(q[[1, 3]].cuda()) if quant == "int4" else q[[1, 3]].cuda(),
+                      s[[1, 3]].cuda()] for q, s in pack]
+            with set_current_vllm_config(VllmConfig()):
+                gpu = fused_experts(x.cuda(), slots[0][0], slots[1][0], topk_weights.cuda(),
+                                    topk_ids.cuda(), activation=MoEActivation.GELU_TANH,
+                                    global_num_experts=table.numel(), expert_map=table,
+                                    quant_config=moe_quant_config(quant, slots[0][1], slots[1][1]))
+            cpu = cpu_experts(x, topk_weights, topk_ids, {1, 3}, lambda e: tuple(
+                cpu_matmul(quant, q[e], s[e]) for q, s in pack))
+            exact = cpu_experts(x, topk_weights, topk_ids, {1, 3}, lambda e: tuple(
+                dequantize(q, s, quant)[e] for q, s in parts))
+            self.assertGreater(exact.norm().item(), 1.0)
+            torch.testing.assert_close(cpu, exact, rtol=3e-2, atol=3e-2, msg=quant)
+            torch.testing.assert_close(gpu.float().cpu(), exact, rtol=3e-2, atol=3e-2, msg=quant)

tests/test_quant.py

@@ -0,0 +1,49 @@
+import unittest
+
+try:
+    import torch
+except ImportError:  # the planning host has no torch
+    torch = None
+
+if torch is not None:
+    from prain.quant import (GROUP, cpu_matmul, dequantize, layout, quantize, to_cpu_int4,
+                             to_gpu_int4)
+
+
+@unittest.skipIf(torch is None, "needs torch")
+class Quant(unittest.TestCase):
+    def setUp(self):
+        torch.manual_seed(0)
+        self.w = torch.randn(64, 256) * 0.02
+
+    def test_layouts_match_what_quantize_returns(self):
+        for quant in ("int8", "int4"):
+            got = [(tuple(t.shape), t.dtype) for t in quantize(self.w, quant)]
+            self.assertEqual(got, layout(quant, 64, 256))
+
+    def test_round_trip_error(self):
+        for quant, bound in (("int8", 0.01), ("int4", 0.11)):
+            err = (dequantize(*quantize(self.w, quant), quant, torch.float32) - self.w).norm()
+            self.assertLess(err / self.w.norm(), bound)
+
+    def test_int4_clipping_beats_plain_round_to_nearest(self):
+        g = self.w.unflatten(-1, (-1, GROUP))
+        scale = (g.abs().amax(-1, keepdim=True) / 7).bfloat16().float()
+        plain = ((g / scale).round() * scale - g).norm()
+        clipped = (dequantize(*quantize(self.w, "int4"), "int4", torch.float32) - self.w).norm()
+        self.assertLess(clipped, plain)
+
+    def test_cpu_layout_is_torch_avx2_int4pack_and_inverts(self):
+        q, _ = quantize(self.w, "int4")
+        nibbles = torch.stack((q & 15, q >> 4), -1).flatten(-2).to(torch.int32)
+        self.assertTrue(torch.equal(to_cpu_int4(q), torch._convert_weight_to_int4pack_for_cpu(nibbles, 1)))
+        self.assertTrue(torch.equal(to_gpu_int4(to_cpu_int4(q)), q))
+
+    def test_cpu_matmul_matches_dequantized_weights(self):
+        x = torch.randn(3, 256, dtype=torch.bfloat16)
+        for quant in ("int8", "int4"):
+            q, scale = quantize(self.w, quant)
+            ref = x.float() @ dequantize(q, scale, quant, torch.float32).t()
+            stored = to_cpu_int4(q) if quant == "int4" else q
+            got = cpu_matmul(quant, stored, scale)(x).float()
+            self.assertLess(((got - ref).norm() / ref.norm()).item(), 0.01, quant)

tests/test_tiered.py

@@ -9,7 +9,7 @@ from prain.config import Miss, Policy
 from prain.policy import SlotCache
 
 if torch is not None:
-    from prain.tiered import (available_experts, cpu_experts, degrade, remap, route_stats,
+    from prain.tiered import (available_experts, cpu_experts, degrade, route_stats, slot_table,
                                 tiered_forward)
 
 HIDDEN, INTER, EXPERTS = 6, 4, 10
@@ -52,8 +52,9 @@ class TieredForward(unittest.TestCase):
             loads.append(expert)
             slots13[slot], slots2[slot] = self.w13[expert], self.w2[expert]
 
-        def kernel(x, w, ids, used):
+        def kernel(x, w, ids, table, used):
-            return naive_moe(x, slots13, slots2, w, ids)
+            self.assertTrue(bool((ids >= 0).all()))  # padding id -1 arrives as num_experts
+            return naive_moe(x, slots13, slots2, w, table[ids.long()])
 
         out = tiered_forward(self.x, self.topk_weights, self.topk_ids, cache, load, kernel,
                              miss=miss, cpu_weights=lambda e: (self.w13[e], self.w2[e]),
@@ -80,6 +81,33 @@ class TieredForward(unittest.TestCase):
         self.assertEqual(cpu_experts(self.x, self.topk_weights, self.topk_ids, set(),
                                      None).abs().sum().item(), 0.0)
 
+    def test_miss_min_weight_skips_missing_experts_for_every_token(self):
+        """[miss] min_weight < 1: every token skips experts that are not in VRAM while it keeps
+        that share of its router weight; experts some token copies anyway serve all tokens."""
+        everyone = torch.ones(6, dtype=torch.bool)
+        resident = torch.zeros(EXPERTS, dtype=torch.bool)
+        resident[[0, 1, 2]] = True
+        _, first = degrade(self.topk_weights, self.topk_ids, everyone, resident, 0.5)
+        free = resident.clone()
+        free[first[first >= 0].long()] = True
+        w, i = degrade(self.topk_weights, self.topk_ids, everyone, free, 0.5)
+        out, loads, _ = self.run_tiered(pinned=[0, 1, 2], cache_slots=3, miss=Miss(min_weight=0.5))
+        torch.testing.assert_close(out, naive_moe(self.x, self.w13, self.w2, w, i))
+        self.assertEqual(sorted(loads), sorted(set(i[i >= 0].tolist()) - {0, 1, 2}))
+        self.assertLess(len(loads), 7)  # computing everything copies all 7 others
+        self.assertEqual(self.seen_experts, list(range(EXPERTS)))  # the trace sees the routing
+        out, loads, _ = self.run_tiered(pinned=[0, 1, 2], cache_slots=3, miss=Miss(min_weight=1.0))
+        torch.testing.assert_close(out, self.expected)
+        self.assertEqual(len(loads), 7)
+
+    def test_cpu_experts_take_matmul_functions(self):
+        """Quantized experts hand the CPU path x -> x @ w.T functions instead of weights."""
+        def matmuls(e):
+            return (lambda v: v @ self.w13[e].t()), (lambda v: v @ self.w2[e].t())
+
+        out = cpu_experts(self.x, self.topk_weights, self.topk_ids, set(range(EXPERTS)), matmuls)
+        torch.testing.assert_close(out, self.expected.float())
+
     def test_route_stats(self):
         stats = route_stats(torch.tensor([[0.5, 0.25], [1.0, 0.0]]),
                             torch.tensor([[2, 1], [2, -1]]), 4, recency_half_life=1.0)
@@ -89,10 +117,10 @@ class TieredForward(unittest.TestCase):
         self.assertAlmostEqual(stats[2][2], 0.5 * 0.5 + 1.0)  # the first token is one half-life old
         self.assertAlmostEqual(stats[1][2], 0.125)
 
-    def test_remap_zeroes_routes_outside_the_round(self):
+    def test_slot_table_maps_the_round_and_skips_the_rest(self):
-        weights, slots = remap(torch.ones(1, 3), torch.tensor([[2, 7, -1]]), {7: 4}, 10)
+        table = slot_table({7: 4, 2: 0}, 10, "cpu")
-        self.assertEqual(weights.tolist(), [[0.0, 1.0, 0.0]])
+        self.assertEqual(table.tolist(), [-1, -1, 0, -1, -1, -1, -1, 4, -1, -1, -1])
-        self.assertEqual(slots.tolist(), [[0, 4, 0]])
+        self.assertEqual(table.dtype, torch.int32)  # one extra row: padding id 10 -> -1
 
     def test_degrade_drops_unavailable_routes_of_draft_tokens_only(self):
         weights = torch.tensor([[0.5, 0.3, 0.2], [0.5, 0.3, 0.2]])
@@ -115,6 +143,8 @@ class TieredForward(unittest.TestCase):
         available[[0, 1, 2, 4, 5, 6, 7, 8]] = True
         self.assertTrue(torch.equal(
             available_experts(self.topk_ids, draft, {0}, EXPERTS), available))
+        _, kept = degrade(self.topk_weights, self.topk_ids, draft, available, 0.5)
+        available[kept[kept >= 0].long()] = True  # experts some row copies anyway serve all rows
         expected = naive_moe(self.x, self.w13, self.w2,
                              *degrade(self.topk_weights, self.topk_ids, draft, available, 0.5))
         torch.testing.assert_close(out, expected)