rewrite cross attention logic

Author: shubhambhokare1

onnxscript/rewriter/ort_fusions/mha.py

@@ -42,7 +42,7 @@ def __init__(
         is_rotary: bool,
         use_mask: bool,
         has_past_present: bool,
-        is_cross_attention_from_past: bool,
+        is_cross_attention: bool,
     ):
         super().__init__(name)
         self._double_transpose = double_transpose
@@ -51,20 +51,14 @@ def __init__(
         self._is_rotary = is_rotary
         self._use_mask = use_mask
         self._has_past_present = has_past_present
-        # Checks for cross-attention pattern when cross
-        # query and key originate from past_key and past_value.
-        self._is_cross_attention_from_past = is_cross_attention_from_past
-        # Store the key/value to check if the cross-attention is
-        # indeed from past_key and past_value.
-        self._k_from_past = None
-        self._v_from_past = None
+        self._is_cross_attention = is_cross_attention
 
     def pattern(
         self,
         op,
         query_BSD,
-        key_BSD,
-        value_BSD,
+        key,
+        value,
         mask,
         past_key,
         past_value,
@@ -83,23 +77,28 @@ def pattern(
         # Transpose from (B, S, H, D/H) to (B, H, S, D/H)
         query_BHSDh = op.Transpose(query_BSHDh, perm=[0, 2, 1, 3])
 
-        # Reshape from (B, S, D) to (B, S, H, D/H)
-        key_BSHDh = op.Reshape(key_BSD, pattern.ANY_VALUE, _outputs=["key_BSHDh"])
-
-        # Possible Transpose patterns for key:
-        # This scenario optimizes the need for a double transpose
-        # 1. (B, S, H, D/H) -> (B, H, D/H, S)
-        # Patterns with double transpose of key
-        # Double transpose should handle this optimization
-        # 2. (B, S, H, D/H) -> (B, H, S, D/H) -> (B, H, D/H, S)
-        # Patterns where key is reshaped to 3D, transposed and reshaped back to 4D
-        # 3. (B, S, H, D/H) -> (B, H, S, D/H) -> R (B, S, D) -> (B, D, S) -> R (B, H, D/H, S)
-        key_BHSDh = op.Transpose(key_BSHDh, perm=key_perm)
-
-        # Reshape from (B, S, D) to (B, S, H, D/H)
-        value_BSHDh = op.Reshape(value_BSD, pattern.ANY_VALUE, _outputs=["value_BSHDh"])
-        # Transpose from (B, S, H, D/H) to (B, H, S, D/H)
-        value_BHSDh = op.Transpose(value_BSHDh, perm=[0, 2, 1, 3])
+        if not self._is_cross_attention:
+            # Reshape from (B, S, D) to (B, S, H, D/H)
+            key_BSHDh = op.Reshape(key, pattern.ANY_VALUE, _outputs=["key_BSHDh"])
+
+            # Possible Transpose patterns for key:
+            # This scenario optimizes the need for a double transpose
+            # 1. (B, S, H, D/H) -> (B, H, D/H, S)
+            # Patterns with double transpose of key
+            # Double transpose should handle this optimization
+            # 2. (B, S, H, D/H) -> (B, H, S, D/H) -> (B, H, D/H, S)
+            # Patterns where key is reshaped to 3D, transposed and reshaped back to 4D
+            # 3. (B, S, H, D/H) -> (B, H, S, D/H) -> R (B, S, D) -> (B, D, S) -> R (B, H, D/H, S)
+            key_BHSDh = op.Transpose(key_BSHDh, perm=key_perm)
+
+            # Reshape from (B, S, D) to (B, S, H, D/H)
+            value_BSHDh = op.Reshape(value, pattern.ANY_VALUE, _outputs=["value_BSHDh"])
+            # Transpose from (B, S, H, D/H) to (B, H, S, D/H)
+            value_BHSDh = op.Transpose(value_BSHDh, perm=[0, 2, 1, 3])
+        else:
+            # For cross-attention, key and value are not reshaped
+            key_BHSDh = key
+            value_BHSDh = value
 
         if self._is_rotary:
             # This is workaround for examples where there is a duplication of Unsqueeze op
@@ -117,9 +116,12 @@ def pattern(
             query_BHSDh_emb = op.RotaryEmbedding(
                 query_BHSDh, position_ids_q, cos, sin, _domain="com.microsoft"
             )
-            key_BHSDh_emb = op.RotaryEmbedding(
-                key_BHSDh, position_ids_k, cos, sin, _domain="com.microsoft"
-            )
+            if not self._is_cross_attention:
+                key_BHSDh_emb = op.RotaryEmbedding(
+                    key_BHSDh, position_ids_k, cos, sin, _domain="com.microsoft"
+                )
+            else:
+                key_BHSDh_emb = key_BHSDh
         else:
             # If rotary embedding is not used, we fuse with positional_embeddings
             query_BHSDh_emb = query_BHSDh
@@ -130,23 +132,13 @@ def pattern(
         if self._has_past_present:
             key_seq = op.Concat(past_key, key_BHSDh_emb, axis=-2)
         else:
-            # For patterns where cross-attention key/value originates from past_key/past_value
-            if self._is_cross_attention_from_past:
-                key_seq = past_key
-                self._k_from_past = key_seq
-            else:
-                key_seq = key_BHSDh_emb
+            key_seq = key_BHSDh_emb
 
         # Concatenate past_value cache and current value
         if self._has_past_present:
             value_seq = op.Concat(past_value, value_BHSDh, axis=-2)
         else:
-            # For patterns where cross-attention key/value originates from past_key/past_value
-            if self._is_cross_attention_from_past:
-                value_seq = past_value
-                self._v_from_past = value_seq
-            else:
-                value_seq = value_BHSDh
+            value_seq = value_BHSDh
 
         # Key/value to be used for dot-product attention computation
         key_seq_to_sdpa = key_seq
@@ -198,8 +190,8 @@ def check(
         self,
         op,
         query_BSD,
-        key_BSD,
-        value_BSD,
+        key,
+        value,
         mask,
         past_key,
         past_value,
@@ -221,97 +213,57 @@ def no_match(val: ir.Value, dims: Sequence[str]) -> bool:
                 f"Shape mismatch: {query_BSD} does not match expected dimensions ['B', 'S', 'D']",
                 query_BSD,
             )
-        # If cross-attention key/value originates from past_key/past_value,
-        # Check if their producer is None, this is done to avoid from the matcher assuming
-        # that if a key/value pattern path does not exist, it is a cross-attention pattern.
-        if self._is_cross_attention_from_past:
-            if self._k_from_past is not None:
-                if self._k_from_past.producer() is not None:
-                    return check_result.fail(
-                        "Key is not from past_key/past_value. This is not a cross-attention pattern.",
-                    )
-            if self._v_from_past is not None:
-                if self._v_from_past.producer() is not None:
-                    return check_result.fail(
-                        "Value is not from past_key/past_value. This is not a cross-attention pattern.",
-                    )
-            # We only consider patterns where,
-            # 1) double_transpose = True, to avoid pattern consuming the transpose of key.
-            # 2) is_rotary = False, as if rotary embeddings are used, the key is not from past_key.
-            # TODO: Determine what parameter conditions would make this pattern full-proof.
-            if not self._double_transpose or self._is_rotary:
-                return check_result.fail(
-                    "Key is not from past_key/past_value. This is not a cross-attention pattern.",
-                )
 
-        """
-        # Check for key transpose values
-        k_perm = _ir_utils.get_singleton_value(key_perm)
-        if k_perm is None or not isinstance(k_perm, list):
-            return check_result.fail(
-                f"Key permutation is not a list.",
-                key_perm,
-            )
-        if len(k_perm) != 4:
+        if no_match(query_BSHDh, ["B", "S", "H", "Dh"]):
             return check_result.fail(
-                f"Key permutation is not of length 4.",
-                key_perm,
+                f"Shape mismatch: {query_BSHDh} does not match expected dimensions ['B', 'S', 'H', 'Dh']",
+                query_BSHDh,
             )
-        if self._double_transpose:
-            if k_perm != [0, 2, 1, 3]:
+        # If cross-attention key/value shapes are 4D
+        if self._is_cross_attention:
+            if no_match(key, ["B", "H", "Skv", "Dh"]):
                 return check_result.fail(
-                    f"Key permutation is not [0, 2, 1, 3].",
-                    key_perm,
+                    f"Shape mismatch: {key} does not match expected dimensions ['B', 'H', 'Skv', 'Dh']",
+                    key,
                 )
-        else:
-            if k_perm != [0, 2, 3, 1]:
+            if no_match(value, ["B", "H", "Skv", "Dv"]):
                 return check_result.fail(
-                    f"Key permutation is not [0, 2, 3, 1].",
-                    key_perm,
+                    f"Shape mismatch: {value} does not match expected dimensions ['B', 'H', 'Skv', 'Dv']",
+                    value,
                 )
-        """
-
-        if not self._is_cross_attention_from_past:
-            if no_match(key_BSD, ["B", "Skv", "D"]):
+            # Ensure that no past_key/past_value is used in cross-attention
+            if past_key is not None:
                 return check_result.fail(
-                    f"Shape mismatch: {key_BSD} does not match expected dimensions ['B', 'Skv', 'D']",
-                    query_BSD,
-                )
-            if no_match(value_BSD, ["B", "Skv", "D"]):
-                return check_result.fail(
-                    f"Shape mismatch: {value_BSD} does not match expected dimensions ['B', 'Skv', 'D']",
-                    value_BSD,
-                )
-
-        if self._has_past_present:
-            if no_match(past_key, ["B", "H", "Spast", "Dh"]):
-                return check_result.fail(
-                    f"Shape mismatch: {past_key} does not match expected dimensions ['B', 'H', 'Spast', 'Dh']",
+                    "past_key should be None in cross-attention.",
                     past_key,
                 )
-            if no_match(past_value, ["B", "H", "Spast", "Dv"]):
+            if past_value is not None:
                 return check_result.fail(
-                    f"Shape mismatch: {past_value} does not match expected dimensions ['B', 'H', 'Spast', 'Dv']",
+                    "past_value should be None in cross-attention.",
                     past_value,
                 )
-
-        if no_match(query_BSHDh, ["B", "S", "H", "Dh"]):
-            return check_result.fail(
-                f"Shape mismatch: {query_BSHDh} does not match expected dimensions ['B', 'S', 'H', 'Dh']",
-                query_BSHDh,
-            )
-
-        if not self._is_cross_attention_from_past:
-            if key_BSHDh and no_match(key_BSHDh, ["B", "S", "H", "Dh"]):
+        else:
+            if no_match(key, ["B", "Skv", "D"]):
                 return check_result.fail(
-                    f"Shape mismatch: {key_BSHDh} does not match expected dimensions ['B', 'S', 'H', 'Dh']",
-                    query_BSHDh,
+                    f"Shape mismatch: {key} does not match expected dimensions ['B', 'Skv', 'D']",
+                    query_BSD,
                 )
-            if value_BSHDh and no_match(value_BSHDh, ["B", "S", "H", "Dh"]):
+            if no_match(value, ["B", "Skv", "D"]):
                 return check_result.fail(
-                    f"Shape mismatch: {value_BSHDh} does not match expected dimensions ['B', 'S', 'H', 'Dh']",
-                    query_BSHDh,
+                    f"Shape mismatch: {value} does not match expected dimensions ['B', 'Skv', 'D']",
+                    value,
                 )
+            if self._has_past_present:
+                if no_match(past_key, ["B", "H", "Spast", "Dh"]):
+                    return check_result.fail(
+                        f"Shape mismatch: {past_key} does not match expected dimensions ['B', 'H', 'Spast', 'Dh']",
+                        past_key,
+                    )
+                if no_match(past_value, ["B", "H", "Spast", "Dv"]):
+                    return check_result.fail(
+                        f"Shape mismatch: {past_value} does not match expected dimensions ['B', 'H', 'Spast', 'Dv']",
+                        past_value,
+                    )
 
         # TODO: mask shape check: ideally, it should be (1 or B, 1 or H, S, St)
         # But this also, unforunately, depends on ORT version.
@@ -326,8 +278,8 @@ def rewrite(
         self,
         op,
         query_BSD,
-        key_BSD,
-        value_BSD,
+        key,
+        value,
         mask,
         past_key,
         past_value,
@@ -353,35 +305,21 @@ def rewrite(
             query_BSD_emb = op.RotaryEmbedding(
                 query_BSD, position_ids, cos, sin, _domain="com.microsoft"
             )
-            key_BSD_emb = op.RotaryEmbedding(
-                key_BSD, position_ids, cos, sin, _domain="com.microsoft"
-            )
+            if not self._is_cross_attention:
+                key_BSD_emb = op.RotaryEmbedding(
+                    key, position_ids, cos, sin, _domain="com.microsoft"
+                )
+            else:
+                key_BSD_emb = key
         else:
             query_BSD_emb = query_BSD
-            key_BSD_emb = key_BSD
+            key_BSD_emb = key
 
         num_outputs = 1 + (2 * self._has_past_present)
-        # Special case for cross-attention that comes from past_key/past_value
-        if self._is_cross_attention_from_past:
-            return op.MultiHeadAttention(
-                query_BSD_emb,
-                past_key,
-                past_value,
-                None,  # bias
-                None,  # key padding mask
-                mask,  # attention mask/bias
-                None,
-                None,
-                num_heads=num_heads,
-                scale=scale,
-                _domain="com.microsoft",
-                _outputs=num_outputs,
-            )
-
         return op.MultiHeadAttention(
             query_BSD_emb,
             key_BSD_emb,
-            value_BSD,
+            value,
             None,  # bias
             None,  # key padding mask
             mask,  # attention mask/bias
@@ -402,7 +340,7 @@ def rewrite(
         "is_rotary": is_rotary,
         "use_mask": use_mask,
         "has_past_present": has_past_present,
-        "is_cross_attention_from_past": is_cross_attention_from_past,
+        "is_cross_attention": is_cross_attention,
     }
     for double_transpose in [False, True]
     for transpose_4d in (
@@ -411,9 +349,9 @@ def rewrite(
     for pre_scale_q in [True, False]
     for is_rotary in [False, True]
     for use_mask in [False, True]
-    # TODO: Avoid this parameter from being order dependent
+    # Enforce has_past_present to be True first, to avoid missing the pattern
     for has_past_present in [True, False]
-    for is_cross_attention_from_past in [False, True]
+    for is_cross_attention in [False, True]
 ]
 
 # Dynamically create the rules
@@ -426,7 +364,7 @@ def rewrite(
             f"{'_Rotary' if params['is_rotary'] else ''}"
             f"{'_Masked' if params['use_mask'] else ''}"
             f"{'_Past' if params['has_past_present'] else ''}"
-            f"{'_CrossAttentionFromPast' if params['is_cross_attention_from_past'] else ''}",
+            f"{'_CrossAttention' if params['is_cross_attention'] else ''}",
             **params,
         )
         for params in parameter_combinations