fPiddlmZddlZddlmZddlmZddlm Z m Z ddl m Z m Z mZmZddlmZeeZGdd ZGd d eZdS) ) getLoggerN)Fusion)AttentionMaskFormat) FusionUtils NumpyHelper) NodeProto TensorProtohelper numpy_helper) OnnxModelcLeZdZdZdefdZdefdZdZdZ de d e d zfd Z d S) AttentionMask: Fuse Attention subgraph into one Attention node. modelc||_i|_i|_t||_t j|_||_ dSN) r mask_indice mask_castedrutilsr MaskIndexEnd mask_formatget_opset_version opset_version)selfrs }/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/onnxruntime/transformers/fusion_attention.py__init__zAttentionMask.__init__sM  '' .;"4466rc||_dSr)r)rrs rset_mask_formatzAttentionMask.set_mask_format s&rcR||jvr||j|ksJ||j|<dSr)r)rmask mask_indexs rset_mask_indicezAttentionMask.set_mask_indice#s< 4# # #!1$!77777!+rcxt|jdksJtt|jSNr)lenrnextiter)rs rget_first_maskzAttentionMask.get_first_mask(s74#$$q((((D)**+++rmask_2dreturnNc ^|jtjkrdS||jvr |j|S|j|r|j|\}}n|j|\}}d}|r ||j |<|jtj kr ||j|<|S|j d}|j dkrwtjd|g|g|j dd}|jtjddgtjd d gnd }|j|=|jtj|t*jdgdgd tjd||g|g|j dd}|jtjd d g|j|||j|<|S)NTr" ReduceSum MaskReduceSuminputsoutputsnameaxeskeepdimsrort_const_1_reduce_sum_axesFr3 data_typedimsvalsraw)rrNoMaskrrfind_graph_inputrcast_graph_input_to_int32cast_input_to_int32rrcreate_node_namerr make_node attributeextendmake_attributeget_initializeradd_initializer make_tensorr INT64add_node)rr*casted input_name _cast_node output_namemask_index_node axes_names r process_maskzAttentionMask.process_mask,sW  29 9 94 d& & &#G, , : & &w / / !%!E!Eg!N!N FJJ%)Z%C%CG%L%L "J F  3(2D W %  2@ @ @(2D W % j11,??   " "$."|$ Z00oNN O  % , ,f.CFQC.P.PRXRghrtuRvRv-w x x x x6Iz)))44< **&&"-"3SS! %."I.$ Z00oNN O  % , ,f.CJPQ.R.R-S T T T O,,,$/!r) __name__ __module__ __qualname____doc__r rrrr#r)strrQrrrrs7i7777'+>'''',,, ,,,8C8C$J888888rrc(XeZdZdZdddddgfdededed edzd ed ed ee ffd Z de de eeffdZ de de eeffdZde fdZde fdZde de de fdZde de de fdZde de dzde dzd e de dzf d!Zd"e d#e d$e de de dzde dzde e e e ffd%Z d5d"e d#e e zdzd$e e zdzde de dzde dzdeded'e d(e de d)ede de d*e d+e d,ede dzf$d-Z d6d.e dzd"e d#e d$e de de de deded/e d'e d0e d1ede de d*e d+e d2edzde dzf&d3Zd4ZxZS)7FusionAttentionrNFSkipLayerNormalizationLayerNormalizationr hidden_size num_headsattention_maskuse_multi_head_attention!disable_multi_head_attention_biassearch_op_typesc|rdnd}t|||||_||_|r|nt ||_||_||_d|_d|_ d|_ d|_ d|_ dS)NMultiHeadAttention AttentionT) superrr\r]rr^r_r`mask_filter_valuenum_heads_warninghidden_size_warning shape_infershape_infer_done) rrr\r]r^r_r`raattention_op_name __class__s rrzFusionAttention.__init__ls5M]00R]  1?CCC&"0>XnnMRWDXDX(@%1R.!%"&#'  $rconcatr+ct|jdkr|j|jd}|j|jd}t |t jrI|jdkr>t |t jr$|jdkr|d|d|dzfS|j|j fS)aU Detect num_heads and hidden_size from Concat node in the following subgraph: SkipLayerNormalization or EmbedLayerNormalization / | MatMul Shape | | Add Gather(indices=0) | | | Unsqueeze | | | Concat (*, -1, 12, 64) | / Reshape | Transpose r5r) r&inputrget_constant_value isinstancenpndarraysizer]r\)rrmr] head_sizes r)get_num_heads_and_hidden_size_from_concatz9FusionAttention.get_num_heads_and_hidden_size_from_concats$ v|   ! ! 55fl1oFFI 55fl1oFFI9bj11 ANa''y"*55(Na'' |Yq\IaL%@@@~t///r reshape_qcT|j|jd}|q|j|d}| |jdkr||St d|jd|j|j fSt|tj r+t|dks|ddks |ddkr)t d ||j|j fS|d}|d}||z}|jdkr:||jkr/|jr(t d |j|d |_|j dkr:||j kr/|jr(t d |j |d |_||fS) zDetect num_heads and hidden_size from a reshape node. Args: reshape_q (NodeProto): reshape node for Q Returns: Tuple[int, int]: num_heads and hidden_size r5NConcatz%s is not initializer.rorprrqzGq_shape_value=%s. Expected value are like [0, 0, num_heads, head_size].z>--num_heads is %d. Detected value is %d. Using detected value.Fz@--hidden_size is %d. Detected value is %d. Using detected value.)rrsrr get_parentop_typeryloggerdebugr]r\rtrurvr&rgwarningrh)rrz q_shape_valuermr]rxr\s rget_num_heads_and_hidden_sizez-FusionAttention.get_num_heads_and_hidden_sizes 55ioa6HII  Z**9a88F!fn&@&@EEfMMM LL19?13E F F F>4#33 3M2:66 4=!!Q&&a A%%q)9Q)>)> LLbdq r r r>4#33 3!!$ !!$ )+ >A  )t~"="=% /TVZVdfo*/&  a  K43C$C$C' 1VX\Xhju,1(+%%radd_qkc|js'|jd|_d|_|jdS|j|jd}|j|jd}||t d|dS||krt d|dS|jdS)NT)updaterr5zone of the inputs of %s is Nonez)the shape of two inputs of %s is not same)rjrinfer_runtime_shaperiget_edge_shaperrrr)rr input_0_shape input_1_shapes rget_add_qk_strzFusionAttention.get_add_qk_strs$ )#z==T=JJD $(D !   #4(77 QHH (77 QHH  M$9 LL:F C C C4 M ) ) LLDf M M M4|Arcdzttfd|j}t|dkrSt|dksJ|jd}t jdfdt|j Dg|d}|j ||j |j |<S)N_maskc&|jdkSr%)output)nodemask_output_names rz0FusionAttention.reshape_add_qk..st{1~AQ/Qrr5rr|cg|]}SrWrW).0_rs r z2FusionAttention.reshape_add_qk..s:::qF:::rr1r2r3axis) listfilter nodes_to_addr&rrAr rBranger]appendthis_graph_namenode_name_to_graph_name)rr concat_nodeconcat_node_nameconcat_add_qk_fp32rs ` @rreshape_add_qkzFusionAttention.reshape_add_qks"G+6"Q"Q"Q"QSWSdeeff {  q # #;1$$$$:66x@@#- ::::E$.$9$9:::%&!       !34449=9M$%56rpast_kpast_vc |jd}|jd}|dzdd}|dzdd}tjd|g|g|dg}tjd|g|g|dg}|j||j||j|j|<|j|j|<|jd} |dd ddd d } tjd||g| g| d } |j| |j|j| <| S) zConcatenate past_k and past_v inputs to create past_kv input. Args: past_k (str): name of past K value past_v (str): name of past V value Returns: kv_output_name (str): name of past KV value Unsqueeze_5d.rr)r1r2r3r4r|z.valuez.kv_value_kvr) rrAreplacer rBrrrr) rrrunsqueeze_k_nameunsqueeze_v_name k_5d_name v_5d_namek_5dv_5drkv_output_name concat_kvs rrzFusionAttention.concat_kvs :66{CC:66{CCe^,,S#66 e^,,S#66  8K!     8K!       &&&   &&&9=9M$%569=9M$%56 :66x@@%88@@cJJRRS[]bcc$ y)#$!       +++9=9M$%56rpresent_k_namepresent_v_namekv_nodec$d\}}|j|}|j|}|Jtjt jdd|}|j||j|Jtjt jdd|}|j||j|jd}|jd} tj d||g|g|d } tj d||g|g| d } |j | |j | |j|j |<|j|j | <dS) a?Split kv_node containing present KV values into separate present K and present V values. Args: present_k_name (str): name of output to store present K value in present_v_name (str): name of output to store present V value in kv_node (str): name of present KV values )index_0index_1Nrint64)dtyper3r5Gatherr)rrFr from_arrayruarrayrGrrAr rBrrr) rrrrk_indexv_indexk_dimv_dim gather_k_name gather_v_name present_k present_vs rsplit_kvzFusionAttention.split_kv8s0 **733 **733 = +BHQg,F,F,FWUUUE J & &ud.B C C C = +BHQg,F,F,FWUUUE J & &ud.B C C C 33H==  33H== $ W%#$    $ W%#$       +++   +++6:6J$]36:6J$]333rq_addk_addv_add name_prefixcD|j|jdp$|j|jd}tj|}t j|}t j|}|^|j|jdp$|j|jd} tj| }|^|j|jdp$|j|jd} tj| }t j|||fd} dt j|j z} |dz} | | |j | g| | S)Nr5rrrq _qkv_biasr3r9r:r;) rrFrrrto_arrayru zeros_likestackprodshaperGr9)rrrrrq_biasqbkbvbk_biasv_biasqkv_bias qkv_bias_dim bias_names rcreate_combined_qkv_biasz(FusionAttention.create_combined_qkv_biases|++EKN;;itz?Y?YZ_ZefgZh?i?i  !& ) ) ]2   ]2    Z// A??m4:C]C]^c^ijk^lCmCmF%f--B  Z// A??m4:C]C]^c^ijk^lCmCmF%f--B8RRLq1112728,,, +-  &    rq_matmulk_matmulv_matmulc |jd}|jd|jdkr|jd|jdksJ|j|jd}|j|jd} |j|jd} t j|} t j| } t j| } | j| jkr| j| jksJ| jd}tj| | | fd |d|zf}|dz}| ||j |jd|jdg||dz}tj d|jd|g|g| }|j|j|<|g}|d z}| |t jdgdgd |d z}| |t jdg|gd |dz}| |t jdgd|zgd |dz}| |t jdgd|zgd |dz}| |t jdgdgd |dz}tj d||||g|g|jd }|j|j|j<|dz}tj d||||g|g|jd }|j|j|j<|dz}tj d||||g|g|jd }|j|j|j<|}|} |}!||||g|jr||j|jdrdnd}"tjt j|j|j|"r8||jd|"z <|}|||j|j|j<||j|jdrdnd}"tjt j|j|j|"r8||jd|"z <|} |||j|j|j<||j|jdrdnd}"tjt j|j|j|"r8||jd|"z <|}!|||j|j|j<|j||| |!fS)aCreate packed QKV MatMul node before MultiHeadAttention node. This is for the scenario where an Attention node should be created but cannot be created because past_key and past_value are separate inputs and not one concatenated input. Args: q_matmul (NodeProto): name of MatMul from Q path - (batch_size, sequence_length, hidden_size) k_matmul (NodeProto): name of MatMul from K path - (batch_size, sequence_length, hidden_size) v_matmul (NodeProto): name of MatMul from V path - (batch_size, sequence_length, hidden_size) q_add (NodeProto): name of Add from Q path k_add (NodeProto): name of Add from K path v_add (NodeProto): name of Add from V path Returns: q_output (NodeProto): Slice node for Q k_output (NodeProto): Slice node for K v_output (NodeProto): Slice node for V MatMulrr5rrq _qkv_weightr_qkv_outr0_q_start_indexFr8_k_start_index_v_start_indexrp_end_of_qkv_index_qkv_last_axis_q_outSlice_k_out_v_out)rrArrrFrrrrurreshaperGr9r rBrrr rIr3rDr`anyrr)#rrrrrrrmatmul_node_nameq_weightk_weightv_weightqwkwvwd qkv_weightqkv_weight_nameqkv_matmul_output qkv_matmul qkv_nodes q_slice_name k_slice_name v_slice_nameend_of_qkv_nameqkv_last_axis_nameq_slice_outputq_slicek_slice_outputk_slicev_slice_outputv_sliceq_outputk_outputv_outputinitializer_inputs# rcreate_packed_qkv_matmul_nodez-FusionAttention.create_packed_qkv_matmul_nodesP4 :66x@@~a HN1$555(.:Kx~^_O`:`:`:``:--hnQ.?@@:--hnQ.?@@:--hnQ.?@@  !( + +  !( + +  !( + +x28##BH(<(<(<< HQKXr2rl333;;QAJGG *]:  ("1%z'7':;    "-z9% N1%7&'!    :>9M$%56L (*::  ,+:KSTRU]^\_ejkkk'*::  ,+:KSTRU]^\_ejkkk'*::  ,+:KSTRU]^ab]b\cinooo*-@@ /[=NVWUX`ade`e_flqrrr-0@@ "4 @QYZX[cebflqrrr)H4" %|\CUV#$,,W55     6:5I$W\2)H4" %|\CUV#$,,W55     6:5I$W\2)H4" %|_FXY#$,,W55     6:5I$W\2'7G4555  1 T )-)C)CEKPQN)S)S$ZAAYZ!6+.tz/I/I%+VgJh/i/ijjkkT9GEK$5 56$H$$U+++?C?SD0< )-)C)CEKPQN)S)S$ZAAYZ!6+.tz/I/I%+VgJh/i/ijjkkT9GEK$5 56$H$$U+++?C?SD0< )-)C)CEKPQN)S)S$ZAAYZ!6+.tz/I/I%+VgJh/i/ijjkkT9GEK$5 56$H$$U+++?C?SD0<   +++8++rrkey_padding_maskunidirectionalrr packed_qkvcz|dksJ|dkr'||zdkrtd||dSd|jjD}|jd}g}|rY|||||||\}}}||jd|jd|jdgn(t|trt|tr{|j r:||jd|jd|jdgn||jd|jd|jdgnt|trlt|trW||vrS||vrO|j r$||jd||gn&||jd||gndS|j s.| ||||}||n|d| r|r|| | | |gn| s| r|| | g| g}|r|r|||gtjd|||}d |_|jtjd || r:|jtjd t)| |d|S) aCreate a MultiHeadAttention node. Args: q_matmul (NodeProto): name of MatMul from Q path - (batch_size, sequence_length, hidden_size) k_matmul (NodeProto): name of MatMul from K path - (batch_size, sequence_length, hidden_size) or (batch_size, num_heads, past_sequence_length, head_size) v_matmul (NodeProto): name of MatMul from V path - (batch_size, sequence_length, hidden_size) or (batch_size, num_heads, past_sequence_length, head_size) q_add (NodeProto): name of Add from Q path k_add (NodeProto): name of Add from K path v_add (NodeProto): name of Add from V path num_heads (int): number of attention heads. If a model is pruned, it is the number of heads after pruning. hidden_size (int): hidden dimension. If a model is pruned, it is the hidden dimension after pruning. output (str): output name of MHA key_padding_mask (str): name of key padding mask add_qk (str): name of add after Q x K' unidirectional (bool): whether to apply causal attention mask automatically or not past_k (str): name of past K value - (batch_size, num_heads, past_sequence_length, head_size) past_v (str): name of past V value - (batch_size, num_heads, past_sequence_length, head_size) present_k (str): name of present K value - (batch_size, num_heads, sequence_length, head_size) present_v (str): name of present V value - (batch_size, num_heads, sequence_length, head_size) packed_qkv (bool): whether to combine MatMuls from Q, K, V paths Note: This is for the scenario where an Attention node should be created but cannot be created because past_key and past_value are separate inputs and not one concatenated input. Returns: Union[NodeProto, None]: the node created or None if failed. r9input hidden size %d is not a multiple of num of heads %dNch|] }|j SrWr)rrs r zBFusionAttention.create_multihead_attention_node..HsLLL4TYLLLrrdrrcr0 com.microsoftr]r)rrrgraphrrrArrDrrtrr`rVrrr rBdomainrCrEintincrease_counter)rrrrrrrr]r\rrrrrrrrrgraph_input_names mha_node_name mha_inputsrr r r mha_outputsmha_nodes rcreate_multihead_attention_nodez/FusionAttention.create_multihead_attention_nodes^1}}}} ?? i 7A== LLTVacl m m m4LL4:3C3C3E3E3KLLL 33K@@   (,(J(J )) %GWg   w~a0'.2CW^TUEVW X X X X ) , , Hi1P1P 5 `!!5<?HOA4F UV"XYYYY!!8?1#5xq7I8?[\K]"^____ x % % 8S)) ------5 L!!5<?Hh"GHHHH!!8?1#5x"JKKKK4 5 "55eUE=YYI   i ( ( ( (   b ! ! !  :f :   /H I I I I  : :   /8 9 9 9h  7 7    95 6 6 6#      *!!&"7 Y"O"OPPP  d   % %f&;Create an Attention node. Args: mask_index (str | None): mask input q_matmul (NodeProto): MatMul node in fully connection for Q k_matmul (NodeProto): MatMul node in fully connection for K v_matmul (NodeProto): MatMul node in fully connection for V q_add (NodeProto): Add bias node in fully connection for Q k_add (NodeProto): Add bias node in fully connection for K v_add (NodeProto): Add bias node in fully connection for V num_heads (int): number of attention heads. If a model is pruned, it is the number of heads after pruning. hidden_size (int): hidden dimension. If a model is pruned, it is the hidden dimension after pruning. first_input (str): first input name output (str): output name add_qk_str (str): name of Add node after Q x K' causal: whether it is uni-directional mask. past_k (str): name of input for past K value past_v (str): name of input for past V value present_k (str): name of output to store present K value present_v (str): name of output to store present V value scale: scale before softmax Returns: Union[NodeProto, None]: the node created or None if failed. rrNTFr5)NNNzl is not an initializer. Please set do_constant_folding=True in torch.onnx.export to unblock attention fusionzInput hidden size (%d) is not same as weight matrix dimension of q,k,v (%d). Please provide a correct input hidden size or pass in 0rrqrdrrrzVMultiHeadAttention does not support relative_position_bias: cannot fuse the attention.rcr0rz.key_keyrrrr]rr'qkv_hidden_sizesrf)!rrrrFrrprintrrrrrur concatenaterrAr_rGr9rrrr rBrrrrrrCrDrErffloat)5rr"rrrrrrr]r\r$rr%r&rrrrr'has_biasrrrrrrrrr qw_in_size kw_in_size vw_in_sizeis_qkv_diff_dims qw_out_size kw_out_size vw_out_sizeqkv_weight_dimrrrrrr q_bias_shape k_bias_shape v_bias_shapeattention_node_nameattention_inputsattention_node past_existspast_kvattention_outputs present_kvs5 rcreate_attention_nodez%FusionAttention.create_attention_nodesJ\1}}}} ?? i 7A== LLTVacl m m m4 =U]u}H:--hnQ.?@@:--hnQ.?@@:--hnQ.?@@!1  Z// A??m4:C]C]^c^ijk^lCmCmFZ// A??m4:C]C]^c^ijk^lCmCmFZ// A??m4:C]C]^c^ijk^lCmCmF  f  t   >!$ggg   4  !( + +  !( + +  !( + +x28####Xa[ Xa[ Xa[ Z5555:555555 ??{j88 NNJ    ! 8rx  #  gbhqrrl++ gbhqrrl++ gbhqrrl++   -R 1===J(;6DNN2r2,Q777J_N &*  0%f--B%f--B%f--B728,,L728,,L728,,L<>>>>;>>>>>>;.... 0>2r2,Q???+l:\I 8RRLq999 </ "j99+FF,   (=8", #n"5"56 !       (;6 *,''( !     (9 /  uvvvt"""#k1   % '' 333#-$'( N  ! !"6 7 7 7 7#m35=E#k112   % '' 3333 ''+++ +VK 1..88 ''000 4"0$++B/// '' 333!'  @Y @&..vr::BB62NNVVWZ\_`` !((444 iJ???#-')( N  ! !+ . . . / '')>{I)V)V(WXXX  Z  $ + +V-BCSUV-W-W,X Y Y Y    $ + +V-B7E-R-R,S T T T    $ + +&'9KVa;bccd     ! -  $ + +V-BCVX]^b^tXuXu-v-v,w x x xrc@|}|}|jdkr#|j|dd}||}ndS|j|gdgd}d}| |\} } } } } n,|j|gdgd}||\} }} } ndSg} t |jD]7\}}||vr ||djdkr"| |8t| dkrdS| d}|j|d d}|p||jd}|6t|d kr#|d}|jdkr|jd}nWdS|!t|d kr|jd}n2dS|jdkr%||}|D]}|jdkr |jd}||}|jd kr%t|jd kr |jd}||}d|D}| ddkrdS|j| gdgd}|t ddS|\} } }}d}d}d}d}gdgdfgdgdfgdgdfgdgdfgdgdfgdgdfd }d} | D]S\}!}"|j| |"d|"d} | 0|!d!krd"}n|!d#krd"}n|!d$krd"}n|!d%krd"}| t d&dSd}#d}$d}%d}&|r| \} }%}$} n%|r| \} }#}%}$n|r| \} } }$n|r | \} }#}$}&} n| \} }#} }$|&p|$}&|j|&gdgd'}'|'>|j|&gd(gd)}'|'t d*dS|'d+}(|'d,})|'d-}*|$}+|rC|j|$d d.gddg},|,t d/dS|,\}+} |j|+gd|rdnddddg}-|->|j|$gd0gd1}-|-t d2dS|-d,}.|-d-}/d}0d3}1|r6|j |%gd4gdfgd5gdfgd6gd7fg|\} }0} n|re|j |%gd8gd7fgd5gdfg|\} }0} |#4||#}1|1t d9|#dSn?|rn<|j |#gd:gd;fgd<gd=fgd>gd?fgd@gdAfg|\} }0} |s|0t dBdS|st|0dkr|j|0d\} }2|28t!|2t"jr|2jdkrt)|2dkrdSt)|2dCkrt)|2|_|jd|kr]|*jd|krM|/jd|kr=|s+|j|0d-jdnd}3|| n| }4||(\}5}6|5dks|6dkrtdDdS||3|*|/||)|.||5|6||4jd|1E }7|7dS|j|7|j|j|7j<||jd}8dF|8z}9|dG|8zt@j!d gdd|5tE|6|5z gdH}:|j#tIj%dI|4jd|:jg|9gdJ|8z|j|9|jd<|j&'|4| | g|j&'| |j&'|j(s|'n |'dd-|j&'|j(s|-n |-dd-|j&'|j(s|n |dd-d"|_)dSdSdSdS)KNr[Addr)rCrReshape Transposer)NNrrr)rCEinsumrEr)r5Nrrr5MulrprZrocg|] }|j SrW)r~)rchilds rrz(FusionAttention.fuse..s>>>E%->>>rrrq)rErDrCr)r5rrNz&fuse_attention: failed to match v pathF)SoftmaxrCDivr)rrNr)rKrCrGr)rKWhererrL)rrrpr)rKrCrMr)rrrrp)rKrLr)rrr)rKrCrrGSqrt)rrNrr5)path1path2path3path4path5sdparQTrRrSrTz'fuse_attention: failed to match qk path)rrrN)rLrErDrCr)rrrrNz&fuse_attention: failed to match q pathrrNz/fuse_attention: failed to match mul sqrt q path)rErErDrCr)r5rrrNz&fuse_attention: failed to match k pathr)ExpandrDEqual)rXrr)CastrWrDrX)rrrr)rYrXrrz6fuse_attention: failed to verify shape inference of %s)rGSubrYrr)Nrr5rr)rGrZrr)Nrr5r)rMrYrZrWrr)Nrrr5rr)rMrYrZrYrWrr)Nrrr5rrrz)fuse_attention: failed to match mask pathizmFailed to detect num_heads and hidden_size for Attention fusion. Please specify those parameters in argument.) r"rrrrrrr]r\r$rr%edge_modified_shape_modified_tensorr8rDreshape_modified_)*r~r match_parentmatch_parent_path enumeraterrrrr&countrritemsmatch_parent_pathsrget_constant_inputrtrurvrwr-rfr^rQrrrArrrr3rGr rIrrJr rBnodes_to_removerDr_ prune_graph);rrinput_name_to_nodesoutput_name_to_nodenormalize_node start_nodeadd_before_layernormr einsum_noder reshape_qkv transpose_qkv matmul_qkv other_inputs_i node_input root_inputmul_before_layernorm mul_childrenlayernorm_nodechildrenrJ parent_nodechildren_typesv_nodesadd_vmatmul_v is_distillis_distill_addis_no_mask_attentionis_sdpaqk_pathsqk_nodeskvr matmul_qkwhere_qkafter_qq_nodesrzadd_qmatmul_qafter_k mul_k_nodesk_nodesadd_kmatmul_k mask_nodesr%mul_valr"attention_last_node q_num_heads q_hidden_sizenew_node unique_indexnew_edge shape_tensors; rfusezFusionAttention.fusevs #  !%9 9 9#':#:#:>5RS#T#T #/1 J00  ? ? ? ! ! !    =F :Q; zz 44DDDoooI$>G;K  ' (899 , ,NB !444Yq\0333    + + + + |   ! ! F!!_  $z66z5!LL  +./C/J1/MNL'C ,=,=,B,B!-a!)-AAA!/!6q!9JJF)c,.?.?1.D.D18;   #'; ; ;*:6H! 1 1=$888!&aJ**5  ": : :s;CU?V?VZ[?[?[$+A.J&z2>>X>>>    ) )Q . . F*..z;d;d;dfufufuvv ? LLA B B B F")Auh $999???K999???K;;;\\\J;;;\\\J222III>@@@BTBTBTU   NN$$  DAqz33J!adKKHG||! g!%g'+$$f    LLB C C C F   1*2 'Q)QQ  1/7 ,Q)) ! 1 ( Q99  119 .Q 7AA(0 %Q9&Y*..w8a8a8acrcrcrss ?j22@@@"""G  EFFFBK  2;  '*66y5&/TUW[S\]]K" NOOO&LWa*.. >>>gATSTVWYZ\`@a   ?j22FFF"""G  EFFF 2;  % #z<<333YYY?888)))D;;;\\\J $   Az11 #z<<@@@,,,O888)))D$    Az1!!0088 %LL!Y[abbbF !  #z<<EEEGYGYGYZ===OQQQShShShiYYY[s[s[st $   Az1$  (: LLD E E E F# 8J!(;(;66z!}EEJAw"7BJ77!!  * *x~a/@J/N/NS[SabcSdhrSrSrZnx,99*R.:Nq:QRRRtxJ1<1D++- )-)K)KI)V)V &Ka=A#5#5C11%!!!%)&*1!4%2  H   $ $X . . .:>:ND ( 7&*03 +l:#330<?)/Q S1L-M-MN 4    ##$!,3A6 8IJ! +l:  ((0 !!$  ' ')@T%U%%%% % &, % #' %,0%c%%%%%%40 0eTWY\T\o0000>,&y,&U3PS8_,&,&,&,&\Y* S    255S5S5555n+Ks+KC+K#+K+K+K+KZ44    T  rs ......11111111============ 8  SSSSSSSSl~$~$~$~$~$f~$~$~$~$~$r