Calibrate SM103 quantile-map replicas in the Adam 8-bit optimizer #9

Problem

The CUDA two-state 8-bit optimizer kernel copies each 256-entry quantile map
into two padded shared-memory rows:

smem_quantiles[2][257]
lane_id = threadIdx.x % 2

ADAM then repeatedly reads these maps while dequantizing both optimizer
states, running the data-dependent quantize_2D search, and checking the
quantized sign. The replica count and 257-float skew have been unchanged since
the original 2021 kernel. A different replica count may alter shared-memory
broadcast/bank-conflict behavior on B300, but more copies also increase map
initialization and shared-memory use. This issue is a measurement-gated
calibration, not a claim that more or fewer replicas are faster.

Scope

• Baseline: bitsandbytes-foundation/bitsandbytes@95f9af309d4d5793847169c39288dcd3fcbdf564.
• Build isolated experimental variants with QMAP_COPIES={1,2,4,8} from the
  same source, compiler, flags, and allocation. Copy count 2 is the baseline.
• Parameterize only the private CUDA
  kOptimizerStatic8bit2StateBlockwise quantile-map copy count. Keep its
  256-thread CTA, 256-element block, one value per thread, grid geometry, CUB
  operations, reductions, optimizer arithmetic, quantile values, packed state
  format, and public interface unchanged.
• If and only if a single copy count passes the calibration, correctness, and
  held-out performance gates below, select it only for exact runtime compute
  capability 10.3 and compile-time optimizer ADAM.
• ADEMAMIX, all one-state optimizers, HIP, and every non-SM103 device must keep
  the exact two-copy path. Use a narrow cached per-current-device runtime check
  and include its overhead in eager timing; do not introduce a general
  architecture-dispatch framework.
• Use official CUDA 13 x64 compatible target input
  75;80;86;89;90;100;120 as the primary deployed-code path. Report native
  103 builds only as separately labeled corroboration.
• If upstream PR Fix decoupled weight decay ordering in the CUDA Adam/AdEMAMix kernels bitsandbytes-foundation/bitsandbytes#2040 lands first, rebase and preserve its weight-decay
  ordering without absorbing that correctness change.

Expected production changes are limited to the private declaration,
definition, selected ADAM instantiations, narrow host launch choice, and
focused tests in csrc/kernels.cuh, csrc/kernels.cu, csrc/ops.cu, and the
nearest optimizer tests. Experimental forcing may remain in validation
scaffolding but must not become a public API.

Correctness oracle

Generate valid deterministic Adam8bit states with a fixed baseline warm-up,
then clone parameter, gradient, state1, state2, absmax1, absmax2, and both
quantile maps before invoking each forced variant.

• Require raw-bit equality of parameter, state1, state2, absmax1, and absmax2
  after one and repeated updates. Require the gradient and quantile maps to
  remain raw-bit unchanged. Raw-byte/integer-view comparison is required for
  NaN controls because torch.equal reports matching NaNs as unequal.
• Cover FP16, BF16, and FP32 parameter/gradient storage; steps 1, 2, 10, and
  1000; fresh-zero and warmed/clumped states; random finite, zero-heavy,
  mixed-magnitude, NaN, and Inf inputs; weight decay 0/nonzero; and
  skip_zeros false/true.
• Cover n=255,256,257,4095,4096,4097,37887,37888,37889,64K,1M and realistic
  large aligned tensors.
• Dispatch controls must prove ADEMAMIX, MOMENTUM/RMSPROP/ADAGRAD/LION,
  non-SM103 devices, and HIP retain copy count 2. The known schedule-sensitive
  ADEMAMIX unaligned-tail result is not part of the candidate oracle; it must
  never enter the candidate path.
• Run focused public Adam8bit/AdamW8bit reference coverage, relevant
  tests/test_optim.py, the official-compatible CUDA build, and full
  pre-commit. All CUDA compilation and GPU execution must run through Slurm.

Fixed calibration and holdout plan

Use one B300/SM103 allocation and interleave all isolated variants. Record the
exact source identity, node/GPU, driver, compiler/toolkit, PyTorch and
bitsandbytes versions, resolved fatbin targets, commands/environment,
warmups/iterations, compiler resource output, build-time/library-size deltas,
raw JSONL/log paths, every sample, median, and p10/p90.

Stage A selects at most one candidate before any held-out result is inspected:

• Compare copies 1, 4, and 8 against baseline 2 using FP16 at
  n=37888, 3*37888, 1M, and 16M, with both fresh and warmed states. This is
  a fixed eight-cell calibration set.
• For each copy count, compute its eight baseline/candidate median ratios.
  It is eligible only if the median of those ratios exceeds 1.02, at least six
  of eight cells favor it, no calibration cell regresses by more than 5%, and
  interleaved dispersion is below the observed median effect.
• Choose the eligible copy count with the highest median ratio; break a tie by
  choosing the smaller shared-memory footprint. If none is eligible, record
  NO-GO. Do not choose by dtype, size, state distribution, profiler result, or
  any Stage-B measurement.

Stage B validates only the frozen Stage-A winner on disjoint held-out cells:

• BF16 and FP32, plus FP16 controls; n=4096,64K,4M,64M or the largest safe
  size, 4096*4096, and 8192*28672; steps 1, 10, and 1000; realistic warmed
  states. Use direct preallocated
  torch.ops.bitsandbytes.optimizer_update_8bit_blockwise calls, at least 20
  warmups, and at least seven interleaved CUDA-event batches with size-scaled
  iterations. Report elements/s and clearly defined effective bytes/s.
• If available in the same allocation, profile one predeclared held-out cell
  with Nsight Compute for shared-load conflict/replay and occupancy evidence.
  Profiling may explain timing but cannot select the copy count or replace it.
• Measure eager Adam8bit and AdamW8bit optimizer.step() over a fixed,
  reproducibly recorded transformer-like resident parameter collection with
  quantized tensors above min_8bit_size=4096 and unchanged small 32-bit
  parameters. Preallocate states and gradients and time only steps. Paged
  controls may be reported separately but are not acceptance evidence.

Proceed only if the frozen winner is bit-exact, improves at least two realistic
held-out direct cells by at least 5% with dispersion below the gain, improves
one representative eager step repeatably by at least 3%, and causes no
greater-than-5% regression in any adjacent held-out dtype/shape. The same
single copy count must pass the official-compatible build and public-path
tests. Native-only, Stage-A-only, direct-only, or profiler-only gains are
NO-GO and must not produce a draft PR.

Non-goals

No ADEMAMIX change; no one-state optimizer change; no persistent grid or CTA
size change; no known tail-layout fix; no bias-correction change; no optimizer
math, quantization algorithm, state format, public API, or CUDA wheel-target
change; no other architecture/backend generalization; and no dependency on
prior fork branches.

Same-address broadcast may already make two copies optimal, and a different
copy count can increase initialization work or lower occupancy. Those are
expected negative outcomes for this bounded experiment, not reasons to widen
its scope.