Optimizing GEMM + Activation on CDNA4 with TLX

Aug 07, 2026

Motivation

Modern large language models spend the majority of their training time inside transformer feed-forward networks (FFNs). While attention layers often receive the spotlight, FFNs account for a substantial fraction of total computation, making even modest improvements in their efficiency translate into meaningful reductions in end-to-end training time.

Recent transformer architectures increasingly rely on gated activation functions from the Gated Linear Unit (GLU) family because they improve model quality with relatively little additional arithmetic. However, the implementation of these layers exposes a classic GPU performance bottleneck: memory movement. After a matrix multiplication produces the intermediate activation, a separate kernel typically reloads that data from high-bandwidth memory (HBM) to apply the gate. The arithmetic is inexpensive, but the extra kernel launch and redundant memory traffic make the operation memory-bound rather than compute-bound.

This blog explores how we eliminate that overhead by fusing the matrix multiplication and GLU activation into a single Triton Low-level Language Extensions (TLX) kernel. TLX is a set of low-level extensions to Triton, maintained in facebookexperimental/triton, that expose direct control over shared memory, asynchronous loads, and warp scheduling. Instead of presenting only the final implementation, we walk through four successive kernel versions, using profiling data to motivate each optimization. Along the way, we'll show how explicit control over shared memory, asynchronous direct-to-LDS transfers, software pipelining, persistent scheduling, and a fused epilogue progressively transform a correct implementation into one that outperforms both torch.compile and even the standalone rocBLAS GEMM on the target workload.

Before diving into the optimization journey, let's first understand why the GLU epilogue becomes the performance bottleneck.

The GLU Block

A GLU block does two things back-to-back: a linear projection, X = A @ B + bias (an add and matrix multiply or addmm), followed by an elementwise gate that mixes X with a second tensor Y. This gate is the block's activation, and in this kernel, it is out = X + X*Y.

At the shape that our kernel targets (M = 1024, N = 21568), total size of X, which is stored in FP16, can be calculated as follows:

M * N * 16 bits * 1 byte / 8 bits ≈ 44,171,264 bytes

Since the gate Y is the same shape as X, it takes up another 44 MB. An unfused path pays for the following transfers:

The GEMM writes the full 44 MB X out to HBM
A second kernel reads those 44 MB back in
Reads the 44 MB of Y
Writes the 44 MB result

This sequence results in roughly 4 * 44 = 176 MB of HBM traffic and two kernel launches for what is, in the end, a single math expression.

The gate is entirely memory-bound as a standalone operation, and the intermediate X is written and re-read for no reason other than the operator boundary. Fusion collapses these redundant HBM transactions such that X is produced in registers by the matmul and consumed by the gate in the same kernel. As a result, fusion drops one full write and one full read of the intermediate (about 88 MB), removes a kernel launch, and lets the gate ride the GEMM epilogue instead of paying for an entirely separate memory-bound pass. Because an FFN block runs in every layer, on every step, and again in the backward pass, these per-call savings compound across the entire training run. The significant time spent in these FFN blocks is what makes a fused addmm + GLU kernel worth hand-tuning.

Implementing the Fusion in TLX

Building this kernel by hand involves shaping the software pipeline directly by determining where tiles live in LDS, when loads are issued, and how compute overlaps memory. Ordinary Triton hides these decisions behind its loop-pipelining compiler pass, which is the right default until it leaves performance on the table. TLX exposes these low-level building blocks so that the kernel writer can construct the pipeline explicitly. This kernel uses the following TLX primitives and techniques:

Explicit LDS (shared memory) management: tlx.local_alloc reserves one or more buffers in LDS, and tlx.local_view / tlx.local_load / tlx.local_store move tiles between LDS and registers. Allocating multiple buffers is how you can multibuffer by hand and keep several tiles along the K dimension resident so that compute on one overlaps the load of the next. A companion tlx.padded_shared_layout_encoding allows you to pad each row of an LDS tile to sidestep bank conflicts.
Asynchronous direct-to-LDS loads: tlx.async_load issues a global-memory read that lands directly in LDS, bypassing vector registers entirely, and returns a token. tlx.async_load_commit_group batches issued loads into a group, and tlx.async_load_wait_group blocks until a chosen number of groups remain in flight, allowing you to keep several transfers moving while the matrix cores stay busy.
Warp pipelining: tlx.warp_pipeline_stage marks explicit stage boundaries inside the K-loop — for example a compute stage and a memory stage. The compiler partitions the loop body at those boundaries and runs one warp group a stage ahead of the other, overlapping load latency with MFMA compute. An optional priority argument is a direct hardware scheduling hint that marks the priority of each stage.
Loop pipelining: The same TLX primitives that build the main-loop pipeline also let you carry that overlap into the epilogue. Issuing the gate's Y operand as an early load hides its memory latency behind the final drain matmuls, so Y is already in registers by the time the gate runs. With that load paid for under compute and the gate itself reduced to a single fused multiply-add, the gate is essentially free.

Baseline: PyTorch with torch.compile

The reference point for performance is the addmm-plus-gate path handed to torch.compile with max-autotune, so Inductor fuses whatever it can:

		def addmm_glu(a, b, bias, y):
    x = torch.addmm(bias, a, b)   # projection
    return x + x * y              # gate

baseline = torch.compile(addmm_glu, mode="max-autotune")

Under the hood, torch.addmm dispatches to hipBLASLt, AMD's tuned GEMM library, which selects a kernel (Cijk_…_Bias_…_MT192x256x64_… on our setup) that already folds the bias into the GEMM epilogue. The projection is therefore a single, well-tuned fused kernel. What the library cannot fuse is the gate, since Y is a full M×N tensor rather than a vector epilogue.

The torch.compile handles the gate by collapsing the multiply, add, and intermediate into a single pointwise kernel, and with max-autotune it also tries a Triton GEMM template with the epilogue fused in. In most cases, it does not fuse the gate into the projection, so X is still written out to HBM by the matmul and read back by the gate kernel, exactly the round-trip our fused kernel eliminates. This compiled path is the strongest baseline available in PyTorch, and it reaches 200, 279, and 414 TFLOPS for K of 256, 512, and 1024. Every version below is measured against it.

Performance is also compared against rocBLAS, AMD's tuned GEMM library, as a second, library-only reference. Since rocBLAS computes only the projection A @ B and performs neither the bias add nor the GLU gate, it does strictly less work than our fused kernel. On this shape it reaches 271, 410, and 578 TFLOPS for K of 256, 512, and 1024.

The complete reference kernels for all four versions below, along with the benchmark harness, are available here.

Version 1: The Starting Point

Full kernel: v1_register_staged.py

The V1 kernel serves as a starting point that yields correct results and leaves a lot of optimization opportunities on the table. Its load path goes from global memory into registers and then into shared memory (tl.load into a register tile, followed by tlx.local_store into LDS). It prefetches tiles ahead in a multi-buffered software pipeline, but every tile still makes the detour through registers before landing in LDS, which adds register pressure and an extra hop along the critical path. The diagram below traces the path a single K-tile takes through the memory hierarchy (top) and how the software pipeline schedules those steps across loop iterations (bottom).

Measured throughput:

Performance was measured with rocprofv3 using kernel time and advanced thread traces (ATT). Every optimization mentioned in the following sections is justified by this data.

K	torch.compile (TFLOPS)	V1 (TFLOPS)	Speedup
256	200	235	1.2x
512	279	354	1.3x
1024	414	463	1.1x

Profiling V1 reveals that the matrix unit is busy about 10% of the time, and roughly 60% of the stall cycles are memory and LDS-staging waits (s_waitcnt, ds_read/ds_write, and global loads). Each tile's global load, LDS staging, and MFMA run in sequence with no overlap, so the matrix unit sits idle waiting on memory. The kernel is memory-staging-bound, not compute-bound, and two things cause it: the register-staging detour on every load, and the fully serial load, stage, and compute chain.

Version 2: Direct to LDS, Prefetching, Swizzling, Autotuning

Full kernel: v2_direct_to_lds.py

V2 attacks exactly the two weaknesses of V1: the register detour and the lack of compute/memory overlap. The main changes in V2 are:

Direct to LDS loads. tlx.async_load copies A and B tiles straight from HBM into LDS, skipping the register detour entirely. The copies are issued asynchronously and tracked with commit groups and wait groups.

Global and local prefetching. The hot loop is split into two warp pipeline stages: an MFMA stage that runs the matrix multiply on the tiles already in registers, and a memory stage that issues the next global prefetch and the local read for the following iteration, so compute and memory overlap.

L2 locality. Three mechanisms keep reused tiles in cache and single-use data out of it:
1. XCD swizzle. Program IDs are remapped so workgroups scheduled together land on the same XCD (gfx950 has 8, each with its own L2 slice), so the reuse resolves within one slice instead of scattering across XCDs.

2. Grouped tile ordering. The kernel sweeps a band of GROUP_SIZE_M row-tiles across all columns before advancing, so the A rows and B columns in that band are reused.

3. Streaming cache hints (.cs). Y and the output are each touched once, so the hint tells the hardware not to retain them, which keeps epilogue traffic from evicting the A and B tiles being reused.

Autotuning. The sweep settled on a 128x128 tile and GROUP_SIZE_M = 8.

Measured throughput:

K	torch.compile (TFLOPS)	V2 (TFLOPS)	Speedup
256	200	303	1.5x
512	279	424	1.5x
1024	414	596	1.4x

Version 3: A Deeper Pipeline and Persistent Scheduling

Full kernel: v3_deep_pipeline_persistent.py

Version 2 was substantially faster than the baseline, but profiling still showed the matrix units spending time idle inside the K reduction. To understand why, the amdgcn assembly of the hot loop needed to be examined. The gate that synchronizes the async prefetch with the MFMA stage told the whole story:

		; Version 2 hot loop, K reduction (software-pipelined; instructions interleaved)
      ...
      v_mfma_f32_16x16x32_f16 v[12:15], v[106:109], v[90:93], v[12:15]
      v_mfma_f32_16x16x32_f16 v[8:11],  v[102:105], v[90:93], v[8:11]
      v_mfma_f32_16x16x32_f16 v[4:7],   v[98:101],  v[90:93], v[4:7]
      v_mfma_f32_16x16x32_f16 v[0:3],   v[94:97],   v[90:93], v[0:3]
      s_setprio 1
  ; wait_asyncmark(0)
      s_waitcnt vmcnt(0) lgkmcnt(0)      ; drain every in-flight global load
      s_barrier
      ...                                ; issue prefetch

The vmcnt(0) is a full drain. vmcnt counts how many global reads are still in flight, and waiting for the count to reach zero stalls the wave until all of them finish. With only two shared memory buffers, the wait in the hot loop is effectively forced to wait on the global read issued immediately before it because the local read that comes right after depends on that global read. The pipeline is too shallow to permit a weaker vmcnt that would let the most recent global read stay in flight.

The fix is to give the pipeline more room. Raising NUM_BUFFERS from 2 to 3 adds a third shared memory slot, so one tile can be computed while a second is already resident and a third is still arriving. As a result, the tile that a local read consumes was fetched by a global read two iterations earlier. This lets the wait relax from "drain everything" to "keep one tile in flight," as illustrated below:

The assembly confirms that the change landed. The same gate in Version 3 no longer drains to zero:

		; Version 3 hot loop, K reduction (software-pipelined; instructions interleaved)
      ...
      ds_read_b128       v[122:125], v79              ; local_load
      ds_read_b64_tr_b16 v[110:111], v79 offset:25312
      ds_read_b64_tr_b16 v[106:107], v79 offset:25376
      s_setprio 0
  ; wait_asyncmark(1)
      s_waitcnt vmcnt(4) lgkmcnt(0)      ; keep one tile in flight
      s_barrier
      ...                                ; back-edge to next K iteration

The wait count changed from wait_asyncmark(0) to wait_asyncmark(1). The loop now tolerates one async group in flight when it hits the barrier. Since a group is one A tile plus one B tile, and each resolves to buffer_load_dwordx4 operations, one group in flight is vmcnt(4) at the machine level. The MFMA stage no longer waits on the prefetch it does not need, and the next tile's loads overlap the current compute instead of blocking it.

The final change in V3 was to implement persistent scheduling. A fixed number of workgroups is launched, sized to the GPU, and each workgroup loops over its share of the output tiles. This persistent scheduling cuts launch overhead and, because each workgroup walks its tiles in the grouped order from Version 2, keeps the L2 reuse working across tiles instead of resetting on every launch.

Measured throughput:

K	torch.compile (TFLOPS)	V3 (TFLOPS)	Speedup
256	200	305	1.5x
512	279	449	1.6x
1024	414	613	1.5x

Profiling this version revealed that the GEMM body was no longer the bottleneck. Around 15% of stall cycles were now spent in the epilogue waiting on the Y load, so the remaining opportunity for optimization was in the epilogue itself.

Version 4: Making the Epilogue Nearly Free

Full kernel: v4_fused_epilogue.py

The final V4 kernel targeted the three costs left in the epilogue.

Fold the bias into the accumulator: Instead of computing A @ B and adding the bias afterward, the accumulator is initialized to the bias (acc = bias + zeros). The K loop then yields X = bias + A@B directly, so the bias add disappears as a separate step.
Prefetch Y into registers before the drain dots: Staging Y through LDS was the obvious idea, but it proved worse: it round-trips Y from HBM to LDS and back for a value used exactly once, and the copy is not hidden, since the drain loop runs only NUM_BUFFERS - 1 iterations. V4 instead issues a streaming Y load straight into registers just before the final drain matmuls, so the load latency overlaps those tail dots and Y is ready when the gate needs it.
Collapse the gate to a single packed fma. With the bias already folded into X and Y sitting in registers, the gate out = X + X*Y is exactly fma(X, Y, X), so it becomes one packed tl.fma(acc, y_regs, acc).

The old epilogue did three things: a bias add, a Y load that stalled the tail, and a two-step gate. V4 folds all three into a single fused multiply add, running on a Y value whose load is already hidden under the tail dots. The V4 kernel is the fastest of the versions and the one the project settled on. Notably, V4 also outpaces rocBLAS at every K, even though rocBLAS computes only the matmul while V4 fuses the bias and the X + X*Y gate on top.

Measured throughput:

K	torch.compile (TFLOPS)	V4 (TFLOPS)	Speedup
256	200	357	1.8x
512	279	502	1.8x
1024	414	619	1.5x

Conclusion

Kernel optimization is rarely about discovering a single breakthrough. Instead, it is the cumulative effect of many targeted improvements, each removing a different bottleneck identified through profiling. In this work, we started with a correct fused kernel and systematically improved it by eliminating unnecessary register staging, introducing direct-to-LDS asynchronous loads, deepening the software pipeline, improving cache locality, adopting persistent scheduling, and finally optimizing the epilogue so the GLU activation became almost free.

The resulting kernel reaches 357, 502, and 619 TFLOPS for K = 256, 512, and 1024, delivering up to 1.8× higher throughput than the torch.compile baseline while also outperforming the standalone rocBLAS GEMM, even though the fused kernel performs additional work by incorporating the bias addition and GLU activation. The results are shown in the table below:

K	torch.compile (TFLOPS)	rocBLAS (TFLOPS)	V1 (TFLOPS)	V2 (TFLOPS)	V3 (TFLOPS)	V4 (TFLOPS)
256	200	271	235	303	305	357
512	279	410	354	424	449	502
1024	414	578	463	596	613	619

Benchmark script: bench.py

More importantly, this optimization illustrates the value of TLX's low-level programming model. By exposing explicit control over asynchronous memory transfers, shared-memory allocation, software pipelining, and synchronization, TLX allows kernel authors to overlap computation and memory movement in ways that are difficult for compiler-generated code to achieve automatically.

The broader lesson extends beyond this specific GLU fusion. As modern AI workloads become increasingly bandwidth-constrained, eliminating unnecessary memory traffic is often more valuable than reducing arithmetic. Keeping intermediate activations on-chip and fusing operations into a single pipeline enables GPUs to spend more time performing useful computation instead of waiting on memory—a principle that applies to many transformer workloads beyond this example.

Version 1: one iteration of the hot loop, captured with ATT. As illustrated by the long green MFMA bars and long orange LDS bars, there are several stalls in this pipeline.

Version 4: the same hot loop iteration on the same time axis. After pipelining and scheduling changes, the iteration is far shorter and the MFMAs run nearly back-to-back, with the LDS and memory stalls largely gone.

Future Work

Several directions could push this kernel further:

Non-power-of-two (NPOT) tile sizes. The target shape (N = 21568) does not divide evenly into the power-of-two tiles these kernels use, so boundary tiles do wasted work. NPOT tiling could fit the blocks to the problem and recover that overhead.
FP8 precision. CDNA4 offers substantially higher matrix throughput in fp8 than in fp16, so moving the projection and potentially the gate to fp8 would also halve the tensor traffic that dominates this shape
Fusing the backward pass. An FFN block runs in both the forward and backward pass of every training step. Applying the same epilogue-fusion approach to the backward computation would extend the per-call savings across the full training loop, not just the forward pass.

Disclaimers

Configuration Details:

All performance data mentioned in this blog was measured on July 31, 2026. Throughput is reported as TFLOP/s computed over the GEMM floating-point operations (2 · M · N · K) of the fused addmm + GLU problem at M = 1024 , N = 21568, fp16, for K ∈ {256, 512, 1024}. Speedups are relative to the torch.compile (max-autotune) PyTorch baseline. All measurements were taken on a single AMD Instinct MI350X GPU (gfx950 / CDNA4), with ROCm 7.2.2, PyTorch 2.10.0+rocm7.0, and a source build of the facebookexperimental/triton fork (branch main, commit c32d385) installed in the system environment.

System Configuration:

AMD Instinct™ MI350X GPU node

System Model: Quanta Grand Teton 1.5
CPU: 2× AMD EPYC 9655 96-Core Processor (384 threads total)
NUMA: 2 NUMA nodes (1 per socket); NUMA auto-balancing enabled
Memory: 3072 GiB DDR5-6400 (Micron Technology)
Disk: 3.4 TB
GPU: 8× AMD Instinct MI350X, 288 GB HBM3E each, 256 CUs
Host OS: Ubuntu 24.04 LTS
System BIOS: F0TC3A09
System BIOS Vendor: American Megatrends International, LLC.
Host GPU Driver: amdgpu 6.16.13 / ROCm 7.2.2