Files
fabolous005 d1ccc5c294 prain: int4 MoE on RDNA2: GEMV-style decode kernel, fp16 rounds, staged expert copies
Measured on the live 16k agent setup (A/B with runtime flags, greedy answers vs reference):
decode 3.49 -> 5.24 tok/s, first token 4.3 -> 3.3 s, 2.3k-token prompt 27.5 -> 20.6 s,
same answers. Dropped after measuring: route stats on the CPU, a Triton router GEMV,
pregate prefetch from disk (all within noise or slower).

The GEMV keeps one config instead of autotuning: Triton benchmarks with a 256 MiB buffer
that torch keeps cached, which left 8 MiB of VRAM and made lazy kernel loads fail
(hipErrorNoBinaryForGpu). prain also returns torch cached blocks after materializing and
whenever the driver has < 128 MiB free. quality.py score --greedy-only skips prompt logprobs,
which OOM the 16k config.
2026-10-29 14:20:23 +01:00

101 lines
5.1 KiB
Python

"""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):
"""(text, token ids, their logprobs) of a greedy answer."""
choice = post(url, "/v1/completions", {
"model": model, "prompt": ids, "max_tokens": tokens, "temperature": 0, "logprobs": 1,
"return_tokens_as_token_ids": True})["choices"][0]
logprobs = choice["logprobs"]
return (choice["text"], [int(t.split(":")[1]) for t in logprobs["tokens"]],
logprobs["token_logprobs"])
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)
parser.add_argument("--greedy-only", action="store_true",
help="score: skip the nll (prompt logprobs need ~0.5 GiB of free VRAM "
"for the logits); compare greedy answers only")
args = parser.parse_args()
if args.action == "make":
with open(args.code) as f:
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, own = [], [], []
for ref in refs:
if not args.greedy_only:
scores.append(nll(args.url, args.model, ref["ids"], ref["answer"]))
text, answer, logprobs = greedy(args.url, args.model, ref["ids"], len(ref["answer"]))
own.append(sum(logprobs) / len(logprobs))
same.append(next((i for i, (a, b) in enumerate(zip(answer, ref["answer"])) if a != b),
min(len(answer), len(ref["answer"]))))
print(json.dumps({"prompt": ref["prompt"][:40],
"nll": round(scores[-1], 3) if scores else None,
"own_lp": round(own[-1], 3), "same_tokens": same[-1],
"answer": text[:100]}), flush=True)
print(json.dumps({"mean_nll": round(sum(scores) / len(scores), 3) if scores else None,
"mean_own_lp": round(sum(own) / len(own), 3), "same_tokens": same,
"identical": sum(s == len(r["answer"]) for s, r in zip(same, refs))}))
if __name__ == "__main__":
main()