Hyper-optimized megakernel-hybrid transformer inference

Headline results

We achieve a 5.76x geomean speedup vs. eager PyTorch and a 2.34x geomean speedup vs. torch.compile on H200 NVL. Additionally, we achieve 1.51x geomean speedup vs. the fastest torch.compile preset on each case.

All implementations pass the numerical gate (atol <= 0.002 OR rtol <= 0.02), across input_scale values spanning seven orders of magnitude from 1e-3 through 1e+3. Case 14 was excluded from the above calculations because materializing full-sequence attention scores in HBM would exhaust all device memory.

Environment

Inference hardware

We optimized our kernels on H100 NVL and H200 NVL. Our team did not have access to any suitably fast GPUs in a sufficiently controlled environment for benchmarking, therefore we opted to use our school's compute cluster instead, which had these cards available.

While many frameworks have been optimized for these Hopper cards, we do not believe that specialized megakernels for these particular problem specifications have been made. Furthermore, our use of production hardware means that the insights from our optimization journey and our software stack are directly applicable to real-world transformer inference.

Runtime stack

• OS: Ubuntu 24.04.4 LTS
• CPU: AMD EPYC 9334 for H100 nodes, AMD EPYC 9355 for H200 node
• Software stack: Python 3.12, CUDA 12.9, Triton 3.7.1, PyTorch 2.13.0+cu129

More details available in README.

Methodology

We decompose the 14 given test cases into three families. Each family is handled by its own megakernel or hybrid architecture. GPU kernels were mostly written in Triton, with a small number in CUDA C++.

Problem decomposition

Sequence-resident megakernel: cases 1, 5, 6, 7, 9, 10 and 11

These cases have short sequences that can fit fully within SMEM. Each CTA owns one sequence. Original 2-CTA cluster implementation for B=64 was scrapped due to poor Triton support; a 1-CTA implementation was used instead, with extra CTAs being used to assist with K projection and attention tiles while resident CTAs computed V projection tiles.

Persistent DAG-scheduled megakernel: cases 2, 3, 4 and 12

Batch size is tiny, so using one CTA per sequence would leave most SMs idle. Therefore we opt for a persistent kernel design. Task graph is created ahead of time for each problem shape; each CTA pops tasks off the queue and executes them until none are left.

Hybrid: cases 8, 13 and 14

These cases are heavily compute-bound, which means launch overhead is insignificant. Tensor-core MMA work is plentiful due to large model dimension (case 8) or seqlen (cases 13, 14). Furthermore, input sequences are too long to fit in SMEM; a sequence-resident megakernel would require extensive cross-CTA sync. Therefore we use a hybrid path instead, which incorporates optimized cuBLAS GEMMs, Flash Attention 3, and cuDNN. A few custom kernels are used: for instance, case 13 hybrid uses a custom Triton QKV matmul into which a head-major transpose is fused.

Benchmarking loop

Agentic optimization was extensively employed during kernel development. Instructions given to the agent for each kernel family were roughly as follows:

1. Identify the binding resource and its theoretical maximum.
   • Compute-bound kernels were optimized against max theoretical FLOP/s.
   • Bandwidth-bound kernels were optimized against max theoretical HBM bandwidth.
   • Latency-bound kernels were optimized against execution time only.
2. Profiling loop: iterate on optimizations, and time and profile each attempt. Numerical correctness is checked for every optimization, and those that fail the error tolerance are rejected.
3. Execute final regression tests and write documentation for current family.

Apart from measuring execution time, the profiling loop collected metrics from Nsight Compute and triton.testing. Instructing the agent to use Nsight metrics to guide optimization increased kernel execution speed by up to 14%.

Agent use strategies

Megakernel prior and problem decomposition

From the beginning, it was obvious that a megakernel would be the fastest strategy for nearly all of the test cases. Therefore, the agent was directly told to evaluate megakernel suitability for each case, and implement megakernels for cases that fit.

Telling the agent to consider a megakernel-based strategy while also allowing it to feedback on megakernel suitability directly enabled it to decompose the problem into sequence-resident megakernel, DAG megakernel, and hybrid specializations. Implementing 14 cases in 14 separate agent turns would have been token- and time-inefficient, especially when the same sequence-resident technique turned out to generalize across 7 of 14 problem shapes.

On the other hand, getting it to implement all kernels in one turn would have been counterproductive: every kernel is bound by different resources, and optimizing solely against full 14-case geomean speedup totally neglects the complexity and inter-case variation of the problem. In particular, shoe-horning clearly compute-bound problem shapes (cases 8, 13, 14) into a megakernel would have duplicated existing libraries' optimization work for practically zero gain. (If our problem is compute-bound, and FLOPs are dominated by attention/GEMM, we should just call into FA3/cuBLAS respectively!)

Therefore, the middle road of grouping the 14 given cases into a small number of families was the optimal approach, rather than forcing the agent to write 14 separate kernels, or "one kernel to rule them all".

Scratchpad folder

The agent was given an llm-scratchpad directory to write Markdown notes for consumption by future agent sessions. This made it possible to avoid re-exploring the entire codebase for each new agent session, saving time and decreasing context window utilization. Each turn was documented in a single file, which included:

• final kernel parameters (CTA count, QK/PV tile sizes, etc.)
• final execution time and speedup vs. eager/compiled
• rejected optimizations
• reproduction scripts

This proved incredibly useful for porting insights across agent sessions. For instance, once Triton support for multi-CTA layouts and DSMEM was found to be poor, later sessions were able to avoid that pitfall.

Instruction to optimize against % of roofline

This technique was inspired by KernelBench. Instead of blindly telling the agent to "make the kernel faster", we instructed the agent to identify bottlenecks by arithmetic intensity calculation and Nsight Compute / Triton profiling, then optimize the utilization of the binding resource against its theoretical maximum. Profiler results provide a strong heuristic for which code changes are of highest value, thereby allowing the agent to massively optimize its traversal of the kernel parameter search space.

For instance, Nsight Compute profiling indicated low occupancy for cases 1 and 10 after pivoting from 2-CTA to 1-CTA per sequence, and so 64 extra blocks were launched to assist with K projection and attention tile calculation. This yielded a 8.6% speedup for case 1 and a 4.1% speedup for case 10 relative to the prior 1-CTA implementation.

Tools used

• GPT 5.6 Sol with xhigh reasoning effort in Codex CLI harness
• Neovim IDE
• Nsight Compute profiler
• Triton benchmarking suite
• Slurm for cluster job submission

Challenges

• Poor Triton support for sm90 thread block clusters blocked originally planned 2-CTA implementation for cases 1, 9, 10.
  • Many multi-CTA-layout ops failed MLIR lowering due to lack of Triton support. Final 2-CTA megakernel implementation had to substitute all these ops for 1-CTA equivalents, but still lost to 1-CTA due to presence of several cross-CTA barriers and duplicated MMA work. 1-CTA per sequence + 64 assisting CTAs remains the production default.
• Low GPU availability on school cluster
  • Resorted to doing numerical correctness testing on less contended H100 47GB MIG nodes. MIG nodes were not used for performance measurements.
• No hardware counter access on H200 node
  • Use triton.testing.Benchmark instead to estimate roofline % of binding resource
• Complying with the numerical gate incurred significant extra effort.
  • Several implementations passed at input_scale == 1.0 but failed at larger and smaller scales. Our final implementations pass at all orders of magnitude from 1e-3 to 1e+3.
  • Lack of numerical parity with provided reference implementation prevented direct application of optimized libraries such as Flash Attention 3, cuDNN, and full-graph torch.compile. Tiny numerical disparities arising from differences in reduction order, rounding behavior, and choice of CUDA numerical primitives compounded across layers, so optimized implementations could at best be used for the final 1 or 2 layers.
  • Several of our LayerNorm implementations had to mirror the exact Welford's algorithm tree used by PyTorch in order to achieve numerical parity. Extensive cross-checking with PyTorch source was performed to achieve this.

Appendix

Sample prompts

Instruction to pivot cases 1, 9, 10 away from 2-CTA approach

Notes on a previous attempt at 2-CTA documented in `llm-scratchpad` indicates that Triton support for this path is seriously lacking. This suggests that we might have to pivot cases 1, 9, 10 (documented in initial_notes.md) either to a 1CTA path or another library/DSL (e.g. CUTLASS, CuTe DSL, or even CUDA C++). Work out a viable plan for step 3 of the initial benchmarking order, and implement it. Remember, the aim is to make cases 1, 9, 10 as fast as possible on the H200 (H100 are almost completely occupied) — use `--gres=gpu:h200-141:1`. Hunt for potential algorithmic optimizations, and iterate until max possible speedup is achieved. Since hardware counters are not available on the H200 node, use `triton.testing.Benchmark` to obtain % of roofline (either compute or memory, whichever binds), and use that figure to optimize speedup. Bit-exactness should not be a priority if the kernel is already under the numerical gate.

Case 14 hybrid implementation, initial prompt

Implement step 8 in the initial benchmarking order described in `llm-scratchpad/initial_notes.md`. Make the relevant test case AS FAST AS POSSIBLE on H200-141 (falling back to H100-96 is acceptable if H200 node is occupied), adhering to the optimization and profiling methodology in AGENTS.md. Iterate and profile until max possible speedup is achieved. Megakernel was previously deemed unsuitable for this configuration; check the justifying logic. For this test case, making use of external tuned frameworks such as cuDNN frontend and Flash Attention 3 is ENCOURAGED; custom kernels may still be written if necessary. H100-47 MIG nodes may be used for numerical correctness testing ONLY; perf numbers from those nodes should not be considered representative because their roofline numbers are different.

Nsight Compute profiler pass prompt

Currently, test cases 5, 6, 1, 9, 10, 2, 3, 4, 12 are working, each handled by one of two megakernel families. Profile all test cases using `/usr/local/cuda-12.9/bin/ncu` on H100-96 cards (not H200 — H100 nodes are the only Hopper nodes with hardware counter access enabled; use `--gres=gpu:h100-96:1`) and hunt for potential algorithmic optimizations. Iterate until max possible speedup is achieved. Remember: we are building MEGAKERNELS here, therefore any device-wide synchronization is not acceptable — a megakernel must overlap different phases of the whole computation.

Repo reorganization

Great work. The repo is currently a little messy; let's reorganize. Benchmarks go in their own folder, as do kernels and profiling scripts. Separate Slurm job outputs from job scripts — outputs go in `job-scripts/outputs`; existing scripts should be changed to output to that folder instead. The optimized transformer should be callable from a single function that handles all (implemented) cases. After the restructuring is complete, regression-test everything to make sure imports still work fine and all cases retain correctness and perf.

Built With

• c++
• cuda
• python
• pytorch
• triton