Fuse LLM.int8 correction operand preparation into one CUDA kernel

[heiheiha798]

Summary

The nonempty LLM.int8 mixed-decomposition path currently prepares correction operands through a chain of eager tensor operations:

subA = A[:, outlier_cols].contiguous()
subB = (
    int8_vectorwise_dequant(CB[:, outlier_cols].contiguous(), SCB)
    .to(A.dtype)
    .t()
)

For current CUDA FP16 use, that means an activation gather, an INT8 weight gather, FP32 row-scale multiplication, FP32 reciprocal-127 multiplication, FP32-to-FP16 conversion, and final layout preparation before the correction addmm. The main INT8 GEMM then runs unchanged.

This issue proposes a measurement-gated experiment: produce the final contiguous subA[M,L] and correction-GEMM subB[L,N] in one private native CUDA kernel, where L = outlier_cols.numel(). Preserve the existing main int8_scaled_mm and output.addmm(subA, subB) operations.

No performance improvement is assumed in advance.

Why this matters on B300

The preparation chain executes for every threshold-enabled Linear8bitLt layer with nonempty outlier columns. It adds serialized launches and transient tensors before a fast main GEMM. At larger N/L, weight gathering and FP32 transforms also add material bandwidth and allocation pressure; at small L, launch latency is the likely target.

The path is public and recurrent in LLM.int8 inference. Upstream issue bitsandbytes-foundation#1867 reports substantial throughput and energy cost in mixed decomposition and explicitly calls for optimizing this pathway. B300/SM103 will determine whether fusing this concrete preparation stage produces a meaningful layer and resident-stack gain.

Bounded method

• Add one private CUDA preparation kernel plus the smallest launcher/C binding needed by the CUDA backend.
• Allocate only the final subA[M,L] FP16 tensor and final subB[L,N] FP16 tensor.
• In one launch, cover two simple linear output regions:
  • gather A[row, outlier_cols[l]] directly into contiguous subA[row,l];
  • gather CB[n,outlier_cols[l]], reproduce the current FP32 dequantization sequence, convert to FP16, and write directly into the transposed subB[l,n] layout consumed by addmm.
• Integrate through a CUDA implementation of the existing bitsandbytes::int8_mixed_scaled_mm operation or a private helper used only by that implementation. Do not add or change a public schema.
• Keep int8_scaled_mm and correction addmm as the existing operations. Do not fuse either GEMM.

The kernel must use a simple bounded grid-stride mapping over the two output regions. Do not introduce a persistent scheduler, shape/architecture dispatch table, multiple production variants, or a general indexed-transform framework.

Exact non-goals

• Do not change thresholding, outlier detection, masks, argwhere, outlier_cols, or the empty/no-outlier path.
• Do not change kInt8VectorQuant, INT8 formats, row/column statistics, the main INT8 GEMM, correction GEMM, bias semantics, autograd, or public API.
• Do not depend on fork issues Evaluate a tiled native detector for LLM.int8 outlier columns #14/Retry tiled LLM.int8 outlier detection with an empty-mask fast path #15 or any earlier branch.
• Keep current upstream FP16 semantics. Do not copy or stack BF16 work from open upstream PR Add BF16 support for int8_vectorwise_quant / LLM.int8 activation quant bitsandbytes-foundation/bitsandbytes#1985.
• Do not absorb the FakeTensor/output-metadata fix from open upstream PR Fix mixed INT8 FakeTensor output metadata bitsandbytes-foundation/bitsandbytes#2044. On the current baseline, runtime work must leave default/fake registrations unchanged. If Fix mixed INT8 FakeTensor output metadata bitsandbytes-foundation/bitsandbytes#2044 lands first, rebase and conform to its contract.

Likely files are the private CUDA kernel/launcher/interface files, bitsandbytes/backends/cuda/ops.py, and focused tests. A tracked benchmark is justified only when it is needed for exact reproducibility.

Exact correctness oracle

Compare isolated baseline and candidate implementations from the same upstream commit.

For the preparation outputs:

• Require raw-bit equality for subA.
• Require raw-bit FP16 equality for subB. Reproduce the current operation order exactly: FP32 CB * SCB, then a separately rounded FP32 multiplication by the existing reciprocal-127 constant, then FP16 conversion. There is no tolerance fallback for this preparation optimization; failure to match bits is a no-go.
• Cover the complete INT8 code domain, zero/normal/subnormal-adjacent/large finite SCB, L={1,2,5,16,64}, first/last/noncontiguous/sorted columns, and repeated indices if accepted by the public operation.

For the existing public operation:

• Require exact subA and baseline-equivalent output over M/N/K boundaries and representative sizes, bias absent/present, threshold-generated outlier patterns, 2-D and 3-D Linear8bitLt, repeated calls, and resident state reuse.
• Keep outlier_cols is None, empty outlier_cols, threshold-zero public paths, and returned metadata behavior unchanged. These are veto controls and must not invoke the fused preparation kernel.
• Run focused op, autograd, module, and Linear8bitLt tests, plus opcheck and applicable compile/dynamic-shape regression coverage.

PR bitsandbytes-foundation#2044 documents a pre-existing FakeTensor metadata defect on the current baseline. Do not repair or duplicate it here. Validation should prove this candidate does not worsen baseline behavior; if bitsandbytes-foundation#2044 merges, rebase and run its complete metadata/opcheck matrix.

Build the official CUDA target set, preserve clean CUDA/HIP registration boundaries, and run full pre-commit. If HIP runtime/tooling is unavailable, disclose that limitation and at least retain shared-source compilation compatibility.

B300 benchmark plan

Use one B300 allocation with separately extracted baseline/candidate builds from the same clean commit and official CUDA 13 x64-compatible target input 75;80;86;89;90;100;120. Assert the actual loaded libraries and private entry point before timing.

Measure both preparation-only and full int8_mixed_scaled_mm for:

• M={1,8,64,512,2048};
• representative K/N={3584,4096,8192,14336};
• L={0,1,5,16,64} plus an approximately 1% control where memory permits.

Use fixed resident inputs, at least 200 warmups, and at least 30 interleaved baseline/candidate samples. Report CUDA-event kernel timing and synchronized host-wall timing, medians and dispersion, not best runs.

Add public threshold-6 Linear8bitLt decode, batched-decode, and prefill cases plus reproducibly defined resident 4-layer and 32-layer stacks. Use generated local tensors only; exclude allocation, weight setup, and weight quantization from steady-state timing.

Profile representative small-L and bandwidth-heavy cells to prove that the eager gather/multiply/cast/layout chain is replaced by one kernel, transient allocations/bytes fall, and strided CB reads do not erase the benefit. Record occupancy, bandwidth, source-read efficiency, final-write efficiency, launches, and allocations.

Proceed to a draft PR only if all correctness/build gates pass and:

• preparation improves by at least 1.5x in representative small- and moderate-L cells;
• full mixed operation improves by at least 10% in several representative cells;
• public Linear8bitLt improves by at least 5% in decode and at least one batched/prefill case;
• the 32-layer resident stack improves by at least 3% in both decode and prefill;
• no representative public, empty/no-outlier, or threshold-zero cell regresses by more than 2%.

Otherwise record a no-go. Do not rescue an isolated microbenchmark win with shape specialization, extra variants, main-GEMM changes, or broader fusion.

Risks

• Indexed CB reads are strided across rows and may make the fused kernel transaction-inefficient.
• Activation and weight output regions differ greatly in size; a poor combined mapping can waste occupancy.
• Small L is launch-sensitive while large L is bandwidth-sensitive, so one simple kernel may not win broadly enough.
• FP32 operation order must be made explicit enough to preserve baseline FP16 bits.
• Open PRs Fix mixed INT8 FakeTensor output metadata bitsandbytes-foundation/bitsandbytes#2044 and Add BF16 support for int8_vectorwise_quant / LLM.int8 activation quant bitsandbytes-foundation/bitsandbytes#1985 create nearby rebase risk even though neither duplicates this runtime preparation fusion.

If the gates pass, this is a contained native system optimization of a documented public LLM.int8 bottleneck with no API, format, detection, or GEMM change.

[heiheiha798]

Final result: NO-GO

The fused preparation kernel is correct and materially faster for nonempty outlier sets on B300, but the candidate does not meet this issue's mandatory no-regression gate for the empty-outlier path. I am therefore not promoting the local candidate commits or creating a draft PR.

Authoritative final-head validation

• Slurm job: 5080 (COMPLETED, exit 0:0, runtime 00:02:18)
• Commit: 5f0bb8f754ef738f004bf04e61d5cce415654a2e, clean local branch b300/int8-mixed-correction-prep
• Node/device: dev-slurm, NVIDIA B300 SXM6 AC, CC 10.3, 148 SMs
• Software: PyTorch 2.13.0+cu130, CUDA 13.0; the log records driver, NVCC, CMake, Ninja, and Python versions
• Build input: official CUDA 13 x64-compatible targets 75;80;86;89;90;100;120
• The benchmark/test process asserted that it loaded the isolated exact-head library from /tmp/.../source-candidate/bitsandbytes/libbitsandbytes_cuda130.so and that the private entrypoint was present
• HIP/ROCm tooling was unavailable, so no HIP compile or runtime claim is made

The exact command recorded in the raw log for the tracked benchmark was:

PYTHONPATH=/tmp/bnb-int8-mixed-final-5080/source-candidate \
  python benchmarking/int8_mixed_scaled_mm_prepare.py \
  --output benchmark-evidence/cycle13/sm103-int8-mixed-final-job-5080.jsonl \
  --full-grid --warmup 200 --samples 30 --prep-repetitions 10

Focused exact-head results:

• tests/test_ops.py::TestLLMInt8Ops: 23 passed
• Selected Linear8bitLt module tests: 4 passed
• test_linear8bitlt_torch_compile: 14 passed, 2 expected xfails
• 125/125 preparation records were raw-bit equal
• All 150 mixed-operation, 3 Linear8bitLt, and 4 synthetic-stack records passed their output oracle

Performance result

The intended nonempty path clears its positive gates:

• Preparation: all 125 cells were at least 1.5x faster; minimum 2.336x, median 3.213x
• Full nonempty mixed op: 80/125 cells were at least 1.10x faster, with no cell more than 2% slower
• Linear8bitLt: decode 1.165x, batched 1.115x, prefill 1.056x
• 32-layer resident stack: decode 1.229x, prefill 1.140x

However, the required empty-outlier veto fails even after the bounded direct-return fix:

• 9/25 L=0 controls were below the allowed 1/1.02 ratio
• Worst ratio: 0.9088x at M=1, K=3584, N=4096
• Median across the 25 empty controls: 0.9881x

The local follow-up moved the empty return before all native eligibility checks, but it did not remove the measured small-shape overhead. Another route would require moving dispatch into a different high-level/default-backend surface or relaxing the explicit veto, both outside the bounded cycle. No further retry or scope expansion is justified.

Supporting evidence and limitations

Job 5070 tested the prior exact commit 80d5620dc79ef79415a78bcab1ab883039ea0b87. It remains supporting evidence only because it predates the final empty-path commit. It supplied 29 focused functional passes, 384 autograd passes, and exact production-kernel NCU evidence. The NCU profile is reusable because 5f0bb8f changes only Python empty-path control flow and leaves the CUDA kernel unchanged; it was intentionally not rerun in job 5080.

Raw workspace evidence:

• benchmark-evidence/cycle13/slurm-5080-int8-mixed-final.out
• benchmark-evidence/cycle13/sm103-int8-mixed-final-job-5080.jsonl
• benchmark-evidence/cycle13/slurm-5070-int8-mixed-candidate.out
• benchmark-evidence/cycle13/sm103-int8-mixed-candidate-job-5070.jsonl
• benchmark-evidence/cycle13/sm103-int8-mixed-candidate-job-5070.ncu.csv

Conclusion: the nonempty fusion mechanism is valid and fast, but the complete public-path acceptance contract is not met. No branch push or draft PR was created. salvage_hint=NONE: the single bounded empty-path fix was exhausted, and the remaining alternatives would expand the implementation surface or weaken the gate.