fix skip_normalization fusion

Author: shubhambhokare1

onnxscript/rewriter/ort_fusions/_core.py

@@ -122,3 +122,13 @@ def optimize_for_ort(
     model, fusion_count = fuse_xformers(model)
     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

@@ -37,7 +37,7 @@ def test_whisper_tiny(self):
         encoder_model.opset_imports["ai.onnxruntime.fusion"] = 1
 
         print(f"Fused {fusion_count_e} ops")
-        self.assertEqual(fusion_count_e["skip_layer_normalization"], 17)
+        self.assertEqual(fusion_count_e["skip_layer_normalization"], 8)
         self.assertEqual(fusion_count_e["sdpa"], 4)
         self.assertEqual(fusion_count_e["mha"], 4)
         self.assertEqual(fusion_count_e["attention"], 4)
@@ -67,7 +67,7 @@ def test_whisper_tiny(self):
         decoder_model.opset_imports["ai.onnxruntime.fusion"] = 1
 
         print(f"Fused {fusion_count_d} ops")
-        self.assertEqual(fusion_count_d["skip_layer_normalization"], 25)
+        self.assertEqual(fusion_count_d["skip_layer_normalization"], 12)
         self.assertEqual(fusion_count_d["sdpa"], 8)
         # 4 self-attention + 4 cross-attention
         self.assertEqual(fusion_count_d["mha"], 8)

onnxscript/rewriter/ort_fusions/attention.py

@@ -198,6 +198,9 @@ def rewrite(
         # attention_bias,
         num_heads,
         # scale,
+        q_mul=None,
+        k_mul=None,
+        v_mul=None,
         **_,
     ):
         # Use bindings to get the values of Dh_q, Dh_k, and Dh_v
@@ -206,6 +209,8 @@ def rewrite(
         # Dh_k = self.bindings.get("Dh_k")
         # Dh_v = self.bindings.get("Dh_v")
         # qkv_hidden_sizes = [Dh_q, Dh_k, Dh_v]
+        if self._no_slice:
+            qkv_weight = op.Concat(q_mul, k_mul, v_mul, axis=0)
 
         if self._has_past:
             attention, present = op.Attention(

onnxscript/rewriter/ort_fusions/skip_normalization.py

@@ -47,35 +47,35 @@ def _skip_layer_norm_pattern(op, input, skip, gamma, beta, epsilon, stash_type):
         epsilon=epsilon,
         stash_type=stash_type,
     )
-    return normalized, skip_sum
+    return normalized
 
 
 def _skip_layer_normalization(op, input, skip, gamma, beta, epsilon, stash_type):
     if stash_type.value != 1:  # FLOAT type
         return None
-    normalized, _mean, _inv_std_var, skip_sum = op.SkipLayerNormalization(
+    normalized, _mean, _inv_std_var = op.SkipLayerNormalization(
         input,
         skip,
         gamma,
         beta,
         epsilon=epsilon,
-        _outputs=4,
+        _outputs=3,
         _domain="com.microsoft",
     )
-    return normalized, skip_sum
+    return normalized
 
 
 # Fusion rule for Add + SkipLayerNormalization
 def _skip_layer_norm_add_bias_pattern(op, input, skip, gamma, beta, bias, epsilon, stash_type):
-    bias_sum = op.Add(input, bias)
-    normalized, _mean, _inv_std_var, skip_sum = op.SkipLayerNormalization(
-        bias_sum,
-        skip,
+    input_with_bias = op.Add(input, bias)
+    skip_sum = op.Add(skip, input_with_bias)
+    normalized = op.LayerNormalization(
+        skip_sum,
         gamma,
         beta,
+        axis=-1,
         epsilon=epsilon,
-        _outputs=4,
-        _domain="com.microsoft",
+        stash_type=stash_type,
     )
     return normalized, skip_sum