ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: split fp8 gemm cta plan by operand layout

Browse Source 
- pass the crosswise operand count from gemm into plan_gemm so congruous and crosswise problems stop sharing one threshold ladder
- congruous grids past one big-cta wave pick big vs narrow by the wave cost ceil(tiles/sm) * T_tile with T_narrow ~= 0.53 * T_big, reproducing every measured crossover
- crosswise problems run the small 64x64 s3 cta up to ~1.5 waves of 128x128 tiles; the narrow cta never wins there (loses to small below the band, to big above it)
- keep the sub-wave congruous ladder and the padding rules unchanged

Benchmark: NVIDIA L20 (92 SMs, sm_89), CUDA 12.8, fp8 e4m3 -> bf16, interleaved A/B against the previous ladder
- NT Mx4096x4096: M=384 114.5 -> 134.5 TF (+17.4%), M=512 152.5 -> 171.7 (+12.6%), M=768 153.5 -> 165.4 (+7.7%); all other NT shapes unchanged
- NN/TN M=64..512 +3.2..+17.4%, 1024^3 +13.7%/+12.8%; M>=640 and 2048^3+ unchanged
- TT 1024^3 +25.5%, M=256 +25.4%; TT M=512 -4.6% at the 1.5-wave boundary that favors TN/NN
This commit is contained in:
ViperEkura
2026-08-28 12:44:34 +08:00
parent c4f7f82725
commit 04a8e2517a
1 changed files with 60 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/kernels/fp8/gemm.cuh
+60 -17
View File
@@ -1088,31 +1088,71 @@ struct Fp8GemmPlan {
bool small_s3; // kSmall64 only: cp.async pipeline depth (2 vs 3 stages)
};
inline Fp8GemmPlan plan_gemm(const FP8Params& p) {
// crosswise_ops: how many operands take the direct crosswise load (A
// ColMajor storage / B RowMajor storage, see Fp8GemmSmem). 0 = the
// dual-congruous NT problem, 1 = TN and the NN swap, 2 = TT. The layout
// shifts the crossovers: the small CTA hides the crosswise LDG+PRMT
// latency far better (more resident CTAs, deeper pipeline), while the big
// CTA's operand reuse mostly buys back load bandwidth the crosswise path
// does not traffic in.
inline Fp8GemmPlan plan_gemm(const FP8Params& p, int crosswise_ops = 0) {
const int64_t sm = device_sm_count();
const int64_t tiles_128 =
(int64_t)p.batch * ((p.m + 127) / 128) * ((p.n + 127) / 128);
const auto small = [&](bool s3) {
return Fp8GemmPlan{Fp8GemmPlan::Cta::kSmall64, s3};
};
// Padding rules first: predication waste beats any wave-fill effect.
if (small_cta_padding(p.m, p.n)) return small(false);
// Past the sub-wave band the 128x128 CTA wins on operand reuse (each A
// element multiplies 128 B columns in-CTA) and per-CTA pipeline depth:
// forcing the 64x64 CTA there measured 2048^3 123->171 TF on L20 and
// 164 vs 308T on sm_120.
if (tiles_128 >= sm * 5 / 8)
const auto big = [] {
return Fp8GemmPlan{Fp8GemmPlan::Cta::kBig128, false};
// Sub-wave band: the 128x64 narrow CTA (8 warps of 32x32) fills the
// wave with N-tiles at full warp depth — measured sm_120, it beats the
// small CTA by +7..77% across the band once the narrow grid passes ~3/8
// of a wave (128x4096 132 vs 123T, 1024^3 174 vs 131T, 4096x384 242 vs
// 147T, 8192x128 233 vs 131T); below that fill the plain 64x64 CTA's
// extra parallelism wins (2048x128: small 99 vs 87T).
};
const auto narrow = [] {
return Fp8GemmPlan{Fp8GemmPlan::Cta::kNarrow128x64, false};
};
// Padding rules first: predication waste beats any wave-fill effect.
if (small_cta_padding(p.m, p.n)) return small(crosswise_ops > 0);
if (crosswise_ops > 0) {
// Crosswise ladder (L20 measured, N=K=4096 band + squares): the
// narrow CTA never wins — below the wave band the small CTA beats
// it (M=128: 82.8 vs 72.3T), above it the big CTA does. The small
// s3 CTA holds ~3/4 of the big CTA's per-SM throughput but tiles 4x
// finer, so it owns the whole sub-wave band and past it: +15% at
// M=256 (129.7 vs 113.1), +17% at M=384, +13% at 1024^3 (107.2 vs
// 94.8). The big CTA takes over once its grid fills ~1.5 waves:
// +18% at 1536^3 (134.4 vs 114.0), +20% at M=640 (161.2 vs 134.4);
// the 1.5-wave boundary itself is within noise either way (M=512:
// 133.8 vs 129.7 small; M=768: 135.8 vs 132.8 small).
if (tiles_128 >= sm * 3 / 2) return big();
return small(true);
}
if (tiles_128 >= sm) {
// Wave band: pick by the wave-quantization cost ceil(tiles/sm) *
// T_tile. The narrow tile carries half the big tile's MMA work at
// ~94% of its per-SM efficiency, i.e. T_narrow ~= 0.53 * T_big
// (L20 measured; integer-scaled by 100 below). This formula
// reproduces every measured crossover: narrow wins the poor-fill
// big grids (M=384: 134.3 vs 114.4, M=512: 171.7 vs 152.5, M=768:
// 164.6 vs 153.7), the big CTA wins the well-filled ones (M=640:
// 187.7 vs 167.3, M=1024: 202.5 vs 178.8, 4096^3: 210.7 vs 192.5).
const int64_t tiles_narrow =
(int64_t)p.batch * ((p.m + 127) / 128) * ((p.n + 63) / 64);
const auto waves = [sm](int64_t tiles) { return (tiles + sm - 1) / sm; };
if (waves(tiles_narrow) * 53 < waves(tiles_128) * 100) return narrow();
return big();
}
// Sub-wave band: the 128x64 narrow CTA fills the wave with N-tiles at
// full warp depth — measured sm_120, it beats the small CTA by +7..77%
// across the band once the narrow grid passes ~3/8 of a wave (128x4096
// 132 vs 123T, 1024^3 174 vs 131T, 4096x384 242 vs 147T, 8192x128 233
// vs 131T); below that fill the plain 64x64 CTA's extra parallelism
// wins (2048x128: small 99 vs 87T). Past ~5/8 of a wave of 128x128
// tiles the big CTA's operand reuse wins instead (each A element
// multiplies 128 B columns in-CTA; L20 M=256: 143.7 vs 135.0 narrow,
// 1024^3: 115.7 vs 113.3; forcing the 64x64 CTA there measured 2048^3
// 123->171 TF on L20 and 164 vs 308T on sm_120).
if (tiles_128 >= sm * 5 / 8) return big();
const int64_t tiles_narrow =
(int64_t)p.batch * ((p.m + 127) / 128) * ((p.n + 63) / 64);
if (tiles_narrow >= sm * 3 / 8)
return Fp8GemmPlan{Fp8GemmPlan::Cta::kNarrow128x64, false};
if (tiles_narrow >= sm * 3 / 8) return narrow();
// Full-ring small CTAs — ONE __syncthreads per k-tile, cuBLAS's barrier
// structure. Two depths by grid shape: the 24KB 3-slot s2 variant keeps
// 4 CTAs/SM while the whole grid stays resident (<= one 3-CTA wave);
@@ -1214,7 +1254,10 @@ inline void canonicalize_gemm(FP8Params& p, bool& trans_a, bool& trans_b) {
template <FP8Format Fmt>
void gemm(FP8Params p, cudaStream_t stream, bool trans_a, bool trans_b) {
canonicalize_gemm(p, trans_a, trans_b);
const Fp8GemmPlan plan = plan_gemm(p);
// Crosswise operand count for the plan: transposed-A storage (ColMajor)
// and plain-B storage (RowMajor) both take the direct crosswise load.
const int crosswise = (trans_a ? 1 : 0) + (trans_b ? 0 : 1);
const Fp8GemmPlan plan = plan_gemm(p, crosswise);
if (trans_a && trans_b)
launch_plan<Fmt, ColMajor, ColMajor, 8>(p, plan, stream);
else if (trans_b)