Tune SM103 fused 4-bit GEMM SIMT dispatch

[heiheiha798]

Summary

• Add an exact SM103/B300 internal dispatch rule that selects the existing SIMT fused 4-bit GEMM only for tall-K M <= 5 shapes from three up to, but not including, four output-tile waves.
• Retain MMA at four and later waves, for adjacent M >= 6 shapes, and on every non-SM103 architecture.
• Add a tracked B300 benchmark CLI that isolates automatic, forced-SIMT, and legal forced-MMA builds below --build-root, supports native or official-wheel-compatible CUDA targets, and records interleaved CUDA-event samples as JSONL.
• Add focused public-op correctness coverage across FP16/BF16, NF4/FP4, nested statistics, bias, block sizes 32/64/128/256, and a non-64-aligned K control.

Fork issue: heiheiha798#3

Problem and design

The internal dispatch in csrc/gemm_4bit.cu had no SM103-specific SIMT/MMA calibration. Direct forced-path measurements on a 148-SM B300 show a stable, dtype-independent SIMT region for public fused tall-K M=4,5 calls at the start of the third output-tile wave. Later-wave validation found a mixed result at exactly four waves and a clear MMA win at eight waves, so the final predicate is deliberately bounded:

cc_maj == 10 && cc_min == 3 && K > N && M <= 5 &&
    mma_blocks >= num_sms * 3 && mma_blocks < num_sms * 4

This changes no kernel math, public API, Python top-level dispatch, wheel target, dependency, or non-SM103 branch. At four and later waves the candidate retains the upstream MMA decision. The benchmark's forced-path edits exist only in temporary source trees and do not ship as an API.

B300 performance

Native-SM103 calibration used baseline job 4821 at 95f9af309d4d5793847169c39288dcd3fcbdf564 and candidate job 4829 at dc7eff3e7cfc2adc61874c82d494da897f09ae9f. It covered the 148/296/444-block boundaries, real projection controls, M=3..8, and FP16/BF16. Those native results established the initial three-wave target region.

Final job 4851 validated commit 7d9c98730d6d483f43e39d3af140a89498943953 using the official CUDA 13 x64 wheel target input 75;80;86;89;90;100;120, which CMake expanded to 75-real;80-real;86-real;89-real;90-real;100-real;120. Each path used 20 warmups followed by 7 rounds of 100 CUDA-event samples, with variant order rotating for every sample. Values below are p10 / median / p90 milliseconds from job 4851; old path is forced MMA.

N x K	dtype	M	Old MMA ms	Candidate auto ms	Median reduction
28416 x 32768	FP16	4	1.4180 / 1.4201 / 1.4222	0.9399 / 0.9420 / 0.9453	33.7%
28416 x 32768	BF16	4	1.4181 / 1.4201 / 1.4222	0.9553 / 0.9593 / 0.9621	32.4%
28416 x 32768	FP16	5	1.4219 / 1.4241 / 1.4262	1.2706 / 1.2717 / 1.2727	10.7%
28416 x 32768	BF16	5	1.4211 / 1.4232 / 1.4254	1.2665 / 1.2677 / 1.2686	10.9%
28480 x 32768	FP16	4	1.4048 / 1.4068 / 1.4096	0.9458 / 0.9481 / 0.9614	32.6%
28480 x 32768	BF16	4	1.4048 / 1.4068 / 1.4096	0.9470 / 0.9500 / 0.9644	32.5%
28480 x 32768	FP16	5	1.4078 / 1.4099 / 1.4120	1.2777 / 1.2788 / 1.2798	9.3%
28480 x 32768	BF16	5	1.4078 / 1.4099 / 1.4120	1.2747 / 1.2765 / 1.2769	9.5%

All seven round-median SIMT/MMA ratios remain below 0.95 for every three-wave target. Across FP16/BF16, the overall ratios are 0.663-0.675 for M=4 and 0.891-0.907 for M=5.

The same job directly checks why the upper bound is required:

Output blocks	Wave point	M	SIMT/MMA median ratio across dtypes	Winner beyond 5%
592	4 waves	4	0.901-0.902	SIMT
592	4 waves	5	1.168-1.171	MMA
1184	8 waves	4	1.595-1.625	MMA
1184	8 waves	5	2.019-2.025	MMA

The final automatic dispatch chooses SIMT at 444/445 blocks and MMA at 592/1184 blocks, with adjacent M=6 controls also on MMA. Automatic median latency matches its selected forced path within 0.11% over all 24 cells. The conservative common upper bound intentionally forgoes the isolated M=4 win at exactly four waves instead of introducing an M-dependent special case.

Job 4837, run at the pre-bound commit 187e4e4faa4d62ba060b0e7ec461ec134b8c2d9c, provided the decision-changing compatible-path evidence that led to this upper bound. Job 4834 is only a pre-build wrapper failure caused by passing already-suffixed CMake targets; it produced no benchmark data and is not performance evidence.

Reproduce

Run the following inside a one-GPU B300/SM103 Slurm allocation from a clean checkout. --build-root must not already exist.

python benchmarking/gemm_4bit_sm103_simt_mma.py \
  --output /tmp/bnb-sm103-compatible-range.jsonl \
  --build-root /tmp/bnb-sm103-compatible-range-build \
  --compute-capability '75;80;86;89;90;100;120' \
  --warmup 20 \
  --repetitions 100 \
  --rounds 7 \
  --case wave3_at_tall:28416:32768:4,5,6 \
  --case wave3_above_tall:28480:32768:4,5,6 \
  --case wave4_tall_k49152:37888:49152:4,5,6 \
  --case wave8_tall_k81920:75776:81920:4,5,6 \
  --dtype fp16 \
  --dtype bf16

Omitting --compute-capability builds native SM103 by default. The default grid also includes the 148/296/444-block boundaries, 11008x4096 and 14336x4096 projections, and a 128256x4096 vocabulary projection. The CLI requires exactly one visible B300 with CC 10.3 and records complete samples, round medians, p10/p90, correctness, commit, job, GPU, and software metadata. All three builds and the automatic package import are isolated below --build-root; the checkout library is not overwritten.

Correctness and checks

Job 4851 identity: NVIDIA B300 SXM6 AC, CC 10.3, 148 SMs, driver 580.126.09, CUDA compiler 13.0.88, PyTorch 2.13.0+cu130, Python 3.12.3.

• Exact-head official-compatible automatic, forced-SIMT, and forced-MMA builds passed from temporary source trees at 7d9c98730d6d483f43e39d3af140a89498943953.
• The exact-head public automatic benchmark ran 24 cells with 700 samples and 7 round medians per path per cell. All 72 path outputs were finite and repository-tolerance reference-close, including 444/445/592/1184-block M=4/5/6 coverage. Maximum absolute error was 0.03125 and maximum relative RMS error was 0.011809.
• The four pytest selections below ran later in job 4851, but their separate processes imported the checkout native library previously linked by native job 4829 at dc7eff3e7cfc2adc61874c82d494da897f09ae9f; they are regression evidence, not exact-head temporary-library validation. The 5 focused cases exercise the 444-block target where dc7eff3 and the final bounded predicate select the same paths. A rerun is unnecessary because the exact-head public automatic benchmark above directly loads and path-checks the final temporary library and covers every changed and restored boundary.
• BNB_TEST_DEVICE=cuda pytest -vv tests/test_ops.py -k gemm_4bit_sm103_tall_k_dispatch: 5 passed using the dc7eff3 checkout native artifact.
• BNB_TEST_DEVICE=cuda pytest -q tests/test_ops.py -k 'test_gemm_4bit and not sm103_tall_k_dispatch': 100 passed using the dc7eff3 checkout native artifact.
• BNB_TEST_DEVICE=cuda pytest -q tests/test_functional.py -k matmul_4bit: 1153 passed using the dc7eff3 checkout native artifact.
• BNB_TEST_DEVICE=cuda pytest -q tests/test_autograd.py -k matmul_4bit: 384 passed using the dc7eff3 checkout native artifact.
• CPU-only tests/test_ops.py collection: 392 tests collected without a CUDA import/collection failure; the 5 SM103 cases skip as intended.
• Submission preflight pre-commit run --all-files: all 10 hooks passed at the exact final head.
• git diff --check passed at the exact final head.

Raw workspace evidence:

• benchmark-evidence/cycle2/slurm-4821-simt-mma.out
• benchmark-evidence/cycle2/sm103-simt-mma-job-4821.jsonl
• benchmark-evidence/cycle2/slurm-4829-candidate.out
• benchmark-evidence/cycle2/sm103-candidate-job-4829.jsonl
• benchmark-evidence/cycle2/slurm-4834-compatible-range.out (pre-build failure, not evidence)
• benchmark-evidence/cycle2/slurm-4837-compatible-range.out
• benchmark-evidence/cycle2/sm103-compatible-range-job-4837.jsonl
• benchmark-evidence/cycle2/slurm-4851-compatible-range.out
• benchmark-evidence/cycle2/sm103-compatible-range-job-4851.jsonl

Limits and risk

• Measurements are from one physical B300. Native-SM103 and official-compatible builds are reported separately; no B200 hardware run was used.
• The exact runtime SM103 guard prevents a dispatch change on SM100 or any other architecture.
• The rule intentionally excludes the mixed exact-four-wave point and all later waves. This leaves one measured M=4 opportunity unused in order to keep one maintainable, dtype-independent rule with no adjacent regression.
• Job 4851's pytest subprocesses used the earlier dc7eff3 checkout native artifact; exact-head binary correctness comes from the path-asserted 24-cell public automatic benchmark, not those pytest counts.
• Fork CI is repository-gated; the full documented pre-commit workflow and focused B300 CUDA selections were run directly instead.

Tracking

• Fork engineering record: Tune SM103 fused 4-bit GEMM SIMT dispatch heiheiha798/bitsandbytes#4