Investigate a 256-thread SM103 path for small 32-bit optimizer tensors #8

Problem

The CUDA 32-bit two-state optimizer path always launches
kOptimizer32bit2State with 1,024 threads and four items per thread. For a
tensor with at most 1,024 elements, only up to 256 threads own valid values,
but all 32 warps still participate in the kernel's CUB load/store phases and
CTA barriers.

This path matters beyond explicitly selected 32-bit optimizers. The default
min_8bit_size=4096 routes small parameters in 8-bit Adam/AdamW/AdEMAMix to
the 32-bit update kernel. Models with many bias, normalization, or adapter
parameters can therefore launch this oversized one-CTA path repeatedly.

The measurement-gated hypothesis is that an SM103-only 256-thread
specialization, still processing four items per thread, can reduce small
parameter update latency enough to remain visible in a model-like optimizer
step. No performance improvement is assumed until the B300 results pass the
gates below.

Scope

• Baseline: bitsandbytes-foundation/bitsandbytes@95f9af309d4d5793847169c39288dcd3fcbdf564.
• Build isolated 1,024-thread baseline and 256-thread candidate variants from
  the same source, toolchain, and flags.
• Parameterize only the private CUDA two-state update kernel's thread count.
  Keep four items per thread, with a 4,096-element baseline tile and a
  1,024-element candidate tile.
• Select the candidate only for exact runtime compute capability 10.3,
  n <= 1024, max_unorm == 0, and ADAM or ADEMAMIX. Retain the existing
  path for n > 1024, max-unorm/LAMB use, every other architecture, and HIP.
• Keep device detection cached per current device and safe for CUDA graph
  capture. Include its host dispatch overhead in performance measurements.
• Use the official CUDA 13 x64 compatible target input
  75;80;86;89;90;100;120 as primary deployed-code evidence. Report a native
  103 build only as a separately labeled corroborating control.
• If open upstream PR Fix decoupled weight decay ordering in the CUDA Adam/AdEMAMix kernels bitsandbytes-foundation/bitsandbytes#2040 lands first, rebase and preserve its
  Adam/AdEMAMix weight-decay ordering without absorbing that correctness fix.

Expected production changes are limited to csrc/kernels.cuh,
csrc/kernels.cu, the narrow launch selection in csrc/ops.cu, and focused
tests. Experimental forcing may live in validation scaffolding but must not
become a public API.

Correctness oracle

For cloned deterministic inputs, require bitwise equality (torch.equal) of
parameter, state1, and state2 between the 1,024-thread and 256-thread variants
after one and repeated updates. The gradient must remain unchanged. Do not
weaken this to approximate equality: the selected path has no cross-CTA
reduction, and changing CTA width must not change optimizer arithmetic.

• Cover ADAM/AdamW and ADEMAMIX; FP16, BF16, and FP32 parameter/gradient
  storage; steps 1, 2, 10, and 1000; zero and nonzero weight decay; and
  skip_zeros false/true where supported.
• Cover zero, random finite, and mixed-magnitude gradients at
  n=1,2,31,32,33,127,128,129,255,256,257,511,512,513,1023,1024.
• Use n=1025,4095,4096,4097, non-SM103 dispatch mocks, and max_unorm > 0
  controls to prove that all non-selected cases retain the existing launch.
• Run the repository optimizer references and focused tests/test_optim.py
  coverage, the official-compatible CUDA build, and full pre-commit. All CUDA
  compilation and GPU execution must run through Slurm.

B300 benchmark and acceptance

Use one B300/SM103 allocation and interleave isolated baseline and candidate
builds. Record the exact source identity, node/GPU, driver, CUDA compiler,
PyTorch/bitsandbytes versions, fatbin target list, commands, environment, raw
JSONL/log paths, warmups, iterations, and every round sample. Report median and
p10/p90 rather than best-only timings.

• Direct preallocated torch.ops.bitsandbytes.optimizer_update_32bit cells:
  ADAM and ADEMAMIX; FP16/BF16/FP32; n=1,32,64,128,256,512,768,1023,1024;
  and 1025/2048/4096 boundary controls. Use at least 20 warmups and at least
  seven interleaved timed batches with enough CUDA-event iterations for stable
  dispersion.
• Eager optimizer.step() collections: fixed, reproducibly recorded
  bias/norm/LoRA-like shapes in the 64-1,024 range mixed with unchanged
  4,096+ tensors. Cover AdamW32bit, AdamW8bit with the default
  min_8bit_size=4096, and AdEMAMix32bit. Preallocate state and gradients and
  time only optimizer steps.
• Include a graph-compatible direct-call collection as a diagnostic control,
  but use eager end-to-end timing for the acceptance decision.
• Record compile-time and native-library size deltas from the additional
  private instantiations.

Proceed only if the bit-exact oracle passes, at least three representative
direct cells in n=128..1024 improve by at least 8% with dispersion below the
gain, and one realistic mixed eager optimizer step improves repeatably by at
least 3%. No boundary, large, or non-selected cell may regress by more than 3%,
and cached host dispatch overhead must not erase the gain. Otherwise record a
no-go without opening a draft PR.

Non-goals

No precondition or reduction-kernel change; no bias-correction broadcast; no
optimizer-math, state-layout, or public API change; no 8-bit state-kernel
change; no CTA/grid tuning outside the 256-thread two-state ADAM/ADEMAMIX
specialization; no CUDA wheel target change; no other architecture/backend
generalization; and no dependency on prior fork branches.

The extra template instantiations increase compile time and binary size, and
launch or Python overhead may dominate small updates. Those are expected
negative outcomes for the experiment, not reasons to widen its scope.