From e92a010d954fbbf9267b9b99a80fc6bbf24f0bf9 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Wed, 5 Aug 2026 14:17:14 -0700 Subject: [PATCH 1/4] Add upper bound validation for DQ block_size in MatMulNBits fusion 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 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> --- .../qdq_transformer/selectors_actions/qdq_actions.cc | 8 +++++++- .../qdq_transformer/selectors_actions/qdq_selectors.cc | 5 ++++- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc index 6bd5e157d8b65..6d3bc5a16cb5d 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc @@ -87,7 +87,13 @@ int64_t GetEffectiveBlockSize(const Node& dq_node, int64_t block_size_for_non_bl const auto& dq_attrs = dq_node.GetAttributes(); const auto bs_iter = dq_attrs.find("block_size"); if (bs_iter != dq_attrs.end() && bs_iter->second.i() > 0) { - return bs_iter->second.i(); + // 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; } // Derive K from the weight input shape if available. Shape information may be missing even diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc index 8bf2b784a5474..a227ffcab5eee 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc @@ -776,7 +776,10 @@ static bool ValidateDQForMatMulNBits(const Graph& graph, const Node& dq_node) { } auto block_size = block_size_iter->second.i(); - if (block_size < 16 || ((block_size - 1) & block_size)) { + // Must be a power-of-two in [16, 256]. Values beyond 256 would overflow int32 in the + // MLAS transpose kernel (static_cast(block_size) == 0 for block_size >= 2^32), + // causing a divide-by-zero (SIGFPE) at CreateSession time. + if (block_size < 16 || block_size > 256 || ((block_size - 1) & block_size)) { return false; } From 4c624cc7c5bbc518979c4bd5598a5eb41e640e5e Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Thu, 6 Aug 2026 10:26:47 -0700 Subject: [PATCH 2/4] Address PR31678 review comments Align block_size comments with actual supported constraints, make GetEffectiveBlockSize fail fast for invalid blockwise values instead of silently clamping, and add regression coverage that rejects block_size > 256 on both CPU and CUDA transformer test paths. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../qdq_transformer/selectors_actions/qdq_actions.cc | 12 +++++++----- .../selectors_actions/qdq_selectors.cc | 5 ++--- .../optimizer/qdq_matmulnbits_transformer_test.cc | 4 ++++ 3 files changed, 13 insertions(+), 8 deletions(-) diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc index 6d3bc5a16cb5d..f50f9fa4de183 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc @@ -87,13 +87,15 @@ int64_t GetEffectiveBlockSize(const Node& dq_node, int64_t block_size_for_non_bl const auto& dq_attrs = dq_node.GetAttributes(); const auto bs_iter = dq_attrs.find("block_size"); if (bs_iter != dq_attrs.end() && bs_iter->second.i() > 0) { - // 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. + // Blockwise DQ feeding MatMulNBits must use a power-of-two block_size in [16, 256]. + // Keep this contract explicit here as a fail-fast defense if selector logic changes. + constexpr int64_t kMinBlockSize = 16; constexpr int64_t kMaxBlockSize = 256; const int64_t block_size = bs_iter->second.i(); - return (block_size <= kMaxBlockSize) ? block_size : kMaxBlockSize; + ORT_ENFORCE(block_size >= kMinBlockSize && block_size <= kMaxBlockSize && ((block_size - 1) & block_size) == 0, + "DQ block_size must be a power-of-two in [", kMinBlockSize, ", ", kMaxBlockSize, + "] for MatMulNBits fusion. Got: ", block_size); + return block_size; } // Derive K from the weight input shape if available. Shape information may be missing even diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc index a227ffcab5eee..5d2e15ad5c3e6 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc @@ -776,9 +776,8 @@ static bool ValidateDQForMatMulNBits(const Graph& graph, const Node& dq_node) { } auto block_size = block_size_iter->second.i(); - // Must be a power-of-two in [16, 256]. Values beyond 256 would overflow int32 in the - // MLAS transpose kernel (static_cast(block_size) == 0 for block_size >= 2^32), - // causing a divide-by-zero (SIGFPE) at CreateSession time. + // Must be a power-of-two in [16, 256], which is the supported blockwise range for + // this MatMulNBits fusion path. if (block_size < 16 || block_size > 256 || ((block_size - 1) & block_size)) { return false; } diff --git a/onnxruntime/test/optimizer/qdq_matmulnbits_transformer_test.cc b/onnxruntime/test/optimizer/qdq_matmulnbits_transformer_test.cc index b53577a81ff4a..8d4fdbdf771aa 100644 --- a/onnxruntime/test/optimizer/qdq_matmulnbits_transformer_test.cc +++ b/onnxruntime/test/optimizer/qdq_matmulnbits_transformer_test.cc @@ -303,6 +303,8 @@ TEST(QDQTransformerTests, DQMatMulNotConvertedToMatMulNBits_ShapeMismatch) { // One representative type combo per rejection scenario (type doesn't affect rejection logic). // block size too small RunDQMatMulNotConverted_TypeShapeMismatch({12, 37}, {37, 12}, 0, 8, 0); + // block size too large + RunDQMatMulNotConverted_TypeShapeMismatch({12, 37}, {37, 12}, 0, 512, 0); // block size not 2's power RunDQMatMulNotConverted_TypeShapeMismatch({12, 37}, {37, 12}, 0, 17, 0); // not axis 0 @@ -315,6 +317,8 @@ TEST(QDQTransformerTests, DQMatMulNotConvertedToMatMulNBits_ShapeMismatch_Cuda) // One representative type combo per rejection scenario. // block size too small RunDQMatMulNotConverted_TypeShapeMismatch({12, 37}, {37, 12}, 0, 8, 0, DefaultCudaExecutionProvider()); + // block size too large + RunDQMatMulNotConverted_TypeShapeMismatch({12, 37}, {37, 12}, 0, 512, 0, DefaultCudaExecutionProvider()); // block size not 2's power RunDQMatMulNotConverted_TypeShapeMismatch({12, 37}, {37, 12}, 0, 17, 0, DefaultCudaExecutionProvider()); // not axis 0 From c480e46a4c955c7ef50e455fc591d17221561739 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Mon, 10 Aug 2026 14:15:26 -0700 Subject: [PATCH 3/4] Harden MatMulNBits block-size replay validation Share the supported block-size contract between selectors and actions, and return a graceful failure for invalid model attributes during minimal-build action replay. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/optimizer/qdq_transformer/qdq_util.h | 12 +++- .../selectors_actions/qdq_actions.cc | 57 ++++++++++--------- .../qdq_selector_action_transformer.cc | 2 +- .../selectors_actions/qdq_selectors.cc | 8 +-- 4 files changed, 45 insertions(+), 34 deletions(-) diff --git a/onnxruntime/core/optimizer/qdq_transformer/qdq_util.h b/onnxruntime/core/optimizer/qdq_transformer/qdq_util.h index 0648a3fc1f188..9275227c8fa2b 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/qdq_util.h +++ b/onnxruntime/core/optimizer/qdq_transformer/qdq_util.h @@ -3,9 +3,10 @@ #pragma once +#include #include -#include #include +#include namespace ONNX_NAMESPACE { class TensorProto; @@ -22,6 +23,15 @@ namespace QDQ { constexpr const char* QOpName = "QuantizeLinear"; constexpr const char* DQOpName = "DequantizeLinear"; +constexpr int64_t kMatMulNBitsMinBlockSize = 16; +constexpr int64_t kMatMulNBitsMaxBlockSize = 256; + +constexpr bool IsValidMatMulNBitsBlockSize(int64_t block_size) { + return block_size >= kMatMulNBitsMinBlockSize && + block_size <= kMatMulNBitsMaxBlockSize && + (block_size & (block_size - 1)) == 0; +} + enum InputIndex : int { INPUT_ID = 0, SCALE_ID = 1, diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc index f50f9fa4de183..865fcb631f835 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc @@ -46,29 +46,25 @@ bool IsDQWeightSigned(int32_t dt_weight) { // Compute the effective block_size for per-tensor/per-channel DQ nodes that lack a block_size attribute. // session_block_size: 0 = default (32), positive = explicit, -1 = min-padding heuristic. int64_t ComputeEffectiveBlockSize(int64_t session_block_size, int64_t K) { - // MatMulNBits CPU kernel currently only supports block_size in [16, 256] correctly. - constexpr int64_t kMinBlockSize = 16; - constexpr int64_t kMaxBlockSize = 256; - if (session_block_size > 0) { - // Explicit block_size — must be power-of-2 and within [kMinBlockSize, kMaxBlockSize]. - ORT_ENFORCE(session_block_size >= kMinBlockSize && - ((session_block_size & (session_block_size - 1)) == 0), - "Explicit qdq_matmulnbits_block_size must be a power-of-2 and >= ", - kMinBlockSize, ", got: ", session_block_size); - ORT_ENFORCE(session_block_size <= kMaxBlockSize, - "Explicit qdq_matmulnbits_block_size must be <= ", - kMaxBlockSize, ", got: ", session_block_size); + ORT_ENFORCE(IsValidMatMulNBitsBlockSize(session_block_size), + "Explicit qdq_matmulnbits_block_size must be a power-of-two in [", + kMatMulNBitsMinBlockSize, ", ", kMatMulNBitsMaxBlockSize, + "], got: ", session_block_size); return session_block_size; } if (session_block_size == -1) { - // Heuristic: largest power-of-2 <= min(K, kMaxBlockSize) that minimizes padding. - // Capped at kMaxBlockSize because CPU EP only supports block_size up to kMaxBlockSize correctly. + // Heuristic: largest supported power-of-two that minimizes padding. // We want ceil(K / B) * B - K to be minimized (least wasted padding). - int64_t best_bs = kMinBlockSize; - int64_t best_padding = (((K + (kMinBlockSize - 1)) / kMinBlockSize) * kMinBlockSize) - K; - for (int64_t bs = kMinBlockSize * 2; bs <= std::min(K, kMaxBlockSize); bs *= 2) { + int64_t best_bs = kMatMulNBitsMinBlockSize; + int64_t best_padding = (((K + (kMatMulNBitsMinBlockSize - 1)) / + kMatMulNBitsMinBlockSize) * + kMatMulNBitsMinBlockSize) - + K; + for (int64_t bs = kMatMulNBitsMinBlockSize * 2; + bs <= std::min(K, kMatMulNBitsMaxBlockSize); + bs *= 2) { int64_t padding = (((K + bs - 1) / bs) * bs) - K; if (padding <= best_padding) { best_padding = padding; @@ -87,15 +83,8 @@ int64_t GetEffectiveBlockSize(const Node& dq_node, int64_t block_size_for_non_bl const auto& dq_attrs = dq_node.GetAttributes(); const auto bs_iter = dq_attrs.find("block_size"); if (bs_iter != dq_attrs.end() && bs_iter->second.i() > 0) { - // Blockwise DQ feeding MatMulNBits must use a power-of-two block_size in [16, 256]. - // Keep this contract explicit here as a fail-fast defense if selector logic changes. - constexpr int64_t kMinBlockSize = 16; - constexpr int64_t kMaxBlockSize = 256; - const int64_t block_size = bs_iter->second.i(); - ORT_ENFORCE(block_size >= kMinBlockSize && block_size <= kMaxBlockSize && ((block_size - 1) & block_size) == 0, - "DQ block_size must be a power-of-two in [", kMinBlockSize, ", ", kMaxBlockSize, - "] for MatMulNBits fusion. Got: ", block_size); - return block_size; + // ProcessNewNode validates model-provided values before using them. + return bs_iter->second.i(); } // Derive K from the weight input shape if available. Shape information may be missing even @@ -649,7 +638,21 @@ Status DQMatMulToMatMulNBitsAction::ProcessNewNode(Graph& graph, Node& replacement_node) const { const auto* dq_node = selected_nodes.Input(0); - int64_t effective_bs = GetEffectiveBlockSize(*dq_node, block_size_for_non_blockwise_); + const auto& dq_attributes = dq_node->GetAttributes(); + const auto block_size_iter = dq_attributes.find("block_size"); + if (block_size_iter != dq_attributes.end()) { + const int64_t block_size = block_size_iter->second.i(); + ORT_RETURN_IF_NOT(IsValidMatMulNBitsBlockSize(block_size), + "DQ block_size must be a power-of-two in [", + kMatMulNBitsMinBlockSize, ", ", kMatMulNBitsMaxBlockSize, + "] for MatMulNBits fusion. Got: ", block_size); + } + + const int64_t effective_bs = GetEffectiveBlockSize(*dq_node, block_size_for_non_blockwise_); + ORT_RETURN_IF_NOT(IsValidMatMulNBitsBlockSize(effective_bs), + "Effective block_size must be a power-of-two in [", + kMatMulNBitsMinBlockSize, ", ", kMatMulNBitsMaxBlockSize, + "] for MatMulNBits fusion. Got: ", effective_bs); TransposedQuantizedTensors transposed; ORT_RETURN_IF_ERROR(TransposeDQWeightsForMatMulNBits( diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.cc index 1e0f633360889..feb8cdf7792fb 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.cc @@ -300,7 +300,7 @@ void DQMatMulToMatMulNBitsRules(SelectorActionRegistry& qdq_selector_action_regi int64_t qdq_matmulnbits_block_size) { // 2 nodes. DQ -> MatMul. DQ is the second input to MatMul. // DQ's weight is 2/4/8-bit int (int2/uint2, int4/uint4, int8/uint8). DQ's scale is float/float16. - // DQ is block-quantized along axis 0, with block_size >= 16 and as 2's power. + // DQ is block-quantized along axis 0, with a power-of-two block_size in [16, 256]. // Also supports per-tensor and per-channel (axis=1) quantized DQ weights by expanding // scales/zero-points to blockwise format using qdq_matmulnbits_block_size. const std::string action_name{"DQMatMulToMatMulNBits"}; diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc index 5d2e15ad5c3e6..63c818d9295da 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc @@ -726,7 +726,7 @@ bool MatMulNodeGroupSelector::Check(const GraphViewer& graph_viewer, const Node& // Validate that a DQ node has the correct structure for MatMulNBits fusion. // Supports three quantization granularities: -// - Blockwise: axis=0, block_size >= 16 and power-of-2, scale/zp rank 2 +// - Blockwise: axis=0, block_size is a power-of-two in [16, 256], scale/zp rank 2 // - Per-tensor: scale is scalar (rank 0), no block_size attribute // - Per-channel (axis=1): scale is 1D with shape [N], weight is 2D [K,N], no block_size attribute // In all cases: weight type is 2/4/8-bit int, scale type is float or float16, @@ -775,10 +775,8 @@ static bool ValidateDQForMatMulNBits(const Graph& graph, const Node& dq_node) { return false; } - auto block_size = block_size_iter->second.i(); - // Must be a power-of-two in [16, 256], which is the supported blockwise range for - // this MatMulNBits fusion path. - if (block_size < 16 || block_size > 256 || ((block_size - 1) & block_size)) { + const auto block_size = block_size_iter->second.i(); + if (!IsValidMatMulNBitsBlockSize(block_size)) { return false; } From a25a8d14cef6eb348e60cc66107b48f96d916ef3 Mon Sep 17 00:00:00 2001 From: Akshay Sonawane Date: Tue, 11 Aug 2026 15:27:20 -0700 Subject: [PATCH 4/4] Only validate model-provided DQ block_size in MatMulNBits fusion Graph::Resolve populates the schema default of 0 for DequantizeLinear nodes that omit block_size, so checking for attribute presence alone also caught per-tensor and per-channel DQ nodes and rejected them. Match the selector by treating only a positive value as blockwise; the derived effective block size is still validated afterwards. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../qdq_transformer/selectors_actions/qdq_actions.cc | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc index 865fcb631f835..44337d38964d0 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_actions.cc @@ -638,9 +638,12 @@ Status DQMatMulToMatMulNBitsAction::ProcessNewNode(Graph& graph, Node& replacement_node) const { const auto* dq_node = selected_nodes.Input(0); + // Graph::Resolve fills in the schema default (block_size = 0) for DQ nodes that do not declare + // one, so presence of the attribute alone does not mean the DQ is blockwise. Only a positive + // value is model-provided; zero means per-tensor/per-channel and the block size is derived below. const auto& dq_attributes = dq_node->GetAttributes(); const auto block_size_iter = dq_attributes.find("block_size"); - if (block_size_iter != dq_attributes.end()) { + if (block_size_iter != dq_attributes.end() && block_size_iter->second.i() > 0) { const int64_t block_size = block_size_iter->second.i(); ORT_RETURN_IF_NOT(IsValidMatMulNBitsBlockSize(block_size), "DQ block_size must be a power-of-two in [",