CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs

• URL: http://arxiv.org/abs/2605.19269v2
• Date: Wed, 20 May 2026 17:38:24 GMT
• Title: CODA: Rewriting Transformer Blocks as GEMM-Epilogue Programs
• Abstract summary: CODA is a GPU kernel abstraction that expresses these computations as GEMM-plus-epilogue programs.<n>We show that CODA offers a practical path toward combining framework-level productivity with hardware-level efficiency.
• Score: 48.961770224590204
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Latent Recurrent Transformer: Architecture Exploration, Training Strategies, and Scaling Behavior [107.2098567818173]
  Latent Recurrent Transformer (LRT) is a lightweight augmentation of autoregressive transformers.<n>LRT reuses a high-level source-layer hidden state from the previous token as recurrent memory for the next token.
  arXiv  Detail & Related papers  (2026-05-26T10:10:26Z)
• EVA: Accelerating LLM Decoding via an Efficient Vector Quantization Architecture [13.194852545534822]
  EVA builds on a simple yet effective insight that combines input-codebook computation with conflict-free memory access.<n> EVA achieves up to 11.17$times$ speedup and 7.17$times$ higher energy efficiency compared with the SOTA lookup-based architecture.
  arXiv  Detail & Related papers  (2026-05-22T19:06:22Z)
• Improving the Performance and Learning Stability of Parallelizable RNNs Designed for Ultra-Low Power Applications [1.226598527858578]
  Bistable Memory Recurrent Unit (BMRU) was introduced to enable hardware-software co-design of ultra-low power RNNs.<n>BMRU performance lags behind parallelizable RNNs on complex sequential tasks.<n>We propose a cumulative update formulation that restores flow while preserving persistent memory, creating skip-connections through time.
  arXiv  Detail & Related papers  (2026-05-12T09:39:33Z)
• LP-GEMM: Integrating Layout Propagation into GEMM Operations [0.20388938295521575]
  This paper introduces LP-GEMM, a decomposition of the GEMM kernel that enables packing- propagation across sequential GEMM operations.<n>We show average speedups of 2.25x over OpenBLA on Intel x86 for sequential GEMMs and competitive gains relative to vendor-optimized libraries.
  arXiv  Detail & Related papers  (2026-04-06T11:24:34Z)
• Echo State Transformer: Attention Over Finite Memories [2.118933003468525]
  We introduce Echo State Transformers (EST), a hybrid architecture that elegantly resolves the challenge of sequential data processing.<n>EST integrates the Transformer attention mechanisms with principles from Reservoir Computing to create a fixed-size window distributed memory system.<n>EST ranks first overall in two of five categories, outperforming strong state-of-the-art baselines on classification and anomaly detection tasks.
  arXiv  Detail & Related papers  (2025-06-25T09:56:25Z)
• UIO-LLMs: Unbiased Incremental Optimization for Long-Context LLMs [111.05657299071648]
  UIO-LLMs is an incremental optimization approach for memory-enhanced transformers under long-context settings.<n>We refine the training process using the Truncated Backpropagation Through Time (TBPTT) algorithm.<n>UIO-LLMs successfully handle long context, such as extending the context window of Llama2-7b-chat from 4K to 100K tokens with minimal 2% additional parameters.
  arXiv  Detail & Related papers  (2024-06-26T08:44:36Z)
• Folded Context Condensation in Path Integral Formalism for Infinite Context Transformers [0.0]
  We present a generalized formulation of the Transformer algorithm by reinterpreting its core mechanisms within the framework of Path Integral formalism.<n>We obtain a more compact and efficient representation, in which the contextual information of a sequence is condensed into memory-like segments.<n>We validate the effectiveness of this approach through the Passkey retrieval task and a summarization task, demonstrating that the proposed method preserves historical information while exhibiting memory usage that scales linearly with sequence length.
  arXiv  Detail & Related papers  (2024-05-07T19:05:26Z)
• Accelerator-driven Data Arrangement to Minimize Transformers Run-time on Multi-core Architectures [5.46396577345121]
  complexity of transformer models in artificial intelligence expands their computational costs, memory usage, and energy consumption. We propose a novel memory arrangement strategy, governed by the hardware accelerator's kernel size, which effectively minimizes off-chip data access. Our approach can achieve up to a 2.8x speed increase when executing inferences employing state-of-the-art transformers.
  arXiv  Detail & Related papers  (2023-12-20T13:01:25Z)
• Blockwise Parallel Transformer for Large Context Models [70.97386897478238]
  Blockwise Parallel Transformer (BPT) is a blockwise computation of self-attention and feedforward network fusion to minimize memory costs. By processing longer input sequences while maintaining memory efficiency, BPT enables training sequences 32 times longer than vanilla Transformers and up to 4 times longer than previous memory-efficient methods.
  arXiv  Detail & Related papers  (2023-05-30T19:25:51Z)
• Linearizing Transformer with Key-Value Memory Bank [54.83663647680612]
  We propose MemSizer, an approach to project the source sequence into lower dimension representation. MemSizer not only achieves the same linear time complexity but also enjoys efficient recurrent-style autoregressive generation. We demonstrate that MemSizer provides an improved tradeoff between efficiency and accuracy over the vanilla transformer.
  arXiv  Detail & Related papers  (2022-03-23T18:10:18Z)
• Gradient Coding with Dynamic Clustering for Straggler Mitigation [57.9123881133818]
  GC-DC regulates the number of straggling workers in each cluster based on the straggler behavior in the previous iteration. We numerically show that GC-DC provides significant improvements in the average completion time (of each iteration) with no increase in the communication load compared to the original GC scheme.
  arXiv  Detail & Related papers  (2020-11-03T18:52:15Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.