Optimize NF4 single-token GEMV on Blackwell

[PiotrPotocki]

Summary

Optimize the fused 4-bit SIMT matrix-vector path for single-token NF4 decoding on NVIDIA Blackwell SM120 GPUs.

On an NVIDIA RTX PRO 4500 Blackwell, the geometric-mean latency across 16 benchmark cases decreased from 44.744 µs to 41.130 µs, an 8.08% latency reduction (1.088× speedup).

Implementation

Changes are contained in csrc/gemm_4bit_simt.cu:

• Use .cg packed-weight loads on SM120 and preload the next K group during single-row computation.
• Hoist reusable second-level quantization scales outside the inner loop for eligible double-quantized NF4 inputs.
• Share nested NF4 scales between adjacent lanes for FP16/BF16 inputs with M=1, N<=4096, K=4096, and block size 64.
• Reorganize FP16/BF16 activation loads and dequantization to improve instruction scheduling.

The device-code optimizations are guarded for SM120. No public API, quantization format, or runtime dependency is added.

Performance

Measured using the production torch.ops.bitsandbytes.gemm_4bit.default operator.

Setting	Value
GPU	NVIDIA RTX PRO 4500 Blackwell
CUDA library	CUDA 13.3
Fixed benchmark clocks	Graphics: 1800 MHz; memory: 14001 MHz
Activation shape	M=1
Activation types	FP16, BF16
Weight format	NF4, block size 64, contiguous packed uint8
Statistics	Ordinary and double-quantized
Double-quantization statistics block size	256
Weight shapes (N, K)	(4096, 14336), (14336, 4096), (12288, 4096), (4096, 4096)

Each case used 100 warm-up calls, followed by 30 batches of 10,000 calls. CUDA events bracketed the operator calls on the current stream, with synchronization after the end event. Batch duration was divided by 10,000, and the median per-call time across batches was recorded.

The aggregate score is the equally weighted geometric mean of the 16 case medians.

Metric	Baseline	Optimized
Geometric-mean latency	44.744 µs	41.130 µs
Latency reduction	—	8.08%
Speedup	—	1.088×

The largest measured reduction was 19.34% for FP16 with ordinary statistics at (N=4096, K=14336). The narrow BF16 double-quantized case was effectively unchanged, with a measured difference of approximately 0.17%.

These are operator-level measurements on one GPU, not end-to-end model decoding results.

Correctness validation

The local validation passed:

• Complete-output comparisons for all 16 benchmark cases against dequantization followed by torch.nn.functional.linear.
• Ten additional configurations covering bias, FP4/NF4, FP16/BF16/FP32, multiple activation rows, different block sizes, non-contiguous activations, and alternative quantization storage dtypes.
• Output shape, dtype, device, and finite-value checks.
• Repeated calls with modified inputs at the same memory address to verify output freshness.
• Invalid-input and unsupported nested-statistics checks.
• Existing repository NF4 test selections, shown below.

BNB_TEST_DEVICE=cuda python3 -m pytest -q \
  tests/test_functional.py::TestQuantize4BitFunctional::test_matmul_4bit \
  -k nf4

BNB_TEST_DEVICE=cuda python3 -m pytest -q \
  tests/test_ops.py::Test4bitBlockwiseQuantOps::test_gemm_4bit \
  -k nf4

Existing repository test tolerances were not changed. Additional local checks used fixed dtype-specific tolerances established against the baseline.

Scope and review notes

• Hardware validation was performed on SM120. Other GPU architectures and the complete repository test suite were not validated in these runs.
• The custom benchmark and additional validation scripts currently reside in the local optimization workspace and are not included in this patch.
• Nsight Compute profiles informed implementation decisions. Reported performance gains come from the operator benchmark; profiler timings used different observed clocks and were not used as the acceptance metric.

Tooling attribution

Developed with assistance from OptiMind, Unikie’s agentic code-performance optimization tool, which proposed and evaluated candidate changes through automated builds, correctness checks, benchmarking, and profiling. OptiMind is not a runtime dependency.