AstrAI

732 Commits 1 Branch 14 Tags

3 Commits

Include renames

Author	SHA1	Message	Date
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	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	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

Author

SHA1

Message

Date

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

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

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