Respect the current CUDA stream in blockwise quantization

[heiheiha798]

Summary

The existing CUDA blockwise quantizers do not receive the PyTorch current stream. The Python backend calls six-argument cquantize_blockwise_* functions, quantizeBlockwise has no stream parameter, and its CUDA kernels launch with <<<grid, threads>>>, which uses stream 0.

That differs from the adjacent dequantization, GEMV, GEMM, and INT8 paths, which pass _get_raw_stream(A.device.index). Under a nonblocking non-default stream, the quantizer can therefore run outside the producer/consumer ordering established by PyTorch. The nested paths are especially relevant: quantize_4bit(..., compress_statistics=True) and nested quantize_blockwise interleave native quantization with current-stream PyTorch operations.

This issue tracks a bounded correctness and systems-performance hypothesis: propagate the current stream through the existing General8/FP4/NF4 quantizers without changing their math, launch geometry, public schemas, or legacy C ABI. B300/SM103 will be used to reproduce the ordering defect and measure whether correct copy/quantization pipelining improves public model-loading workloads. No performance gain is assumed in advance.

Related precedent: bitsandbytes-foundation#1308 reported incorrect CUDA-graph behavior, and bitsandbytes-foundation#1330 fixed the same missing-stream class for dequantization/GEMV. The existing quantization wrappers were not covered by that change.

Proposed scope

• Add stream-aware native entry points for the existing FP16/BF16/FP32 General8, FP4, and NF4 quantizers.
• Keep every existing six-argument cquantize_blockwise_* symbol and its stream-0 behavior unchanged for direct C consumers. New stream-aware symbols may use a clear suffix such as _with_stream and append bnb_stream_t/cudaStream_t.
• Thread the supplied stream through quantizeBlockwise and all existing blocksize 32 through 4096 launches using <<<..., 0, stream>>>.
• Bind the new symbols in bitsandbytes/backends/cuda/ops.py and pass _get_raw_stream(A.device.index) from both registered quantize operations.
• Preserve all dtype/format dispatch, block geometry, packing, reductions, rounding, stochastic behavior, output allocation, and public torch.library schemas.
• Keep the shared HIP source buildable. B300 supplies the performance evidence; this issue does not claim a ROCm performance improvement.

Likely files are bitsandbytes/backends/cuda/ops.py, csrc/ops.cu, csrc/ops.cuh, csrc/pythonInterface.cpp, and focused quantization tests. A tracked benchmark is warranted only if it is needed to make the copy/quantization pipeline reproducible.

Correctness plan

Default-stream parity

Require bitwise-identical packed output and absmax before and after the change for:

• FP16, BF16, and FP32 inputs;
• General8, FP4, and NF4;
• all supported block sizes, odd tails, all-zero blocks, and representative finite/extreme values;
• nested quantize_blockwise and quantize_4bit(compress_statistics=True);
• allocating and caller-provided output/state behavior already supported by the public APIs.

For normal Python calls, assert that the stream-aware symbol is actually loaded. Also require all legacy six-argument symbols to remain resolvable with their old signatures.

Deterministic non-default-stream dependency test

Do not use wait_event() on an event that has never been recorded; CUDA treats that as an already-completed dependency and it cannot prove ordering.

Instead:

1. Initialize A to a stale value and synchronize.
2. On a separate nonblocking blocker stream, enqueue a bounded long-running operation and then record event gate. The event is recorded but still pending.
3. On test stream S, wait for gate, enqueue A.fill_(known_value), and call the public quantizer.
4. Compare against a serialized reference produced from known_value.

The current stream-0 quantizer can execute while S is blocked and observe stale input. A correct implementation remains ordered behind the wait and fill. Cover direct General8, direct FP4/NF4, and the nested/compressed-statistics functional paths.

Attempt fixed-shape CUDA-graph capture/replay after warmup as supporting evidence. If existing output allocation prevents capture, report that limitation; do not expand this issue into new output schemas or a memory-management redesign.

Run the relevant test_ops.py and functional quantization suites, pre-commit, the official CUDA multi-architecture build, and HIP/hipCUB compile validation. CUDA compilation and GPU validation must run through the established Slurm workflow.

B300 validation and benchmark

Use isolated baseline and candidate builds from the same upstream commit in one B300 allocation. Record the loaded library, source commit, GPU/driver/toolkit/PyTorch versions, commands, raw logs, and resolved CUDA targets.

Measure:

• per-op default-stream latency for General8 and NF4/FP4 from 1 MiB through 1 GiB;
• NF4/FP4 with compressed statistics;
• a correct two-stream model-load proxy using pinned FP16/BF16 weights shaped like 4096x4096, 11008x4096, and 4096x11008, repeated to at least 8 GiB when memory permits;
• host wall time, CUDA-event kernel time, Nsight Systems stream placement, copy/quantization overlap, and default-stream contamination.

Compare the candidate pipeline with a correctly serialized baseline in which the copy completes before the stream-0 quantizer. Do not benchmark against the known-racy execution. Use interleaved runs, warmups, at least 100 timed pipeline repetitions where practical, medians and dispersion/confidence intervals rather than best runs.

Acceptance

The change is upstream-worthy as a correctness fix if the non-default-stream failure is reproduced on the baseline, fixed by the candidate, default-stream results remain bitwise identical, the old C ABI remains intact, and the relevant build/test matrix passes.

Claim a B300 performance improvement only if the correct copy/quantization pipeline improves median model-load proxy wall time by at least 5% with a positive confidence interval and profiler-confirmed overlap, while default-stream median kernel latency regresses by no more than 2%. If neither a correctness/capture failure nor a stable pipeline benefit can be reproduced, record a no-go rather than broadening the scope.

Non-goals and coordination

• No quantization arithmetic, tile, blocksize, packing, stochastic, or architecture-dispatch changes.
• No changes to dequantization, GEMM/GEMV, INT8 matmul, optimizers, public schemas, or output-buffer APIs.
• No device-wide synchronization or fallback added to mask ordering bugs.
• Do not stack on the unmerged k-bit work in Add k-bit blockwise quantization (K=2-5) with warp-level CUDA kernels bitsandbytes-foundation/bitsandbytes#1858. Rebase if it lands and keep this change limited to the existing General8/FP4/NF4 quantizers.
• Do not change old exported C function signatures; additive stream-aware symbols are the compatibility boundary for this investigation.