🤖 machine learning

Tile-Level Activation Overlap for Efficient LLM Inference

This paper introduces two specialized CUTLASS-based SM90 kernels that fuse SwiGLU activation with matrix multiplication at the tile level to eliminate intermediate tensor materialization overhead, achieving up to 2.47x speedup and 79.5% peak utilization on NVIDIA H100 GPUs while outperforming both PyTorch and cuBLAS in efficiency and numerical accuracy.

Original authors: Abhinav Jangda, Tyler Sorensen, Sebastian Burckhardt, Jianlan YE, Chaoyin Li, Atul Gupta

Published 2026-07-07

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Abhinav Jangda, Tyler Sorensen, Sebastian Burckhardt, Jianlan YE, Chaoyin Li, Atul Gupta

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are running a high-speed factory that builds complex robots (these robots are Large Language Models like the ones powering chatbots). To build each robot, your factory has a specific assembly line called the MLP (a core part of the robot's brain).

For a long time, this assembly line has had a major inefficiency. Here is the problem and the solution presented in this paper, explained simply.

The Problem: The "Middleman" Bottleneck

In modern robot factories, a specific step called SwiGLU is used to make the robots smarter. Think of SwiGLU as a quality check that requires two separate calculations:

Calculation A: Measure the robot's current state.
Calculation B: Measure the robot's potential.
The Glue: Combine these two measurements to get the final result.

The Old Way (The Bottleneck):
In the standard factory setup (like the one used by PyTorch), after doing Calculation A, the workers have to write the results down on a giant whiteboard in the hallway (High-Bandwidth Memory). Then, they have to walk over, read that whiteboard, do Calculation B, write that result on a second whiteboard, and finally walk back to read both whiteboards to do the "Glue" step.

The paper found that for smaller robots (smaller AI models), this "walking to the whiteboard" takes up 30% to 37% of the total time. It's like a chef spending half their time walking to the pantry to grab a single spice, rather than cooking.

The Solution: Two New "Super-Kitchens"

The authors built two new, custom-designed kitchens (called Kernels) that eliminate the whiteboards entirely. Instead of writing results down and walking back, the workers keep the ingredients in their hands (in the Registers) and do everything in one continuous motion.

They created two different strategies for different factory sizes:

1. The "Ping-Pong" Kitchen (Kernel-1)

How it works: Imagine a relay race. While one worker is fetching the next batch of ingredients (loading data), another worker is already mixing the current batch.
The Trick: They use a "Ping-Pong" schedule. While the machine is busy fetching the second set of ingredients, the workers use that exact same time to do the "Glue" math on the first set.
Best For: Factories making huge robots (Large Models) or running many robots at once (Large Batches). It's like a massive assembly line where you have enough workers to keep the big machines fully busy.

2. The "Interleaved" Kitchen (Kernel-2)

How it works: Imagine you are packing boxes. Instead of packing all the red items, then all the blue items, you pack a red item, then a blue item, then a red one, alternating them perfectly.
The Trick: They mix the two weight matrices (the "recipes" for the calculations) together before starting. This allows the workers to grab a "red" and a "blue" ingredient simultaneously and process them together immediately.
Best For: Factories making smaller robots or running fewer robots at a time (Small Batches). This method is so efficient that it keeps the factory floor fully packed with workers, preventing anyone from standing around waiting.

The Results: Speed and Accuracy

The authors tested these new kitchens on NVIDIA H100 chips (the most powerful AI processors available) using various robot sizes (from tiny 0.5B models to massive 72B models).

Massive Speedups: For the smaller robots, the new kitchens were 2.47 times faster than the old standard. That's like cutting a 10-minute task down to 4 minutes.
Shifting the Bottleneck: The old system was "Memory-Bound" (spending too much time walking to the whiteboard). The new system is "Compute-Bound" (spending all its time actually doing the math). They reached 79.5% of the chip's maximum theoretical speed.
The Compiler Can't Do It: The authors tried using the standard "auto-pilot" software (PyTorch's torch.compile) to fix this automatically. It failed. The auto-pilot was 3 to 7 times slower than their custom kitchens. This proves that for this specific problem, you need a human expert to hand-craft the solution; the computer can't figure it out on its own yet.
Better Accuracy: Surprisingly, the new custom kitchens were also more accurate than the standard method. The standard method had small math errors in 4.5% to 11% of the results, while the new kitchens had zero errors.

Summary

The paper shows that by reorganizing how the factory floor works—specifically by stopping the workers from constantly walking to the whiteboard to read/write intermediate notes—they can make AI models run significantly faster, especially for smaller models used on edge devices (like phones or laptops). They proved that standard software tools cannot replicate this efficiency, requiring specialized, hand-written code to achieve these results.

Technical Summary: Tile-Level Activation Overlap for Efficient LLM Inference

Problem Statement

Modern Large Language Models (LLMs), including Qwen-2.5, LLaMA, Mistral, and Gemma, predominantly utilize the SwiGLU activation function. This gated MLP structure computes $Y = \text{SiLU}(A \times W_1) \odot (A \times W_2)$ , requiring two independent matrix multiplications (GeMMs) followed by an element-wise gated activation.

In standard implementations, this process materializes two full intermediate tensors ($Gate$ and $Up$) in High-Bandwidth Memory (HBM) between the GeMM and activation stages. Profiling on NVIDIA H100 GPUs reveals that this intermediate materialization and the associated activation computation consume 9–37% of total MLP execution time. The overhead is most severe for smaller models (e.g., Qwen-2.5 0.5B), where SwiGLU accounts for over 30% of MLP time. As tensor core arithmetic density increases via quantization (FP8, INT4), the relative cost of these memory-bound operations grows, creating a significant bottleneck.

Existing compiler infrastructure, specifically PyTorch's torch.compile, fails to address this. It cannot fuse operations across two separate GeMMs with different weight matrices due to fundamental limitations in graph-level fusion passes. Experiments show torch.compile achieves only 34–94% of eager PyTorch performance for this pattern, with explicit fusion hints providing negligible improvement.

Methodology

The authors present two complementary, hand-crafted CUTLASS-based SM90 kernels that fuse SwiGLU into GeMM at the tile level, eliminating intermediate HBM materialization. Both kernels leverage the Tensor Memory Accelerator (TMA) and warp-specialized scheduling available on the H100 (SM90).

1. Kernel-1: Sync SwiGLU via Pingpong Schedule Overlap

Mechanism: This kernel utilizes the CUTLASS Pingpong warp-specialized schedule, which partitions warps into producers (loading data via TMA) and consumers (executing MMA).
Strategy: It creates an $[M, N]$ threadblock grid. Within the consumer phase, after the Gate tile is computed but before the Up tile data arrives, there is a brief window where consumer warps are idle. The kernel exploits this window to compute the Swish activation ( $\text{SiLU}(Gate)$ ) on the register-resident Gate accumulator.
Overlap: The Swish computation overlaps temporally with the producer's TMA loads for the Up weight tiles. The final element-wise multiplication ( $\text{SiLU}(Gate) \odot Up$ ) and store occur in the epilogue.
Optimization: This design is optimized for large batch sizes ( $M \ge 2048$ ) where the $[M, N]$ grid sufficiently saturates the GPU's Streaming Multiprocessors (SMs).

2. Kernel-2: Interleave SwiGLU via Custom Epilogue Visitor Tree

Mechanism: This kernel introduces a custom Epilogue Visitor Tree (EVT) node called PairMulStore.
Strategy: It concatenates weight matrices $W_1$ and $W_2$ into a single fused matrix $W_{fused} \in \mathbb{R}^{K \times 2N}$ with interleaved columns (even columns for Gate, odd for Up). The kernel launches a standard GeMM over $[M, 2N, K]$ , creating an $[M, 2N]$ threadblock grid.
Fusion: The PairMulStore node operates on the $[M, 2N]$ accumulator within each threadblock. It reads adjacent column pairs $(D[m, 2n], D[m, 2n+1])$ , applies $\text{SiLU}(D[m, 2n]) \odot D[m, 2n+1]$ , and stores the result.
Overlap: This design interleaves SwiGLU computation with tile stores. The arithmetic for the current tile overlaps with the asynchronous TMA stores of the previous tile, providing compute-store overlap.
Optimization: The $2\times$ increase in threadblock count significantly improves SM occupancy for small batch sizes ( $M \le 2048$ ), making it the preferred default for most configurations.

Key Contributions

Tile-Level Fusion: The first implementation to fuse SwiGLU into GeMM at the tile level, eliminating 4 intermediate HBM operations (read/write of Gate and Up) per inference step.
Dual-Kernel Strategy: A complementary approach using Kernel-1 for large batches (sequential Gate→Swish→Up) and Kernel-2 for small batches (interleaved weights with custom EVT), addressing occupancy and scheduling constraints differently.
Numerical Superiority: The fused kernels achieve zero mismatches (0% elements with >1% relative error) compared to an FP32 reference, whereas the baseline cuBLAS implementation exhibits 4.5–11% mismatches due to aggressive BF16 partial accumulation across tiles.
Compiler Limitation Validation: Empirical evidence demonstrating that torch.compile cannot replicate this fusion, remaining 3–7× slower than the hand-crafted kernels.

Results

Evaluated on Qwen-2.5 models (0.5B to 72B) on NVIDIA H100:

Speedup: The kernels achieve up to 2.47× speedup over PyTorch eager (cuBLAS + separate SwiGLU). Kernel-2 wins 13/20 benchmarked configurations, while Kernel-1 wins 7/20 (primarily for large models at high batch sizes).
Regime Shift: Roofline analysis confirms that fusion shifts workloads from memory-bound to compute-bound. For small models, the arithmetic intensity increases by up to 247%, moving the operating point across the roofline ridge.
Utilization: Peak BF16 utilization reaches 79.5% on H100, nearly doubling the utilization of the baseline for small models.
Efficiency: Fusion efficiency (actual speedup vs. theoretical memory savings) is near 100% for small models. For Kernel-2, efficiency exceeds 100% in several configurations due to compute-store overlap benefits.
Scalability: Speedup diminishes with model size (2.47× for 0.5B vs. ~1.0× for 72B at high batch sizes) because the GeMM computation dominates total time, and the relative fraction of intermediate memory traffic decreases.

Significance

The paper argues that for modern LLM inference, particularly on edge-deployment and mobile-sized models (0.5B–14B), the bottleneck has shifted from raw compute to memory traffic caused by intermediate tensor materialization. The authors demonstrate that hand-crafted, hardware-aware tile-level design is necessary to overcome limitations in current compiler infrastructure. By eliminating intermediate HBM accesses and overlapping activation computation with data movement, these kernels significantly reduce inference latency and improve compute utilization, offering a critical optimization for production serving systems where latency and throughput are constrained.