Validate GGUF tensor dimensions

[roshaninfordham]

Fixes #4244 (thanks @perparimmjeku for the report).

The bug

GGUF stores tensor dimensions as 64-bit values, while MLX represents dimensions with ShapeElem, which is int32_t.

Previously, get_shape() copied GGUF dimensions directly into the MLX shape, allowing a 64-bit dimension to narrow to a different 32-bit value. In the quantized loader, this could make the allocation size smaller than the number of elements subsequently processed, resulting in a heap-buffer-overflow during tensor extraction.

The dimension product could also wrap in 64-bit arithmetic, causing the MLX shape and the tensor's element count to describe different sizes.

Reproducing it

The regression tests cover:

• {32} as a valid quantized tensor dimension.
• (1ull << 32) + 32 to test a dimension larger than the ShapeElem range.
• 1ull << 31 to test a value that would become negative when represented as signed int32_t.
• {96, 384307168202282326ull} to produce a dimension product that wraps in 64-bit arithmetic.
• Three large dimensions whose product wraps to zero.
• {} to test a quantized tensor with no dimensions.

On unmodified main, the crafted dimension-product case produces:

==82611==ERROR: AddressSanitizer: heap-buffer-overflow on address 0x6020000008fc
WRITE of size 2 at 0x6020000008fc thread T0
    #0 mlx::core::extract_q8_0_data(...) gguf_quants.cpp:90
    #1 mlx::core::gguf_load_quantized(...) gguf_quants.cpp:145
    #2 mlx::core::load_arrays(gguf_ctx*) gguf.cpp:340
    #3 mlx::core::load_gguf(...) gguf.cpp:367

0x6020000008fc is located 0 bytes after 12-byte region [0x6020000008f0,0x6020000008fc)
allocated by thread T0 here:
    #1 mlx::core::allocator::CommonAllocator::malloc(unsigned long) allocator.cpp:102
    #3 mlx::core::gguf_load_quantized(...) gguf_quants.cpp:138

The allocation size is derived from a 32-bit accumulator, while the extraction loop operates on the full shape-derived element count. This allows the loop to write beyond the allocated buffer.

The same underlying issue is present in the Q4_0 and Q4_1 extraction paths.

The fix

mlx/io/gguf.cpp

• Added <limits>.
• Added tensor_error() to centralize tensor validation errors.
• get_shape() now validates every GGUF dimension against the maximum representable ShapeElem value before narrowing.
• Added overflow-safe multiplication using uint64_t maximum divided by the next dimension.
• Reused tensor_error() for the existing tensor file-boundary validation.

The validation is performed before narrowing or multiplying, so malformed dimensions cannot silently produce a different MLX shape or a wrapped element count.

mlx/io/gguf_quants.cpp

• Added an explicit check for an empty shape before the quantized loader indexes the last dimension. A zero-dimensional tensor is valid elsewhere in MLX, so this check is kept in the quantized path where a last dimension is required.
• Changed both std::accumulate() initial values from 1 to size_t{1}.

The std::accumulate() change is important because its accumulator type is deduced from the initial value, not from the binary operation. With the literal 1, the accumulator is an int, even though the operation is std::multiplies<size_t>. This causes intermediate products to be converted back to 32-bit int. Using size_t{1} keeps the accumulation in the intended width.

I did not add allocator-failure checks in this change. Once GGUF dimensions are validated, an allocation of genuinely enormous size requires a correspondingly enormous input file; handling general out-of-memory conditions is outside the scope of this fix.

Why not compare against tensor.num_weights?

tensor.num_weights is the 64-bit product maintained by gguflib, but comparing a newly calculated product against it would not solve the representation problem.

MLX ultimately stores dimensions as ShapeElem, which is int32_t, so every individual dimension must first be proven representable in that type. The multiplication must then be proven safe before calculating the product.

The fix therefore validates the GGUF inputs before narrowing and before multiplication, rather than relying on a potentially narrowed shape or a wrapped product.

Prior art

• Bound GGUF tensor data offsets against the file mapping #4179 validates GGUF tensor data offsets and byte extents against the mapped file.
• Bound GGUF metadata string/array values against the file mapping #4212 validates GGUF metadata string and array lengths against the mapped file.

This PR addresses a different issue: malformed tensor dimensions and dimension-product arithmetic can cause the MLX shape and element count to disagree, leading to unsafe allocation and indexing behavior in the quantized loader.

Testing

Added TEST_CASE("test gguf tensor dimension validation") in tests/load_tests.cpp, covering six cases.

Built CPU-only with AddressSanitizer:

cmake -S . -B build-asan -DCMAKE_BUILD_TYPE=Debug -DMLX_BUILD_METAL=OFF \
  -DMLX_BUILD_TESTS=ON \
  -DCMAKE_CXX_FLAGS="-fsanitize=address -fno-omit-frame-pointer -g" \
  -DCMAKE_EXE_LINKER_FLAGS="-fsanitize=address"

cmake --build build-asan -j10 --target tests

On unmodified main, the new dimension-product test triggers the heap-buffer-overflow shown above.

With this change, the whole C++ suite passes under ASAN:

[doctest] test cases:  252 |  252 passed | 0 failed | 0 skipped
[doctest] assertions: 3366 | 3366 passed | 0 failed |

Python load and quantized tests:

51 passed, 3 skipped, 3291 subtests passed

uvx pre-commit run --all is clean.

The reporter's four PoC files are also rejected by the loader with:

[load_gguf] Tensor 'w.weight' has a dimension that is too large. Perhaps an incomplete download or corrupt file?

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)
• I understand it is strictly prohibited to use AI to write PR description
• 100% responsible for every line, however it was produced, I reviewed the change, understand it, and take responsibility.
• AI usage disclosure: AI was used to:
  • read and analyse the GGUF load path
  • derive the smaller reproducer, including the constraint solve that produced dim[0] = 96, dim[1] = 384307168202282326
  • write the change to mlx/io/gguf.cpp
  • write the change to mlx/io/gguf_quants.cpp
  • write the six-subcase regression test
  • run the ASAN build, the C++ suite, the Python tests, and pre-commit
  • confirm the test fails on unmodified main, and that the PoC files are rejected
  • typeset my PR description's markdown — no wording changed