gkapur

TMA Transport for Matmul and Flash Kernels. cuBLAS and Flash Attention kernels still use the same cp.async transport on consumer Blackwell dies (in fact, most cuBLAS GEMM kernels on consumer Blackwell, including the FP16 tensorop path, are forward-ported Ampere-era cutlass_80_* kernels). Swapping cp.async with TMA allows us to reduce the number of instructions kernels need to issue, and the TMA’s swizzle drops shared-memory bank conflicts for free.

Hybrid FP16/FP32 Accumulation. The default matmul path on the production stack uses FP16 tensor cores with FP32 accumulation. However, on consumer dies, the FP16-input/FP32-accumulate HMMA runs at exactly half the rate of FP16-input/FP16-accumulate. To work around this, I use the fast atom mma_m16n8k16_f16_f16, but keep accuracy in check by promoting the FP16 partials into the FP32 registers and thus doing global accumulation accurately in FP32. So you get the FP16 tensor-core speed with an FP32 accumulation instead of paying the FP32-accumulate tax on every single mma. We measured the error of C = A@B using FP16 inputs drawn from N(0,1), comparing each accumulation strategy against an FP64 reference over the identical FP16-rounded operands. All configurations land within one standard error of each other, so the hybrid accumulation does not degrade task quality, matching the kernel-level error analysis.

The compiler approach allowed us to introduce both of these optimizations into all GEMM and Attention kernels used in the model rather than requiring hand optimization.

We were able to demonstrate that the compiler approach can match or even outperform vendor libraries. On Gemma 4 12B, a small open-weight model, running on RTX 5090, compiler-generated kernels achieved up to 1.6× speedup over cuBLASLt, and a 1.30× geometric-mean speedup over PyTorch.

Primarily, two optimizations allowed the Emmy compiler to outperform cuBLAS on common GEMM and FA shapes on RTX 5090 and RTX 4090: TMA transport (Blackwell-only) and leveraging full FP16 tensor cores with FP32 shadow accumulation registers:

TMA Transport for Matmul and Flash Kernels. cuBLAS and Flash Attention kernels still use the same cp.async transport on consumer Blackwell dies (in fact, most cuBLAS GEMM kernels on consumer Blackwell, including the FP16 tensorop path, are forward-ported Ampere-era cutlass_80_* kernels). Swapping cp.async with TMA allows us to reduce the number of instructions kernels need to issue, and the TMA’s swizzle drops shared-memory bank conflicts for free.

Hybrid FP16/FP32 Accumulation. The default matmul path on the production stack uses FP16 tensor cores with FP32 accumulation. However, on consumer dies, the FP16-input/FP32-accumulate HMMA runs at exactly half the rate of FP16-input/FP16-accumulate. To work around this, I use the fast atom mma_m16n8k16_f16_f16, but keep accuracy in check by promoting the FP16 partials into the FP32 registers and thus doing global accumulation accurately in FP32. So you get the FP16 tensor-core speed with an FP32 accumulation instead of paying the FP32-accumulate tax on every single mma. We measured the error of C = A@B using FP16 inputs drawn from N(0,1), comparing each accumulation strategy against an FP64 reference over the identical FP16-rounded operands. All configurations land within one standard error of each other, so the hybrid accumulation does not degrade task quality, matching the kernel-level error analysis. The compiler approach allowed us to introduce both of these optimizations into all GEMM and Attention kernels used in the model.

The kernel reaches parity with FA-2 (206us on RTX5090 with batch=1, heads=8, seq_len=4096, head_dim=64), but unfortunately, FA-3/4 optimizations are either not applicable or not helpful on consumer cards. It looks like FA-2 is the ceiling.

In summary: - Faster tensor-core instructions (WGMMA) are the main lever behind FA-3, but they are not available on RTX GPUs. - TMA (tensor memory accelerator) is available on sm_120 (RTX 5XXX). It helps on paper (LSU drops), but the transport isn't the bottleneck, so the final number barely moves. - Warp specialization is also available. However, it is mainly a scheduling optimization, i.e., it helps eliminate pipeline bubbles and better utilize tensor cores, but without asynchronous tensor core instructions. The result is negative: 213 vs 206 us. - In FA-4, they also simulated exp using FMA instructions because the tensor cores on the B200 are so fast that the whole pipeline became SFU-bound (special functions unit). RTX 5090 is tensor-core bound, so no point in this optimization either. In fact, even a conventional optimization of using faster exp2f instead of expf for softmax doesn't move the number.

I have tried a handful of other optimizations that could potentially work on consumer silicon, such as a deeper pipeline and register ping-pong. No luck. Since the whole pipeline is tensor-pipe-bound, I believe the FA-2 is the ceiling, and that all meaningful levers will require sacrificing some accuracy to leverage faster, lower-precision tensor cores.

Note that this is an exploration of the regular attention that dominates the prefill- and compute-bound regimes. Decoding against a large KV cache is a different, memory-bound story where split-KV/Flash-Decoding matters more than any of the above.