Skip to content

Beyond kernels: Letting agents optimize the whole AI job

OpenPerfAgent (Cao, Hu, Li) releases /who-ate-my-flops, a Claude Code / Codex plugin that points a coding agent at end-to-end training or inference jobs — not individual kernels — and iteratively speeds them up under a numerical-agreement correctness contract. The harness provides workload context (parsed computation-structure JSON), profiling evidence (torch.profiler + Nsight Systems query scripts), a parity tool that records/compares intermediate tensor values across runs at declared tolerances, and a proactive-questioning init step that clarifies user goals and allowed change scope. Applied by Fable 5.1 to six OSS repos (FunASR, FastVideo, gsplat, Ultralytics, Unsloth, SGLang), the agent found and upstreamed unnecessary CPU↔GPU transfers, repeated attention plan setup, missing fast-kernel installs, and other whole-job bottlenecks — measured speedups up to 3.6× on tested workloads, per-run cost ~$6–15 in agent time / 98–214M tokens per repo, all merged upstream.

  • Whole-job optimization differs from kernel optimization on four axes: optimization goal (total runtime vs kernel latency), change scope (spans model code + data loading + config + kernels vs a single operator), correctness (per-checkpoint tensor tolerance across multiple steps vs one operator-output comparison), and inputs (repository + launch command + hardware vs reference impl + test inputs) [§From kernel optimization to whole-job optimization, comparison table].
  • The agent operates on two grounding contexts: a computation-structure JSON (layers, modules, ops, scheduling of compute/data-movement/sync, with file:line source pointers) that describes what the job is supposed to do, and runtime profiling views (time-by-op, busy/idle per rank, communication overlap) computed via canned scripts over torch.profiler / Nsight Systems traces that describe what actually happened [§How we designed the harness, “Two key contexts”].
  • Correctness is handled by a companion tool parity that records intermediate values (input checksum, loss, gradient norm) with declared rtol at named checkpoints during a baseline run, then compares tagged values from an optimized run — reporting per-check pass/fail so the agent can localize regressions instead of relying on a single output comparison [§Correctness checks, code block].
  • Case study — FastVideo MiniMax H3 video generation on 4×B200: agent identified that VAE decoding stage spent 5.87s of 8.87s on unnecessary CPU↔GPU transfers (10.4 GB of VAE weights loaded to GPU then copied back to CPU to free memory, though inference didn’t modify them). Fix: keep CPU copies pinned, release GPU copies after decode. End-to-end median time 15.20s → 6.34s (2.40× speedup), merged as FastVideo PR #1867 [§FastVideo case study].
  • Case study — FunASR Fun-ASR-Nano fine-tuning on 1×B200: 59 of 382 steps were ~1s (vs ~98ms baseline) because cuDNN was rebuilding attention execution plans for new batch shapes (444ms forward + 554ms backward CPU-side). Fix: select alternate SDPA backends + compile the decoder. Mean step time (ex-first) 243ms → 68ms (3.6× speedup), shipped opt-in in FunASR PR #3705 [§FunASR case study].
  • Six merged-upstream results across FunASR, FastVideo, gsplat, Ultralytics, Unsloth, and SGLang, all on B200 single-node, run via Fable 5.1. Per-workload cost: init 3m31s–14m08s, optimize 1h13m–2h50m, tokens 98M–214.8M (mostly cache reads), excluding human review and PR revisions [§Experiments table].
  • Human judgment remains load-bearing at three points: (a) correctness — humans decide what checks matter and how strict the tolerances are; (b) trade-offs — a speedup on one workload may hurt maintainability, memory use, or other callers; (c) PR shaping — the agent tends to bundle unrelated optimizations into one branch, so users still decide what belongs together and what should be opt-in [§Where human judgment still matters].
  • Related-work positioning: distinct from the LLM-inference-serving-only end-to-end agent cluster (InferenceBench, ISO-Bench, SGLang/TensorRT-LLM/Hyperloom internal skill packs — all bound to a single inference engine) and from ASAP (LLM-training sharding search that doesn’t modify the codebase or verify correctness); target is arbitrary user training/inference jobs across application domains [§Related work].

The plugin exposes three top-level commands to the coding agent. /who-ate-my-flops:init asks the user proactive questions (background, allowed change scope, benchmark-repeatability preferences) and produces a workload spec + constraints file. /who-ate-my-flops:diagnose runs the profiling scripts and emits a computation-structure JSON + a set of ranked performance-issue candidates with file:line back-links. /who-ate-my-flops:optimize runs the diagnose→edit→verify loop: propose a change, run the modified job, replay parity checkpoints against baseline (fail hard if any tagged value drifts beyond its declared rtol), measure end-to-end runtime, keep or revert. Helper scripts convert an arbitrary job into a repeatable benchmark (reduce dataset size / iteration count while preserving representative behavior) and track experiment provenance. parity is released separately as a standalone tool at openperfagent/parity — users insert parity.record(tag, tensor, rtol=…) calls at the checkpoints they care about, then parity compare baseline.json optimized.json produces a per-tag delta report.

Six upstream-merged PRs across FunASR, FastVideo, gsplat, Ultralytics, Unsloth, SGLang, all on B200 single-node with Fable 5.1 [§Experiments table]. FastVideo: 2.40× E2E speedup on MiniMax H3 (15.20s → 6.34s median, 4×B200) via pinning CPU copies + eliminating post-decode GPU→CPU returns [§FastVideo case study]. FunASR: 3.6× E2E speedup on Fun-ASR-Nano fine-tuning (243ms → 68ms mean step, 1×B200) via SDPA-backend selection + decoder compilation [§FunASR case study]. Table lists PRs merged to the four remaining repos (gsplat PR #1062, Ultralytics PR #26180, Unsloth PR #11238, SGLang PR #41223) without per-workload speedup breakdowns in the visible portion of the post. Per-run cost across the six: init 3m31s–14m08s, optimize 1h13m–2h50m, 98M–214.8M tokens (mostly cache reads) [§What a run costs].

Fits the AI-for-AI Research cluster’s deployment-stack granularity — the same slot occupied by Sol Video Inference Engine: Agent-Native Full-Stack Acceleration Framework for Efficient Video Generation (autoresearching five acceleration levers for video-DiT inference) and Speeding up end-to-end inference with self-improving agents (agent-optimizing the vLLM stack for DeepSeek v4 Pro + GLM 5.2), but broader in workload scope than either: Sol Video Inference Engine and Asari both target production inference-serving stacks the authors control, whereas /who-ate-my-flops is a portable plugin pointed at arbitrary upstream OSS training + inference repos. The correctness contract is a middle point on the strictness axis this concept has been mapping: Asari’s token-level distribution matching on AllenAI C4 sits at one end (production-shippability gate), Sol-Engine’s perceptual human-validator loop (Sol-H3: Recursive Self-Improvement for MiniMax-H3 Inference Acceleration on Sol-Engine across Cloud and Edge sharpens the operator-layer variant) at the other, and parity’s per-tag rtol on intermediate tensors sits between — stricter than perceptual, cheaper than distribution-matching.

Also a concrete Agentic Software Engineering datapoint on the “harness > model” thesis surfaced in Benchmarking Coding Agents on Databricks' Multi-Million Line Codebase (per-task cost moves >2× at fixed model+quality across harnesses) and AutoDesign: Meta-Harness Optimization for Long-Horizon Agentic Design (learned harness transfers across seven code-agent × model pairs): the whole plugin is a specialized harness — profiling-query scripts, computation-structure parser, parity, proactive init questionnaire — that turns a generic coding agent (Fable 5.1) into an end-to-end performance engineer without model retraining. Sibling to AI Performance Engineering — curated resource list for GPU perf engineering and production inference on the human-facing side (curated resource list for the same skillset); the OpenPerfAgent contribution is that the same skillset now installs as a plugin the agent uses autonomously. Positioned by the authors as complementary to kernel-design agents like BackendBench: Ship Correct and Fast LLM Kernels to PyTorch rather than competing — the kernel agents optimize a single op, this harness optimizes how the ops fit together.