Support Gemma3 with Clip fused attention (microsoft#24280)

### Description
<!-- Describe your changes. -->

Essentially, the vision model is traced differently (this time it's
without mask.), and the input indices of op.Add and op.MatMul can be
different. Also, fp16 and fp32 need different tracing patterns
(op.Cast).

1. Add another traced pattern to CLIP attention to cover no
attention_mask case
2. Accept different index of input on op.Add and op.MatMul (be more
general)
3. fp16 and fp32 shows different pattern (op.Cast after op.Softmax)
4. Refactor test_fastgelu.py to cover torch.onnx.export(...,
dynamo=True)
5. Add gemma3 vision attention (SigLip) test to cover both fp16 and fp32

### Motivation and Context
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->

To optimize Gemma3 multi-modal model, the changes are needed.
https://huggingface.co/google/gemma-3-4b-it

NOTE: some related follow-ups (upstream optimizations to
onnxscript-optimizer):
microsoft/onnxscript#2158
microsoft/onnxscript#2156

Author: titaiwangms

onnxruntime/python/tools/transformers/fusion_attention_clip.py

@@ -126,7 +126,10 @@ def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
                 if node_before_layer_norm is None:
                     continue
                 child = self.model.find_first_child_by_type(
-                    node_before_layer_norm, "LayerNormalization", input_name_to_nodes, False
+                    node_before_layer_norm,
+                    "LayerNormalization",
+                    input_name_to_nodes,
+                    False,
                 )
                 if child is None:
                     continue
@@ -146,19 +149,26 @@ def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
             qkv_nodes = self.model.match_parent_path(
                 normalize_node,
                 ["Add", "MatMul", "Reshape", "Transpose", "MatMul"],
-                [1, 1, 0, 0, 0],
+                [1, None, 0, 0, 0],
             )
             if qkv_nodes is None:
                 logger.debug("fuse_attention: failed to match qkv path")
                 return
-
-        reshape_qkv, transpose_qkv, matmul_qkv = qkv_nodes[2], qkv_nodes[3], qkv_nodes[-1]
+        reshape_qkv, transpose_qkv, matmul_qkv = (
+            qkv_nodes[2],
+            qkv_nodes[3],
+            qkv_nodes[-1],
+        )
 
         v_nodes = self.model.match_parent_path(
-            matmul_qkv, ["Reshape", "Transpose", "Reshape", "Add", "MatMul"], [1, 0, 0, 0, None]
+            matmul_qkv,
+            ["Reshape", "Transpose", "Reshape", "Add", "MatMul"],
+            [1, 0, 0, 0, None],
         )
         if v_nodes is None:
-            v_nodes = self.model.match_parent_path(matmul_qkv, ["Transpose", "Reshape", "Add", "MatMul"], [1, 0, 0, 1])
+            v_nodes = self.model.match_parent_path(
+                matmul_qkv, ["Transpose", "Reshape", "Add", "MatMul"], [1, 0, 0, None]
+            )
             if v_nodes is None:
                 logger.debug("fuse_attention: failed to match v path")
                 return
@@ -182,17 +192,30 @@ def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
             )
             if qk_nodes is None:
                 qk_nodes = self.model.match_parent_path(matmul_qkv, ["Softmax", "Add", "Mul", "MatMul"], [0, 0, 0, 0])
-                if qk_nodes is None:
-                    qk_nodes = self.model.match_parent_path(
-                        matmul_qkv, ["Cast", "Cast", "Softmax", "Add", "Mul", "MatMul"], [0, 0, 0, 0, 0, 0]
-                    )
-                    if qk_nodes is None:
-                        logger.debug("fuse_attention: failed to match qk path")
-                        return
-                    else:
-                        add_mask = qk_nodes[3]
-                else:
+                if qk_nodes is not None:
                     add_mask = qk_nodes[1]
+                else:
+                    # If attention mask is not used, we can still match the qk path.
+                    qk_nodes = self.model.match_parent_path(matmul_qkv, ["Softmax", "Mul", "MatMul"], [0, 0, 0])
+                    if qk_nodes is None:
+                        # Cast nodes are added in the model for fp16.
+                        qk_nodes = self.model.match_parent_path(
+                            matmul_qkv,
+                            ["Cast", "Cast", "Softmax", "Add", "Mul", "MatMul"],
+                            [0, 0, 0, 0, 0, 0],
+                        )
+                        if qk_nodes is not None:
+                            add_mask = qk_nodes[3]
+                        else:
+                            # If attention mask is not used, we can still match the qk path.
+                            qk_nodes = self.model.match_parent_path(
+                                matmul_qkv,
+                                ["Cast", "Cast", "Softmax", "Mul", "MatMul"],
+                                [0, 0, 0, 0, 0],
+                            )
+                            if qk_nodes is None:
+                                logger.debug("fuse_attention: failed to match qk path")
+                                return
         else:
             assert len(add_mask_indices) == 1
             causal_mask_input_index = 1 - add_mask_indices[0]
@@ -201,10 +224,14 @@ def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
         matmul_qk = qk_nodes[-1]
 
         q_nodes = self.model.match_parent_path(
-            matmul_qk, ["Reshape", "Transpose", "Reshape", "Mul", "Add", "MatMul"], [0, 0, 0, 0, None, None]
+            matmul_qk,
+            ["Reshape", "Transpose", "Reshape", "Mul", "Add", "MatMul"],
+            [0, 0, 0, 0, None, None],
         )
         if q_nodes is None:
-            q_nodes = self.model.match_parent_path(matmul_qk, ["Transpose", "Reshape", "Add", "MatMul"], [0, 0, 0, 1])
+            q_nodes = self.model.match_parent_path(
+                matmul_qk, ["Transpose", "Reshape", "Add", "MatMul"], [0, 0, 0, None]
+            )
             if q_nodes is None:
                 logger.debug("fuse_attention: failed to match q path")
                 return
@@ -216,10 +243,14 @@ def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
         add_q, matmul_q = q_nodes[-2], q_nodes[-1]
 
         k_nodes = self.model.match_parent_path(
-            matmul_qk, ["Transpose", "Reshape", "Transpose", "Reshape", "Add", "MatMul"], [1, 0, 0, 0, 0, None]
+            matmul_qk,
+            ["Transpose", "Reshape", "Transpose", "Reshape", "Add", "MatMul"],
+            [1, 0, 0, 0, 0, None],
         )
         if k_nodes is None:
-            k_nodes = self.model.match_parent_path(matmul_qk, ["Transpose", "Reshape", "Add", "MatMul"], [1, 0, 0, 1])
+            k_nodes = self.model.match_parent_path(
+                matmul_qk, ["Transpose", "Reshape", "Add", "MatMul"], [1, 0, 0, None]
+            )
             if k_nodes is None:
                 logger.debug("fuse_attention: failed to match k path")
                 return
@@ -242,7 +273,17 @@ def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
             # 4D Add after Q x K'
             add_qk_nodes = self.model.match_parent_path(
                 add_mask,
-                ["Where", "Sub", "Cast", "Expand", "Unsqueeze", "Unsqueeze", "Reshape", "Reshape", "Cast"],
+                [
+                    "Where",
+                    "Sub",
+                    "Cast",
+                    "Expand",
+                    "Unsqueeze",
+                    "Unsqueeze",
+                    "Reshape",
+                    "Reshape",
+                    "Cast",
+                ],
                 [1, 2, 1, 0, 0, 0, 0, 0, 0],
             )
             if add_qk_nodes is not None:

onnxruntime/python/tools/transformers/fusion_fastgelu.py

@@ -177,13 +177,12 @@ def fuse_2(self, tanh_node, input_name_to_nodes: dict, output_name_to_node: dict
             return
         mul_after_mul_half = children[0]
 
+        # root_node could be None when root_input is graph input
         root_node = self.model.get_parent(
             mul_after_mul_half,
             0 if mul_after_mul_half.input[1] == mul_half.output[0] else 1,
             output_name_to_node,
         )
-        if root_node is None:
-            return
 
         mul_before_tanh = self.model.match_parent(tanh_node, "Mul", 0, output_name_to_node)
         if mul_before_tanh is None:
@@ -197,7 +196,13 @@ def fuse_2(self, tanh_node, input_name_to_nodes: dict, output_name_to_node: dict
         if add_before_tanh is None:
             return
 
-        mul_after_pow = self.model.match_parent(add_before_tanh, "Mul", None, output_name_to_node, exclude=[root_node])
+        mul_after_pow = self.model.match_parent(
+            add_before_tanh,
+            "Mul",
+            None,
+            output_name_to_node,
+            exclude=[root_node] if root_node else [],
+        )
         if mul_after_pow is None:
             return
 
@@ -212,7 +217,9 @@ def fuse_2(self, tanh_node, input_name_to_nodes: dict, output_name_to_node: dict
         if not self.model.has_constant_input(pow, 3.0):
             return
 
-        if pow.input[0] != root_node.output[0]:
+        root_input = mul_after_mul_half.input[0 if mul_after_mul_half.input[1] == mul_half.output[0] else 1]
+
+        if pow.input[0] != root_input:
             return
 
         subgraph_nodes = [
@@ -236,7 +243,7 @@ def fuse_2(self, tanh_node, input_name_to_nodes: dict, output_name_to_node: dict
         self.nodes_to_remove.extend(subgraph_nodes)
         fused_node = helper.make_node(
             "FastGelu",
-            inputs=[root_node.output[0]],
+            inputs=[root_input],
             outputs=mul_after_mul_half.output,
             name=self.model.create_node_name("FastGelu"),
         )

onnxruntime/test/python/transformers/test_data/models/gemma3-vision-attention_fp16.onnx

@@ -3,6 +3,7 @@
 import unittest
 
 import torch
+from parameterized import parameterized
 from parity_utilities import find_transformers_source
 
 if find_transformers_source():
@@ -43,16 +44,6 @@ def forward(self, x):
         return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x)))
 
 
-test_cases = [
-    ("huggingface", "Gelu", HuggingfaceGelu),
-    ("huggingface", "FastGelu", HuggingfaceFastGelu),
-    ("huggingface", "QuickGelu", HuggingfaceQuickGelu),
-    ("huggingface", "FastGelu", HuggingfaceTorchGeluTanh),
-    ("megatron", "Gelu", MegatronGelu),
-    ("megatron", "FastGelu", MegatronFastGelu),
-]
-
-
 class TestGeluFusions(unittest.TestCase):
     def verify_node_count(self, bert_model, expected_node_count, test_name):
         for op_type, count in expected_node_count.items():
@@ -62,25 +53,47 @@ def verify_node_count(self, bert_model, expected_node_count, test_name):
                     print(f"{op}: {len(bert_model.get_nodes_by_op_type(op))} expected={counter}")
             self.assertEqual(len(bert_model.get_nodes_by_op_type(op_type)), count)
 
-    def test_fusions(self):
-        for test_case in test_cases:
-            source, operator, model_class = test_case
-            model = model_class()
-            dummy_input = torch.ones(3, dtype=torch.float32)
-            test_name = f"{operator}_{source}"
-            onnx_path = f"{test_name}.onnx"
-            torch.onnx.export(
-                model,
-                (dummy_input),
-                onnx_path,
-                input_names=["input"],
-                output_names=["output"],
-            )
-            optimizer = optimize_model(onnx_path, "bert")
-            # optimizer.save_model_to_file(f"{operator}_{source}_opt.onnx")
-            os.remove(onnx_path)
-            expected_node_count = {operator: 1}
-            self.verify_node_count(optimizer, expected_node_count, test_name)
+    @parameterized.expand(
+        [
+            (("huggingface", "Gelu", HuggingfaceGelu), True),
+            (("huggingface", "FastGelu", HuggingfaceFastGelu), True),
+            (("huggingface", "QuickGelu", HuggingfaceQuickGelu), True),
+            (("huggingface", "FastGelu", HuggingfaceTorchGeluTanh), True),
+            (("megatron", "Gelu", MegatronGelu), True),
+            (("megatron", "FastGelu", MegatronFastGelu), True),
+            (("huggingface", "Gelu", HuggingfaceGelu), False),
+            (("huggingface", "FastGelu", HuggingfaceFastGelu), False),
+            (("huggingface", "QuickGelu", HuggingfaceQuickGelu), False),
+            (("huggingface", "FastGelu", HuggingfaceTorchGeluTanh), False),
+            (("megatron", "Gelu", MegatronGelu), False),
+            (("megatron", "FastGelu", MegatronFastGelu), False),
+        ]
+    )
+    def test_fusions(self, test_case, dynamo):
+        source, operator, model_class = test_case
+        model = model_class()
+        dummy_input = torch.ones(3, dtype=torch.float32)
+        test_name = f"{operator}_{source}"
+        onnx_path = f"{test_name}.onnx"
+        torch.onnx.export(
+            model,
+            (dummy_input,),
+            onnx_path,
+            input_names=["input"],
+            output_names=["output"],
+            dynamo=dynamo,
+            optimize=True,  # Only meaningful when dynamo is True
+        )
+        optimizer = optimize_model(onnx_path, "bert")
+        # optimizer.save_model_to_file(f"{operator}_{source}_opt.onnx")
+        os.remove(onnx_path)
+        # Remove the associated .data file (dynamo)
+        data_path = onnx_path + ".data"
+        if os.path.exists(data_path):
+            os.remove(data_path)
+        expected_node_count = {operator: 1}
+
+        self.verify_node_count(optimizer, expected_node_count, test_name)
 
 
 if __name__ == "__main__":