perf: fp8 batched gemm and measured dispatch table
- mm_fp8 accepts 3D operands through the same signature: grid.z slices by batch strides, size-1 batches broadcast (stride 0), inner .t() views fold into the layout tag at zero copy - fix _LinearFp8 backward crash on 3D [B,L,d] training inputs (flatten before mm_fp8, reduce grad_b over leading dims) - expose kRasterGroup/kStreamOut as template knobs; drop the 64x128 mid CTA and staged crosswise-B path from dispatch (direct wins everywhere re-measured, including DRAM-streamed B) - dispatch thresholds grounded in fresh sweeps: m<=64 -> 64x64 CTA (+27% at 64x8192x2048), small-CTA crossover at SM*14/3 total tiles (+13% at 96 tiles), threshold counts batch x per-matrix tiles (+31% at 64x512^3 bmm, +25% at 8x1024x2048) - remove scripts/tools/bench_fp8_gemm.py (superseded by csrc/tests/fp8_sweep.cu for kernel-level tuning) Benchmark: NVIDIA L20, E4M3, NT pre-quantized, median of 100-200 iters - 64x8192x2048: 29.1 -> 22.8 us (94 TF/s) - 1024x1536x2048: 67.4 -> 59.6 us (108 TF/s) - bmm 64x512^3: 139.8 -> 106.7 us; bmm 8x1024x2048: 186 TF/s - regression-free: 4096^3 192 TF/s, 8192^3 200 TF/s, 512^3 unchanged
This commit is contained in:
@@ -56,6 +56,12 @@ takes the combined dequant scale (`sa * sb`); the strategy layer passes
|
||||
`scale.reciprocal()` / `sa * sb` respectively. `amax` is always returned in
|
||||
the original input domain.
|
||||
|
||||
`mm_fp8` also accepts 3D (batched) operands through the same signature:
|
||||
`grid.z` slices the operands by their batch strides, a size-1 batch
|
||||
broadcasts (stride 0), and inner-transposed views (e.g. `x.t()`) fold into
|
||||
the kernel's layout tag at zero copy — only genuinely strided operands pay
|
||||
a `.contiguous()` copy.
|
||||
|
||||
Python layer (two levels): `astrai/extension/ops/fp8.py` provides stateless
|
||||
primitives (`quantize` / `mm_fp8`) via `torch.library.custom_op`, and
|
||||
`astrai/extension/fp8.py` is the strategy layer (`fp8_autocast`, delayed /
|
||||
|
||||
Reference in New Issue
Block a user