Use the current CUDA stream for non-paged optimizer updates

[heiheiha798]

Summary

• add stream-aware optimizer entry points while preserving every legacy C symbol and its default-stream behavior;
• pass PyTorch's current raw stream through both 32-bit and blockwise 8-bit optimizer launcher families, including async max_unorm reset and precondition kernels;
• keep updates for actual managed/paged state on stream 0 so they remain ordered after the existing stream-0 UVM prefetch;
• remove per-parameter device synchronization only for non-paged NVIDIA CUDA optimizers; paged, HIP/ROCm, XPU, CPU, and other non-selected high-level paths keep their existing synchronization;
• add focused event-gated stream, paged-ordering, max-unorm, ABI, and synchronization-policy tests plus a reproducible interleaved benchmark CLI.

The baseline launches registered optimizer updates on stream 0. With a pending event on a nondefault caller stream, a non-paged optimizer can read the stale gradient before the caller-stream producer runs. It also calls torch.cuda.synchronize() once per active parameter. This patch orders ordinary non-paged NVIDIA CUDA work on the caller's current stream and relies on the caller's normal synchronization boundary.

Paged managed state is deliberately excluded from that routing. Its existing cprefetch ABI submits migration on stream 0, so an actual state tensor carrying is_paged=True passes a null/default stream to the additive optimizer wrapper and preserves prefetch-before-update FIFO ordering plus the existing post-update synchronizations.

Fork issue: heiheiha798#25

Fork engineering record: heiheiha798#26

Scope

The patch does not change optimizer arithmetic, state formats, kernel geometry, public schemas, UVM prefetch behavior, or the legacy direct C ABI. The tracked benchmark loads the legacy symbols from a separately built exact-baseline library and the stream-aware symbols from the exact-head candidate library in one process.

B300 validation

Environment: one NVIDIA B300 SXM6 AC, CC 10.3, 148 SMs; driver 580.126.09; CUDA/nvcc 13.0.88; PyTorch 2.13.0+cu130. Baseline was 95f9af309d4d5793847169c39288dcd3fcbdf564; final exact head was b6530a80713f7ab6da29f7587b6398234bcce515. Isolated builds used CUDA targets 75;80;86;89;90;100;120 and resolved to 75-real;80-real;86-real;89-real;90-real;100-real;120.

• Job 5223: baseline mechanism reproduction. The event was still pending before the direct call; gradient matched the serialized reference while parameter/state1/state2 did not, with max parameter error 0.0009999722. A four-parameter AdamW32bit step made four device-wide synchronizations.
• Job 5224: setup/integration failure only. Both builds succeeded, but the first candidate test exposed missing unmangled extern-C exports for the new helpers. No broad test, sanitizer, or performance evidence from this job is used.
• Job 5225: production-launch validation at parent commit 63d3785f81a1ea59c7505fc2f854e71feec2a65f. Candidate/baseline provenance and legacy/new symbols were asserted. Focused tests: 5 passed. Relevant optimizer matrix: 192 passed, 12 supported skips, 48 deselected. Racecheck: 0 hazards/errors/warnings. Synccheck: 0 errors. Every benchmarked final parameter, state tensor, and gradient was bitwise equal. Its Nsight Systems trace supplies the native launch/synchronization mechanism evidence because the review fix changes only Python routing for actual paged state and tests.
• Job 5238: final exact-head validation. Independent official-compatible baseline/candidate builds and the loaded candidate library were asserted. All 7 focused CUDA tests passed. The new actual-managed-state case ran PagedAdamW32bit from a nondefault caller, observed both state prefetch submissions before the update, and verified the update received stream 0; the event-gated max_unorm > 0 case matched its serialized reference. The unchanged full benchmark matrix remained bitwise equal and passed every performance gate.
• Full pre-commit run --all-files: passed at the final head.

Final exact-head completed wall-time ratios, baseline/current-stream (greater than 1 is faster):

Active parameters	AdamW8/32 FP16/BF16 geomean	Cell range
32	1.441x	1.394x-1.490x
256	1.488x	1.420x-1.567x
1024	1.567x	1.494x-1.638x

The resident 32-layer, 64-trainable-parameter AdamW8 PEFT-style workload improved from 4.3408 ms to 3.9394 ms, or 1.102x. All four final-head single-large controls improved; the smallest ratio was BF16 AdamW8 at 1.013x.

Job 5225 Nsight Systems timed ranges showed the expected unchanged non-paged mechanism with identical optimizer launch counts:

• 32-parameter inventory: legacy 32 cudaDeviceSynchronize calls, candidate 0;
• single-large control: legacy 1, candidate 0;
• four-layer profile topology: legacy 8, candidate 0.

The official-compatible candidate library grew by 23,080 bytes.

Reproduction

After building the baseline and candidate in separate source trees with the target list above, run the tracked CLI against their exact libraries:

python benchmarking/optimizer_current_stream.py \
  --baseline-library /path/to/baseline/bitsandbytes/libbitsandbytes_cuda130.so \
  --expected-library /path/to/candidate/bitsandbytes/libbitsandbytes_cuda130.so \
  --output optimizer-stream.jsonl \
  --warmups 10 \
  --rounds 7 \
  --inventories 32,256,1024 \
  --dtypes fp16,bf16 \
  --bits 8,32

The fork engineering record linked above retains the complete job and raw-evidence inventory.

Limits

Performance and runtime ordering were measured on one B300 and this CUDA/PyTorch stack. The other listed CUDA targets were compile-only. ROCm/hipcc was unavailable on both login and Slurm environments, so shared-source HIP compilation and runtime behavior were not validated here and must be covered by downstream CI. No CUDA graph-capture support is claimed. Fork CI currently provides no substitute for upstream multi-platform CI.