From 5fcc005b52530116ba0b122ab866a698c66fea0c Mon Sep 17 00:00:00 2001 From: yashb98 Date: Fri, 14 Aug 2026 09:41:48 +0100 Subject: [PATCH] Add sm121 (GB10) to the wide_n_simt 4-bit GEMM dispatch rule wide_n_simt already routes sm120 with >=48 SMs to the SIMT kernel at M<=6 with wide N. GB10 (sm121, 48 SMs) clears the same SM threshold but was excluded by the minor-version check, so it fell through to MMA. Measured on GB10: all 8 rerouted cells are faster with SIMT, 1.03x to 1.86x. Enumeration over 78,936 and 177,540 dispatch decisions confirms the change is confined to cc_maj==12 && cc_min==1 at M<=6, False->True only, with no other architecture affected. --- csrc/gemm_4bit.cu | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/csrc/gemm_4bit.cu b/csrc/gemm_4bit.cu index beb814829..ff2f011a2 100644 --- a/csrc/gemm_4bit.cu +++ b/csrc/gemm_4bit.cu @@ -87,11 +87,12 @@ static void gemm_4bit( const bool undersubscribed = (M <= 8 && mma_blocks * 3 <= num_sms * 2) || (hbm_arch && M == 4 && mma_blocks <= num_sms); - // sm89 (>=60 SMs) and sm120 (>=48 SMs): at M<=6 with wide N, SIMT saturates + // sm89 (>=60 SMs) and sm120/sm121 (>=48 SMs): at M<=6 with wide N, SIMT saturates // bandwidth more efficiently than blocked MMA. const bool wide_n_simt = M <= 6 && mma_blocks >= num_sms && - ((cc_maj == 8 && cc_min == 9 && num_sms >= 60) || (cc_maj == 12 && cc_min == 0 && num_sms >= 48)); + ((cc_maj == 8 && cc_min == 9 && num_sms >= 60) || (cc_maj == 12 && cc_min == 0 && num_sms >= 48) || + (cc_maj == 12 && cc_min == 1 && num_sms >= 48)); // GDDR tall-K (K>N): K-loop too long relative to output tile at small M. const bool tall_k_simt = gddr_arch && K > N && M <= 17 && mma_blocks * 3 < num_sms;