I Built an Agent That Optimizes GPU Kernels While the Model Keeps Serving

I created this for the purposes of entering the All Things Agentic Hackathon.

---

Writing optimized GPU kernels follows a loop that gets kind of boring after the third time. You profile an operation. You figure out why it's slow — usually memory bandwidth, sometimes compute. You write a fused kernel that keeps data in the GPU's fast on-chip memory instead of bouncing it back to main memory. You test it. You plug it back into the model. Same steps, every time.

I kept doing this loop manually while working on inference infrastructure, and I kept thinking: this whole cycle is mechanical enough that an agent should be able to own it. Not just generate the code — any LLM can spit out a Triton kernel. The hard parts are deciding what to optimize (the profiling), verifying the speedup is real (not gaming the benchmark), and deploying it without breaking the running model (the hot-swap). That's the stuff that takes a human in the loop.

So I built gpuyantra. It's an autonomous agent tree that profiles a running model, identifies the bottleneck, writes an optimized Triton kernel, verifies it's correct and genuinely faster, and hot-swaps it into a live inference server — while the model keeps serving. No restart, no weight reinitialization. The model just gets faster.

The result: 7.24× faster than stock PyTorch, 1.39× faster than PyTorch's own compiler, verified across 15 correctness checks. On a single $0.71/hour GPU.

The Loop the Agent Runs

Here's the workflow, broken down into what actually happens on the GPU.

Step 1 — Profile. The Profiler agent measures the target operation on the GPU and places it on the L4's roofline model — a chart that shows whether an operation is limited by how fast data can move (memory-bound) or how fast the chip can do math (compute-bound). For RMSNorm on Qwen2.5-1.5B, it measures an arithmetic intensity of about 1.25 calculations per byte of data moved. The L4's crossover point — where math speed equals memory speed — sits at 101 calculations per byte. So RMSNorm is roughly 81× below that crossover, meaning the GPU's compute units are sitting idle waiting for data to arrive from memory. That gap is the headroom a fused kernel can recover.

The key insight: the profiler outputs a bottleneck fingerprint, and that fingerprint contains no model name and no operation name. Just the physics: op=norm mem_bound=True ai=1.2 tile=1024 hw=L4. This is what makes cross-model skill transfer possible, which I'll get to.

Step 2 — Retrieve. The Supervisor takes that fingerprint, embeds it with gemini-embedding-001, and runs a Firestore vector search against the skill library. Skills are indexed by why an operation was slow — the memory-access pattern and the hardware characteristics — not by what it's called. A bandit algorithm (think slot machines: mostly pick what worked before, occasionally try something new) selects which retrieved skill to start from.

Step 3 — Write. The Coder agent writes a Triton kernel, plus something I'm particularly proud of: it generates the deployment contract — an adapter_mapping that declares how the kernel's parameters connect to the model's module attributes. Every published kernel deployment system I found (HuggingFace kernels, FlashInfer-Bench, Kernel Contracts from arXiv:2604.22032) uses human-authored deployment bridges. gpuyantra is the first where the agent writes the bridge and a deterministic verifier validates it.

Step 4 — Judge. The Judge agent calls the verifier, which runs a four-layer trust anchor:

Layer 1 — AST checker: seven syntactic rules that catch reward hacking before any code executes. It catches things like calling PyTorch's own implementation and claiming the speedup, returning the input unchanged, having a @triton.jit function that's never actually called, returning uninitialized memory, and more.

Layer 2 — Contract validation: checks that every attribute in the adapter mapping actually exists on the target module, using a torch.device("meta") instance (not hasattr on the class, which misses attributes assigned in __init__).

Layer 3 — Subprocess sandbox: 5 random seeds × 3 input shapes = 15 correctness checks, all within 1% tolerance. Runs in a scrubbed environment with only four env vars, kills with SIGKILL on timeout, and probes GPU health afterward.

Layer 4 — Reward computation: crucially, recomputed in-process from the verified numbers. The subprocess's own reward field is discarded. A kernel that prints {"reward": 3} in its stdout still scores −1 if it's wrong. This is a trust boundary.

The reward ladder: −1 (wrong), +1 (correct but not faster), +2 (beats eager by >5%), +3 (beats both eager AND torch.compile by >5%).

Step 5 — Deploy. If the reward is +2 or higher, the Supervisor calls hotswap_kernel, which POSTs to the inference server's /swap endpoint. The swap uses types.MethodType to rebind the forward method on every matching module instance — 57 Qwen2RMSNorm layers in Qwen2.5's case. Nothing is copied. The original weights stay in GPU memory in the right dtype on the right device. The server doesn't restart. A parity check against the live weights runs before the swap commits, and it auto-rolls back if anything is off.

One critical constraint: the served model is never torch.compile'd. A compiled graph bakes the current forward into its dispatch; a later types.MethodType patch silently does nothing. We'd report a successful swap and unchanged throughput. The only torch.compile in the repo is inside measure_baselines(), where the compiled reference is timed for comparison and then discarded.

Making Sure the Speedup Is Real

This is the part I think matters most, and it's the part most hackathon projects skip.

If you time a kernel against a deliberately slow baseline, everything looks amazing. The initial KernelBench leaderboard had entries claiming 3×+ speedups that collapsed to under 1.5× when measured properly. One prominent result (Sakana AI's CUDA Agent, arXiv:2602.24286) reported a 3.13× headline that dropped to 1.49× under honest baselines.

gpuyantra does three things to prevent this.

Honest baselines. Every timing comparison enables TF32 (a hardware-accelerated precision mode that makes PyTorch's own math faster) and also measures torch.compile(mode="reduce-overhead"). This is the KernelBench-Verified protocol. The number to beat is PyTorch at its best, not PyTorch handicapped.

Determinism off for timing. We discovered that torch.use_deterministic_algorithms(True) penalizes eager PyTorch by about 23% (it forces slower cuBLAS codepaths) while leaving Triton kernels completely unaffected — Triton generates its own PTX and never consults the determinism flag. With the flag on, the speedup was 8.52×. With it off: 7.24×. We went with 7.24× — the honest number. We have a context manager that saves the flag state, turns it off for timing only, and restores it in a finally block.

The AST checker. Seven rules, purely syntactic, running before any execution. The most satisfying one catches "decoy kernels" — a @triton.jit function that exists in the source but is never called. The actual compute path just calls the library function. The Sakana AI paper documented exactly this pattern.

The Skill Transfer Discovery

This is the "wow" moment of the project, and I'm not going to pretend I planned it.

We ran the same profiler across three architecturally diverse models: Qwen2.5-1.5B (RMSNorm, a decoder LLM), GPT-2 (LayerNorm, an older decoder), and ResNet-50 (BatchNorm, a vision model). In every case, the normalization layer came back as the top optimization target.

Then we pointed the retrieval system at GPT-2's LayerNorm — a different operation on a different model — and it retrieved the RMSNorm kernel that was written for Qwen2.5. Three skills came back, with vector distances of 0.012–0.015. Because the bottleneck fingerprint contained no model name and no operation name — just op_family=norm, hardware=L4, is_memory_bound=True — the system found the match on its own.

A name-keyed cache can't make that jump. A system that stores kernels by "rmsnorm on qwen2.5" would never retrieve anything for "layernorm on gpt2." But a system that stores them by "memory-bound normalization at this arithmetic intensity on this hardware" can. That's the whole mechanism, and it's one line of code in the fingerprint builder.

Three Bugs That Passed All Unit Tests

These are the most useful things I learned. All three had full rows of green checks while being completely wrong.

Bug 1: Hot-swap couldn't load Triton kernels. The loader used exec(compile(source, ...)), but @triton.jit calls inspect.getsourcelines at decoration time. With no file on disk, it raises ValueError: @jit functions should be defined in a Python file. Every single kernel this system produces is a Triton kernel, so 100% of hot-swaps were failing. Unit tests used pure-torch stand-in kernels with no @triton.jit, so they all passed. The fix: write to a real file, import via importlib, register in sys.modules before exec_module (so inspect can resolve it while decorators run), and never delete the file (Triton re-reads the source when specializing for new shapes).

Bug 2: dict[str, str] made the LLM output empty. The adapter mapping — the deployment contract, the novel contribution — was typed as dict[str, str], which compiles to a JSON schema with no named properties. Gemini's structured output had nothing to anchor on. It emitted {} on every draft — 0/3 trials. The generic adapter never ran once. It silently fell back to the human-written bridge. The fix: list[AdapterBinding] where AdapterBinding has two named string fields. 3/3 trials filled correctly. General rule for the codebase going forward: never put a free-form dict[str, str] on a boundary an LLM has to fill.

Bug 3: The baseline was rigged and we didn't know. measure_baselines() ran with torch.use_deterministic_algorithms(True) still on. That made the speedup 8.52×. The real number was 7.24×. We caught this during integration testing, and the fix is the context manager I described above. This one is worth telling because it moves the headline number down.

The Stack

Everything runs on Google Cloud, all in-process:

• Agent framework: Google ADK 2.7.1 (in-process only, no A2A, no separate services)
• Models: Gemini 3.7 Flash (all four agents), Gemma 4 26B (kernel explainer), gemini-embedding-001 (768-dim, manually L2-normalized)
• Memory: Firestore Native with Vector(768) COSINE index and composite pre-filter on op_family + hardware
• Compute: single g2-standard-4 SPOT instance with 1× NVIDIA L4 (24GB, 300 GB/s bandwidth)
• Serving: FastAPI + uvicorn, Qwen2.5-1.5B-Instruct
• Dashboard: Streamlit with live event streaming and JSONL replay
• Deployment: two Cloud Run services (dashboard + React explorer), both scale to zero
• Tests: 716 total (698 unit + 18 integration on the L4)
• Cost: ~$53 total against $150 hackathon credits. Fresh reproduction: under $1

The entire agent tree uses temperature=0 and seed=42 on every model call. Every dependency is ==-pinned and the full 127-package transitive closure is locked by uv.lock. make demo reproduces the run on a fresh L4 VM for under a dollar.

The Numbers

Metric	Value
Speedup vs eager PyTorch (TF32)	7.24×
Speedup vs torch.compile	1.39×
Verifier reward	+3 (maximum)
Correctness checks passed	15/15
Iterations to converge	1 (warm skill library)
Modules hot-swapped	57 (all Qwen2RMSNorm layers)
Inference before swap	529ms for 3 tokens
Inference after swap	87.5ms for 2 tokens
do_bench parameters	warmup=150, rep=200, median

For LayerNorm on GPT-2 (tested separately): 7.45× vs eager, +3 reward. The hot-swap correctly refused — Qwen2.5 has no LayerNorm modules, and the system refused rather than reporting a zero-module success. That's the anti-fake-speedup guard firing on a real trace.

What I'd Build Next

The obvious extension is broader operation coverage. Right now, the system does normalization layers (RMSNorm, LayerNorm) well. SiLU/SwiGLU activation and RoPE positional encoding are in the OP_REGISTRY but haven't been optimized in the demo. The architecture supports them — it's the same profile→write→verify→swap loop — but each new operation family needs its own reference implementation in the registry and appropriate correctness shapes.

The more interesting extension is the bandit. Right now it's UCB1 with a cold start. With enough runs across enough models, it could learn hardware-specific strategies: "on an L4, for memory-bound norms, always start from the single-pass fused pattern with BLOCK_SIZE = next_power_of_2(hidden_size)." That's a genuine learning-to-optimize trajectory, not just code generation.

Try It

• Explorer: gpuyantra-explorer on Cloud Run — three-model audit comparison, winning kernel source, Gemma's explanation
• Dashboard: gpuyantra-dashboard on Cloud Run — replay a real optimization run step by step
• Code: github.com/KaustubhUp025/gpuyantra

Built with Google ADK, Gemini 3.7 Flash, Gemma 4, Firestore, Vertex AI, Compute Engine, and Cloud Run. 716 tests. The whole thing reproduces for under a dollar on a fresh VM.

---

I created this for the purposes of entering the All Things Agentic Hackathon. #AllThingsAgenticHackathon