fix(rocm): avoid SIMT workgroup-count resonance (#2012)

* fix(rocm): avoid SIMT workgroup-count resonance

Add one bounds-checked workgroup when a large HIP SIMT grid is an exact multiple of 256, avoiding scheduler underutilization observed on RDNA3 and RDNA4.\n\nCo-Authored-By: Claude <noreply@anthropic.com>

* refactor(rocm): add explicit RDNA and CDNA macros

Co-Authored-By: Claude <noreply@anthropic.com>

* refactor(rocm): detect RDNA2 architectures

Co-Authored-By: Claude <noreply@anthropic.com>

* fix(rocm): exclude RDNA2 from dot2 path

Co-Authored-By: Claude <noreply@anthropic.com>

---------

Co-authored-by: Claude <noreply@anthropic.com>

Author: sstamenk

csrc/common.cuh

@@ -4,10 +4,26 @@
 
 #include "compat.cuh"
 
+// AMD GPU architecture detection. Use explicit device macros rather than the
+// broad __GFX9__/__GFX11__/__GFX12__ macros, which span distinct architectures.
+#define IS_RDNA2                                                                                                       \
+    (defined(__gfx1030__) || defined(__gfx1031__) || defined(__gfx1032__) || defined(__gfx1033__) ||                   \
+     defined(__gfx1034__) || defined(__gfx1035__) || defined(__gfx1036__))
+#define IS_RDNA3 (defined(__gfx1100__) || defined(__gfx1101__) || defined(__gfx1102__) || defined(__gfx1103__))
+#define IS_RDNA3_5 (defined(__gfx1150__) || defined(__gfx1151__) || defined(__gfx1152__) || defined(__gfx1153__))
+#define IS_RDNA4 (defined(__gfx1200__) || defined(__gfx1201__))
+#define IS_RDNA (IS_RDNA2 || IS_RDNA3 || IS_RDNA3_5 || IS_RDNA4)
+
+#define IS_CDNA1 (defined(__gfx908__))
+#define IS_CDNA2 (defined(__gfx90a__))
+#define IS_CDNA3 (defined(__gfx942__))
+#define IS_CDNA4 (defined(__gfx950__))
+#define IS_CDNA (IS_CDNA1 || IS_CDNA2 || IS_CDNA3 || IS_CDNA4)
+
 // Warp size
 
 #if BNB_HIP
-#if defined(__GFX9__)
+#if IS_CDNA
 #define BNB_WARP_SIZE 64 // CDNA
 #else
 #define BNB_WARP_SIZE 32 // RDNA

csrc/gemm_4bit_simt.cu

@@ -4,10 +4,11 @@
 #include <cstdint>
 #include <type_traits>
 
+#include "common.cuh"
 #include "gemm_4bit_common.cuh"
 #include "gemm_4bit_simt.cuh"
 
-#if defined(__GFX9__)
+#if IS_CDNA
 // gfx9/CDNA tuning:
 // - use fmaf for fp32 accumulation
 // - accumulate fp16/bf16 pairs in fp32 via fused multiply-add
@@ -26,9 +27,9 @@
 #define BNB_SIMT_FP16_PACKED_ACCUM 0
 #endif
 
-// RDNA3/RDNA4 tuning:
+// RDNA3/RDNA3.5/RDNA4 tuning:
 // - use native packed bf16/fp16 dot2 instructions with fp32 accumulation
-#if defined(__GFX11__) || defined(__GFX12__)
+#if IS_RDNA3 || IS_RDNA3_5 || IS_RDNA4
 #define BNB_HIP_BF16_VDOT2 1
 #define BNB_HIP_FP16_DOT2 1
 #else
@@ -490,7 +491,14 @@ void launch_gemm_4bit_simt(
     bnb_stream_t stream           // CUDA/HIP stream
     // clang-format on
 ) {
-    const int n_blocks = (N + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK;
+    int n_blocks = (N + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK;
+#if BNB_HIP
+    // Large HIP SIMT dispatches can underutilize the GPU when the workgroup
+    // count is an exact multiple of 256. One additional bounds-checked block
+    // breaks the scheduler resonance without changing useful work.
+    if (n_blocks >= 512 && n_blocks % 256 == 0)
+        n_blocks++;
+#endif
 
     // M=1..8: M_BLOCK == M, so the inner m-loop fully unrolls at compile time.
     // M>8: M_BLOCK=8, ceil(M/8) grid rows.