| """Native vLLM draft_model speculative smoke + benchmark. |
| Target = Qwen3-VL bf16, draft = AWQ int4 self-twin (vision embeds shared from target). |
| --mode spec : speculative engine, records token ids + timings + spec metrics |
| --mode vanilla : same engine settings without speculation (identity + speed reference) |
| Outputs one jsonl per run; identity check compares the two files offline. |
| """ |
| import argparse, io, json, time |
| from PIL import Image |
| import pyarrow.parquet as pq |
|
|
|
|
| def load_samples(parquet, pids_from, n, img_max_side): |
| pids = [] |
| for line in open(pids_from): |
| d = json.loads(line) |
| if "pid" in d: |
| pids.append(d["pid"]) |
| if len(pids) >= n: |
| break |
| tbl = pq.read_table(parquet, columns=["pid", "query", "decoded_image"]) |
| rowmap = {r["pid"]: r for r in tbl.to_pylist()} |
| out = [] |
| for pid in pids: |
| row = rowmap[pid] |
| img = Image.open(io.BytesIO(row["decoded_image"]["bytes"])).convert("RGB") |
| if max(img.size) > img_max_side: |
| sc = img_max_side / max(img.size) |
| img = img.resize((int(img.width * sc), int(img.height * sc))) |
| out.append((pid, row["query"], img)) |
| return out |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--mode", choices=["spec", "vanilla"], required=True) |
| ap.add_argument("--model", required=True) |
| ap.add_argument("--draft", default="") |
| ap.add_argument("--gamma", type=int, default=6) |
| ap.add_argument("--parquet", required=True) |
| ap.add_argument("--pids-from", required=True) |
| ap.add_argument("--n", type=int, default=3) |
| ap.add_argument("--max-tokens", type=int, default=256) |
| ap.add_argument("--gpu-mem-util", type=float, default=0.35) |
| ap.add_argument("--max-model-len", type=int, default=6144) |
| ap.add_argument("--img-max-side", type=int, default=1024) |
| ap.add_argument("--min-tokens", type=int, default=0) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--enforce-eager", action="store_true") |
| ap.add_argument("--no-image", action="store_true") |
| ap.add_argument("--tp", type=int, default=1) |
| args = ap.parse_args() |
|
|
| samples = load_samples(args.parquet, args.pids_from, args.n, args.img_max_side) |
| from stage0_gonogo import build_inputs |
| from transformers import AutoProcessor |
| proc = AutoProcessor.from_pretrained(args.model) |
| prompts = [] |
| for pid, query, img in samples: |
| if args.no_image: |
| img = None |
| enc = build_inputs(proc, query, img) |
| prompts.append((pid, enc["input_ids"][0].tolist(), img)) |
| print(f"[init] mode={args.mode} n={len(prompts)}", flush=True) |
|
|
| from vllm import LLM, SamplingParams |
| kw = {} |
| if args.mode == "spec": |
| kw["speculative_config"] = { |
| "method": "draft_model", |
| "model": args.draft, |
| "num_speculative_tokens": args.gamma, |
| "draft_tensor_parallel_size": args.tp, |
| } |
| llm = LLM(model=args.model, tensor_parallel_size=args.tp, |
| disable_custom_all_reduce=True, |
| gpu_memory_utilization=args.gpu_mem_util, |
| max_model_len=args.max_model_len, enable_prefix_caching=False, |
| enforce_eager=args.enforce_eager, disable_log_stats=False, |
| max_num_batched_tokens=2048, max_num_seqs=4, |
| limit_mm_per_prompt={"image": 1}, **kw) |
| sp = SamplingParams(temperature=0, max_tokens=args.max_tokens, |
| min_tokens=args.min_tokens) |
|
|
| def make_input(ids, img): |
| if img is None: |
| return {"prompt_token_ids": ids} |
| return {"prompt_token_ids": ids, "multi_modal_data": {"image": img}} |
|
|
| |
| llm.generate([make_input(*prompts[0][1:])], |
| SamplingParams(temperature=0, max_tokens=8), use_tqdm=False) |
|
|
| fout = open(args.out, "w") |
| for pid, ids, img in prompts: |
| t0 = time.time() |
| out = llm.generate([make_input(ids, img)], sp, use_tqdm=False)[0] |
| dt = time.time() - t0 |
| toks = list(out.outputs[0].token_ids) |
| rec = dict(pid=pid, mode=args.mode, gamma=args.gamma if args.mode == "spec" else 0, |
| in_len=len(ids), gen_len=len(toks), t=round(dt, 3), |
| tok_s=round(len(toks) / dt, 2), gen_ids=toks) |
| fout.write(json.dumps(rec) + "\n") |
| fout.flush() |
| print(f"[{pid}] gen={len(toks)} t={dt:.2f}s tok_s={len(toks)/dt:.1f}", flush=True) |
| fout.close() |
|
|
| try: |
| for m in llm.get_metrics(): |
| if "spec" in m.name: |
| val = getattr(m, "value", None) |
| if val is None: |
| val = getattr(m, "values", None) |
| print(f"[metric] {m.name} = {val}", flush=True) |
| except Exception as e: |
| print(f"[metric] unavailable: {e}", flush=True) |
| print("[done]", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
