Problem

Transformers are the compute bill. Projections, attention, softmax, normalization, feed-forward — the same operations, millions of times a day, in every serving stack. The task hands us that layer, fixed, and asks for GPU kernels that run it faster than the reference without changing what it computes.

Solution

Applying loop engineering to the problem with agentic research and kernel optimization techniques given a set of constraints that the loops respect. Final solution is multiple kernels.

Context and Reasoning

Firstly, we realized this problem exist within the context of data centers and optimizing the transformer layer for a specific GPU. Hence, we started up a virtual private server with a data center-grade GPU, more specifically a A100 SXM GPU on RunPod.

We also realized that - Different input shapes have different bottlenecks, one candidate can perform well for one case but very poorly for others, hence the decision to push for one global generalist and a few per-shape specialist. The identified Shape → Bottleneck :

Shape regime Dominant bottleneck What wins
Small / batch=1 CUDA launch + host overhead Launch elimination, graph replay
Large batch, wide model Throughput and bandwidth Tensor-core-friendly GEMM/FFN paths
Long sequence Attention memory traffic Blocked, FlashAttention-style attention
Case 14, seq=100000 Basic feasibility Never materialize the N×N score matrix

Implementation

Instead of hand-writing one optimized layer, we built an autonomous research loop in which an AI model is the researcher and the repository is the laboratory. Every turn of that loop produces one new candidates one immutable evidence record. The loop runs a closed cycle: read evidence → state one falsifiable hypothesis → write one candidate kernel → verify correctness → measure on real A100 → promote or reject → record

The layers

Judgement

Layer 1 · Contract A config file promoted to a constitution. It fixes the target hardware and dtype, the case scope, the promotion mathematics, the ownership split, and an explicit prohibited list. The hash is the load-bearing part: a human confirms one specific hash, and a previous session's approval is void the moment the contract changes.

Layer 2 · Research - Codex subagents Four roles - lead, researcher, skeptic, synthesizer — up to four in parallel, with Exa as the only approved source connector. The important mechanism is the packet validator's rejection of benchmark fields: a validated research packet is structurally forbidden from containing a number. No scores, no per-case timings, no promotion outcomes. External evidence can motivate a hypothesis but will never become evidence for a candidate. Ingestion fails loudly rather than quietly letting a citation stand in for a measurement.

Layer 3 · Generation - candidates/vNNN_*.py Candidates are append-only in the same way the ledger is, which is why v000 — the identity implementation — is still present as a live noise control rather than a remembered number.

Deterministic mechanism

Layer 4 · Routing Before a hypothesis is formed, the router computes each shape's gap to speed-of-light - measured latency over the larger of that shape's memory-bound and compute-bound floors — and routes the next move on that gap:

Gap to SOL Mode What it means
< 1.5× done No tuning headroom left on this shape
1.5× – 4× TUNE The structure is right and the parameters are not: block sizes, num_warps, num_stages, vectorization, occupancy
> 4× STRUCTURAL Latency-bound. Neither bandwidth nor FLOPs is the constraint, so tile sizes cannot help — the move is to eliminate launches, fuse operations, remove synchronization, or replace the algorithm
every shape < 1.5× STOP The search is finished, and says so

The underlying argument is one line: autotuning moves a kernel toward the roofline it is already on; it cannot move the roofline. So gap-to-SOL decides tuning versus restructuring, and gives the loop a genuine stopping criterion instead of "keep going until the budget runs out."

The router does not merely pick a mode — it hands the researcher the reasoning: which cases sit furthest off their roofline, whether each is memory- or compute-bound, and why the rejected mode would not have helped. The model is prompted with an argument from hardware limits rather than an instruction to try harder.

Two subordinate rules bound the search's own failure modes, and we added both after catching ourselves in them:

  • Width. Beam and elite widths scale with the evidence actually available. A beam of three over four candidates is keep everything wearing a costume.
  • Bias. Screening on a hypothesis-chosen profile selects for candidates that happen to suit that profile. So the loop forces a neutral profile every third cycle and measures how often a screen winner fails to survive the full sweep. It instruments its own selection bias rather than assuming it away.

Layer 5 · Tuning The mechanical autotuner: enumerate the valid launch-configuration grid, measure every variant, freeze the best configuration per shape. Because it costs no model tokens, routing exhausts tuning before paying for a new idea — and correctly reports nothing to tune when the reigning champion is pure PyTorch with no launch parameters to sweep. Every variant runs inside one pod process so the comparison stays paired (same clocks, same allocator state, same device), and each variant-and-shape pair is re-checked for correctness, because a bad block size is fast and wrong.

This is a Halide-style algorithm / schedule split: the model searches algorithms, the autotuner searches schedules for whatever algorithm it is handed.

Layer 6 · Execution - RunPod MCP Ran on RunPod instance, connected the agents via MCP so the harness includes tools to start and stop with guardrails in place since the instance is a paid service.

Evaluation, the Oracle

Layer 7 · Evaluation The hardened measurement path described above: five input distributions, dual event-and-wall timing with the ratio gate, environment snapshots either side of the run, and a fingerprint of the model's own weights. A challenger is first screened on the case family its hypothesis actually targets, and only the strongest survivor earns the expensive full 13-case sweep — but losing screen runs still enter the archive, which is precisely how a global loser becomes a per-shape specialist. The paired bootstrap is implemented from scratch rather than pulled from a statistics package, so no dependency can silently change a number between runs.

Layer 8 · Selection The portfolio memory. It maintains a global beam of the best statistically promoted candidates, which become the parents the next hypothesis mutates, and alongside it a per-shape archive of the fastest valid implementation for every official case. Selection is deliberately not winner-take-all: an implementation that loses the overall sweep can still hold a shape outright, and the archive is what preserves it.

The scope guards here matter as much as the ranking. Every result carries the hardware, dtype, and case set it was measured under, and the layer refuses to rank across them — which is what stops a three-case half-precision result from ever being compared against a thirteen-case single-precision one.

Layer 9 · Persistence and assembly The canonical history is a single append-only event log with sequence validation and file locking: nothing is ever rewritten, and every promotion, rejection, and measurement in this description can be replayed from it. The human-readable result and status pages are projections regenerated from that log.

The final step is mechanical. The submission file is compiled from the ledger rather than assembled by hand: the generalist and its proven specialists are resolved into a single model-construction entry point, so what the judges run is a direct function of the recorded evidence.

Development tools used

Tool Role
Claude Code (CLI) — Opus 5 and Sonnet 5 Harness, event store, dispatcher, controller and test-suite engineering; nightly health checks. Commits carry Co-Authored-By trailers, so the split of work is auditable in git log
OpenAI Codex (CLI) Lead researcher: evidence review, hypothesis generation, candidate authoring, result interpretation, direction selection. Runs as native Codex subagents in four roles — lead, researcher, skeptic, synthesizer — up to 4 in parallel
RunPod All paid measurement. A single NVIDIA A100-SXM4 pod, driven through the connected RunPod MCP server for lifecycle and through our own tools/podrun, podput, podget, podsync scripts for execution and transfer (the pod's SSH proxy is interactive-only, so podrun feeds a base64 payload over an interactive bash in sub-4096-byte chunks)
Git / GitHub The evidence trail. Each loop iteration is one commit naming the candidate and the resulting champion score
nvidia-smi and NVIDIA profiling on the pod Bottleneck attribution
Custom CLI tooling — tools/controller.py, race.py, harness.py, report.py The deterministic half of the system: everything the AI is not allowed to decide

APIs used

API Use
Codex-native subagents (Responses + Chat Completions) Powered the autonomous loop which called the model directly for each candidate with configurable reasoning effort.
Exa The single approved connector for web search and source fetching during evidence investigations. Every external source is returned with its URL, retrieval time, quality classification, and the specific claim it supports
RunPod API / RunPod MCP GPU pod lifecycle and control-plane operations
Anthropic API, via Claude Code Engineering of the deterministic infrastructure

Libraries and frameworks used

Library Use
PyTorch The whole substrate. torch.nn, scaled_dot_product_attention, nn.MultiheadAttention fast paths, nn.TransformerEncoderLayer, TorchScript (torch.jit.script / optimize_for_inference), torch.compile, and manual CUDA graph capture and replay (torch.cuda.CUDAGraph) — the technique behind the current champion
Triton Custom fused kernels: LayerNorm+residual fusion, fused FFN, fused attention-output projection. Explored across ~22 candidates; retained in the evidence archive, absent from the final portfolio, because these shapes turned out to be launch-bound rather than kernel-bound
Python standard library, for everything else statistics, random, json, hashlib, importlib, fcntl for ledger locking. The paired bootstrap CI is implemented from scratch in tools/search_stats.py — no NumPy, SciPy or pandas anywhere in the project. The measurement path has no dependency that could silently change a number between runs
pytest tests/ covers the portfolio logic, report generation, event-store migration, and orchestration

Built With

  • agentic-ai
  • cuda
  • cuda-graphs
  • pi-exa
  • pytest
  • pytorch
  • torchscript
  • triton
Share this project:

Updates

Submission history