refactor custom_scale usage

Author: shubhambhokare1

onnxscript/rewriter/ort_fusions/sdpa.py

@@ -13,7 +13,7 @@ def __init__(self, name: str, *, use_mask: bool, pre_scale: bool, use_mul: bool)
         self._use_mask = use_mask
         self._pre_scale = pre_scale
         self._use_mul = use_mul
-        self._custom_scale = False
+        self._scale: float | None = None
 
     def pattern(
         self, op, query, key_transposed, value, mask, query_scale, key_scale, qk_scale
@@ -60,14 +60,18 @@ def check(self, op, query, key_transposed, value, mask, query_scale, key_scale,
             # Check if query_scale and key_scale are scalars == sqrt(expected_scaling_factor)
             # If they are scalars but != sqrt(expected_scaling_factor), a custom scale is being used.
             sqrt_scaling_factor = math.sqrt(expected_scaling_factor)
-
+            # Calculate the scaling factor for query
             if _ir_utils.get_singleton_value(query_scale) is None:
                 return check_result.fail(
                     "Query scale is not a scalar.",
                     query_scale,
                 )
             if not _ir_utils.is_singleton_value(query_scale, sqrt_scaling_factor, rtol=1e-3):
-                self._custom_scale = True
+                query_scale_value = _ir_utils.get_singleton_value(query_scale)
+                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 _ir_utils.get_singleton_value(key_scale) is None:
                 return check_result.fail(
                     "Key scale is not a scalar.",
@@ -81,7 +85,6 @@ def check(self, op, query, key_transposed, value, mask, query_scale, key_scale,
                         "Query and key scales are not equal.",
                         query_scale,
                     )
-                self._custom_scale = True
         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
@@ -91,22 +94,19 @@ def check(self, op, query, key_transposed, value, mask, query_scale, key_scale,
                     qk_scale,
                 )
             if not _ir_utils.is_singleton_value(qk_scale, expected_scaling_factor, rtol=1e-3):
-                self._custom_scale = True
+                self._scale = _ir_utils.get_singleton_value(qk_scale)
+            else:
+                self._scale = expected_scaling_factor
 
         # check ranks/shapes
 
         return check_result
 
-    def rewrite(
-        self, op, query, key_transposed, value, mask, query_scale, key_scale, qk_scale, **_
-    ):
+    def rewrite(self, op, query, key_transposed, value, mask, **_):
         sdpa_args = [query, key_transposed, value]
         if self._use_mask:
             sdpa_args.append(mask)
-        if self._custom_scale:
-            scale = _ir_utils.get_singleton_value(query_scale if self._pre_scale else qk_scale)
-            return op.SDPA(*sdpa_args, scale=scale, _domain="ai.onnxruntime.fusion")
-        return op.SDPA(*sdpa_args, _domain="ai.onnxruntime.fusion")
+        return op.SDPA(*sdpa_args, scale=self._scale, _domain="ai.onnxruntime.fusion")
 
 
 # Rules for SDPA without mask