ViperEkura
69f35c46e0
fix: quantize amax from raw values, not scaled fp8 values
...
- amax for delayed scale was the quantized max (always ~448), so scale collapsed to 1
- this made fp8 gradients diverge (cosine 0.05) and training stall
- stop w/x transpose-quantize amax from polluting the grad scale
2026-08-14 14:26:16 +08:00
ViperEkura
5244f1a8fc
feat: add te-style scaled fp8 training via fp8_autocast
...
- per-tensor scales applied inside cublasLt via A_SCALE/B_SCALE
- delayed scaling: weight amax history ring, refresh every 16 steps
- quantize kernels emit atomic amax, device-side scale updates
- fp8_autocast context toggles aten::linear dispatch like torch.autocast
- fallback to bf16 when M/N not 16-aligned (fp8 gemm constraint)
- x/g scales delayed one step, reuse free atomic amax (no abs/max reduce)
2026-08-14 12:14:04 +08:00
ViperEkura
5104638447
perf: use fp8 tensor-core gemm in linear backward
...
- dX/dW run as fp8 cublasLt gemms via fused transpose-cast
- shared (m,k,n) algo cache for fwd/bwd, mutex-protected
- bias add in-place on bf16 output, drop output copy
2026-08-14 10:43:03 +08:00
ViperEkura
a711d9f478
perf: eliminate gemm output transpose via A/B swap
...
- pass w as param A (op=T) and x as param B (op=N) so the col-major [N,M] output storage is row-major C[M,N] directly, zero copy
- transpose_bias_cast kernel becomes a plain bias+write kernel
- fp8 e2e now beats bf16: 1.09x at M=4096, 1.06x at M=8192 (was 0.88x)
2026-08-14 01:42:09 +08:00
ViperEkura
15862d4b56
perf: fuse fp8 linear fwd and bwd into single kernel calls
...
- fp8_linear_forward: cast + cublasLt GEMM + transpose + bias in one call
- fp8_linear_backward: scale-free, dtype derived from input tensor
- drops per-op Python dispatch (was ~6-8 launches per linear) and amax syncs
- 1024x1024 linear: 6.8x slow -> 0.67x (36.7us vs 24.8us bf16)
- small-model e2e still 1.71x slow; 15bt estimate ~0.78x (linear-heavy)
2026-08-14 01:24:37 +08:00
ViperEkura
f9efb705b8
perf: output fp8 gemm in bf16 instead of fp32
...
- cublasLt C layout and buffer switched to CUDA_R_16BF, halving output bandwidth
- downstream ops (RMSNorm etc.) keep matching bf16 dtype, fused kernels stay
- numeric error unchanged (0.19% vs fp32 ref on quantized inputs)
2026-08-14 01:08:26 +08:00
ViperEkura
a5b238dd86
feat: add fp8 training via cublasLt dispatch
...
- fp8_mm kernel (csrc): cublasLt fp8 e4m3 gemm, TN layout mapped zero-copy
- custom::fp8_mm custom op: meta/cuda/cpu kernels + scale-corrected bf16 autograd
- aten::linear and linear_backward dispatch on CUDA key, zero model changes
- per-tensor scale or raw cast; single-GPU smoke loss matches bf16
2026-08-14 00:39:49 +08:00
ViperEkura