Fuse nested 4-bit scale reconstruction on SM103

[heiheiha798]

Summary

Default compressed 4-bit state reconstructs per-block FP32 scales in three CUDA launches: General8 scale dequantization, an offset add, and NF4/FP4 value dequantization. This change adds a private nested NF4/FP4 dequantization kernel that performs the same separately rounded FP32 multiply and add while decoding values, eliminating the intermediate FP32 scale tensor and two launches.

The production path is intentionally narrow:

• exact runtime compute capability 10.3 only;
• nested quantization state with state2.blocksize == 256;
• existing NF4/FP4 and FP16/BF16/FP32 output combinations;
• standalone functional.dequantize_4bit in eager mode and the existing CUDA dequantize-plus-linear fallback.

Non-SM103 CUDA devices, HIP/non-CUDA backends, nonnested states, unsupported nested metadata, and torch.compile tracing retain the legacy registered-op chain. This does not change GEMM dispatch, fallback selection, launch geometry, quantization formats, serialization, any existing torch-op schema, or any public Python API. It adds two internal torch-op schemas, dequantize_4bit_nested and its .out variant, as dispatch plumbing for the nested path.

Implementation

• Add private CUDA nested-dequant kernel/launcher instantiations for NF4/FP4 and FP16/BF16/FP32.
• Preserve the legacy FP32 operation sequence with __fmul_rn followed by __fadd_rn; no FMA contraction or relaxed correctness oracle is used.
• Add private ctypes entry points and route production calls through them only on exact SM103.
• Preserve the established public .out shape/dtype validation before any native call.
• Keep the registered-op chain while torch.compiler.is_compiling() is true so fullgraph backward remains traceable.
• Add focused dispatch/public-path tests and a configurable B300 benchmark CLI.

The CUDA library built with the official-compatible target list grew from 3,922,064 to 4,053,448 bytes, an increase of 131,384 bytes.

Validation

Kernel and performance validation was recorded in Slurm job 5140 at commit 4c7e1f54cab5dffeb9e308ab958e6aa38d051813. Review fixes were validated in Slurm job 5157 at fe35f68b46c9a285e96cf9488513c7b15a17f056. Both jobs ran on one NVIDIA B300 SXM6 AC (CC 10.3, 148 SMs), driver 580.126.09, CUDA 13.0.88, and PyTorch 2.13.0+cu130. Candidate builds used CUDA targets 75-real;80-real;86-real;89-real;90-real;100-real;120, and each process asserted that it loaded the library from its isolated source tree.

• The private C-symbol oracle passed 630 raw-bit baseline/candidate cases with input and metadata immutability. The matrix covered NF4/FP4, FP16/BF16/FP32, outer block sizes 32 through 4096, full/odd/tail nested groups, scale/value code sweeps, offset controls, and repetitions.
• Public allocating and .out behavior and architecture guards were covered separately by focused pytest.
• compute-sanitizer --tool racecheck: 0 hazards, 0 errors, 0 warnings.
• compute-sanitizer --tool synccheck: 0 errors.
• Job 5140 focused routes: 23 CUDA passes and 23 corresponding CPU skips.
• Job 5157 focused routes: 24 CUDA passes, including wrong-shape and wrong-dtype .out rejection before the native kernel.
• Job 5157 existing fullgraph compile case: one compressed-statistics NF4/FP32/no-bias/default-mode/batch-1 case passed, including input gradients.
• Job 5140 relevant existing tests: 312 ops passes (156 CUDA and 156 CPU), 24 compressed-statistics functional passes (12 CUDA and 12 CPU), and 768 MatMul4Bit autograd passes (384 CUDA and 384 CPU).

The post-submission follow-up at exact head 23aaa291eb5c9d917209f268076e8d3f218f666e changes only tests/test_functional.py. Its forced-native test monkeypatches the architecture predicate only inside the test, so it exercises the architecture-generic native candidate without widening production dispatch. Slurm job 7303 built for SM103 and passed all 19 selected CUDA cases (18 dtype/format/block-size combinations plus the non-SM103 production-guard case). CPU targeted collection produced the expected 20 CUDA skips, while the nested torch-op CPU tests passed 7 cases. Fresh exact-head uvx --from pre-commit pre-commit run --all-files --show-diff-on-failure and git diff --check upstream/main...HEAD both pass. Jobs 7301 and 7302 stopped during setup (shell pipefail, then missing CMake) and contain no test result.

The official PR CUDA workflow builds targets 75;80;89 and schedules the full test suite on T4, A10, and L40S runners across CUDA 11.8, 12.6, and 13.0. The forced-native correctness test is reachable in that matrix, but those GitHub Actions jobs are currently awaiting workflow approval and have not run; no pass claim is made for those GPUs.

Job 5126 stopped before configuration because CMake was absent from the inherited PATH and is setup-failure evidence only. Job 5127 is a successful baseline-transform pilot and supporting evidence only. Job 5140 is the authoritative committed kernel/build/correctness/sanitizer/performance/profile run. Job 5157 is the authoritative .out, eager-route, and fullgraph compatibility run. Job 7303 is the authoritative exact-current-head test-only follow-up run.

B300 Performance

Both variants used the same allocation and inputs. Timings used CUDA events, 20 alternating warmups, seven interleaved timed batches, and size-scaled repetitions of up to 100 per batch. Every recorded cell includes all round samples plus median, p10, and p90.

The 24-cell direct matrix covers four realistic matrix shapes, NF4/FP4, and FP16/BF16/FP32 at block size 64:

Shape	Cells	Geomean baseline/candidate	Minimum	Maximum
4096 x 4096	6	1.4200x	1.3979x	1.4426x
11008 x 4096	6	1.1313x	1.1077x	1.1552x
4096 x 11008	6	1.1338x	1.1134x	1.1573x
8192 x 8192	6	1.1301x	1.1095x	1.1550x
All direct cells	24	1.1978x	1.1077x	1.4426x

All direct outputs were raw-bit equal. Every direct cell improved by more than 10%; per-round ratio dispersion remained below the observed effect. An independently executed tracked CLI reproduced a 1.1976x geomean and 1.1099x minimum across the same 24 cells.

Public-path controls, with quantization/setup excluded from steady-state timing:

Path	Shape	Baseline median	Candidate median	Ratio
recurrent grad_A	M=1, N=4096, K=4096	33.098 us	26.735 us	1.2380x
recurrent grad_A	M=8, N=4096, K=4096	32.863 us	24.748 us	1.3279x
asserted fallback forward	M=32, N=4096, K=4096	32.903 us	24.752 us	1.3293x
asserted fallback forward	M=32, N=11008, K=4096	68.804 us	61.726 us	1.1147x

The fallback cells assert the selected runtime path rather than inferring it from M alone. The later review and test-only fixes do not change the native kernel or eager selected implementation, so the native performance matrix was not rerun.

Nsight Systems shows the legacy General8 scale-dequantization, FP32 offset-add, and 4-bit value-dequantization sequence becoming one kDequantizeBlockwiseNested launch, with the full FP32 scale temporary removed. NCU on the candidate kernel reports 33.95 us duration, 69.15% compute-memory throughput, 65.08% SM throughput, 66.17% active warps, 40 registers per thread, and 2.66 KiB static shared memory.

Reproduction

After building the candidate with CUDA 13 and target input 75;80;86;89;90;100;120, load that isolated package and run the benchmark on one B300:

BNB_CUDA_VERSION=130 python benchmarking/nested_4bit_dequant.py \
  --output /tmp/nested-4bit-dequant.jsonl \
  --cases square4096:4096:4096,wide11008:11008:4096,tall11008:4096:11008,square8192:8192:8192 \
  --dtypes fp16,bf16,fp32 \
  --formats nf4,fp4 \
  --blocksize 64 \
  --warmup 20 \
  --rounds 7 \
  --repetitions 100

The CLI gates on one B300 SM103 with 148 SMs, records the loaded native library and environment in JSONL, alternates measurement order, reports complete samples/medians/p10/p90/effective bandwidth, and checks raw-bit equality against the legacy launch chain.

Limits

• Production dispatch remains exact SM103; forcing the native predicate exists only in correctness tests.
• Runtime performance and sanitizer evidence comes from one physical B300 and one CUDA/toolchain stack.
• The architecture-generic native candidate was compiled for the official-compatible CUDA targets, but runtime correctness on T4/A10/L40S is pending the official PR workflow described above.
• HIP/hipCUB compile and runtime validation were unavailable because neither the login environment nor the Slurm environment provided ROCm/hipcc. This PR makes no HIP behavior or performance claim; downstream HIP compilation is required before merge.
• Native-SM103-only measurements are not substituted for official multi-architecture CI.
• Fork CI does not substitute for upstream multi-architecture CI.

Tracking

• Fork issue: Fuse nested scale reconstruction into SM103 4-bit dequantization heiheiha798/bitsandbytes#21
• Fork engineering record: Fuse nested 4-bit scale reconstruction on SM103 heiheiha798/bitsandbytes#22