reorder checks

Author: shubhambhokare1

onnxscript/rewriter/ort_fusions/sdpa.py

@@ -66,10 +66,6 @@ def check(self, op, query, key_transposed, value, mask, query_scale, key_scale,
                     "Query scale is not a scalar.",
                     query_scale,
                 )
-            if not math.isclose(query_scale_value, sqrt_scaling_factor, rel_tol=1e-3):
-                self._scale = query_scale_value * query_scale_value
-            else:
-                self._scale = expected_scaling_factor
             # Ensure the scaling factor for key is the same as for query
             if (key_scale_value := _ir_utils.get_singleton_value(key_scale)) is None:
                 return check_result.fail(
@@ -81,6 +77,11 @@ def check(self, op, query, key_transposed, value, mask, query_scale, key_scale,
                     "Query and key scales are not equal.",
                     query_scale,
                 )
+            if not math.isclose(query_scale_value, sqrt_scaling_factor, rel_tol=1e-3):
+                self._scale = query_scale_value * query_scale_value
+            else:
+                # Pass no scaling factor to SDPA, SDPA will use the default scaling factor
+                self._scale = None
         else:
             # Check if qk_scale is a scalar == expected_scaling_factor)
             # If it is a scalar but != sqrt(expected_scaling_factor), a custom scale is being used
@@ -92,7 +93,8 @@ def check(self, op, query, key_transposed, value, mask, query_scale, key_scale,
             if not math.isclose(qk_scale_value, expected_scaling_factor, rel_tol=1e-3):
                 self._scale = qk_scale_value
             else:
-                self._scale = expected_scaling_factor
+                # Pass no scaling factor to SDPA, SDPA will use the default scaling factor
+                self._scale = None
 
         # check ranks/shapes
 

onnxscript/rewriter/ort_fusions/sdpa_test.py

@@ -303,20 +303,20 @@ def test_sdpa_fusion(self, name, script_func):
         # Ensure that the scale of the SDPA node is set correctly
         sdpa_node = next(n for n in model.graph if n.op_type == "SDPA")
         self.assertEqual(sdpa_node.op_type, "SDPA")
-        self.assertIsNotNone(sdpa_node.attributes.get("scale"))
 
-        scale_factor = sdpa_node.attributes["scale"].value
-        self.assertIsNotNone(scale_factor)
         if "custom" in name:
+            self.assertIsNotNone(sdpa_node.attributes.get("scale"))
+            scale_factor = sdpa_node.attributes["scale"].value
+            self.assertIsNotNone(scale_factor)
             if "pre" in name:
                 self.assertEqual(scale_factor, CUSTOM_SCALE_FACTOR * CUSTOM_SCALE_FACTOR)
             elif "post" in name:
                 self.assertEqual(scale_factor, CUSTOM_SCALE_FACTOR)
         else:
-            if "div" in name:
-                self.assertEqual(scale_factor, SCALE_FACTOR)
-            elif "mul" in name:
-                self.assertEqual(scale_factor, MUL_SCALE_FACTOR)
+            # These tests are for the default scaling factors, no scale factor is passed to SDPA
+            # pattern rewriting check functions should be sufficient to check if expected value
+            # of scale_factor (is =default_scaling_factor)
+            self.assertIsNone(sdpa_node.attributes.get("scale"))
 
         # new_outputs = ort_run("optimized", model, inputs)
         # assert_allclose(new_outputs, original_outputs)