Add MHA fusion for Nemotron speech conformer encoder (#27764)

### Description

This PR updates the pattern matchings to perform multi-head attention
fusion for the conformer encoder inside [Nemotron
speech](https://huggingface.co/nvidia/nemotron-speech-streaming-en-0.6b).

<img width="550" height="976" alt="image"
src="https://github.com/user-attachments/assets/a194308e-ce69-4128-9389-aae2a64b312f"
/>

### Motivation and Context

These changes allow the `MultiHeadAttention` op to appear in the encoder
ONNX model.

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>

Author: kunal-vaishnavi

onnxruntime/python/tools/transformers/fusion_conformer_attention.py

@@ -32,8 +32,14 @@ def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
             [1, None, 0, 0, 0],
         )
         if qkv_nodes is None:
-            logger.debug("fuse_conformer_attention: failed to match qkv path")
-            return
+            qkv_nodes = self.model.match_parent_path(
+                normalize_node,
+                ["MatMul", "Reshape", "Transpose", "MatMul"],
+                [1, 0, 0, 0],
+            )
+            if qkv_nodes is None:
+                logger.debug("fuse_conformer_attention: failed to match qkv path")
+                return
 
         reshape_qkv, transpose_qkv, matmul_qkv = qkv_nodes[-3], qkv_nodes[-2], qkv_nodes[-1]
 
@@ -50,33 +56,50 @@ def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
                 [1, 0, 0, 0],
             )
             if v_nodes is None:
-                logger.debug("fuse_conformer_attention: failed to match v path")
-                return
+                v_nodes = self.model.match_parent_path(
+                    matmul_qkv,
+                    ["Transpose", "Reshape", "MatMul"],
+                    [1, 0, 0],
+                )
+                if v_nodes is None:
+                    logger.debug("fuse_conformer_attention: failed to match v path")
+                    return
         else:
             concat_v = v_nodes[0]
             concat_parent = self.model.get_parent(concat_v, 0, None)
             present_v = concat_v.output[0]
             past_v = concat_parent.output[0]
 
-        add_v, matmul_v = v_nodes[-2], v_nodes[-1]
+        add_v = v_nodes[-2] if len(v_nodes) >= 2 and v_nodes[-2].op_type == "Add" else None
+        matmul_v = v_nodes[-1]
 
         attn_mask = ""
         qk_nodes = self.model.match_parent_path(
             matmul_qkv,
             ["Softmax", "Add", "MatMul"],
             [0, 0, 0],
         )
+        where_qk = None
         if qk_nodes is None:
             qk_nodes = self.model.match_parent_path(
                 matmul_qkv,
                 ["Where", "Softmax", "Where", "Add", "MatMul"],
                 [0, 2, 0, 2, 0],
             )
             if qk_nodes is None:
-                logger.debug("fuse_conformer_attention: failed to match qk path")
-                return
+                qk_nodes = self.model.match_parent_path(
+                    matmul_qkv,
+                    ["Where", "Softmax", "Where", "Div", "Add", "MatMul"],
+                    [0, 2, 0, 2, 0, 0],
+                )
+                if qk_nodes is None:
+                    logger.debug("fuse_conformer_attention: failed to match qk path")
+                    return
+                where_qk = qk_nodes[2]
+            else:
+                where_qk = qk_nodes[2]
 
-            where_qk = qk_nodes[2]
+        if where_qk is not None:
             mask_nodes = self.model.match_parent_path(
                 where_qk,
                 ["Equal", "Unsqueeze", "Cast"],
@@ -99,20 +122,46 @@ def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
                 [0, 0, 0, 0, 0],
             )
             if q_nodes is None:
-                logger.debug("fuse_conformer_attention: failed to match q path")
-                return
+                q_nodes = self.model.match_parent_path(
+                    matmul_qk,
+                    ["Transpose", "Add", "Reshape", "MatMul"],
+                    [0, 0, 0, 1],
+                )
+                if q_nodes is None:
+                    q_nodes = self.model.match_parent_path(
+                        matmul_qk,
+                        ["Transpose", "Add", "Reshape", "MatMul"],
+                        [0, 0, 0, 0],
+                    )
+                    if q_nodes is None:
+                        logger.debug("fuse_conformer_attention: failed to match q path")
+                        return
 
-        reshape_q, add_q, matmul_q = q_nodes[-3], q_nodes[-2], q_nodes[-1]
+        reshape_q = next((node for node in q_nodes if node.op_type == "Reshape"), None)
+        add_q = next((node for node in q_nodes if node.op_type == "Add"), None)
+        matmul_q = next((node for node in reversed(q_nodes) if node.op_type == "MatMul"), None)
+        if reshape_q is None or add_q is None or matmul_q is None:
+            logger.debug("fuse_conformer_attention: failed to identify q reshape/add/matmul nodes")
+            return
 
         extra_q_nodes = self.model.match_parent_path(
             add_qk,
             ["Reshape", "Transpose", "MatMul", "Transpose", "Reshape", "Div"],
             [1, 0, 0, 0, 0, 0],
         )
-        if extra_q_nodes is not None and q_nodes[0] != extra_q_nodes[-1]:
+        if extra_q_nodes is not None and q_nodes[0].op_type in ["Div", "Mul"] and q_nodes[0] != extra_q_nodes[-1]:
             logger.debug("fuse_conformer_attention: failed to match extra q path")
             return
 
+        if extra_q_nodes is None:
+            nemotron_extra_q_nodes = self.model.match_parent_path(
+                add_qk,
+                ["Slice", "Reshape", "Slice", "Reshape", "Pad", "MatMul", "Transpose", "Add"],
+                [1, 0, 0, 0, 0, 0, 0, 0],
+            )
+            if nemotron_extra_q_nodes is not None:
+                extra_q_nodes = nemotron_extra_q_nodes
+
         past_k, present_k = "", ""
         k_nodes = self.model.match_parent_path(
             matmul_qk,
@@ -132,24 +181,50 @@ def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
                     [1, 0, 0, 0],
                 )
                 if k_nodes is None:
-                    logger.debug("fuse_conformer_attention: failed to match k path")
-                    return
+                    k_nodes = self.model.match_parent_path(
+                        matmul_qk,
+                        ["Transpose", "Reshape", "MatMul"],
+                        [1, 0, 0],
+                    )
+                    if k_nodes is None:
+                        logger.debug("fuse_conformer_attention: failed to match k path")
+                        return
         else:
             concat_k = k_nodes[1]
             concat_parent = self.model.get_parent(concat_k, 0, None)
             past_k = concat_parent.output[0]
             present_k = concat_k.output[0]
 
-        add_k, matmul_k = k_nodes[-2], k_nodes[-1]
+        add_k = k_nodes[-2] if len(k_nodes) >= 2 and k_nodes[-2].op_type == "Add" else None
+        matmul_k = k_nodes[-1]
 
         num_heads, hidden_size = self.get_num_heads_and_hidden_size(reshape_q)
         if num_heads <= 0 or hidden_size <= 0 or (hidden_size % num_heads) != 0:
             logger.debug("fuse_conformer_attention: failed to detect num_heads or hidden_size")
             return
 
+        # Validate attention_bias: the Attention and MultiHeadAttention kernels require a 4-D
+        # tensor with shape [batch_size or 1, num_heads or 1, sequence_length, total_sequence_length].
+        # Scalar or 1-D initializers (e.g. a plain QK scaling constant) must not be forwarded as
+        # attention_bias. Non-initializer values (computed positional-bias outputs) are kept as-is.
+        attention_bias = add_qk.input[1]
+        bias_init = self.model.get_initializer(attention_bias)
+        if bias_init is not None and len(bias_init.dims) != 4:
+            logger.debug(
+                "fuse_conformer_attention: skipping attention_bias %s with dims %s (expected 4-D)",
+                attention_bias,
+                list(bias_init.dims),
+            )
+            attention_bias = ""
+
         new_node = None
         use_packed_attention_op = (
-            matmul_q.input[0] == matmul_k.input[0] and matmul_k.input[0] == matmul_v.input[0] and extra_q_nodes is None
+            matmul_q.input[0] == matmul_k.input[0]
+            and matmul_k.input[0] == matmul_v.input[0]
+            and extra_q_nodes is None
+            and add_q is not None
+            and add_k is not None
+            and add_v is not None
         )
         if use_packed_attention_op:
             # Self-attention, use Attention op
@@ -165,7 +240,7 @@ def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
                 hidden_size=hidden_size,
                 first_input=matmul_q.input[0],
                 output=reshape_qkv.output[0],
-                add_qk_str=add_qk.input[1],
+                add_qk_str=attention_bias,
                 past_k=past_k,
                 past_v=past_v,
                 present_k=present_k,
@@ -183,7 +258,7 @@ def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
                 hidden_size=hidden_size,
                 output=reshape_qkv.output[0],
                 key_padding_mask=attn_mask,
-                add_qk=add_qk.input[1],
+                add_qk=attention_bias,
                 past_k=past_k,
                 past_v=past_v,
                 present_k=present_k,