GPUyantra

Built an agent that watches your model run, writes a faster GPU kernel, proves it's not cheating, and swaps it in live — no restart. 7.24× speedup on a single L4.

Comment

Inspiration

I'm a Systems Engineer with a background in distributed systems and while experimenting with GPU inference and optimization — profiling models, understanding roofline models, reading PMPP and CS:APP — I stumbled into writing custom Triton kernels. And immediately noticed the pattern: profile the operation, figure out whether it's memory-bound or compute-bound, check the arithmetic intensity against the hardware's ridge point, write a fused kernel that keeps data in SRAM instead of bouncing it to HBM, test correctness across shapes and seeds, benchmark against honest baselines, and finally wire it back into the model. Every time. The same manual loop. That's when it clicked — this loop is mechanical enough for an agent to own. Not just the code generation part (any LLM can spit out a Triton kernel), but the entire cycle: profiling the bottleneck, deciding what to optimize, writing the kernel, judging whether it's correct and genuinely faster, learning from the result, and deploying it live. An explore → decide → code → judge → learn loop that improves with each iteration. Building developer tools that automate tedious expert work is the thread that connects all my projects. gpuyantra is that same idea applied to GPU kernel engineering.

What it does

gpuyantra is an autonomous, in-process Google ADK agent tree that profiles a running GPU model, identifies the computational bottleneck, writes an optimized Triton kernel, verifies it through a multi-layer trust anchor, and hot-swaps it into a live inference server — while the model keeps serving requests. No restart, no weight reinitialization, no human in the loop.

The agent runs a 5-step autonomous workflow:

  1. Profile — the Profiler agent measures the target operation on the NVIDIA L4 GPU using triton.testing.do_bench (warmup=150, rep=200, median mode). It computes the arithmetic intensity and places the operation on the L4's roofline model. For RMSNorm on Qwen2.5-1.5B, the measured arithmetic intensity is ~1.25 FLOP/byte against the L4's ridge point of ~101 FLOP/byte — meaning the GPU's compute units sit idle ~81× longer than necessary, just waiting for data from HBM. That gap is the optimization headroom.

  2. Retrieve — the Supervisor embeds the profiler's output as a bottleneck fingerprint using gemini-embedding-001 (768 dimensions, manually L2-normalized) and runs a Firestore find_nearest vector search with a composite pre-filter on op_family + hardware. A UCB1 bandit selects which retrieved skill to seed the Coder with — balancing exploitation of proven strategies against exploration of new approaches.

  3. Write — the Coder agent (LlmAgent with output_schema for structured generation) produces a complete Triton kernel AND an adapter_mapping — the deployment contract that declares how kernel parameters map to model module attributes. This is the novel contribution: every published kernel deployment system (HuggingFace kernels, FlashInfer-Bench, Kernel Contracts from arXiv:2604.22032) uses human-authored deployment bridges. gpuyantra is the first where the agent generates the bridge and a deterministic verifier validates it.

  4. Verify — the Judge agent calls verify_kernel, which runs a 4-layer trust anchor:

    • Layer 1 — AST Checker: 7 syntactic anti-hack rules (chmod 444 protected). Catches reward gaming before any code executes: library function calls disguised as custom kernels, decoy @triton.jit functions never invoked, returning input unchanged, returning uninitialized GPU memory.
    • Layer 2 — Contract Validator: checks every attribute in the adapter_mapping against a torch.device("meta") module instance. Catches mapping errors without allocating GPU memory.
    • Layer 3 — Subprocess Sandbox: 5 random seeds × 3 input shapes = 15 correctness checks at atol=rtol=1e-2. Runs in a scrubbed environment (4 env vars only), SIGKILL on timeout, GPU health probe after execution.
    • Layer 4 — Reward Computation: recomputed in-process from verified measurements. The subprocess's own reward field is discarded — a kernel that prints {"reward": 3} in stdout still scores −1 if it's wrong. Reward ladder: −1 (wrong) / +1 (correct, not faster) / +2 (beats eager >5%) / +3 (beats both eager AND torch.compile >5%).

  5. Deploy — the Supervisor calls hotswap_kernel, which POSTs to the inference server's /swap endpoint. The swap uses types.MethodType to rebind forward on every matching module — 57 Qwen2RMSNorm layers in Qwen2.5's case. Weights stay in GPU memory, same dtype, same device. A parity check against live weights runs before commit, with automatic rollback on mismatch.

The Result

Metric Value
Speedup vs eager PyTorch (TF32 enabled) 7.24×
Speedup vs torch.compile(mode="reduce-overhead") 1.39×
Verifier reward +3 (maximum)
Correctness checks passed 15/15 (5 seeds × 3 shapes)
Iterations to converge 1 (warm skill library)
Modules hot-swapped 57 (all Qwen2RMSNorm layers)
Inference latency before swap 529 ms (5.7 tok/s)
Inference latency after swap 87.5 ms (22.9 tok/s)
do_bench config warmup=150, rep=200, median

The Twist: Bottleneck Fingerprint Indexing

Skills in the Firestore library are NOT indexed by operation name or model name. They're indexed by a bottleneck fingerprint: op_family=norm, is_memory_bound=True, arithmetic_intensity=1.2, tile_size=1024, hardware=L4. No model name. No operation name. Just the physics of why something is slow.

This enables cross-model skill transfer. We proved it: pointed the retrieval system at GPT-2's LayerNorm (a completely different operation on a different model architecture) and it retrieved 3 Qwen2.5 RMSNorm skills at vector distances of 0.012–0.015. Because both operations have the same bottleneck shape — memory-bound normalization at similar arithmetic intensity on the same hardware.

A name-keyed cache (rmsnorm_qwen2.5 → kernel) can never make that jump. A fingerprint-keyed one (memory_bound_norm_L4 → kernel) can. That's the architectural innovation, and it's one field in the fingerprint builder.

Multi-Model Audit

We validated generalization by running the profiler across three architecturally diverse models:

  • Qwen2.5-1.5B (RMSNorm, decoder-only LLM) — 7.24× speedup, +3 reward
  • GPT-2 (LayerNorm, older decoder) — 7.45× speedup, +3 reward. Hot-swap correctly refused (Qwen2.5 has no LayerNorm modules — the system refused rather than reporting a zero-module success)
  • ResNet-50 (BatchNorm2d, vision CNN) — profiled as optimization target, BatchNorm2d identified as audit-only (different op family)

In every case, the normalization layer came back as the top optimization target. Cross-model transfer was confirmed.

How I built it

Agent framework: Google ADK 2.7.1, fully in-process (no A2A, no separate services, no LangGraph). Four LlmAgents (Supervisor, Profiler, Coder, Judge) plus a deterministic EscalationChecker (subclass of BaseAgent, no LLM call — uses reward >= 3 as a pure Python exit condition). The RefinementLoop is an ADK LoopAgent with max_iterations=6. The Supervisor runs a 7-step resumable protocol with tools: retrieve_skills, upsert_skill, hotswap_kernel, explain_kernel.

Models: gemini-3.7-flash (GA on Vertex AI global endpoint) for all four agents at temperature=0. gemma-4-26b-a4b-it-maas as a kernel explainer agent that generates plain-English descriptions of optimized kernels for the explorer page. gemini-embedding-001 at 768 dimensions — truncated vectors require manual L2 normalization before cosine similarity produces correct results (undocumented behavior we discovered during integration).

Memory: Firestore Native (us-central1) with Vector(768) field, COSINE distance, flat index. Composite pre-filter on op_family + hardware (Firestore vector search doesn't support inequality pre-filters — we designed around this constraint). Three collections: skills (the kernel library), runs (agent execution logs), traces (JSONL event streams for dashboard replay).

Compute: Single Compute Engine g2-standard-4 SPOT instance with 1× NVIDIA L4 (24 GB VRAM, 300 GB/s HBM bandwidth, 30.3 TFLOPS non-tensor FP16). --instance-termination-action=STOP to survive SPOT preemptions. Total project cost: ~$53 against $150 hackathon credits. Fresh make demo reproduction: under $1.

Serving: FastAPI + uvicorn serving Qwen2.5-1.5B-Instruct in FP16. The served model is never torch.compile'd — compiled graph dispatch bakes the current forward into its execution graph; a later types.MethodType patch silently no-ops. The only torch.compile in the codebase runs inside measure_baselines(), where the compiled reference is timed for comparison and discarded.

Dashboard: Streamlit with live ADK event streaming (SSE) and JSONL replay mode for the hosted Cloud Run version. Plain-English event descriptions, tokens/s metric, iteration counter, three-bar speedup chart, narrative captions. The dashboard serves as both a development monitoring tool and a judge-facing demo surface.

Deployment: Two Cloud Run services, both --allow-unauthenticated, both scale to zero:

  • gpuyantra-dashboard — Streamlit replay (512 Mi)
  • gpuyantra-explorer — nginx-served React SPA with three-model audit comparison, winning kernel source, Gemma 4's kernel explanation (256 Mi)

Reproducibility: temperature=0 + seed=42 on every Gemini call. torch.manual_seed(42), numpy, Python random all seeded. CUBLAS_WORKSPACE_CONFIG=:4096:8. All 127 dependencies ==-pinned, transitive closure locked by uv.lock. make demo reproduces the full optimization run on a fresh L4 VM.

Tests: 716 total — 698 unit tests + 18 integration tests that require the L4 GPU. The integration tests cover the full pipeline: profiling, kernel generation, verification, hot-swap, and inference measurement.

Architecture Diagram

Challenges I ran into

Three integration 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 in production.

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

Bug 2: dict[str, str] made the LLM output empty. The adapter_mapping field in the Coder's output schema was typed as dict[str, str], which compiles to a JSON schema with additionalProperties: {type: string} and no named properties. Gemini's structured output had nothing to anchor on — it emitted {} on 3/3 trials. The agent-generated deployment contract (the novel contribution) never ran once. The system silently fell back to a human-written bridge. Fix: list[AdapterBinding] where AdapterBinding has two named string fields (module_attr and kernel_param). 3/3 trials filled correctly. General rule for the codebase: never put a free-form dict[str, str] on a boundary an LLM has to fill.

Bug 3: The baseline was secretly rigged. measure_baselines() ran with torch.use_deterministic_algorithms(True) still active. This forces slower cuBLAS codepaths for eager PyTorch (~23% penalty) 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\times\). With it off: \(7.24\times\). Fix: a context manager that saves the flag state, turns it off for timing only, and restores it in a finally block. We went with \(7.24\times\) — the honest number.

The honest benchmarking problem

The initial KernelBench leaderboard had entries claiming \(3\times\)+ speedups that collapsed to under \(1.5\times\) when measured properly (one prominent result went from \(3.13\times\) to \(1.49\times\) under honest baselines). We adopted the KernelBench-Verified protocol: TF32 enabled via torch.set_float32_matmul_precision("high") and torch.compile(mode="reduce-overhead") also measured. The number to beat is PyTorch at its best, not PyTorch handicapped.

The AST checker catches seven reward-hacking patterns statically, before any code executes. The file is chmod 444 protected so the agent can't modify the checker to pass its own kernels. The most satisfying rule catches "decoy kernels" — a @triton.jit function that exists in the source but is never called from the forward path. The actual compute just calls F.rms_norm. This exact pattern was documented in Sakana AI's CUDA Agent paper.

What I learned

  • Bottleneck-fingerprint indexing enables cross-model transfer — the most important architectural discovery. Name-keyed caches are dead ends for generalization. Physics-based indexing (arithmetic intensity + memory access pattern + hardware) is the right key.
  • torch.compile and hot-swap are fundamentally incompatible. Compiled graph dispatch bakes the current forward into its execution graph; subsequent MethodType rebinding is silently ignored. The served model must never be compiled.
  • LLM-driven loop exit is fragile. The EscalationChecker uses deterministic Python conditions, not LLM judgment. Relying on the LLM to decide "should I stop?" introduces non-determinism and occasionally infinite loops.
  • gemini-embedding-001 truncated vectors require manual L2 normalization before cosine similarity works correctly. This is undocumented — without normalization, cosine distances are meaningless.
  • Integration bugs are the real killers. All three critical bugs above had full unit test coverage. They only surfaced during end-to-end pipeline testing on the GPU. Unit tests with mock stand-ins can't catch interface-boundary failures.
  • do_bench has a documented ~30% underestimation bug at low warmup/rep values. The fix is warmup=150, rep=200, return_mode="median".
  • Judges need narrative, not raw JSON. The dashboard's original developer-debug output (raw JSON events, stuck counters, missing tok/s metrics) was a losing presentation. Plain-English event descriptions with context changed the entire demo story.

What's next

  • Broader operation coverage — SiLU/SwiGLU activation and RoPE positional encoding are in the OP_REGISTRY but haven't been optimized in the demo. The architecture supports them — same profile → write → verify → swap loop.
  • Hardware-specific bandit strategies — with enough runs across models and GPUs, the UCB1 bandit could learn hardware-specific optimization patterns: "on an L4, for memory-bound norms, always start from the single-pass fused pattern with BLOCK_SIZE = next_power_of_2(hidden_size)."
  • KernelBench-Verified evaluation — running the full KernelBench L1/L2 suite with honest baselines to benchmark the agent against the public leaderboard.
  • Multi-GPU and multi-hardware — extending the fingerprint to capture hardware topology (NVLink, PCIe bandwidth) for distributed inference scenarios.

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

Built With

Share this project:

Updates

Submission history