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]

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.

[yashb98]

Sorry, I resolved that thread before actually making the change. Trimmed in b09a98c: the sm121 block now just states the calibration scope, with the history and the cross arch comparison dropped.

The SIMT/MMA follow up offer from my earlier comment still stands whenever you want it.

The push put the workflow runs back into action_required, so the checks are held. Could you approve the run when you get a chance?

[matthewdouglas]

Thanks for the details and the data. I think this is reasonable to do as a good first calibration step as it seems like a strict improvement over what was shipped in v0.50.0.

Within the MMA kernel, there's additionally a tile selection mechanism, which can choose from about ~10 different tiling configurations. It's possible there's room left to improve for sm121 within that too, not just from adjusting the MMA-vs-SIMT decision.

Would be happy to look at follow-ups here. For one of those a simple change to the wide_n_simt rule as suggested seems reasonable.

[yashb98]

Thanks for merging it, and glad it reads as a strict improvement over v0.50.0.

wide_n_simt for sm121

Agreed this is the concrete next step. It matches what the SIMT/MMA measurement from my last comment showed: 32 of 260 cells mispicked, all at M<=16, and the root cause note there already flagged wide_n_simt (along with highbw_gddr) as one of the clauses that enumerate sm86/sm89/sm120 but stop short of sm121. Looking at the code now, the sm120 branch in wide_n_simt gates on cc_maj == 12 && cc_min == 0 && num_sms >= 48, and GB10 already clears the SM count on its own, so the change on the table is adding a cc_min == 1 case alongside it, not a restructure.

I don't want to hand you a diff off just that though. The wider three-tier version of this PR got thrown out for skipping exactly that step, so before I propose anything here I'll enumerate the decisions the sm121 addition changes across archs, dtypes, and shapes, confirm nothing outside wide_n_simt moves, and benchmark the flipped cells specifically since the 260 cells I already have were measured against the stock heuristic, not a modified one. I'll bring it as its own follow-up PR once that's done.

Tile selection in the MMA path

New to me, hadn't looked at that layer before your comment. My SIMT/MMA measurement only covers the family choice, not the tiling once MMA is picked, so I don't have anything to say about it yet. I'll go find where the roughly 10 configs get chosen and treat it as a separate follow-up, same enumerate-and-benchmark approach, rather than bundling it with wide_n_simt. If you already have a pointer to where that selection lives, it would save me some searching, but I can track it down either way.

[yashb98]

I went and measured the tile selection layer. Short version: of the three shape clusters GB10
actually lands in, one is already optimal on every cell I tested, one has a consistent gap, and
one produced a candidate rule that looked good on the first pass and did not survive rechecking.
The nulls are most of the result, so they are first.

Where it lives

select_tile() in csrc/gemm_4bit_sm80.cu, dispatching through the LAUNCH_SM80 table. Of its
11 live entries, 10 are legal on GB10: 32x256-128 traps under __CUDA_ARCH__ == 1210 and its
smem exceeds the 99 KB opt-in cap anyway, and select_tile() only emits it under hbm_arch,
which GB10 never is.

GB10 gets no special casing anywhere in the selector. The only cc_major == 12 branch is
high_sm_sm120, gated at num_sms >= 150; GB10 has 48 SMs so it never fires regardless of minor
version. Everything here runs the generic path.

I first walked the selector in pure Python over a 3,380 cell M/N/K grid to find where GB10 lands.
Three clusters cover 56% of it: 32x128-128 (M<=32, wide N), 128x128-64 (large M),
64x32-128 (small to mid M, narrow N).

All 28 measured cells, grouped by cluster:

Method

I added a BNB_FORCE_TILE env override immediately before the LAUNCH_SM80 dispatch, deliberately
after the KC=128 to KC=64 remap, so the remap cannot silently rewrite a forced tile. One build then
runs any tile on any shape. 14 shapes x 2 dtypes, every legal config each.

Both sides of every comparison below are measured through that same forced path, so the numbers
isolate the tile choice rather than the dispatch mechanism. That matters: forcing the tile the
selector would have picked anyway does not reproduce the native unforced timing exactly, and on
one cell the two differ by 24.7%.

Clocks are not locked on this box (nvidia-smi -lgc needs root), so drift is handled by
interleaving and by rechecking. A screening sweep (3 rounds, 8 warmup + 25 reps) picked out
contested cells, and every contested cell was then re-run 7 independent passes at 20 warmup + 51
reps, alternating processes with the order flipped on odd passes. Only recheck-confirmed numbers
are quoted below, and I have marked the two rows that the recheck never had to visit. The screen
alone carries run to run drift of order 10 to 25% here, enough to invent gaps that are not there,
and it did: three apparent wins evaporated on recheck, one of which I had already written up as a
result.

Two notes for anyone else poking at this. LAUNCH_SM80 is an if/else-if chain with no else, so a
tuple with no call site in the table launches nothing at all and returns, leaving the output buffer
untouched rather than erroring. Arch-invalid tuples that do have a call site are caught instead by
the __trap() guards at the top of the kernel, so the silent case is only the missing-else one. My
harness adds an explicit abort on both chains and a compile-time whitelist, so nothing silent got
through; on top of that I pre-filled every output with a sentinel and checked against an fp32
dequant reference before trusting any timing. All 308 (config, shape, dtype) triples passed, no
sentinel survived, max NRMSE 0.0017.

Result 1: 64x32-128 is already right, 10/10 cells

M=32 to 44, N=896 to 1536. Stock picks the fastest available tile on every cell, both dtypes,
confirmed gap 0.0%. Nothing to do here.

For symmetry with result 2, the screen did throw up one apparent counter-example in this cluster:
M=44, N=1024, K=4096 showed 32x64-128 ahead by 1.4%. Seven rechecks put it behind in 7 passes
out of 7. If you compute medians straight from the raw sweep you will find that cell, so I am
naming it rather than letting the recheck quietly delete it.

Result 2: 128x128-64 has no rule I can defend

64x64-64 beats stock on two of the five shapes, confirmed:

M	N x K	confirmed winner	faster by (bf16 / fp16)	7-pass recheck
128	8192 x 8192	64x64-64	20.8% / 21.2%	yes
2048	4096 x 14336	64x64-64	15.6% / 19.1%	yes
384	5120 x 13824	stock	stock optimal	yes
512	4096 x 4096	stock	stock optimal	screen only
1024	14336 x 4096	stock	stock optimal	screen only

The M=384 row is the one I want to flag explicitly. The screening sweep had it as a 14.2% win for
64x64-64 and I had it written up as such. Seven rechecks put the ratio at 0.999, with the
challenger ahead in 2 passes out of 7. It was drift, not a win.

The last two rows are marked screen only because they were never contested, so the recheck never
visited them. The M=512 cell is stable across all its rounds; the M=1024 cell is not, and its
per-round numbers straddle zero, so treat it as "no evidence stock is beaten here" rather than as
a measured null.

Side by side, a cell where the screen held up next to the one where it did not:

So: two wins, three cells where stock is already best, and no monotonic relationship with M
(the wins are at M=128 and M=2048, the nulls at M=384, M=512 and M=1024). A selector can only act
on a rule, and I do not have one here. Reporting it as a non result rather than quoting the two
wins on their own.

Result 3: 32x128-128 has a consistent gap on all 8 cells

This is the one clean positive, and every cell is recheck-confirmed:

M	N x K	confirmed winner	faster by (bf16 / fp16)
4	14336 x 4096	32x64-128	35.0% / 34.2%
6	11008 x 4096	32x64-128	15.3% / 18.8%
16	14336 x 4096	32x64-128	34.9% / 38.3%
32	8960 x 1536	32x64-64	78.6% / 79.0%

Stock picks NT=128 on all four; the measured winner is NT=64 on all four. Consistent across both
dtypes, no counter-example in the set.

On the interaction with #2039: if that lands, the M=4 and M=6 rows stop reaching select_tile()
at all, since wide_n_simt diverts them to SIMT. The M=16 and M=32 rows do not, because
wide_n_simt requires M<=6. The largest gap in the table by a wide margin (M=32, 78.6%) is one of
the retained rows, so the overlap shrinks the affected shape space without removing the effect.

On mechanism, briefly

I did look for one, and I do not have it. Wave quantization plus a tall-K term fit the screening
numbers cleanly, and then the M=384 recheck removed the cell that story depended on. I would
rather say I have no mechanism than quote one that a rerun already contradicted.

What I am asking

Is sm121 tile calibration worth pursuing from your side? I am not proposing a patch. One box, 14
shapes and one clean cluster is not enough for a calibration change, and result 2 is a live
demonstration of why: a rule that looked reasonable on the first pass, and a 14% "win" inside it
that turned out to be measurement noise.

If it is worth pursuing, the obvious next step is a designed sweep to check whether the
32x128-128 gap holds outside the shapes I happened to pick, since that is the only result here
carrying its own weight. If it is not worth it right now, that is a completely fine answer and I
will leave it here.

Raw per cell JSON, the recheck passes, the correctness gate results and the harness are available
if useful.

Environment: GB10 (DGX Spark), sm_121, 48 SMs, driver 580.142, torch 2.12.1+cu130, CUDA 13.0.88.