Add upper bound validation for DQ block_size in MatMulNBits fusion - #31678
Open
apsonawane wants to merge 1 commit into
Open
Add upper bound validation for DQ block_size in MatMulNBits fusion#31678apsonawane wants to merge 1 commit into
apsonawane wants to merge 1 commit into
Conversation
The DQ->MatMulNBits blockwise fusion selector now rejects models with block_size > 256, mirroring the existing [16, 256] range enforced by ComputeEffectiveBlockSize for the session-option path. Without this cap, a model-supplied block_size >= 2^32 truncates to 0 via static_cast<int> in the MLAS transpose kernel, causing an integer divide-by-zero (SIGFPE) at CreateSession time before any inference. Changes: - qdq_selectors.cc: extend the existing power-of-two check to also reject block_size > 256; invalid models skip the fusion cleanly. - qdq_actions.cc: GetEffectiveBlockSize caps the raw model attribute at kMaxBlockSize (256) as defense-in-depth, so no oversized value can reach the MLAS kernel even if the selector is bypassed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
This PR hardens MatMulNBits fusion in the QDQ transformer by tightening validation of the DequantizeLinear block_size attribute so that only kernel-supported values reach the MLAS transpose path, preventing session-creation crashes from unsafe attributes.
Changes:
- Enforced
block_sizeto be a power-of-two within[16, 256]in MatMulNBits fusion eligibility checks. - Added an additional “safety guard” in block size computation to avoid unsafe values reaching the MLAS transpose call.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.
| File | Description |
|---|---|
| onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc | Tightens selector validation for blockwise DQ block_size and documents the safety rationale. |
| onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc | Adds extra guarding when deriving the effective block size used by the fusion action/MLAS transpose. |
Comment on lines
+779
to
+781
| // Must be a power-of-two in [16, 256]. Values beyond 256 would overflow int32 in the | ||
| // MLAS transpose kernel (static_cast<int>(block_size) == 0 for block_size >= 2^32), | ||
| // causing a divide-by-zero (SIGFPE) at CreateSession time. |
Comment on lines
+90
to
+96
| // Mirror the kMaxBlockSize guard from ComputeEffectiveBlockSize. Without this cap, a | ||
| // model-supplied block_size >= 2^32 truncates to 0 in the MLAS int32 transpose argument, | ||
| // causing a divide-by-zero at CreateSession. The selector rejects such models, but guard | ||
| // here too so any future selector bypass cannot reach MLAS with an unsafe value. | ||
| constexpr int64_t kMaxBlockSize = 256; | ||
| const int64_t block_size = bs_iter->second.i(); | ||
| return (block_size <= kMaxBlockSize) ? block_size : kMaxBlockSize; |
Comment on lines
+782
to
783
| if (block_size < 16 || block_size > 256 || ((block_size - 1) & block_size)) { | ||
| return false; |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces additional validation and safety checks for the
block_sizeattribute in quantization-related code to prevent unsafe values from causing runtime errors. The main changes ensure thatblock_sizeis always within the safe range supported by the underlying MLAS kernel, avoiding potential divide-by-zero errors.Validation and safety improvements for
block_size:kMaxBlockSize = 256) inGetEffectiveBlockSizeto cap user-supplied values and prevent unsafe block sizes from reaching the MLAS kernel, mirroring existing checks elsewhere.ValidateDQForMatMulNBitsto require thatblock_sizeis a power of two in the range [16, 256], explicitly rejecting values above 256 to avoid integer overflow and possible divide-by-zero errors.