onnxruntime-directml 1.20.0__cp313-cp313-win_amd64.whl

onnxruntime/transformers/fusion_bart_attention.py

@@ -0,0 +1,640 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+import logging
+
+import numpy as np
+from fusion_attention import AttentionMask, FusionAttention
+from onnx import TensorProto, helper
+from onnx_model import OnnxModel
+
+logger = logging.getLogger(__name__)
+
+
+class FusionBartAttention(FusionAttention):
+    """
+    Fuse Bart Attention subgraph into one Attention node.
+    """
+
+    def __init__(
+        self,
+        model: OnnxModel,
+        hidden_size: int,
+        num_heads: int,
+        attention_mask: AttentionMask,
+    ):
+        super().__init__(model, hidden_size, num_heads, attention_mask)
+
+    def check_runtime_shape_path(
+        self,
+        reshape_qkv_2,
+        reshape_qkv_1,
+        reshape_q_2,
+        reshape_k_2,
+        reshape_v_2,
+        root_input,
+    ):
+        concat_qkv_2_path = self.model.match_parent_path(reshape_qkv_2, ["Concat"], [1])
+        if concat_qkv_2_path is None:
+            return False
+        concat_qkv_2 = concat_qkv_2_path[0]
+
+        reshape_qkv_2_path_1 = self.model.match_parent_path(concat_qkv_2, ["Unsqueeze", "Gather", "Shape"], [0, 0, 0])
+        reshape_qkv_2_path_2 = self.model.match_parent_path(concat_qkv_2, ["Unsqueeze", "Gather", "Shape"], [1, 0, 0])
+        if reshape_qkv_2_path_1 is None or reshape_qkv_2_path_2 is None:
+            return False
+
+        _, gather_1, shape_1 = reshape_qkv_2_path_1
+        _, gather_2, shape_2 = reshape_qkv_2_path_2
+
+        if shape_1.input[0] != root_input or shape_2.input[0] != root_input:
+            return False
+
+        reshape_qkv_1_path_1 = self.model.match_parent_path(reshape_qkv_1, ["Concat", "Unsqueeze", "Gather"], [1, 0, 0])
+        reshape_qkv_1_path_2 = self.model.match_parent_path(reshape_qkv_1, ["Concat", "Unsqueeze", "Gather"], [1, 2, 0])
+        if reshape_qkv_1_path_1 is None or reshape_qkv_1_path_2 is None:
+            return False
+        if reshape_qkv_1_path_1[-1].name != gather_1.name or reshape_qkv_1_path_2[-1].name != gather_2.name:
+            return False
+
+        reshape_q_2_path = self.model.match_parent_path(reshape_q_2, ["Concat", "Unsqueeze", "Mul"], [1, 0, 0])
+        reshape_k_2_path = self.model.match_parent_path(reshape_k_2, ["Concat", "Unsqueeze", "Mul"], [1, 0, 0])
+        reshape_v_2_path = self.model.match_parent_path(reshape_v_2, ["Concat", "Unsqueeze", "Mul"], [1, 0, 0])
+        if reshape_q_2_path is None or reshape_k_2_path is None or reshape_v_2_path is None:
+            return False
+
+        mul_q = reshape_q_2_path[-1]
+        mul_k = reshape_k_2_path[-1]
+        mul_v = reshape_v_2_path[-1]
+
+        gather_1_out = gather_1.output[0]
+        if mul_q.input[0] != gather_1_out or mul_k.input[0] != gather_1_out or mul_v.input[0] != gather_1_out:
+            return False
+
+        return True
+
+    def check_runtime_shape_path_openai(
+        self,
+        reshape_qkv_2,
+        matmul_qkv,
+        add_qk,
+        matmul_qk,
+        add_q,
+    ):
+        reshape_qkv_2_path = self.model.match_parent_path(
+            reshape_qkv_2, ["Concat", "Slice", "Gather", "Shape"], [1, 0, 0, 0]
+        )
+        if reshape_qkv_2_path is None:
+            return False
+        else:
+            if reshape_qkv_2_path[-1].input[0] != matmul_qkv.output[0]:
+                return False
+
+        matmul_qk_path_1 = self.model.match_parent_path(
+            matmul_qk, ["Mul", "Pow", "Cast", "Div", "Gather", "Shape"], [0, 1, 0, 0, 0, 0]
+        )
+        matmul_qk_path_2 = self.model.match_parent_path(
+            matmul_qk, ["Mul", "Pow", "Cast", "Div", "Gather", "Shape"], [1, 1, 0, 0, 0, 0]
+        )
+        if matmul_qk_path_1 is None or matmul_qk_path_2 is None:
+            return False
+
+        mul_1 = matmul_qk_path_1[0]
+        mul_2 = matmul_qk_path_2[0]
+        if mul_1.input[1] != mul_2.input[1]:
+            return False
+        if matmul_qk_path_1[-1].input[0] != add_q.output[0] and matmul_qk_path_2[-1].input[0] != add_q.output[0]:
+            return False
+
+        # For decoder attentions only
+        if add_qk is not None:
+            add_qk_path = self.model.match_parent_path(add_qk, ["Slice"], [1])
+            if add_qk_path is None:
+                return False
+            slice_q_path_1 = self.model.match_parent_path(
+                add_qk_path[0], ["Slice", "Unsqueeze", "Gather", "Shape"], [0, 2, 0, 0]
+            )
+            slice_q_path_2 = self.model.match_parent_path(add_qk_path[0], ["Unsqueeze", "Gather", "Shape"], [2, 0, 0])
+            if slice_q_path_1 is None and slice_q_path_2 is None:
+                return False
+            _, unsqueeze_1, _, _ = slice_q_path_1
+            unsqueeze_2, _, _ = slice_q_path_2
+            if unsqueeze_1.input[0] != unsqueeze_2.input[0]:
+                return False
+            if slice_q_path_1[-1].input[0] != add_q.output[0] and slice_q_path_2[-1].input[0] != add_q.output[0]:
+                return False
+
+        return True
+
+    def fuse(self, normalize_node, input_name_to_nodes, output_name_to_node):
+        # Track if fusion is occurring for OpenAI implementation of Whisper
+        model_impl_openai = False
+
+        # SkipLayerNormalization has two inputs, and one of them is the root input for attention.
+        qkv_nodes = self.model.match_parent_path(
+            normalize_node,
+            ["Add", "MatMul", "Reshape", "Transpose", "Reshape", "MatMul"],
+            [1, 1, 0, 0, 0, 0],
+        )
+        qkv_nodes_openai = self.model.match_parent_path(
+            normalize_node,
+            ["Add", "MatMul", "Reshape", "Transpose", "MatMul"],
+            [1, 1, 0, 0, 0],
+        )
+        if qkv_nodes is not None:
+            (
+                add_out,
+                matmul_out,
+                reshape_qkv_2,
+                transpose_qkv,
+                reshape_qkv_1,
+                matmul_qkv,
+            ) = qkv_nodes
+        elif qkv_nodes_openai is not None:
+            qkv_nodes = qkv_nodes_openai
+            (
+                add_out,
+                matmul_out,
+                reshape_qkv_2,
+                transpose_qkv,
+                matmul_qkv,
+            ) = qkv_nodes
+            # Set model implementation to openai
+            model_impl_openai = True
+        else:
+            return
+
+        other_inputs = []
+        for input in normalize_node.input:
+            if input not in output_name_to_node:
+                continue
+            if input == qkv_nodes[0].output[0]:
+                continue
+            other_inputs.append(input)
+        if len(other_inputs) != 1:
+            return
+        root_input = other_inputs[0]
+
+        # Sometimes the input name to the attention MatMul nodes does not match the input name to the end
+        # SkipLayerNormalization node (name saved in root_input). We find the true input name to the MatMul
+        # nodes by getting the initial SkipLayerNormalization node and checking how many MatMul nodes are
+        # children nodes for each of its output names.
+        """
+                                        root_input
+                    +---------------------------------------------------+
+                    |                                                   |
+                    |                                                   |
+        SkipLayerNormalization --> Attention --> MatMul --> SkipLayerNormalization
+        """
+        skip_layernorm = output_name_to_node[root_input]
+        # For some attention blocks, the end SkipLayerNormalization node may point to an Add node whose
+        # child is the LayerNormalization node.
+        if skip_layernorm.op_type == "Add":
+            skip_layernorm = self.model.get_children(skip_layernorm)[0]
+        for output in skip_layernorm.output:
+            if not output:
+                continue
+            children = input_name_to_nodes[output]
+            children_types = [child.op_type for child in children]
+            if children_types.count("MatMul") >= 1:
+                root_input = output
+                break
+
+        graph_input_names = set([node.name for node in self.model.graph().input])
+        graph_output_names = set([node.name for node in self.model.graph().output])
+
+        v_nodes = self.model.match_parent_path(
+            matmul_qkv,
+            ["Reshape", "Transpose", "Reshape", "Add", "MatMul"],
+            [1, 0, 0, 0, None],
+        )
+        v_nodes_openai = self.model.match_parent_path(
+            matmul_qkv,
+            ["Transpose", "Reshape", "Add", "MatMul"],
+            [1, 0, 0, None],
+        )
+        v_nodes_with_past_self_attn = self.model.match_parent_path(
+            # Decoder attention with past value concatenated before MatMul
+            matmul_qkv,
+            ["Reshape", "Concat", "Transpose", "Reshape", "Add", "MatMul"],
+            [1, 0, 1, 0, 0, None],
+        )
+        v_nodes_with_past_cross_attn = self.model.match_parent_path(
+            # Decoder attention with past value directly used in MatMul
+            matmul_qkv,
+            ["Reshape"],
+            [1],
+        )
+        v_nodes_with_past_cross_attn_openai = self.model.match_parent_path(
+            matmul_qkv,
+            ["Transpose", "Reshape", "Reshape", "Transpose"],
+            [1, 0, 0, 0],
+        )
+        past_v, present_v = "", ""
+        reshape_v_2, add_v = None, None
+        if v_nodes is not None:
+            (reshape_v_2, transpose_v, reshape_v_1, add_v, matmul_v) = v_nodes
+            # For initial pass through encoder-decoder_with_past to get starting past values (beam search)
+            present_v = transpose_v.output[0]
+        elif v_nodes_openai is not None:
+            v_nodes = v_nodes_openai
+            (transpose_v, reshape_v_1, add_v, matmul_v) = v_nodes
+            # For initial pass through encoder-decoder_with_past to get starting past values (beam search)
+
+            # Find the child path to access the correct present_v values
+            # Openai impl provides present/past v values in 3D format
+            # whereas ort MultiHeadAttention expects v values in 4D, hence the
+            # additional Reshape and Transpose nodes are added
+            # For encoder attention types
+            # Add -> Reshape -> Transpose -> Present_V
+            reshape_path = self.model.match_child_path(
+                add_v,
+                ["Reshape", "Transpose"],
+                exclude=[reshape_v_1],
+            )
+            # For decoder attention types
+            # add_v_node                     Reshape <- Transpose <-Past_V
+            #           \                  /
+            #             \              /
+            #               -> Concat <-
+            #                    |
+            #                    |--> Reshape -> Transpose -> Present_V
+            concat_path = self.model.match_child_path(add_v, ["Concat", "Reshape", "Transpose"])
+            if reshape_path is not None:
+                (_, transpose_add_v) = reshape_path
+                if transpose_add_v.output[0] in graph_output_names:
+                    present_v = transpose_add_v.output[0]
+            if concat_path is not None:
+                (concat_v, _, transpose_concat_v) = concat_path
+                if transpose_concat_v.output[0] in graph_output_names:
+                    present_v = transpose_concat_v.output[0]
+                concat_nodes = self.model.match_parent_path(concat_v, ["Reshape", "Transpose"], [0, 0])
+                _, transpose_concat_v_in = concat_nodes
+                past_v = transpose_concat_v_in.input[0]
+        elif v_nodes_with_past_self_attn is not None:
+            (reshape_v_2, concat_v, transpose_v, reshape_v_1, add_v, matmul_v) = v_nodes_with_past_self_attn
+            v_nodes = v_nodes_with_past_self_attn
+            past_v = concat_v.input[0]
+            present_v = concat_v.output[0]
+        elif (
+            v_nodes_with_past_cross_attn is not None and v_nodes_with_past_cross_attn[-1].input[0] in graph_input_names
+        ):
+            v_nodes = v_nodes_with_past_cross_attn
+            past_v = v_nodes[-1].input[0]
+            present_v = v_nodes[-1].output[0]
+            if present_v not in graph_output_names:
+                identity_node_v = list(
+                    filter(lambda node: node.op_type == "Identity", self.model.input_name_to_nodes()[past_v])
+                )
+                present_v = identity_node_v[0].output[0] if len(identity_node_v) == 1 else ""
+        elif (
+            v_nodes_with_past_cross_attn_openai is not None
+            and v_nodes_with_past_cross_attn_openai[-1].input[0] in graph_input_names
+        ):
+            v_nodes = v_nodes_with_past_cross_attn_openai
+            past_v = v_nodes[-1].input[0]
+            present_v = v_nodes[-1].output[0]
+            if present_v not in graph_output_names:
+                identity_node_v = list(
+                    filter(lambda node: node.op_type == "Identity", self.model.input_name_to_nodes()[past_v])
+                )
+                present_v = identity_node_v[0].output[0] if len(identity_node_v) == 1 else ""
+        else:
+            logger.debug("fuse_attention: failed to match v path")
+            return
+        past_v = past_v if past_v in graph_input_names else ""
+        present_v = present_v if present_v in graph_output_names else ""
+
+        qk_nodes_1 = self.model.match_parent_path(matmul_qkv, ["Softmax", "MatMul"], [0, 0])
+        qk_nodes_2 = self.model.match_parent_path(
+            matmul_qkv, ["Softmax", "Reshape", "Add", "Reshape", "MatMul"], [0, 0, 0, 0, 0]
+        )
+        qk_nodes_2_openai = self.model.match_parent_path(matmul_qkv, ["Softmax", "Add", "MatMul"], [0, 0, 0])
+        add_qk = None
+        if qk_nodes_1 is not None:
+            _, matmul_qk = qk_nodes_1
+            qk_nodes = qk_nodes_1
+        elif qk_nodes_2 is not None:
+            _, _, add_qk, _, matmul_qk = qk_nodes_2
+            qk_nodes = qk_nodes_2
+        elif qk_nodes_2_openai is not None:
+            _, add_qk, matmul_qk = qk_nodes_2_openai
+            qk_nodes = qk_nodes_2_openai
+        else:
+            return
+
+        q_nodes = self.model.match_parent_path(
+            matmul_qk,
+            ["Reshape", "Transpose", "Reshape", "Mul", "Add", "MatMul"],
+            [0, 0, 0, 0, 0, 1],
+        )
+        q_nodes_openai = self.model.match_parent_path(
+            matmul_qk,
+            ["Mul", "Transpose", "Reshape", "Add", "MatMul"],
+            [0, 0, 0, 0, 1],
+        )
+        reshape_q_2 = None
+        if q_nodes is not None:
+            reshape_q_2, transpose_q, reshape_q_1, mul_q, add_q, matmul_q = q_nodes
+        elif q_nodes_openai is not None:
+            q_nodes = q_nodes_openai
+            mul_q, transpose_q, reshape_q_1, add_q, matmul_q = q_nodes
+        else:
+            return
+
+        k_nodes_with_bias = self.model.match_parent_path(
+            matmul_qk,
+            ["Transpose", "Reshape", "Transpose", "Reshape", "Add", "MatMul"],
+            [1, 0, 0, 0, 0, 1],
+        )
+        k_nodes_with_bias_openai = self.model.match_parent_path(
+            matmul_qk,
+            ["Mul", "Transpose", "Reshape", "MatMul"],
+            [1, 0, 0, 0],
+        )
+        k_nodes_no_bias = self.model.match_parent_path(
+            matmul_qk,
+            ["Transpose", "Reshape", "Transpose", "Reshape", "MatMul"],
+            [1, 0, 0, 0, 0],
+        )
+        k_nodes_no_bias_with_past_self_attn = self.model.match_parent_path(
+            # Decoder attention with past key concatenated before MatMul
+            matmul_qk,
+            ["Transpose", "Reshape", "Concat", "Transpose", "Reshape", "MatMul"],
+            [1, 0, 0, 1, 0, 0],
+        )
+        k_nodes_no_bias_with_past_cross_attn = self.model.match_parent_path(
+            # Decoder attention with past key directly used in MatMul
+            matmul_qk,
+            ["Transpose", "Reshape"],
+            [1, 0],
+        )
+        k_nodes_no_bias_with_past_cross_attn_openai = self.model.match_parent_path(
+            # Decoder attention with past key directly used in MatMul
+            matmul_qk,
+            ["Mul", "Transpose", "Reshape", "Reshape", "Transpose"],
+            [1, 0, 0, 0, 0],
+        )
+        past_k, present_k = "", ""
+        reshape_k_2, reshape_k_1, matmul_k = None, None, None
+        if k_nodes_with_bias is not None:
+            _, reshape_k_2, transpose_k_1, reshape_k_1, add_k, matmul_k = k_nodes_with_bias
+            k_nodes = k_nodes_with_bias
+        elif k_nodes_with_bias_openai is not None:
+            mul_k, transpose_k_1, reshape_k_1, matmul_k = k_nodes_with_bias_openai
+            k_nodes = k_nodes_with_bias_openai
+            present_k = matmul_k.output[0]
+
+            # Find the child path to access the correct present_k values
+            # Openai impl provides present/past k values in 3D format
+            # whereas ort MultiHeadAttention expects k values in 4D, hence the
+            # additional Reshape and Transpose nodes are added
+            # For encoder attention types
+            # Matmul -> Reshape -> Transpose -> Present_K
+            reshape_path = self.model.match_child_path(
+                matmul_k,
+                ["Reshape", "Transpose"],
+                exclude=[reshape_k_1],
+            )
+            # For decoder attention types
+            # matmul_k_node                  Reshape <- Transpose <- Past_K
+            #           \                  /
+            #             \              /
+            #               -> Concat <-
+            #                    |
+            #                    |--> Reshape -> Transpose -> Present_K
+            concat_path = self.model.match_child_path(matmul_k, ["Concat", "Reshape", "Transpose"])
+            if reshape_path is not None:
+                (_, transpose_matmul_k) = reshape_path
+                if transpose_matmul_k.output[0] in graph_output_names:
+                    present_k = transpose_matmul_k.output[0]
+            if concat_path is not None:
+                (concat_k, _, transpose_concat_k) = concat_path
+                if transpose_concat_k.output[0] in graph_output_names:
+                    present_k = transpose_concat_k.output[0]
+                concat_nodes = self.model.match_parent_path(concat_k, ["Reshape", "Transpose"], [0, 0])
+                _, transpose_concat_k_in = concat_nodes
+                past_k = transpose_concat_k_in.input[0]
+        elif k_nodes_no_bias is not None:
+            _, reshape_k_2, transpose_k_1, reshape_k_1, matmul_k = k_nodes_no_bias
+            k_nodes = k_nodes_no_bias
+            # For initial pass through encoder-decoder_with_past to get starting past values (beam search)
+            present_k = transpose_k_1.output[0]
+        elif k_nodes_no_bias_with_past_self_attn is not None:
+            _, reshape_k_2, concat_k, _, reshape_k_1, matmul_k = k_nodes_no_bias_with_past_self_attn
+            k_nodes = k_nodes_no_bias_with_past_self_attn
+            past_k = concat_k.input[0]
+            present_k = concat_k.output[0]
+        elif (
+            k_nodes_no_bias_with_past_cross_attn is not None
+            and k_nodes_no_bias_with_past_cross_attn[-1].input[0] in graph_input_names
+        ):
+            k_nodes = k_nodes_no_bias_with_past_cross_attn
+            past_k = k_nodes[-1].input[0]
+            present_k = k_nodes[-1].output[0]
+            if present_k not in graph_output_names:
+                identity_node_k = list(
+                    filter(lambda node: node.op_type == "Identity", self.model.input_name_to_nodes()[past_k])
+                )
+                present_k = identity_node_k[0].output[0] if len(identity_node_k) == 1 else ""
+        elif (
+            k_nodes_no_bias_with_past_cross_attn_openai is not None
+            and k_nodes_no_bias_with_past_cross_attn_openai[-1].input[0] in graph_input_names
+        ):
+            k_nodes = k_nodes_no_bias_with_past_cross_attn_openai
+            past_k = k_nodes[-1].input[0]
+            present_k = k_nodes[-1].output[0]
+            if present_k not in graph_output_names:
+                identity_node_k = list(
+                    filter(lambda node: node.op_type == "Identity", self.model.input_name_to_nodes()[past_k])
+                )
+                present_k = identity_node_k[0].output[0] if len(identity_node_k) == 1 else ""
+        else:
+            return
+        past_k = past_k if past_k in graph_input_names else ""
+        present_k = present_k if present_k in graph_output_names else ""
+
+        if k_nodes in (k_nodes_with_bias_openai, k_nodes_no_bias, k_nodes_no_bias_with_past_self_attn):
+            # Create empty Add node for attention graph
+            bias_dim = self.model.get_initializer(add_v.input[0]).dims[0]
+            empty_bias_name = "empty_bias"
+            empty_tensor = self.model.get_initializer(empty_bias_name)
+            if empty_tensor is None:
+                self.add_initializer(
+                    empty_bias_name,
+                    TensorProto.FLOAT,
+                    dims=[bias_dim],
+                    vals=np.array([0.0] * bias_dim, dtype=np.float32),
+                )
+
+            add_name = self.model.create_node_name("Add")
+            add_k = helper.make_node("Add", [empty_bias_name, matmul_k.output[0]], [reshape_k_1.name], add_name)
+
+        if (
+            model_impl_openai
+            and not past_k
+            and not self.check_runtime_shape_path_openai(
+                reshape_qkv_2,
+                matmul_qkv,
+                add_qk,
+                matmul_qk,
+                add_q,
+            )
+        ):
+            return
+        elif (
+            not model_impl_openai
+            and not past_k
+            and not self.check_runtime_shape_path(
+                reshape_qkv_2,
+                reshape_qkv_1,
+                reshape_q_2,
+                reshape_k_2,
+                reshape_v_2,
+                root_input,
+            )
+        ):
+            return
+
+        three_root_inputs = past_k and past_v and matmul_k is None and "matmul_v" not in locals()
+        one_root_input = (
+            not three_root_inputs
+            and matmul_k.input[0] == root_input
+            and matmul_q.input[0] == root_input
+            and matmul_v.input[0] == root_input
+        )
+        two_root_inputs = (
+            not three_root_inputs
+            and matmul_q.input[0] == root_input
+            and matmul_k.input[0] == matmul_v.input[0]
+            and matmul_k.input[0] != matmul_q.input[0]
+        )
+
+        # There are 5 types of attention:
+        # 1) Encoder attention with one_root_input=True and qk_nodes=qk_nodes_1
+        # 2) Decoder attention with one_root_input=True and qk_nodes=qk_nodes_2
+        # 3) Decoder attention with past with one_root_input=True and qk_nodes=qk_nodes_1 and past_k=past_decoder_key and past_v=past_decoder_value
+        # 4) Decoder cross attention with two_root_inputs=True and qk_nodes=qk_nodes_1
+        # 5) Decoder cross attention with past with three_root_inputs=True and qk_nodes=qk_nodes_1
+        encoder_attention = one_root_input and qk_nodes == qk_nodes_1
+        decoder_attention = one_root_input and qk_nodes in (qk_nodes_2, qk_nodes_2_openai)
+        decoder_attention_with_past = (
+            (encoder_attention if not model_impl_openai else decoder_attention) and past_k and past_v
+        )
+        decoder_cross_attention = two_root_inputs and qk_nodes == qk_nodes_1
+        decoder_cross_attention_with_past = three_root_inputs and qk_nodes == qk_nodes_1
+
+        # For decoder_attention, the attention mask needs to be included in the attention node
+        mask_index = None
+        if decoder_attention:
+            mask_nodes_bart = self.model.match_parent_path(
+                add_qk,
+                ["Where"],
+                [1],
+            )
+            mask_nodes_whisper = self.model.match_parent_path(
+                add_qk,
+                ["Expand", "Unsqueeze", "Unsqueeze", "Where"],
+                [1, 0, 0, 0],
+            )
+            if mask_nodes_whisper is not None:
+                mask_index = mask_nodes_whisper[0].output[-1]
+            elif mask_nodes_bart is not None:
+                mask_index = mask_nodes_bart[0].output[-1]
+
+        if (
+            encoder_attention
+            or decoder_attention
+            or decoder_attention_with_past
+            or decoder_cross_attention
+            or decoder_cross_attention_with_past
+        ):
+            attention_last_node = reshape_qkv_2
+            num_heads, hidden_size = self.get_num_heads_and_hidden_size(reshape_q_1)
+
+            if num_heads <= 0 or hidden_size <= 0 or (hidden_size % num_heads) != 0:
+                logger.debug("fuse_attention: failed to detect num_heads or hidden_size")
+                return
+
+            new_node = None
+            if decoder_attention_with_past or decoder_cross_attention or decoder_cross_attention_with_past:
+                # Note: Decoder attention with past key and past value is fused as multihead attention
+                # rather than attention because multihead attention supports separate past key and past
+                # value whereas attention supports concatenated past key and past value.
+                new_node = (
+                    self.create_multihead_attention_node(
+                        matmul_q,
+                        matmul_k if decoder_cross_attention or decoder_attention_with_past else past_k,
+                        matmul_v if decoder_cross_attention or decoder_attention_with_past else past_v,
+                        add_q,
+                        add_k if decoder_cross_attention or decoder_attention_with_past else None,
+                        add_v if decoder_cross_attention or decoder_attention_with_past else None,
+                        num_heads,
+                        hidden_size,
+                        attention_last_node.output[0],
+                        past_k=past_k if decoder_attention_with_past else "",
+                        past_v=past_v if decoder_attention_with_past else "",
+                        present_k=present_k,
+                        present_v=present_v,
+                        packed_qkv=decoder_attention_with_past,
+                    )
+                    if self.use_multi_head_attention
+                    else None
+                )
+            else:
+                # Temporarily set multihead attention flag to false
+                use_multi_head_attention_ground_truth = self.use_multi_head_attention
+                self.use_multi_head_attention = False
+                new_node = self.create_attention_node(
+                    None,
+                    matmul_q,
+                    matmul_k,
+                    matmul_v,
+                    add_q,
+                    add_k,
+                    add_v,
+                    num_heads,
+                    hidden_size,
+                    root_input,
+                    attention_last_node.output[0],
+                    add_qk_str=mask_index if decoder_attention else None,
+                    past_k=past_k,
+                    past_v=past_v,
+                    present_k=present_k,
+                    present_v=present_v,
+                )
+                self.use_multi_head_attention = use_multi_head_attention_ground_truth
+            if new_node is None:
+                return
+
+            self.nodes_to_add.append(new_node)
+            self.node_name_to_graph_name[new_node.name] = self.this_graph_name
+
+            self.nodes_to_remove.extend([attention_last_node, transpose_qkv, matmul_qkv])
+            self.nodes_to_remove.extend(qk_nodes)
+
+            # When using multihead attention, keep MatMul nodes in original graph
+            if decoder_attention_with_past or decoder_cross_attention or decoder_cross_attention_with_past:
+                if q_nodes[-1].op_type == "MatMul":
+                    q_nodes.pop()
+                if k_nodes[-1].op_type == "MatMul":
+                    k_nodes.pop()
+                if v_nodes[-1].op_type == "MatMul":
+                    v_nodes.pop()
+                if self.disable_multi_head_attention_bias and (
+                    decoder_cross_attention or decoder_cross_attention_with_past
+                ):
+                    if q_nodes[-1].op_type == "Add":
+                        q_nodes.pop()
+                    if k_nodes[-1].op_type == "Add":
+                        k_nodes.pop()
+                    if v_nodes[-1].op_type == "Add":
+                        v_nodes.pop()
+
+            self.nodes_to_remove.extend(q_nodes)
+            self.nodes_to_remove.extend(k_nodes)
+            self.nodes_to_remove.extend(v_nodes)
+
+            # Use prune graph to remove mask nodes since they are shared by all attention nodes.
+            self.prune_graph = True