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 wherever the device has no
e4m3fn _scaled_mm (pre-sm89 CUDA, pre-CDNA4 ROCm); the capability gate and
the FP8 constants are shared with fp8_config so the inference and QAT paths
cannot drift.
fastvideo.layers.fp8linear.fp8_linear_forward
fp8_linear_forward(x: torch.Tensor) -> torch.TensorSource
def fp8_linear_forward(self, x: torch.Tensor) -> torch.Tensor:
# pass config **positionally**; autograd.Function.apply ignores kwargs
return _LinearFWD8BWD16Fn.apply(x, self.weight, self.bias, "tensor"), None