Mixed Precision & Tensor Cores

Mixed-precision practice: store/compute in FP16/BF16, accumulate in FP32, use loss scaling for FP16 training.

27.1 Reduced-precision types

__half / __half2 16 (5-bit exponent) inference, graphics 16 (8-bit exponent) training (FP32 range) 19 (in Tensor Cores) automatic in cuBLAS/cuDNN FP32 GEMMs on Ampere+ __nv_fp8_e4m3, __nv_fp8_e5m2 Hopper/Blackwell training & inference FP6 / FP4 (e2m1)

Mixed-precision practice: store/compute in FP16/BF16, accumulate in FP32, use loss scaling for FP16 training.

#include <cuda_fp16.h>
__global__ void haxpy(int n, __half a, const __half2* x, __half2* y) {
    int i = blockIdx.x * blockDim.x + threadIdx.x;
    if (i < n / 2) y[i] = __hfma2(__half2half2(a), x[i], y[i]);   // 2 FP16 ops per instruction
}
__float2half(f), __half2float(h), __hadd(a, b), __hmul2(a, b)

27.2 Tensor Cores with WMMA (warp matrix multiply-accumulate)

Lower-level PTX mma.sync / Hopper wgmma (warpgroup MMA) and Blackwell tcgen05 instructions give more control; in practice use cuBLAS/cuBLASLt, CUTLASS/CuTe, cuDNN, or Triton to exploit Tensor Cores. Structured sparsity (2:4) doubles throughput on Ampere+.

#include <mma.h>
using namespace nvcuda;

// Each warp computes one 16×16 tile of C = A×B (A, B in half, C in float)
__global__ void wmmaGemm(const half* A, const half* B, float* C, int M, int N, int K) {
    int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / warpSize;
    int warpN = blockIdx.y * blockDim.y + threadIdx.y;

    wmma::fragment<wmma::matrix_a, 16, 16, 16, half, wmma::row_major> aFrag;
    wmma::fragment<wmma::matrix_b, 16, 16, 16, half, wmma::col_major> bFrag;
    wmma::fragment<wmma::accumulator, 16, 16, 16, float> cFrag;
    wmma::fill_fragment(cFrag, 0.0f);

    for (int k = 0; k < K; k += 16) {
        int aRow = warpM * 16, bCol = warpN * 16;
        if (aRow < M && bCol < N) {
            wmma::load_matrix_sync(aFrag, A + aRow * K + k, K);
            wmma::load_matrix_sync(bFrag, B + bCol * K + k, K);
            wmma::mma_sync(cFrag, aFrag, bFrag, cFrag);       // Tensor Core operation
        }
    }
    int cRow = warpM * 16, cCol = warpN * 16;
    if (cRow < M && cCol < N)
        wmma::store_matrix_sync(C + cRow * N + cCol, cFrag, N, wmma::mem_row_major);
}