Commit Graph
10 Commits
Author SHA1 Message Date
ViperEkura 7da1439c9e 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
2026-08-24 19:25:23 +08:00
ViperEkura 4b10d3ca37 perf: vectorize fp8 quantize and swizzle gemm smem 2026-08-23 20:31:44 +08:00
ViperEkura 2bc4d2b8a8 refactor: unify kernel module loading and packaging
- loader.py: lazy/cached import; is_available defers the actual load; get_module raises on unavailable
- ops/{attention,rotary,fp8}: use get_module instead of touching private _modules or their own _mod() cache
- package-data: ship astrai.extension.lib *.so in built wheels (non-editable installs previously lost every kernel)
2026-08-23 15:57:04 +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 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 3d3ea47d37 refactor: standardize packed 3d inference
- keep training attention on dense 4d tensors
- use packed 3d tensors with KV cache for inference
- extend CUDA rotary embedding to packed 3d inputs
- adapt torch, CUDA and FlashAttention backend dispatch

Benchmark: NVIDIA L20, BF16, 1B model, paged KV cache, CUDA Graph, prompt 512, generation 128 (median of 3 alternating runs)
- batch 1: 234.5 -> 242.6 tok/s (1.034x, +3.4%)
- batch 8: 1243.1 -> 1286.6 tok/s (1.035x, +3.5%)
2026-08-19 00:36:53 +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
ViperEkura 1bcd8f53ab perf: precompute ragged Q tile scheduling 2026-08-16 23:32:46 +08:00
ViperEkura 6ac3b51496 refactor: separate extension ops and backends 2026-08-16 21:15:52 +08:00