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

[heiheiha798]

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.

[heiheiha798]

NO-GO: the public AdEMAMix path fails the bit-exact oracle

I ran the measurement-gated experiment on one B300, but the 256-thread specialization does not meet the issue's correctness and performance gates. I am recording the result without making a production change or opening a draft PR.

Environment and method

• Slurm job: 4924 on dev-slurm
• Source: clean 95f9af309d4d5793847169c39288dcd3fcbdf564 on b300/sm103-32bit-optimizer-small-cta
• GPU: NVIDIA B300 SXM6 AC, compute capability 10.3, 148 SMs
• Driver: 580.126.09
• CUDA compiler: 13.0.88
• Python: 3.12.3
• PyTorch: 2.13.0+cu130
• Primary compatible build input: 75;80;86;89;90;100;120
• Separate native control build input: 103
• Four isolated source/build trees: compatible/native x 1,024-thread baseline/256-thread candidate
• Direct timings: 20 warmups, 7 interleaved rounds, 200 CUDA-event samples per variant per round
• Raw local evidence: benchmark-evidence/cycle6/slurm-4924-opt32-small-cta.out and benchmark-evidence/cycle6/sm103-opt32-small-cta-job-4924.jsonl

The job completed 240 boundary correctness records, 24 nonzero-weight-decay/skip_zeros controls, four max_unorm controls, and 88 direct performance cells before reaching the mixed eager workloads. All completed forced baseline-versus-candidate direct comparisons were bit-exact for parameter, state1, and state2, and left gradients unchanged.

Positive but insufficient performance evidence

The official-compatible build showed a coherent ADEMAMIX direct-kernel benefit for several selected cells:

dtype	n	baseline / candidate	gain
BF16	128	1.108x	10.8%
BF16	512	1.114x	11.4%
BF16	768	1.141x	14.1%
BF16	1024	1.156x	15.6%
FP32	512	1.128x	12.8%
FP32	1023	1.187x	18.7%
FP32	1024	1.145x	14.5%

The 1025/2048/4096 non-selected controls stayed within 2.1%. The fixed mixed eager collections for AdamW32bit and AdamW8bit both remained exact and improved by 4.56% and 4.55%, respectively, over seven rounds.

The additional instantiations increased the official-compatible library by 106,128 bytes and its measured build by 1.88 seconds. The native library increased by 24,200 bytes and its measured build by 0.40 seconds.

Correctness failure and decision

The fixed AdEMAMix32bit eager collection failed on its first update. The first parameter (n=64, FP16, weight_decay=0.01) differed between the 1,024-thread baseline and 256-thread candidate under torch.equal. The harness stopped immediately at:

AssertionError: initialization/ademamix32: candidate parameter 0 differs

This is a required public-path, first-update bit-exactness failure. Job 4924 is therefore an incomplete calibration after an oracle failure, not a passing benchmark: it did not produce the AdEMAMix eager timing, graph diagnostic, completion record, or PASS marker.

An ADAM-only narrowing does not rescue the proposal. Its selected direct cells ranged roughly from a 1% loss to a 4% gain, so it does not supply the required three representative n=128..1024 cells with at least 8% improvement. Keeping ADEMAMIX would weaken the mandatory bit-exact oracle; removing it would miss the direct performance gate.

Result: NO-GO for this cycle. No tracked source was changed, no commit was created, no branch was pushed, and no draft PR was opened. The issue remains open as a record of the B300 evidence.