prain: int8/int4 experts quantized once on the GPU, rounds through vLLM's expert_map, [miss] min_weight skips experts not in VRAM; scripts: tune_moe, quality (answer NLL on identical token ids), bench --context, measure --quant

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):
-    text = f"[{uuid.uuid4().hex[:8]}] {prompt}"
+def run(base_url, model, prompt, tokens, context=""):
+    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
-    host = torch.empty(expert_bytes // 2, dtype=torch.bfloat16).pin_memory()
+    expert_bytes = size_of(HIDDEN, INTER, args.quant)
+    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))