Commit Graph
8 Commits
Author SHA1 Message Date
ViperEkura 74e694921c perf: speed up fp8 gemm tiles and scheduling
- K tile 32->64 (new default): fewer barriers, more MMA per stage; generalize tile_at swizzle and load_operand_tile accordingly
- 64x128 small-M CTA for m<=64 (2x at 64x4096x4096)
- L2 rasterization for crosswise-A layouts (+6..21%)
- micro-bench: NT 4096^3 +35%; linear fwd 1.24-1.76x, bwd 1.71-2.27x vs bf16
- add csrc/tests/fp8_test.cu (single MMA demo + GEMM layouts x K-tiles vs CPU reference)
2026-08-24 18:29:55 +08:00
ViperEkura d5067af064 refactor: harden param PODs and CUTLASS-style fp8 layout tags
- NSDMI null/-1 defaults for AttentionParams/FP8Params pointer+flag members: partially packed structs can no longer hold garbage non-null pointers that gate optional paths (root cause class of the paged test bug); still aggregates, still trivially copyable
- move per-lane ldmatrix wrappers (ldsm_x2/x4) from fp8/gemm.cuh to common/mma.cuh as ldmatrix_x2_lane/x4_lane, next to the single-address variants
- DEVICE_FORCEINLINE macro in common/mma.cuh (matches layout_policies.cuh, internal linkage)
- frag_addr now delegates to tile_at: the swizzle math has one source
- operand layouts as CUTLASS-style RowMajor/ColMajor tags threaded from launch_fp8_gemm through the kernel to load_operand_tile; B's operand view via transpose_layout_t; call sites read <Fmt, false, RowMajor, ColMajor> instead of <Fmt, false, false, true>
2026-08-24 15:27:52 +08:00
ViperEkura 31ca357c61 refactor: namespace csrc kernels and extract common helpers
- attention family -> astrai::attention; fp8 family -> astrai::fp8
- new common/reduce.cuh (warp/group reductions, atomic_max_float)
- new common/cp_async.cuh (predicated cp_async_16, commit/wait group)
- move MAX_SPLITS into attention/common.h; delete warp_utils.cuh
- .cu bindings and pure C tests open family namespaces via using
2026-08-24 14:49:55 +08:00
ViperEkura 34471252ab perf: pipeline fp8 gemm fragment loads and pack bf16 epilogue
- software-pipeline A-fragment ldmatrix: row mt+1 loads hide behind row mt MMAs
- bf16 epilogue packs two columns into one bfloat162 store (half the stores)
- fp8 vs bf16 linear: fwd 1.15x@512, 1.5x@2048, 2.8x@4096; bwd up to 2.7x, peak 35 TFLOPS
2026-08-23 21:32:20 +08:00
ViperEkura aa08479285 perf: widen fp8 gemm tile and load fragments with ldmatrix
- 128x128 CTA of 8 warps x 64x32 warp tiles: 16 mma.sync per warp per K-segment (was 8)
- ldmatrix.x4/x2 with per-lane swizzled addresses replaces 36 scalar LDS per warp-tile step
- __launch_bounds__(256, 2) caps registers at 124 so two CTAs fit per SM
- fp8 linear vs bf16 cuBLAS: fwd 1.07x->2.65x, bwd 1.41x->2.65x by size, peak 31-33 TFLOPS
2026-08-23 21:13:09 +08:00
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 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