perf: fp8 rings, lean autocast, gemm staging
- Finalize scale rings inside the quantize kernels: a last-block epilogue (threadfence + counter elect) folds amax into hist, reduces the window and publishes the next scale on device, zero extra launches; _ScaleRing packs [hist | scale | counter] into one CUDA buffer. - Split FP8QuantizeParams out of FP8Params so each operator owns its fields; linear_forward/backward_fp8 take optional ring arguments. - Drop the inference weight-quantization cache; the optimizer bumps the weight version every step, so a cache would miss anyway. - Zero amax scratch via empty + cudaMemsetAsync instead of torch::zeros, cutting a ~50us fill_ dispatch per quantize. - Stage crosswise-B operands K-major with cp.async (contract >= 8192) and PRMT-transpose per k_seg region in smem, interleaved with the MMAs; the sync LDG + byte-scatter path it replaces was long-scoreboard bound (ncu 4.6 vs 0.4 stalls/issue). - Load crosswise-A direct with an in-register PRMT transpose; its operands are typically L2-resident and the staging round trip measured as a net loss. - Enable grouped rasterization for the congruous NT forward (shared B stripe keeps the weight operand hot in L2) and make the smem budget layout-aware (Fp8GemmSmem) while holding two CTAs per SM. - Annotate ops/fp8.py return types; drop weight-cache and decorator tests, hoist their imports to module level. e2e 12L/dim1024/B4xT512 fused AdamW: fp8 137.8ms/step vs bf16 210.3ms, 1.53x. Kernel vs cuBLASLt _scaled_mm: fwd 1.03-1.09x, dX 1.33-1.47x, dW 1.30-1.39x (from 1.10/1.42-1.49/1.52-1.56x), before the pre-transposed copies cuBLASLt needs for dX/dW. fp8 train step vs bf16: 1.34x at 2048 tokens (was 1.25x), 1.08x at 512.
This commit is contained in:
@@ -61,6 +61,12 @@ struct Fp8GemmTraits {
|
||||
static constexpr __nv_fp8_interpretation_t kNvFormat =
|
||||
kIsE5M2 ? __NV_E5M2 : __NV_E4M3;
|
||||
static constexpr float kFp8Max = kIsE5M2 ? 57344.0f : 448.0f;
|
||||
|
||||
// Derived launch geometry: 64x32 warp tiles give the CTA thread count.
|
||||
// The shared-memory budget is layout-aware (crosswise operands add K-
|
||||
// major staging + a canonical buffer), so it lives in Fp8GemmSmem in
|
||||
// gemm.cuh together with the resident-CTA hint for __launch_bounds__.
|
||||
static constexpr int kCtaThreads = (BlockM / 64) * (BlockN / 32) * 32;
|
||||
};
|
||||
|
||||
// Quantize-kernel parameter POD: BF16 -> FP8 with fused amax and optional
|
||||
|
||||
Reference in New Issue
Block a user