Investigate SM103 bias-correction broadcast in 32-bit optimizers #7

Problem

The CUDA 32-bit two-state optimizer kernels recompute CTA-invariant bias-correction scalars in every thread.

• kOptimizer32bit2State launches 1,024 threads per CTA and each thread evaluates the same powf(beta1, step), powf(beta2, step), sqrtf, and derived step size before processing a 4,096-element tile.
• The ADAM instantiation of kPreconditionOptimizer32bit2State, used when max_unorm > 0, launches 512 threads and each thread evaluates the same two reciprocal powf corrections.

All operands are CTA-uniform. Computing each correction once in thread 0, publishing the float bit patterns through shared memory, and consuming them after an already-required first CTA barrier could remove redundant special-function work on B300. However, the compiler or scheduler may already hide that cost, and shared publication can change resource use. This issue is a measurement-gated hypothesis, not a performance claim.

Scope

• Baseline: bitsandbytes-foundation/bitsandbytes@95f9af309d4d5793847169c39288dcd3fcbdf564.
• Build isolated baseline and scalar-broadcast variants from the same source/toolchain. The broadcast variant changes only the bias-correction scalars in kOptimizer32bit2State and the two-state ADAM precondition kernel.
• Reuse the kernels' existing first CTA barrier to publish the shared scalars if safe; do not add another synchronization merely for this optimization.
• If and only if B300 evidence clears the gates below, select private baseline/broadcast kernel variants through a cached per-device runtime check for exact compute capability 10.3. Every non-SM103 device and HIP must keep the current expression path.
• Validate the official CUDA 13 x64 compatible target input 75;80;86;89;90;100;120 as the primary deployed-code path, and report native SM103 builds separately. Do not rely solely on __CUDA_ARCH__ == 1030: official wheels do not currently include an SM103 image.
• Keep the existing C/Python API and optimizer arithmetic unchanged. Open upstream PR Fix decoupled weight decay ordering in the CUDA Adam/AdEMAMix kernels bitsandbytes-foundation/bitsandbytes#2040 changes Adam/AdEMAMix weight-decay ordering in the same source file; rebase and preserve that ordering if it lands, but do not absorb the correctness fix here.

The expected production surface is limited to csrc/kernels.cu, the narrow host launch selection in csrc/ops.cu, and focused validation/benchmark coverage. No broad architecture-dispatch helper is required.

Correctness oracle

For identical deterministic inputs, require bitwise-identical parameter, state1, state2, and unorm buffers between baseline and broadcast variants after one and multiple steps. Thread 0 must evaluate the same float expressions and shared publication must preserve their bits; approximate equality is not an acceptable replacement for this launch-only optimization.

• Cover Adam32bit/AdamW32bit and AdEMAMix32bit with max_unorm=0, plus LAMB32bit or the corresponding two-state ADAM path with max_unorm>0 to exercise both precondition and update kernels.
• Cover FP16, BF16, and FP32 parameter/gradient storage where registered; steps 1, 2, 10, and 1000; weight_decay=0 plus one nonzero control; and skip_zeros false/true where existing semantics support it.
• Include n=1, 255/256/257, 4095/4096/4097, 8191/8192/8193, 64K, 1M, and realistic parameter counts. Include zero, random finite, and mixed large/small finite gradients.
• Keep the repository's public optimizer-reference tolerances as a secondary oracle and run the focused tests/test_optim.py matrix, CUDA build, and full pre-commit suite. All CUDA compilation and GPU tests run only through Slurm.

Preserve whichever weight-decay ordering is present after rebasing. Do not change max-unorm, skip_zeros, NaN/Inf, paging, or state-layout semantics.

B300 benchmark and acceptance

Use one B300/SM103 allocation and record the exact commit/build identity, node/GPU, driver, CUDA compiler, PyTorch/bitsandbytes versions, target list, commands, environment, and raw JSONL/log paths. Interleave baseline and broadcast variants built with identical flags.

• Use the official-compatible target list as the primary matrix so the runtime SM103 selection is exercised in the same form as distributed CUDA 13 x64 binaries. Use a native 103 build as a separately labeled corroborating control, not as a substitute.
• Measure direct torch.ops.bitsandbytes.optimizer_update_32bit calls with preallocated buffers at n=256, 4096, 64K, 1M, 16M, 64M or the largest memory-safe size, 4096x4096, and 8192x28672; FP16/BF16/FP32; Adam, AdEMAMix, and a max_unorm/LAMB control; steps 1, 10, and 1000.
• Measure end-to-end optimizer.step() for AdamW32bit, AdEMAMix32bit, and LAMB32bit on representative large tensors and a model-like collection containing many sub-4096 parameters. Exclude allocation and state construction from timing.
• Use at least 20 warmups where practical and at least 7 interleaved batches per cell. Scale CUDA-event iterations by problem size so each batch is stable. Report all round samples, median, p10/p90, elements/s, and a clearly defined effective-bandwidth estimate; never report best-only values.

Proceed only if the bit-exact oracle passes, the official-compatible broadcast path produces at least a 5% median direct-op gain on at least two realistic Adam/AdEMAMix shapes with dispersion below the gain, and at least one representative end-to-end optimizer step improves repeatably by 3% or more without a greater-than-5% regression in any adjacent large-shape/dtype cell. If only the native build or LAMB benefits, the compiler already eliminates the cost, or gains disappear end to end, record a no-go.

Non-goals

No optimizer-math or weight-decay-ordering change; no broadcast of unrelated invariant expressions; no CTA-size/grid tuning; no persistent optimizer grid; no 8-bit optimizer or state-layout work; no max-unorm/skip_zeros semantic fix; no public API; no CUDA wheel target change; no other architecture/backend generalization; no dependency on prior fork cycles.