Update 01_benchmark.py

01_benchmark.py

@@ -0,0 +1,129 @@
+#!/usr/bin/env python3
+"""SimpleTool multi-head parallel decode — vLLM, v1/v2, external prompts
+python 01_benchmark.py --version v2                    # v2 default model
+python 01_benchmark.py --version v1                    # v1 default model
+python 01_benchmark.py --version v2 --n-args 3         # fixed three arg heads 
+python 01_benchmark.py --version v1 --model /my/model  # customed model path
+"""
+import argparse, json, time, os
+from pathlib import Path
+
+DIR = Path("./prompts")
+HEADS = [("function","<function>","</function>")] + [(f"arg{i}",f"<arg{i}>",f"</arg{i}>") for i in range(1,7)]
+STOPS = ["</function>"] + [f"</arg{i}>" for i in range(1,7)] + ["</content>","<|null|>","<|im_end|>"]
+MODELS = {"v1":"./models/RT-Qwen3-4B-AWQ", "v2":"./models/RT-Qwen3-4B-AWQ-v2"}
+
+def load_scenarios():
+    scs = json.loads((DIR/"scenarios.json").read_text())
+    for sc in scs:
+        sc["tools"] = (DIR/sc["tools_file"]).read_text().strip()
+    return scs
+
+def max_tool_params(tools_str):
+    m = 0
+    for l in tools_str.strip().split("\n"):
+        try: m = max(m, len(json.loads(l)["function"]["parameters"]["properties"]))
+        except: pass
+    return m
+
+def build_prompt(sc, ver):
+    t = sc["tools"]
+    if ver == "v1":
+        v1sys = (DIR/"v1_system.txt").read_text()
+        return (f"<|im_start|>system\n{v1sys}\n## Available Tools:\n\n{t}<|im_end|>\n"
+                f"<|im_start|>user\nenvironment: []\nhistory: {sc['history']}\n\n{sc['system']}\n\n{sc['query']}<|im_end|>\n"
+                f"<|im_start|>assistant\n")
+    return (f"<|im_start|>system\n{sc['system']}\n\n## Available Tools:\n\n{t}<|im_end|>\n"
+            f"<|im_start|>user\nhistory: {sc['history']}\n\n{sc['query']}<|im_end|>\n"
+            f"<|im_start|>assistant\n")
+
+def clean(t):
+    t = t.strip()
+    return "<|null|>" if "<|null|>" in t or t == "" else t.split("</")[0].strip()
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", default=None)
+    ap.add_argument("--version", default="v2", choices=["v1","v2"])
+    ap.add_argument("--n-args", default="auto")
+    ap.add_argument("--gpu", type=int, default=0)
+    ap.add_argument("--max-model-len", type=int, default=4096)
+    a = ap.parse_args()
+    a.model = a.model or MODELS[a.version]
+    os.environ["CUDA_VISIBLE_DEVICES"] = str(a.gpu)
+    from vllm import LLM, SamplingParams
+
+    SC = load_scenarios()
+    print(f"\n{'='*60}\n  {a.version} | {a.model}\n{'='*60}")
+    llm = LLM(model=a.model, trust_remote_code=True, dtype="auto", gpu_memory_utilization=0.80,
+              max_model_len=a.max_model_len, max_num_seqs=8, enable_prefix_caching=True)
+    sp = SamplingParams(temperature=0.0, max_tokens=128, stop=STOPS, include_stop_str_in_output=True)
+    na = [min(max_tool_params(s["tools"]),6) if a.n_args=="auto" else max(1,min(6,int(a.n_args))) for s in SC]
+    for s,n in zip(SC,na): print(f"  {s['name']:<35} heads={1+n}")
+
+    def run(sc, n):
+        hd = HEADS[:1+n]; base = build_prompt(sc, a.version)
+        t0 = time.perf_counter()
+        outs = llm.generate([base+op for _,op,_ in hd], sp)
+        ms = (time.perf_counter()-t0)*1000
+        raw, toks, full = {}, {}, {}
+        for j,(nm,_,_) in enumerate(hd):
+            if j<len(outs) and outs[j].outputs:
+                o = outs[j].outputs[0]; full[nm]=o.text; raw[nm]=clean(o.text); toks[nm]=len(o.token_ids)
+            else: raw[nm],toks[nm],full[nm] = "<|null|>",0,""
+        return raw, toks, full, ms, hd
+
+    # Cold
+    print(f"\n{'='*60}\n  COLD START\n{'='*60}")
+    cold = []
+    for i,s in enumerate(SC): _,_,_,ms,_=run(s,na[i]); cold.append(ms); print(f"  {s['name']:<35} {ms:7.1f}ms")
+
+    # Hot x3
+    print(f"\n{'='*60}\n  HOT WARMUP (3 rounds)\n{'='*60}")
+    hot = [[] for _ in SC]
+    for r in range(3):
+        for i,s in enumerate(SC): _,_,_,ms,_=run(s,na[i]); hot[i].append(ms)
+        print(f"  Round {r+1}: "+"  ".join(f"{hot[j][-1]:6.1f}ms" for j in range(len(SC))))
+
+    # Test
+    print(f"\n{'='*60}\n  PARALLEL TEST ({a.version})\n{'='*60}\n")
+    res = []
+    for i,s in enumerate(SC):
+        raw,toks,full,ms,hd = run(s,na[i]); mt=max(toks.values()) if toks else 0
+        ok = raw.get("function","") == s["expected"]; res.append((s,raw,toks,full,ms,mt,hd,ok))
+        print(f"─── {s['name']} ───\n{'PASS' if ok else 'FAIL'}  {s['desc']}")
+        for nm,_,_ in hd:
+            v,tc = raw.get(nm,""),toks.get(nm,0); d=v if len(v)<=43 else v[:43]+"…"
+            st = ("OK" if ok else f"WRONG({v})") if nm=="function" else ("NULL" if v=="<|null|>" else "FILL")
+            print(f"  {nm:<10} {d:<45} {tc:<4} {st}")
+        print(f"  e2e={ms:.1f}ms  max_tok={mt}\n")
+
+    # Summary
+    N=len(res); np_=sum(r[7] for r in res); ae=sum(r[4] for r in res)/N; amt=sum(r[5] for r in res)/N
+    print(f"{'='*60}\n  SUMMARY ({a.version})\n{'='*60}")
+    print(f"  Accuracy       : {np_}/{N}\n  Cold start avg : {sum(cold)/N:.1f}ms\n  Hot prefill avg: {sum(sum(h) for h in hot)/sum(len(h) for h in hot):.1f}ms")
+    print(f"  E2E avg (hot)  : {ae:.1f}ms\n  Max head tokens: {amt:.1f} avg\n  E2E / max_tok  : {ae/amt:.1f}ms/tok (decode bottleneck)\n")
+    print(f"  {'Scenario':<35} {'Cold':>7} {'Hot':>7} {'E2E':>7} {'MaxTk':>6} {'ms/tk':>6}\n  {'─'*70}")
+    for i,(s,_,_,_,ms,mt,_,_) in enumerate(res):
+        print(f"  {s['name']:<35} {cold[i]:6.1f}  {sum(hot[i])/3:6.1f}  {ms:6.1f}  {mt:>5}  {ms/mt if mt else 0:5.1f}")
+
+    # Example dump
+    s,raw,toks,full,ms,mt,hd,ok = res[0]; base=build_prompt(s,a.version)
+    print(f"\n{'='*60}\n  EXAMPLE ({a.version}): {s['name']}\n{'='*60}")
+    print(f"\n┌─ Shared Prefix ({len(base)} chars) ────────────────────")
+    for ln in base.split("\n"): print(f"│ {ln}")
+    print(f"└──────────────────────────────────────────────────")
+    print(f"\n┌─ Per-Head Trigger Tokens ─────────────────────────")
+    for nm,op,_ in hd: print(f"│  {nm:<10} → {op}")
+    print(f"└──────────────────────────────────────────────────")
+    print(f"\n┌─ Decode Output (all tokens, incl. stop) ──────────")
+    for nm,op,_ in hd: print(f"│  {nm:<10} [{toks.get(nm,0):>2} tok]  {op}{full.get(nm,'')}")
+    print(f"└──────────────────────────────────────────────────")
+    print(f"\n  Reconstructed multi-head response:")
+    for nm,op,cl in hd:
+        if raw.get(nm,"")=="<|null|>": print(f"    {op}<|null|>")
+        else:
+            ft=full.get(nm,""); print(f"    {op}{ft}" if any(ft.rstrip().endswith(x) for x in STOPS) else f"    {op}{ft}{cl}")
+    print()
+
+if __name__ == "__main__": main()