Make non-paged CUDA optimizer updates current-stream correct and asynchronous

[heiheiha798]

Summary

The CUDA optimizer path currently uses device-wide synchronization as its ordering mechanism. Optimizer8bit.step() calls sync_gpu(p) after every updated parameter, while both native optimizer launcher families omit a stream argument and launch their kernels on the implicit default stream. The 32-bit max_unorm path also uses a synchronous, streamless device memset.

This issue proposes a bounded, measurement-gated change: make the existing non-paged NVIDIA CUDA optimizer updates obey PyTorch's current stream throughout the Python -> native ABI -> launcher chain, then remove the per-parameter device synchronization only from that now stream-ordered path.

The expected benefit is lower host serialization and fewer launch gaps for model-like optimizers with many tensors, especially LoRA/bias/norm-heavy training on B300. This is a hypothesis to validate, not a performance claim.

Baseline: upstream/main@95f9af309d4d5793847169c39288dcd3fcbdf564.

Current behavior

• bitsandbytes/optim/optimizer.py performs prefetch_state -> update_step -> sync_gpu for every parameter with a gradient. On CUDA, sync_gpu calls torch.cuda.synchronize().
• The CUDA implementations of bitsandbytes::optimizer_update_32bit and bitsandbytes::optimizer_update_8bit_blockwise invoke native C functions without passing PyTorch's current raw stream.
• optimizer32bit and optimizerStatic8bitBlockwise launch their kernels with the implicit default stream.
• The max_unorm precondition path resets its reduction buffer with a synchronous, streamless device memset.

The per-parameter barriers serialize the host with all device work. They also hide the fact that direct optimizer ops do not follow the current-stream convention used by modern bitsandbytes CUDA operations.

Bounded design

1. Add stream-aware optimizer C entry points for all existing 32-bit and blockwise 8-bit dtype/optimizer instantiations. Keep every existing exported optimizer symbol and signature as a legacy default-stream wrapper; use additive, clearly suffixed stream-aware symbols for the modern backend so direct C/ctypes users are not broken.
2. Pass bnb_stream_t through the two native launcher families and all explicit instantiations. Launch every precondition/update kernel on that stream.
3. Add the existing CUDA/HIP-compatible asynchronous memset abstraction and reset unorm on the same stream before the precondition kernel.
4. In bitsandbytes/backends/cuda/ops.py, route the registered optimizer ops through the new symbols and pass the gradient device's current raw stream.
5. In Optimizer8bit.step(), omit per-parameter sync_gpu only for non-paged NVIDIA CUDA execution (device.type == "cuda" and not ROCm). Preserve the existing synchronization behavior for paged optimizers, ROCm/HIP, XPU, and all other paths.

The change must be atomic: synchronization is not removed until every native operation in the selected path is stream ordered.

Explicit non-goals

• No optimizer math, bias-correction, state format, launch geometry, persistent grid, multi-tensor batching, or parameter grouping changes.
• No change to paged/UVM prefetching or synchronization; the prior paged lookahead experiment remains independent.
• No FSDP/DTensor behavior or distributed semantics.
• No CUDA graph support claim. Optimizer step is a Python integer today, so capturable/replay-safe bias correction is a separate problem.
• No asynchronous high-level behavior claim for ROCm/HIP or XPU without runtime validation.
• No deletion or signature change of existing native optimizer symbols.

Correctness and compatibility gates

Before broad implementation, reproduce at least one baseline defect/mechanism on B300: a trace must place a direct optimizer kernel on the wrong stream, or a deterministic event-gated producer/optimizer/consumer test must expose ordering dependence. If neither wrong-stream placement nor serialization is reproduced, stop.

For the candidate:

• Use event-gated non-default-stream tests for direct 32-bit and blockwise 8-bit optimizer ops. Produce gradients, update parameters/state, and consume the result on the same non-default stream without a device synchronization; compare with a serialized reference.
• Add a two-independent-stream test and verify each native kernel plus async memset runs on its caller's stream. Synchronize only recorded completion events.
• Require bitwise equality of parameters and complete state for deterministic paths across AdamW, Lion, Momentum/SGD, RMSprop, Adagrad, and AdEMAMix; 8-bit and 32-bit states; FP16/BF16/FP32 gradients where supported; one- and two-state paths; small 32-bit fallback and large blockwise tensors; aligned/tail sizes; weight decay; skip_zeros; and one/repeated steps.
• Exercise LAMB/LARS and other max_unorm > 0 paths for stream ordering and repository-reference correctness. Do not require cross-run bitwise equality for the existing atomic reduction, whose accumulation order is not deterministic.
• Prove all legacy C symbols remain resolvable with their old signatures and that the registered Python backend selects the new stream-aware symbols.
• Preserve paged optimizer behavior and tests, including its existing synchronization. Preserve HIP/XPU high-level synchronization; require shared-source HIP compile compatibility and disclose unavailable HIP runtime validation.
• Run focused optimizer tests, compute-sanitizer coverage for representative current-stream/max-unorm paths, the official CUDA architecture build, and full pre-commit.

B300 benchmark and profiling plan

Use isolated baseline and candidate builds from the exact baseline in one B300/SM103 Slurm allocation. Record source/library identity, GPU/driver/CUDA/PyTorch/bitsandbytes versions, build targets, commands, raw samples, and job/log paths.

Primary workloads:

• Non-paged AdamW8bit and AdamW32bit with model-like inventories of 32, 256, and 1024 tensors spanning LoRA, bias, norm, and transformer projection sizes, in FP16 and BF16.
• A single-large-tensor regression control.
• A small PEFT-style training loop whose forward/backward consumes updated parameters on the same stream.

Use at least 7 interleaved rounds after warmup. Optimizer-only timing must report completed host wall time and completed device time; synchronization occurs once at the measurement boundary, never once per candidate step. End-to-end iteration timing must likewise measure completed work rather than enqueue latency. Report median and dispersion.

Capture an Nsight Systems comparison showing:

• the baseline's device-wide synchronization calls and host launch gaps;
• zero intermediate device-wide synchronization in the selected non-paged NVIDIA CUDA path;
• optimizer kernels and async memset on the caller's stream;
• no hidden synchronization or changed kernel math/count beyond the removed barriers and stream placement.

Go / no-go

Proceed toward a draft PR only if all correctness and compatibility gates pass and all of the following hold:

• at least 1.25x median optimizer-section improvement for each 32/256/1024-tensor inventory;
• at least 5% median end-to-end improvement in one representative PEFT-style workload, with no statistically meaningful regression in the others;
• no more than 2% regression in the single-large-tensor control;
• the profile confirms the intended stream placement and barrier removal.

If the gain is only enqueue-time, only a cherry-picked tensor count, or cannot clear the public end-to-end gate, record a no-go rather than weakening the gate or widening into multi-tensor/paged work.

Duplicate and mergeability notes

Current upstream issue/PR and source/history searches found no optimizer current-stream or per-parameter-synchronization implementation. Merged PR bitsandbytes-foundation#1330 established stream propagation for other CUDA operations but did not cover optimizer launchers. Active optimizer PR bitsandbytes-foundation#2040 changes optimizer math in csrc/kernels.cu; this proposal deliberately leaves kernel math untouched. Nearby open work may create textual rebase conflicts, but no active PR owns these semantics.

The intended patch is a coherent system change across existing ownership boundaries, not a new kernel subsystem: additive native symbols, stream plumbing, same-stream memset/launches, a narrow non-paged NVIDIA CUDA synchronization change, focused tests, and reproducible B300 evidence.