Numerics discipline for low-precision inference

Most “my kernel is broken” bugs are actually numerics bugs. The kernel runs, produces plausible-looking numbers, and quietly diverges from the reference at long sequence lengths or specific input distributions. This document is the checklist that prevents 90% of those bugs.

The dtype landscape

Dtype

Bits

Range

Precision

Where it’s used

FP32

32

~1e-38 to 3.4e38

7 decimal digits

Accumulators, softmax state, master weights

TF32

19

FP32 range

10-bit mantissa

Ampere+ tensor-core input (implicit)

FP16

16

~6e-5 to 65k

3-4 decimal digits

Legacy training; still common at inference

BF16

16

FP32 range

2-3 decimal digits

Modern default for LLMs. Same exponent as fp32

FP8 (E4M3)

8

~2e-4 to 448

~2 digits

H100+ tensor cores, weights + activations

FP8 (E5M2)

8

~6e-5 to 57344

~1 digit

Gradients (wider range, less precision)

NVFP4

4

~2e-4 to 6

tiny

Blackwell weights (block scale 16 + FP32 second-level)

MXFP4

4

~2e-4 to 6

tiny

OCP standard, block scale 32

INT8

8

-128 to 127

linear

Weight-only quant, activations in some setups

INT4

4

-8 to 7 (or 0-15)

linear

Weight-only quant (GPTQ, AWQ, Marlin)

The one rule that solves half of numerics bugs

Accumulate in FP32. Always. Inputs bf16, outputs bf16, accumulator fp32.

Tensor cores physically implement this: bf16×bf16 → fp32 accumulator, then cast back to bf16 on write. In Triton, tl.dot(a, b) defaults to fp32 accumulator when a and b are bf16. Don’t override this.

Corollaries:

  • Softmax state (m, l) in fp32.

  • Running output o in fp32 until the final store.

  • Any sum / mean / var reduction in fp32, even if inputs are bf16.

  • LayerNorm/RMSNorm reduction in fp32.

  • Only the final tensor written to HBM should be the target dtype.

Tolerances — what “close enough” means

Operation

Dtype pair

atol

rtol

GEMM (single)

bf16 in, bf16 out, fp32 acc

1e-2

1e-2

GEMM chain (attention)

bf16 in, bf16 out

5e-3

5e-3

RMSNorm

bf16 in, bf16 out

1e-3

1e-3

Softmax

bf16 in, bf16 out

1e-3

1e-3

Attention forward

bf16 in, bf16 out

1e-2

1e-2

Attention forward

fp8 in, bf16 out

5e-2

5e-2

FP32 references

fp32

1e-5

1e-5

These are guidelines, not absolutes. Longer sequences and larger tensors accumulate more rounding error. If you’re inside these bounds you probably have no numerics bug. If you’re outside them, you almost certainly do.

PyTorch verification pattern

import torch

def verify(ours, ref, atol=1e-2, rtol=1e-2, name="kernel"):
    diff = (ours.float() - ref.float()).abs()
    max_abs = diff.max().item()
    max_rel = (diff / (ref.float().abs() + 1e-6)).max().item()
    print(f"[{name}] max_abs={max_abs:.4e}  max_rel={max_rel:.4e}")
    torch.testing.assert_close(ours, ref, atol=atol, rtol=rtol)

Always print max_abs and max_rel even when the test passes — they’re your regression signal.

Failure smells — what NaN patterns tell you

Symptom

Likely cause

All-NaN output

0 * inf, or inf - inf in a mask/subtract, or division by 0

First row NaN, rest OK

Fully-masked first row in causal attention (esp. sliding window at pos 0)

NaN appears at N > 4k, fine at N < 2k

fp32 accumulator missing; bf16 overflow in sum(exp(...))

Output all zeros

Wrong scale applied; mask value 0 instead of -inf; softmax denominator = inf

Correct at atol=1e-1 but not 1e-2

Softmax scale ordering (post-multiply instead of pre-multiply Q); loss of a ULP

Correct on H100 but wrong on A100

TF32 vs FP32 mismatch: A100’s tensor cores use TF32 by default for fp32 matmul; force fp32 or match TF32 explicitly

Correct with is_causal=False but wrong with is_causal=True

Mask direction wrong or mask value = 0

Nondeterministic (differs run-to-run)

Some reduction in bf16 accumulator (bf16 add is non-associative and warp scheduling is nondeterministic)

The FP8 rules (Hopper/Blackwell)

FP8 is the first dtype where naive casts give bad results. Key rules:

  1. Use E4M3 for weights & activations, E5M2 for gradients (which you don’t care about at inference). E4M3 has more precision at the expense of range; use it where clipping isn’t a problem.

  2. Always block-quantize. Global scale × global tensor loses too much accuracy. Per-tile or per-channel scaling is the norm. FP8 activations in FA3 use block quantization along the sequence dim.

  3. Delayed scaling for activations (during training) — use a moving max. For inference, static per-tensor scale calibration works.

  4. Incoherent processing / Hadamard rotation (from FA3): pre-rotate Q, K, V by a random orthogonal transform. Spreads outliers, reducing quant error 2–3×. See FA3 §3.

  5. Accumulator stays fp32. FP8 tensor cores emit fp32; you cast to bf16 or fp16 on the final write.

FP8 tolerance: atol=5e-2, rtol=5e-2 is typical. If you can hold 1e-2 you’re doing something impressive.

The FP4 rules (Blackwell)

FP4 (NVFP4 or MXFP4) is currently weight-only. Activations stay bf16 / fp8.

  1. NVFP4: block size 16, per-block E4M3 scale, plus per-tensor FP32 scale. Tighter block → better accuracy.

  2. MXFP4: block size 32, OCP standard. More portable but coarser.

  3. Dequantize on-chip. FP4 weights are loaded from HBM as packed uint8 (2 weights per byte), dequantized in registers with the per-block scale, consumed directly by tensor cores.

  4. Model accuracy loss should be < 1% on standard LM eval tasks if calibration is done properly. If you’re losing > 3%, your calibration set is bad or quantization is applied to layers that shouldn’t be quantized (e.g., LM head, layernorm weights).

Reference: NVIDIA NVFP4 blog: https://developer.nvidia.com/blog/introducing-nvfp4-for-efficient-and-accurate-low-precision-inference

Determinism trade-off

BF16 add is non-associative: (a + b) + c ≠ a + (b + c) at the last bit. Combined with nondeterministic warp scheduling, kernel outputs can differ by 1–2 ULP run-to-run.

When determinism matters (regression tests, reproducibility research):

  • Set torch.use_deterministic_algorithms(True)

  • Force CUBLAS_WORKSPACE_CONFIG=:4096:8

  • Accept the (usually small) perf hit

When it doesn’t (production serving): allow nondeterminism. The variance is smaller than the model’s inherent noise.

The 5 numerics debugging tactics that actually work

  1. Bisect on shape. If broken at N=8192, works at N=512, the bug scales with N. Almost certainly accumulator precision.

  2. Bisect on input. Replace random inputs with all-ones. If bug persists, it’s algorithmic. If bug vanishes, it’s input-distribution-sensitive (outliers, extreme values).

  3. Print intermediate tensors. In Triton use tl.static_print (compile-time) or TRITON_INTERPRET=1 (runtime) to get printable state inside kernels.

  4. Reference in fp32. Compute the reference in fp32, cast to bf16 for comparison. Isolates “my kernel is wrong” from “my kernel is fine but bf16 rounds differently.”

  5. Turn off masking. If bug vanishes without the causal mask, mask value or comparison direction is wrong.

Standard test harness for a new kernel

def standard_kernel_test(ours_fn, ref_fn, shapes, dtype=torch.bfloat16):
    for shape in shapes:
        for seed in [0, 1, 42]:
            torch.manual_seed(seed)
            args = [torch.randn(*s, dtype=dtype, device="cuda") for s in shape]
            ours = ours_fn(*args)
            ref  = ref_fn(*args)
            verify(ours, ref)

    # Adversarial
    args = [torch.full(shapes[0][0], 1e4, dtype=dtype, device="cuda")]
    ours = ours_fn(*args); ref = ref_fn(*args)
    assert not torch.isnan(ours).any(), "NaN under extreme inputs"

    # Edge: all-zero
    args = [torch.zeros(*shapes[0][0], dtype=dtype, device="cuda")]
    ours = ours_fn(*args); ref = ref_fn(*args)
    verify(ours, ref)

Every new kernel gets this treatment. It takes 5 minutes to write once and saves days of debugging.

References