Two GPU memory-safety defects in bitsandbytes' 4-bit path

Summary. bitsandbytes (the library behind 4-bit and 8-bit quantization in Hugging Face Transformers and PEFT) has two buffer-size defects on its 4-bit path. quantize_4bit allocates its packed output 1 to 7 bytes too small whenever quant_storage is wider than uint8, so the CUDA kernel writes past the end of it. dequantize_4bit never checks that the packed tensor is large enough for the shape stored in its QuantState, so a corrupt or crafted state makes the kernel read far past the buffer. Both were reproduced on a real GPU, confirmed with NVIDIA's compute-sanitizer, fixed with regression tests that fail before the fix and pass after, and filed upstream. Neither is merged. Status checked live on 2026-10-25: PR #2105, issue #2106 and PR #2107 are all open, with no maintainer reply.

Neither is a security vulnerability by the project's own definition. Section 4 explains why, with measurements.

Environment. RTX 5060 (sm_120), CUDA 12.8, torch 2.11.0+cu128, bitsandbytes 0.50.0 binary, Windows.

1. Background

4-bit quantization stores two 4-bit values per byte. quantize_4bit returns a packed tensor and a QuantState (shape, dtype, block size, per-block absmax, code table). dequantize_4bit reverses it. The packed tensor's storage dtype is set by quant_storage. The default is uint8, but Transformers exposes it as bnb_4bit_quant_storage and documents bfloat16 for FSDP/QLoRA training, so wider storage is a real configuration, not a corner case.

2. Defect 1: quantize_4bit under-allocates (PR #2105)

Mechanism. bitsandbytes/backends/cuda/ops.py:392 sizes the packed buffer as:

out = torch.empty(((n + 1) // (quant_storage.itemsize * 2), 1), ...)

That is a floor division. The buffer must hold ceil(n/2) bytes, which is ceil(ceil(n/2) / itemsize) storage elements. For uint8 (itemsize 1) the two expressions are identical, so the default path is always correct and the bug stays hidden. For wider dtypes the buffer is short whenever n is not a multiple of 2 × itemsize. The same expression is in the register_fake meta implementation (_ops.py:235), so the fake and the real kernel agree with each other and both disagree with what the kernel writes. One line above, the same function already uses the correct ceiling idiom for blocks.

Evidence. - Zero-element case: quantize_4bit on a 2-element tensor with quant_storage=float16 raises CUDA error: an illegal memory access and poisons the CUDA context. - Larger sizes: PyTorch's caching allocator rounds blocks up, so the overrun lands in slack and compute-sanitizer reports nothing. With PYTORCH_NO_CUDA_MEMORY_CACHING=1 it reports Invalid __global__ write of size 1 bytes in kQuantizeBlockwiseSmall at numel=2050, bfloat16 storage.

Fix. Use the ceiling idiom in both places. For uint8 it is numerically identical to the old expression, so the default path is unchanged. After the fix the sanitizer reports 0 errors (513 elements, 1026 bytes allocated for 1025 required).

Verification. - New test test_quantize_4bit_packed_buffer_holds_all_bytes over {uint8, float16, bfloat16, int32} × {2, 6, 12, 2050}. Before the fix: uint8 passes, the first wide-storage CUDA case fails, and the process aborts on the illegal access. After the fix: 132 passed, 60 skipped. - Regression check by failure-ID diff, with and without the patch: tests/test_ops.py 0 failures both ways. tests/test_linear4bit.py has 116 failures both ways, with identical sets. I did not diagnose those 116 individually; they appear to come from running main's tests against the 0.50.0 binary and CPU torch.compile here, and they are unchanged by the patch. - The repo's own pre-commit suite passes.

3. Defect 2: dequantize_4bit trusts quant_state.shape (issue #2106, PR #2107)

Mechanism. dequantize_4bit (functional.py) validates block size, quant type and dtype, but not that the packed tensor holds prod(shape) values. QuantState.from_dict and Params4bit.from_prequantized, the documented loader path for prequantized models, copy shape and absmax verbatim from the stored state with no consistency check. The kernel then reads prod(shape) values from a buffer that may be far smaller.

Reproduction through public APIs. A 64-element tensor is honestly quantized (32 packed bytes). A forged QuantState claims shape=(4096, 4096) (16,777,216 elements) with absmax and quant_map sized for the claim, so the packed data is the only undersized buffer. Params4bit.from_prequantized accepts the mismatch, and dequantize_4bit returns a 4096×4096 tensor from 32 bytes.

compute-sanitizer:

Invalid __global__ read of size 1 bytes
  at kDequantizeBlockwise<__half, 512, 64, 8, 2>
  Access ... is out of bounds and is 225 bytes after the nearest allocation of size 32 bytes

Fix. Raise ValueError if A.numel() * A.element_size() * 2 < prod(shape) (capacity in 4-bit slots, correct for any quant_storage width), and also if absmax has fewer entries than ceil(prod(shape) / blocksize).

Verification. After the guard the sanitizer reports 0 errors and the call raises a clear error. New test test_dequantize_4bit_rejects_shape_larger_than_packed: 4 fail without the guard, and with it the surrounding cases pass (40 passed). The 4-bit op suite is identical with and without the patch (637 passed, 57 skipped, 60 xfailed). The repo's pre-commit suite passes.

4. Severity, measured rather than assumed

I tested whether either defect escalates beyond a correctness bug. - Overrun length (Defect 1). The overrun is capped at itemsize − 1 bytes, at most 7 for int64, and does not grow with numel. I checked this exhaustively for numel 1 to 5000 and on about 200,000 random values up to 50 million. A bounded 7-byte overrun cannot be groomed into a larger overflow. - Neighbouring data (Defect 1). A canary tensor placed after the under-allocated buffer was never modified: 0 of 800 trials, under both the default allocator and expandable_segments:True. - Reachability. numel and quant_storage are application configuration set at model preparation, not attacker input at inference. - Defect 2 is a read, not a write. Nothing is corrupted, and the output buffer is correctly sized. Worst cases are a crash or adjacent GPU memory landing in the dequantized weight.

Disclosure channel. I first planned a private security advisory for Defect 2. The project's own agents/security_guide.md (Tier 7) classifies GPU-side out-of-bounds reads and writes as "correctness and reliability issues rather than traditional security exploits", because CUDA kernels cannot make syscalls or exfiltrate. An advisory would have overstated the issue by the maintainers' own standard, so I filed a public issue and PR.

5. Process, and mistakes caught

• Method. Equivalence oracles (does the output match an independent reference?) plus compute-sanitizer, rather than static pattern matching.
• A harness bug. My first relative-error metric subtracted a (2,1) tensor from a (1,2) one, which broadcast to (2,2) and made a correct result look broken. Shape-checking the harness exposed a real, separate shape issue in dequantize_4bit, which led to the allocation bug underneath it.
• A retracted observation. I noticed a tail value differing by 3.14 at numel=2050, which looked like data corruption. The reference buffer was allocated adjacent to the test buffer, so the difference was not isolated from allocator reuse. Given the 0-of-800 canary result, I dropped it as a corruption claim.
• Duplicates caught before filing. Two earlier leads (Lightning's batch-size scaler; accelerate's get_balanced_memory with a "disk" key) were already open PRs (#22055 and #4417), so nothing was filed.

6. What I did not verify

• Fixes were tested against the installed 0.50.0 binary. Both fixes are pure Python and the relevant lines match main, but I did not compile main's CUDA kernels.
• No real checkpoint load. I demonstrated Defect 2 by building the QuantState dictionary directly through the public API. I did not craft a safetensors file and load it with from_pretrained.
• No cross-process leakage. The out-of-bounds read exposes memory within the process; I did not demonstrate reading another process's data.
• The main inference path is unchecked. gemm_4bit and gemv_4bit, used by Linear4bit at inference, take the shape as a separate argument and do not pass through the new guard. I did not test whether they are affected, and said so in a comment on PR #2107.
• Noted but not fixed. A backward-compatibility shim in dequantize_4bit transposes the output when the packed tensor has one row (functional.py:1075), and the CPU backend raises instead of overwriting on these sizes. Both are documented for the maintainers.

7. Process disclosure

The investigation, fixes and verification were done with an AI assistant (Claude Code), the same disclosure as in the pull requests: reviewed before submitting, not merged, and no maintainer has responded.

Reproduce

Scripts are in repro/: repro_bnb_underalloc.py, repro_dequant_oob_read.py, severity_canary.py, isolate_case.py. They need a CUDA GPU and bitsandbytes 0.50.0. Sanitizer runs need PYTORCH_NO_CUDA_MEMORY_CACHING=1.

Source and scripts: tritsystem/bitsandbytes-4bit-memory-safety. Written 2026-10-25; PR and issue status was re-checked that day and may have changed since.