Spaces:
Running
Running
| """Baseline: frozen model on the held-out KernelBench split, multi-turn, via the kernel_env server. | |
| python baseline_eval.py --env-url http://localhost:8000 --model Qwen/Qwen2.5-Coder-7B-Instruct --backend vllm | |
| Each eval task is one episode of up to `max_turns` submissions; after each submission the model | |
| sees the env feedback (stage, errors, timings) and may refine. Reports compile %, correct %, | |
| mean/median speedup (best correct submission per episode) and a stage histogram. | |
| Also holds the helpers grpo_kernel.py shares (code extraction, feedback message, metrics). | |
| """ | |
| import argparse | |
| import json | |
| import os | |
| import re | |
| import shlex | |
| import statistics | |
| import subprocess | |
| import threading | |
| import time | |
| import urllib.request | |
| from collections import Counter | |
| from concurrent.futures import ThreadPoolExecutor | |
| from kernel_env import KernelAction, KernelEnv | |
| from kernel_env.server.tasks import load_tasks | |
| STAGE_ORDER = ["format", "compile", "runtime", "timeout", "correctness", "timed"] | |
| _FENCE = re.compile(r"```(\w*)[^\n]*\n(.*?)```", re.DOTALL) | |
| def extract_code(text: str) -> str: | |
| """Last ```python block, else last fenced block, else the raw text (prompt asks for bare source).""" | |
| blocks = _FENCE.findall(text) | |
| py = [body for lang, body in blocks if lang.lower() in ("python", "py")] | |
| if py or blocks: | |
| return (py or [body for _, body in blocks])[-1].strip() | |
| # unterminated fence (generation cut off): drop the opening line | |
| if "```" in text: | |
| return text.split("```", 1)[1].split("\n", 1)[-1].strip() | |
| return text.strip() | |
| def feedback_message(obs, truncated: bool = False) -> str: | |
| """User turn shown to the model after a submission.""" | |
| msg = f"Submission {obs.turn}/{obs.max_turns}: stage={obs.stage}." | |
| if truncated: | |
| msg = ("Your previous reply hit the token limit and was cut off before the module was complete. " | |
| "Reply with the module source only, no commentary.\n" + msg) | |
| if obs.correct and obs.speedup is not None: | |
| msg += f" Correct, speedup {obs.speedup:.2f}x (ref {obs.t_ref_ms:.3f} ms, yours {obs.t_new_ms:.3f} ms)." | |
| if obs.feedback: | |
| msg += f"\n{obs.feedback}" | |
| if obs.turn < obs.max_turns: | |
| msg += "\nImprove it (fix errors first, then make it faster). Reply with the complete module source only." | |
| return msg | |
| def summarize(episodes: list) -> dict: | |
| """episodes: dicts with `turns` = [{"stage", "correct", "speedup", "reward"}, ...].""" | |
| n = max(len(episodes), 1) | |
| furthest = [max((STAGE_ORDER.index(t["stage"]) for t in ep["turns"]), default=0) for ep in episodes] | |
| best = [max((t["speedup"] for t in ep["turns"] if t["correct"]), default=None) for ep in episodes] | |
| sp = [s for s in best if s is not None] | |
| roof = [max((t["roofline_frac"] for t in ep["turns"] if t["correct"] and t.get("roofline_frac")), default=None) | |
| for ep in episodes] | |
| roof = [r for r in roof if r is not None] | |
| return { | |
| "episodes": len(episodes), | |
| "compile_pct": 100 * sum(f >= STAGE_ORDER.index("runtime") for f in furthest) / n, | |
| "correct_pct": 100 * len(sp) / n, | |
| "correct_turn1_pct": 100 * sum(bool(ep["turns"] and ep["turns"][0]["correct"]) for ep in episodes) / n, | |
| "faster_than_ref_pct": 100 * sum(s > 1.0 for s in sp) / n, | |
| "mean_speedup": statistics.fmean(sp) if sp else None, | |
| "median_speedup": statistics.median(sp) if sp else None, | |
| "median_roofline_frac": statistics.median(roof) if roof else None, | |
| "stage_hist_all_turns": dict(Counter(t["stage"] for ep in episodes for t in ep["turns"])), | |
| "stage_hist_furthest": dict(Counter(STAGE_ORDER[f] for f in furthest)), | |
| "suspicious": sum(bool(h.get("suspicious")) for ep in episodes for h in ep.get("history", [])), | |
| **_pass_at_k(episodes, best), | |
| } | |
| def _pass_at_k(episodes, best) -> dict: | |
| """With several samples per task (--samples k): pass@k = share of tasks with any correct episode.""" | |
| by_task = {} | |
| for ep, b in zip(episodes, best): | |
| by_task.setdefault(ep.get("task_id"), []).append(b is not None) | |
| k = max((len(v) for v in by_task.values()), default=1) | |
| if k < 2: | |
| return {} | |
| return {"tasks": len(by_task), "samples_per_task": k, | |
| "pass@1_pct": 100 * sum(sum(v) / len(v) for v in by_task.values()) / len(by_task), | |
| f"pass@{k}_pct": 100 * sum(any(v) for v in by_task.values()) / len(by_task)} | |
| def split_task_ids(split: str, levels, kernelbench_dir: str = "") -> list: | |
| """Same deterministic split as the server (hash of task id), enumerated from a local KernelBench copy.""" | |
| tasks = load_tasks(kernelbench_dir, tuple(levels)) | |
| return sorted(t for t, task in tasks.items() if task.split == split) | |
| def run_episodes(generate, env_url: str, task_ids: list, split: str = "eval", env_timeout_s: float = 900.0, | |
| on_done=None, on_turn=None) -> list: | |
| """Turn-synchronous batched episodes: one generate() call per turn for all live episodes. | |
| on_done(ep) fires as each episode finishes (for incremental writes); on_turn(turn, stats) after each turn. | |
| """ | |
| envs = [KernelEnv(base_url=env_url, message_timeout_s=env_timeout_s).sync() for _ in task_ids] | |
| pool = ThreadPoolExecutor(max_workers=max(len(envs), 1)) | |
| try: | |
| obs = list(pool.map(lambda e, t: e.reset(task_id=t, split=split).observation, envs, task_ids)) | |
| convs = [[{"role": "user", "content": o.problem}] for o in obs] | |
| eps = [{"task_id": t, "turns": [], "completions": []} for t in task_ids] | |
| live = list(range(len(envs))) | |
| turn = 0 | |
| while live: | |
| turn += 1 | |
| t0 = time.monotonic() | |
| texts = generate([convs[i] for i in live]) | |
| gen_s = time.monotonic() - t0 | |
| def step(i, code): | |
| t = time.monotonic() | |
| return envs[i].step(KernelAction(code=code)), time.monotonic() - t | |
| results = list(pool.map(step, live, [extract_code(t) for t in texts])) | |
| for i, text, (r, step_s) in zip(live, texts, results): | |
| o = r.observation | |
| eps[i]["completions"].append(text) | |
| eps[i]["turns"].append({"stage": o.stage, "correct": bool(o.correct), "speedup": o.speedup or 0.0, | |
| "reward": r.reward or 0.0, "feedback": o.feedback, | |
| "t_ref_ms": o.t_ref_ms, "t_new_ms": o.t_new_ms, "t_min_ms": o.t_min_ms, | |
| "bound": o.bound, "roofline_frac": o.roofline_frac, | |
| "ref_roofline_frac": o.ref_roofline_frac, | |
| "gen_s": gen_s, "step_s": step_s}) | |
| convs[i] += [{"role": "assistant", "content": text}, {"role": "user", "content": feedback_message(o)}] | |
| if r.done: | |
| st = envs[i].state() | |
| eps[i].update(episode_id=st.episode_id, history=st.history) | |
| if on_done: | |
| on_done(eps[i]) | |
| if on_turn: | |
| on_turn(turn, {"live": len(live), "gen_s": gen_s, | |
| "stages": dict(Counter(eps[i]["turns"][-1]["stage"] for i in live)), | |
| "max_step_s": max(s for _, s in results)}) | |
| live = [i for i, (r, _) in zip(live, results) if not r.done] | |
| return eps | |
| finally: | |
| for e in envs: | |
| e.close() | |
| pool.shutdown() | |
| def run_episodes_pipelined(chat, env_url: str, task_ids: list, split: str = "eval", env_timeout_s: float = 3600.0, | |
| on_done=None, on_step=None) -> list: | |
| """One independent loop per episode (generate -> step -> ...), all concurrent. | |
| Against a vLLM server, continuous batching overlaps one episode's generation with another's | |
| kernel eval, instead of every episode waiting for the slowest one each turn. An episode that | |
| raises is kept with its error so one failure never loses the run. | |
| """ | |
| def one(task_id): | |
| ep = {"task_id": task_id, "turns": [], "completions": [], "reasoning": []} | |
| env = KernelEnv(base_url=env_url, message_timeout_s=env_timeout_s).sync() | |
| try: | |
| o = env.reset(task_id=task_id, split=split).observation | |
| conv = [{"role": "user", "content": o.problem}] | |
| while True: | |
| t0 = time.monotonic() | |
| text, finish, reasoning, ntok = chat(conv) | |
| gen_s = time.monotonic() - t0 | |
| t1 = time.monotonic() | |
| r = env.step(KernelAction(code=extract_code(text))) | |
| step_s = time.monotonic() - t1 | |
| o, truncated = r.observation, finish == "length" | |
| ep["completions"].append(text) | |
| ep["reasoning"].append(reasoning) | |
| ep["turns"].append({"stage": o.stage, "correct": bool(o.correct), "speedup": o.speedup or 0.0, | |
| "reward": r.reward or 0.0, "feedback": o.feedback, | |
| "t_ref_ms": o.t_ref_ms, "t_new_ms": o.t_new_ms, "t_min_ms": o.t_min_ms, | |
| "bound": o.bound, "roofline_frac": o.roofline_frac, | |
| "ref_roofline_frac": o.ref_roofline_frac, "truncated": truncated, | |
| "completion_tokens": ntok, "gen_s": gen_s, "step_s": step_s}) | |
| if on_step: | |
| on_step(ep) | |
| if r.done: | |
| break | |
| conv += [{"role": "assistant", "content": text}, {"role": "user", "content": feedback_message(o, truncated)}] | |
| st = env.state() | |
| ep.update(episode_id=st.episode_id, history=st.history) | |
| except Exception as e: # keep partial episode | |
| ep["error"] = repr(e) | |
| finally: | |
| env.close() | |
| if on_done: | |
| on_done(ep) | |
| return ep | |
| with ThreadPoolExecutor(max_workers=max(len(task_ids), 1)) as pool: | |
| return list(pool.map(one, task_ids)) | |
| def start_vllm_server(model: str, port: int, max_model_len: int, tp: int, spec_mtp: int, extra: str): | |
| """`vllm serve` subprocess; returns (proc, chat_factory).""" | |
| cmd = ["vllm", "serve", model, "--port", str(port), "--max-model-len", str(max_model_len), | |
| "--tensor-parallel-size", str(tp), *shlex.split(extra)] | |
| if spec_mtp: # model's own multi-token-prediction layer drafts tokens; greedy output is identical to plain decoding | |
| cmd += ["--speculative-config", json.dumps({"method": "mtp", "num_speculative_tokens": spec_mtp})] | |
| print("starting:", shlex.join(cmd), flush=True) | |
| proc = subprocess.Popen(cmd) | |
| deadline = time.monotonic() + 1800 | |
| while True: | |
| if proc.poll() is not None: | |
| raise RuntimeError(f"vllm serve exited with {proc.returncode}") | |
| try: | |
| urllib.request.urlopen(f"http://127.0.0.1:{port}/health", timeout=2) | |
| return proc | |
| except OSError: | |
| if time.monotonic() > deadline: | |
| proc.kill() | |
| raise RuntimeError("vllm serve not healthy after 30 min") | |
| time.sleep(3) | |
| def make_chat(base_url: str, model: str, temperature: float, max_new_tokens: int, enable_thinking: bool): | |
| from openai import OpenAI | |
| client = OpenAI(base_url=base_url, api_key="EMPTY", timeout=7200) | |
| def chat(conv): | |
| r = client.chat.completions.create( | |
| model=model, messages=conv, temperature=temperature, max_tokens=max_new_tokens, | |
| extra_body={"chat_template_kwargs": {"enable_thinking": enable_thinking}}) | |
| c = r.choices[0] | |
| return (c.message.content or "", c.finish_reason, getattr(c.message, "reasoning_content", None), | |
| r.usage.completion_tokens if r.usage else None) | |
| return chat | |
| def make_generate(backend: str, model: str, temperature: float, max_new_tokens: int, tp: int, batch_size: int, | |
| enable_thinking: bool = False, max_model_len: int = 32768): | |
| # templates without a thinking switch (Qwen2.5) ignore the kwarg | |
| tmpl = {"enable_thinking": enable_thinking} | |
| if backend == "vllm": | |
| from vllm import LLM, SamplingParams | |
| llm = LLM(model=model, tensor_parallel_size=tp, max_model_len=max_model_len) | |
| params = SamplingParams(temperature=temperature, max_tokens=max_new_tokens) | |
| return lambda convs: [o.outputs[0].text for o in llm.chat(convs, params, use_tqdm=False, chat_template_kwargs=tmpl)] | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| tok = AutoTokenizer.from_pretrained(model, padding_side="left") | |
| lm = AutoModelForCausalLM.from_pretrained(model, dtype="auto", device_map="auto") | |
| def generate(convs): | |
| out = [] | |
| for k in range(0, len(convs), batch_size): | |
| enc = tok.apply_chat_template(convs[k:k + batch_size], add_generation_prompt=True, padding=True, | |
| return_tensors="pt", return_dict=True, **tmpl).to(lm.device) | |
| with torch.no_grad(): | |
| ids = lm.generate(**enc, max_new_tokens=max_new_tokens, do_sample=temperature > 0, | |
| temperature=temperature or None, pad_token_id=tok.pad_token_id or tok.eos_token_id) | |
| out += tok.batch_decode(ids[:, enc["input_ids"].shape[1]:], skip_special_tokens=True) | |
| return out | |
| return generate | |
| def main(): | |
| p = argparse.ArgumentParser() | |
| p.add_argument("--env-url", default="http://localhost:8000") | |
| p.add_argument("--model", default="Qwen/Qwen2.5-Coder-7B-Instruct") | |
| p.add_argument("--backend", choices=["vllm-server", "openai", "vllm", "hf"], default="vllm-server", | |
| help="vllm-server: start `vllm serve` and pipeline episodes (fastest); openai: use --base-url; " | |
| "vllm/hf: in-process, turn-synchronous batches") | |
| p.add_argument("--base-url", default="http://127.0.0.1:8001/v1", help="OpenAI-compatible endpoint (openai backend)") | |
| p.add_argument("--port", type=int, default=8001, help="vllm-server port") | |
| p.add_argument("--max-model-len", type=int, default=65536) | |
| p.add_argument("--spec-mtp", type=int, default=0, help="speculative decoding with N MTP draft tokens (Qwen3.5: 2)") | |
| p.add_argument("--vllm-args", default="", help="extra `vllm serve` flags, e.g. '--language-model-only'") | |
| p.add_argument("--tp", type=int, default=1, help="vLLM tensor parallel size") | |
| p.add_argument("--split", default="eval") | |
| p.add_argument("--levels", type=int, nargs="+", default=[1]) | |
| p.add_argument("--kernelbench-dir", default=os.environ.get("KERNELBENCH_DIR", "")) | |
| p.add_argument("--task-ids", nargs="*", help="override task list") | |
| p.add_argument("--limit", type=int, default=0) | |
| p.add_argument("--samples", type=int, default=1, help="episodes per task (use temperature > 0) for pass@k") | |
| p.add_argument("--temperature", type=float, default=0.0) | |
| p.add_argument("--max-new-tokens", type=int, default=4096) | |
| p.add_argument("--enable-thinking", action="store_true", help="Qwen3-style <think>; raise --max-new-tokens with it") | |
| p.add_argument("--batch-size", type=int, default=8, help="HF backend generate batch") | |
| p.add_argument("--out", default="baseline_results.jsonl") | |
| p.add_argument("--wandb", metavar="PROJECT", help="log progress, summary and an episode table to W&B") | |
| args = p.parse_args() | |
| task_ids = args.task_ids or split_task_ids(args.split, args.levels, args.kernelbench_dir) | |
| task_ids = task_ids[: args.limit] if args.limit else task_ids | |
| task_ids = [t for t in task_ids for _ in range(args.samples)] | |
| print(f"{len(task_ids)} {args.split} tasks, model {args.model} ({args.backend})") | |
| server = None | |
| if args.backend == "vllm-server": | |
| server = start_vllm_server(args.model, args.port, args.max_model_len, args.tp, args.spec_mtp, args.vllm_args) | |
| args.base_url = f"http://127.0.0.1:{args.port}/v1" | |
| if args.backend in ("vllm-server", "openai"): | |
| chat = make_chat(args.base_url, args.model, args.temperature, args.max_new_tokens, args.enable_thinking) | |
| else: | |
| generate = make_generate(args.backend, args.model, args.temperature, args.max_new_tokens, args.tp, | |
| args.batch_size, args.enable_thinking, args.max_model_len) | |
| run = None | |
| if args.wandb: | |
| import wandb | |
| run = wandb.init(project=args.wandb, config=vars(args), job_type="baseline") | |
| done, lock = [], threading.Lock() | |
| out = open(args.out, "w") | |
| t_start = time.monotonic() | |
| def on_done(ep): | |
| with lock: | |
| done.append(ep) | |
| out.write(json.dumps(ep) + "\n") | |
| out.flush() | |
| best = max((t["speedup"] for t in ep["turns"] if t["correct"]), default=None) | |
| print(f"[{len(done)}/{len(task_ids)} done, {time.monotonic() - t_start:.0f}s] {ep['task_id']}: " | |
| f"{' > '.join(t['stage'] for t in ep['turns'])}" + (f", best {best:.2f}x" if best else "") | |
| + (f", ERROR {ep['error']}" if ep.get("error") else ""), flush=True) | |
| def on_step(ep): | |
| t = ep["turns"][-1] | |
| tok = t.get("completion_tokens") | |
| rate = f", {tok / t['gen_s']:.0f} tok/s" if tok and t["gen_s"] else "" | |
| print(f" {ep['task_id']} turn {len(ep['turns'])}: {t['stage']}" | |
| + (f" {t['speedup']:.2f}x" if t["correct"] else "") + (" [TRUNCATED]" if t.get("truncated") else "") | |
| + f" | gen {t['gen_s']:.0f}s ({tok} tok{rate}) | env step {t['step_s']:.0f}s", flush=True) | |
| if run: | |
| run.log({"gen_s": t["gen_s"], "step_s": t["step_s"], "completion_tokens": tok or 0, | |
| "truncated": int(bool(t.get("truncated"))), "correct": int(t["correct"])}) | |
| def on_turn(turn, stats): | |
| s = summarize(done) if done else {} | |
| print(f"[turn {turn}] {stats['live']} live, gen {stats['gen_s']:.0f}s, slowest step {stats['max_step_s']:.0f}s, " | |
| f"stages {stats['stages']}; done {len(done)}, correct {s.get('correct_pct', 0):.0f}%", flush=True) | |
| if run: | |
| run.log({"turn": turn, "live": stats["live"], "gen_s": stats["gen_s"], "max_step_s": stats["max_step_s"], | |
| **{f"stage/{k}": v for k, v in stats["stages"].items()}}) | |
| try: | |
| if args.backend in ("vllm-server", "openai"): | |
| eps = run_episodes_pipelined(chat, args.env_url, task_ids, args.split, on_done=on_done, on_step=on_step) | |
| else: | |
| eps = run_episodes(generate, args.env_url, task_ids, args.split, on_done=on_done, on_turn=on_turn) | |
| finally: | |
| out.close() | |
| if server: | |
| server.terminate() | |
| print(f"wall time {time.monotonic() - t_start:.0f}s") | |
| summary = summarize(eps) | |
| print(json.dumps(summary, indent=2)) | |
| if run: | |
| run.summary.update({k: v for k, v in summary.items() if not isinstance(v, dict)}) | |
| cols = ["task_id", "episode_id", "stages", "best_speedup", "turns"] | |
| rows = [[ep["task_id"], ep.get("episode_id"), " > ".join(t["stage"] for t in ep["turns"]), | |
| max((t["speedup"] for t in ep["turns"] if t["correct"]), default=None), len(ep["turns"])] for ep in eps] | |
| run.log({"episodes": wandb.Table(columns=cols, data=rows)}) | |
| run.finish() | |
| if __name__ == "__main__": | |
| main() | |