add quality check

Author: shubhambhokare1

onnxscript/rewriter/ort_fusions/sdpa.py

@@ -74,6 +74,13 @@ def check(self, op, query, key_transposed, value, mask, query_scale, key_scale,
                     key_scale,
                 )
             if not _ir_utils.is_singleton_value(key_scale, sqrt_scaling_factor, rtol=1e-3):
+                if _ir_utils.get_singleton_value(query_scale) != _ir_utils.get_singleton_value(
+                    key_scale
+                ):
+                    return check_result.fail(
+                        "Query and key scales are not equal.",
+                        query_scale,
+                    )
                 self._custom_scale = True
         else:
             # Check if qk_scale is a scalar == expected_scaling_factor)

onnxscript/rewriter/ort_fusions/sdpa_test.py

@@ -98,6 +98,19 @@ def _custom_scale_pre_mul_sdpa_script(query, key, value):
     return attn_output
 
 
+@script()
+def _custom_multi_scale_pre_mul_sdpa_script(query, key, value):
+    key_transposed = op.Transpose(key, perm=[0, 1, 3, 2])
+    multiplier_q = op.Constant(value_float=0.5)
+    multiplier_k = op.Constant(value_float=0.5)
+    scaled_query = op.Mul(query, multiplier_q)
+    scaled_key = op.Mul(key_transposed, multiplier_k)
+    attn_score = op.MatMul(scaled_query, scaled_key)
+    attn_weight = op.Softmax(attn_score, axis=-1)
+    attn_output = op.MatMul(attn_weight, value)
+    return attn_output
+
+
 @script()
 def _custom_scale_post_div_sdpa_script(query, key, value):
     key_transposed = op.Transpose(key, perm=[0, 1, 3, 2])
@@ -265,6 +278,7 @@ class TestSDPAFusion(unittest.TestCase):
             ("custom_scale_post_div_masked", _custom_scale_post_div_sdpa_script),
             ("custom_scale_pre_mul_masked", _custom_scale_pre_mul_sdpa_script),
             ("custom_scale_pre_div_masked", _custom_scale_pre_div_sdpa_script),
+            (_custom_multi_scale_pre_mul_sdpa_script, _custom_multi_scale_pre_mul_sdpa_script),
         ]
     )
     def test_sdpa_fusion(self, name, script_func):