Add k-bit blockwise quantization (K=2-5) with warp-level CUDA kernels

[TimDettmers]

Summary

• Add k-bit blockwise quantization/dequantization CUDA kernels for K=2, 3, 4, 5 with blocksize=32, using warp-level intrinsics (__ballot_sync packing, __shfl_sync codebook lookup)
• Add E4M4 uint8 absmax format (bias=11) that reduces per-block scale storage from 4 bytes to 1 byte, bringing K=4 to 4.25 bits/element
• New Python API: quantize_kbit(), dequantize_kbit(), create_normal_float_codebook(), encode_absmax_e4m4(), decode_absmax_e4m4()
• 222 tests passing, including correctness, error analysis, NF4 cross-validation, and performance benchmarks

Architecture

Each quantization block of 32 elements maps to exactly one CUDA warp. This enables:

• Bit-plane packing via __ballot_sync: K calls produce K uint32 words with zero bit waste for any K value. No word-boundary issues for odd K (unlike sequential packing).
• Codebook lookup via __shfl_sync: Each lane holds one codebook entry (up to 2^K=32 entries fit in warp width). Lookup is register-to-register (~5 cycles), no shared memory needed.
• Absmax reduction via __shfl_down_sync: 5 reduction steps, no shared memory, no __syncthreads().
• Multi-block dequant: Each warp processes 4 quantization blocks, amortizing codebook load.

Zero shared memory used. Zero warp divergence in the hot path. Templated on output type (fp16/bf16/fp32) and absmax format (E4M4 uint8/fp16).

E4M4 absmax format

uint8 micro-float with 4-bit exponent, 4-bit mantissa, bias=11, IEEE-style subnormals. Range [6.1e-5, 31.0]. Mean encode/decode relative error: 1.1%, 95th percentile: 2.4%. SQNR degradation vs fp32 absmax: <0.4 dB across all K values. Decode uses direct IEEE 754 bit construction (__uint_as_float) for zero overhead on the dequant hot path.

Benchmarks (RTX 4090, 67M elements, E4M4 absmax)

Dequant kernel throughput

K	bits/elem	fp16 (us)	bf16 (us)	fp32 (us)	GB/s (fp16)	% peak BW
2	2.25	205	208	394	781	78%
3	3.25	215	215	416	786	78%
4	4.25	244	246	420	729	72%
5	5.25	271	270	428	689	68%

Comparison with existing NF4 (fp16 output)

Method	bits/elem	Dequant (us)	GB/s	% peak BW	vs NF4
NF4 bs=64	4.50	239	754	75%	ref
kbit K=2	2.25	215	745	74%	1.11x
kbit K=3	3.25	229	737	73%	1.04x
kbit K=4	4.25	246	722	72%	0.97x
kbit K=5	5.25	271	689	68%	0.88x

K=4 is at parity with NF4 in kernel throughput (0.97x) while using 0.25 fewer bits/element. K=2 and K=3 are faster due to less data to read. Both kernels are bandwidth-bound at 68-78% of peak HBM.

Quality (SQNR, 1M elements, normal distribution)

Method	SQNR (dB)	MSE	bits/elem
NF4 bs=64	20.74	0.0085	4.50
kbit K=2	7.43	0.181	2.25
kbit K=3	14.99	0.032	3.25
kbit K=4	21.09	0.0078	4.25
kbit K=5	25.95	0.0026	5.25

K=4 achieves comparable quality to NF4 at 0.25 fewer bits/element (4.25 vs 4.50). K=3 offers 3.25 bits/element (4.9x compression) with 15 dB SQNR. K=2 provides 7.1x compression for extreme quantization.

Storage comparison

Method	bits/elem	Compression vs fp16
NF4 bs=64	4.50	3.6x
kbit K=5	5.25	3.0x
kbit K=4	4.25	3.8x
kbit K=3	3.25	4.9x
kbit K=2	2.25	7.1x

Key design decisions

Decision	Choice	Rationale
Blocksize	32 (fixed)	Maps exactly to CUDA warp width for __ballot_sync/__shfl_sync
K range	2-5	Codebook fits in warp lanes (2^K <= 32)
Packing	Bit-plane via __ballot_sync	Zero waste for any K; sequential packing has word-boundary issues for odd K
Codebook lookup	__shfl_sync from lane registers	Register-to-register, no shared memory
Absmax format	E4M4 uint8 (default), fp16 (option)	1 byte/block vs 2-4 bytes; <0.4 dB SQNR loss
Codebook	Symmetric normal-float default, arbitrary user-provided	Maximum flexibility; NF codebooks precomputed in Python
Output types	fp16, bf16, fp32 (templated)	Feature parity with existing kernels
Target hardware	CC 7.0+ (Volta)	Required for _sync warp primitives

API

from bitsandbytes.functional import quantize_kbit, dequantize_kbit

# Quantize (default: symmetric normal-float codebook, E4M4 absmax)
packed, absmax, codebook = quantize_kbit(A, k=4)

# Dequantize
recovered = dequantize_kbit(packed, absmax, codebook, k=4, n=A.numel(), dtype=torch.float16)

# Custom codebook
my_cb = torch.tensor([...], dtype=torch.float32, device="cuda")
packed, absmax, cb = quantize_kbit(A, k=3, codebook=my_cb)

The default codebook is a symmetric normal-float codebook (quantiles of N(0,1), normalized to [-1, 1]). Unlike the existing NF4 codebook which is asymmetric (7 negative, 1 zero, 8 positive), this codebook has equal representation on both sides (8 negative, 8 positive, no explicit zero). Custom codebooks of any shape can be passed via the codebook parameter.

Test plan

222 tests covering:

• Pure Python reference implementation
• CUDA quantize/dequant correctness against reference
• Error analysis (analytical bounds, MSE scaling, dtype consistency)
• Cross-validation against existing NF4 at K=4
• Performance benchmarks (bandwidth utilization, scaling, NF4 comparison)
• E4M4 absmax (encode/decode round-trip, SQNR degradation, edge cases)
• Output dtype correctness (bf16/fp32 match fp16 within precision)
• Asymmetric codebooks (all-positive, all-negative, skewed, non-uniform, duplicate entries)

Files changed

• csrc/ops.cu — CUDA kernels and launchers (+229 lines)
• csrc/pythonInterface.cpp — C interface wrappers (+125 lines)
• bitsandbytes/functional.py — Python API (+194 lines)
• bitsandbytes/_ops.py — torch.library op definitions (+44 lines)
• bitsandbytes/backends/cuda/ops.py — CUDA backend dispatch (+90 lines)
• tests/test_kbit_quantization.py — Test suite (+1372 lines)

[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.

[TimDettmers]

PR Review: #1858 — Add k-bit blockwise quantization (K=2-5) with warp-level CUDA kernels

Adds a complete k-bit blockwise quantization subsystem: CUDA kernels using __ballot_sync bit-plane packing and __shfl_sync codebook lookup (blocksize=32), E4M4 uint8 absmax format, Python API (quantize_kbit/dequantize_kbit), torch.library op definitions, CUDA backend dispatch, and 222 tests covering correctness, error analysis, NF4 cross-validation, and benchmarks.

No blocking issues.

The CUDA kernel architecture is sound — warp-level primitives for packing/lookup/reduction eliminate shared memory usage entirely, the multi-block dequant amortizes codebook load, and the E4M4 decode uses direct IEEE 754 bit construction for zero overhead. Template instantiations cover all type combinations. The Python-level API follows existing patterns (torch.library.define, register_kernel, register_fake) and integrates cleanly with the existing quantization stack.

Suggestion: The create_normal_float_codebook function requires scipy at runtime (for norm.ppf), which is only in the [test] extra. If kbit quantization is intended for production use, the codebooks could be hardcoded (like NF4's get_4bit_type) to remove the scipy dependency, or scipy could be moved to a non-test dependency group.

Suggestion: The dequant CUDA backend dispatch (backends/cuda/ops.py:140-143) silently converts fp32 absmax to E4M4 before passing to the kernel. This is convenient but means a user who explicitly passes fp32 absmax gets E4M4 precision without opting in. A warning or explicit parameter might be clearer.

• Security: Clear
• Downstream impact: None (additive — new API, no changes to existing functions)
• Tests: Adequate (222 tests including reference cross-validation, error bounds, edge cases, benchmarks, asymmetric codebooks)
• CI: All pass (Lint, CPU x4 platforms, CUDA T4+L40S x3 CUDA versions, Windows CUDA, ROCm builds, XPU builds)
• torch.compile: Compatible (uses torch.library registration with register_fake shape functions)
• Serialization: N/A (new format, no backward compat concerns yet)

[TimDettmers]

The create_normal_float_codebook function imports scipy at runtime. Since scipy is only in the [test] extra, calling quantize_kbit() without a pre-built codebook in a production install will raise ImportError. Consider hardcoding the codebooks (like get_4bit_type does for NF4) or documenting the scipy requirement.

[TimDettmers]

When the user passes fp32 absmax, the dequant dispatch silently encodes it to E4M4 before calling the kernel. This is a lossy conversion the caller may not expect — they passed fp32 precision but get E4M4 precision. Consider either warning or documenting this behavior.