fabolous005/prain
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

prain: int4 MoE on RDNA2: GEMV-style decode kernel, fp16 rounds, staged expert copies

Browse Source 
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.
This commit is contained in:
fabolous005
2026-10-29 14:20:23 +01:00
parent 0f3911afb5
commit d1ccc5c294
5 changed files with 380 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
scripts/quality.py
+19 -8
View File
@@ -33,10 +33,13 @@ def post(url, path, body):
def greedy(url, model, ids, tokens):
"""(text, token ids, their logprobs) of a greedy answer."""
choice = post(url, "/v1/completions", {
"model": model, "prompt": ids, "max_tokens": tokens, "temperature": 0, "logprobs": 1,
"return_tokens_as_token_ids": True})["choices"][0]
return choice["text"], [int(t.split(":")[1]) for t in choice["logprobs"]["tokens"]]
logprobs = choice["logprobs"]
return (choice["text"], [int(t.split(":")[1]) for t in logprobs["tokens"]],
logprobs["token_logprobs"])
def nll(url, model, ids, answer):
@@ -56,6 +59,9 @@ def main():
parser.add_argument("--tokens", type=int, default=48, help="answer tokens")
parser.add_argument("--code", default="prain/policy.py", help="file for the code question")
parser.add_argument("--code-chars", type=int, default=900)
parser.add_argument("--greedy-only", action="store_true",
help="score: skip the nll (prompt logprobs need ~0.5 GiB of free VRAM "
"for the logits); compare greedy answers only")
args = parser.parse_args()
if args.action == "make":
with open(args.code) as f:
@@ -65,7 +71,7 @@ def main():
ids = post(args.url, "/tokenize", {"model": args.model, "add_generation_prompt": True,
"messages": [{"role": "user", "content": prompt}]})
ids = ids["tokens"]
text, answer = greedy(args.url, args.model, ids, args.tokens)
text, answer, _ = greedy(args.url, args.model, ids, args.tokens)
refs.append({"prompt": prompt, "ids": ids, "answer": answer, "text": text})
print(json.dumps({"prompt": prompt[:60], "answer": text[:120]}), flush=True)
with open(args.ref, "w") as f:
@@ -73,15 +79,20 @@ def main():
return
with open(args.ref) as f:
refs = json.load(f)
scores, same = [], []
scores, same, own = [], [], []
for ref in refs:
scores.append(nll(args.url, args.model, ref["ids"], ref["answer"]))
text, answer = greedy(args.url, args.model, ref["ids"], len(ref["answer"]))
if not args.greedy_only:
scores.append(nll(args.url, args.model, ref["ids"], ref["answer"]))
text, answer, logprobs = greedy(args.url, args.model, ref["ids"], len(ref["answer"]))
own.append(sum(logprobs) / len(logprobs))
same.append(next((i for i, (a, b) in enumerate(zip(answer, ref["answer"])) if a != b),
min(len(answer), len(ref["answer"]))))
print(json.dumps({"prompt": ref["prompt"][:40], "nll": round(scores[-1], 3),
"same_tokens": same[-1], "answer": text[:100]}), flush=True)
print(json.dumps({"mean_nll": round(sum(scores) / len(scores), 3), "same_tokens": same,
print(json.dumps({"prompt": ref["prompt"][:40],
"nll": round(scores[-1], 3) if scores else None,
"own_lp": round(own[-1], 3), "same_tokens": same[-1],
"answer": text[:100]}), flush=True)
print(json.dumps({"mean_nll": round(sum(scores) / len(scores), 3) if scores else None,
"mean_own_lp": round(sum(own) / len(own), 3), "same_tokens": same,
"identical": sum(s == len(r["answer"]) for s, r in zip(same, refs))}))