LLM-Driven Kernel Composition: Synthesizing Fused GPU Training Passes for Transformer Blocks

Abstract

Training a transformer at scale is a memory bandwidth problem as much as a compute problem. Libraries like QuACK provide hand-optimised GPU kernels for every operation in a transformer block, but using them at full efficiency requires non-obvious wiring decisions: which primitives to call, in which order, and with which arguments to eliminate intermediate HBM round-trips. Those decisions require hardware knowledge that does not appear at the Python level, and making them correctly across an entire transformer block is what separates a fast implementation from an optimal one. This thesis asks whether a large language model can make those decisions automatically. The task is kernel composition: given the QuACK API and a reference PyTorch implementation of a LLaMA-3 8B transformer block, synthesise a fused torch.autograd.Function using only existing library primitives, with no new kernel code written. We propose a two-stage approach that separates fusion planning from code generation. In Stage 1, the LLM produces a kernel breakdown table showing every operation with its HBM reads, writes, and register contents. In Stage 2, it implements the function from the agreed table. This separation makes fusion decisions inspectable before any code is written and gives the human a natural point to intervene at the planning level. The resulting implementation achieves a 1.14× speedup over PyTorch eager on a single H100 at LLaMA-3 training shapes, with the advantage growing linearly with layer count. The central finding is that fusion decisions determinable from the API surface are handled automatically by the LLM; decisions that depend on runtime performance consequences require human steering. The approach reduces the barrier to building hardware-optimal fused kernels, but does not eliminate the need for someone who can interpret what the profiler is telling them.