Use self._use_mask in sdpa rewrite call (microsoft#2135)

Co-authored-by: G. Ramalingam <[email protected]>

Author: 2 people

onnxscript/rewriter/ort_fusions/sdpa.py

@@ -60,7 +60,10 @@ def check(self, op, query, key_transposed, value, mask, query_scale, key_scale,
         return True
 
     def rewrite(self, op, query, key_transposed, value, mask, **_):
-        return op.SDPA(query, key_transposed, value, mask, _domain="ai.onnxruntime.fusion")
+        if self._use_mask:
+            return op.SDPA(query, key_transposed, value, mask, _domain="ai.onnxruntime.fusion")
+        else:
+            return op.SDPA(query, key_transposed, value, _domain="ai.onnxruntime.fusion")
 
 
 masked_pre_mul_sdpa_rule = SDPA.rule("masked_pre_mul_sdpa", use_mask=True, pre_scale=True)

onnxscript/rewriter/ort_fusions/sdpa_test.py

@@ -29,7 +29,7 @@
 
 
 @script()
-def _unmasked_pre_div_sdpa_script(query, key, value, mask):
+def _unmasked_pre_div_sdpa_script(query, key, value):
     key_transposed = op.Transpose(key, perm=[0, 1, 3, 2])
     divisor = op.Constant(value_float=SQRT_SCALE_FACTOR)
     scaled_query = op.Div(query, divisor)
@@ -41,7 +41,7 @@ def _unmasked_pre_div_sdpa_script(query, key, value, mask):
 
 
 @script()
-def _unmasked_pre_mul_sdpa_script(query, key, value, mask):
+def _unmasked_pre_mul_sdpa_script(query, key, value):
     key_transposed = op.Transpose(key, perm=[0, 1, 3, 2])
     multiplier = op.Constant(value_float=SQRT_MUL_SCALE_FACTOR)
     scaled_query = op.Mul(query, multiplier)
@@ -53,7 +53,7 @@ def _unmasked_pre_mul_sdpa_script(query, key, value, mask):
 
 
 @script()
-def _unmasked_post_div_sdpa_script(query, key, value, mask):
+def _unmasked_post_div_sdpa_script(query, key, value):
     key_transposed = op.Transpose(key, perm=[0, 1, 3, 2])
     divisor = op.Constant(value_float=SCALE_FACTOR)
     attn_score = op.MatMul(query, key_transposed)
@@ -64,7 +64,7 @@ def _unmasked_post_div_sdpa_script(query, key, value, mask):
 
 
 @script()
-def _unmasked_post_mul_sdpa_script(query, key, value, mask):
+def _unmasked_post_mul_sdpa_script(query, key, value):
     key_transposed = op.Transpose(key, perm=[0, 1, 3, 2])
     multiplier = op.Constant(value_float=MUL_SCALE_FACTOR)
     attn_score = op.MatMul(query, key_transposed)