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

refactor: fp8 gemm policy layering with swap-NN and narrow-N ctas

Browse Source 
Kernel restructured CUTLASS-style: Fp8GemmPolicy as the kernel's single template parameter (traits + operand layouts + scheduling knobs), the body split into Fp8GemmTileScheduler / Fp8CollectiveMainloop / Fp8CollectiveEpilogue collectives, and the entry split into canonicalize_gemm -> plan_gemm -> launch_plan behind fp8::gemm.

- NN (dual-N-contiguous) problems run as their transpose: the swap in canonicalize_gemm plus an out-transposed epilogue removes one kernel instantiation per (format, tile config)
- new 128x64 narrow CTA (8 warps of 32x32) serves the sub-wave band once its grid passes ~3/8 of a wave: +7..77% there (128x4096x4096 116->131T, 1024^3 131->174T, 4096x384x4096 147->242T, 8192x128x4096 131->233T); decode, the padding band and multi-wave shapes unchanged
- launch_with_smem no longer swallows cudaFuncSetAttribute failures
- fp8_test: GPU-side fp32 reference (O(m*n) compare instead of O(m*n*k) host loop), production-dispatch cases for the NN swap and the plan selection; dead transpose_layout trait removed

Device: NVIDIA RTX 6000D (sm_120, 156 SMs), CUDA 13.1, torch 2.11.0+cu130. Kernel-only bench vs CUTLASS 4.8.0 sm120 dense fp8: ahead up to 1.68x below one wave (512^3 44 vs 26T, 64x4096x4096 95 vs 62T), within ~7% in the DRAM-streaming regime (8192^3 248 vs 266T).
This commit is contained in:
ViperEkura
2026-08-28 01:21:55 +08:00
parent bbb2d95256
commit fac9d07542
5 changed files with 716 additions and 469 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/fp8/ops.cu
+5 -21
View File
@@ -70,25 +70,9 @@ void pack_gemm(FP8Params& p, const void* a, const void* b, void* output,
p.b_ld = static_cast<int>(b_ld);
}
template <FP8Format Fmt, int Variant>
void launch_variant(const FP8Params& p, cudaStream_t stream) {
using LayoutA = std::conditional_t<(Variant & 2) != 0, ColMajor, RowMajor>;
using LayoutB = std::conditional_t<(Variant & 1) != 0, ColMajor, RowMajor>;
launch_fp8_gemm<Fmt, LayoutA, LayoutB>(p, stream);
}
template <FP8Format Fmt>
void dispatch_gemm(const FP8Params& p, cudaStream_t stream, bool trans_a,
bool trans_b) {
const int variant = (static_cast<int>(trans_a) << 1) |
static_cast<int>(trans_b);
switch (variant) {
case 0: launch_variant<Fmt, 0>(p, stream); break;
case 1: launch_variant<Fmt, 1>(p, stream); break;
case 2: launch_variant<Fmt, 2>(p, stream); break;
case 3: launch_variant<Fmt, 3>(p, stream); break;
}
}
// Layout dispatch (the NN swap in canonicalize_gemm) and launch planning
// (plan_gemm/launch_plan) live in gemm.cuh behind fp8::gemm — pure CUDA,
// shared with the C test suite.
// Inner-layout resolution for one GEMM operand. The user flag names the
// math (0 = tensor's last two dims are [rows][contract], 1 = transposed);
@@ -223,9 +207,9 @@ torch::Tensor mm_fp8(torch::Tensor a, torch::Tensor b, torch::Tensor scale,
p.b_batch_stride = (batch_b == 1 && batch > 1) ? 0 : b_bstride;
p.out_batch_stride = m * n;
if (a.scalar_type() == torch::kFloat8_e4m3fn)
dispatch_gemm<FP8Format::E4M3>(p, stream.stream(), tag_a, tag_b);
gemm<FP8Format::E4M3>(p, stream.stream(), tag_a, tag_b);
else
dispatch_gemm<FP8Format::E5M2>(p, stream.stream(), tag_a, tag_b);
gemm<FP8Format::E5M2>(p, stream.stream(), tag_a, tag_b);
C10_CUDA_CHECK(cudaGetLastError());
return output;
}