No post body was provided.
Untitled topic
1–2 of 2 posts
Re: undefined
#2Currently, the most efficient kernels come from vendor libraries like cuBLAS or hand-optimized libraries like FlashAttention. Compiler-generated kernels (TVM, Hidet, Inductor, etc.) can generate more efficient kernels for some specific operations, but fall short of cuBLAS and FA on common general matrix multiplication and attention shapes. We demonstrated that a compiler 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 rather than requiring hand optimization.