perf: single-wave big cta takes a stage deeper pipeline

- dispatch 128x128 CTAs at kStages=3 when the grid fits one wave (tiles <= SM count): with no second wave to overlap the drain, latency hiding comes only from the pipeline depth
- multi-wave grids keep kStages=2 — the shorter prologue wins once retiring CTAs overlap (measured 4096^3: s2 196T vs s3 175T)
- geometry sweep across the mid band (128x64, 64x128, kK=128, 256-row CTAs) measured and rejected: all lose to the 128x128 fast loop; the remaining mid-band gap concentrates in the 1.0-1.4 wave dip (1280^3-class shapes, ~105T vs cuBLAS 186T), which is a scheduling problem (split-K), not a geometry one

Benchmark: L20 (sm_89), CUDA-graph e2e: 1024^3 106.3->107.9T, 1152^3 133.3->137.1T, others unchanged (512^3 54.2T, 2048^3 177.9T, 4096^3 192.9T, 8192^3 198.2T). C++ four-layout suite and 114 targeted pytests pass.
This commit is contained in:
2026-08-26 15:02:00 +08:00
parent 1d57588d27
commit d4534be8ca
+14
View File
@@ -830,6 +830,20 @@ void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
// single generic body — no dead second loop in their I-cache.
constexpr bool kBigFast = !std::is_same_v<LayoutA, ColMajor> &&
!std::is_same_v<LayoutB, RowMajor>;
// Single-wave grids (tiles <= SM count) take one stage deeper: with no
// second wave to overlap the drain, latency hiding comes only from the
// pipeline (measured, L20, in-wave band: 1024^3 +1.5%, 1152^3 +1.2%,
// K=4096 rects +2%); multi-wave grids flip back — the shorter prologue
// wins once retiring CTAs overlap (4096^3: s2 196T vs s3 175T).
if (tiles_128 <= (int64_t)device_sm_count()) {
using TraitsS3 = Fp8GemmTraits<Fmt, 128, 128, kK, Stages + 1>;
launch_with_smem<
fp8_gemm_kernel<TraitsS3, LayoutA, LayoutB, GroupRaster, false,
false, kBigFast>>(
Fp8GemmSmem<TraitsS3, LayoutA, LayoutB, false>::kBytes, grid,
dim3(TraitsS3::kCtaThreads), stream, p);
return;
}
launch_with_smem<
fp8_gemm_kernel<Traits, LayoutA, LayoutB, GroupRaster, false, false,
kBigFast>>(