Files
AstrAI/csrc/kernels/fp8
ViperEkura 36e39496d4 perf: vectorize tiled fp8 transpose quantize and arm amax via memset
- tiled transpose quantize becomes one 64x32-tile kernel: native pair loads (128B warp reads) with in-kernel scalar fallback at unaligned or ragged rows, so odd widths and misaligned bases no longer route to a separate kernel
- the old 32x32 scalar tiled kernel and its launcher correctness branch are gone; grid sizing simplifies to 1 + total / (vec * threads) since both elementwise loops are grid-stride
- quantize arms the amax buffer with cudaMemsetAsync instead of the zeros() fill kernel, dropping one tensor-op dispatch and kernel launch per call
- byte-exact parity holds over 1404 golden records (13 shapes x 3 dtypes x 3 scales x 2 formats x 3 layouts x aligned/misaligned) and tests/extension passes 65/65
- elementwise quantize kernel left unchanged: 16B-store pairing, __ldcs streaming hints and amax tree reduction all measured neutral at its ~52% DRAM ceiling and were reverted

Benchmark: L20 (sm_89), profiler kernel time with L2 flushed between calls.
- transposed quantize (layout 1): 230 -> 294 GB/s on 2048x1536 (+28%), 245 -> 299 on 2048x1536 weights (+22%); dual-layout (layout 2) 248 -> 329 (+33%) on the same shapes
- DRAM-saturated sizes (~10.6M elements) regress ~5% (404 -> 384 GB/s on 8192x1536), ~0.02% of a training step; accepted for the single-kernel shape after scalar-path and geometry variants both measured the same
- amax init fill kernel 3.0us -> memset 0.9us; quantize call CPU wall 18.5 -> 13.4us on 128x1536
2026-08-31 14:24:51 +08:00
..