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 new 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, schemas, serialization, or public APIs.

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 only eligible exact-SM103 eager calls through them.
• 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 the final head 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 final-head 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).
• CPU-only targeted collection skipped all 19 SM103-only cases cleanly.
• Submission preflight at repaired head 2883b3873ffd32953c1c32620fdd8d3f6e0bb600: uvx --from pre-commit pre-commit run --all-files --show-diff-on-failure and git diff --check upstream/main...HEAD pass.

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 exact-final-head .out, eager-route, and fullgraph compatibility 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 final review fix changes only invalid-.out validation and compile-time routing; the native kernel and eager selected implementation are unchanged, so the native performance matrix was not rerun after that fix.

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

• Runtime performance and sanitizer evidence comes from one physical B300 and one CUDA/toolchain stack.
• Other official-compatible CUDA targets were compiled into the library but were not executed on their hardware.
• 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 the official-compatible CUDA 13 evidence above.
• 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