add decoder model

Author: shubhambhokare1

onnxscript/rewriter/ort_fusions/_core.py

@@ -123,12 +123,13 @@ def optimize_for_ort(
     rewrite(model, ORT_PATTERN_REWRITE_RULES)
     return model, fusion_count
 
-'''
+
+"""
 from onnxscript import ir, rewriter
 import onnxscript.rewriter.ort_fusions as ort_fusions
 model_ir = ir.serde.deserialize_model(model)
 model_ir, count = ort_fusions.optimize_for_ort(model_ir)
 print("Applied fusions", count)
 print("\n\n\n\n\n\n\n\n\n\n\n")
 model = ir.serde.serialize_model(model_ir)
-'''
+"""

onnxscript/rewriter/ort_fusions/_whisper_tiny.py

@@ -58,8 +58,7 @@ def create_model(self, with_past=False):
 
         @script()
         def model_with_mha(input, weight, bias):
-            qkv_no_bias = op.MatMul(input, weight)
-            qkv = op.Add(qkv_no_bias, bias)
+            qkv = op.MatMul(input, weight)
 
             query_BSDh = op.Slice(qkv, [0], [160], [2])
             key_BSDh = op.Slice(qkv, [160], [320], [2])
@@ -69,14 +68,18 @@ def model_with_mha(input, weight, bias):
                 query_BSDh,
                 key_BSDh,
                 value_BSDh,
+                bias,
+                None,
+                None,
+                None,
+                None,
                 num_heads=self.num_heads,
             )
             return mha
 
         @script()
         def model_with_mha_past(input, weight, bias, past):
-            qkv_no_bias = op.MatMul(input, weight)
-            qkv = op.Add(qkv_no_bias, bias)
+            qkv = op.MatMul(input, weight)
 
             query_BSDh = op.Slice(qkv, [0], [160], [2])
             key_BSDh = op.Slice(qkv, [160], [320], [2])
@@ -91,7 +94,7 @@ def model_with_mha_past(input, weight, bias, past):
                 query_BSDh,
                 key_BSDh,
                 value_BSDh,
-                None,
+                bias,
                 None,
                 None,
                 past_key,

onnxscript/rewriter/ort_fusions/attention_test.py

@@ -27,7 +27,7 @@ def test_fuse_xformers(self):
         self.assertEqual(fusion_count["partial_rotary_embedding"], 0)
         self.assertEqual(fusion_count["cos_sin_cache"], 2)
         self.assertEqual(fusion_count["sdpa"], 1)
-        self.assertEqual(fusion_count["mha"], 0)
+        self.assertEqual(fusion_count["mha"], 1)
         self.assertEqual(fusion_count["attention"], 0)
         self.assertEqual(fusion_count["gqa"], 0)
         self.assertEqual(fusion_count["gelu"], 0)

onnxscript/rewriter/ort_fusions/fuse_xformers_test.py

@@ -8,8 +8,11 @@
 
 import onnxscript.optimizer
 import onnxscript.rewriter.ort_fusions._core as xformers
+from onnxscript.ir.passes.common import shape_inference
 from onnxscript.rewriter.ort_fusions._test_utils import ORT_VERSION, assert_allclose, ort_run
 from onnxscript.rewriter.ort_fusions.models._smollm_2 import smollm_test_2
+from onnxscript.rewriter.ort_fusions.models._whisper_decoder import whisper_decoder_test
+from onnxscript.rewriter.ort_fusions.models._whisper_encoder import whisper_encoder_test
 
 
 class TestMultiHeadAttention(unittest.TestCase):
@@ -40,6 +43,32 @@ def test_smollm(self):
             new_outputs = ort_run("optimized", model, inputs)
             assert_allclose(new_outputs, original_outputs)
 
+    def test_whisper_encoder(self):
+        # Generate model
+        whisper_encoder = whisper_encoder_test()
+        model = whisper_encoder.get_onnx_model()
+        onnxscript.optimizer.optimize(model)
+
+        # Fuse SDPA and MHA
+        sdpa_count = xformers.fuse_sdpa(model)
+        self.assertGreater(sdpa_count, 0)
+        model = shape_inference.infer_shapes(model)
+        mha_count = xformers.fuse_mha(model)
+        self.assertGreater(mha_count, 0)
+
+    def test_whisper_decoder(self):
+        # Generate model
+        whisper_decoder = whisper_decoder_test()
+        model = whisper_decoder.get_onnx_model()
+        onnxscript.optimizer.optimize(model)
+
+        # Fuse SDPA and MHA
+        sdpa_count = xformers.fuse_sdpa(model)
+        self.assertGreater(sdpa_count, 0)
+        model = shape_inference.infer_shapes(model)
+        mha_count = xformers.fuse_mha(model)
+        self.assertGreater(mha_count, 0)
+
 
 if __name__ == "__main__":
     unittest.main()