Mask the absmax store in the Triton 4-bit quantize kernels

[truong-v]

The Triton 4-bit quantize grid is sized in block pairs (grid = cdiv(blocks, 4)), but each program unconditionally stores a whole group of PAIRED_SPLIT_NUM_BLOCKS = 8 absmax entries, so the grid writes 8 * ceil(blocks / 4) entries into an absmax tensor that holds 2 * blocks. Whenever blocks % 4 != 0 the last program writes 6, 4 or 2 entries past the end — torch.randn(1000, 1000) at blocksize=64 overruns by 6 float32.

Fixes #2042

Changes

• quantize_nf4_blockwise_kernel / quantize_fp4_blockwise_kernel: hoist the absmax offsets into a variable and mask the store with absmax_offsets < absmax_elements, matching the input load and the packed store, which are already masked.
• quantize_4bit_blockwise_triton: pass absmax_elements=absmax.numel(). Taking the length from the tensor rather than recomputing it keeps the mask correct for a caller-supplied absmax and preserves the current behaviour for the one padding entry the host allocates when the block count is odd (it is still written, as a zero).

Output is unchanged: over 24 configurations (nf4/fp4 x fp32/fp16/bf16 x four shapes, including block counts that do and do not overrun) the packed bytes and the returned absmax are bit-for-bit identical before and after.

Two more copies of the same unmasked store are left alone, both deliberately. quantize_8bit_blockwise_kernel (kernels_8bit_quant.py:103) is in bounds only because split_num_blocks is 1 there. quantize_4bit_blockwise_kernel, further down this same file, has the identical statement and the identical arithmetic, but it has no caller anywhere in the repo and its @triton.autotune decorator is commented out, so it is unreachable today. Happy to guard either or both if you would rather have them covered.

Performance

triton.testing.do_bench on the full backends.triton.ops.quantize_4bit call, before/after interleaved over three rounds on one B200:

case	before (µs)	after (µs)
nf4, 4096x4096 bf16, blocksize 64	48.63 – 48.97	49.37 – 49.54
fp4, 4096x4096 bf16, blocksize 64	40.38 – 40.70	39.86 – 39.92
nf4, 4096x4096 bf16, blocksize 256	39.17 – 39.54	38.55 – 38.62
fp4, 4096x4096 bf16, blocksize 256	31.71 – 31.79	33.07 – 33.12
nf4/fp4, 1000x1000 fp32, blocksize 64	8.20 – 8.21	8.20 – 8.21

Within ±1.4 µs, two configurations faster and two slower — the mask is one compare and one select over 8 lanes, outside the loop. For scale, the same binary measured 49.5 µs and 62.2 µs for the first row in two different rounds of process launches, so these differences sit inside the run-to-run spread.

Tests

The shipped cases land on the two in-bounds block counts. test_quantize_4bit uses 1024x1024, whose block count is a multiple of 8 at every blocksize tested; test_quantize_4bit_not_divisible_by_blocksize uses (7, blocksize - 1), which pads to exactly 7 blocks. And where a shape does overrun, the extra entries land in allocator padding, so an assertion on returned values cannot see it.

test_quantize_4bit_triton_absmax_stays_in_bounds calls the Triton launcher with absmax pointing into a larger buffer, so entries written past its end are visible, and asserts both directions: nothing is written past the end, and nothing inside is left unwritten. 129 and 130 blocks cover odd and even block counts; 24 cases in total (2 quant types x 2 blocksizes x 3 dtypes x 2 block counts). It is skipped where the Triton backend is not used. Reverting the kernel change alone fails all 24.

tests/test_ops.py on this machine: 657 passed / 57 skipped / 60 xfailed at 2b6cfb7, and 681 passed / 57 skipped / 60 xfailed with this PR — the 24 new cases, no other change. pre-commit run --all-files is clean.

Environment

• bitsandbytes main at 2b6cfb7, installed with pip install -e .
• NVIDIA B200 (sm_100), driver 595.71.05 — no Intel GPU on hand, so the kernels were exercised through bitsandbytes.backends.triton on CUDA
• torch 2.13.0+cu130, triton 3.7.1, Python 3.12

[matthewdouglas]

Thanks for the PR. In general this looks sound.

I would ask to remove the regression test. I don't think it adds any value. It's not useful at all to run it on NVIDIA/AMD hardware since that hardware won't ever hit this path. Besides that, it's unlikely anyone's going to change this kernel in the future in such a way that reintroduces the issue. So in my view, it seems like bloat.

Separate from this issue if you want to introduce further cleanup PRs for the other kernels, I'm open to looking at that. quantize_4bit_blockwise_kernel could probably just be removed entirely if we have no internal callers.

[github-actions]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[truong-v]

Thanks. I removed the test and force-pushed. The PR is now just the kernel change

[matthewdouglas]

LGTM, thanks!