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:
@@ -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>>(
|
||||
|
||||
Reference in New Issue
Block a user