perf: dispatch non-128-divisible shapes to the 64x64 cta

- 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.
This commit is contained in:
2026-08-26 19:02:50 +08:00
parent 76aa4edc9f
commit 4c82d5d84b
+14 -6
View File
@@ -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 // 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 // (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 // 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 // below 5/8 of a wave, the narrow-M/N edge, and the divisibility rule below.
// rows); inside [5/8, 1] waves the big CTA was already the measured winner inline bool prefer_small_cta(int64_t tiles_128, int64_t m, int64_t n) {
// (63-tile rect +8%). if (m <= 64 || n <= 64) return true;
inline bool prefer_small_cta(int64_t tiles_128, int64_t m) { const bool big_div = (m % 128 == 0) && (n % 128 == 0);
if (m <= 64) return true; 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; 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). // per-matrix tiles (see prefer_small_cta).
const int64_t tiles_128 = const int64_t tiles_128 =
(int64_t)p.batch * ((p.m + 127) / 128) * ((p.n + 127) / 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); dim3 grid((p.n + 63) / 64, (p.m + 63) / 64, p.batch);
// Full-ring small CTAs — ONE __syncthreads per k-tile, cuBLAS's // Full-ring small CTAs — ONE __syncthreads per k-tile, cuBLAS's
// barrier structure (the lean ring traded a second barrier for a // barrier structure (the lean ring traded a second barrier for a