Skip to content

fp8linear ΒΆ

FP8 quantization-aware training for linear layers.

Mirror of fp4linear.py but for FP8 (e4m3). The forward pass quantizes both activations and weights to FP8 and runs torch._scaled_mm; the backward pass is a bf16 straight-through estimator so the high-precision master weights stay trainable. Falls back to a bf16 fake-quant forward on GPUs older than sm89.