ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: pipeline fp8 gemm fragment loads and pack bf16 epilogue

Browse Source 
- 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
This commit is contained in:
ViperEkura
2026-08-23 21:32:20 +08:00
parent aa08479285
commit 34471252ab
1 changed files with 29 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/fp8/gemm.cuh
+27 -10
View File
@@ -411,16 +411,25 @@ __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) {
frag_addr<T8, kK>(b_smem[stage], row, frag_addr<T8, kK>(b_smem[stage], row,
k_seg * 2 + rh8)); k_seg * 2 + rh8));
} }
// Software-pipelined A fragments: the ldmatrix.x4 for row mt+1
// is issued before the MMAs consuming row mt, so the LDS fixed
// latency hides behind tensor-pipe work (cuts the `wait` stall,
// ~2.3 cycles/issue before this). Costs 4 extra registers.
unsigned a_frag[5][4];
ldsm_x4(a_frag[0],
frag_addr<T8, kK>(a_smem[stage], a_row0 + rh8 * 8 + r7,
k_seg * 2 + rh16));
#pragma unroll #pragma unroll
for (int mt = 0; mt < 4; ++mt) { for (int mt = 0; mt < 4; ++mt) {
unsigned a_frag[4]; if (mt < 3)
const int row = a_row0 + mt * 16 + rh8 * 8 + r7; ldsm_x4(a_frag[mt + 1],
ldsm_x4(a_frag, frag_addr<T8, kK>(
frag_addr<T8, kK>(a_smem[stage], row, a_smem[stage],
a_row0 + (mt + 1) * 16 + rh8 * 8 + r7,
k_seg * 2 + rh16)); k_seg * 2 + rh16));
#pragma unroll #pragma unroll
for (int nt = 0; nt < 4; ++nt) for (int nt = 0; nt < 4; ++nt)
astrai::mma_sync<T8>(acc[nt][mt], a_frag, b_frag[nt], astrai::mma_sync<T8>(acc[nt][mt], a_frag[mt], b_frag[nt],
acc[nt][mt]); acc[nt][mt]);
} }
} }
@@ -439,8 +448,9 @@ __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) {
for (int nt = 0; nt < 4; ++nt) { for (int nt = 0; nt < 4; ++nt) {
const int64_t col = output_col + nt * 8; const int64_t col = output_col + nt * 8;
// Per-row store: FP8 packs two adjacent columns into one 16-bit // Per-row store: FP8 packs two adjacent columns into one 16-bit
// write; the BF16 path writes two scalars. Boundary columns fall // write, BF16 into one 32-bit __nv_bfloat162 (single cvt+pack
// back to a scalar convert so the pack never crosses the row edge. // instruction); boundary or unaligned columns fall back to scalar
// converts so a pack never crosses the row edge or misaligns.
auto store_out = [&](int64_t row, float v0, float v1) { auto store_out = [&](int64_t row, float v0, float v1) {
if (row >= m) return; if (row >= m) return;
if constexpr (OutFp8) { if constexpr (OutFp8) {
@@ -453,10 +463,17 @@ __global__ void __launch_bounds__(kWarps * 32, 2) fp8_gemm_kernel(FP8Params p) {
out_fp8[row * n + col] = __nv_fp8_e4m3(v0 * o8_scale); out_fp8[row * n + col] = __nv_fp8_e4m3(v0 * o8_scale);
} }
} else { } else {
out_bf16[row * n + col] = __float2bfloat16(v0 * output_scale); auto* dst = out_bf16 + row * n + col;
if (col + 1 < n &&
(reinterpret_cast<uintptr_t>(dst) & 3) == 0) {
*reinterpret_cast<__nv_bfloat162*>(dst) =
__floats2bfloat162_rn(v0 * output_scale,
v1 * output_scale);
} else {
dst[0] = __float2bfloat16(v0 * output_scale);
if (col + 1 < n) if (col + 1 < n)
out_bf16[row * n + col + 1] = dst[1] = __float2bfloat16(v1 * output_scale);
__float2bfloat16(v1 * output_scale); }
} }
}; };
#pragma unroll #pragma unroll