Ë
    cQêjâg  ã                   ó¦  — d dl Z d dlZd dlZd dlZd dlZd dlZd dlZd dlZd dlZ	d dl
Z
d dlZd dlZd dlmZmZ d dlmZmZ d dlmZmZmZmZmZmZ d dlmZ d dlmZmZmZ d dl m!Z! d dl"m#Z#m$Z$m%Z% d dl&Z' ejP                  e)«      Z*d	„ Z+d
e jX                  de-fd„Z.d
e jX                  fd„Z/d„ Z0d„ Z1d„ Z2d„ Z3d„ Z4d„ Z5dd„Z6d„ Z7e)dk(  r e7«        yy)é    N)Úmeasure_memoryÚsetup_logger)Úget_rankÚget_size)Úadd_io_bindings_as_ortvaluesÚ%get_merged_sample_with_past_kv_inputsÚget_msft_sample_inputsÚget_sample_inputsÚget_sample_with_past_kv_inputsÚverify_ort_inputs)ÚORTModelForCausalLM)ÚProfilerActivityÚprofileÚrecord_function)Útrange)Ú
AutoConfigÚAutoModelForCausalLMÚAutoTokenizerc                 óø   — | j                   dv ry| j                   dk(  r	 t        |j                  «      S t        |j                  «       «      S # t        $ r" t        |j                  j
                  «      cY S w xY w)N¾   úhf-pt-eagerúhf-pt-compiler   úhf-ort)Úbenchmark_typeÚlenÚinputs_namesÚ	ExceptionÚdecoderÚinput_namesÚ
get_inputs)ÚargsÚmodels     út/home/oem/CyberBrain26-unified/.venv/lib/python3.12/site-packages/onnxruntime/transformers/models/llama/benchmark.pyÚget_ort_model_inputs_lenr$   (   st   € Ø×ÑÐ>Ñ>ØØ×Ñ˜hÒ&ð	2ä�u×)Ñ)Ó*Ð*ô ˆu×ÑÓ!Ó"Ð"øô ò 	2ä�u—}‘}×0Ñ0Ó1Ò1ð	2ús    A Á(A9Á8A9r!   Úort_model_inputs_lenc                 óÖ  — d\  }}| j                   dk(  rdn| j                  j                  }| j                   dv rt        | j                  | j                  | j
                  | j                  d¬«      }t        | j                  | j                  | j
                  | j                  | j                  d¬«      }||fS | j                   dv �r*|d	k(  rt        | j                  | j                  | j
                  | j                  d¬«      }t        | j                  | j                  | j
                  | j                  | j                  d¬«      }||fS t        | j                  | j                  | j
                  | j                  d
|| j                  | j                  dd¬«
      }t        | j                  | j                  | j
                  d| j                  || j                  | j                  dd¬«
      }||fS | j                   dk(  r¼t        | j                  | j                  | j
                  | j                  d
|| j                  | j                  dd| j                  ¬«      }t        | j                  | j                  | j
                  d| j                  || j                  | j                  dd| j                  ¬«      }||fS | j                   dk(  r“|dkD  }t        | j                  | j
                  d
| j                  || j                  | j                  |¬«      }t        | j                  | j
                  | j                  d|| j                  | j                  |¬«      }||fS t        d«      ‚)N©NNúort-msfti   r   T)Úreturn_dict)Úuse_fp16r)   >   r   é   r   Úpt)Úseq_lenÚpast_seq_lenÚmax_seq_lenr*   Úuse_buffer_shareÚenginer)   é   úort-convert-to-onnxÚort)r-   r.   r/   r*   r0   r1   r)   Ú
world_sizeé   )r.   r-   r/   r*   r0   Úsplit_kvz/Unable to auto-detect inputs for provided model)r   ÚconfigÚmax_position_embeddingsr
   Útarget_deviceÚ
batch_sizeÚsequence_lengthr   r*   r   r0   r5   r	   r   )r!   r%   Úinit_inputsÚiter_inputsr/   r7   s         r#   r    r    5   s>  € Ø)Ñ€K�ð
 ×-Ñ-°Ò;‘$ÀÇÁ×AdÑAd€Kà×ÑÐ>Ñ>Ü'Ø�K‰KØ×ÑØ�O‰OØ× Ñ Øô
ˆô 5Ø�K‰KØ×ÑØ�O‰OØ× Ñ Ø—]‘]Øô
ˆð^ ˜Ð#Ð#ðM 
×	Ñ	 
Ò	*Ø 1Ò$ä+Ø—‘Ø×"Ñ"Ø—‘Ø×$Ñ$Ø ôˆKô 9Ø—‘Ø×"Ñ"Ø—‘Ø×$Ñ$ØŸ™Ø ôˆKðx ˜Ð#Ð#ôe @Ø—‘Ø×"Ñ"Ø—‘Ø×,Ñ,ØØ'ØŸ™Ø!%×!6Ñ!6ØØ ôˆKô @Ø—‘Ø×"Ñ"Ø—‘ØØ!×1Ñ1Ø'ØŸ™Ø!%×!6Ñ!6ØØ ôˆKðL ˜Ð#Ð#ðs 
×	Ñ	Ð 5Ò	5ä;Ø�K‰KØ×ÑØ�O‰OØ×(Ñ(ØØ#Ø—]‘]Ø!×2Ñ2ØØØ—‘ô
ˆô <Ø�K‰KØ×ÑØ�O‰OØØ×-Ñ-Ø#Ø—]‘]Ø!×2Ñ2ØØØ—‘ô
ˆðT ˜Ð#Ð#ð9 
×	Ñ	 
Ò	*à'¨!Ñ+ˆä,Ø�K‰KØ�O‰OØØ×(Ñ(Ø#Ø—]‘]Ø!×2Ñ2Øô	
ˆô -Ø�K‰KØ�O‰OØ×-Ñ-ØØ#Ø—]‘]Ø!×2Ñ2Øô	
ˆð ˜Ð#Ð#ô ÐIÓJÐJó    c                 óÜ  — d\  }}d\  }}| j                   dv rí| j                  r| j                  n| j                  }t        j                  «       }t	        j
                  || j                  rt        j                  nt        j                  | j                  | j                  d| j                  ¬«      j                  | j                  «      }t        j                  «       }| j                   dk(  r|t        j                  |«      }nf| j                   dv r@t        j                   «       }| j"                  |_        | j&                  r'd|_        d|_        nt-        d| j                   › �«      ‚| j                   d	k(  �rt/        | j0                  «      t2        u r| j0                  d
   n| j0                  }t/        | j0                  «      t2        u r| j0                  d   nd }d }d }	t5        j6                  | j8                  «      D ])  }
d|
vsd|
v sd|
v rŒd|
v s|
dk(  r|
}d|
v r|
}	d|
v sŒ&|
}|
}	Œ+ t        j                  «       }t;        j
                  | j8                  ||	| j                  | j                  d|dk(  rdnd |||¬«
      }t        j                  «       }| j                   dv r©t<        j?                  d| j@                  jC                  | jD                  «      › �«       t        j                  «       }t        jF                  | j@                  jC                  | jD                  «      || j0                  g¬«      }t        j                  «       }t<        j?                  d||z
  › d�«       |S )Nr'   r   T)Útorch_dtypeÚuse_auth_tokenÚtrust_remote_codeÚ	use_cacheÚ	cache_dirr   >   r   r(   r3   r2   úCannot recognize r   r   z.onnxz
.onnx_dataz
.onnx.dataÚdecoder_modelz
model.onnxÚdecoder_with_past_modelÚdecoder_merged_model)	Údecoder_file_nameÚdecoder_with_past_file_namerB   rC   Úuse_io_bindingÚ
use_mergedÚproviderÚprovider_optionsÚsession_options¾   r(   r3   zLoading model from )Ú	providerszLoaded model in ú s)$r   Úhf_pt_dir_pathÚ
model_nameÚtimer   Úfrom_pretrainedr*   ÚtorchÚfloat16Úfloat32ÚauthrE   Útor:   Úcompiler4   ÚSessionOptionsr   Úenable_profilingÚverboseÚlog_verbosity_levelÚlog_severity_levelr   ÚtypeÚexecution_providerÚtupleÚosÚlistdirÚhf_ort_dir_pathr   ÚloggerÚinfoÚort_model_pathÚformatÚrankÚInferenceSession)r!   r"   Úsess_optionsÚ
start_timeÚend_timeÚsourcerN   rO   rJ   rK   Úfilenames              r#   Ú	get_modelrt   ·   s  € Ø$Ñ€Eˆ<Ø%Ñ€J�ð ×ÑÐ>Ñ>Ø(,×(;Ò(;�×$Ò$ÀÇÁˆÜ—Y‘Y“[ˆ
Ü$×4Ñ4ØØ)-¯ªœŸš¼E¿M¹MØŸ9™9Ø"Ÿi™iØØ—n‘nô
÷ ‰"ˆT×ÑÓ
 ð 	ô —9‘9“;ˆà×Ñ /Ò1Ü—M‘M %Ó(‰Eà	×	Ñ	Ð MÑ	MÜ×)Ñ)Ó+ˆØ(,¯©ˆÔ%Ø�<Š<Ø/0ˆLÔ,Ø./ˆLÕ+ô Ð+¨D×,?Ñ,?Ð+@ÐAÓBÐBà×Ñ˜hÓ&ä15°d×6MÑ6MÓ1NÔRWÑ1W�4×*Ñ*¨1Ò-Ð]a×]tÑ]tˆÜ9=¸d×>UÑ>UÓ9VÔZ_Ñ9_˜4×2Ñ2°1Ò5ÐeiÐà ÐØ&*Ð#ÜŸ
™
 4×#7Ñ#7Ó8ò 		7ˆHØ˜hÑ&¨,¸(Ñ*BÀlÐV^ÑF^ØØ (Ñ*¨h¸,Ò.FØ$,Ð!Ø(¨HÑ4Ø.6Ð+Ø%¨Ò1Ø$,Ð!Ø.6Ñ+ð		7ô —Y‘Y“[ˆ
Ü#×3Ñ3Ø× Ñ Ø/Ø(CØŸ9™9Ø"Ÿi™iØØ 1°\Ò A™ÀtØØ-Ø(ô
ˆô —9‘9“;ˆà×ÑÐAÑAä�‰Ð)¨$×*=Ñ*=×*DÑ*DÀTÇYÁYÓ*OÐ)PÐQÔRÜ—Y‘Y“[ˆ
Ü×$Ñ$Ø×Ñ×&Ñ& t§y¡yÓ1ØØ×.Ñ.Ð/ô
ˆô
 —9‘9“;ˆä
‡K�KÐ" 8¨jÑ#8Ð"9¸Ð<Ô=Ø€Lr?   c                 óô  ‡ — ‰ j                   dv rt        ‰ j                  «      n%t        ‰ j                  t        j
                  d¬«      }‰ j                  r ||«      }t        j                  |«       ˆ fd„}ˆ fd„}|D ]  } |«         ||«        |«        Œ d}‰ j                   dv rt        ‰ j                  «      n%t        ‰ j                  t        j
                  d¬«      }	|	D ]H  } |«        t        j                  «       }
 ||«        |«        t        j                  «       }|||
z
  z  }ŒJ ‰ j                   dvrt        j                  d«       |‰ j                  z  }‰ j                  |z  }‰ j                  dk(  rvt        j                  d	‰ j                  › �«       t        j                  d
‰ j                  › �«       t        j                  d|› d�«       t        j                  d|› d�«       y )NrQ   zWarm up)ÚfileÚdescc                  ó|   •— ‰j                   dk7  r(‰j                  dv r‰j                  j                  «       S ˆfd„S )NÚcpurQ   c                  ó    •— ‰j                   dk7  r<t        j                  j                  «       rt        j                  j	                  «       S d„ S )Nry   c                   ó   — y ©N© ©Úkwargss    r#   ú<lambda>z=time_fn.<locals>.<lambda>.<locals>.<lambda>.<locals>.<lambda>  ó   � r?   ©ÚdevicerX   ÚcudaÚis_availableÚsynchronize©r   r!   s    €r#   r€   z+time_fn.<locals>.<lambda>.<locals>.<lambda>  ó?   ø€ à�{‰{˜eÒ#¬¯
©
×(?Ñ(?Ô(Aô �J‰J×"Ñ"Ó$ð ñ &ð r?   )rƒ   r   Ú
io_bindingÚsynchronize_inputsr‡   s    €r#   r€   ztime_fn.<locals>.<lambda>  s>   ø€ à�;‰;˜%Ò D×$7Ñ$7Ð;^Ñ$^ð 	�‰×*Ñ*Ó,ð ó
ð r?   c                  ó|   •— ‰j                   dk7  r(‰j                  dv r‰j                  j                  «       S ˆfd„S )Nry   rQ   c                  ó    •— ‰j                   dk7  r<t        j                  j                  «       rt        j                  j	                  «       S d„ S )Nry   c                   ó   — y r|   r}   r~   s    r#   r€   z=time_fn.<locals>.<lambda>.<locals>.<lambda>.<locals>.<lambda>(  r�   r?   r‚   r‡   s    €r#   r€   z+time_fn.<locals>.<lambda>.<locals>.<lambda>%  rˆ   r?   )rƒ   r   r‰   Úsynchronize_outputsr‡   s    €r#   r€   ztime_fn.<locals>.<lambda>"  s>   ø€ à�;‰;˜%Ò D×$7Ñ$7Ð;^Ñ$^ð 	�‰×+Ñ+Ó-ð ó
ð r?   r   Ú	BenchmarkÚ zBatch Size: zSequence Length: z	Latency: rS   zThroughput: z tps)r   ÚrangeÚwarmup_runsr   ÚsysÚstdoutr`   ri   rj   Únum_runsrV   r;   rm   r<   )r!   ÚfnÚinputsÚwarmup_rangeÚoutputsÚ
input_syncÚoutput_syncÚ_Ú
total_timeÚbench_rangerp   rq   ÚlatencyÚ
throughputs   `             r#   Útime_fnr¡     s¶  ø€ ð ×ÑÐ"EÑEô 	ˆd×ÑÔä�D×$Ñ$¬3¯:©:¸IÔFð ð ‡|‚|Ù�V“*ˆÜ�‰�GÔó€Jó€Kð ò ˆÙŒÙ
ˆ6Œ
Ù�ðð €Jð ×ÑÐ"EÑEô 	ˆd�m‰mÔä�D—M‘M¬¯
©
¸ÔEð ð
 ò 	,ˆÙŒÜ—Y‘Y“[ˆ
á
ˆ6Œ
áŒÜ—9‘9“;ˆà�h Ñ+Ñ+‰
ð	,ð ×ÑÐ"EÑEÜ�‰�BŒà˜4Ÿ=™=Ñ(€GØ—‘ 7Ñ*€Jà‡y�y�A‚~Ü�‰�l 4§?¡?Ð"3Ð4Ô5Ü�‰Ð'¨×(<Ñ(<Ð'=Ð>Ô?Ü�‰�i ˜y¨Ð+Ô,Ü�‰�l : ,¨dÐ3Ô4Ø
r?   c                 óf  — d| j                   › d| j                  › d| j                  j                  «       › d| j                  › d| j
                  › d|j                  j                  dd«      › d|› dt        j                  j                  «       d›�}d }| j                  dv rÝt        t        j                  t        j                  gdd¬«      5 }t        d	«      5   ||«       d d d «       d d d «       j                  d
¬«      j!                  | j"                  | j$                  ¬«      }t&        j(                  j+                  | j,                  |› d�«      }t/        |d«      5 }|j1                  |«       d d d «       |S  ||«       |› d�}|S # 1 sw Y   Œ¬xY w# 1 sw Y   Œ°xY w# 1 sw Y   |S xY w)NÚbÚ_srœ   ú-z%Y-%m-%d_%H:%M:%Sr   T)Ú
activitiesÚrecord_shapesÚprofile_memoryÚmodel_inferencer6   )Úgroup_by_stack_n)Úsort_byÚ	row_limitz.logÚwz.json)r;   r<   r   ÚlowerÚ	precisionrƒ   Ú__name__ÚreplaceÚdatetimeÚnowr   r   ÚCPUÚCUDAr   Úkey_averagesÚtableÚpt_filter_byÚpt_num_rowsrf   ÚpathÚjoinÚ
log_folderÚopenÚwrite)	r!   r–   r—   Úinputs_typeÚprefixrs   ÚprofÚ	prof_dataÚfs	            r#   Ú
profile_fnrÄ   R  s  € ð �—‘Ð!  D×$8Ñ$8Ð#9¸¸4×;NÑ;N×;TÑ;TÓ;VÐ:WÐWXÐY]×YgÑYgÐXhÐhiÐjn×juÑjuÐivÐvwÐxz÷  yDñ  yD÷  yLñ  yLð  MPð  RUó  yVð  xWð  WXð  Ydð  Xeð  efô  go÷  gxñ  gx÷  g|ñ  g|ó  g~ð  Pð  fQð  R€FØ€Hà×ÑÐ>Ñ>äÜ(×,Ñ,Ô.>×.CÑ.CÐDÐTXÐimô
ð 	àÜ Ð!2Ó3ñ Ù�6”
÷÷	ð
 ×%Ñ%°qÐ%Ó9×?Ñ?È×HYÑHYÐei×euÑeuÐ?Óvˆ	ä—7‘7—<‘< §¡°F°8¸4°ÓAˆÜ�(˜CÓ ð 	 AØ�G‰G�IÔ÷	ð €Oñ 	ˆ6Œ
ð �X˜UÐ#ˆà€O÷ð ú÷	ð 	ú÷	ð €Oús0   ÃFÃ	FÃ#FÅ#F&ÆF	ÆFÆF#Æ&F0c                 ó  ‡‡— t        j                  «       }t        j                  |«      }|j	                  d¬«        ‰‰«       | j
                  dk(  r@t        j                  d|j	                  d ¬«      t        j                  d¬«      z  › d�«       t        j                  «        t        j                  j                  «        t        | j                  dk7  ˆˆfd	„¬
«       t         j"                  j%                  «        y )Ngš™™™™™¹?)Úintervalr   zCPU usage: F)Úlogicalú%ry   c                  ó   •—  ‰ ‰«      S r|   r}   )r–   r—   s   €€r#   r€   zmeasure_fn.<locals>.<lambda>|  s   ø€ ¹rÀ&»z€ r?   )Úis_gpuÚfunc)rf   ÚgetpidÚpsutilÚProcessÚcpu_percentrm   ri   rj   Ú	cpu_countÚgcÚcollectrX   r„   Úempty_cacher   rƒ   r“   r”   Úflush)r!   r–   r—   ÚpidÚprocesss    ``  r#   Ú
measure_fnr×   o  s¶   ù€ ä
�)‰)‹+€CÜ�n‰n˜SÓ!€GØ×Ñ ÐÔ%á€v„JØ‡y�y�A‚~Ü�‰�k '×"5Ñ"5¸tÐ"5Ó"DÄv×GWÑGWÐ`eÔGfÑ"fÐ!gÐghÐiÔjô ‡J�J„LÜ	‡J�J×ÑÔÜ˜4Ÿ;™;¨%Ñ/Ô7IÕJô ‡J�J×ÑÕr?   c                 ó’  ‡— ˆfd„}|}| j                   dk(  r ||«        ||«       | j                  �r5t        | ||d«      }| j                   dk(  r}‰j                  j                  j                  «       }t        j                  d|› d|› �«       t        j                  |t        j                  j                  | j                  |«      «       t        | ||d«      }| j                   dk(  r}‰j                  j                  j                  «       }t        j                  d|› d|› �«       t        j                  |t        j                  j                  | j                  |«      «       y t        j                  d«       t        | ||«       t!        | ||«       t        j                  d	«       t        | ||«       t!        | ||«       y )
Nc                 ó   •—  ‰di | ¤Ž}|S )Nr}   r}   ©r—   r™   r"   s     €r#   Ú
get_logitsz$run_hf_inference.<locals>.get_logits„  s   ø€ á‘/˜&‘/ˆØˆr?   r   Úpromptr   ú	Renaming ú to Útokenú7
Evaluating `model(inputs)` step to get past_key_valuesú5
Evaluating `model(inputs)` step with past_key_values)r   r   rÄ   r   ÚsessionÚend_profilingri   Úwarningrf   Úrenamerº   r»   r¼   Údecoder_with_pastrj   r¡   r×   )r!   r=   r>   r"   rÛ   Úgenerate_fnÚnew_lognameÚold_lognames      `    r#   Úrun_hf_inferencerê   ‚  se  ø€ ôð4 €Kà×Ñ˜oÒ-á�KÔ Ù�KÔ à‡|ƒ|Ü   {°KÀÓJˆØ×Ñ (Ò*àŸ-™-×/Ñ/×=Ñ=Ó?ˆKÜ�N‰N˜Y { m°4¸°}ÐEÔFÜ�I‰I�k¤2§7¡7§<¡<°·±ÀÓ#MÔNä   {°KÀÓIˆØ×Ñ (Ò*à×1Ñ1×9Ñ9×GÑGÓIˆKÜ�N‰N˜Y { m°4¸°}ÐEÔFÜ�I‰I�k¤2§7¡7§<¡<°·±ÀÓ#MÔNàô ‡K�KÐJÔKÜˆD�+˜{Ô+Üˆt�[ +Ô.ä
‡K�KÐHÔIÜˆD�+˜{Ô+Üˆt�[ +Õ.r?   c                 ó|  ‡ ‡— ˆ ˆfd„}ˆfd„}ˆfd„}‰ j                   dk7  r|n|}i }‰ j                  �r |||«      \  }	}t        ‰ ||	d«      }
‰j                  «       }t        j                  d|› d|
› �«       t        j                  |t        j                  j                  ‰ j                  |
«      «       t        ‰ «      Š |||«      \  }}t        ‰ ||d«      }
‰j                  «       }t        j                  d|› d|
› �«       t        j                  |t        j                  j                  ‰ j                  |
«      «       y t        j                  d	«        |||«      \  }	}t        ‰ ||	«       t        ‰ ||	«       t        j                  d
«        |||«      \  }}t        ‰ ||«       t        ‰ ||«       y )Nc                 óØ   •— t        ‰| «      } ‰j                  dk7  rKt        ‰| ‰j                  t        ‰j                  «      ‰j
                  |«      \  }}t        ‰d|«       ||fS | |fS )Nry   r‰   )r   rƒ   r   Úintrm   r0   Úsetattr)r—   Úkv_cache_ortvaluesr‰   r!   r"   s      €€r#   Úprepare_ort_inputsz-run_ort_inference.<locals>.prepare_ort_inputsÁ  ss   ø€ ä" 5¨&Ó1ˆð �;‰;˜%ÒÜ-IØ�v˜tŸ{™{¬C°·	±	«N¸D×<QÑ<QÐSeó.Ñ*ˆJÐ*ô �D˜,¨
Ô3ØÐ1Ð1Ð1àÐ)Ð)Ð)r?   c                 ó(   •— ‰j                  | «       y r|   )Úrun_with_iobinding)r‰   r"   s    €r#   Úwith_io_bindingz*run_ort_inference.<locals>.with_io_bindingÏ  s   ø€ à× Ñ  Õ,r?   c                 ó,   •— ‰j                  d | «      }|S r|   )ÚrunrÚ   s     €r#   Úwithout_io_bindingz-run_ort_inference.<locals>.without_io_bindingÓ  s   ø€ à—)‘)˜D &Ó)ˆØˆr?   ry   rÜ   rÝ   rÞ   rß   rà   rá   )rƒ   r   rÄ   rã   ri   rä   rf   rå   rº   r»   r¼   rt   rj   r¡   r×   )r!   r=   r>   r"   rð   ró   rö   rç   rï   Úort_init_inputsrè   ré   Úort_iter_inputss   `  `         r#   Úrun_ort_inferencerù   À  sŽ  ù€ õ*ô-ôð
 &*§[¡[°EÒ%9‘/Ð?Q€KØÐà‡|ƒ|Ù.@ÀÐN`Ó.aÑ+ˆÐ+Ü   {°OÀXÓNˆð ×)Ñ)Ó+ˆÜ�‰˜ ; -¨t°K°=ÐAÔBÜ
�	‰	�+œrŸw™wŸ|™|¨D¯O©O¸[ÓIÔJô ˜$“ˆÙ.@ÀÐN`Ó.aÑ+ˆÐ+Ü   {°OÀWÓMˆð ×)Ñ)Ó+ˆÜ�‰˜ ; -¨t°K°=ÐAÔBÜ
�	‰	�+œrŸw™wŸ|™|¨D¯O©O¸[ÓIÔJØô ‡K�KÐJÔKÙ*<¸[ÐJ\Ó*]Ñ'€OÐ'ÜˆD�+˜Ô/Üˆt�[ /Ô2ä
‡K�KÐHÔIÙ*<¸[ÐJ\Ó*]Ñ'€OÐ'ÜˆD�+˜Ô/Üˆt�[ /Õ2r?   c                 ó¦   — | j                   dv rt        | |||«       y | j                   dv rt        | |||«       y t        d| j                   › �«      ‚)N>   r   r   r   rQ   rF   )r   rê   rù   r   )r!   r=   r>   r"   s       r#   Úrun_inferencerû   û  sV   € Ø×ÑÐHÑHÜ˜˜{¨K¸Õ?Ø	×	Ñ	Ð CÑ	CÜ˜$ ¨[¸%Õ@äÐ+¨D×,?Ñ,?Ð+@ÐAÓBÐBr?   c           	      ó4  — t        j                  «       }|j                  ddt        dg d¢¬«       |j                  ddt        dd¬	«       |j                  d
dddd¬«       |j                  dddt        dg d¢d¬«       |j                  dt        dd¬«       |j                  dt        dd¬«       |j                  dt        dd¬«       |j                  ddd ¬!«       |j                  d"d#d$¬!«       |j                  d%d&t        t        j
                  j                  «       rd'nd(d(d'g¬)«       |j                  d*d+t        d,¬-«       |j                  d.d/t        d0¬-«       |j                  d1d2t        d3¬-«       |j                  d4t        d5¬-«       |j                  d6t        d7¬-«       |j                  d8t        d9¬-«       |j                  d:dd¬;«       |j                  d<t        d=d>¬«       |j                  d?t        d@dA¬«       |j                  dBdd¬;«       |j                  dCt        t        j                  j                  dD«      dE¬«       |j                  dFt        ddGdH¬I«       |j                  «       }t        j                  j                  |j                  «       t	        j                  |j                  «       dJ|j                   v rLt#        |dK|j$                  j'                  «       › dL�«       |j(                  dMk(  r|j(                  dN| if|_        |j                   dOk(  r|j*                  sJ dP«       ‚|j                   dQv r|j,                  sJ dR«       ‚|j.                  j1                  dS«      |_        |j2                  j1                  dS«      |_        |j4                  dTv s|j4                  dUk(  r|j$                  d(k(  rdndV|_        |j6                  r7t9        |j.                  «      d9k(  rt9        |j2                  «      d9k(  sJ dW«       ‚|S )XNz-btz--benchmark-typeT)r   r   r   r(   r3   )rc   ÚrequiredÚchoicesz-mz--model-namez<Hugging Face name of model (e.g. 'meta-llama/Llama-2-7b-hf'))rc   rý   Úhelpz-az--authFÚ
store_truez5Use Hugging Face authentication token to access model)ÚdefaultÚactionrÿ   z-pz--precisionÚfp32)Úint4Úint8Úfp16r  zePrecision for model. For ONNX models, the model's precision should be set before running this script.)rý   rc   r  rþ   rÿ   z--hf-pt-dir-pathr�   zNPath to directory containing all PyTorch files (e.g. tokenizer, PyTorch model))rc   r  rÿ   z--hf-ort-dir-pathzhPath to directory containing all ONNX files (e.g. tokenizer, decoder_merged, decoder, decoder_with_past)z--ort-model-pathzPath to ONNX modelz-bz--batch-sizesz1 2)r  z-sz--sequence-lengthsz32 64 128 256 512z-dz--devicer„   ry   )rc   r  rþ   z-idz--device-idr   )rc   r  z-wz--warmup-runsr6   z-nz
--num-runsé
   z--seedé   z--max-lengthé    z--num-return-sequencesr2   z	--profile)r  r  z--pt-filter-byÚself_cpu_time_totalz"What to filter PyTorch profiler byz--pt-num-rowsiè  z.Number of rows for PyTorch profiler to displayz	--verbosez--log-folderú.zFolder to cache log filesz--cache-dirz./model_cachez-Cache dir where Hugging Face files are stored)rc   rý   r  rÿ   r4   rd   ÚExecutionProviderÚCUDAExecutionProviderÚ	device_idr   z,Please specify a path to `--hf-ort-dir-path`rQ   z+Please specify a path to `--ort-model-path`ú >   r  r  r  r  zOPlease provide only one (batch_size, sequence_length) combination for profiling)ÚargparseÚArgumentParserÚadd_argumentÚstrrX   r„   r…   rí   rf   rº   r»   Ú
parse_argsÚnpÚrandomÚseedÚmanual_seedr   rî   rƒ   Úupperrd   rh   rk   Úbatch_sizesÚsplitÚsequence_lengthsr¯   r   r   )rm   Úparserr!   s      r#   Úget_argsr    s6  € Ü×$Ñ$Ó&€FØ
×ÑØØÜØò
ð ô ð ×ÑØØÜØØKð ô ð ×ÑØˆh ¨lÐAxð ô ð
 ×ÑØØØÜØÚ0Øtð ô ð ×ÑØÜØØ]ð	 ô ð ×ÑØÜØØwð	 ô ð ×ÑØÜØØ!ð	 ô ð ×ÑØØØð ô ð
 ×ÑØØØ#ð ô ð
 ×ÑØØÜÜŸ*™*×1Ñ1Ô3‘¸Ø˜�ð ô ð ×Ñ˜˜}´3ÀÐÔBØ
×Ñ˜˜o´CÀÐÔCØ
×Ñ˜˜l´¸bÐÔAØ
×Ñ˜¤s°AÐÔ6ð ×Ñ˜¬S¸"ÐÔ=Ø
×ÑÐ0´sÀAÐÔFð ×Ñ˜¨U¸<ÐÔHØ
×ÑØœsÐ,AÐHlð ô ð ×Ñ˜¬c¸4ÐFvÐÔwØ
×Ñ˜¨U¸<ÐÔHØ
×Ñ˜¬S¼"¿'¹'¿,¹,ÀsÓ:KÐRmÐÔnØ
×ÑØÜØØØ<ð ô ð ×ÑÓ€Dô ‡I�I‡N�N�4—9‘9ÔÜ	×Ñ�d—i‘iÔ ð �×#Ñ#Ñ#Ü�Ð*¨t¯{©{×/@Ñ/@Ó/BÐ.CÐCTÐ,UÔVØ×"Ñ"Ð&=Ò=Ø'+×'>Ñ'>ÀÈdÐ@SÐ&TˆDÔ#ð ×Ñ˜hÒ&Ø×#Ò#ÐSÐ%SÓSÐ#Ø×ÑÐAÑAØ×"Ò"ÐQÐ$QÓQÐ"à×'Ñ'×-Ñ-¨cÓ2€DÔØ ×1Ñ1×7Ñ7¸Ó<€DÔð —.‘.Ð$4Ñ4¸¿¹È6Ò9QÐVZ×VaÑVaÐejÒVj‰Ðqwð 	„Nð
 ‡|‚|Ü�4×#Ñ#Ó$¨Ò)¬c°$×2GÑ2GÓ.HÈAÒ.Mð 	
Ø]ó	
ÐMð €Kr?   c                  ó¼  — t        «       } t        «       }t        | «      }t        |j                  «       t
        j                  |j                  «       dt        j                  j                  _        | |_        ||_        t        j                  |j                   |j"                  |j$                  |j$                  ¬«      }t'        j                  |j                   |j"                  |j$                  |j$                  ¬«      }|j(                  dk7  rd|j                  › �n|j(                  }|j*                  dk(  }t-        |d|«       t-        |d|«       t-        |d|«       t-        |d	|«       t/        |«      }t1        ||«      }|j2                  d
v r•t5        j6                  |j8                  j;                  |j                  «      d¬«      }	t=        t?        d„ |	j@                  jB                  «      «      }
|xr tE        |
«      dkD  xr |j(                  dk7  }t-        |d|«       nt-        |dd«       tG        jH                  |jJ                  |jL                  «      D ]y  \  }}|j                  dk(  rt
        j                  d|› d|› d�«       t-        |dtO        |«      «       t-        |dtO        |«      «       tQ        ||«      \  }}tS        ||||«       Œ{ y )NT)rE   rB   rC   ry   zcuda:r  Ú	tokenizerr8   r:   r*   rQ   F)Úload_external_datac                 ó    — | j                   dk(  S )NÚGroupQueryAttention)Úop_type)Únodes    r#   r€   zmain.<locals>.<lambda>©  s   € ¨T¯\©\Ð=RÑ-R€ r?   r   r0   z
Batch size = z and sequence length = z...r;   r<   )*r   r   r  r   r`   ri   rj   Ú__dict__rX   ÚbackendsÚcudnnÚ	benchmarkrm   r5   r   rW   rU   rE   r[   r   rƒ   r¯   rî   rt   r$   r   ÚonnxÚ
load_modelrk   rl   ÚlistÚfilterÚgraphr%  r   Ú	itertoolsÚproductr  r  rí   r    rû   )rm   r5   r!   r   r8   r:   r*   r"   r%   Ú
onnx_modelÚ	gqa_nodesr0   r;   r<   r=   r>   s                   r#   Úmainr3  ‰  sU  € Ü‹:€DÜ“€Jä�D‹>€DÜ�—‘ÔÜ
‡K�K�—‘ÔØ%)„E‡N�N×ÑÔ"à€D„IØ €D„OÜ×-Ñ-Ø�‰ 4§>¡>À$Ç)Á)Ð_c×_hÑ_hô€Iô ×'Ñ'Ø�‰ 4§>¡>À$Ç)Á)Ð_c×_hÑ_hô€Fð ,0¯;©;¸%Ò+?�e˜DŸI™I˜;Ñ'ÀTÇ[Á[€MØ�~‰~ Ñ'€HäˆD�+˜yÔ)ÜˆD�(˜FÔ#ÜˆD�/ =Ô1ÜˆD�*˜hÔ'ô �d‹O€EÜ3°D¸%Ó@Ðð ×ÑÐAÑAÜ—_‘_ T×%8Ñ%8×%?Ñ%?ÀÇ	Á	Ó%JÐ_dÔeˆ
ÜœÑ RÐT^×TdÑTd×TiÑTiÓjÓkˆ	à#ÒS¬¨I«¸Ñ(:ÒS¸t¿{¹{ÈeÑ?SÐÜ�Ð(Ð*:Õ;ä�Ð(¨%Ô0ô (1×'8Ñ'8¸×9IÑ9IÈ4×K`ÑK`Ó'aò =Ñ#ˆ
�OØ�9‰9˜Š>Ü�K‰K˜/¨*¨Ð5LÈ_ÐL]Ð]`ÐaÔbÜ��l¤C¨
£OÔ4Ü�Ð'¬¨_Ó)=Ô>ä#-¨dÐ4HÓ#IÑ ˆ�[Ü�d˜K¨°eÕ<ñ=r?   Ú__main__)r   )8r  r²   rÑ   r/  Úloggingrf   r“   rV   Únumpyr  r*  rÍ   rX   Úbenchmark_helperr   r   Údist_settingsr   r   Úllama_inputsr   r   r	   r
   r   r   Úoptimum.onnxruntimer   Útorch.profilerr   r   r   Útqdmr   Útransformersr   r   r   Úonnxruntimer4   Ú	getLoggerr°   ri   r$   Ú	Namespacerí   r    rt   r¡   rÄ   r×   rê   rù   rû   r  r3  r}   r?   r#   ú<module>rA     sß   ðó Û Û 	Û Û Û 	Û 
Û ã Û Û Û ß 9ß ,÷÷ õ 4ß EÑ EÝ ß HÑ Hã à	ˆ×	Ñ	˜8Ó	$€ò
#ð$�X×'Ñ'ð $¸só $ðDR�H×&Ñ&ó RòjCòLò:ò&;/ò|83òvCóBòJ/=ðd ˆzÒÙ…Fð r?   