From 35e54e62bca51012386564d16c460e6a5a1914c4 Mon Sep 17 00:00:00 2001 From: yashb98 Date: Fri, 31 Jul 2026 10:41:02 +0100 Subject: [PATCH 1/2] Add sm121 (GB10) tier to the 4-bit GEMM dispatch heuristic --- bitsandbytes/backends/cuda/ops.py | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py index 1f3ce82e0..dd7fb92aa 100644 --- a/bitsandbytes/backends/cuda/ops.py +++ b/bitsandbytes/backends/cuda/ops.py @@ -592,6 +592,10 @@ def _gemm_4bit_use_custom_cuda(device_index, dtype, M, N, K): sm90 (H100/H200, HBM3/HBM3e): dequant+linear is much faster; thresholds are tight. sm100 (B200/B300, HBM3e): exits early at top of function. sm120 (RTX 5000, GDDR7): dedicated block; medium-N tiers differ from sm89. + sm121 (GB10 DGX Spark, LPDDR5X): dedicated block at >=1 wave only; the low + memory bandwidth keeps the custom kernel ahead + through M=256 there. Sub-wave shapes use the + sm89 tiers below. """ if M <= _GEMM_4BIT_CUSTOM_FLOOR_M: return True @@ -623,6 +627,7 @@ def _gemm_4bit_use_custom_cuda(device_index, dtype, M, N, K): is_sm86 = major == 8 and minor == 6 is_sm90 = major == 9 is_sm120 = major == 12 and minor == 0 + is_sm121 = major == 12 and minor == 1 is_hbm = is_sm80 or is_sm90 # sm100 already returned above tall_k_2xn = K > N * 2 @@ -753,10 +758,17 @@ def _gemm_4bit_use_custom_cuda(device_index, dtype, M, N, K): return M <= 16 return False + if is_sm121: + # GB10 (DGX Spark): unified LPDDR5X, far less bandwidth than the sm89 GDDR6X + # parts this used to fall through to, so dequant+F.linear stays expensive + # much further up the M range. Calibrated on GB10 at >=1 wave only; below + # one wave the crossover is strongly K-dependent, so those shapes keep using + # the shared tiers below (which already branch on tall-K). + if n_blocks >= num_sms: + return M <= 256 + if is_sm120: # GDDR7 (~1-1.8 TB/s). Medium-N threshold tiers differ from sm89. - # sm121 (DGX Spark) has a different bandwidth/SM profile; uses sm89 - # fallback below until validated. if n_blocks >= num_sms * 3: return M <= 256 if n_blocks >= num_sms * 2: @@ -773,7 +785,7 @@ def _gemm_4bit_use_custom_cuda(device_index, dtype, M, N, K): return M <= 16 return M <= 8 - # Fallback: sm89 (4090, L40S, L4), sm121 (DGX Spark), unrecognized arches. + # Fallback: sm89 (4090, L40S, L4), sm121 below one wave, unrecognized arches. # GDDR bandwidth makes dequant relatively expensive so custom wins at higher M. if n_blocks >= num_sms * 3: return M <= 256 From b09a98c09c6861f0e6985e3c4356c21ae4af3c77 Mon Sep 17 00:00:00 2001 From: yashb98 Date: Wed, 5 Aug 2026 08:40:09 +0100 Subject: [PATCH 2/2] Trim the sm121 dispatch comment --- bitsandbytes/backends/cuda/ops.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/bitsandbytes/backends/cuda/ops.py b/bitsandbytes/backends/cuda/ops.py index dd7fb92aa..a0d9ffe83 100644 --- a/bitsandbytes/backends/cuda/ops.py +++ b/bitsandbytes/backends/cuda/ops.py @@ -759,11 +759,9 @@ def _gemm_4bit_use_custom_cuda(device_index, dtype, M, N, K): return False if is_sm121: - # GB10 (DGX Spark): unified LPDDR5X, far less bandwidth than the sm89 GDDR6X - # parts this used to fall through to, so dequant+F.linear stays expensive - # much further up the M range. Calibrated on GB10 at >=1 wave only; below - # one wave the crossover is strongly K-dependent, so those shapes keep using - # the shared tiers below (which already branch on tall-K). + # GB10 (DGX Spark): unified LPDDR5X. Calibrated at >=1 wave only; below one + # wave the crossover is strongly K-dependent, so those shapes use the shared + # tiers below (which already branch on tall-K). if n_blocks >= num_sms: return M <= 256