FlipFloppers GPU Kernel Optimization w Agentic Research Loop

• 

  Results

• 

  Optimization Path

• 

  Agentic Workflow

• 

  Development Environment

At a Glance

• We built a hardware-aware optimization agent that researches, implements and validates GPU kernel optimizations on its own. Everything below is what it produced on a Transformer layer targeting a single consumer RTX 4090.

• A kernel-level rewrite achieving 1.93x-31.76x per shape, ~7.7x geometric mean, 6.3x aggregate latency reduction, within a rtol<0.02 or atol<0.002 accuracy budget, with the compute-bound path running at 94-96% of the accuracy-legal hardware roofline.

• We run a sequence length the reference implementation cannot. At S=100,000 the provided benchmark dies at OOM while our automatically sequence-chunked version completes in 9952ms on 24 GB of VRAM.

• Optimizations are escalated through a tier system, with lossy techniques weighed against an accuracy-performance tradeoff heuristic rather than accepted on speed alone.

• Fairness is enforced through Slurm resource reservation, Apptainer environment pinning, GPU clock persistence and thermal monitoring. An anti-gaming checker verifies that candidates actually compute the answer instead of caching a previous one, because an agent optimizing against a benchmark will exploit it if you let it.

• Final optimization employs shape-based arbiter strategies, with algorithmic fusions, hybrid reduced precision, memory manipulation and low level optimizations at the CUDA / PTX level.

Inspiration

GPU kernel optimization can only go so far without hardware awareness. Deep optimizations come from hardware intrinsics and real-world measurements on actual hardware in controlled environments, which is exactly the kind of knowledge that does not survive being written down once and reused. Every new architecture invalidates part of it.

So we did not set out to hand-write the fastest Transformer layer we could. We set out to build the thing that writes it, and then used this workload as proof that it works: targeting the dominant GEMM operations while balancing precision loss and reducing bottlenecks in attention, memory movement and kernel launch on a single-node RTX 4090.

Results

For the 14th test shape, sequence chunking was introduced to run the 100,000-token sequence within the RTX 4090's 24 GB of VRAM, processing the sequence against an incrementally filled FP16 KV cache. This is the one row where the comparison is not a speedup but a capability: the reference cannot produce an answer at all.

On compute-bound workloads, our isolated GEMMs reached 94-96% of the accuracy-limited hardware roofline. Memory-bound paths sit at 96% of their measured ceiling (22.7 of 23.6). Both numbers are against the accuracy-legal roofline, not the datasheet peak: the FP16/FP32-accumulate throttling on this card and the atol budget both move the ceiling down, and we measured where it actually lands rather than quoting the spec sheet.

Agentic Auto Research System

The agent is the actual deliverable. It summons multiple subagents with different roles and limitations, each handling one part of the research to implement to validate cycle, and manual research can be injected as a new candidate at any point. Every candidate is applied as one diff per run and evaluated against the official harness, with anti-gaming and bias controls built in. Refer to the figure for more information.

• Stage Based. Optimization effort escalates through stages, categorising techniques by type and risk level, so cheap safe wins land before anything touches PTX.

• Auto Research with Manual Input. Facilitates wide exploration of optimization techniques without locking us out of the loop when we have a hunch.

• Comprehensive Validation Loop. Testing on the official harness, with anti-gaming mechanisms to prevent the agent from gaming the benchmark. This matters more than it sounds: a large share of the candidates it generated were fast and wrong, and the loop is what makes the 6.3× trustworthy rather than just reported.

• Token Conscious. A single Sonnet orchestrator summons subagents on models appropriate to the task assigned, rather than spending frontier tokens on bookkeeping.

The rejection log is as informative as the acceptance log. Candidates thrown out for being too lossy included inline-PTX mma.sync FP16 GEMMs, degree-7 GELU polynomials, and BF16/FP8/INT8 precision reductions — all of which are "wins" if you only look at the clock.

Challenges

• Hardware Architecture: We were limited to a single-node RTX 4090 (Ada Lovelace), meaning no distributed compute and no access to modern Hopper intrinsics like WGMMA, TMA, or DSMEM. The agent had to find its speedups inside a smaller instruction budget than a datacenter part would give it.

• VRAM & Performance Locks: FP16 Tensors with FP32 Accumulation is artificially throttled on the RTX 4090. Strict VRAM constraints also ruled out processing long sequences directly, forcing overhead-heavy sequence chunking.

• Strict Accuracy Budgets: The project demanded an absolute tolerance (atol) of 0.002. Many generated optimizations were rejected for being too lossy, and our tightest passing shape lands at max_abs = 0.00195, 97.5% of the budget consumed. There was very little room left.

• Toolchain and Dependencies: The original test environment had to be rebuilt to support everything the agentic system needed: Slurm, CUDA, Apptainer and UV.

What's Next

• Expanding to Full Decode. Instead of just prefill, we plan to extend optimizations to the full autoregressive generation loop through KV-cache work like Multi-Headed Latent Attention.

• Broadening Support. The agent's search strategy, the tier system, the shape-based arbiter and the fusion candidates are architecture-agnostic and should port directly. What does not port is the bottom tier, the PTX-level work and the hybrid precision choices, both of which were tuned around Ada's specific throttling behaviour. Moving to Hopper means giving the agent a new intrinsic vocabulary (WGMMA, TMA, DSMEM) and re-measuring the roofline, not rewriting the agent.

• Scaling for TikTok. The reason to automate this is that the knowledge expires. Each new GPU generation resets which optimizations are legal and which are fastest, and hand-tuning every kernel again is not a strategy that scales with the hardware release cadence. A hardware-aware agent that re-derives the optimizations from measurements on the target card is, and that is what we would want pointed at real-time inference workloads at TikTok's volume.

Built With

• apptainer
• claude
• cublas
• cuda
• cutlass
• ptx
• pytoch
• slurm