cs.LGMay 19, 2026

CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs

Authors: Han Guo, Jack Zhang, Arjun Menon, Driss Guessous, Vijay Thakkar, Yoon Kim, Tri Dao

Organizations: 1Massachusetts Institute of Technology · *Part of this work was done at Together AI. · 2Princeton University · 4Meta · 3Together AI

Abstract

Transformer training systems are built around dense linear algebra, yet a nontrivial fraction of end-to-end time is spent on surrounding memory-bound operators. Normalization, activations, residual updates, reductions, and related computations repeatedly move large intermediate tensors through global memory while performing little arithmetic, making data movement an increasingly important bottleneck in otherwise highly optimized training stacks. We introduce CODA, a GPU kernel abstraction that expresses these computations as GEMM-plus-epilogue programs. CODA is based on the observation that many Transformer operators exposed as separate framework kernels can be algebraically reparameterized to execute while a GEMM output tile remains on chip, before it is written to memory. The abstraction fixes the GEMM mainloop and exposes a small set of composable epilogue primitives for scaling, reductions, pairwise transformations, and accumulation. This constrained interface preserves the performance structure of expert-written GEMMs while remaining expressive enough to cover nearly all non-attention computation in the forward and backward pass of a standard Transformer block. Across representative Transformer workloads, both human- and LLM-authored CODA kernels achieve high performance, suggesting that GEMM-plus-epilogue programming offers a practical path toward combining framework-level productivity with hardware-level efficiency.

Explore similar work

CardsList
Jul 21, 2026cs.LG

MoA-Structured Decode Attention DNF Derivation, KV-Cache Accumulation, GQA/MQA, and OpenACC Kernel

We derive four memory-optimal inference artifacts for transformer attention using the Mathematics of Arrays (MoA), each following directly from the forward-pass Denotational Normal Form (DNF) of with the query-row index fixed to the current decode step. The artifacts are: (1)~a single-query decode DNF in which the ψψ-reduction eliminates the K⊤K^\top buffer algebraically, achieving (dk+ndk+ndv+dv)×4 B(d_k + nd_k+ nd_v+ d_v)\times4\,{B} Dynamic Random Access Memory (DRAM) traffic result numerically verified to ∥err∥≤2×10−7\|{err}\|_\leq2\times10^{-7}; (2)~a C/OpenACC Graphics Processing Unit (GPU) kernel with Operational Normal Form (ONF) stride arithmetic and hardware-coalesced memory access, verified to ∥err∥∞=0\|\mathrm{err}\|_\infty=0 (exact IEEE-754 floating-point arithmetic); (3)~a multi-step KV-cache with O(dk+dv)O(d_k+d_v) per-step append via MoA concatenation #\#; and (4)~Grouped-Query Attention (GQA) and Multi-Query Attention (MQA) derived via ψψ-selection, achieving a proven hqhkv\frac {h_q} { h_{kv} } reduction in KV traffic. All programs are verified against PyTorch scaled_dot_product_attention.
Lenore Mulin, Gaetan Hains
Quadratic AttentionTransformer AttentionKey-Value Cache QuantizationGraphics Processing Unit KernelsVram
AskPDF
Apr 16, 2026cs.PL

Nautilus: An Auto-Scheduling Tensor Compiler for Efficient Tiled GPU Kernels

We present Nautilus, a novel tensor compiler that moves toward fully automated math-to-kernel optimization. Nautilus compiles a high-level algebraic specification of tensor operators into efficient tiled GPU kernels. Nautilus's successive lowering design allows high-level optimizations, expression rewrites, and tile optimizations to be jointly applied in a single end-to-end system. Nautilus presents a novel auto-scheduler that discovers sequences of high-level optimizations, while preserving the regular program structure needed by tile optimizers. Nautilus's auto-scheduler captures complex interactions and trade-offs in the high-level optimizations, including aggressive global transformations like advanced reduction fusion. Nautilus is the first end-to-end tensor compiler capable of starting from a math-like description of attention and automatically discovering FlashAttention-3-like kernels, offloading the entire burden of optimization from the programmer to the compiler. Across five transformer-based models and 150 evaluation configurations on NVIDIA GH200 and RTX 5090 GPUs, Nautilus achieves up to 23% higher throughput than state-of-the-art compilers on GH200 and up to 42% on RTX 5090, while matching or exceeding manually written cuDNN kernels on many long-sequence configurations.
Yifan Zhao, Yuchen Yang, Matei Budiu +1
Graphics Processing Unit KernelsCode Optimization
AskPDF
May 28, 2026cs.CL

HTAM: Hierarchical Transition-Attended Memory for Operator Optimization

High-performance GPU kernels are essential for efficient LLM deployment, yet optimizing them remains expertise-intensive. Recent LLM-based code generation makes automatic GPU operator generation promising, but operator optimization remains a hardware-aware search problem. Existing LLM-based methods face a granularity mismatch: coarse hints are reusable but hard to execute, whereas detailed memories are actionable but enlarge the search space and obscure optimization bottlenecks. The key challenge is therefore to organize optimization experience at an appropriate granularity. To address this issue, this paper proposes HTAM (Hierarchical Transition-Attended Memory), a coarse-to-fine framework for LLM-based operator optimization. HTAM builds a two-level Hierarchical Transition Graph (HTG) to organize coarse global directions, detailed local strategies, and transition experience between optimization steps. During each evolution step, HTAM selects a global direction from the current state and recent optimization history, retrieves the corresponding local strategy memory, and uses it to guide concrete CUDA code generation. Experiments on the full KernelBench suite demonstrate that HTAM consistently improves correctness, fast-solution rate, and speedup over LLM-based baselines, while backend and Robust-KBench studies indicate transferable benefits from structured memory.
Yining Zhang, Mingyang Yi, Chen Wang +5
Code OptimizationLarge Language Model ServingBottlenecksHierarchical
AskPDF