Validate DQ scale/zero_point shapes before QGemm fusion

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<const Node*>& dq_nodes,
                                   const std::vector<const Node*>& 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.

onnxruntime/test/optimizer/qdq_transformer_test.cc

@@ -902,6 +902,71 @@ TEST(QDQTransformerTests, Gemm_S8S8U8) {
   QDQTransformerGemmTests<int8_t, int8_t, uint8_t, uint8_t>();
 }
 
+// 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<float>({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<uint8_t>(input_a, .039f, 128, q1_output);
+    builder.AddDequantizeLinearNode<uint8_t>(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<uint8_t>({4, 6}, 0, 255);
+    auto* dq2_output = builder.MakeIntermediate();
+
+    // Create rank-2 scale and zero_point with shape [2, 6]
+    std::vector<float> scales_data(2 * 6, 0.04f);
+    std::vector<uint8_t> zp_data(2 * 6, 128);
+    NodeArg* scale_arg = builder.MakeInitializer<float>({2, 6}, scales_data);
+    NodeArg* zp_arg = builder.MakeInitializer<uint8_t>({2, 6}, zp_data);
+
+    std::vector<NodeArg*> 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<uint8_t>(gemm_output, .039f, 128, q_out);
+    builder.AddDequantizeLinearNode<uint8_t>(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<QDQSelectorActionTransformer>(QDQIsInt8Allowed()));
+}
+
 // Runs a test case that checks if Q/DQ nodes are dropped from DQ -> Gather -> Q.
 template <typename QuantType>
 static void RunGatherDropQDQTestCase(const std::vector<int64_t>& input1_shape,