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
+2 -2
View File
@@ -75,16 +75,16 @@ struct FP8Params {
// the pre-quantized path. Scales are quantization steps (device scalars).
const void* __restrict__ a_ptr = nullptr;
const void* __restrict__ b_ptr = nullptr;
const void* __restrict__ bias = nullptr;
const float* __restrict__ scale_a = nullptr;
const float* __restrict__ scale_b = nullptr;
const float* __restrict__ bias_scale = nullptr;
// Output: BF16 or FP8 (E4M3). out_scale is the output quantization step
// (FP8 output only).
void* __restrict__ out_ptr = nullptr;
const float* __restrict__ out_scale = nullptr;
// Fused forward extras: bias (may be null) and amax slots (may be null).
const __nv_bfloat16* __restrict__ bias = nullptr;
float* __restrict__ amax_a = nullptr;
float* __restrict__ amax_b = nullptr;