[tritsystem]

What / why

Fixes #2106.

dequantize_4bit trusts quant_state.shape (and the absmax length) to decide how many 4-bit values to read from the packed tensor, but never checks the packed tensor is large enough. A QuantState whose shape claims more elements than the packed weight holds — e.g. reconstructed from a corrupt or crafted checkpoint via Params4bit.from_prequantized / QuantState.from_dict — makes kDequantizeBlockwise read past the end of the packed buffer.

compute-sanitizer (with PYTORCH_NO_CUDA_MEMORY_CACHING=1):

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

Per agents/security_guide.md Tier 7, this is framed as a correctness/reliability issue (out-of-bounds read → wrong results / possible crash), not a traditional exploit.

Fix

Add a size check in dequantize_4bit: if the packed tensor can't hold prod(quant_state.shape) 4-bit values (or absmax is shorter than the implied block count), raise a clear ValueError instead of reading out of bounds. The capacity is computed as A.numel() * A.element_size() * 2, which is correct regardless of quant_storage width. The honest round-trip path is unchanged.

Test

tests/test_ops.py::Test4bitBlockwiseQuantOps::test_dequantize_4bit_rejects_shape_larger_than_packed builds a 32-byte packed tensor with a QuantState claiming (4096, 4096) and asserts a ValueError is raised.

• Without the fix: the test fails (on CUDA the kernel reads out of bounds).
• With the fix: passes; the existing 4-bit op/round-trip tests in test_ops.py are unaffected (verified against the installed 0.50.0 binary — pure-Python change, so no rebuild needed).

The repo's pre-commit suite passes on both changed files.

---

I used Claude Code as my assistant to investigate, fix, and verify this issue; I reviewed the analysis and the fix before submitting.

[tritsystem]

A scope note for reviewers: this PR adds the packed-size check in dequantize_4bit, which also covers the deprecated [K, N] branch of matmul_4bit, since that branch routes through dequantize_4bit. The main inference entry points, gemm_4bit and gemv_4bit, take quant_state.shape as a separate argument and don't go through this check. I haven't verified whether those paths are affected, so I've kept this PR narrow.

If you'd prefer, the check could live where the shape first arrives from serialized data (QuantState.from_dict / Params4bit.from_prequantized), so every consumer is covered in one place. Happy to rework the PR that way, or leave it as is. Your call on where validation belongs.