Commit Graph
4 Commits
Author SHA1 Message Date
ViperEkura 6354dbe8bc perf: interleave prefetch into the mma phase
- our 128x128 fast loop padded 26 NOPs between the 32 QMMAs while all four LDGSTS sat bunched at the loop tail: ptxas had no independent instructions to fill the tensor-pipe issue gaps, the exact structure the decompiled cuBLAS loop (166 i, NOP=0) and CUTLASS MmaMultistage avoid by issuing cp.async in small groups inside the MMA phase (copy_tiles_and_advance per warp-tile batch)
- the steady-state prefetch is now a loop-carried register pair per congruous operand (PrefetchCarry: swizzled stage offset + global source, constructed once from the same (r, c0) mapping as the interior loader), whose chunks ride after the first and last k_seg MMA batches — SASS: 282 -> 110 instructions, 0 branches, 0 UIMAD.WIDE magic-divisions, LDGSTS interleaved inside the QMMA range, 26 -> 19 NOPs, still 128 regs (2 CTAs/SM)
- the wait-count dispatch ladder (16 instructions of ISETP/SEL picking DEPBAR immediates) and the per-k-tile (tile % ring) * stage_bytes recomputation (UIMAD.WIDE by 0x55555555) are gone: the prologue commits unconditionally so the wait_group<kStages-1> immediate is valid for every iteration, and both read and write stage addresses advance as carried pointers with an equality wrap
- cp_async.cuh splits the emitter from its policies: one raw PTX site (cp_async_16_raw) plus wrappers for unconditional/predicated and pointer/offset destinations, and the now-unused wait_group_dispatch ladder is deleted; the dispatch flip: with the stall gone the big CTA wins the whole former dip band (1280^3 kernel-level 98.2->104.4T), so prefer_small_cta keeps only the sub-5/8-wave band and the single-wave s3 variant is retired

Benchmark: L20 (sm_89), kernel-level sweep 128s2ff: 1024^3 104.2->114.0T (cuBLAS 151.2), 1152^3 130.4->145.9T (157.2), 1280^3 98.2->104.7T (163.8), 1536^3 141.4->150.4T (182.3), 2048^3 178.1->189.6T (202.8), 4096^3 197.6->208.8T (223.7), 8192^3 207.6->219.1T (227.5). CUDA-graph e2e: 1024^3 108.2->118.4T, 1152^3 137.7->150.7T, 1536^3 143.4->153.7T, 2048^3 179.4->191.9T, 8192^3 199.0->208.9T, 1280^3 108.3->106.6T (old dip-band rule re-measured 106.1T — within noise). Four-layout C++ suite and 596 pytests pass.
2026-08-26 17:47:24 +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 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