From 4c82d5d84b604772a76264c41e2194c71050b64c Mon Sep 17 00:00:00 2001
From: ViperEkura <3081035982@qq.com>
Date: Wed, 26 Aug 2026 19:02:50 +0800
Subject: [PATCH] perf: dispatch non-128-divisible shapes to the 64x64 cta
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- a 128x128 CTA that is not exactly tiled (m or n not a multiple of 128) runs its edge tiles on the predicated generic path, and with a single in-flight wave the runtime is the slowest CTA — the edge tiles drag the whole shape down, so prefer_small_cta now also takes n and returns true when 64 divides both dims but 128 does not (measured sweep: 1088^3 big-CTA 76T vs 64x64 93T)
- the 64x64 grid tiles exactly on such shapes and overlaps waves (CTA count jumps 64->100->144 in the 1024-1536 band; the non-divisible points like 1088/1216 sit in deep sawtooth valleys that the divisibility rule lifts to the flat ~100-111T plateau)
- double-non-divisible shapes (e.g. 1000^3) stay on the big CTA: measured 64x64 36.7T vs 128x128 40.0T — both grids carry edge tiles there and the big CTA's efficiency wins
- m <= 64 or n <= 64 also takes the small CTA (a 128-wide CTA wastes more than half its columns on narrow N)

Benchmark: L20, e2e CUDA graph (GPU 7, same GPU as all comparisons): 960^3 60.2->100.1T, 1088^3 77.6->101.6T, 1216^3 69.7->103.5T, 1344^3 84.9->111.1T; 128-divisible shapes unchanged within noise (1024^3 118.4, 1152^3 150.6, 1280^3 106.8, 1536^3 153.7, 2048^3 192.2). 596 pytests pass.
---
 csrc/kernels/fp8/gemm.cuh | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/csrc/kernels/fp8/gemm.cuh b/csrc/kernels/fp8/gemm.cuh
index d0c7bdb..2026ee9 100644
--- a/csrc/kernels/fp8/gemm.cuh
+++ b/csrc/kernels/fp8/gemm.cuh
@@ -942,11 +942,19 @@ void launch_with_smem(int smem_bytes, dim3 grid, dim3 block,
 // that stall, and the trade flipped across the whole measured band: 1280^3
 // (1.09 waves) big 163T vs small 100T, 4096x512x4096 (1.4 waves) big 139T
 // vs small 113T. The small CTA now only serves the genuinely sub-wave band
-// below 5/8 of a wave (and m <= 64, where a 128-row CTA wastes half its
-// rows); inside [5/8, 1] waves the big CTA was already the measured winner
-// (63-tile rect +8%).
-inline bool prefer_small_cta(int64_t tiles_128, int64_t m) {
-    if (m <= 64) return true;
+// below 5/8 of a wave, the narrow-M/N edge, and the divisibility rule below.
+inline bool prefer_small_cta(int64_t tiles_128, int64_t m, int64_t n) {
+    if (m <= 64 || n <= 64) return true;
+    const bool big_div = (m % 128 == 0) && (n % 128 == 0);
+    const bool small_div = (m % 64 == 0) && (n % 64 == 0);
+    // Divisibility: a 128x128 CTA that is NOT exactly tiled (m or n not a
+    // multiple of 128) runs its edge tiles on the predicated generic path,
+    // and with a single in-flight wave the runtime is the slowest CTA — the
+    // edge tiles drag the whole shape down (1088^3: 76T vs 93T with the
+    // 64x64 CTA, whose grid tiles exactly and overlaps waves; measured
+    // sweep, perf 5.1). When 64 divides both dims, the 64x64 small CTA
+    // wins the non-128-divisible band by 23..67%.
+    if (!big_div && small_div) return true;
     return tiles_128 < device_sm_count() * 5 / 8;
 }
 
@@ -962,7 +970,7 @@ void launch_fp8_gemm(const FP8Params& p, cudaStream_t stream) {
     // per-matrix tiles (see prefer_small_cta).
     const int64_t tiles_128 =
         (int64_t)p.batch * ((p.m + 127) / 128) * ((p.n + 127) / 128);
-    if (prefer_small_cta(tiles_128, p.m)) {
+    if (prefer_small_cta(tiles_128, p.m, p.n)) {
         dim3 grid((p.n + 63) / 64, (p.m + 63) / 64, p.batch);
         // Full-ring small CTAs — ONE __syncthreads per k-tile, cuBLAS's
         // barrier structure (the lean ring traded a second barrier for a