Validate MLAS blockwise QDQ index ranges

[Akshay Sonawane (apsonawane)]

This pull request introduces important safety checks to the quantization code in ONNX Runtime to ensure that blockwise quantization shapes do not exceed the valid int32_t index range, preventing potential overflows and runtime errors. The main changes include the addition of a validation function, integration of this check into quantization and transpose routines, and new unit tests to verify the behavior.

Shape validation and enforcement:

• Added MlasQDQBlockwiseShapeIsValid function in mlas_q4.h to validate that quantization shape parameters fit within the int32_t index domain, guarding against arithmetic overflows.
• Updated MlasQDQQuantizeBlockwise and MlasQDQTransposeBlockwiseQuantized in q4_dq.cpp to enforce this validation using ORT_ENFORCE, throwing an exception if the shape is invalid. [1] [2]
• Added a similar check in TransposeDQWeightsForMatMulNBits to return an error if the shape is invalid.

Testing and code hygiene:

• Added a new unit test RejectsShapesOutsideInt32IndexDomain in test_blockq4.cpp to verify that invalid shapes are correctly rejected and exceptions are thrown as expected.
• Included missing headers <cstdint> and <limits> in mlas_q4.h to support the new validation logic.

[Copilot]

Pull request overview

This PR adds safety validation around MLAS blockwise QDQ quantization/transposition to ensure internal int32 index arithmetic cannot overflow, and wires those checks into both the MLAS kernels and the QDQ transformer path.

Changes:

• Introduces MlasQDQBlockwiseShapeIsValid to validate blockwise QDQ shapes against the MLAS int32 index domain.
• Enforces the validation in MLAS blockwise quantize/transpose routines and returns an error in the QDQ transformer when the shape is invalid.
• Adds a unit test to verify invalid shapes are rejected.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/core/mlas/inc/mlas_q4.h	Adds the shared shape validation helper used to guard int32 index arithmetic.
onnxruntime/core/mlas/lib/q4_dq.cpp	Enforces the new validation at entry points for blockwise QDQ quantize/transpose.
onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc	Adds an early shape validity check in the MatMulNBits weight transpose path.
onnxruntime/test/mlas/unittest/test_blockq4.cpp	Adds a unit test ensuring out-of-domain shapes are rejected (including exception expectations).

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.