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 2e9d46656b514..7b913414de22a 100644 --- a/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc +++ b/onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc @@ -797,6 +797,41 @@ bool DQMatMulNodeGroupSelector::Check(const GraphViewer& graph_viewer, const Nod return ValidateDQForMatMulNBits(graph, *weight_dq); } +// Check if a QDQ node's scale/zero_point shape is scalar or 1D with a compatible size. +// When expected_n <= 0 (default), only accepts scalar or 1D with size 1 (per-tensor quantization). +// When expected_n > 0, also accepts 1D with size equal to expected_n (per-column quantization). +// Rejects when shape or dim is unknown to avoid incorrect fusion. +static bool IsScalarOr1DWithSizeOneOrN(const NodeArg* arg, int64_t expected_n = -1) { + if (!arg || !arg->Exists()) { + return true; // absent optional input is OK + } + + const auto* shape = arg->Shape(); + if (!shape) { + return false; // unknown shape, reject to avoid incorrect fusion + } + + int rank = shape->dim_size(); + if (rank == 0) { + return true; // scalar + } + + if (rank == 1) { + const auto& dim = shape->dim(0); + if (!dim.has_dim_value()) { + return false; // unknown dim, reject to avoid incorrect fusion + } + + if (expected_n > 0) { + return dim.dim_value() == 1 || dim.dim_value() == expected_n; + } else { + return dim.dim_value() == 1; + } + } + + return false; // rank > 1 +} + bool GemmNodeGroupSelector::Check(const GraphViewer& graph_viewer, const Node& node, const Node* redundant_clip_node, const std::vector& dq_nodes, const std::vector& q_nodes) const { @@ -805,6 +840,61 @@ bool GemmNodeGroupSelector::Check(const GraphViewer& graph_viewer, const Node& n return false; } + // Validate scale/zero_point shapes are compatible with QGemm. + // QGemm requires: + // A: scale/zero_point must be scalar or 1D with size 1 + // B: scale/zero_point must be scalar or 1D with size 1 or N, scale and zero_point must have the same shape size + // Y: scale/zero_point must be scalar or 1D with size 1 + // + // QGemm requires a_zero_point and b_zero_point as mandatory inputs. + // DQ nodes must have zero_point (input 2) that actually exists. + const auto& a_defs = dq_nodes[0]->InputDefs(); + if (a_defs.size() <= 2 || !a_defs[2]->Exists()) { + return false; // DQ_A missing zero_point, QGemm requires it + } + + if (!IsScalarOr1DWithSizeOneOrN(a_defs[1]) || + !IsScalarOr1DWithSizeOneOrN(a_defs[2])) { + return false; + } + + const auto& b_defs = dq_nodes[1]->InputDefs(); + if (b_defs.size() <= 2 || !b_defs[2]->Exists()) { + return false; // DQ_B missing zero_point, QGemm requires it + } + + // Try to get N for validating per-column scale/zero_point size. + // N depends on transB: if transB=1, N=B.shape[0]; otherwise N=B.shape[1]. + int64_t expected_n = -1; + const auto* b_weight_shape = b_defs[0]->Shape(); + if (b_weight_shape && b_weight_shape->dim_size() == 2) { + const auto* trans_b_attr = graph_utils::GetNodeAttribute(node, "transB"); + bool trans_b = trans_b_attr && trans_b_attr->i() != 0; + int n_dim = trans_b ? 0 : 1; + if (b_weight_shape->dim(n_dim).has_dim_value()) { + expected_n = b_weight_shape->dim(n_dim).dim_value(); + } + } + + if (!IsScalarOr1DWithSizeOneOrN(b_defs[1], expected_n) || + !IsScalarOr1DWithSizeOneOrN(b_defs[2], expected_n)) { + return false; + } + + // B's scale and zero_point must have the same rank and the same dim value. + const auto* b_scale_shape = b_defs[1]->Shape(); + const auto* b_zp_shape = b_defs[2]->Shape(); + if (b_scale_shape && b_zp_shape) { + if (b_scale_shape->dim_size() != b_zp_shape->dim_size()) { + return false; + } + + if (b_scale_shape->dim_size() == 1 && + b_scale_shape->dim(0).dim_value() != b_zp_shape->dim(0).dim_value()) { + return false; + } + } + // input and output types need to be same int32_t dt_A = dq_nodes[0]->InputDefs()[0]->TypeAsProto()->tensor_type().elem_type(); int32_t dt_B = dq_nodes[1]->InputDefs()[0]->TypeAsProto()->tensor_type().elem_type(); @@ -820,6 +910,13 @@ bool GemmNodeGroupSelector::Check(const GraphViewer& graph_viewer, const Node& n if (dt_A != dt_Y) { // activation and output must be same type return false; } + + // QGemm requires Y's scale and zero_point to be scalar or 1D with size 1. + const auto& y_defs = q_nodes[0]->InputDefs(); + if (!IsScalarOr1DWithSizeOneOrN(y_defs[1]) || + (y_defs.size() > 2 && !IsScalarOr1DWithSizeOneOrN(y_defs[2]))) { + return false; + } } // 16-bit int types must be explicitly allowed. diff --git a/onnxruntime/test/optimizer/qdq_transformer_test.cc b/onnxruntime/test/optimizer/qdq_transformer_test.cc index bdbd2c488584d..a585640ced161 100644 --- a/onnxruntime/test/optimizer/qdq_transformer_test.cc +++ b/onnxruntime/test/optimizer/qdq_transformer_test.cc @@ -902,6 +902,71 @@ TEST(QDQTransformerTests, Gemm_S8S8U8) { QDQTransformerGemmTests(); } +// Verify that DQ nodes with rank > 1 scale/zero_point are not fused into QGemm. +TEST(QDQTransformerTests, Gemm_NoQGemmFusionWithHighRankScaleZp) { + auto build_test_case = [](ModelTestBuilder& builder) { + // Shapes: A is [2, 4], B is [4, 6] + auto* input_a = builder.MakeInput({2, 4}, -1.f, 1.f); + auto* output_arg = builder.MakeOutput(); + + // DQ for A: scalar scale/zp (compatible with QGemm) + auto* q1_output = builder.MakeIntermediate(); + auto* dq1_output = builder.MakeIntermediate(); + builder.AddQuantizeLinearNode(input_a, .039f, 128, q1_output); + builder.AddDequantizeLinearNode(q1_output, .039f, 128, dq1_output); + + // DQ for B: rank-2 scale and zero_point (blockwise quantization, NOT compatible with QGemm) + // Weight shape [4, 6], block_size=2 along axis 0 => scale/zp shape [2, 6] + auto* weight_b = builder.MakeInitializer({4, 6}, 0, 255); + auto* dq2_output = builder.MakeIntermediate(); + + // Create rank-2 scale and zero_point with shape [2, 6] + std::vector scales_data(2 * 6, 0.04f); + std::vector zp_data(2 * 6, 128); + NodeArg* scale_arg = builder.MakeInitializer({2, 6}, scales_data); + NodeArg* zp_arg = builder.MakeInitializer({2, 6}, zp_data); + + std::vector dq_inputs = {weight_b, scale_arg, zp_arg}; + NodeAttributes dq_attrs; + ONNX_NAMESPACE::AttributeProto axis_attr; + axis_attr.set_name("axis"); + axis_attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + axis_attr.set_i(0); + dq_attrs["axis"] = axis_attr; + ONNX_NAMESPACE::AttributeProto block_size_attr; + block_size_attr.set_name("block_size"); + block_size_attr.set_type(ONNX_NAMESPACE::AttributeProto_AttributeType_INT); + block_size_attr.set_i(2); + dq_attrs["block_size"] = block_size_attr; + builder.AddNode("DequantizeLinear", dq_inputs, {dq2_output}, "", &dq_attrs); + + // Gemm node + auto* gemm_output = builder.MakeIntermediate(); + builder.AddNode("Gemm", {dq1_output, dq2_output}, {gemm_output}); + + // Q for output + auto* q_out = builder.MakeIntermediate(); + builder.AddQuantizeLinearNode(gemm_output, .039f, 128, q_out); + builder.AddDequantizeLinearNode(q_out, .039f, 128, output_arg); + }; + + auto check_graph = [](InferenceSessionWrapper& session) { + auto op_to_count = CountOpsInGraph(session.GetGraph()); + // QGemm should not be created because DQ B has rank-2 scale/zero_point + EXPECT_EQ(op_to_count["com.microsoft.QGemm"], 0); + EXPECT_EQ(op_to_count["Gemm"], 1); + }; + + TransformerTester(build_test_case, + check_graph, + TransformerLevel::Level1, + TransformerLevel::Level2, + 21 /*opset_version*/, + 0.01 /*per_sample_tolerance*/, + 0.01 /*relative_per_sample_tolerance*/, + std::make_unique(QDQIsInt8Allowed())); +} + // Runs a test case that checks if Q/DQ nodes are dropped from DQ -> Gather -> Q. template static void RunGatherDropQDQTestCase(const std::vector& input1_shape,