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}))