Evaluate a tiled native detector for LLM.int8 outlier columns

[heiheiha798]

Summary

The CUDA LLM.int8 threshold path currently discovers outlier columns with an eager PyTorch pipeline before launching the native vector quantizer:

1. materialize A.abs() for the full activation tensor;
2. materialize the full >= threshold boolean mask;
3. reduce it with outliers.any() and convert that result to a Python boolean;
4. for nonempty inputs, reduce the mask again with outliers.any(dim=0) and call argwhere.

The native vector quantizer then rereads A to compute row statistics and quantized output. For batched/prefill inputs, outlier discovery therefore adds several full-tensor passes, two O(rows * K) temporary tensors, multiple launches, and a host-visible synchronization before quantization.

This issue proposes a measurement-gated experiment: replace only that eager discovery pipeline with a private tiled native detector that writes an O(K) int32 column mask directly from FP16 A, then retain the existing argwhere step and public return contract.

Why B300

The path runs for every positive-threshold Linear8bitLt layer. It matters for batch-1 decode, where launch and synchronization overheads are exposed, and for batched/prefill shapes, where the eager pipeline scans and materializes the full activation mask. The repository already benchmarks threshold-enabled M=1 LLM.int8, and upstream issue bitsandbytes-foundation#1867 reports substantial user-visible cost in the mixed LLM.int8 path.

B300/SM103 will be used to determine whether eliminating these full-tensor intermediates and reductions produces a meaningful public-path benefit. No speedup is assumed in advance.

Bounded implementation

• Add one private FP16 CUDA detector in csrc/kernels.cu / csrc/kernels.cuh.
• Use a 2-D launch where 256 threads own contiguous columns and each CTA scans a bounded row tile, initially 32 rows.
• Each thread sets its column's int32 flag to 1 with an idempotent atomic only when an outlier is observed in its row tile.
• Zero the O(K) mask once, invoke the detector, and use the existing argwhere behavior to produce sorted CUDA int64 column indices.
• Add the smallest private launcher/C entry point needed by the CUDA backend; do not add a public torch.library schema.
• If the initial row tile misses the performance gate, at most compare fixed row tiles 16, 32, and 64. Do not add architecture tables or shape-specific production dispatch.

Keep these paths unchanged:

• threshold == 0 behavior;
• the existing kInt8VectorQuant kernel and quantization math;
• row statistics, output-column zeroing, mixed matmul, and public return metadata;
• default/non-CUDA backend behavior;
• op schemas and autograd;
• all prior fork branches and unmerged work.

Current upstream accepts FP16 in this CUDA path. Do not copy or stack the unmerged BF16 changes from upstream PR bitsandbytes-foundation#1985. If that PR lands before this work is ready, rebase and reassess compatibility explicitly rather than silently broadening the experiment.

Correctness gates

Compare isolated baseline and candidate builds from the same upstream commit. Require exact equality for quantized rows, row statistics, and ordered outlier-column indices, including empty tensor dtype/device/rank.

Cover:

• rows around the row tile and realistic prefill sizes: 1, 2, 15, 16, 17, 31, 32, 33, 63, 64, 65, 128, and 2048;
• K around the column tile plus 3584, 4096, 8192, and 14336;
• thresholds 0.1, 2.0, 3.0, and 6.0;
• no, one, repeated-column, sparse, and dense outliers;
• positive/negative values exactly at threshold and immediately below it;
• first/last rows and columns, NaN, and Inf according to current PyTorch comparison semantics.

threshold == 0 must still return outlier_cols is None. Scalar threshold conversion must match the current A.abs() >= threshold behavior exactly. Run the focused vectorwise/COO tests, threshold-enabled Linear8bitLt tests, and applicable fullgraph/dynamic-output compile tests. Build the supported CUDA targets and preserve HIP/ROCm source compatibility; B300 provides the performance decision, not permission to break other builds.

B300 measurement plan

Use one B300 allocation and isolated baseline/candidate builds with identical software, target inputs, and inputs. Record the commit, loaded native-library paths, GPU/driver/toolkit/PyTorch identity, commands, raw samples, medians, and dispersion.

Benchmark the public CUDA quantizer with host wall time because dynamic output discovery is synchronization-sensitive. Interleave baseline and candidate over FP16 rows {1, 8, 64, 512, 2048}, K {3584, 4096, 8192, 14336}, and no/one/five/1%-dense outlier patterns. Include dense controls. Use at least 200 warmups and 30 interleaved samples with enough calls for stable samples.

Profile representative decode (1, 3584) and prefill (2048, 8192) cells to verify that the candidate actually removes the O(rows * K) abs/bool intermediates and Python-bool reduction, and to measure launches, synchronization, temporary allocation, bandwidth, and atomic behavior.

For required public-path evidence, benchmark threshold-enabled Linear8bitLt and a reproducible resident synthetic stack of such layers for both decode and prefill/batched shapes. Preallocate weights, activations, and states; exclude setup/weight quantization; record the exact stack definition. The repository's model-based int8-decomp benchmark is optional corroboration only when its artifacts are already available; this experiment must not depend on downloads.

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

• profiler evidence confirms the intended intermediate/synchronization removal;
• isolated threshold quantization improves by at least 10% in at least four representative decode/prefill cells, with no primary cell regressing more than 3%;
• threshold-enabled Linear8bitLt improves by at least 5% in at least one decode and one batched/prefill shape, with no primary cell regressing more than 2%;
• the reproducible synthetic public-path stack improves by at least 3% in at least two workload regimes, with no regime regressing more than 2%.

Otherwise record a no-go rather than adding fusion, public API changes, architecture-specific dispatch, or a broader INT8 redesign.

Risks

• argwhere remains dynamic and may still synchronize; this experiment removes preceding redundant work rather than claiming an asynchronous path.
• No-outlier and small-row inputs may lose to mask initialization plus the detector, so they are veto cells.
• Small row tiles can increase CTA/atomic overhead; large tiles can underutilize the GPU. Tuning is deliberately bounded.
• Dense or repeated outliers may contend on the idempotent column flags and must be measured.
• Upstream PR Add BF16 support for int8_vectorwise_quant / LLM.int8 activation quant bitsandbytes-foundation/bitsandbytes#1985 touches nearby native and Python INT8 code. This issue is semantically independent but has real rebase risk.

If the gates pass, the result is a contained native-system optimization with no public API or quantization-format change and direct relevance to a documented LLM.int8 user bottleneck.

[heiheiha798]

B300 calibration result: NO-GO for the fixed tile-32 candidate

I calibrated the proposed private FP16 detector on one NVIDIA B300 SXM6 AC (SM103, 148 SMs) in Slurm job 5044. The worktree was clean at the issue baseline 95f9af309d4d5793847169c39288dcd3fcbdf564; no production source was edited or committed.

Reproduction and provenance

• Job/node: 5044 / dev-slurm
• Software: Python 3.12.3, PyTorch 2.13.0+cu130, CUDA 13.0, driver 580.126.09
• Builds: separately extracted baseline and candidate source trees, both built with 75;80;86;89;90;100;120; build times were 60.03 s and 59.30 s.
• Candidate: 256 columns per CTA, 32-row tiles, and an int32[K] mask updated with idempotent atomicExch.
• Loaded-library provenance was asserted against the temporary candidate tree. Baseline and candidate cint8_vector_quant entry points were distinct.
• Timing: 200 warmups and 30 interleaved host-wall samples per cell, with a post-call CUDA synchronization so dynamic-output synchronization was included.
• Command:

python benchmark-evidence/cycle11/sm103_int8_outlier_mask.py \
  --output benchmark-evidence/cycle11/sm103-int8-outlier-mask-job-5044.jsonl \
  --build-root "$SLURM_TMPDIR/bnb-int8-outlier-mask-5044/benchmark-build" \
  --compute-capability '75;80;86;89;90;100;120' \
  --warmup 200 \
  --samples 30

Raw evidence:

• benchmark-evidence/cycle11/slurm-5044-int8-outlier-mask.out
• benchmark-evidence/cycle11/sm103-int8-outlier-mask-job-5044.jsonl
• benchmark-evidence/cycle11/sm103-int8-outlier-mask-job-5044-decode.ncu.csv
• benchmark-evidence/cycle11/sm103-int8-outlier-mask-job-5044-prefill.ncu.csv

Correctness and mechanism

All 320 exact comparison records passed. Baseline and candidate had torch.equal outputs for packed rows, row statistics, and ordered CUDA int64 outlier columns, and the FP16 inputs remained raw-bit unchanged. Coverage included row/tile boundaries, realistic K values through 14336, thresholds 0.1/2/3/6, empty/one/five/sparse/dense patterns, exact-threshold neighbors, NaN/Inf, and the actual candidate public op.

Nsight Compute successfully profiled the detector in both required regimes. The decode launch was a (14,1,1) grid and took 6.336 us; prefill was (32,64,1), took 45.248 us, reached 741.8 GB/s measured memory throughput, and achieved 85.29% occupancy. The prefill torch-profiler peak allocation fell from 67,125,760 to 16,842,752 bytes. For decode, however, the fixed int32[K] mask is slightly larger than the eager FP16-abs/bool temporaries (20,480 versus 14,848 observed peak bytes), which is part of the small-shape tradeoff rather than a hidden benefit.

Performance

The core tiled scan clearly works when outliers exist:

• All 80 one/five/sparse/dense direct cells favored the candidate by at least 23.6%; their median ratios by pattern were 1.460x, 1.473x, 1.462x, and 1.431x respectively.
• Five-outlier Linear8bitLt controls improved from 0.2192 to 0.1805 ms for decode (1.215x), 0.3052 to 0.2599 ms for the batched case (1.174x), and 0.5459 to 0.5001 ms for prefill (1.092x).
• The four-layer resident decode stack improved from 0.4973 to 0.4690 ms (1.060x).

The complete acceptance gate nevertheless failed:

• The no-outlier pattern had a median ratio of 0.925x. Fifteen of 20 cells regressed by more than 3%; the worst was rows=64, K=8192 at 0.802x (0.0570 versus 0.0711 ms).
• The four-layer resident prefill stack was flat/slightly slower at 0.9965x (1.0965 versus 1.1004 ms), so only one of the two required synthetic-stack regimes cleared 3%.

The regression is localized and reproducible: the eager baseline stops after its Python-bool any() when there are no outliers, while the fixed candidate must still execute dynamic argwhere(mask) to preserve the public return contract. Changing the detector row tile to 16 or 64 cannot remove that empty-argwhere host/dynamic-output overhead, so I did not spend another allocation on a width change that cannot address the failed gate. A dtype/shape/architecture dispatch would also not predict this data-dependent condition and is outside the issue scope.

Disposition

The fixed tile-32 implementation does not satisfy the no-regression and two-regime resident-stack gates. No production candidate, commit, push, or draft PR was created. The issue remains open as a measured NO-GO record rather than presenting the outlier-only wins as a general speedup.