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.