From d4534be8ca421b010310d8087938837673cf4454 Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Wed, 26 Aug 2026 15:00:02 +0800
Subject: [PATCH] perf: single-wave big cta takes a stage deeper pipeline
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- 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.
---
 csrc/kernels/fp8/gemm.cuh | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/csrc/kernels/fp8/gemm.cuh b/csrc/kernels/fp8/gemm.cuh
index eb4cfdd..e6cdfab 100644
--- a/csrc/kernels/fp8/gemm.cuh
+++ b/csrc/kernels/fp8/gemm.cuh
@@ -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>>(