Skip to content

Add sm121 (GB10) tier to the 4-bit GEMM dispatch heuristic #2030

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yashb98 wants to merge 2 commits into
base: main
Choose a base branch
from
+13 −3
Open

Add sm121 (GB10) tier to the 4-bit GEMM dispatch heuristic #2030

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
35e54e6
Add sm121 (GB10) tier to the 4-bit GEMM dispatch heuristic
Jul 31, 2026
b09a98c
Trim the sm121 dispatch comment
Aug 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions bitsandbytes/backends/cuda/ops.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -753,10 +758,15 @@ 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. 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

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:
Expand All @@ -773,7 +783,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
Expand Down