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

[yashb98]

What this does

Adds an sm121 branch to _gemm_4bit_use_custom_cuda so GB10 (DGX Spark) stops falling through to the sm89 tiers, and updates the two comments that say sm121 is unvalidated.

Right now sm121 is explicitly routed to the sm89 fallback with a comment saying "uses sm89 fallback below until validated". This is that validation, for the part of the space I could measure reliably.

Why sm121 needs its own tier

GB10 is unified LPDDR5X, so it has far less bandwidth than the GDDR6X parts sm89 was calibrated on. Dequant plus F.linear stays expensive much further up the M range, which means the custom kernel keeps winning past the sm89 caps.

The change is one branch:

if is_sm121:
    if n_blocks >= num_sms:
        return M <= 256

Everything below one wave keeps using the shared tiers, because there the crossover is strongly K-dependent and I did not have clean enough data to justify a separate rule.

Safety

This cannot change behaviour on any other architecture, and that is checkable rather than argued. I enumerated the dispatch decisions of the stock and patched functions over a cross product of 10 architecture profiles (sm75/sm80/sm86/sm89/sm90/sm100/sm120/sm121 plus an unrecognized arch), 3 dtypes, and 23 x 11 x 8 M x N x K shapes:

60,720 decisions enumerated
    216 changed
    216 False -> True
      0 True  -> False

All 216 changes land on the sm121 profile, at M between 48 and 256. No other arch is touched anywhere in the grid. Since it is exhaustive over the enumerated space rather than sampled, it is a proof of no-regression for those shapes, not evidence of one.

Tests

Run against the released 0.50.0 wheel with only this function transplanted, since the source tree ships no compiled library. The function source is byte-identical between the wheel and main (verified by direct comparison), so the substitution is faithful.

suite	stock	patched
test_functional.py -k 4bit	1977 passed, 864 skipped	identical
test_linear4bit.py + test_autograd.py	2156 passed, 1 failed	identical

The one failure is test_fsdp_state_dict_save_4bit, and it fails on stock too, so it is pre-existing on this machine and not caused by this change.

Honest limits

Three things I want to be upfront about, because they bound what this PR claims.

The cap is conservative, not exact. I originally thought M=256 was the crossover. It is not: across five measurement sessions, M=384 still favours the custom kernel on 10 of 14 shapes (cell median about 1.2x) and M=512 on 7 to 8 of 14. Within the at-or-above-one-wave region the only M=384 loser is the smallest tier-A shape, (3072, 3072), at 0.86x to 0.94x across five sessions. So M=256 leaves real headroom, and I chose it deliberately rather than pushing to the measured edge on one shape.

Only the at-or-above-one-wave region is calibrated. I had a wider three-tier version of this patch. I threw it away: enumerating it the same way showed it changing decisions at N=2560 in the tall-K shapes, a region where my own measurements do not support a clean rule (the custom kernel swings from about 3x ahead at low M to 0.75x behind at M=384, across four seeds). Two of its three tiers rested on a single measured shape each. This PR is what survived that.

No real model forward was run. All timings are synthetic stacks of real projection shapes, not an actual HuggingFace model, and the machine had an unrelated job resident throughout, so absolute microseconds are upper bounds. The dispatch decisions above are exact; the performance motivation behind them is measured but not production-validated.

Environment: GB10 (DGX Spark), sm_121, aarch64, driver 580.142, torch 2.13.0+cu130, CUDA 13.0.

Happy to add a regression test for the dispatch table, or to widen the cap if someone with a second GB10 can confirm the M=384 numbers.

[matthewdouglas]

Some of this comment seems a bit unnecessary. We don't need to really talk about how it /used to/ fall through to sm89 or how it compares.

[matthewdouglas]

Thanks for the PR. In general it looks good. I would expect the custom kernels to do well considering the lower memory bandwidth on GB10 so the result makes sense in general.

I'm wondering if you can share your benchmark figures and exact methodology?

The other part I've wondered here, is there's an additional layer of dispatch decisions that I was not able to test on the real hardware. Besides deciding to use one of our fused kernels, there's also decisions on the C++ side on which kernel to use: SIMT, or MMA with different configurations. It's possible we're not selecting the best and leaving some performance on the table since this also hasn't been calibrated for GB10 and there are arch-specific decisions in there. But this doesn't really block merging this if we're confident there's no significant performance regressions being introduced with it.

[github-actions]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[yashb98]

Thanks for the quick look. Happy to share everything, and I also went ahead and measured the SIMT/MMA layer you mentioned, since I have the hardware. Methodology and figures first, then that data.

Methodology

Harness: for each (shape, M, dtype) cell I time the two paths directly, bypassing dispatch entirely, so the measurement is the kernels rather than the heuristic:

• custom: _gemm_4bit_kernel_impl
• fallback: _dequant_linear_fallback (dequant + F.linear)

The two are timed interleaved in the same loop (custom, fallback, custom, fallback, ...) so clock drift and thermal state hit both sides equally. 8 warmup iterations, 25 timed iterations per side, per cell. Reported as median with min/max.

Setup: nf4, blocksize 64, bf16 and fp16 activations, 14 (N, K) shapes taken from real projection dimensions (Llama2-7B, Llama3-8B, 13B-class, Qwen2.5-1.5B), M swept 1 to 1536. bitsandbytes 0.50.0, torch 2.13.0+cu130, driver 580.142, GB10 (48 SMs). Five independent sessions over two days (four full sweeps plus one confirmation run). Numerical agreement between the two paths: max abs diff 0.0625 across all cells, i.e. bf16 rounding, with relative mean diff ~0.

Figures

The region this PR changes (>= 1 wave, M in [48, 256], both dtypes, 120 cells):

• all 120 cells favour the custom kernel
• median speedup 2.52x, min 1.10x, max 5.01x

Per-shape medians over that region (bf16):

shape	N	K	median speedup
3B-class square	3072	3072	1.79x
Llama3-8B o_proj	4096	4096	2.04x
Llama3-8B down_proj	4096	14336	2.32x
13B down_proj	5120	13824	2.54x
fused qkv	6144	4096	2.96x
large square	8192	8192	2.48x
Qwen2.5-1.5B gate/up	8960	1536	2.24x
Llama2-7B gate/up	11008	4096	3.21x
Llama3-8B gate/up	14336	4096	2.93x
Llama3-8B lm_head	128256	4096	3.57x

Why the cap stayed at 256: at M=384 the custom kernel still wins 18 of 20 wave-eligible cells (median 1.34x), and at M=512 only 15 of 20 (median 1.07x). The consistent loser is the smallest tier shape, 3072x3072 at M=384: 0.86x, 0.93x, 0.93x, 0.94x, 0.93x across the five sessions. So there is real headroom above 256, but it is shape-dependent, and I would rather a second GB10 confirm that before widening

The SIMT / MMA layer

I measured this too, same rig. Built libbitsandbytes_cuda130.so from main (a2b90e6) natively for sm_121, three times: unmodified, forced SIMT (use_simt = true), and forced MMA (use_simt = false on K%64==0 cells). Timed the production fused path (_gemm_4bit_kernel_impl -> cgemm_4bit_*) on 10 (N,K) shapes x M in {4..512} x {bf16, fp16} = 260 cells per variant. nf4, blocksize 64, 8 warmup + 25 reps, medians, identical seeds and buffers across variants. Results only, no proposal.

1. On this arch the dispatcher picks the slower kernel in 32 of 260 cells (12%).

2. The largest losses are at M=4 on large weights, where stock takes MMA but SIMT wins:

shape	dtype	SIMT	MMA (stock pick)	MMA is slower by
8192x8192	bf16	193us	363us	1.88x
8192x8192	fp16	191us	365us	1.91x
4096x14336	bf16	190us	304us	1.60x
14336x4096	bf16	169us	255us	1.51x
5120x13824	bf16	215us	306us	1.42x
4096x4096	bf16	60us	84us	1.41x
11008x4096	bf16	135us	177us	1.31x
8960x1536	bf16	54us	69us	1.27x

(fp16 rows sit within a few percent of bf16 throughout.)

The root cause looks like exactly what you suspected: at M=4 the SIMT route needs highbw_gddr or wide_n_simt, and both enumerate sm86/sm89/sm120 but not sm121, so a GB10 falls through to MMA even where SIMT is up to ~1.9x faster. (Worst cell re-confirmed on a fresh seed with 50 reps: 189us vs 365us.)

3. The reverse direction exists but is smaller: on 2048x2048 at M=6-8 the undersubscribed rule sends GB10 to SIMT while MMA is 1.2-1.6x faster.

4. Everywhere else, 228 of 260 cells (88%), the current choice is the measured winner, including the whole M>=48 region this PR touches. And to be explicit about the regression question: all 32 mispicked cells sit at M<=16, where this PR changes no routing at all, so they behave identically on main with or without this change.

Raw JSONs (all 780 cells, min/med/max per cell) and the harness are available if useful; happy to rerun any cell or grid shape under different conditions, and happy to look at what calibrated sm121 rules would look like as a follow-up if there is appetite.