§
    kŠtjRŠ  ã                   ó0  — d dl mZ d dlZd dlmZ d dlmZ d dlm	Z	m
Z
 d dlmZmZ d dlmZ d dlmZmZmZmZmZ d d	lmZ  ee¦  «        Z G d
„ d¦  «        Z G d„ d¦  «        Z G d„ d¦  «        Z G d„ d¦  «        Z G d„ d¦  «        Z G d„ d¦  «        Z G d„ d¦  «        Z  G d„ d¦  «        Z! G d„ de¦  «        Z" G d„ de¦  «        Z# G d„ de"¦  «        Z$ G d „ d!e"¦  «        Z% G d"„ d#e"¦  «        Z& G d$„ d%e"¦  «        Z' G d&„ d'e¦  «        Z(dS )(é    )Ú	getLoggerN)ÚDynamoOnnxHelper)ÚFusion)ÚAttentionOpTypeÚFusionOptions)Ú FusionBiasSkipLayerNormalizationÚFusionSkipLayerNormalization)ÚNumpyHelper)Ú
ModelProtoÚ	NodeProtoÚTensorProtoÚhelperÚnumpy_helper)Ú	OnnxModelc                   ó   — e Zd Zd„ ZdS )ÚProcessGemmWFuncc                 ó,   — t          j        |d¦  «        S )N©é   r   )ÚnpÚ	transpose©ÚselfÚxs     úe/var/www/html/CA-Chatbot/venv/lib/python3.11/site-packages/onnxruntime/transformers/onnx_model_phi.pyÚ__call__zProcessGemmWFunc.__call__   s   € ÝŒ|˜A˜vÑ&Ô&Ð&ó    N©Ú__name__Ú
__module__Ú__qualname__r   © r   r   r   r      s#   € € € € € ð'ð 'ð 'ð 'ð 'r   r   c                   ó   — e Zd Zd„ ZdS )ÚProcessMatMulQFuncc                 ó`   — t          j        t          j        |dd¦  «        d         d¦  «        S )Né   r   r   ©r   r   Úsplitr   s     r   r   zProcessMatMulQFunc.__call__   ó'   € ÝŒ|�BœH Q¨¨1Ñ-Ô-¨aÔ0°&Ñ9Ô9Ð9r   Nr   r"   r   r   r$   r$      ó#   € € € € € ð:ð :ð :ð :ð :r   r$   c                   ó   — e Zd Zd„ ZdS )ÚProcessMatMulKFuncc                 ó`   — t          j        t          j        |dd¦  «        d         d¦  «        S )Nr&   r   r   r   r'   r   s     r   r   zProcessMatMulKFunc.__call__   r)   r   Nr   r"   r   r   r,   r,      r*   r   r,   c                   ó   — e Zd Zd„ ZdS )ÚProcessMatMulVFuncc                 ó`   — t          j        t          j        |dd¦  «        d         d¦  «        S )Nr&   r   é   r   r'   r   s     r   r   zProcessMatMulVFunc.__call__$   r)   r   Nr   r"   r   r   r/   r/   #   r*   r   r/   c                   ó   — e Zd Zd„ ZdS )ÚProcessBiasQFuncc                 ó>   — t          j        |dd¦  «        d         }|S )Nr&   éÿÿÿÿr   ©r   r(   r   s     r   r   zProcessBiasQFunc.__call__)   ó   € ÝŒH�Q˜˜2ÑÔ˜qÔ!ˆØˆr   Nr   r"   r   r   r3   r3   (   ó#   € € € € € ðð ð ð ð r   r3   c                   ó   — e Zd Zd„ ZdS )ÚProcessBiasKFuncc                 ó>   — t          j        |dd¦  «        d         }|S )Nr&   r5   r   r6   r   s     r   r   zProcessBiasKFunc.__call__/   r7   r   Nr   r"   r   r   r:   r:   .   r8   r   r:   c                   ó   — e Zd Zd„ ZdS )ÚProcessBiasVFuncc                 ó>   — t          j        |dd¦  «        d         }|S )Nr&   r5   r1   r6   r   s     r   r   zProcessBiasVFunc.__call__5   r7   r   Nr   r"   r   r   r=   r=   4   r8   r   r=   c                   ó   — e Zd Zd„ ZdS )ÚProcessRotCacheFuncc                 óx   — t          |j        ¦  «        dk    sJ ‚|j        d         dk    r|d d …dd…f         S |S )Nr1   r   é    r   é   )ÚlenÚshaper   s     r   r   zProcessRotCacheFunc.__call__;   sG   € å�1”7‰|Œ|˜qÒ Ð Ð Ð ØŒ7�1Œ:˜ÒÐØ�Q�Q�Q˜˜"˜�W”:ÐØˆr   Nr   r"   r   r   r@   r@   :   s#   € € € € € ðð ð ð ð r   r@   c                   óD  ‡ — e Zd Zdedee         fˆ fd„Zdefd„Zd„ Z	d„ Z
d„ Zd	„ Zd#d„Zd„ Zd„ Zd„ Zdee         dedee         fd„Zd$dee         dee         defd„Zd$dee         dee         defd„Zd%dee         dee         defd„Zd$dee         dee         defd„Zd$dee         dee         defd„Zd&dee         dee         defd„Zd&dee         dee         defd„Zd&dee         dee         defd„Z	 	 	 	 d'dee         dee         defd"„Zˆ xZS )(ÚFissionÚmodelÚnodes_to_findc                 óN   •— t          ¦   «                              |d|¦  «         d S )NÚDONOTUSE©ÚsuperÚ__init__)r   rH   rI   Ú	__class__s      €r   rN   zFission.__init__E   s'   ø€ õ
 	‰Œ×Ò˜ 
¨MÑ:Ô:Ð:Ð:Ð:r   Úattn_op_typec                 ó   — || _         d S ©N)rP   )r   rP   s     r   Úset_attention_op_typezFission.set_attention_op_typeL   s   € Ø(ˆÔÐÐr   c                 ó,   — |dz   t          |¦  «        z   S )NÚ_)Ústr)r   Úlayer_idÚnames      r   Ú	get_unamezFission.get_unameO   s   € Ø�c‰z�C ™MœMÑ)Ð)r   c                 óš   — |D ]6}||k    s*|                      |¦  «        s|                     |¦  «        r|c S Œ7t          d|› d�¦  «        ‚)NzEdge z
 not found)ÚendswithÚ
startswithÚ
ValueError)r   ÚedgesrX   Úedges       r   Úget_edge_by_namezFission.get_edge_by_nameR   sb   € Øð 	ð 	ˆDØ�tŠ|ˆ|˜tŸ}š}¨TÑ2Ô2ˆ|°d·o²oÀdÑ6KÔ6Kˆ|Ø���ð åÐ1 Ð1Ð1Ð1Ñ2Ô2Ð2r   c                 ó8   — |                       |j        |¦  «        S rR   )r`   Úinput©r   ÚnoderX   s      r   Úget_input_by_namezFission.get_input_by_nameX   s   € Ø×$Ò$ T¤Z°Ñ6Ô6Ð6r   c                 ó8   — |                       |j        |¦  «        S rR   )r`   Úoutputrc   s      r   Úget_output_by_namezFission.get_output_by_name[   s   € Ø×$Ò$ T¤[°$Ñ7Ô7Ð7r   Nc                 óh  — | j                              |¦  «        }t          j        |¦  «        } ||¦  «        }t	          j        |€|dz   n|t          j        |j        | 	                    ¦   «          
                    ¦   «         d¬¦  «        }| j                              || j        ¦  «         |j        S )NÚ
_processedT©Ú	data_typeÚdimsÚvalsÚraw)rH   Úget_initializerr
   Úto_arrayr   Úmake_tensorr   ÚFLOATrE   ÚflattenÚtobytesÚadd_initializerÚthis_graph_namerX   )r   Úinitializer_nameÚfunctorÚcustom_nameÚiÚ
i_np_arrayÚprocessed_i_np_arrayÚ
new_tensors           r   Úprocess_initializerzFission.process_initializer^   s±   € ØŒJ×&Ò&Ð'7Ñ8Ô8ˆÝ Ô)¨!Ñ,Ô,ˆ
Ø&˜w zÑ2Ô2ÐÝÔ'Ø/:Ð/BÐ˜|Ñ+Ð+ÈÝ!Ô'Ø%Ô+Ø%×-Ò-Ñ/Ô/×7Ò7Ñ9Ô9Øð
ñ 
ô 
ˆ
ð 	Œ
×"Ò" :¨tÔ/CÑDÔDÐDØŒÐr   c                 óª   — | j                              ¦   «         j                             ¦   «         }||_        t
          j        |j        j        _	        d S rR   )
rH   ÚgraphÚ
value_infoÚaddrX   r   rs   ÚtypeÚtensor_typeÚ	elem_type©r   rX   Únew_value_infos      r   Úadd_fp32_value_infozFission.add_fp32_value_infol   óC   € Øœ×)Ò)Ñ+Ô+Ô6×:Ò:Ñ<Ô<ˆØ"ˆÔÝ4?Ô4EˆÔÔ'Ô1Ð1Ð1r   c                 óª   — | j                              ¦   «         j                             ¦   «         }||_        t
          j        |j        j        _	        d S rR   )
rH   r�   r‚   rƒ   rX   r   ÚINT64r„   r…   r†   r‡   s      r   Úadd_int64_value_infozFission.add_int64_value_infoq   rŠ   r   c                 ój  — | j                              ¦   «         j        D ]@}|j        |k    r3| j                              ¦   «         j                             |¦  «          nŒAt          j        |t          j        |¬¦  «        }| j                              ¦   «         j         	                    |g¦  «         d S )N©r†   rE   )
rH   r�   r‚   rX   Úremover   Úmake_tensor_value_infor   rs   Úextend)r   rX   rE   r‚   rˆ   s        r   Úreplace_fp32_value_infozFission.replace_fp32_value_infov   s³   € Øœ*×*Ò*Ñ,Ô,Ô7ð 	ð 	ˆJØŒ $Ò&Ð&Ø”
× Ò Ñ"Ô"Ô-×4Ò4°ZÑ@Ô@Ð@Ø�ð 'õ  Ô6ØÝ!Ô'Øð
ñ 
ô 
ˆð
 	Œ
×ÒÑÔÔ%×,Ò,¨nÐ-=Ñ>Ô>Ð>Ð>Ð>r   Úsubgraph_nodesrW   Úlayer_known_edges_namesc                 ó<  — |D �]}t          |j        ¦  «        D ]N\  }}|dk    rŒ||vr>|                      ||¦  «        |j        |<   |                      |j        |         ¦  «         ŒOt          |j        ¦  «        D ]N\  }}|dk    rŒ||vr>|                      ||¦  «        |j        |<   |                      |j        |         ¦  «         ŒO|                      ||j        ¦  «        |_        | j                             |¦  «         | j        | j	        |j        <   �Œd S )NÚ )
Ú	enumeraterb   rY   r‰   rg   rX   Únodes_to_addÚappendrw   Únode_name_to_graph_name)r   r”   rW   r•   Únew_noder{   rX   s          r   Úset_unique_name_and_add_nodesz%Fission.set_unique_name_and_add_nodes‚   s>  € ð 'ð 	Oñ 	OˆHÝ$ X¤^Ñ4Ô4ð @ð @‘��4Ø˜2’:�:ØØÐ!8Ð8Ð8Ø(,¯ª°xÀÑ(FÔ(F�H”N 1Ñ%Ø×,Ò,¨X¬^¸AÔ->Ñ?Ô?Ð?øÝ$ X¤_Ñ5Ô5ð Að A‘��4Ø˜2’:�:ØØÐ!8Ð8Ð8Ø)-¯ª¸À$Ñ)GÔ)G�H”O AÑ&Ø×,Ò,¨X¬_¸QÔ-?Ñ@Ô@Ð@øØ ŸNšN¨8°X´]ÑCÔCˆHŒMØÔ×$Ò$ XÑ.Ô.Ð.Ø:>Ô:NˆDÔ(¨¬Ñ7Ñ7ð	Oð 	Or   r—   ÚinputsÚoutputsÚprefixc                 ó”   — t          |¦  «        dk    sJ ‚t          |¦  «        dk    sJ ‚t          j        d|||dz   d¬¦  «        }|gS )Nr&   r   ÚLayerNormalizationÚ_LayerNormalizationg   €µøä>)rž   rŸ   rX   Úepsilon©rD   r   Ú	make_node©r   rž   rŸ   r    rd   s        r   Ú	layernormzFission.layernorm–   se   € Ý�6‰{Œ{˜aÒÐÐÐÝ�7‰|Œ|˜qÒ Ð Ð Ð ÝÔØ ØØØÐ/Ñ/Ø)ð
ñ 
ô 
ˆð ˆvˆr   c                 ó  — t          |¦  «        dk    sJ ‚t          |¦  «        dk    sJ ‚t          j        d|d         |d         g|dz   g|dz   ¬¦  «        }t          j        d|dz   |d         g||d	z   ¬¦  «        }||gS )
Nr&   r   ÚMatMulr   Ú
matmul_out©rž   rŸ   rX   ÚAddr1   ÚBiasr¥   )r   rž   rŸ   r    Úmatmulrƒ   s         r   ÚgemmzFission.gemm¢   s°   € Ý�6‰{Œ{˜aÒÐÐÐÝ�7‰|Œ|˜qÒ Ð Ð Ð ÝÔ!ØØ˜1”I˜v aœyÐ)Ø˜lÑ*Ð+Ø˜(Ñ"ð	
ñ 
ô 
ˆõ ÔØØ˜\Ñ)¨6°!¬9Ð5ØØ˜&‘ð	
ñ 
ô 
ˆð ˜ˆ}Ðr   rB   c           	      ó˜   — t          |¦  «        dk    sJ ‚t          |¦  «        dk    sJ ‚t          j        d|||dz   d||¬¦  «        }|gS )Né   r   ÚRotaryEmbeddingúcom.microsoft)rž   rŸ   rX   ÚdomainÚrotary_embedding_dimÚ	num_headsr¥   )r   rž   rŸ   r    Úrot_dimr·   rd   s          r   ÚrotaryzFission.rotary³   sk   € Ý�6‰{Œ{˜aÒÐÐÐÝ�7‰|Œ|˜qÒ Ð Ð Ð ÝÔØØØØÐ+Ñ+Ø"Ø!(Øð
ñ 
ô 
ˆð ˆvˆr   c                 ó”   — t          |¦  «        dk    sJ ‚t          |¦  «        dk    sJ ‚t          j        d|||dz   d¬¦  «        }|gS )Nr   ÚFastGelur´   )rž   rŸ   rX   rµ   r¥   r§   s        r   ÚfastgeluzFission.fastgeluÁ   sd   € Ý�6‰{Œ{˜aÒÐÐÐÝ�7‰|Œ|˜qÒ Ð Ð Ð ÝÔØØØØ˜*Ñ$Ø"ð
ñ 
ô 
ˆð ˆvˆr   c                 ó’   — t          |¦  «        dk    sJ ‚t          |¦  «        dk    sJ ‚t          j        d|||dz   ¬¦  «        }|gS )Nr1   r   r­   r¬   r¥   r§   s        r   rƒ   zFission.addÍ   s`   € Ý�6‰{Œ{˜aÒÐÐÐÝ�7‰|Œ|˜qÒ Ð Ð Ð ÝÔØØØØ˜%‘ð	
ñ 
ô 
ˆð ˆvˆr   c           	      ó˜   — t          |¦  «        dk    sJ ‚t          |¦  «        dk    sJ ‚t          j        d|||dz   d|d¬¦  «        }|gS )Né   r&   ÚMultiHeadAttentionr´   r   )rž   rŸ   rX   rµ   r·   Úunidirectionalr¥   ©r   rž   rŸ   r    r·   rd   s         r   ÚmhazFission.mhaØ   sk   € Ý�6‰{Œ{˜aÒÐÐÐÝ�7‰|Œ|˜qÒ Ð Ð Ð ÝÔØ ØØØÐ.Ñ.Ø"ØØð
ñ 
ô 
ˆð ˆvˆr   c           	      ó˜   — t          |¦  «        dk    sJ ‚t          |¦  «        dk    sJ ‚t          j        d|||dz   d||¬¦  «        }|gS )Né   r&   ÚGroupQueryAttentionr´   )rž   rŸ   rX   rµ   r·   Úkv_num_headsr¥   rÂ   s         r   ÚgqazFission.gqaæ   sk   € Ý�6‰{Œ{˜aÒÐÐÐÝ�7‰|Œ|˜qÒ Ð Ð Ð ÝÔØ!ØØØÐ/Ñ/Ø"ØØ"ð
ñ 
ô 
ˆð ˆvˆr   c                 óœ   — t          |¦  «        dk    sJ ‚t          |¦  «        dk    sJ ‚t          j        d|||dz   d|ddd¬¦	  «	        }|gS )Né   r1   Ú	Attentionr´   r   rB   )rž   rŸ   rX   rµ   r·   rÁ   Ú	do_rotaryr¶   r¥   rÂ   s         r   Ú	attentionzFission.attentionô   sp   € Ý�6‰{Œ{˜aÒÐÐÐÝ�7‰|Œ|˜qÒ Ð Ð Ð ÝÔØØØØ˜+Ñ%Ø"ØØØØ!#ð

ñ 

ô 

ˆð ˆvˆr   éP   ç   À%Ÿ¼?c                 óœ   — t          |¦  «        dk    sJ ‚t          |¦  «        dk    sJ ‚t          j        d|||dz   d||||¬¦	  «	        }|gS )Né   r   ÚPagedAttentionzvllm.ort.ext)rž   rŸ   rX   rµ   r·   Únum_kv_headsÚ	head_sizeÚscaler¥   )r   rž   rŸ   r    r·   rÔ   rÕ   rd   s           r   Ú
paged_attnzFission.paged_attn  ss   € õ �6‰{Œ{˜aÒÐÐÐÝ�7‰|Œ|˜qÒ Ð Ð Ð ÝÔØØØØÐ*Ñ*Ø!ØØ"ØØð

ñ 

ô 

ˆð ˆvˆr   rR   )r—   )r—   rB   rB   )r—   rB   )r—   rB   rÎ   rÏ   )r   r    r!   r   ÚlistrV   rN   r   rS   rY   r`   re   rh   r   r‰   r�   r“   r   Úintr�   r¨   r°   r¹   r¼   rƒ   rÃ   rÈ   rÍ   rÖ   Ú__classcell__©rO   s   @r   rG   rG   D   s)  ø€ € € € € ð;àð;ð ˜C”yð;ð ;ð ;ð ;ð ;ð ;ð)°/ð )ð )ð )ð )ð*ð *ð *ð3ð 3ð 3ð7ð 7ð 7ð8ð 8ð 8ðð ð ð ðFð Fð Fð
Fð Fð Fð

?ð 
?ð 
?ðOØ" 9œoðOØ9<ðOØW[Ð\_ÔW`ðOð Oð Oð Oð(
ð 
  S¤	ð 
°D¸´Ið 
Àsð 
ð 
ð 
ð 
ðð ˜4 œ9ð ¨t°C¬yð À#ð ð ð ð ð"ð ˜T #œYð °°c´ð ÀCð ð ð ð ð
ð 
˜t Cœyð 
°4¸´9ð 
Àcð 
ð 
ð 
ð 
ð	ð 	˜$˜sœ)ð 	¨d°3¬ið 	Àð 	ð 	ð 	ð 	ðð ˜$˜sœ)ð ¨d°3¬ið Àð ð ð ð ðð ˜$˜sœ)ð ¨d°3¬ið Àð ð ð ð ðð   S¤	ð °D¸´Ið Àsð ð ð ð ð( ØØØ!ðð à�S”	ðð �c”ðð ð	ð ð ð ð ð ð ð r   rG   c                   óV   ‡ — e Zd Zdededefˆ fd„Zdefd„Zd„ Zde	fd	„Z
de	fd
„Zˆ xZS )ÚPhi2PreProcessorrH   r·   Úhidden_sizec                 ó‚   •— t          ¦   «                              |¦  «         d| _        || _        || _        d| _        d S )NrB   Úmodeling_phi_PhiModel_model_1)rM   rN   Únum_hidden_layersÚnum_attention_headsrÝ   Ú	func_name©r   rH   r·   rÝ   rO   s       €r   rN   zPhi2PreProcessor.__init__  s>   ø€ Ý‰Œ×Ò˜ÑÔÐØ!#ˆÔØ#,ˆÔ Ø&ˆÔà8ˆŒˆˆr   Úreturnc                 óD  — i }d|d<   d|d<   d|d<   d|d<   t          d	| j        d	¦  «        D ]0}d
|› �|d|› �<   d|› �|d|› �<   d|› �|d|› d�<   d|› �|d|› d�<   Œ1d„ | j        j        j        D ¦   «         }d|v rd|v rd|d<   d|d<   nd|v rd|v sJ ‚d|d<   d|d<   |S )NÚlogitsÚ	lm_head_1Ú	input_idsÚl_input_ids_Ú
past_key_0Ú
key_statesÚpast_value_0Úvalue_statesr   Ú	past_key_Úkey_states_Úpast_value_Úvalue_states_Úpresent_key_Úmodel_layers_Ú_1Úpresent_value_Ú_1_1c                 ó   — g | ]	}|j         ‘Œ
S r"   ©rX   )Ú.0Úos     r   ú
<listcomp>z7Phi2PreProcessor.get_phi2_edge_dict.<locals>.<listcomp>2  s   € Ð;Ð;Ð;˜a�1”6Ð;Ð;Ð;r   Úmodel_layers_0_1_1Úmodel_layers_0_1_2Úpresent_key_0Úpresent_value_0Úmodel_layers_0_1)Úrangerà   rH   r�   rg   )r   Ú	edge_dictr{   rŸ   s       r   Úget_phi2_edge_dictz#Phi2PreProcessor.get_phi2_edge_dict&  sD  € Øˆ	Ø!)ˆ	�+ÑØ$/ˆ	�.Ñ!Ø".ˆ	�,ÑØ$2ˆ	�.Ñ!Ý�q˜$Ô0°!Ñ4Ô4ð 	Fð 	FˆAØ+:°q¨?¨?ˆIÐ' AÐ'Ð'Ñ(Ø->¸1Ð->Ð->ˆIÐ) aÐ)Ð)Ñ*Ø/A¸aÐ/AÐ/AˆIÐ+ aÐ+Ð+Ð+Ñ,Ø1EÀ!Ð1EÐ1EˆIÐ- aÐ-Ð-Ð-Ñ.Ð.à;Ð; 4¤:Ô#3Ô#:Ð;Ñ;Ô;ˆØ 7Ð*Ð*Ð/CÀwÐ/NÐ/NØ.=ˆIÐ*Ñ+Ø.?ˆIÐ*Ñ+Ð+à%¨Ð0Ð0Ð5IÈWÐ5TÐ5TÐ5TÐTØ,;ˆIÐ(Ñ)Ø.?ˆIÐ*Ñ+ØÐr   c                 óš   — d}| j         j        j        D ]6}|j                             |¦  «        }|dk    r|j        |d …         |_        Œ7d S )NÚ)modeling_phi_PhiDecoderLayer_model_layersr5   )rH   r�   rd   Úop_typeÚfind)r   Úphi2_transformer_layer_namerd   Úindexs       r   Úsimplify_phi2_op_typez&Phi2PreProcessor.simplify_phi2_op_type<  s[   € Ø&QÐ#Ø”JÔ$Ô)ð 	4ð 	4ˆDØ”L×%Ò%Ð&AÑBÔBˆEØ˜Š{ˆ{Ø#œ|¨E¨F¨FÔ3�”øð	4ð 	4r   rP   c                 óŽ  — |t           j        k    | _        |t           j        k    | _        | j        j        }g }|j        D �]�}d|j        v rüt          j
        |j        | j        st          j        nt          j        ddg¬¦  «        }t          j
        dt          j        dg¬¦  «        }t          j
        dt          j        ddg¬¦  «        }t          j
        dt          j        ddg¬¦  «        }t          j
        d	t          j        dg¬¦  «        }	| j        s|                     |||g¦  «        n|                     |||	g¦  «         | j        rwd
|j        v rlt          j
        |j                             d
d¦  «        |j        j        j        dd| j        d| j        | j        z  g¬¦  «        }
|                     |
g¦  «         �Œ†| j        ršd
|j        v rCt          j
        |j        |j        j        j        g d¢¬¦  «        }
|                     |
g¦  «         d|j        v rCt          j
        |j        |j        j        j        g d¢¬¦  «        }
|                     |
g¦  «         �Œ'd
|j        v s	d|j        v rWt          j
        |j        |j        j        j        d| j        d| j        | j        z  g¬¦  «        }
|                     |
g¦  «         �Œ’|                     d¦  «         |j                             |¦  «         g }t-          |j        ¦  «        D ]þ\  }}|dk    r|                     |g¦  «         Œ"| j        rvd|j        v rlt          j
        |j                             dd¦  «        |j        j        j        dd| j        d| j        | j        z  g¬¦  «        }
|                     |
g¦  «         ŒŸ| j        rŒ§t          j
        |j        |j        j        j        d| j        d| j        | j        z  g¬¦  «        }
|                     |
g¦  «         Œÿ|                     d¦  «         |j                             |¦  «         d S )Nrè   Ú
batch_sizeÚseq_lenr�   Ústepr   Úposition_idsÚattention_maskÚinput_metadataÚpast_keyÚpastr1   Úpast_seq_len)Ú
num_blocksr·   Úhead_size_xÚ
block_sizeÚblock_xÚ
past_value)r  r·   rÔ   r  rb   r   Úpresent_keyÚpresentÚtotal_seq_lenrg   )r   rË   Úuse_attnrÒ   Úuse_vllmrH   r�   rb   rX   r   r‘   r   ÚINT32rŒ   r’   Úreplacer„   r…   r†   rá   rÝ   Ú
ClearFieldr˜   rg   )r   rP   r�   Ú
new_inputsÚviÚvi_iidÚvi_stepÚvi_pidÚvi_maskÚvi_metaÚvi_cacheÚnew_outputsr{   s                r   Úprocess_graph_ioz!Phi2PreProcessor.process_graph_ioC  s·  € Ø$­Ô(AÒAˆŒØ$­Ô(FÒFˆŒØ”
Ô ˆØˆ
Ø”+ð N	2ñ N	2ˆBØ˜bœgÐ%Ð%ÝÔ6Ø”GØ7;´}Ð[�kÔ/Ð/Í+ÔJ[Ø'¨Ð3ðñ ô �õ
 !Ô7ØÝ)Ô/Ø˜#ðñ ô �õ
  Ô6Ø"Ý)Ô/Ø'¨Ð3ðñ ô �õ
 !Ô7Ø$Ý)Ô/Ø'¨Ð3ðñ ô �õ
 !Ô7Ø$Ý)Ô/Ø˜#ðñ ô �ð  œ=ðF�J×%Ò% v¨w¸Ð&@ÑAÔAÐAà#×*Ò*¨F°F¸GÐ+DÑEÔEøàŒ}ð .2Ø ¤Ð(Ð(Ý%Ô<ØœŸš¨
°FÑ;Ô;Ø"$¤'Ô"5Ô"?àØ(Ø Ô4Ø*Ø Ô,°Ô0HÑHðð
 ñ 
 ô 
 �Hð ×%Ò% x jÑ1Ô1Ð1ùØ”ð  2Ø ¤Ð(Ð(Ý%Ô<ØœØ"$¤'Ô"5Ô"?ØaÐaÐað ñ  ô  �Hð
 ×%Ò% x jÑ1Ô1Ð1Ø 2¤7Ð*Ð*Ý%Ô<ØœØ"$¤'Ô"5Ô"?ðð ð ð	 ñ 	 ô 	 �Hð ×%Ò% x jÑ1Ô1Ð1ùà ¤Ð(Ð(¨L¸B¼GÐ,CÐ,CÝ%Ô<ØœØ"$¤'Ô"5Ô"?à(Ø Ô4Ø*Ø Ô,°Ô0HÑHð	ð	 ñ 	 ô 	 �Hð ×%Ò% x jÑ1Ô1Ð1ùà×Ò˜Ñ!Ô!Ð!ØŒ×Ò˜:Ñ&Ô&Ð&àˆÝ˜uœ|Ñ,Ô,ð 	3ð 	3‰EˆAˆrØ�AŠvˆvØ×"Ò" B 4Ñ(Ô(Ð(Ð(à”=ð 3Ø$¨¬Ð/Ð/Ý#)Ô#@ØœGŸOšO¨M¸9ÑEÔEØ&(¤gÔ&9Ô&Cà !Ø ,Ø $Ô 8Ø /Ø $Ô 0°DÔ4LÑ Lð#ð
$ñ 
$ô 
$˜ð $×*Ò*¨H¨:Ñ6Ô6Ð6øØ”]ð 3Øå%Ô<ØœØ"$¤'Ô"5Ô"?à(Ø Ô4Ø+Ø Ô,°Ô0HÑHð	ð	 ñ 	 ô 	 �Hð  ×&Ò&¨ zÑ2Ô2Ð2Ð2à×Ò˜Ñ"Ô"Ð"ØŒ×Ò˜KÑ(Ô(Ð(Ð(Ð(r   c                 óº  — d }| j         j        D ]*}|j                             | j        ¦  «        r	|j        } nŒ+|€J ‚|                      |¦  «         |                      |                      ¦   «         ¦  «         |                      ¦   «          |  	                    ¦   «          |t          j        k    r|                      ¦   «          |                      |¦  «         d S rR   )rH   Ú	functionsrX   r[   râ   Úunroll_functionÚupdate_edgesr  r
  Úremove_dropout_layerr   rÒ   Úremove_lm_head_layerr+  )r   rP   Úfunction_nameÚfuncs       r   Úpreprocess_onnxz Phi2PreProcessor.preprocess_onnxÀ  sæ   € ØˆØ”JÔ(ð 	ð 	ˆDØŒy×!Ò! $¤.Ñ1Ô1ð Ø $¤	�Ø�ðð Ð(Ð(Ð(Ø×Ò˜]Ñ+Ô+Ð+Ø×Ò˜$×1Ò1Ñ3Ô3Ñ4Ô4Ð4Ø×"Ò"Ñ$Ô$Ð$Ø×!Ò!Ñ#Ô#Ð#Ø�?Ô9Ò9Ð9Ø×%Ò%Ñ'Ô'Ð'Ø×Ò˜lÑ+Ô+Ð+Ð+Ð+r   )r   r    r!   r   rØ   rN   Údictr  r
  r   r+  r4  rÙ   rÚ   s   @r   rÜ   rÜ     s¸   ø€ € € € € ð9˜jð 9°Sð 9Àsð 9ð 9ð 9ð 9ð 9ð 9ð Dð ð ð ð ð,4ð 4ð 4ð{)¨_ð {)ð {)ð {)ð {)ðz,¨Oð ,ð ,ð ,ð ,ð ,ð ,ð ,ð ,r   rÜ   c                   ó*   ‡ — e Zd Zdefˆ fd„Zd„ Zˆ xZS )ÚFissionTransformerEmbeddingPhirH   c                 óN   •— t          ¦   «                              |dg¦  «         d S )NÚ6torch_nn_modules_sparse_Embedding_model_embed_tokens_1rL   ©r   rH   rO   s     €r   rN   z'FissionTransformerEmbeddingPhi.__init__Ñ  s)   ø€ õ 	‰Œ×Ò˜Ð!YÐ ZÑ[Ô[Ð[Ð[Ð[r   c                 óÀ  — t                                d|j        ¦  «         t          |j        ¦  «        dk    sJ ‚t          |j        ¦  «        dk    sJ ‚|j        d         }|j        d         }|                      |d¦  «        }|||g}t          j        d||g|gd¬¦  «        g}|  	                    |d|¦  «         | j
                             |¦  «         d	| _        d S )
NúOptimizing %s...r1   r   r   zembed_tokens.weightÚGatherÚEmbedding_Gatherr¬   T)ÚloggerÚinforX   rD   rb   rg   re   r   r¦   r�   Únodes_to_removerš   Úprune_graph)	r   rd   Úinput_name_to_nodesÚoutput_name_to_noderb   rg   Ú	embeddingr•   r”   s	            r   Úfusez#FissionTransformerEmbeddingPhi.fuse×  s÷   € Ý�ŠÐ&¨¬	Ñ2Ô2Ð2å�4”:‰Œ !Ò#Ð#Ð#Ð#Ý�4”;ÑÔ 1Ò$Ð$Ð$Ð$à”
˜1”ˆØ”˜Q”ˆà×*Ò*¨4Ð1FÑGÔGˆ	à#(¨&°)Ð"<Ðõ ÔØØ! 5Ð)Ø˜Ø'ð	ñ ô ð
ˆð 	×*Ò*¨>¸1Ð>UÑVÔVÐVØÔ×#Ò# DÑ)Ô)Ð)ØˆÔÐÐr   ©r   r    r!   r   rN   rF  rÙ   rÚ   s   @r   r7  r7  Ð  s[   ø€ € € € € ð\àð\ð \ð \ð \ð \ð \ð ð  ð  ð  ð  ð  ð  r   r7  c                   ó*   ‡ — e Zd Zdefˆ fd„Zd„ Zˆ xZS )ÚFissionTransformerLayerNormPhirH   c                 óN   •— t          ¦   «                              |dg¦  «         d S )NÚ@torch_nn_modules_normalization_LayerNorm_model_final_layernorm_1rL   r:  s     €r   rN   z'FissionTransformerLayerNormPhi.__init__ó  s)   ø€ õ 	‰Œ×Ò˜Ð!cÐ dÑeÔeÐeÐeÐer   c                 óv  — t                                d|j        ¦  «         t          |j        ¦  «        dk    sJ ‚t          |j        ¦  «        dk    sJ ‚|j        d         }|j        d         }|                      |d¦  «        }|                      |d¦  «        }||||g}g }	|	                     |                      |||g|gd¦  «        ¦  «         |  	                    |	d|¦  «         |  
                    |g d	¢¦  «         |  
                    |g d	¢¦  «         | j                             |¦  «         d
| _        d S )Nr<  r&   r   r   zfinal_layernorm.weightzfinal_layernorm.biasÚFinaléc   ©r  r  rÝ   T)r?  r@  rX   rD   rb   rg   re   r’   r¨   r�   r“   rA  rš   rB  )
r   rd   rC  rD  rb   rg   Ú	ln_weightÚln_biasr•   r”   s
             r   rF  z#FissionTransformerLayerNormPhi.fuseù  sI  € Ý�ŠÐ&¨¬	Ñ2Ô2Ð2å�4”:‰Œ !Ò#Ð#Ð#Ð#Ý�4”;ÑÔ 1Ò$Ð$Ð$Ð$à”
˜1”ˆØ”˜Q”ˆà×*Ò*¨4Ð1IÑJÔJˆ	Ø×(Ò(¨Ð/EÑFÔFˆà#(¨&°)¸WÐ"EÐàˆØ×Ò˜dŸnšn¨e°YÀÐ-HÈ6È(ÐT[Ñ\Ô\Ñ]Ô]Ð]à×*Ò*¨>¸2Ð?VÑWÔWÐWà×$Ò$ UÐ,TÐ,TÐ,TÑUÔUÐUØ×$Ò$ VÐ-UÐ-UÐ-UÑVÔVÐVàÔ×#Ò# DÑ)Ô)Ð)ØˆÔÐÐr   rG  rÚ   s   @r   rI  rI  ò  s[   ø€ € € € € ðfàðfð fð fð fð fð fð ð  ð  ð  ð  ð  ð  r   rI  c                   ó*   ‡ — e Zd Zdefˆ fd„Zd„ Zˆ xZS )Ú!FissionTransformerCausalLMHeadPhirH   c                 óN   •— t          ¦   «                              |dg¦  «         d S )NÚ(torch_nn_modules_linear_Linear_lm_head_1rL   r:  s     €r   rN   z*FissionTransformerCausalLMHeadPhi.__init__  s)   ø€ õ 	‰Œ×Ò˜Ð!KÐ LÑMÔMÐMÐMÐMr   c                 ó¶  — t                                d|j        ¦  «         t          |j        ¦  «        dk    sJ ‚t          |j        ¦  «        dk    sJ ‚|j        d         }|j        d         }|                      |                      |d¦  «        t          ¦   «         ¦  «        }|                      |d¦  «        }||||g}g }	|	 	                    |  
                    |||g|gd¦  «        ¦  «         |                      |	d	|¦  «         |                      |g d
¢¦  «         |                      |g d¢¦  «         | j                             |¦  «         d| _        d S )Nr<  rÊ   r   r1   r   zlm_head.weightzlm_head.biasÚLMHead_rN  rO  )r  r  i È  T)r?  r@  rX   rD   rb   rg   r   re   r   r’   r°   r�   r“   rA  rš   rB  )
r   rd   rC  rD  rb   rg   Ú	fc_weightÚfc_biasr•   r”   s
             r   rF  z&FissionTransformerCausalLMHeadPhi.fuse  s^  € Ý�ŠÐ&¨¬	Ñ2Ô2Ð2å�4”:‰Œ !Ò#Ð#Ð#Ð#Ý�4”;ÑÔ 1Ò$Ð$Ð$Ð$à”
˜1”ˆØ”˜Q”ˆà×,Ò,¨T×-CÒ-CÀDÐJZÑ-[Ô-[Õ]mÑ]oÔ]oÑpÔpˆ	Ø×(Ò(¨¨~Ñ>Ô>ˆà#(¨&°)¸WÐ"EÐàˆØ×Ò˜dŸiši¨°	¸7Ð(CÀfÀXÈyÑYÔYÑZÔZÐZà×*Ò*¨>¸2Ð?VÑWÔWÐWà×$Ò$ UÐ,TÐ,TÐ,TÑUÔUÐUØ×$Ò$ VÐ-MÐ-MÐ-MÑNÔNÐNàÔ×#Ò# DÑ)Ô)Ð)ØˆÔÐÐr   rG  rÚ   s   @r   rS  rS    s[   ø€ € € € € ðNàðNð Nð Nð Nð Nð Nð ð  ð  ð  ð  ð  ð  r   rS  c                   ó@   ‡ — e Zd Zdedefˆ fd„Zd„ Zd„ Zd„ Zd„ Z	ˆ xZ
S )ÚFissionTransformerBlockPhirH   r·   c                 óÞ   •— || _         d}i | _        g }t          |¦  «        D ]'}d|› d�}|                     |¦  «         || j        |<   Œ(t	          ¦   «                              ||¦  «         d S )NrB   Ú*modeling_phi_PhiDecoderLayer_model_layers_rô   )r·   Úfunc_to_layer_idr  rš   rM   rN   )r   rH   r·   Úmax_num_layersrI   Úlayerrâ   rO   s          €r   rN   z#FissionTransformerBlockPhi.__init__5  s‰   ø€ ð
 #ˆŒØˆØ "ˆÔØˆÝ˜>Ñ*Ô*ð 	5ð 	5ˆEØNÀUÐNÐNÐNˆIØ× Ò  Ñ+Ô+Ð+Ø/4ˆDÔ! )Ñ,Ð,å‰Œ×Ò˜ Ñ.Ô.Ð.Ð.Ð.r   c                 ó&   — | j         |j                 S rR   )r^  r  )r   rd   s     r   Úget_layer_idz'FissionTransformerBlockPhi.get_layer_idE  s   € ØÔ$ T¤\Ô2Ð2r   c                 ó²  — t          j        ddgdgdt          j        ¬¦  «        t          j        dddgdgd	¬
¦  «        t          j        dddgdgd¬
¦  «        t          j        ddgdgdt          j        ¬¦  «        t          j        ddgdgd¬
¦  «        t          j        dddgdgdd¬¦  «        t          j        ddgdgdt          j        ¬¦  «        g}|S )NÚCastr  Ú
mask_int64ÚCast_gqa_aux_0)rž   rŸ   rX   ÚtoÚ	ReduceSumÚoneÚmask_row_sumsÚReduceSum_gqa_auxr¬   ÚSubÚseqlens_k_int64ÚSub_gqa_auxÚ	seqlens_kÚCast_gqa_aux_1ÚShapeÚ
mask_shapeÚShape_gqa_aux_0r=  Útotal_seq_len_int64ÚGather_gqa_aux_0r   )rž   rŸ   rX   ÚaxisÚtotal_sequence_lengthÚCast_gqa_aux_2)r   r¦   r   rŒ   r  )r   Úgqa_aux_nodess     r   Úget_gqa_aux_nodesz,FissionTransformerBlockPhi.get_gqa_aux_nodesH  s<  € åÔØØ(Ð)Ø%˜Ø%ÝÔ$ðñ ô õ ÔØØ$ eÐ,Ø(Ð)Ø(ð	ñ ô õ ÔØØ'¨Ð/Ø*Ð+Ø"ð	ñ ô õ ÔØØ)Ð*Ø$˜Ø%ÝÔ$ðñ ô õ Ô˜W¨l¨^ÀlÀ^ÐZkÐlÑlÔlÝÔØØ$ eÐ,Ø.Ð/Ø'Øðñ ô õ ÔØØ-Ð.Ø0Ð1Ø%ÝÔ$ðñ ô ðG*
ˆðV Ðr   c	                 ó8  — | j                              |¦  «        }	| j                              |¦  «        }
| j                              |¦  «        }t          j        t	          j        |	¦  «        d¦  «        }t          j        t	          j        |
¦  «        d¦  «        }t          j        t	          j        |¦  «        d¦  «        }t          j        |||fd¬¦  «        }| j                              |¦  «        }| j                              |¦  «        }| j                              |¦  «        }t	          j        |¦  «        }t	          j        |¦  «        }t	          j        |¦  «        }t          j        |||fd¬¦  «        }|j        d         }t          j	        |t          j        ||dz  g|                     ¦   «                              ¦   «         d¬¦  «        }| j                              || j        ¦  «         t          j	        |t          j        |dz  g|                     ¦   «                              ¦   «         d¬¦  «        }| j                              || j        ¦  «         |                      |j        ¦  «         |                      |j        ¦  «         ||fS )Nr   r   )rv  r   r&   Trk   )rH   rp   r   r   r
   rq   ÚstackrE   r   rr   r   rs   rt   ru   rv   rw   r‰   rX   )r   Úq_wÚk_wÚv_wÚq_bÚk_bÚv_bÚweight_nameÚ	bias_nameÚq_weightÚk_weightÚv_weightÚqwÚkwÚvwÚ
qkv_weightÚq_biasÚk_biasÚv_biasÚqbÚkbÚvbÚqkv_biasrÝ   ÚweightÚbiass                             r   Úpack_qkv_gemmz(FissionTransformerBlockPhi.pack_qkv_gemmv  sH  € Ø”:×-Ò-¨cÑ2Ô2ˆØ”:×-Ò-¨cÑ2Ô2ˆØ”:×-Ò-¨cÑ2Ô2ˆÝŒ\�+Ô.¨xÑ8Ô8¸&ÑAÔAˆÝŒ\�+Ô.¨xÑ8Ô8¸&ÑAÔAˆÝŒ\�+Ô.¨xÑ8Ô8¸&ÑAÔAˆÝ”X˜r 2 r˜l°Ð3Ñ3Ô3ˆ
à”×+Ò+¨CÑ0Ô0ˆØ”×+Ò+¨CÑ0Ô0ˆØ”×+Ò+¨CÑ0Ô0ˆÝÔ! &Ñ)Ô)ˆÝÔ! &Ñ)Ô)ˆÝÔ! &Ñ)Ô)ˆÝ”8˜R  R˜L¨qÐ1Ñ1Ô1ˆà Ô& qÔ)ˆåÔ#ØÝ!Ô'Ø˜{¨Q™Ð/Ø×#Ò#Ñ%Ô%×-Ò-Ñ/Ô/Øð
ñ 
ô 
ˆð 	Œ
×"Ò" 6¨4Ô+?Ñ@Ô@Ð@åÔ!ØÝ!Ô'Ø ‘/Ð"Ø×!Ò!Ñ#Ô#×+Ò+Ñ-Ô-Øð
ñ 
ô 
ˆð 	Œ
×"Ò" 4¨Ô)=Ñ>Ô>Ð>à× Ò  ¤Ñ-Ô-Ð-Ø× Ò  ¤Ñ+Ô+Ð+à˜IÐ%Ð%r   c                 ó°  — t                                d|j        ¦  «         t                                d| j        › �¦  «         |                      |¦  «        }|j        d         }|                      |d¦  «        }|                      |d¦  «        }|j        d         }|                      |d¦  «        }	|                      |d¦  «        }
|                      |d	¦  «        }|                      |d
¦  «        }d\  }}}}}}d\  }}d\  }}| j        t          j
        k    �rQ|                      |                      |d¦  «        t          ¦   «         ¦  «        }|                      |                      |d¦  «        t          ¦   «         ¦  «        }|                      |                      |d¦  «        t          ¦   «         ¦  «        }|                      |d¦  «        }|                      |d¦  «        }|                      |d¦  «        }|                      |                      |d¦  «        t          ¦   «         ¦  «        }|                      |                      |d¦  «        t          ¦   «         ¦  «        }n¿|                      |                      |d¦  «        |                      |d¦  «        |                      |d¦  «        |                      |d¦  «        |                      |d¦  «        |                      |d¦  «        |                      |d¦  «        |                      |d¦  «        ¦  «        \  }}|                      |                      |d¦  «        t          ¦   «         ¦  «        }|                      |d¦  «        }|                      |                      |d¦  «        t          ¦   «         ¦  «        }|                      |                      |d¦  «        t          ¦   «         ¦  «        }|                      |d¦  «        }|                      |d¦  «        }g }|                     |||g¦  «         |                     ||	|
g¦  «         |                     ||g¦  «         | j        t          j
        k    r|                     ||||||||g¦  «         n|                     ||g¦  «         |                     ||||||g¦  «         |                     g d¢¦  «         g }|                     |                      |||gdg¦  «        ¦  «         |                     |                      d||gd gd!¦  «        ¦  «         |                     |                      d||gd"gd#¦  «        ¦  «         |                     |                      d"gd$g¦  «        ¦  «         |                     |                      d$||gd%gd&¦  «        ¦  «         |                     |                      d d%gd'gd(¦  «        ¦  «         |                     |                      |d'g|gd)¦  «        ¦  «         | j        t          j
        k    �rq|                     |                      d||gd*gd+¦  «        ¦  «         |                     |                      d||gd,gd-¦  «        ¦  «         |                     |                      d||gd.gd/¦  «        ¦  «         | j        t          j        k    rd0nd1}|                     |                      d*|||gd2gd+¦  «        ¦  «         |                     |                      d,|||gd3gd-¦  «        ¦  «         | j        t          j        k    r6|                     |                      d2d3d.d4d5d4||gd|	|
g¦  «        ¦  «         �nv| j        t          j        k    rÊ|                     |                      d2d3d.||d6d7gd|	|
g¦  «        ¦  «         |dk    r�|                      ¦   «         } | D ]0}!| j                             |!¦  «         | j        | j        |!j        <   Œ1| j          !                    tE          j#        tI          j%        d8gd9¬:¦  «        d;¬<¦  «        | j        ¦  «         n—| j        t          j        k    r0|                     |  &                    d2d3d.||d=gdg¦  «        ¦  «         nQd>|› �}"d?|› �}#|                     |"|#g¦  «         |                     |  '                    d||d5|"gd|#g¦  «        ¦  «         |  (                    |||¦  «         |  )                    |g d@¢¦  «         |  )                    |g d@¢¦  «         | j*                             |¦  «         dA| _+        d S )BNr<  zAttentionOpType: r   r  r  r5   r  Úpresent_valuezinput_layernorm.weightzinput_layernorm.bias)NNNNNN)NNzself_attn.q_proj.weightzself_attn.k_proj.weightzself_attn.v_proj.weightzself_attn.q_proj.biaszself_attn.k_proj.biaszself_attn.v_proj.biaszrotary_emb.cos_cachedzrotary_emb.sin_cachedÚattn_qkv_weightÚattn_qkv_biaszself_attn.dense.weightzself_attn.dense.biaszmlp.fc1.weightzmlp.fc2.weightzmlp.fc1.biaszmlp.fc2.bias)r  r  ro  rw  r  r  Úln_outÚattn_outÚattn_add_outÚOutProj_Úfc1_outÚFC1_Úgelu_outÚfc2_outÚFC2_Úresidual_1_outÚ
Residual_1Ú
Residual_2ÚqueryÚQ_ÚkeyÚK_ÚvalueÚV_r  r  Ú	query_rotÚkey_rotr—   r  ro  rw  r   Úint64)Údtyperi  rø   r  Úpast_Úpresent_rO  T),r?  r@  rX   rP   rb  rb   re   rg   rh   r   rË   r   r   r@   r•  rY   r’   r¨   r°   r¼   rƒ   rÒ   r¹   rÀ   rÃ   rÆ   rÈ   rz  r™   rš   rw   r›   rH   rv   r   Ú
from_arrayr   ÚarrayrÖ   rÍ   r�   r“   rA  rB  )$r   rd   rC  rD  rW   Úi_hidden_statesÚi_key_cacheÚi_value_cacheÚo_hidden_statesÚo_key_cacheÚo_value_cacherP  rQ  Úattn_q_weightÚattn_q_biasÚattn_k_weightÚattn_k_biasÚattn_v_weightÚattn_v_biasr˜  r™  Ú	cos_cacheÚ	sin_cacheÚattn_out_weightÚattn_out_biasÚmlp_fc1_weightÚmlp_fc2_weightÚmlp_fc1_biasÚmlp_fc2_biasr•   r”   Úpos_ids_namery  rœ   Ú	past_nameÚpresent_names$                                       r   rF  zFissionTransformerBlockPhi.fuse   sþ	  € õ 	�ŠÐ&¨¬	Ñ2Ô2Ð2å�ŠÐ;¨Ô(9Ð;Ð;Ñ<Ô<Ð<à×$Ò$ TÑ*Ô*ˆàœ* Qœ-ˆØ×,Ò,¨T°:Ñ>Ô>ˆØ×.Ò.¨t°\ÑBÔBˆàœ+ bœ/ˆØ×-Ò-¨d°MÑBÔBˆØ×/Ò/°°oÑFÔFˆà×*Ò*¨4Ð1IÑJÔJˆ	Ø×(Ò(¨Ð/EÑFÔFˆð^
ÑZˆ�{ M°;ÀÈ{ð *4Ñ&ˆ˜Ø)Ñˆ	�9àÔ¥Ô 9Ò9Ñ9Ø ×4Ò4Ø×&Ò& tÐ-FÑGÔGÕIYÑI[ÔI[ñô ˆMð !×4Ò4Ø×&Ò& tÐ-FÑGÔGÕIYÑI[ÔI[ñô ˆMð !×4Ò4Ø×&Ò& tÐ-FÑGÔGÕIYÑI[ÔI[ñô ˆMð ×0Ò0°Ð7NÑOÔOˆKØ×0Ò0°Ð7NÑOÔOˆKØ×0Ò0°Ð7NÑOÔOˆKà×0Ò0Ø×&Ò& tÐ-DÑEÔEÕGZÑG\ÔG\ñô ˆIð ×0Ò0Ø×&Ò& tÐ-DÑEÔEÕGZÑG\ÔG\ñô ˆIˆIð .2×-?Ò-?Ø×&Ò& tÐ-FÑGÔGØ×&Ò& tÐ-FÑGÔGØ×&Ò& tÐ-FÑGÔGØ×&Ò& tÐ-DÑEÔEØ×&Ò& tÐ-DÑEÔEØ×&Ò& tÐ-DÑEÔEØ—’˜xÐ):Ñ;Ô;Ø—’˜x¨Ñ9Ô9ñ	.ô 	.Ñ*ˆO˜]ð ×2Ò2Ø×"Ò" 4Ð)AÑBÔBÕDTÑDVÔDVñ
ô 
ˆð ×.Ò.¨tÐ5KÑLÔLˆà×1Ò1°$×2HÒ2HÈÐO_Ñ2`Ô2`ÕbrÑbtÔbtÑuÔuˆØ×1Ò1°$×2HÒ2HÈÐO_Ñ2`Ô2`ÕbrÑbtÔbtÑuÔuˆØ×-Ò-¨d°NÑCÔCˆØ×-Ò-¨d°NÑCÔCˆà"$ÐØ×&Ò&¨¸ÀmÐ'TÑUÔUÐUØ×&Ò&¨¸ÀmÐ'TÑUÔUÐUØ×&Ò&¨	°7Ð';Ñ<Ô<Ð<ØÔ¥Ô 9Ò9Ð9Ø#×*Ò*à!ØØ!ØØ!ØØØð	ñô ð ð ð $×*Ò*¨O¸]Ð+KÑLÔLÐLØ×&Ò&Ø˜m¨^¸\È>Ð[gÐhñ	
ô 	
ð 	
ð 	 ×&Ò&ØnÐnÐnñ	
ô 	
ð 	
ð ˆØ×Ò˜dŸnšn¨o¸yÈ'Ð-RÐU]ÐT^Ñ_Ô_Ñ`Ô`Ð`Ø×Ò˜dŸiši¨°_ÀmÐ(TÐWeÐVfÐhrÑsÔsÑtÔtÐtØ×Ò˜dŸiši¨°>À<Ð(PÐS\ÐR]Ð_eÑfÔfÑgÔgÐgØ×Ò˜dŸmšm¨Y¨K¸*¸ÑFÔFÑGÔGÐGØ×Ò˜dŸiši¨°^À\Ð(RÐU^ÐT_ÐagÑhÔhÑiÔiÐiØ×Ò˜dŸhšh¨¸	Ð'BÐEUÐDVÐXdÑeÔeÑfÔfÐfØ×Ò˜dŸhšh¨Ð9IÐ'JÈ_ÐL]Ð_kÑlÔlÑmÔmÐmØÔ¥Ô 9Ò9Ñ9Ø×!Ò! $§)¢)¨X°}ÀkÐ,RÐU\ÐT]Ð_cÑ"dÔ"dÑeÔeÐeØ×!Ò! $§)¢)¨X°}ÀkÐ,RÐUZÐT[Ð]aÑ"bÔ"bÑcÔcÐcØ×!Ò! $§)¢)¨X°}ÀkÐ,RÐU\ÐT]Ð_cÑ"dÔ"dÑeÔeÐeà-1Ô->Å/ÔB`Ò-`Ð-`˜>˜>ÐflˆLØ×!Ò! $§+¢+¨w¸ÀiÐQZÐ.[Ð^iÐ]jÐlpÑ"qÔ"qÑrÔrÐrØ×!Ò! $§+¢+¨u°lÀIÈyÐ.YÐ\eÐ[fÐhlÑ"mÔ"mÑnÔnÐnØÔ ¥OÔ$FÒFÐFØ×%Ò%Ø—H’HØ$ i°¸"Ð>NÐPRÐT_ÐanÐoØ# [°-Ð@ñô ñô ð ñ ð Ô"¥oÔ&IÒIÐIØ×%Ò%Ø—H’Hà'Ø%Ø#Ø'Ø)Ø'Ø3ðð $ [°-Ð@ñô ñô ð ð ˜q’=�=Ø$(×$:Ò$:Ñ$<Ô$<�MØ$1ð [ð [˜ØÔ)×0Ò0°Ñ:Ô:Ð:ØFJÔFZ˜Ô4°X´]ÑCÐCØ”J×.Ò.Ý$Ô/µ´¸!¸ÀGÐ0LÑ0LÔ0LÐSXÐYÑYÔYÐ[_Ô[oñô ð øð Ô"¥oÔ&DÒDÐDØ×%Ò%Ø—O’OØ$ i°¸+À}ÐVfÐgØ#˜ñô ñô ð øð + Ð*Ð*ˆIØ0 hÐ0Ð0ˆLØ#×*Ò*¨I°|Ð+DÑEÔEÐEØ×!Ò!Ø—’Ø˜°Ð?OÐQZÐ[Ð^hÐjvÐ]wñô ñô ð ð 	×*Ò*¨>¸8ÐE\Ñ]Ô]Ð]à×$Ò$ _Ð6^Ð6^Ð6^Ñ_Ô_Ð_Ø×$Ò$ _Ð6^Ð6^Ð6^Ñ_Ô_Ð_àÔ×#Ò# DÑ)Ô)Ð)ØˆÔÐÐr   )r   r    r!   r   rØ   rN   rb  rz  r•  rF  rÙ   rÚ   s   @r   r[  r[  4  s”   ø€ € € € € ð/àð/ð ð/ð /ð /ð /ð /ð /ð 3ð 3ð 3ð,ð ,ð ,ð\(&ð (&ð (&ðTk ð k ð k ð k ð k ð k ð k r   r[  c                   óV   ‡ — e Zd Zdededefˆ fd„Zddedz  defˆ fd	„Zd
„ Z	dd„Z
ˆ xZS )ÚPhiOnnxModelrH   r·   rÝ   c                 ó"  •— t          ¦   «                              |¦  «         t          | j        ||¦  «        | _        t          | |¦  «        | _        t          | ¦  «        | _        t          | ¦  «        | _
        t          | ¦  «        | _        d S rR   )rM   rN   rÜ   rH   Úphi2_preprocessorr[  Úfission_transformer_blockrS  Úfission_causal_lm_headrI  Úfission_transformer_layernormr7  Úfission_transformer_embeddingrã   s       €r   rN   zPhiOnnxModel.__init__O  sy   ø€ Ý‰Œ×Ò˜ÑÔÐÝ!1°$´*¸iÈÑ!UÔ!UˆÔÝ)CÀDÈ)Ñ)TÔ)TˆÔ&Ý&GÈÑ&MÔ&MˆÔ#Ý-KÈDÑ-QÔ-QˆÔ*Ý-KÈDÑ-QÔ-QˆÔ*Ð*Ð*r   NFÚoptionsÚadd_dynamic_axesc                 óB  •— |€J ‚|j         }| j                             |¦  «         | j                             |¦  «         | j                             ¦   «          | j                             ¦   «          | j                             ¦   «          | j                             ¦   «          t          ¦   «          
                    ¦   «          t          | ¦  «        | _        t          | ¦  «        | _        | j                             ¦   «          | j                             ¦   «          d S rR   )Úattention_op_typerÏ  rS   rÎ  r4  ÚapplyrÑ  rÐ  rÒ  rM   rB  r	   Úfuse_slnr   Úfuse_bias_sln)r   rÓ  rÔ  rP   rO   s       €r   ÚoptimizezPhiOnnxModel.optimizeW  sÿ   ø€ ØÐ"Ð"Ð"ØÔ0ˆàÔ&×<Ò<¸\ÑJÔJÐJàÔ×.Ò.¨|Ñ<Ô<Ð<àÔ&×,Ò,Ñ.Ô.Ð.ØÔ*×0Ò0Ñ2Ô2Ð2ØÔ#×)Ò)Ñ+Ô+Ð+ØÔ*×0Ò0Ñ2Ô2Ð2å‰Œ×ÒÑÔÐõ 5°TÑ:Ô:ˆŒÝ=¸dÑCÔCˆÔØŒ×ÒÑÔÐØÔ× Ò Ñ"Ô"Ð"Ð"Ð"r   c                 ó¤   — i }g d¢}|D ])}|                       |¦  «        }t          |¦  «        ||<   Œ*t                               d|› �¦  «         |S )z8
        Returns node count of fused operators.
        )	rË   rÀ   rÆ   rÒ   ÚGeluÚBiasGelur»   r¢   ÚSkipLayerNormalizationzOptimized operators: )Úget_nodes_by_op_typerD   r?  r@  )r   Úop_countÚopsÚopÚnodess        r   Úget_fused_operator_statisticsz*PhiOnnxModel.get_fused_operator_statisticsl  so   € ð ˆð

ð 

ð 

ˆð ð 	&ð 	&ˆBØ×-Ò-¨bÑ1Ô1ˆEÝ˜u™:œ:ˆH�R‰LˆLå�ŠÐ6¨HÐ6Ð6Ñ7Ô7Ð7Øˆr   c                 ó  ‡— ‰€|                       ¦   «         Šdt          fˆfd„} |d¦  «         |d¦  «        z    |d¦  «        z    |d¦  «        z   } |d¦  «         |d	¦  «        z    |d
¦  «        z   } |d¦  «         |d¦  «        z   }|dk    o||k    o||k    }|dk    rt                               d¦  «         |dk    rt                               d¦  «         |dk    rt                               d¦  «         |S )zA
        Returns True when the model is fully optimized.
        NÚop_namec                 ó2   •— ‰                      | ¦  «        pdS )Nr   )Úget)ræ  Úfused_op_counts    €r   rà  z1PhiOnnxModel.is_fully_optimized.<locals>.op_countŠ  s   ø€ Ø!×%Ò% gÑ.Ô.Ð3°!Ð3r   rË   rÀ   rÆ   rÒ   rÜ  rÝ  r»   r¢   rÞ  r   zLayer Normalization not fusedzGelu (or FastGelu) not fusedz+Attention (or MultiHeadAttention) not fused)rä  rV   r?  ÚdebugÚwarning)r   ré  rà  rÍ   ÚgeluÚ
layer_normÚ
is_perfects    `     r   Úis_fully_optimizedzPhiOnnxModel.is_fully_optimizedƒ  sW  ø€ ð Ð!Ø!×?Ò?ÑAÔAˆNð	4�cð 	4ð 	4ð 	4ð 	4ð 	4ð 	4ð ˆH�[Ñ!Ô!ØˆhÐ+Ñ,Ô,ñ-àˆhÐ,Ñ-Ô-ñ.ð ˆhÐ'Ñ(Ô(ñ)ð 	ð ˆx˜ÑÔ ( (¨:Ñ"6Ô"6Ñ6¸¸À*Ñ9MÔ9MÑMˆØ�XÐ2Ñ3Ô3°h°hÐ?WÑ6XÔ6XÑXˆ
à !’mÐZ¨)°tÒ*;ÐZÀ*ÐPYÒBYˆ
à˜Š?ˆ?Ý�LŠLÐ8Ñ9Ô9Ð9à�1Š9ˆ9Ý�LŠLÐ7Ñ8Ô8Ð8à˜Š>ˆ>Ý�NŠNÐHÑIÔIÐIàÐr   )NFrR   )r   r    r!   r   rØ   rN   r   ÚboolrÚ  rä  rï  rÙ   rÚ   s   @r   rÌ  rÌ  N  s´   ø€ € € € € ðR˜jð R°Sð RÀsð Rð Rð Rð Rð Rð Rð#ð # °Ñ 4ð #Ètð #ð #ð #ð #ð #ð #ð*ð ð ð.ð ð ð ð ð ð ð r   rÌ  ))Úloggingr   Únumpyr   Údynamo_onnx_helperr   Úfusion_baser   Úfusion_optionsr   r   Úfusion_skiplayernormr   r	   Úfusion_utilsr
   Úonnxr   r   r   r   r   Ú
onnx_modelr   r   r?  r   r$   r,   r/   r3   r:   r=   r@   rG   rÜ   r7  rI  rS  r[  rÌ  r"   r   r   ú<module>rú     so  ðð Ð Ð Ð Ð Ð à Ð Ð Ð Ø /Ð /Ð /Ð /Ð /Ð /Ø Ð Ð Ð Ð Ð Ø 9Ð 9Ð 9Ð 9Ð 9Ð 9Ð 9Ð 9Ø _Ð _Ð _Ð _Ð _Ð _Ð _Ð _Ø $Ð $Ð $Ð $Ð $Ð $Ø IÐ IÐ IÐ IÐ IÐ IÐ IÐ IÐ IÐ IÐ IÐ IÐ IÐ IØ  Ð  Ð  Ð  Ð  Ð  à	ˆ�8Ñ	Ô	€ð'ð 'ð 'ð 'ð 'ñ 'ô 'ð 'ð
:ð :ð :ð :ð :ñ :ô :ð :ð
:ð :ð :ð :ð :ñ :ô :ð :ð
:ð :ð :ð :ð :ñ :ô :ð :ð
ð ð ð ð ñ ô ð ðð ð ð ð ñ ô ð ðð ð ð ð ñ ô ð ðð ð ð ð ñ ô ð ðVð Vð Vð Vð Vˆfñ Vô Vð Vðrp,ð p,ð p,ð p,ð p,Ð'ñ p,ô p,ð p,ðf ð  ð  ð  ð   Wñ  ô  ð  ðD ð  ð  ð  ð   Wñ  ô  ð  ðB ð  ð  ð  ð  ¨ñ  ô  ð  ðBW ð W ð W ð W ð W  ñ W ô W ð W ðtSð Sð Sð Sð S�9ñ Sô Sð Sð Sð Sr   