Respect the current CUDA stream in blockwise quantization

[heiheiha798]

Summary

Blockwise CUDA quantization currently launches on the legacy default stream even when PyTorch has selected a different current stream. This can let quantization run before a producer on that current stream has finished writing its input.

This change:

• adds stream-aware native entry points for General8, FP4, and NF4 blockwise quantization;
• passes PyTorch's current raw CUDA stream from both registered CUDA quantization operations;
• threads that stream through every existing block-size launch from 32 through 4096;
• preserves every existing six-argument C symbol and its stream-0 behavior for ABI compatibility; and
• adds a focused ordering regression test and a reproducible B300 benchmark CLI.

The public schema, Python API, output allocation, launch geometry, quantization math, and existing C ABI are unchanged.

Tracking:

• Fork issue: Respect the current CUDA stream in blockwise quantization heiheiha798/bitsandbytes#18
• Fork draft PR and engineering record: Respect the current CUDA stream in blockwise quantization heiheiha798/bitsandbytes#19

Reproduced failure

Slurm job 5099 built the exact baseline 95f9af309d4d5793847169c39288dcd3fcbdf564 for the official CUDA 13 x64 target input 75;80;86;89;90;100;120 and ran on one NVIDIA B300 SXM6 AC (CC 10.3, 148 SMs).

The reproducer recorded a pending blocker event before entering the test stream, made that stream wait, filled the input there, and immediately called the public quantizer. All four baseline paths read stale data while the final input itself was correct:

• direct General8 blockwise quantization;
• nested General8 blockwise quantization;
• direct NF4 quantization; and
• compressed-statistics NF4 quantization.

Archived evidence is listed in the linked fork draft PR.

Correctness and safety

Slurm job 5101 built production/test commit 86aabee18f42d54448e26cb281aea8fe083b8b16 with the same official-compatible target list and asserted the loaded temporary library path.

• 345/345 raw-bit parity cells passed between the preserved legacy entry points and the new stream-aware entry points.
• Coverage included FP16/BF16/FP32, General8/FP4/NF4, every supported block size from 32 through 4096, full blocks, and tails.
• 73 focused ops cases passed with 30 known storage-dtype opcheck xfails.
• 144 focused functional cases passed.
• compute-sanitizer --tool racecheck: 0 hazards, 0 errors, 0 warnings.
• compute-sanitizer --tool synccheck: 0 errors.

The final head 56d8d94082568e7d6a8c0788e4247583cdbb65d5 changes only the tracked benchmark after that production/test commit. Job 5104 separately built and loaded that exact final head and rechecked bitwise equality throughout the benchmark matrix.

CPU collection also remains healthy: 33 passed, 22 CUDA-only skips, 345 deselected. Full pre-commit run --all-files passes.

B300 measurement

Job 5104 built isolated baseline and candidate source trees with CUDA 13.0.88 and the official-compatible targets, asserted both native library paths, and measured on the same B300 allocation. The tracked CLI interleaves variants after 20 warmups for 7 rounds, retaining every sample, per-round medians, p10/p90, bootstrap ratio intervals, and input-byte effective bandwidth.

Default-stream direct latency covered 36 cells: FP16/BF16/FP32 x General8/FP4/NF4 x n={524288,8388608,67108864,536870912} elements. Those counts are 1 MiB, 16 MiB, 128 MiB, and 1 GiB for FP16/BF16 inputs, and 2 MiB, 32 MiB, 256 MiB, and 2 GiB for FP32 inputs. The worst baseline/candidate ratio was 0.98310, equivalent to a 1.72% candidate slowdown, so the <=2% regression gate passed.

The correct two-stream pinned-H2D plus compressed-NF4 pipeline covered six cells. The baseline explicitly waits for both copies before stream-0 quantization; the candidate chains each copy and public quantizer on its current stream.

dtype	shape	baseline ms	candidate ms	ratio	95% ratio interval
FP16	4096x4096	1.6109	1.5937	1.0108x	[1.0091, 1.0121]
FP16	11008x4096	4.0615	4.0386	1.0057x	[1.0047, 1.0063]
FP16	4096x11008	4.0521	4.0332	1.0047x	[1.0041, 1.0053]
BF16	4096x4096	1.6333	1.6175	1.0098x	[1.0086, 1.0106]
BF16	11008x4096	4.0723	4.0497	1.0056x	[1.0049, 1.0060]
BF16	4096x11008	4.1249	4.1045	1.0050x	[1.0045, 1.0065]

These gains are 0.47%-1.08%, below the 5% performance-claim threshold. This PR is therefore a correctness fix and makes no B300 speedup claim. Nsight Systems confirms H2D and quantization activity in the serialized-baseline and current-stream-candidate ranges; its one-shot range durations include setup and are not used as performance claims.

Archived evidence is listed in the linked fork draft PR.

Reproduction

After building baseline and candidate into isolated source trees with the CUDA 13 x64 target list, run:

export PYTHONPATH=/tmp/bnb-quant-stream/source-candidate
export BENCHMARK_BASELINE_REVISION=95f9af309d4d5793847169c39288dcd3fcbdf564
export BENCHMARK_GIT_REVISION=56d8d94082568e7d6a8c0788e4247583cdbb65d5
export BENCHMARK_CUDA_TARGETS='75;80;86;89;90;100;120'

python /tmp/bnb-quant-stream/source-candidate/benchmarking/quantize_blockwise_stream.py \
  --baseline-library /tmp/bnb-quant-stream/source-baseline/bitsandbytes/libbitsandbytes_cuda130.so \
  --candidate-library /tmp/bnb-quant-stream/source-candidate/bitsandbytes/libbitsandbytes_cuda130.so \
  --output /tmp/quantize-blockwise-stream.jsonl \
  --warmup 20 \
  --rounds 7 \
  --direct-repetitions 100 \
  --pipeline-repetitions 15

The exact build and invocation transcript is listed in the linked fork draft PR.

Limitations

• Fixed-shape CUDA graph capture was attempted, but the existing public path copies the quantization code map from CPU to CUDA during capture. PyTorch rejects that pre-existing allocation/copy surface. This PR does not add an out schema or memory-management redesign.
• Neither the login environment nor the Slurm B300 environment provides ROCm/hipcc, so HIP compile/runtime validation was unavailable. The shared source remains syntactically scoped through the existing CUDA/HIP stream compatibility layer, but downstream HIP CI is required.
• CUDA targets other than CC10.3 were compiled but not executed on their hardware.
• The fork does not run upstream GitHub checks; upstream CI remains required before merge.