Skip to content

Validate storage/shape block count in QTensor::new - #3716

Open
v-code01 wants to merge 1 commit into
huggingface:mainfrom
v-code01:qtensor-validate-storage-shape
Open

v-code01 wants to merge 1 commit into
huggingface:mainfrom
v-code01:qtensor-validate-storage-shape

Conversation

@v-code01

@v-code01 v-code01 commented Jul 5, 2026

Copy link
Copy Markdown

What

QTensor::new runs only check_shape, which verifies the last dimension is divisible by the block size but never checks that the storage actually holds as many blocks as the shape requires:

pub fn new<S: Into<Shape>>(storage: QStorage, shape: S) -> Result<Self> {
    let shape = shape.into();
    check_shape(&shape, storage.block_size())?;   // last-dim % block_size only
    Ok(Self { storage, shape, repacked_qs: OnceLock::new() })
}

So a mismatched pair — e.g. a one-block Q4_0 storage (32 elements) with shape (64,) (which needs two blocks) — is accepted, and then dequantize sizes its output from the shape and reads the storage past its end:

  • legacy quants (Q4_0/Q4_1/Q5_0/Q5_1/Q8_0): out-of-bounds index → panic;
  • k-quants: silently truncated / zero-padded tensor (wrong result, no panic).

QTensor::new and QStorage are public and GgmlDType::cpu_zeros hands out storage, so this is reachable through the public API:

let storage = QStorage::Cpu(GgmlDType::Q4_0.cpu_zeros(32)); // 1 block
QTensor::new(storage, (64,))?;                              // needs 2 blocks — accepted today

Fix

Compare the storage block count against shape.elem_count() / block_size and return an error on mismatch:

let storage_blocks = storage.size_in_bytes() / storage.dtype().type_size();
let expected_blocks = shape.elem_count() / block_size;
if storage_blocks != expected_blocks {
    crate::bail!("quantized storage holds {storage_blocks} blocks but shape {shape:?} requires {expected_blocks} for dtype {:?}", storage.dtype())
}

The in-tree constructors (QTensor::quantize and the GGUF/GGML loaders) already size storage from the shape, so they're unaffected.

Tests

Adds qtensor_new_rejects_storage_shape_mismatch. Existing quantized_tests (39), gguf_tests (10) stay green; cargo fmt/clippy clean.

cc @ivarflakstad @EricLBuehler

@v-code01

v-code01 commented Jul 5, 2026

Copy link
Copy Markdown
Author

cc @LaurentMazare @ivarflakstad @EricLBuehler — QTensor::new never validated the storage block count against the shape, so a mismatched (storage, shape) pair constructs fine and then reads out of bounds in dequantize (panic for legacy quants, silent truncation for k-quants). Reachable via the public API. In-tree constructors already stay consistent, so they're unaffected.

`QTensor::new` only ran `check_shape`, which verifies the last dimension
is divisible by the block size but never checks that the storage actually
holds as many blocks as the shape requires. A mismatched pair such as a
one-block Q4_0 storage with shape `(64,)` (which needs two blocks) is
accepted, then `dequantize` sizes its output from the shape and indexes
the storage past its end: an out-of-bounds panic for the legacy quants,
and a silently truncated/zero-padded tensor for the k-quants.

`QTensor::new`/`QStorage` are public and `GgmlDType::cpu_zeros` hands out
storage, so this is reachable through the public API.

Compare the storage block count against `shape.elem_count() / block_size`
and return an error on mismatch. The in-tree constructors (quantize and
the GGUF/GGML loaders) already size storage from the shape, so they are
unaffected. Adds a regression test.
@v-code01
v-code01 force-pushed the qtensor-validate-storage-shape branch from 569f243 to 0082709 Compare July 10, 2026 18:40
astorise added a commit to astorise/candle that referenced this pull request Jul 21, 2026
astorise added a commit to astorise/candle that referenced this pull request Jul 21, 2026
Upstream huggingface#3716: validate storage/shape block count in QTensor::new

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant