Mixed-Precision Kernels

A single A100 GPU delivers 312 TFLOPS in FP16/BF16 but only 77 TFLOPS in FP32. That 4x gap is not free: you earn it by convincing every kernel in your training or inference stack to operate at reduced precision without corrupting the model. That gap between available compute and what naive code actually uses is the central tension mixed-precision engineering resolves.

The Precision Hierarchy and What Each Format Costs

Modern GPU kernels work across at least four floating-point formats. Understanding their bit layouts explains which operations break at low precision and which do not.

BF16 has the same exponent width as FP32, so it never overflows where FP32 does not. FP16's 5-bit exponent caps at 65504, which is why gradient norms routinely overflow it during training without intervention. FP8 pushes further; the E4M3 sub-format preserves more mantissa bits for forward-pass activations, E5M2 preserves more range for backward-pass gradients.

The arithmetic throughput ratio is roughly linear in mantissa bit width for matrix multiplications, because the hardware tensor cores natively accumulate in a wider format. An A100's FP16 tensor core multiplies two FP16 inputs but accumulates into FP32; the final write can be downcast. This accumulate-in-higher-precision pattern is the foundation all mixed-precision kernels exploit.

The Three-Tier Storage Pattern

Micikevicius et al. (ICLR 2018) codified the canonical pattern for mixed-precision training: store weights in FP32 as a "master copy", cast them to FP16 for each forward and backward pass, accumulate gradients in FP32, and update the master copy. The memory cost of carrying both copies is offset by the halved bandwidth of every activation tensor.

Master weights  (FP32, on DRAM)
       |
   cast to FP16
       |
Forward pass    (FP16 compute, FP32 accumulation in tensor cores)
       |
Activations     (FP16, checkpointed or recomputed)
       |
Backward pass   (FP16 compute)
       |
Gradients       (FP16) ─── loss scaling ──► FP32 grads ──► FP32 weight update

Loss scaling multiplies the loss by a large constant S (commonly 2^8 to 2^15) before the backward pass, shifting gradient magnitudes into the representable FP16 range, then divides the accumulated FP32 gradients by S before the weight update. Dynamic loss scaling adjusts S up when no inf/NaN appears for N consecutive steps, and halves S immediately on any inf/NaN detection.

PyTorch's torch.cuda.amp.autocast context manager and GradScaler implement exactly this pattern. autocast routes eligible ops (matmul, conv, attention) to FP16 or BF16 tensor cores; GradScaler handles the scaling bookkeeping. With BF16, loss scaling is unnecessary because overflow is not a concern, which is why Ampere and later hardware (A100, H100) train comfortably with dtype=torch.bfloat16 and no scaler.

Kernel-Level Mechanics: Cast Placement and Fusion

The performance of a mixed-precision kernel depends not just on which dtype is used, but on where casts are inserted relative to memory transactions.

Consider layer normalisation: the mean and variance computation requires accumulation across a row of activations. If those activations are FP16, the partial sums in a naive kernel overflow or lose precision. The correct approach is to load FP16 activations, immediately widen to FP32 for the running sum, compute the normalised value in FP32, and write back in FP16. The cast happens in registers; no FP32 tensor is materialised in DRAM.

// Pseudocode: mixed-precision layer norm kernel
__global__ void layernorm_fp16(half* x, half* y, float* gamma, float* beta, int N) {
    float sum = 0.f, sq_sum = 0.f;
    for (int i = ...) {
        float v = __half2float(x[i]);  // widen on load
        sum    += v;
        sq_sum += v * v;
    }
    float mean = sum / N;
    float rstd = rsqrtf(sq_sum / N - mean * mean + 1e-5f);
    for (int i = ...) {
        float v = __half2float(x[i]);
        y[i] = __float2half(gamma[i] * (v - mean) * rstd + beta[i]);  // narrow on store
    }
}

This "load-narrow, compute-wide, store-narrow" idiom is the template for writing numerically safe mixed-precision kernels. The key insight: registers on modern GPUs hold FP32 with no throughput penalty; the savings come from halved DRAM bandwidth, not from running all arithmetic in FP16.

Kernel fusion compounds this benefit. A fused attention kernel (as in FlashAttention) keeps the entire softmax-weighted reduction in registers or SRAM, never writing intermediate FP32 tensors to global memory. The operational intensity (FLOPs per byte of DRAM traffic) rises sharply, which is why fused implementations see 2-4x wall-clock speedups even at the same arithmetic precision.

Quantised Kernels: INT8 and FP8 Inference

At inference time, the master-weights pattern is unnecessary because no backward pass occurs. Weights can be quantised to INT8 or FP8 and stay there. The practical pipeline for static quantisation:

1. Calibrate: run a small representative dataset; record per-tensor max absolute values.
2. Compute scale factors: scale = max_abs / 127.0 for INT8 symmetric quantisation.
3. Store quantised weights; dequantise in the kernel fused with the matmul.
4. Accumulate in INT32 (for INT8) or FP32 (for FP8), then downcast the output activation.

The dequantise-inside-matmul approach avoids a separate dequantise kernel and keeps DRAM traffic at INT8 rates. NVIDIA's cuBLAS and cuDNN expose INT8 tensor-core GEMMs via cublasGemmEx; at batch sizes beyond ~32, INT8 throughput (4x that of FP16 on A100) becomes the dominant term.

FP8 support on H100 via the transformer_engine library extends this to training: E4M3 for forward activations, E5M2 for backward gradients, with per-tensor or per-channel scaling. FP8 training is more numerically fragile than BF16 and typically requires careful tuning of scaling granularity.

When It Falls Down

Gradient underflow in FP16 without scaling. Gradients in later layers of deep networks routinely have magnitudes below 2^-14, the minimum normal FP16 value. Without loss scaling they flush to zero, and the model fails to train. BF16 avoids this; FP16 does not.

Accumulation error in long reductions. Summing 8192 FP16 values naively can accumulate rounding error proportional to sqrt(N). Widening to FP32 mid-kernel is cheap but easy to forget. Poorly written custom CUDA kernels silently degrade model quality this way.

Calibration distribution shift for INT8. If the calibration set does not match the deployment distribution, scale factors are wrong, and outlier activations clip hard. This is the dominant failure mode for LLM INT8 quantisation; methods like SmoothQuant migrate outliers from activations to weights before quantising to mitigate it.

Tensor-core alignment requirements. On A100, FP16 tensor-core GEMMs require M, N, K dimensions to be multiples of 8 (multiples of 16 for BF16, multiples of 32 for INT8). Odd-sized tensors fall back to CUDA cores and lose most of the throughput advantage. Padding is the standard workaround; this matters acutely for vocabulary-size matmuls with prime-number vocab sizes.

FP8 instability in fine-tuning. FP8 training is sensitive to learning rate schedule and weight initialisation in ways BF16 is not. Short fine-tuning runs that work in BF16 can diverge in FP8 when per-tensor scaling granularity is too coarse for the weight distribution after many gradient steps.

Mixed-precision and non-determinism. Reducing a vector in different thread orderings produces slightly different FP16/BF16 sums. This is expected and acceptable for training but can cause reproducibility headaches in unit tests that compare floating-point outputs exactly.

Further Reading

• Mixed Precision Training - Micikevicius et al., ICLR 2018
• PyTorch Automatic Mixed Precision Examples
• Triton Fused Softmax Tutorial
• NVIDIA TensorRT Developer Guide - Precision Control