Merge branch 'rama/qwen-gqa' into rama/qwen-gqa2

Author: gramalingam

onnxscript/rewriter/ort_fusions/gqa.py

@@ -163,6 +163,13 @@ def pattern(
     ):
         # Reshape query from (B, S, D) to (B, S, H, D/H)
         query_BSHDh = op.Reshape(query_BSD, pattern.ANY_VALUE, _outputs=["query_BSHDh"])
+        # Qwen variant uses normalization of query/key before rotary embedding:
+        # The normalization can happen before (eg., Qwen) or after the Transpose (eg., Gemma).
+        query_BSHDh_normalized = op.SimplifiedLayerNormalization(
+            query_BSHDh, pattern.ANY_VALUE, axis=-1, _outputs=["query_BSHDh_normalized"]
+        )
+        query_BSHDh = pattern.OrValue([query_BSHDh, query_BSHDh_normalized])
+
         # Transpose from (B, S, H, D/H) to (B, H, S, D/H)
         query_BHSDh = op.Transpose(query_BSHDh, perm=[0, 2, 1, 3])
 
@@ -174,6 +181,11 @@ def pattern(
 
         # Reshape key from (B, S, Dkv) to (B, S, Hkv, D/H)
         key_BSHkvDh = op.Reshape(key_BSDkv, pattern.ANY_VALUE, _outputs=["key_BSHkvDh"])
+        key_BSHkvDh_normalized = op.SimplifiedLayerNormalization(
+            key_BSHkvDh, pattern.ANY_VALUE, axis=-1, _outputs=["key_BSHkvDh_normalized"]
+        )
+        key_BSHkvDh = pattern.OrValue([key_BSHkvDh, key_BSHkvDh_normalized])
+
         # Transpose from (B, S, Hkv, D/H) to (B, Hkv, S, D/H)
         key_BHkvSDh = op.Transpose(key_BSHkvDh, perm=[0, 2, 1, 3])
 
@@ -258,8 +270,23 @@ def check(
         query_BSHDh,
         key_BSHkvDh,
         mask,
+        query_BSHDh_normalized=None,
+        query_BHSDh_normalized=None,
+        key_BSHkvDh_normalized=None,
+        key_BHkvSDh_normalized=None,
         **_,
     ):
+        result = pattern.MatchResult()
+        if query_BSHDh_normalized is not None and query_BHSDh_normalized is not None:
+            return result.fail(
+                "Query normalized twice",
+                [query_BSHDh_normalized, query_BHSDh_normalized],
+            )
+        if key_BSHkvDh_normalized is not None and key_BHkvSDh_normalized is not None:
+            return result.fail(
+                "Key normalized twice",
+                [key_BSHkvDh_normalized, key_BHkvSDh_normalized],
+            )
         bindings: dict[str, Dim] = {}
 
         def no_match(val: ir.Value, dims: Sequence[str]) -> bool:
@@ -282,7 +309,7 @@ def no_match(val: ir.Value, dims: Sequence[str]) -> bool:
         # and bindings["H"] * bindings["Dh"] == bindings["H*Dh"]:
         # or check Reshape's shape-input value
 
-        result = pattern.MatchResult()
+        
         num_heads = _ir_utils.get_dim(query_BSHDh, 2)
         kv_num_heads = _ir_utils.get_dim(key_BSHkvDh, 2)
         if not isinstance(num_heads, int):
@@ -334,7 +361,9 @@ def rewrite(
         mask,
         query_BSHDh,
         key_BSHkvDh,
+        query_BSHDh_normalized=None,
         query_BHSDh_normalized=None,
+        key_BSHkvDh_normalized=None,
         key_BHkvSDh_normalized=None,
         **_,
     ):
@@ -356,9 +385,10 @@ def rewrite(
         max_seq_length = op.ReduceMax(seqlens_k, zero_int64_1d, keepdims=0)
         total_seq_length_int32 = op.Add(max_seq_length, one_int32_0d)
 
-        if query_BHSDh_normalized is not None:
+        normalized_query = query_BHSDh_normalized or query_BSHDh_normalized
+        if normalized_query is not None:
             # We apply normalization without the transpose, which is fused into GQA
-            norm_node = query_BHSDh_normalized.producer()
+            norm_node = normalized_query.producer()
             norm_attrs = norm_node.attributes
             norm_scale = norm_node.inputs[1]
             query_BSHDh_normalized = op.SimplifiedLayerNormalization(
@@ -367,9 +397,10 @@ def rewrite(
             reshape_BSHDh_to_BSD = op.Constant(value_ints=[0, 0, -1])
             query_BSD = op.Reshape(query_BSHDh_normalized, reshape_BSHDh_to_BSD)
 
-        if key_BHkvSDh_normalized is not None:
+        normalized_key = key_BHkvSDh_normalized or key_BSHkvDh_normalized
+        if normalized_key is not None:
             # We apply normalization without the transpose, which is fused into GQA
-            norm_node = key_BHkvSDh_normalized.producer()
+            norm_node = normalized_key.producer()
             norm_attrs = norm_node.attributes
             norm_scale = norm_node.inputs[1]
             key_BSHkvDh_normalized = op.SimplifiedLayerNormalization(