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