onnxruntime-directml 1.24.1__cp314-cp314-win_amd64.whl

onnxruntime/transformers/onnx_model_phi.py

@@ -0,0 +1,929 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from logging import getLogger
+
+import numpy as np
+from dynamo_onnx_helper import DynamoOnnxHelper
+from fusion_base import Fusion
+from fusion_options import AttentionOpType, FusionOptions
+from fusion_skiplayernorm import FusionBiasSkipLayerNormalization, FusionSkipLayerNormalization
+from fusion_utils import NumpyHelper
+from onnx import ModelProto, NodeProto, TensorProto, helper, numpy_helper
+from onnx_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+class ProcessGemmWFunc:
+    def __call__(self, x):
+        return np.transpose(x, (1, 0))
+
+
+class ProcessMatMulQFunc:
+    def __call__(self, x):
+        return np.transpose(np.split(x, 3, 0)[0], (1, 0))
+
+
+class ProcessMatMulKFunc:
+    def __call__(self, x):
+        return np.transpose(np.split(x, 3, 0)[1], (1, 0))
+
+
+class ProcessMatMulVFunc:
+    def __call__(self, x):
+        return np.transpose(np.split(x, 3, 0)[2], (1, 0))
+
+
+class ProcessBiasQFunc:
+    def __call__(self, x):
+        x = np.split(x, 3, -1)[0]
+        return x
+
+
+class ProcessBiasKFunc:
+    def __call__(self, x):
+        x = np.split(x, 3, -1)[1]
+        return x
+
+
+class ProcessBiasVFunc:
+    def __call__(self, x):
+        x = np.split(x, 3, -1)[2]
+        return x
+
+
+class ProcessRotCacheFunc:
+    def __call__(self, x):
+        # half rotary embedding
+        assert len(x.shape) == 2
+        if x.shape[1] == 32:
+            return x[:, 0:16]
+        return x
+
+
+# TODO: move to a separate file
+class Fission(Fusion):
+    def __init__(
+        self,
+        model: OnnxModel,
+        nodes_to_find: list[str],
+    ):
+        super().__init__(model, "DONOTUSE", nodes_to_find)
+
+    def set_attention_op_type(self, attn_op_type: AttentionOpType):
+        self.attn_op_type = attn_op_type
+
+    def get_uname(self, layer_id, name):
+        return name + "_" + str(layer_id)
+
+    def get_edge_by_name(self, edges, name):
+        for edge in edges:
+            if edge == name or edge.endswith(name) or edge.startswith(name):
+                return edge
+        raise ValueError(f"Edge {name} not found")
+
+    def get_input_by_name(self, node, name):
+        return self.get_edge_by_name(node.input, name)
+
+    def get_output_by_name(self, node, name):
+        return self.get_edge_by_name(node.output, name)
+
+    def process_initializer(self, initializer_name, functor, custom_name=None):
+        i = self.model.get_initializer(initializer_name)
+        i_np_array = NumpyHelper.to_array(i)
+        processed_i_np_array = functor(i_np_array)
+        new_tensor = helper.make_tensor(
+            initializer_name + "_processed" if custom_name is None else custom_name,
+            data_type=TensorProto.FLOAT,
+            dims=processed_i_np_array.shape,
+            vals=processed_i_np_array.flatten().tobytes(),
+            raw=True,
+        )
+        self.model.add_initializer(new_tensor, self.this_graph_name)
+        return new_tensor.name
+
+    def add_fp32_value_info(self, name):
+        new_value_info = self.model.graph().value_info.add()
+        new_value_info.name = name
+        new_value_info.type.tensor_type.elem_type = TensorProto.FLOAT
+
+    def add_int64_value_info(self, name):
+        new_value_info = self.model.graph().value_info.add()
+        new_value_info.name = name
+        new_value_info.type.tensor_type.elem_type = TensorProto.INT64
+
+    def replace_fp32_value_info(self, name, shape):
+        for value_info in self.model.graph().value_info:
+            if value_info.name == name:
+                self.model.graph().value_info.remove(value_info)
+                break
+        new_value_info = helper.make_tensor_value_info(
+            name,
+            elem_type=TensorProto.FLOAT,
+            shape=shape,
+        )
+        self.model.graph().value_info.extend([new_value_info])
+
+    def set_unique_name_and_add_nodes(
+        self, subgraph_nodes: list[NodeProto], layer_id: int, layer_known_edges_names: list[str]
+    ):
+        for new_node in subgraph_nodes:
+            for i, name in enumerate(new_node.input):
+                if name == "":
+                    continue
+                elif name not in layer_known_edges_names:
+                    new_node.input[i] = self.get_uname(layer_id, name)
+                    self.add_fp32_value_info(new_node.input[i])
+            for i, name in enumerate(new_node.output):
+                if name == "":
+                    continue
+                elif name not in layer_known_edges_names:
+                    new_node.output[i] = self.get_uname(layer_id, name)
+                    self.add_fp32_value_info(new_node.output[i])
+            new_node.name = self.get_uname(layer_id, new_node.name)
+            self.nodes_to_add.append(new_node)
+            self.node_name_to_graph_name[new_node.name] = self.this_graph_name
+
+    def layernorm(self, inputs: list[str], outputs: list[str], prefix: str = ""):
+        assert len(inputs) == 3
+        assert len(outputs) == 1
+        node = helper.make_node(
+            "LayerNormalization",
+            inputs=inputs,
+            outputs=outputs,
+            name=prefix + "_LayerNormalization",
+            epsilon=9.999999747378752e-06,
+        )
+        return [node]
+
+    def gemm(self, inputs: list[str], outputs: list[str], prefix: str = ""):
+        assert len(inputs) == 3
+        assert len(outputs) == 1
+        matmul = helper.make_node(
+            "MatMul",
+            inputs=[inputs[0], inputs[1]],
+            outputs=[prefix + "matmul_out"],
+            name=prefix + "MatMul",
+        )
+        add = helper.make_node(
+            "Add",
+            inputs=[prefix + "matmul_out", inputs[2]],
+            outputs=outputs,
+            name=prefix + "Bias",
+        )
+        return [matmul, add]
+
+    def rotary(self, inputs: list[str], outputs: list[str], prefix: str = "", rot_dim=32, num_heads=32):
+        assert len(inputs) == 4
+        assert len(outputs) == 1
+        node = helper.make_node(
+            "RotaryEmbedding",
+            inputs=inputs,
+            outputs=outputs,
+            name=prefix + "RotaryEmbedding",
+            domain="com.microsoft",
+            rotary_embedding_dim=rot_dim,
+            num_heads=num_heads,
+        )
+        return [node]
+
+    def fastgelu(self, inputs: list[str], outputs: list[str], prefix: str = ""):
+        assert len(inputs) == 1
+        assert len(outputs) == 1
+        node = helper.make_node(
+            "FastGelu",
+            inputs=inputs,
+            outputs=outputs,
+            name=prefix + "FastGelu",
+            domain="com.microsoft",
+        )
+        return [node]
+
+    def add(self, inputs: list[str], outputs: list[str], prefix: str = ""):
+        assert len(inputs) == 2
+        assert len(outputs) == 1
+        node = helper.make_node(
+            "Add",
+            inputs=inputs,
+            outputs=outputs,
+            name=prefix + "Add",
+        )
+        return [node]
+
+    def mha(self, inputs: list[str], outputs: list[str], prefix: str = "", num_heads=32):
+        assert len(inputs) == 8
+        assert len(outputs) == 3
+        node = helper.make_node(
+            "MultiHeadAttention",
+            inputs=inputs,
+            outputs=outputs,
+            name=prefix + "MultiHeadAttention",
+            domain="com.microsoft",
+            num_heads=num_heads,
+            unidirectional=1,
+        )
+        return [node]
+
+    def gqa(self, inputs: list[str], outputs: list[str], prefix: str = "", num_heads=32):
+        assert len(inputs) == 7
+        assert len(outputs) == 3
+        node = helper.make_node(
+            "GroupQueryAttention",
+            inputs=inputs,
+            outputs=outputs,
+            name=prefix + "GroupQueryAttention",
+            domain="com.microsoft",
+            num_heads=num_heads,
+            kv_num_heads=num_heads,
+        )
+        return [node]
+
+    def attention(self, inputs: list[str], outputs: list[str], prefix: str = "", num_heads=32):
+        assert len(inputs) == 5
+        assert len(outputs) == 2
+        node = helper.make_node(
+            "Attention",
+            inputs=inputs,
+            outputs=outputs,
+            name=prefix + "Attention",
+            domain="com.microsoft",
+            num_heads=num_heads,
+            unidirectional=1,
+            do_rotary=1,
+            rotary_embedding_dim=32,
+        )
+        return [node]
+
+    def paged_attn(
+        self,
+        inputs: list[str],
+        outputs: list[str],
+        prefix: str = "",
+        num_heads=32,
+        head_size=80,
+        scale=0.11180339753627777,
+    ):
+        assert len(inputs) == 6
+        assert len(outputs) == 1
+        node = helper.make_node(
+            "PagedAttention",
+            inputs=inputs,
+            outputs=outputs,
+            name=prefix + "PagedAttention",
+            domain="vllm.ort.ext",
+            num_heads=num_heads,
+            num_kv_heads=num_heads,
+            head_size=head_size,
+            scale=scale,
+        )
+        return [node]
+
+
+class Phi2PreProcessor(DynamoOnnxHelper):
+    def __init__(self, model: ModelProto, num_heads: int, hidden_size: int):
+        super().__init__(model)
+        self.num_hidden_layers = 32
+        self.num_attention_heads = num_heads
+        self.hidden_size = hidden_size
+
+        self.func_name = "modeling_phi_PhiModel_model_1"
+
+    def get_phi2_edge_dict(self) -> dict:
+        edge_dict = {}
+        edge_dict["lm_head_1"] = "logits"
+        edge_dict["l_input_ids_"] = "input_ids"
+        edge_dict["key_states"] = "past_key_0"
+        edge_dict["value_states"] = "past_value_0"
+        for i in range(1, self.num_hidden_layers, 1):
+            edge_dict[f"key_states_{i}"] = f"past_key_{i}"
+            edge_dict[f"value_states_{i}"] = f"past_value_{i}"
+            edge_dict[f"model_layers_{i}_1"] = f"present_key_{i}"
+            edge_dict[f"model_layers_{i}_1_1"] = f"present_value_{i}"
+
+        outputs = [o.name for o in self.model.graph.output]
+        if "model_layers_0_1_1" in outputs and "model_layers_0_1_2" in outputs:
+            edge_dict["model_layers_0_1_1"] = "present_key_0"
+            edge_dict["model_layers_0_1_2"] = "present_value_0"
+        else:
+            assert "model_layers_0_1" in outputs and "model_layers_0_1_1" in outputs
+            edge_dict["model_layers_0_1"] = "present_key_0"
+            edge_dict["model_layers_0_1_1"] = "present_value_0"
+        return edge_dict
+
+    def simplify_phi2_op_type(self):
+        phi2_transformer_layer_name = "modeling_phi_PhiDecoderLayer_model_layers"
+        for node in self.model.graph.node:
+            index = node.op_type.find(phi2_transformer_layer_name)
+            if index != -1:
+                node.op_type = node.op_type[index:]
+
+    def process_graph_io(self, attn_op_type: AttentionOpType):
+        self.use_attn = attn_op_type == AttentionOpType.Attention
+        self.use_vllm = attn_op_type == AttentionOpType.PagedAttention
+        graph = self.model.graph
+        new_inputs = []
+        for vi in graph.input:
+            if "input_ids" in vi.name:
+                vi_iid = helper.make_tensor_value_info(
+                    vi.name,
+                    elem_type=TensorProto.INT32 if not self.use_vllm else TensorProto.INT64,
+                    shape=["batch_size", "seq_len"],
+                )
+                vi_step = helper.make_tensor_value_info(
+                    "step",
+                    elem_type=TensorProto.INT64,
+                    shape=[1],
+                )
+                vi_pid = helper.make_tensor_value_info(
+                    "position_ids",
+                    elem_type=TensorProto.INT64,
+                    shape=["batch_size", "seq_len"],
+                )
+                vi_mask = helper.make_tensor_value_info(
+                    "attention_mask",
+                    elem_type=TensorProto.INT32,
+                    shape=["batch_size", "seq_len"],
+                )
+                vi_meta = helper.make_tensor_value_info(
+                    "input_metadata",
+                    elem_type=TensorProto.INT64,
+                    shape=[1],
+                )
+                (
+                    new_inputs.extend([vi_iid, vi_step, vi_mask])
+                    if not self.use_vllm
+                    else new_inputs.extend([vi_iid, vi_pid, vi_meta])
+                )
+            if self.use_attn:
+                if "past_key" in vi.name:
+                    vi_cache = helper.make_tensor_value_info(
+                        vi.name.replace("past_key", "past"),
+                        elem_type=vi.type.tensor_type.elem_type,
+                        shape=[
+                            2,
+                            "batch_size",
+                            self.num_attention_heads,
+                            "past_seq_len",
+                            self.hidden_size // self.num_attention_heads,
+                        ],
+                    )
+                    new_inputs.extend([vi_cache])
+            elif self.use_vllm:
+                if "past_key" in vi.name:
+                    vi_cache = helper.make_tensor_value_info(
+                        vi.name,
+                        elem_type=vi.type.tensor_type.elem_type,
+                        shape=["num_blocks", "num_heads", "head_size_x", "block_size", "block_x"],
+                    )
+                    new_inputs.extend([vi_cache])
+                if "past_value" in vi.name:
+                    vi_cache = helper.make_tensor_value_info(
+                        vi.name,
+                        elem_type=vi.type.tensor_type.elem_type,
+                        shape=[
+                            "num_blocks",
+                            "num_heads",
+                            "head_size",
+                            "block_size",
+                        ],
+                    )
+                    new_inputs.extend([vi_cache])
+            else:
+                if "past_key" in vi.name or "past_value" in vi.name:
+                    vi_cache = helper.make_tensor_value_info(
+                        vi.name,
+                        elem_type=vi.type.tensor_type.elem_type,
+                        shape=[
+                            "batch_size",
+                            self.num_attention_heads,
+                            "past_seq_len",
+                            self.hidden_size // self.num_attention_heads,
+                        ],
+                    )
+                    new_inputs.extend([vi_cache])
+
+        graph.ClearField("input")
+        graph.input.extend(new_inputs)
+
+        new_outputs = []
+        for i, vi in enumerate(graph.output):
+            if i == 0:
+                new_outputs.extend([vi])
+            else:
+                if self.use_attn:
+                    if "present_key" in vi.name:
+                        vi_cache = helper.make_tensor_value_info(
+                            vi.name.replace("present_key", "present"),
+                            elem_type=vi.type.tensor_type.elem_type,
+                            shape=[
+                                2,
+                                "batch_size",
+                                self.num_attention_heads,
+                                "total_seq_len",
+                                self.hidden_size // self.num_attention_heads,
+                            ],
+                        )
+                        new_outputs.extend([vi_cache])
+                elif self.use_vllm:
+                    pass
+                else:
+                    vi_cache = helper.make_tensor_value_info(
+                        vi.name,
+                        elem_type=vi.type.tensor_type.elem_type,
+                        shape=[
+                            "batch_size",
+                            self.num_attention_heads,
+                            "total_seq_len",
+                            self.hidden_size // self.num_attention_heads,
+                        ],
+                    )
+                    new_outputs.extend([vi_cache])
+
+        graph.ClearField("output")
+        graph.output.extend(new_outputs)
+
+    def preprocess_onnx(self, attn_op_type: AttentionOpType):
+        function_name = None
+        for func in self.model.functions:
+            if func.name.endswith(self.func_name):
+                function_name = func.name
+                break
+        assert function_name is not None
+        self.unroll_function(function_name)
+        self.update_edges(self.get_phi2_edge_dict())
+        self.simplify_phi2_op_type()
+        self.remove_dropout_layer()
+        if attn_op_type == AttentionOpType.PagedAttention:
+            self.remove_lm_head_layer()
+        self.process_graph_io(attn_op_type)
+
+
+class FissionTransformerEmbeddingPhi(Fission):
+    def __init__(
+        self,
+        model: OnnxModel,
+    ):
+        super().__init__(model, ["torch_nn_modules_sparse_Embedding_model_embed_tokens_1"])
+
+    def fuse(self, node, input_name_to_nodes, output_name_to_node):
+        logger.info("Optimizing %s...", node.name)
+
+        assert len(node.input) == 2
+        assert len(node.output) == 1
+
+        input = node.input[0]
+        output = node.output[0]
+
+        embedding = self.get_input_by_name(node, "embed_tokens.weight")
+
+        layer_known_edges_names = [input, output, embedding]
+
+        subgraph_nodes = [
+            helper.make_node(
+                "Gather",
+                inputs=[embedding, input],
+                outputs=[output],
+                name="Embedding_Gather",
+            ),
+        ]
+
+        self.set_unique_name_and_add_nodes(subgraph_nodes, 0, layer_known_edges_names)
+        self.nodes_to_remove.append(node)
+        self.prune_graph = True
+
+
+class FissionTransformerLayerNormPhi(Fission):
+    def __init__(
+        self,
+        model: OnnxModel,
+    ):
+        super().__init__(model, ["torch_nn_modules_normalization_LayerNorm_model_final_layernorm_1"])
+
+    def fuse(self, node, input_name_to_nodes, output_name_to_node):
+        logger.info("Optimizing %s...", node.name)
+
+        assert len(node.input) == 3
+        assert len(node.output) == 1
+
+        input = node.input[0]
+        output = node.output[0]
+
+        ln_weight = self.get_input_by_name(node, "final_layernorm.weight")
+        ln_bias = self.get_input_by_name(node, "final_layernorm.bias")
+
+        layer_known_edges_names = [input, output, ln_weight, ln_bias]
+
+        subgraph_nodes = []
+        subgraph_nodes.extend(self.layernorm([input, ln_weight, ln_bias], [output], "Final"))
+
+        self.set_unique_name_and_add_nodes(subgraph_nodes, 99, layer_known_edges_names)
+
+        self.replace_fp32_value_info(input, ["batch_size", "seq_len", "hidden_size"])
+        self.replace_fp32_value_info(output, ["batch_size", "seq_len", "hidden_size"])
+
+        self.nodes_to_remove.append(node)
+        self.prune_graph = True
+
+
+class FissionTransformerCausalLMHeadPhi(Fission):
+    def __init__(
+        self,
+        model: OnnxModel,
+    ):
+        super().__init__(model, ["torch_nn_modules_linear_Linear_lm_head_1"])
+
+    def fuse(self, node, input_name_to_nodes, output_name_to_node):
+        logger.info("Optimizing %s...", node.name)
+
+        assert len(node.input) == 5
+        assert len(node.output) == 1
+
+        input = node.input[2]
+        output = node.output[0]
+
+        fc_weight = self.process_initializer(self.get_input_by_name(node, "lm_head.weight"), ProcessGemmWFunc())
+        fc_bias = self.get_input_by_name(node, "lm_head.bias")
+
+        layer_known_edges_names = [input, output, fc_weight, fc_bias]
+
+        subgraph_nodes = []
+        subgraph_nodes.extend(self.gemm([input, fc_weight, fc_bias], [output], "LMHead_"))
+
+        self.set_unique_name_and_add_nodes(subgraph_nodes, 99, layer_known_edges_names)
+
+        self.replace_fp32_value_info(input, ["batch_size", "seq_len", "hidden_size"])
+        self.replace_fp32_value_info(output, ["batch_size", "seq_len", 51200])
+
+        self.nodes_to_remove.append(node)
+        self.prune_graph = True
+
+
+class FissionTransformerBlockPhi(Fission):
+    def __init__(
+        self,
+        model: OnnxModel,
+        num_heads: int,
+    ):
+        self.num_heads = num_heads
+        max_num_layers = 32
+        self.func_to_layer_id = {}
+        nodes_to_find = []
+        for layer in range(max_num_layers):
+            func_name = f"modeling_phi_PhiDecoderLayer_model_layers_{layer}_1"
+            nodes_to_find.append(func_name)
+            self.func_to_layer_id[func_name] = layer
+
+        super().__init__(model, nodes_to_find)
+
+    def get_layer_id(self, node):
+        return self.func_to_layer_id[node.op_type]
+
+    def get_gqa_aux_nodes(self):
+        gqa_aux_nodes = [
+            helper.make_node(
+                "Cast",
+                inputs=["attention_mask"],
+                outputs=["mask_int64"],
+                name="Cast_gqa_aux_0",
+                to=TensorProto.INT64,
+            ),
+            helper.make_node(
+                "ReduceSum",
+                inputs=["mask_int64", "one"],
+                outputs=["mask_row_sums"],
+                name="ReduceSum_gqa_aux",
+            ),
+            helper.make_node(
+                "Sub",
+                inputs=["mask_row_sums", "one"],
+                outputs=["seqlens_k_int64"],
+                name="Sub_gqa_aux",
+            ),
+            helper.make_node(
+                "Cast",
+                inputs=["seqlens_k_int64"],
+                outputs=["seqlens_k"],
+                name="Cast_gqa_aux_1",
+                to=TensorProto.INT32,
+            ),
+            helper.make_node("Shape", inputs=["mask_int64"], outputs=["mask_shape"], name="Shape_gqa_aux_0"),
+            helper.make_node(
+                "Gather",
+                inputs=["mask_shape", "one"],
+                outputs=["total_seq_len_int64"],
+                name="Gather_gqa_aux_0",
+                axis=0,
+            ),
+            helper.make_node(
+                "Cast",
+                inputs=["total_seq_len_int64"],
+                outputs=["total_sequence_length"],
+                name="Cast_gqa_aux_2",
+                to=TensorProto.INT32,
+            ),
+        ]
+        return gqa_aux_nodes
+
+    def pack_qkv_gemm(self, q_w, k_w, v_w, q_b, k_b, v_b, weight_name, bias_name):
+        q_weight = self.model.get_initializer(q_w)
+        k_weight = self.model.get_initializer(k_w)
+        v_weight = self.model.get_initializer(v_w)
+        qw = np.transpose(NumpyHelper.to_array(q_weight), (1, 0))
+        kw = np.transpose(NumpyHelper.to_array(k_weight), (1, 0))
+        vw = np.transpose(NumpyHelper.to_array(v_weight), (1, 0))
+        qkv_weight = np.stack((qw, kw, vw), axis=1)
+
+        q_bias = self.model.get_initializer(q_b)
+        k_bias = self.model.get_initializer(k_b)
+        v_bias = self.model.get_initializer(v_b)
+        qb = NumpyHelper.to_array(q_bias)
+        kb = NumpyHelper.to_array(k_bias)
+        vb = NumpyHelper.to_array(v_bias)
+        qkv_bias = np.stack((qb, kb, vb), axis=0)
+
+        hidden_size = qkv_weight.shape[0]
+
+        weight = helper.make_tensor(
+            weight_name,
+            data_type=TensorProto.FLOAT,
+            dims=[hidden_size, hidden_size * 3],
+            vals=qkv_weight.flatten().tobytes(),
+            raw=True,
+        )
+        self.model.add_initializer(weight, self.this_graph_name)
+
+        bias = helper.make_tensor(
+            bias_name,
+            data_type=TensorProto.FLOAT,
+            dims=[hidden_size * 3],
+            vals=qkv_bias.flatten().tobytes(),
+            raw=True,
+        )
+        self.model.add_initializer(bias, self.this_graph_name)
+
+        self.add_fp32_value_info(weight.name)
+        self.add_fp32_value_info(bias.name)
+
+        return weight_name, bias_name
+
+    def fuse(
+        self,
+        node,
+        input_name_to_nodes,
+        output_name_to_node,
+    ):
+        logger.info("Optimizing %s...", node.name)
+
+        logger.info(f"AttentionOpType: {self.attn_op_type}")
+
+        layer_id = self.get_layer_id(node)
+
+        i_hidden_states = node.input[0]
+        i_key_cache = self.get_input_by_name(node, "past_key")
+        i_value_cache = self.get_input_by_name(node, "past_value")
+
+        o_hidden_states = node.output[-1]
+        o_key_cache = self.get_output_by_name(node, "present_key")
+        o_value_cache = self.get_output_by_name(node, "present_value")
+
+        ln_weight = self.get_input_by_name(node, "input_layernorm.weight")
+        ln_bias = self.get_input_by_name(node, "input_layernorm.bias")
+
+        attn_q_weight, attn_q_bias, attn_k_weight, attn_k_bias, attn_v_weight, attn_v_bias = (
+            None,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+        attn_qkv_weight, attn_qkv_bias = None, None
+        cos_cache, sin_cache = None, None
+
+        if self.attn_op_type != AttentionOpType.Attention:
+            attn_q_weight = self.process_initializer(
+                self.get_input_by_name(node, "self_attn.q_proj.weight"), ProcessGemmWFunc()
+            )
+            attn_k_weight = self.process_initializer(
+                self.get_input_by_name(node, "self_attn.k_proj.weight"), ProcessGemmWFunc()
+            )
+            attn_v_weight = self.process_initializer(
+                self.get_input_by_name(node, "self_attn.v_proj.weight"), ProcessGemmWFunc()
+            )
+            attn_q_bias = self.get_input_by_name(node, "self_attn.q_proj.bias")
+            attn_k_bias = self.get_input_by_name(node, "self_attn.k_proj.bias")
+            attn_v_bias = self.get_input_by_name(node, "self_attn.v_proj.bias")
+
+            cos_cache = self.process_initializer(
+                self.get_input_by_name(node, "rotary_emb.cos_cached"), ProcessRotCacheFunc()
+            )
+            sin_cache = self.process_initializer(
+                self.get_input_by_name(node, "rotary_emb.sin_cached"), ProcessRotCacheFunc()
+            )
+        else:
+            attn_qkv_weight, attn_qkv_bias = self.pack_qkv_gemm(
+                self.get_input_by_name(node, "self_attn.q_proj.weight"),
+                self.get_input_by_name(node, "self_attn.k_proj.weight"),
+                self.get_input_by_name(node, "self_attn.v_proj.weight"),
+                self.get_input_by_name(node, "self_attn.q_proj.bias"),
+                self.get_input_by_name(node, "self_attn.k_proj.bias"),
+                self.get_input_by_name(node, "self_attn.v_proj.bias"),
+                self.get_uname(layer_id, "attn_qkv_weight"),
+                self.get_uname(layer_id, "attn_qkv_bias"),
+            )
+
+        attn_out_weight = self.process_initializer(
+            self.get_input_by_name(node, "self_attn.dense.weight"), ProcessGemmWFunc()
+        )
+        attn_out_bias = self.get_input_by_name(node, "self_attn.dense.bias")
+
+        mlp_fc1_weight = self.process_initializer(self.get_input_by_name(node, "mlp.fc1.weight"), ProcessGemmWFunc())
+        mlp_fc2_weight = self.process_initializer(self.get_input_by_name(node, "mlp.fc2.weight"), ProcessGemmWFunc())
+        mlp_fc1_bias = self.get_input_by_name(node, "mlp.fc1.bias")
+        mlp_fc2_bias = self.get_input_by_name(node, "mlp.fc2.bias")
+
+        layer_known_edges_names = []
+        layer_known_edges_names.extend([i_hidden_states, i_key_cache, i_value_cache])
+        layer_known_edges_names.extend([o_hidden_states, o_key_cache, o_value_cache])
+        layer_known_edges_names.extend([ln_weight, ln_bias])
+        if self.attn_op_type != AttentionOpType.Attention:
+            layer_known_edges_names.extend(
+                [
+                    attn_q_weight,
+                    attn_q_bias,
+                    attn_k_weight,
+                    attn_k_bias,
+                    attn_v_weight,
+                    attn_v_bias,
+                    cos_cache,
+                    sin_cache,
+                ]
+            )
+        else:
+            layer_known_edges_names.extend([attn_qkv_weight, attn_qkv_bias])
+        layer_known_edges_names.extend(
+            [attn_out_weight, attn_out_bias, mlp_fc1_weight, mlp_fc1_bias, mlp_fc2_weight, mlp_fc2_bias]
+        )
+        layer_known_edges_names.extend(
+            ["attention_mask", "step", "seqlens_k", "total_sequence_length", "input_metadata", "position_ids"]
+        )
+
+        subgraph_nodes = []
+        subgraph_nodes.extend(self.layernorm([i_hidden_states, ln_weight, ln_bias], ["ln_out"]))
+        subgraph_nodes.extend(self.gemm(["attn_out", attn_out_weight, attn_out_bias], ["attn_add_out"], "OutProj_"))
+        subgraph_nodes.extend(self.gemm(["ln_out", mlp_fc1_weight, mlp_fc1_bias], ["fc1_out"], "FC1_"))
+        subgraph_nodes.extend(self.fastgelu(["fc1_out"], ["gelu_out"]))
+        subgraph_nodes.extend(self.gemm(["gelu_out", mlp_fc2_weight, mlp_fc2_bias], ["fc2_out"], "FC2_"))
+        subgraph_nodes.extend(self.add(["attn_add_out", "fc2_out"], ["residual_1_out"], "Residual_1"))
+        subgraph_nodes.extend(self.add([i_hidden_states, "residual_1_out"], [o_hidden_states], "Residual_2"))
+        if self.attn_op_type != AttentionOpType.Attention:
+            subgraph_nodes.extend(self.gemm(["ln_out", attn_q_weight, attn_q_bias], ["query"], "Q_"))
+            subgraph_nodes.extend(self.gemm(["ln_out", attn_k_weight, attn_k_bias], ["key"], "K_"))
+            subgraph_nodes.extend(self.gemm(["ln_out", attn_v_weight, attn_v_bias], ["value"], "V_"))
+            # vllm engine requires full position ids as the input
+            pos_ids_name = "position_ids" if self.attn_op_type == AttentionOpType.PagedAttention else "step"
+            subgraph_nodes.extend(self.rotary(["query", pos_ids_name, cos_cache, sin_cache], ["query_rot"], "Q_"))
+            subgraph_nodes.extend(self.rotary(["key", pos_ids_name, cos_cache, sin_cache], ["key_rot"], "K_"))
+            if self.attn_op_type == AttentionOpType.MultiHeadAttention:
+                subgraph_nodes.extend(
+                    self.mha(
+                        ["query_rot", "key_rot", "value", "", "attention_mask", "", i_key_cache, i_value_cache],
+                        ["attn_out", o_key_cache, o_value_cache],
+                    )
+                )
+            elif self.attn_op_type == AttentionOpType.GroupQueryAttention:
+                subgraph_nodes.extend(
+                    self.gqa(
+                        [
+                            "query_rot",
+                            "key_rot",
+                            "value",
+                            i_key_cache,
+                            i_value_cache,
+                            "seqlens_k",
+                            "total_sequence_length",
+                        ],
+                        ["attn_out", o_key_cache, o_value_cache],
+                    )
+                )
+                if layer_id == 0:
+                    gqa_aux_nodes = self.get_gqa_aux_nodes()
+                    for new_node in gqa_aux_nodes:
+                        self.nodes_to_add.append(new_node)
+                        self.node_name_to_graph_name[new_node.name] = self.this_graph_name
+                    self.model.add_initializer(
+                        numpy_helper.from_array(np.array([1], dtype="int64"), name="one"), self.this_graph_name
+                    )
+            elif self.attn_op_type == AttentionOpType.PagedAttention:
+                subgraph_nodes.extend(
+                    self.paged_attn(
+                        ["query_rot", "key_rot", "value", i_key_cache, i_value_cache, "input_metadata"],
+                        ["attn_out"],
+                    )
+                )
+        else:
+            past_name = f"past_{layer_id}"
+            present_name = f"present_{layer_id}"
+            layer_known_edges_names.extend([past_name, present_name])
+            subgraph_nodes.extend(
+                self.attention(
+                    ["ln_out", attn_qkv_weight, attn_qkv_bias, "attention_mask", past_name], ["attn_out", present_name]
+                )
+            )
+
+        self.set_unique_name_and_add_nodes(subgraph_nodes, layer_id, layer_known_edges_names)
+
+        self.replace_fp32_value_info(i_hidden_states, ["batch_size", "seq_len", "hidden_size"])
+        self.replace_fp32_value_info(o_hidden_states, ["batch_size", "seq_len", "hidden_size"])
+
+        self.nodes_to_remove.append(node)
+        self.prune_graph = True
+
+
+class PhiOnnxModel(OnnxModel):
+    def __init__(self, model: ModelProto, num_heads: int, hidden_size: int):
+        super().__init__(model)
+        self.phi2_preprocessor = Phi2PreProcessor(self.model, num_heads, hidden_size)
+        self.fission_transformer_block = FissionTransformerBlockPhi(self, num_heads)
+        self.fission_causal_lm_head = FissionTransformerCausalLMHeadPhi(self)
+        self.fission_transformer_layernorm = FissionTransformerLayerNormPhi(self)
+        self.fission_transformer_embedding = FissionTransformerEmbeddingPhi(self)
+
+    def optimize(self, options: FusionOptions | None = None, add_dynamic_axes: bool = False):
+        assert options is not None
+        attn_op_type = options.attention_op_type
+
+        self.fission_transformer_block.set_attention_op_type(attn_op_type)
+
+        self.phi2_preprocessor.preprocess_onnx(attn_op_type)
+
+        self.fission_transformer_block.apply()
+        self.fission_transformer_layernorm.apply()
+        self.fission_causal_lm_head.apply()
+        self.fission_transformer_embedding.apply()
+
+        super().prune_graph()
+
+        # SLN ctor is placed here intentionally to delay the symbolic shape inference
+        self.fuse_sln = FusionSkipLayerNormalization(self)
+        self.fuse_bias_sln = FusionBiasSkipLayerNormalization(self)
+        self.fuse_sln.apply()
+        self.fuse_bias_sln.apply()
+
+    def get_fused_operator_statistics(self):
+        """
+        Returns node count of fused operators.
+        """
+        op_count = {}
+        ops = [
+            "Attention",
+            "MultiHeadAttention",
+            "GroupQueryAttention",
+            "PagedAttention",
+            "Gelu",
+            "BiasGelu",
+            "FastGelu",
+            "LayerNormalization",
+            "SkipLayerNormalization",
+        ]
+        for op in ops:
+            nodes = self.get_nodes_by_op_type(op)
+            op_count[op] = len(nodes)
+
+        logger.info(f"Optimized operators: {op_count}")
+        return op_count
+
+    def is_fully_optimized(self, fused_op_count=None):
+        """
+        Returns True when the model is fully optimized.
+        """
+        if fused_op_count is None:
+            fused_op_count = self.get_fused_operator_statistics()
+
+        def op_count(op_name: str):
+            return fused_op_count.get(op_name) or 0
+
+        attention = (
+            op_count("Attention")
+            + op_count("MultiHeadAttention")
+            + op_count("GroupQueryAttention")
+            + op_count("PagedAttention")
+        )
+        gelu = op_count("Gelu") + op_count("BiasGelu") + op_count("FastGelu")
+        layer_norm = op_count("LayerNormalization") + op_count("SkipLayerNormalization")
+
+        is_perfect = (attention > 0) and (attention == gelu) and (layer_norm >= attention)
+
+        if layer_norm == 0:
+            logger.debug("Layer Normalization not fused")
+
+        if gelu == 0:
+            logger.debug("Gelu (or FastGelu) not fused")
+
+        if attention == 0:
+            logger.warning("Attention (or MultiHeadAttention) not fused")
+
+        return is_perfect