Beyond kernels: Letting agents optimize the whole AI job

AI agents can now write and optimize individual GPU kernels. But there is more to a job’s performance than the speed of its kernels. Data transfers, CPU work, and synchronization can also slow it down.

Recently, we have been trying agents to improve the end-to-end performance of AI jobs. We built /who-ate-my-flops, a small plugin for Claude Code and Codex, to provide profiling tools, workload context, and correctness checks.

Performance improvements made with the tool were merged into six open-source repositories: FunASR, FastVideo, gsplat, Ultralytics, Unsloth, and SGLang, with measured speedups of up to 3.6× on the tested workloads. This blog describes what we learned. We hope the tool can complement kernel-design agents and help extend the use of AI agents to end-to-end optimization of AI workloads.

From kernel optimization to whole-job optimization

A kernel-design agent (KDA, K-Search) typically starts with a reference implementation and test inputs. It then follows an optimization loop: use profiling tools to identify bottlenecks, revise the kernel, verify correctness, and measure performance.

The same loop can also be used to optimize a whole training or inference job. But the optimization target is no longer a single operation. Changes may span the codebase, and each job has its own constraints and correctness requirements.

Optimizing a whole job means reasoning about how its parts interact. A faster kernel may expose a data-loading bottleneck; overlapping communication with computation may require changes to scheduling and memory use. Evaluating an optimization requires checking both its effect on end-to-end runtime and whether it preserves the job’s behavior.

The table below summarizes the axes along which whole-job optimization differs from kernel-only optimization.

Where a whole job loses time

The patterns below illustrate some common sources of lost time when looking beyond individual kernels.

> /who-ate-my-flops

To address these challenges, we built /who-ate-my-flops.

To get started, give the agent a repository, a launch command, and access to an idle GPU environment. Run /who-ate-my-flops:init to agree on the goal and constraints. Then choose /who-ate-my-flops:diagnose to find performance issues, or /who-ate-my-flops:optimize to iteratively improve end-to-end performance.

The walkthrough below shows this setup for YOLO training with Ultralytics.

Get the plugin and setup instructions ↗

How we designed the harness

We rely on the underlying foundation model to reason about performance and propose optimizations. Our harness provides (1) workload context and profiling evidence to ground that reasoning, (2) correctness checks during optimization, and (3) guidance for the agent to ask users questions and align on goals and constraints.

Two key contexts: computation structure and runtime profiling

Correctness checks

Checking whether a change preserves an AI job’s behavior is more complex than checking a single kernel’s outputs. So we built a small tool, parity, to help with these checks. It records intermediate values and compares them across runs. For training, these should include inputs, loss, and gradient norms:

for step, batch in enumerate(loader):
    # Record an input checksum.
    parity.record(f"batch_step{step}", batch.sum())

    loss = model(batch)
    # Record the forward-pass loss.
    parity.record(f"loss_step{step}", loss, rtol=1e-3)

    loss.backward()
    # Record the gradient norm.
    parity.record(f"gradnorm_step{step}", grad_norm(model), rtol=1e-3)

    optimizer.step()  # Update the model weights.

The checkpoints are saved to baseline.json and optimized.json. Running parity compare baseline.json optimized.json compares matching tags using the declared tolerances. The report shows which checks failed and by how much, helping the agent investigate the cause:

== 3 tag(s) compared, 2 pass, 1 fail ==

FAIL
   gradnorm_step0
      expected 3.0   actual 3.012
      |diff| 0.012   budget 0.003   4x over budget

With parity, users only need to reason about what to check and how much numerical error is acceptable. They can leave the performance work to the agent without having to work through all the ML systems details. In practice, the agent can suggest where to add checks, but we sometimes request additional checks for greater confidence in correctness. We fix random seeds and data order, then run the baseline multiple times to measure the natural variation in the recorded values. This helps the agent distinguish natural variation from differences introduced by an optimization.

Asking users questions proactively

Users may forget to mention some requirements. Those without an ML systems background may not know which details matter. So we let the agent proactively ask questions during init to clarify the user’s goals, constraints, and preferences before optimization begins.

For example, the agent asks, “Are you an ML systems engineer, or an algorithm engineer less familiar with systems?” This helps it tailor its explanations to your background. It also asks what changes are allowed: can it adjust configurations, install faster backends, change Python code, or modify Triton/CUDA kernels? Some users prefer smaller changes that are easier to understand, review, and maintain.

In /who-ate-my-flops, we also provide helpers to turn a job into a repeatable benchmark, reduce file sizes and experiment runtime, and track changes and measurements across experiments. As foundation models improve, agents might need less of this guidance. For now, it helps avoid wasted attempts. The implementation details are in the repository.

Optimization case studies

The agent found several kinds of performance problems, including unnecessary data transfers, repeated initialization, and frequent CPU–GPU synchronization. The two case studies below show how profiling evidence led to concrete changes; the following section summarizes speedups and costs across the evaluated workloads.

Not every finding required this much investigation. In our own workloads, the agent sometimes found a missing fast kernel library or a performance-related flag that had not been enabled. These requirements may be missing from the installation instructions, so a coding agent can overlook them when setting up the environment. /who-ate-my-flops can catch and fix these setup issues too.

Experiments: speedups and costs

We explored training and inference workloads from repositories spanning several application domains. Here, we highlight six examples whose optimizations were merged upstream. For each workload, we ran init to clarify the goal and constraints, followed by optimize. We used Fable 5.1, and all runs used B200 GPUs within a single node. The table below reports the measured improvements for each workload. Expand a repository row to see the problems found and how they were addressed.

What a run costs

The charts below show the time spent on init and optimize, along with total token usage for these six workloads.

Most recorded tokens were cache reads. The costs shown exclude human review and later PR revisions.

Where human judgment still matters in end-to-end optimization

Correctness. The agent can add correctness checks to the codebase and run them on its own. Users can help ensure the checks cover what matters for the workload and add stricter ones based on their goals and tolerance for risk.

Tradeoffs. A change that makes one workload faster may not suit other uses of the codebase. It may use more memory, break another use case, or make the code harder to maintain. Human review is needed to understand these tradeoffs and intervene when necessary.

From experiments to PRs. The agent might bundle several unrelated optimizations into one branch. Together, they may give the largest speedup, but they also make the PR harder to review. Users still need to decide which changes belong together in a PR, which should stay as local changes, and whether an optimization should be opt-in rather than enabled by default.

Related work

Agents are now widely believed to design standalone kernels very well. And there is increasing attention on end-to-end performance. Almost all of that attention points at one workload: LLM inference serving. InferenceBench gives an agent a model, one bare H100, and two hours to produce a fast OpenAI-compatible server, using whatever stack it chooses. ISO-Bench points the agent at a bottleneck in vLLM or SGLang that a maintainer has already fixed, and scores its patch against theirs. Both are confined to LLM inference, and both measure one ability: optimizing an inference engine.

ASAP optimizes LLM training instead: it searches for sharding configurations for distributed LLM training on TPU clusters. However, it does not modify the codebase, and it does not verify correctness, on the assumption that changing the parallelization does not change output. Its three experiments are scored by whether the proposal matches what a human expert chose, rather than by a measured A/B.

Many open-source projects have built their own skills for internal performance work. BBuf's collection and the skills in SGLang's own repository cover torch.profiler triage, cross-framework benchmarking and more, with an autonomous loop that ties them together. All of it is for an inference engine, and most of it is specific to SGLang. NVIDIA ships a similar family inside TensorRT-LLM: twenty skills and ten agents covering kernel writing, bottleneck classification, and an optimization loop with an explicit revert_file step. It makes a great deal of internal performance practice portable, and it is scoped to TensorRT-LLM. AMD's Hyperloom agent optimizes model inference on AMD Instinct.

Our goal is to make performance agents easy to use across users’ own training and inference jobs.

Try it on your own workload

The plugin is available for Claude Code and Codex. Parity can also be used separately to record and compare values across runs. We’d like to hear where it helps, where it fails, and what still needs human intervention. Give it a try!

References

1. Z.ai. GLM built its inference infrastructure.
2. OpenAI. How GPT-5.6 fuses frontier intelligence with frontier efficiency.
3. NVIDIA. Kernel Design Agents (KDA).
4. Shiyi Cao et al. K-Search: LLM Kernel Generation via Co-Evolving Intrinsic World Model.
5. Jehyeok Yeon et al. InferenceBench.
6. Ayush Nangia et al. ISO-Bench.
7. Yuran Ding et al. ASAP.
8. BBuf. AI-Infra-Auto-Driven-SKILLS.
9. SGLang. Performance skills.
10. NVIDIA TensorRT-LLM. Skills and agents.
11. AMD. Hyperloom.