Deep Dive Into 4-Wave Interleave FP8 GEMM

6 min read. | 1485 total words.

Our previous two posts in this GEMM optimization series covered Matrix Core instructions and 8-wave ping-pong FP8 GEMM design. Here we discuss another algorithm design introduced by HipKittens - 4-wave interleave, which further improves the performance of the 8-wave ping-pong implementation. For the most complete understanding, we recommend reading this post alongside the source code.

To understand why 4-Wave Interleave exists, recall how the 8-Wave kernel worked: it placed two waves on each SIMD unit — one handling MFMA, one handling memory loads. These two waves alternate between memory and MMA instructions (“ping-pong”), and the hardware scheduler overlaps them because they use different hardware units.

4-Wave Interleave takes the opposite approach. It places one wave per SIMD, and that single wave is responsible for issuing both MFMA and memory instructions but in a very carefully hand-crafted order.

With only one wave on the SIMD, that wave gets all 512 VGPRs compared to the 8-Wave kernel where the register budget was half that. Doubling the register file means the wave can hold a complete 128×128 output register tile at once, compared to 64×128 in the 8-wave case.

The 4-Wave and 8-Wave designs differ in their implementation complexity. The 8-Wave kernel requires creating alternating wave behavior via conditional __builtin_amdgcn_s_barrier(), whereas the 4-Wave employs a finer-grained software pipeline that overlaps memory and MFMA instructions.

In this section, we discuss the code and the implementation in detail.

For the 4-Wave kernel we will keep the same tile sizes as we used in the 8-Wave example from our previous article that is: 256x256x128. Also, we will only consider input matrices with dimensions that are multiples of the blocking parameters.

Each wave ends up handling 4 64x64 tiles of the output matrix C, and each one of these is made up of a sequence of smaller, 16x16, sub-tiles. Using these smaller sub-tiles is the key that enables the interleaving: instead of issuing bulk MFMA instructions to compute, e.g., an entire 64x64 tile, we will issue smaller 16x16 MFMA instructions and mix them with memory instructions, creating custom software pipeline schemes. The idea is that by the time we finish processing an entire 64x64 tile of C, we will have already loaded the operands for the next one.

The core of the kernel is interleaved_block. Each call does three things:

Issues 16 MFMA instructions (a 4x4 grid of 16x16x128 MFMA) computing a 64x64 part of C.

Issues 16 MFMA instructions (a 4x4 grid of 16x16x128 MFMA) computing a 64x64 part of C.

Issues 8 LDS → register loads (loading a 64x128 fragment of A/B) to prepare the operands for the next call.

Issues 8 LDS → register loads (loading a 64x128 fragment of A/B) to prepare the operands for the next call.

Issues 4 global → LDS loads loading a 128x128 sub-tile of A/B that will be used in the next call to bring the operands from LDS to registers.

Issues 4 global → LDS loads loading a 128x128 sub-tile of A/B that will be used in the next call to bring the operands from LDS to registers.

The key mechanism is __builtin_amdgcn_sched_barrier(x), which prevents the compiler from reordering instructions across the barrier. The parameter x is a mask that tells the compiler which types of instructions can cross the intrinsic in the compiled schedule. x=0 means no instruction can cross the barrier. Without it, the compiler might legally batch all memory instructions after all MFMA, eliminating the overlap. With it, memory instructions are placed between MFMA:

The prologue loads both cur and next tile pairs into LDS before the loop starts, seeding the double buffer. The last two iterations are manually unrolled (the epilogue) because there is no further tile to prefetch — the cluster calls are replaced by plain mfma calls interleaved with load_rt for the remaining register loads.

In the epilogue, compared to the original HipKittens 4-wave kernel, we modified how the accumulator is stored: instead of waiting for the entire c fragment to be computed and then storing it, we start storing parts c[i][j] as soon as they are available. From our observations, this improves performance across all test cases.

The AMD Instinct™ MI355X GPU (CDNA4) has 64 banks, each 4 bytes wide. Given an address, a bank number can be computed as bank = (addr / 4) % 64. This means that the period is 256 bytes (i.e. addresses that are multiple of 256 all map to bank 0).

To load A/B from LDS into registers, we use the ds_read_b128 instruction. This instruction allows us to load 128 consecutive bits (16 FP8 values) per thread. Since we need a 16x128 fragment of A/B for each MFMA we’ll have to issue two ds_read_b128 per thread in order to load one fragment: (16*128) / (16*64) = 2.

The ds_read_b128 instruction is carried out in 4 phases, during which only a subset of threads is active. For the instruction to be bank conflict-free, all these phases must be conflict-free.

Unfortunately, the way we currently access A/B isn’t conflict-free as shown in figure 4.

Each lane needs to access elements by column, but since the stride is 128 bytes, lanes along the same column will always hit the same two banks. This results in an 8-way bank conflict.

To eliminate bank conflicts, we need to change the way in which threads access data in LDS. We can achieve this by applying this swizzle pattern:

Figure 5 shows how swizzling changes the access pattern. As you can see, we are now bank conflict-free during all four phases of the ds_read_b128 instruction.

Together with LDS swizzling, we’ll also apply a “chiplet-aware” swizzle in order to better optimize cache hit rates. The algorithm is detailed in Appendix A: Optimizing Cache Reuse. The idea is that we can change the way we compute the indices block_m and block_n shown in figure 2 by using some insights about the hardware to better utilize the cache hierarchy of our GPU.

All tests were conducted on AMD Instinct™ MI355X GPU using ROCm 7.2.2. The workload is FP8 GEMM with BF16 output and FP32 accumulation.

Figure 6 shows a performance comparison of the two algorithms. For benchmarking, we used 1000 warm-up iterations, 1000 benchmark iterations, and rotating buffers along with normally distributed data. We report the average FLOP/s over 1000 benchmark iterations.

One wave per SIMD gives each wave the full 512-VGPR budget, enabling a larger 128×128 output register tile instead of 64×128.

One wave per SIMD gives each wave the full 512-VGPR budget, enabling a larger 128×128 output register tile instead of 64×128.

Fine-grained software pipelining via __builtin_amdgcn_sched_barrier interleaves MFMA instructions with global→LDS and LDS→register loads at the sub-tile level, hiding memory latency without relying on the hardware scheduler to overlap two waves.

Fine-grained software pipelining via __builtin_amdgcn_sched_barrier interleaves MFMA instructions with global→LDS and LDS→register loads at the sub-tile level, hiding memory latency without relying on the hardware scheduler to overlap two waves.

Early accumulator store — storing completed c[i][j] sub-tiles as soon as they are available rather than waiting for the full fragment — yields measurable performance gains across all test cases.

Early accumulator store — storing completed c[i][j] sub-tiles as soon as they are available rather than waiting for the full fragment — yields measurable performance gains across all test cases.

LDS swizzling eliminates 8-way bank conflicts that arise from the column-major access pattern required by the 16×16×128 MFMA instruction.

LDS swizzling eliminates 8-way bank conflicts that arise from the column-major access pattern required by the 16×16×128 MFMA instruction.

Chiplet-aware grid swizzling (Appendix A) remaps thread block IDs so each XCD works on a compact rectangular region of the output matrix, maximizing both L2 and LLC hit rates on the MI355X’s 8-XCD topology.

Chiplet-aware grid swizzling (Appendix A) remaps thread block IDs so each XCD works on a compact rectangular region of the output matrix, maximizing both L2 and LLC hit rates on the MI355X’s 8-XCD topology.

The 4-wave design also proved more robust across compiler versions: unlike the 8-wave kernel, it required no manual #pragma unroll tuning to avoid register spilling and delivered consistent performance across different ROCm releases.

In an upcoming post, we will show how all of the above — the software pipeline, LDS swizzling, and grid scheduling — can be expressed concisely in FlyDSL, a Python DSL that generates high-performance HIP kernels while hiding much of the low-level scheduling complexity.

Data measured by AMD by the authors on May 11th, 2026. All the experiments were conducted on the AMD Instinct™ MI355X platform. System configuration:

GPU: MI355X-DLC 1.4KW

GPU: MI355X-DLC 1.4KW

ROCm Version: 7.2.2

ROCm Version: 7.2.2

OS: Linux (Ubuntu 22.04)

OS: Linux (Ubuntu 22.04)

In addition to the two schedules we discussed in this article, the HipKittens paper also introduced an algorithm to improve cache reuse on modern GPUs.

The key insight is that modern GPUs are moving towards a chiplet architecture rather than a monolithic one. For example, the AMD Instinct™ MI355X is a cluster of 8 chiplets called accelerated complex die, or XCD, each one with 32 compute units. Each XCD has an L2 cache and shares a Last Level Cache with the others. Ideally, we’d like to optimize for both cache levels:

Make thread blocks assigned to the same XCD reuse the same rows of A and the same columns of B as much as possible to optimize for the L2 cache.

Make thread blocks assigned to the same XCD reuse the same rows of A and the same columns of B as much as possible to optimize for the L2 cache.

Make multiple XCDs work on common regions of the input matrices so that shared data stays in the LLC as much as possible.

Make multiple XCDs work on common regions of the input matrices so that shared data stays in the LLC as much as possible.

Consider the case where the number of XCDs is 3. With the naive grid schedule, tiles are assigned to XCDs in round-robin order as shown in the figure below:

At the end, each XCD will own a WINDOW_SIZE x CHUNK_SIZE / WINDOW_SIZE contiguous tile block.

By assigning a rectangular region of the output matrix to each XCD and assigning contiguous regions to contiguous XCDs, we can maximize both the L2 and LLC hit rates.

It’s worth pointing out that this algorithm requires hand-tuning in order to find the best schedule for a given GPU. The parameters WINDOW_SIZE and CHUNK_SIZE must be tuned based on the hardware characteristics of your target GPU as well as the specific problem size you wish to optimize for.

Third-party content is licensed to you directly by the third party that owns the content and is not licensed to you by AMD. ALL LINKED THIRD-PARTY CONTENT IS PROVIDED “AS IS” WITHOUT A WARRANTY OF ANY KIND. USE OF SUCH THIRD-PARTY CONTENT IS DONE AT YOUR SOLE DISCRETION AND UNDER NO CIRCUMSTANCES WILL AMD BE LIABLE TO YOU FOR ANY THIRD-PARTY CONTENT. YOU ASSUME ALL RISK AND ARE SOLELY RESPONSIBLE FOR ANY DAMAGES THAT MAY ARISE FROM YOUR USE OF THIRD-PARTY CONTENT.