ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity
767 Commits 1 Branch 14 Tags
31ca357c61a39c6606d39fc265ccc98ad58a63c3
Commit Graph
6 Commits
Author SHA1 Message Date
ViperEkura 4b10d3ca37 perf: vectorize fp8 quantize and swizzle gemm smem 2026-08-23 20:31:44 +08:00
ViperEkura 4244df2785 perf: pure FP8 fwd/bwd and lean non-transposed GEMM 
- drop the fused kernel; forward/backward are quantize + a pre-quantized GEMM
- rename module fp8_mm -> fp8_ops (mm.cu -> ops.cu)
- kernels/launchers fp8_gemm_kernel / launch_fp8_gemm; drop PqTraits/gather_trans/pack_fp8x4_vector
- remove the in-kernel transposed-operand branches (TransA/TransB)
- backward: quantize g once (amax_g here), explicit fp8 transposes, fast non-transposed GEMMs (dX = g@w^T, dW = g^T@x^T)
- each pass uses a single FP8 format (E4M3 fwd / E5M2 bwd)
 2026-08-23 15:38:30 +08:00
ViperEkura 75304d084d refactor: unify CUDA skip guards in tests 
- add skip_no_fp8 (CUDA + fp8_mm kernel + cc 8.9+) to tests/conftest.py
- use skip_no_cuda / skip_no_kernel / skip_no_fp8 directly in test modules
- drop _GPU alias and tests.extension.conftest re-exports
- remove unused imports (Union in hf_adapter, make_grpo_config in data conftest)
 2026-08-22 21:08:06 +08:00
ViperEkura 16a55bb474 refactor: reorganize CUDA kernels into per-family directories 
- move attention kernels to csrc/kernels/attention/ and rotary to rotary/
- add shared common/mma.cuh (mma_sync, ldmatrix) and device.cuh (sm checks)
- split fp8_mm into three-layer fp8/common.h, gemm.cuh, mm.cu
- fix fused FP8 GEMM ldmatrix lane indexing to fix OOB shared reads
- update extension ops, loader, and kernel tests
 2026-08-22 20:40:31 +08:00
ViperEkura 7580d80d45 perf: accelerate FP8 backward with fused fast kernel 
- route dX/dW through the fused 128x64 fast kernel via contiguous transposes
- drop the legacy 64x64 kernel, cutting dX 1.55->0.38 ms and dW 1.28->0.26 ms
- sync all threads after cp.async.wait_group to fix sporadic NaN in large GEMMs
- add fp8_mm_prequant_fp8 custom op for FP8-in/FP8-out GEMM
 2026-08-18 23:46:43 +08:00
ViperEkura cb51a3587b perf: optimize fused FP8 GEMM kernel 2026-08-18 19:56:15 +08:00