Second Deep Dive post. The WGMMA post covered how a warpgroup's output tile lands in registers — split across 128 threads in a structured, Layout-describable pattern. What it didn't cover: what happens when that output needs to become the input to another WGMMA, which is exactly what fused kernels like FlashAttention do, and exactly where the FP8 path in FlashAttention-3 gets more expensive than it looks.

The fused-GEMM setup

FlashAttention's inner loop is two matrix multiplies back to back: S = Q·K^T, then (after softmax) O = P·V. The output of the first WGMMA — the accumulator holding P, in registers — needs to become operand A of the second WGMMA. Not copied out to shared memory and reloaded — used directly, register to register, which is the whole point of fusing the two GEMMs in the first place.

Why FP16 gets this for free

For FP16, this just works: the layout the FP32 accumulator naturally sits in after the first WGMMA is already the layout WGMMA expects for an input operand. No data movement, no extra instructions — the compiler can hand the same registers straight to the second WGMMA call.

Why FP8 doesn't

FP8 WGMMA expects its operand data arranged differently: interleaved across pairs of threads, rather than each thread simply owning its own contiguous chunk. Colfax Research's writeup on adding FP8 to FlashAttention-2 gives the concrete pattern — after the reshuffle, thread 0 needs to hold, in order: T0d0, T0d1, T1d0, T1d1, T0d2, T0d3, T1d2, T1d3. Two of those eight values (T1d0, T1d1) originally belonged to thread 1, not thread 0.

Diagram comparing the natural FP32 accumulator layout, where each thread holds its own four contiguous values, against the FP8 WGMMA operand layout, which interleaves values from two threads together

That's the actual problem: it's not that FP8 needs more data movement in some vague sense, it's that the accumulator's layout and the next operand's required layout are two different, specific arrangements — and turning one into the other requires values to physically cross between threads, which registers can't do on their own. A thread can't reach into another thread's registers directly; the only way to move a value from thread 1's registers into thread 0's is through an explicit instruction.

The fix: two kinds of movement, because there are two kinds of mismatch

Some of the reshuffling is within a single thread's own registers — reordering bytes it already owns. That doesn't need cross-thread communication at all, just __byte_perm:

auto upper0 = __byte_perm(upper, lower, 0x7654);
auto lower0 = __byte_perm(upper, lower, 0x3210);

The rest genuinely needs to move between threads, which is what __shfl_sync is for — each thread pulls a value from a specific other thread's registers, using a fixed permutation map within groups of 4 threads:

int upper_map[4] = {0, 3, 1, 2};
int lower_map[4] = {1, 2, 0, 3};
 
upper0 = __shfl_sync(uint32_t(-1), upper0, upper_map[threadIdx.x % 4], 4);
lower0 = __shfl_sync(uint32_t(-1), lower0, lower_map[threadIdx.x % 4], 4);

Both are warp-level intrinsics — no shared memory round-trip, no synchronization barrier across the whole block, just direct register-to-register exchange within a small group of threads. That's about as cheap as fixing a layout mismatch can be, but it's still real instructions on the critical path that FP16 simply never has to issue.

Why this is worth knowing

It's easy to read "FP8 is faster" as a flat statement — lower precision, higher Tensor Core throughput, done. The actual picture, once you're fusing GEMMs instead of running them in isolation, is that FP8 trades some of that throughput gain against a layout-conformance cost FP16 doesn't have. FlashAttention-3 still comes out ahead with FP8 (the numbers in the previous post show that clearly) — but "ahead" is a net result that already has this shuffle cost baked into it, not a number you get by assuming FP8 is a drop-in replacement for FP16 inside a fused kernel.

References

Layout comparison diagram above is original artwork made for this post.