Warp Matrix Multiply and Accumulate (WMMA) Architecture

The Physics of NVIDIA Tensor Cores

Standard CUDA cores execute scalar FMA operations per thread. Tensor Cores execute tiled matrix multiply-accumulate (D = A × B + C) instructions cooperatively across warps or warp-groups. Across modern NVIDIA generations (CUDA 12.8–13.1), Tensor Core PTX ISA evolved through three distinct hardware paradigms:

• Ampere / Ada (sm_80 / sm_89): Warp-level (32 threads) wmma.mma.sync and fine-grained mma.sync.aligned.m16n8k16 holding fragments in thread registers.
• Hopper (sm_90a, H100): 4th-Gen Tensor Cores introduce wgmma.mma_async (Warp-Group MMA across 4 warps = 128 threads), streaming tiles asynchronously from Shared Memory via the hardware TMA (Tensor Memory Accelerator).
• Blackwell (sm_100a, B200 / GB200): 5th-Gen Tensor Cores introduce tcgen05.mma backed by dedicated on-chip Tensor Memory (TMEM) for accumulators and native microscaled FP8, FP6, and FP4 (MXFP4) formats.

WMMA Fragment Geometry

All 32 threads in a warp cooperate to hold the matrix fragments in their register files:

// CUDA PTX WMMA instruction format accessed via Rust FFI
// Ampere/Ada, Hopper & Blackwell Tensor Core PTX instructions:
// Warp-level WMMA (32 threads): wmma.mma.sync.aligned.m16n16k16.row.col.f32.f16.f16.f32
// Hopper Warp-Group (128 threads): wgmma.mma_async.sync.aligned.m64n128k16.f32.f16.f16
// Blackwell 5th-Gen (TMEM): tcgen05.mma.cta_group::1.kind::f16

Accessing Tensor Cores from Rust directly or through cuBLAS/CUTLASS achieves up to 989 TFLOPS of dense FP16/BF16 and 1,979 TFLOPS of dense FP8 compute on H100 SXM5 (compared to ~67 TFLOPS on standard CUDA FP32 cores), and scales even higher with Blackwell tcgen05 microscaled MXFP4.

FlashAttention SRAM Tiling

Standard self-attention materializes an N × N score matrix in slow HBM, making attention memory-bound at long context lengths. FlashAttention loads Q, K, and V blocks into fast on-chip Shared Memory (SMEM / TMEM), computes tiled Tensor Core GEMMs (Sij = QiKjT), and maintains running row-max mi and denominator ℓi statistics via online softmax rescaling—never writing the O(N2) intermediate matrix to HBM.