Triton 4-bit quantization stores absmax past the end of the tensor unless the block count is a multiple of 8

[truong-v]

System Info

Affects the Triton backend (bitsandbytes/backends/triton/), which is registered for XPU (Intel GPU) and used as the fallback when no native library is available. Both quantize_fp4_blockwise_kernel and quantize_nf4_blockwise_kernel in kernels_4bit.py are affected.

The grid is sized in block pairs, but each program unconditionally stores a whole group of 8 absmax entries:

# bitsandbytes/backends/triton/ops.py:85-87, quantize_4bit
blocks = -(n // -(blocksize * 2))                  # number of block PAIRS
absmax = torch.empty((blocks * 2,), ...)           # one entry per block

# bitsandbytes/backends/triton/kernels_4bit.py:158-160, quantize_4bit_blockwise_triton
split_num_blocks = 4
grid = (triton.cdiv(blocks, split_num_blocks),)

# bitsandbytes/backends/triton/kernels_4bit.py:109 (and :42 for fp4)
# PAIRED_SPLIT_NUM_BLOCKS = SPLIT_NUM_BLOCKS * 2 = 8, block_start_idx = program_id(0) * 8
tl.store(absmax_ptr + block_start_idx + tl.arange(0, PAIRED_SPLIT_NUM_BLOCKS), absmax)   # no mask

So the grid writes 8 * ceil(blocks / 4) entries into a tensor that holds 2 * blocks — an overrun of 8 - 2 * (blocks % 4) entries (6, 4 or 2) whenever blocks % 4 != 0. In terms of the block count nb = numel / blocksize, the store is in bounds only when nb % 8 is 0 or 7. torch.randn(1000, 1000) at blocksize=64 gives nb = 15625 and writes 6 entries past the end. Weight shapes that are powers of two land on nb % 8 == 0, the in-bounds case, which is why the test suite never sees this.

Every other memory operation in these kernels is masked (the input load at :102, the packed store at :154); this store is not. The same unmasked store is in quantize_8bit_blockwise_kernel (kernels_8bit_quant.py:103), where it stays in bounds today only because split_num_blocks is 1 there.

The kernels were added in #1692.

Reproduction

absmax is allocated inside quantize_4bit, so the overrun normally lands in allocator padding. Pointing absmax at a view of a larger buffer makes it visible:

import torch
from bitsandbytes.backends.triton.kernels_4bit import quantize_4bit_blockwise_triton

device, blocksize, n = "xpu", 64, 1000 * 1000   # 15625 blocks; absmax holds 15626
blocks = -(n // -(blocksize * 2))

A = torch.randn(n, device=device)
buffer = torch.full((blocks * 2 + 8,), -1.0, device=device)   # absmax values are non-negative
absmax = buffer[: blocks * 2]
out = torch.empty((n - n // 2, 1), device=device, dtype=torch.uint8)

quantize_4bit_blockwise_triton(A, blocksize, "nf4", blocks, absmax, num_elements=n, quantized_out=out)
print("entries written past the end of absmax:", int((buffer[blocks * 2 :] != -1.0).sum()))

entries written past the end of absmax: 6

Confirmed independently with compute-sanitizer on the unmodified public entry point (backends.triton.ops.quantize_4bit(A, 64, "nf4", torch.uint8) on the same input, PYTORCH_NO_CUDA_MEMORY_CACHING=1 so each tensor is its own allocation):

========= Invalid __global__ write of size 4 bytes
=========     at quantize_nf4_blockwise_kernel+0xba0 in kernels_4bit.py:109
=========     by thread (2,0,0) in block (1953,0,0)
=========     and is 1 bytes after the nearest allocation at 0x7f2b36a00000 of size 62504 bytes
[... threads 3-7 of the same block, 5/9/13/17/21 bytes after the same allocation ...]
=========         Host Frame: quantize_4bit_blockwise_triton in kernels_4bit.py:171
=========         Host Frame: quantize_4bit in ops.py:92
========= ERROR SUMMARY: 10 errors

Six invalid writes from the last program of the grid (block 1953 of 1954), then the context dies with unspecified launch failure.

Expected behavior

The kernels should write only the blocks * 2 entries that absmax holds; the store should be masked like every other memory operation in the file. Today the overrun happens to land in PyTorch's 512-byte allocation padding, so no user data is corrupted — but it is an out-of-bounds write, it is fatal under compute-sanitizer, and it becomes a real corruption under any allocator that packs tensors tightly.

Environment

• bitsandbytes main at 2b6cfb7, installed with pip install -e .
• NVIDIA B200 (sm_100), driver 595.71.05 (no XPU on hand; the Triton kernels and the index arithmetic are device-agnostic)
• torch 2.13.0+cu130, triton 3.7.1, Python 3.12
• compute-sanitizer 2025.4.1 (CUDA 13.1)