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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user