"""The kernel agent loop: the given model writes and iterates a ModelNew with feedback from the NON-protected local check (compile + rules + correctness on one public input, no timing), then submits once to the protected evaluator. Durable: every turn is appended to /attempts.jsonl, so an in-episode crash (or a real kill) resumes from the saved attempts, and the final /eval uses a deterministic request id (never evaluated twice). """ from __future__ import annotations import hashlib import json import os import re import tempfile import time import uuid MAX_TURNS = int(os.environ.get("KERNEL_MAX_TURNS", "5")) MAX_TOKENS = int(os.environ.get("KERNEL_MAX_TOKENS", "32000")) EPISODE_WALL_S = float(os.environ.get("KERNEL_EPISODE_WALL_S", "2400")) SYSTEM = """You are an expert GPU kernel engineer. You are given a PyTorch reference module `Model` from KernelBench. Write `ModelNew`: a drop-in replacement (same __init__ arguments, same forward inputs and outputs) that is faster because it uses your own GPU kernels. Environment (fixed): - One NVIDIA A100-SXM4-40GB (sm_80, 108 SMs, 40 MB L2, ~1.5 TB/s HBM). CUDA 12.9, PyTorch 2.9.1, Triton 3.5.1, Python 3.10. - Your kernels: CUDA C++ through torch.utils.cpp_extension.load_inline, and/or Triton (@triton.jit). At least one of your own kernels must actually run in every forward call. You may still call torch ops (cuBLAS/cuDNN through torch) for the parts you do not replace. - Numerics are strict float32: outputs must match the reference within atol=rtol=1e-4 on fresh random inputs and fresh random weights (initialised like the reference: build parameters in the same order with the same nn layers, e.g. self.conv = nn.Conv2d(...), so the evaluator's seeded init gives identical weights). TF32, fp16 and bf16 shortcuts usually fail this tolerance; cuBLAS/cuDNN TF32 is disabled and global flags cannot be changed. - Score: speedup = reference eager time / your time, measured by a protected evaluator (cold L2, fresh inputs on every call, device-synchronised timing interleaved with the reference). Correct and faster than eager with 95% confidence counts as a win; the size of the speedup is your score. Rules (enforced; a violation rejects the submission): - Imports only from: torch, torch.nn, torch.nn.functional, torch.utils.cpp_extension (load_inline only), triton, triton.language, math, numpy, typing, functools, itertools, operator, collections, dataclasses, enum. - No try/except, no global statements, no streams, events, CUDA graphs, torch.compile, torch.jit, torch.backends, profiler, file/OS/process access, exec/eval/open/setattr/getattr-with-variables, no private (_x) or dunder attributes except __init__, no assignment to attributes other than self., no class named Model. - Do not cache results across calls, do not modify the input tensors in place, do not return a tensor subclass. - Native code: no Python C API or pybind11 types, no stream/graph/event APIs, no host system calls or files. Launch kernels on the current stream (at::cuda::getCurrentCUDAStream()) or the default stream. - load_inline compile time: about 35 s when cuda_sources include ATen/torch headers, about 5 s when cuda_sources are plain CUDA (#include ) with launcher functions that take raw pointers and a cudaStream_t, while the torch::Tensor wrapper lives in cpp_sources (declare the launcher there). Either works. Reply format: think as needed, then give exactly ONE ```python code block containing the complete module (imports, kernels, class ModelNew). End with a final line that is either ACTION: CHECK (run the local check: compile, rules and correctness on one public input; no timing) or ACTION: SUBMIT (final: the protected evaluator scores this code once). """ FEEDBACK_HEAD = "Local check result for your last code (correctness only, public input, no timing):" def _sha(s: str) -> str: return hashlib.sha256(s.encode()).hexdigest() def extract_code(text: str) -> str | None: blocks = re.findall(r"```(?:python|py)?[ \t]*\n(.*?)```", text, flags=re.DOTALL) blocks = [b for b in blocks if "ModelNew" in b] or blocks return blocks[-1].strip() + "\n" if blocks else None def extract_action(text: str) -> str: m = re.findall(r"ACTION:\s*(CHECK|SUBMIT)", text) return m[-1] if m else "CHECK" def render_check(r: dict) -> str: st = r.get("status") lines = [FEEDBACK_HEAD, f"status: {st}"] if st == "ok": lines.append("compiled, passed the rules, your kernel ran, and the output matched the reference on the public " "input (max abs diff %s)." % r.get("max_abs_diff")) else: if r.get("error"): lines.append("error: " + str(r["error"])[:2500]) if r.get("compile_log"): lines.append("compiler output (tail):\n" + str(r["compile_log"])[-3500:]) if r.get("detail") and st in ("runtime_error",): lines.append("detail (tail):\n" + str(r["detail"])[-2000:]) for k in ("max_abs_diff", "max_rel_diff", "mismatch_frac"): if r.get(k) is not None: lines.append(f"{k}: {r[k]}") return "\n".join(lines) def compact_summary(attempts: list[dict], why: str) -> str: """Rebuild progress from saved attempts (used for in-episode compaction and crash resume).""" if not attempts: return f"[{why}: no attempts yet]" last = attempts[-1] passing = [a for a in attempts if (a.get("check") or {}).get("status") == "ok"] out = [f"[{why}. Progress so far: {len(attempts)} attempt(s); " f"{len(passing)} passed the local check.]"] hist = "; ".join(f"turn {a['turn']}: {(a.get('check') or {}).get('status', a.get('action'))}" for a in attempts) out.append("History: " + hist) best = passing[-1] if passing else last out.append(("Latest code that passed the local check:" if passing else "Latest code (did not pass):") + f"\n```python\n{best.get('code', '')}```") if last.get("check") and last is not best: out.append("Latest attempt's check: " + render_check(last["check"])[:2500]) elif last.get("check"): out.append(render_check(last["check"])[:2500]) return "\n\n".join(out) class KernelAgent: def __init__(self, service, llm_factory): self.svc = service self.llm_factory = llm_factory # model id -> lineage LLM def run(self, pid: str, task_text: str, model: str, context: str, *, lifecycle=None, run_dir=None, purpose="stream") -> dict: t0 = time.time() own_tmp = None if not run_dir: own_tmp = tempfile.mkdtemp(prefix="kernel-ep-") run_dir = own_tmp os.makedirs(run_dir, exist_ok=True) att_path = os.path.join(run_dir, "attempts.jsonl") meta_path = os.path.join(run_dir, "episode.json") final_path = os.path.join(run_dir, "final.json") if os.path.exists(final_path): # a completed episode re-entered after a kill: same result, no new work fin = json.load(open(final_path)) fin["resumed_completed"] = True return fin if os.path.exists(meta_path): meta = json.load(open(meta_path)) else: meta = {"episode": uuid.uuid4().hex, "problem": pid, "model": model, "purpose": purpose, "started": time.time()} json.dump(meta, open(meta_path, "w")) attempts = [json.loads(ln) for ln in open(att_path)] if os.path.exists(att_path) else [] first = (f"{context.strip()}\n\n---\n\n" if context and context.strip() else "") + task_text tokens_in = tokens_out = calls = 0 events = [dict(e) for e in (lifecycle or [])] applied = [] at = {} for e in events: k = 1 + (int(e.get("at_step", 1)) % max(1, MAX_TURNS - 1)) at.setdefault(k, []).append(e) if attempts: # real resume after a process kill messages = [{"role": "user", "content": first + "\n\n" + compact_summary(attempts, "Resumed after a restart") + "\n\nContinue: improve or fix the code."}] applied.append({"type": "resume_from_disk", "turns": len(attempts)}) else: messages = [{"role": "user", "content": first}] submit_code, submit_turn, err = None, None, "" start_turn = len(attempts) if attempts and attempts[-1].get("check") is None and attempts[-1].get("code"): # killed after deciding to submit: re-send the same submission (same request id -> no second eval) submit_code, submit_turn = attempts[-1]["code"], attempts[-1]["turn"] start_turn = MAX_TURNS for turn in range(start_turn, MAX_TURNS): if time.time() - t0 > EPISODE_WALL_S: err = f"episode wall clock exceeded {EPISODE_WALL_S}s" break use_model = model for e in at.get(turn, []): typ = e.get("type") if typ == "compaction": messages = [{"role": "user", "content": first + "\n\n" + compact_summary(attempts, "Context compacted") + "\n\nContinue: improve or fix the code."}] applied.append({"type": "compaction", "turn": turn}) elif typ == "crash": # simulated crash: drop all in-memory state, then resume from what is on disk messages = None attempts = [json.loads(ln) for ln in open(att_path)] if os.path.exists(att_path) else [] messages = [{"role": "user", "content": first + "\n\n" + compact_summary(attempts, "Resumed after a crash") + "\n\nContinue: improve or fix the code."}] applied.append({"type": "crash", "turn": turn, "restored_attempts": len(attempts)}) elif typ == "model_swap": use_model = e.get("to") or model applied.append({"type": "model_swap", "turn": turn, "to": use_model, "back_to": model}) if turn == MAX_TURNS - 1: messages[-1]["content"] += "\n\nThis is your last reply: it will be submitted to the protected evaluator." llm = self.llm_factory(use_model) try: res = llm.chat(messages, system=SYSTEM, max_tokens=MAX_TOKENS) except Exception as e: # noqa: BLE001 err = f"llm error ({use_model}): {type(e).__name__}: {str(e)[:300]}" calls += getattr(e, "llm_attempts", 1) break calls += 1 + max(0, ((res.raw or {}).get("attempts", 1) if isinstance(res.raw, dict) else 1) - 1) tokens_in += res.input_tokens tokens_out += res.output_tokens text = res.text or "" code = extract_code(text) action = extract_action(text) rec = {"turn": turn, "model": use_model, "action": action, "code": code or "", "ts": time.time(), "tokens_in": res.input_tokens, "tokens_out": res.output_tokens, "stop": (res.raw or {}).get("stop_reason") if isinstance(res.raw, dict) else None} messages.append({"role": "assistant", "content": text if text.strip() else "(empty reply)"}) if code is None: rec["check"] = {"status": "no_code", "error": "no ```python code block found in the reply"} fb = "No ```python code block was found in your reply. Reply with one complete module." elif action == "SUBMIT" or turn == MAX_TURNS - 1: submit_code, submit_turn = code, turn rec["check"] = None _append(att_path, rec) attempts.append(rec) break else: try: chk = self.svc.check(pid, code, client=f"{purpose}:{meta['episode']}") except Exception as e: # noqa: BLE001 err = f"local check unavailable: {type(e).__name__}: {str(e)[:300]}" rec["check"] = {"status": "service_error", "error": err} _append(att_path, rec) attempts.append(rec) break rec["check"] = {k: chk.get(k) for k in ("status", "error", "compile_log", "detail", "max_abs_diff", "max_rel_diff", "mismatch_frac", "wall_s")} fb = render_check(chk) left = MAX_TURNS - 1 - turn fb += (f"\n\nYou have {left} more repl{'y' if left == 1 else 'ies'}. Either SUBMIT this code, or " f"improve it (speed or correctness) and CHECK again.") _append(att_path, rec) attempts.append(rec) messages.append({"role": "user", "content": fb}) # choose what to submit if submit_code is None: passing = [a for a in attempts if (a.get("check") or {}).get("status") == "ok" and a.get("code")] if passing: submit_code, submit_turn = passing[-1]["code"], passing[-1]["turn"] report, status = None, "" if submit_code is not None and not err.startswith("local check unavailable"): rid = f"kernel|{meta['episode']}|{_sha(submit_code)[:24]}" try: report = self.svc.evaluate(pid, submit_code, rid, client=f"{purpose}:{meta['episode']}") status = report.get("status", "") except Exception as e: # noqa: BLE001 err = (err + "; " if err else "") + f"protected evaluator unavailable: {type(e).__name__}: {str(e)[:300]}" status = "service_error" elif submit_code is None: status = "no_submission" err = err or "no code passed the local check and none was submitted" out = {"problem": pid, "model": model, "purpose": purpose, "episode": meta["episode"], "status": status, "report": report, "submit_turn": submit_turn, "submitted_source": submit_code, "attempts": attempts, "error": err, "tokens_in": tokens_in, "tokens_out": tokens_out, "calls": calls, "wall_s": round(time.time() - t0, 2), "events": events, "applied": applied} if status not in ("service_error",) and not err.startswith(("llm error", "local check unavailable")): tmp = final_path + ".tmp" json.dump(out, open(tmp, "w"), default=str) os.replace(tmp, final_path) return out def _append(path: str, rec: dict) -> None: with open(path, "a") as f: f.write(json.dumps(rec, default=str) + "\n") f.flush() os.fsync(f.fileno())