Fuse nested scale reconstruction into SM103 4-bit dequantization

[heiheiha798]

Problem

The default Linear4bit configuration uses compressed statistics. On CUDA, standalone dequantization of that nested state currently reconstructs the per-block scales in separate work before dequantizing the packed weights:

1. dequantize the uint8 scale codes to FP32;
2. add the nested offset in a separate elementwise operation;
3. launch the NF4/FP4 dequantization kernel using the materialized FP32 scale buffer.

bitsandbytes.functional.dequantize_4bit uses this sequence, including the path used by MatMul4Bit.backward for grad_A. The CUDA dequantize-plus-linear fallback also reconstructs the nested scales before materializing the full weight and calling F.linear, with an additional FP32 temporary in the current implementation.

On B300 (SM103), the main dequantization pass is fast enough that the extra launches, allocator work, and FP32 scale traffic may be material in recurrent backward and fallback paths.

Proposed experiment

Add a private CUDA nested 4-bit dequantization variant that reconstructs each scale inside the existing value-dequantization pass:

scale = scale_codebook[scale_code] * nested_absmax[nested_group] + offset
value = decode_4bit(packed_value) * scale

The kernel should consume the packed NF4/FP4 values, quantized scale bytes, nested FP32 absmax/codebook metadata, and scalar offset, and write the final dequantized output without materializing a full FP32 scale tensor.

Route this path only when all of the following hold:

• the runtime device is exactly compute capability 10.3;
• the state is nested/compressed;
• state2.blocksize == 256, matching the supported default nested layout;
• the format and output dtype are already supported by the current CUDA 4-bit dequantizer.

Use it from both functional.dequantize_4bit and the CUDA standalone dequantize-plus-linear fallback. All non-SM103 CUDA devices, HIP and other backends, nonnested states, unsupported metadata, and the existing quantization formats must retain their current path unchanged.

Scope and non-goals

• Do not change fused 4-bit GEMM kernels or their dispatch thresholds.
• Do not depend on any unmerged fork branch or prior SM103 dispatch experiment.
• Do not alter outer dequantization launch geometry, NF4/FP4 decode tables, serialization, QuantState, public schemas, or public APIs.
• Do not fold in the prior dequantization-barrier or SIMT nested-scale experiments.
• Prefer one dedicated internal op/helper over a broad cross-backend API redesign.
• Do not add dtype-, matrix-, or shape-specific dispatch beyond the exact-SM103 and supported-nested-state guard.

Correctness requirements

Baseline and candidate must be built independently from the same upstream commit. The fused result must be bitwise identical to the existing nested dequantization result for:

• NF4 and FP4;
• FP16, BF16, and FP32 output;
• outer block sizes 32 through 4096, with block size 64 primary;
• odd lengths, packed-value tails, partial outer blocks, and partial nested groups;
• both signs and representative magnitudes of the nested offset;
• every scale-code value and every 4-bit code value;
• allocating and .out forms, repeated launches, and transposed-output behavior already supported by the op.

Preserve the current FP32 scale reconstruction rounding sequence. In particular, prevent compiler contraction from changing the separately rounded multiply and add; do not weaken the bitwise oracle to accommodate an FMA difference.

Also require:

• current-stream ordering remains correct;
• nonnested and non-SM103 dispatch remains unchanged;
• fixed-input MatMul4Bit backward and verified fallback-forward results match baseline, using repository GEMM tolerances only after the dequantized tensor itself passes bitwise equality;
• focused functional, native-op, autograd, and Linear4bit tests pass;
• official CUDA multi-architecture build, full pre-commit, and focused compute-sanitizer checks pass.

CUDA graph capture is a compatibility check, not permission to expand this issue into output-buffer or schema redesign. If an existing allocation or metadata operation prevents capture, record that boundary while retaining direct preallocated-path coverage where feasible.

B300 evidence plan

Use one B300 allocation and isolated baseline/candidate builds from the same upstream baseline. The official CUDA 13 compatible target list is the primary deployment evidence; a native SM103 build may be reported separately as corroboration. Record the exact source identity, loaded libraries, compiler/runtime/GPU metadata, commands, raw logs, and profiler output.

Interleave baseline and candidate with CUDA events, at least 20 warmups, and at least seven timed batches with enough inner iterations for stable medians. Report all batch samples, median, p10, and p90 rather than a best-only result.

Primary standalone matrix:

• NF4 and FP4;
• FP16 and BF16, with FP32 controls;
• default block size 64;
• representative weights such as 4096x4096, 11008x4096, 4096x11008, 8192x8192, and 28672x8192, subject to allocation limits;
• an odd-tail control.

Public-path controls:

• MatMul4Bit.backward / grad_A for representative training shapes;
• CUDA dequantize-plus-linear fallback cells whose runtime path is explicitly asserted. Do not assume M > 32 alone proves fallback selection;
• steady-state Linear4bit forward/backward with quantization and setup excluded from timing.

Nsight Systems should confirm that nested standalone dequantization changes from scale-dequantization plus offset plus value-dequantization work to one fused kernel and removes the full FP32 scale temporary. Nsight Compute should report effective bandwidth, memory transactions, register use, occupancy, and metadata-load behavior.

Proceed to a draft PR only if:

• every correctness/build/sanitizer gate passes;
• standalone nested dequantization improves by at least 10% geomean across the primary large-matrix set;
• no primary standalone cell regresses by more than 3%;
• at least one representative recurrent grad_A cell and one verified fallback-forward cell improve by at least 3%, with dispersion below the effect;
• no primary public-path cell regresses by more than 3%.

If the gain is isolated, allocator-dependent, disappears in the public operations, or requires broader dispatch, record a no-go instead.

Baseline

upstream/main@95f9af309d4d5793847169c39288dcd3fcbdf564

This issue records a B300 measurement hypothesis. It does not claim a performance improvement before the controlled evidence is complete.