Research map / Architectures and efficiency / Attention and transformers

CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Sparse efficient transformers for long context · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes CODA, a GPU kernel abstraction that fixes a GEMM mainloop and exposes composable epilogue primitives so that non-attention Transformer forward and backward operators (RMSNorm, RoPE, SwiGLU, cross-entropy) can be fused into GEMM epilogues.

Problem

Memory-bound operators surrounding GEMMs in Transformer training (normalization, activations, residual updates, reductions) move large intermediate tensors through global memory, and framework operator boundaries prevent fusion across forward and backward passes.

Approach

Algebraically reparameterize Transformer computations (e.g., delaying the RMSNorm row scale past the next GEMM, accumulating RMSNorm backward statistics from the following GEMM's output, pairing adjacent features for RoPE/SwiGLU) so they become GEMM-plus-epilogue programs built from five restricted primitive classes, with only lightweight tile-partial reductions remaining outside the GEMM.

Key results

• CODA kernels, both LLM- and human-authored, achieve high performance relative to cuBLAS with torch.compile across reparameterized Transformer forward and backward kernels.
• Kernel benchmarks use M=N=K in {4096, 8192}, cross-entropy vocabulary 32768, and report means and standard deviations over 30 runs with Triton do_bench.
• The paper reports numerical error of CODA and QuACK normalized by the standard PyTorch path for BF16 GEMM-RMSNorm-GEMM on Llama-3 8B layers.

Surprising findings

• LLM-authored (Claude Code) CODA kernels reach performance comparable to human-written ones when given the primitive abstraction and curated examples.

Limitations

• Targets a common Transformer architecture; broader model families are future work.
• Currently single-GPU kernels; no distributed execution support.
• Reparameterization can obscure module boundaries and algorithmic semantics, complicating framework integration.

Experimental setup

Benchmarks: Transformer++ forward/backward kernel suite (GEMM-RMSNorm-GEMM, pairwise activations, cross-entropy) LLaMA-style block and layer benchmarks at hidden sizes 2048/4096/8192

Metrics: kernel speedup relative to cuBLAS with torch.compile relative numerical error vs FP32 reference block-level speedup

Baselines: cuBLAS with torch.compile Liger Kernel FlashInfer QuACK PyTorch/cuBLAS raw GEMM

Builds on: Epilogue Visitor Trees CUTLASS QuACK Cut Cross-Entropy CuTeDSL

Closest papers by meaning

• Linear Transformers Implicitly Discover Unified Numerical Algorithms (NeurIPS 2025 · similarity 0.88)
• Dense Structural Compression of Transformers via Gauge-Correct Channel Removal (NeurIPS 2026 · similarity 0.87)
• Learning Rate Transfer for Hybrid Transformer-SSM Architectures (NeurIPS 2026 · similarity 0.86)
• Residual Matrix Transformers: Scaling the Size of the Residual Stream (ICML 2025 · similarity 0.86)
• Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models (NeurIPS 2024 · similarity 0.86)
• Ladder-Residual: Parallelism-Aware Architecture for Accelerating Large Model Inference with Communication Overlapping (ICML 2025 · similarity 0.85)