Benchmark stream-ordered lookahead prefetch for CUDA paged optimizers on B300

[heiheiha798]

Summary

This issue tracks a measurement-gated B300 experiment for CUDA paged optimizers. The current optimizer loop executes each active parameter as prefetch_state -> update_step -> device synchronize, while the native cudaMemPrefetchAsync wrapper always uses stream 0. That serializes managed-state migration and optimizer compute even though paged state is independent across parameters.

The hypothesis is that a one-parameter, event-ordered lookahead pipeline can overlap migration of parameter i + 1 with the unchanged optimizer kernel for parameter i, reducing steady-state paged-optimizer step time on B300 without changing optimizer math or the completed-step contract.

This is a hypothesis, not a measured performance claim.

Current behavior

• Optimizer8bit.step() calls prefetch_state(p), update_step(...), and sync_gpu(p) for every active parameter.
• functional.prefetch_tensor() exposes no stream argument.
• cprefetch() invokes cudaMemPrefetchAsync(..., stream=0) for both pre-CUDA-13 and CUDA-13 runtime signatures.
• Optimizer kernels run on PyTorch's current CUDA stream.

As a result, the next parameter's managed state cannot be deliberately prefetched on an independent stream while the current parameter update is running.

Proposed bounded experiment

1. Keep the existing cprefetch symbol and behavior intact. Add an internal stream-aware CUDA prefetch entry point that accepts an explicit raw stream and handles both CUDA <13 and CUDA >=13 signatures.
2. For a single-device CUDA paged optimizer, lazily own one prefetch stream per device. Enqueue state1 and, when present, state2 for the active parameter, then record one event after both requests.
3. Prime the first active paged parameter. Before each update, make the user current stream wait for that parameter's prefetch event. Launch the unchanged optimizer update, then enqueue the next active paged parameter's state on the prefetch stream so migration can overlap current compute.
4. Preserve the final synchronization before step() returns. Preserve the existing sequential path for non-paged optimizers, non-paged state tensors, unsupported concurrent-managed-access devices/backends, one-active-parameter cases, capture-unsafe execution, and multi-device parameter groups not covered by the experiment.
5. Keep lookahead depth exactly one. Any minimum-size or capability guard must be selected from complete B300 evidence rather than assumed.

The implementation must not add multi-tensor optimizer kernels, change optimizer arithmetic or state formats, add eviction policy, alter checkpoint/FSDP behavior, or broaden into non-paged synchronization cleanup.

Expected scope

• bitsandbytes/functional.py: internal stream-aware paged prefetch wrapper while retaining the existing wrapper.
• bitsandbytes/optim/optimizer.py: paged-only one-item lookahead and per-device stream/event ownership.
• csrc/pythonInterface.cpp: additive raw-stream prefetch wrapper with CUDA-version handling and unchanged legacy ABI.
• Focused paged optimizer tests and a reproducible benchmark/profile harness.

No optimizer CUDA kernel or public constructor/API change is expected.

Correctness oracle

• Compare upstream and candidate from identical seeded parameters and gradient sequences for paged Adam/AdamW, Lion, and AdEMAMix, covering 8-bit and 32-bit state where supported, one-state/two-state/double-buffer layouts, FP16/BF16/FP32 parameters, and 5-20 steps. Require bitwise equality of parameters and all state tensors.
• Cover heterogeneous parameter lists, the 100k-element paging boundary, absent/reappearing gradients, one active parameter, and mixed paged/non-paged state.
• Exercise a non-default current stream with gradients produced immediately before step(). Verify event ordering and immediate visibility after step() returns.
• Verify that the event follows all state prefetch requests and that the update stream waits before dereferencing managed state.
• Preserve the existing path for non-paged optimizers, unsupported concurrent-managed-access devices/backends, capture-unsafe execution, and unsupported multi-device groups.
• Run focused paged optimizer tests, relevant functional tests, available CUDA 12/13 builds, compute-sanitizer checks where applicable, and full pre-commit. All CUDA/GPU work runs through Slurm.

B300 benchmark and profiler gate

Use one B300 allocation with isolated upstream/candidate builds in the same software environment. Exclude allocation and first-step initialization, warm up at least 10 steady-state steps, and time at least 50 steps with CUDA events plus host-wall confirmation. Report median and p10/p90.

Primary cases:

• PagedAdamW8bit and PagedAdamW32bit; PagedLion8bit as a one-state control.
• 8-64 heterogeneous parameters with per-state working sets around 2, 8, 32, 64, and 256 MiB.
• Fully resident managed state as a regression control.
• Reproducibly host-resident state established before timing by CPU page touch.
• An oversubscribed working set only when runtime HBM and host-memory limits permit it.
• A Transformers-style heterogeneous parameter distribution as the public integration control.

Nsight Systems must show prefetch of state i + 1 overlapping optimizer kernel i, correct event waits, and removal of intermediate device-wide barriers only in the selected paged path. Record managed-memory HtoD bytes, effective migration bandwidth, overlap fraction, page faults, and peak HBM.

Proceed only if:

• all correctness and ordering checks pass;
• host-resident PagedAdamW 8/32-bit cells show at least 15% geomean steady-state step improvement;
• the public heterogeneous workload retains a repeatable gain;
• no primary paging cell regresses by more than 5%;
• resident-path regression is at most 3%; and
• peak memory and completed-step semantics are unchanged.

If gains come only from removing synchronization without demonstrated migration/compute overlap, only from tiny tensors, or disappear in the public-style workload, record a no-go rather than expand scope.

Risks and non-goals

• Migration may contend with optimizer kernels for HBM or interconnect bandwidth.
• One-step lookahead can evict useful pages under severe oversubscription.
• Streams/events must be per device and long-lived; incorrect device ownership is a correctness bug.
• CUDA graph capture and non-CUDA backends may require retaining the existing path.
• This does not change non-paged synchronization, optimizer kernels/math, state formats, FSDP/checkpoint behavior, multi-device scheduling, or UVM eviction policy.

No open upstream issue or PR was found that adds a raw-stream cudaMemPrefetchAsync wrapper or an event-ordered lookahead scheduler for the existing CUDA paged optimizers.

[heiheiha798]

NO-GO: one-parameter UVM prefetch lookahead misses the 32-bit gate

The fixed one-parameter event-ordered lookahead was evaluated on a B300, but it does not meet this issue's mandatory host-resident gate for both PagedAdamW8bit and PagedAdamW32bit. No production candidate, commit, push, or draft PR is being promoted.

Evidence roles

• Job 5171 is setup-failure evidence only. It passed the B300 identity gate, then stopped before compilation/workload because the ignored wrapper had not created its scratch directory before linking the managed-access probe.
• Job 5174 is the exact-baseline capability and serialization run at 95f9af309d4d5793847169c39288dcd3fcbdf564.
  • NVIDIA B300 SXM6 AC, CC 10.3, 148 SMs, driver 580.126.09, CUDA 13.0.88, PyTorch 2.13.0+cu130.
  • The runtime probe reported concurrent_managed_access=1.
  • The official-compatible build used 75;80;86;89;90;100;120, resolving to 75-real;80-real;86-real;89-real;90-real;100-real;120.
  • The AdamW32 baseline profile recorded 32 cudaMemPrefetchAsync calls, 36 cudaDeviceSynchronize calls, and about 5.09 ms of UVM HtoD activity, confirming the serialized prefetch/update/synchronize mechanism.
• Job 5176 is the authoritative supporting prototype run. It independently extracted the exact baseline, applied only the three-file uncommitted prototype diff (functional.py, optimizer.py, pythonInterface.cpp), built the same official-compatible targets, asserted the loaded candidate library, and interleaved a forced legacy optimizer with the event-ordered candidate in one process. It is supporting evidence rather than commit-addressable final validation.

Correctness and timing

Each cell used identical FP16 parameters/gradients, 10 warmups, 7 rounds x 8 interleaved repetitions, synchronized host-wall and CUDA-event timings, and raw-bit comparison of every parameter and optimizer-state tensor after repeated updates. All reported cells remained bitwise equal.

Optimizer	Layout	Paged state	Legacy median	Candidate median	Legacy / candidate	Round-ratio range
AdamW8	8 x 1,048,576, host-resident	16 MiB	1.703 ms	1.604 ms	1.061x	1.031x-1.070x
AdamW8	8 x 4,194,304, host-resident	64 MiB	3.393 ms	2.823 ms	1.202x	1.191x-1.208x
AdamW32	8 x 1,048,576, host-resident	64 MiB	2.654 ms	2.747 ms	0.966x	0.966x-0.974x
AdamW32	8 x 4,194,304, host-resident	256 MiB	6.766 ms	6.897 ms	0.981x	0.975x-0.994x
Lion8	8 x 4,194,304, host-resident	32 MiB	1.568 ms	1.381 ms	1.135x	1.132x-1.138x
AdEMAMix32	8 x 1,048,576, host-resident	64 MiB	2.531 ms	2.641 ms	0.958x	0.950x-0.964x

The 8-bit large AdamW cell shows that the overlap mechanism can work, but the required 32-bit result is a stable regression across every round. The resident AdamW32 candidate update was about 0.397 ms versus about 6.897 ms for the host-resident candidate step, so a depth-one pipeline cannot hide the migration-dominated 32-bit time sufficiently to reach the required 15% host-resident gain. Increasing tensor size does not repair that ratio; deeper lookahead, an 8-bit-only policy, or a different migration strategy would change the frozen design or acceptance scope.

The required public Transformers-style topology, oversubscribed matrix, and broader validation were therefore intentionally skipped once the mandatory direct AdamW8/32 gate had already failed. This is a performance NO-GO for the selected depth-one design, not a correctness failure.

Reproduction and raw evidence

Baseline wrapper:

sbatch benchmark-evidence/cycle18/slurm_baseline_pilot.sh

Candidate wrapper and tracked benchmark invocation are recorded in:

sbatch benchmark-evidence/cycle18/slurm_candidate_pilot.sh

Persistent evidence:

• benchmark-evidence/cycle18/slurm-5171-baseline-pilot.out
• benchmark-evidence/cycle18/slurm-5174-baseline-pilot.out
• benchmark-evidence/cycle18/paged-prefetch-baseline-job-5174.jsonl
• benchmark-evidence/cycle18/nsys-paged-prefetch-baseline-job-5174.nsys-rep
• benchmark-evidence/cycle18/nsys-paged-prefetch-baseline-job-5174-stats.csv
• benchmark-evidence/cycle18/slurm-5176-candidate-pilot.out
• benchmark-evidence/cycle18/paged-prefetch-candidate-job-5176.jsonl
• benchmark-evidence/cycle18/nsys-paged-prefetch-candidate-job-5176.nsys-rep
• benchmark-evidence/cycle18/nsys-paged-prefetch-candidate-job-5176-stats.csv

salvage_hint=NONE: the required 32-bit mechanism does not provide threshold-level value, and an 8-bit-only or greater-depth follow-up would change the frozen acceptance/scope without complete public-path evidence.