Investigate a persistent SM103 grid for blockwise 8-bit optimizers

[heiheiha798]

Problem

The CUDA one-state and two-state blockwise 8-bit optimizer kernels already use grid-stride loops, but the host launch currently creates one 256-thread CTA for every 256 parameters and each kernel sets its loop bound from gridDim.x * 256. As a result, each CTA executes exactly one state block.

Before the loop, every CTA copies one or two fixed 256-entry quantization maps into shared memory, initializes CUB storage, and synchronizes. For large optimizer tensors this setup and CTA scheduling are repeated thousands of times per step. On a 148-SM B300, capping the launch to a measured SM-wave multiple could let each CTA process several independent state blocks while reusing its shared quantization maps.

This is a measurement-gated hypothesis, not an assumed performance win. A persistent grid can reduce load balancing or latency hiding, and the maps may already be cache-resident. No source dispatch change is justified unless direct and end-to-end B300 measurements clear the acceptance gate below.

Scope

• Baseline: bitsandbytes-foundation/bitsandbytes@95f9af309d4d5793847169c39288dcd3fcbdf564.
• Compare the current uncapped launch with a small forced matrix of exact-SM103 grid caps using isolated builds in the same B300 allocation.
• Keep forced cap choices in benchmark/build scaffolding. Do not expose a public forcing API.
• Only if evidence identifies a coherent winner, cap the exact-SM103 launch grid and make the existing grid-stride loops cover the full n; retain the current uncapped launch on every other CUDA architecture and on HIP.
• Use an existing cached device-property mechanism or equally narrow per-device cache. Do not add an uncached device query to every optimizer step and do not introduce a broad dispatch framework.
• Preserve all optimizer arithmetic and current public/Python behavior. 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 kernel file; this issue must preserve whichever ordering is present after any required rebase and must not absorb that correctness fix.

The intended production surface is limited to the one-state and two-state blockwise 8-bit CUDA optimizer launch/loop bounds plus focused validation. No paging, FSDP, optimizer refactor, or state-format change belongs here.

Correctness oracle

For deterministic inputs, compare forced uncapped and capped variants after one and multiple steps.

• Cover AdamW8bit/Adam8bit and AdEMAMix8bit on the two-state path, plus Lion8bit, Momentum8bit, RMSprop8bit, and Adagrad8bit on the one-state path where supported.
• Cover FP16, BF16, and FP32 parameters/gradients; weight_decay=0 as the primary launch-only comparison; one nonzero-decay control that preserves upstream ordering; and skip_zeros true/false where existing support permits.
• Require byte-identical packed state buffers and bitwise-identical absmax arrays. Require bitwise-identical parameter storage between forced variants when the only difference is CTA-to-block assignment; retain existing public-reference tolerances as a secondary oracle.
• Include n below one block, at 255/256/257, immediately below/at/above the selected cap * 256, across multiple grid-stride rounds, and with a final partial block. Include realistic parameter counts.
• Run focused tests/test_optim.py coverage, relevant public optimizer-step regressions, the CUDA build, and the full pre-commit suite. All CUDA compilation and GPU tests run only through Slurm.

Do not use this work to change existing skip_zeros, FSDP, paging, or weight-decay semantics.

B300 benchmark and acceptance

Use one B300/SM103 allocation and record exact source/build identity, job/node/GPU, driver, CUDA compiler, PyTorch/bitsandbytes versions, commands, environment, and raw logs/results. Baseline and forced variants must use the same compiler/flags and be interleaved in the same allocation.

• Compare current uncapped launch with 1x, 3x, 6x, and 12x SM caps; narrow around a winner only if the first matrix supports it.
• Measure direct blockwise 8-bit optimizer updates with preallocated state for 4K, 64K, 1M, 16M, and 64M or the largest memory-safe size, plus representative counts such as 4096x4096 and 8192x28672.
• Cover FP16/BF16/FP32 and both one-state and two-state kernels.
• Measure end-to-end optimizer.step() for AdamW8bit, AdEMAMix8bit, and Lion8bit on representative large tensors. Exclude allocation/state initialization from the timed region.
• Use at least 20 warmups where practical and at least 7 interleaved batches per cell. Choose enough CUDA-event iterations to make each batch stable, scaling the count down for very large tensors and end-to-end steps rather than forcing 100 expensive calls everywhere. Report complete/round-level samples, median, p10/p90, elements/s, and an explicitly defined effective-bandwidth estimate; never report best-only results.

Proceed only if one cap produces at least a 5% median gain across at least two large realistic shapes, the gain is larger than interleaved dispersion and remains visible in end-to-end optimizer.step(), correctness is bit-exact between launch variants, and no adjacent large shape regresses by more than 5%. A gain confined to the direct kernel, a cherry-picked size, or small launch-bound tensors is a no-go.

Non-goals

No optimizer math change; no weight-decay-ordering fix; no skip_zeros, paging, FSDP, or state-format change; no 32-bit optimizer tuning; no General8bit/4-bit quantization work; no public API; no other architecture/backend generalization; no dependency on prior fork branches or issues.

[heiheiha798]

B300 calibration result: NO-GO

The persistent-grid hypothesis does not meet the required correctness gate on the independent baseline 95f9af309d4d5793847169c39288dcd3fcbdf564. No source change, commit, push, or draft PR was produced.

Environment and method

• GPU: NVIDIA B300 SXM6 AC, compute capability 10.3, 148 SMs
• Node: dev-slurm
• Driver: 580.126.09
• CUDA toolkit: 13.0.88
• PyTorch: 2.13.0+cu130
• Python: 3.12.3
• Native build target: SM103 (-DCOMPUTE_CAPABILITY=103)
• Isolated variants: upstream uncapped launch and 148/444/888/1776-CTA caps (1x/3x/6x/12x SM)
• Planned timing method: 20 warmups followed by seven interleaved batches with size-scaled CUDA-event samples. Timing was intentionally not reached because exactness ran first and failed.

Job 4876: CTA-reassignment mismatch

Job 4876 passed the SM103 identity gate, built all five isolated variants, and completed 59 exactness cells. It then failed for AdEMAMix FP16 at n=37,889 (148*256+1), step 1: the 148-CTA cap changed absmax1 relative to the uncapped launch.

The cause is the existing AdEMAMix tail layout in kOptimizerStatic8bit2StateBlockwise: the appended state's scale uses absmax1[(n+i)/256]. When n is not divisible by 256, that index range overlaps the first state's tail scale. Reassigning multiple logical blocks to one CTA changes the ordering of writes to the overlapping location.

Raw output: benchmark-evidence/cycle4/slurm-4876-optimizer-grid.out

Job 4878: same-upstream-path nondeterminism

The approved replacement retained the exact uncapped upstream entrypoint for every non-256-aligned control and exercised caps only on aligned counts. It again passed the SM103 identity gate and built all five isolated variants.

Before the failure, it recorded 78 successful exactness cells:

• 75 Adam cells: FP16, BF16, and FP32 across small/tail, aligned cap-plus-one, multiple grid-stride, weight-decay, and skip_zeros controls
• 3 initial AdEMAMix FP16 cells

It then failed for AdEMAMix FP16 at n=37,887 (148*256-1), step 2: absmax1 differed between buffers even though each nominal variant called the same uncapped upstream function. This demonstrates that the overlapping tail-scale writes are already schedule-nondeterministic; it is not merely a mismatch introduced by the cap.

Raw output: benchmark-evidence/cycle4/slurm-4878-optimizer-grid.out
Partial JSONL: benchmark-evidence/cycle4/sm103-optimizer-grid-job-4878.jsonl

Decision

The accepted scope requires byte-identical packed states, bitwise-identical absmax arrays, and bitwise-identical parameters for Adam/AdamW and AdEMAMix two-state paths. AdEMAMix cannot satisfy that oracle at required tail counts on the current upstream implementation, even when repeating the same uncapped path.

Repairing the appended-state tail indexing would change optimizer state/absmax semantics and belongs to a separate correctness issue. That work is explicitly outside this issue's launch-only scope, as are dropping AdEMAMix coverage and weakening the bit-exact oracle. The performance gate was therefore not evaluated, and there is no B300 speedup claim to submit upstream from this cycle.