feat: static fp8 weights and bias with fused epilogue

- linear_forward_fp8 accepts pre-quantized w8 (matching fmt) and skips the weight quantize; amax_w returns 0 on that path since no bf16 values are seen
- bias is now fused into the GEMM epilogue for both dtypes, replacing the separate torch-level add (one elementwise kernel per linear removed)
- FP8Params.bias becomes void* with a new bias_scale slot: null scale = raw bf16 bias, non-null = fp8 storage dequantized in the epilogue after the operand scaling and before any output quantization
- ops/fp8.py relaxes the w dtype check to bf16-or-fp8 and passes bias_scale through
- regression test covers w8/b8, w8/bf16-bias and the amax_w = 0 contract vs an explicit quantization reference
This commit is contained in:
2026-08-24 19:25:23 +08:00
parent 29e5f571af
commit 7da1439c9e
5 changed files with 114 additions and 44 deletions
+30
View File
@@ -146,6 +146,36 @@ def test_linear_backward_e5m2_gradients():
torch.testing.assert_close(amax_g, grad.abs().amax().float().reshape(1))
@skip_no_fp8
def test_fp8_linear_static_fp8_weight_and_bias():
"""Static fp8 inference: pre-quantized w8/b8 + their scales take the GEMM
directly (no weight quantize, amax_w = 0); the bias is fused in the
epilogue (bf16 and fp8 bias share the fused path)."""
torch.manual_seed(9)
m, n, k = 67, 45, 129
x = torch.randn(m, k, device="cuda", dtype=torch.bfloat16)
weight = torch.randn(n, k, device="cuda", dtype=torch.bfloat16) * 0.5
bias = torch.randn(n, device="cuda", dtype=torch.bfloat16) * 0.5
sx, sw, sb = _scale(x), _scale(weight), _scale(bias)
w8, _ = quantize_bf16(weight, sw, "e4m3")
b8, _ = quantize_bf16(bias, sb, "e4m3")
out, amax_x, amax_w = linear_forward_fp8(x, w8, b8, sx, sw, "e4m3", sb)
qx = _quantize(x, sx)
qw = _quantize(weight, sw)
qb = _quantize(bias, sb)
expected = (qx @ qw.t() * sx * sw + qb * sb).to(torch.bfloat16)
torch.testing.assert_close(out, expected, atol=0.125, rtol=0.01)
torch.testing.assert_close(amax_x, x.abs().amax().float().reshape(1))
assert amax_w.item() == 0.0 # nothing measured on the static path
# bf16 bias stays bf16 on the same fused-epilogue path
out_bf16bias, _, _ = linear_forward_fp8(x, w8, bias, sx, sw, "e4m3")
expected_b = (qx @ qw.t() * sx * sw + bias.float()).to(torch.bfloat16)
torch.testing.assert_close(out_bf16bias, expected_b, atol=0.125, rtol=0.01)
@skip_no_fp8
def test_fp8_linear_backward_outside_autocast():
"""aten::linear records an fp8 autograd node inside fp8_autocast; the