+
    UV-jJ\  ã                   óv  € ^ RI HtHtHtHtHtHt ^ RIHt	 ^ RI
Ht ^RIHtHtHtHt R R ltR R ltR R	 ltR
 R ltR R ltR R ltR R ltR R ltR R ltR R ltR+R R lltR R ltR+R R lltR R ltR RR!RR"]	P>                  R#R$/R% R& llt R RR!RR"]	P>                  R'RR(RR#R$/R) R* llt!R# ),é    )ÚAnyÚCallableÚ	GeneratorÚListÚOptionalÚTupleN)Ú_batch_cache_left_paddingÚ_record_speculative_roundÚ_SpeculativeSamplerRNGÚgeneration_streamc                ó¨   € V ^8„  d   QhR\         P                  R\         P                  R\        R\        \        \        \        ,          3,          /# )é   Údraft_tokensÚtarget_tokensÚbudgetÚreturn)ÚmxÚarrayÚintr   r   )Úformats   "Úk/Volumes/fast/ai/experiments/ui-tars-smoke/.venv/lib/python3.14/site-packages/mlx_vlm/speculative/eagle3.pyÚ__annotate__r      sE   € ÷ )ñ )Ü—(‘(ð)ä—8‘8ð)ô ð)ô Œ3””S•	ˆ>Õñ	)ó    c                 óJ  € \        V P                  ^,          4      pV P                  R4      P                  4       RV pVP                  R4      P                  4       pTp\	        \        WE4      4       F  w  pw  r‰W‰8w  g   K  Tp M	  VRV WVV^,            ,           p
WjRV 3# )é   Néÿÿÿÿ)r   ÚshapeÚreshapeÚtolistÚ	enumerateÚzip)r   r   r   Ún_draftÚ	draft_rowÚ
target_rowÚacceptedÚiÚ	draft_tokÚ
target_tokÚ
new_tokenss   &&&        r   Ú_eagle3_walkr*      s®   € ô
 �,×$Ñ$ QÕ'Ó(€GØ×$Ñ$ RÓ(×/Ñ/Ó1°(°7Ð;€IØ×&Ñ& rÓ*×1Ñ1Ó3€Jà€HÜ&/´°IÓ0JÖ&KÑ"ˆÑ"ˆIØÖ"ØˆHÙñ 'Lð
 ˜9˜HÐ%¨
¸hÈ½lÐ(KÕK€JØ  Ð(Ð(Ð(r   c                óê   € V ^8„  d   QhR\         P                  R\         P                  R\        \        ,          R\        \        \        ,          \        \        \        ,          ,          3,          /# )r   r   r   Úbudgetsr   ©r   r   r   r   r   )r   s   "r   r   r   !   sR   € ÷ *ñ *Ü—(‘(ð*ä—8‘8ð*ô ”#�Yð*ô Œ4”�9”dœ4¤�9•oÐ%Õ&ñ	*r   c                 óô  € \        V P                  ^ ,          4      p\        V P                  ^,          4      pV P                  4       pVP                  4       p. p. p\        V4       F‹  p	Tp
\	        \        WY,          RV Wi,          4      4       F  w  pw  rÍWÍ8w  g   K  Tp
 M	  VP                  V
4       WY,          RV
 Wi,          Wª^,            ,           pVP                  VRW),           4       K�  	  Wx3# )r   N)r   r   r   Úranger    r!   Úappend)r   r   r,   ÚBr"   Ú
draft_rowsÚtarget_rowsÚaccepted_listÚnew_tokens_listr&   r%   Újr'   r(   r)   s   &&&            r   Ú_eagle3_walk_batchr7   !   sô   € ô
 	ˆL×Ñ˜qÕ!Ó"€AÜ�,×$Ñ$ QÕ'Ó(€GØ×$Ñ$Ó&€JØ×&Ñ&Ó(€KØ€MØ')€OÜ�1ŽXˆØˆÜ*3Ü�
•˜h˜wÐ'¨­Ó8ö+
Ñ&ˆAÑ&�	ð Ö&Ø�Ùñ+
ð 	×Ñ˜XÔ&Ø•] 9 HÐ-°µ¸xÐUVÍ,Ð0WÕWˆ
Ø×Ñ˜z¨,¨G­JÐ7Ö8ñ ð Ð)Ð)r   c                óx   € V ^8„  d   QhR\         P                  R\         P                  R\         P                  /# )r   r   r   r   )r   r   )r   s   "r   r   r   :   s-   € ÷ 	ñ 	¬¯©ð 	ÄÇÁð 	ÌbÏhÉhñ 	r   c           
      óÜ  € \        V P                  ^,          4      pWRRV13,          8g  p\        P                  ! VP	                  \        P
                  4      ^R7      p\        P                  ! V^R7      p\        P                  ! VV\        P                  ! \        V P                  ^ ,          4      3V\        P
                  R7      4      P	                  \        P
                  4      # )r   ºNNNN©Úaxis©Údtype)	r   r   r   ÚargmaxÚastypeÚint32ÚanyÚwhereÚfull)r   r   r"   Ú
mismatchesÚfirst_mismatchÚhas_mismatchs   &&    r   Ú_eagle3_accept_countsrH   :   s©   € Ü�,×$Ñ$ QÕ'Ó(€GØ¨q°(°7°(¨{Õ!;Ñ;€JÜ—Y’Y˜z×0Ñ0´·±Ó:ÀÔC€NÜ—6’6˜*¨1Ô-€LÜ�8Š8ØØÜ
�Š”�\×'Ñ'¨Õ*Ó+Ð-¨w¼b¿h¹hÔGó÷ �fŒR�X‰XÓð	r   c                ó  € V ^8„  d   QhR\         P                  R\         P                  R\        \        ,          R\        \        ,          R\        \        \        ,          \        \        \        ,          ,          3,          /# )r   r   r   r4   r,   r   r-   )r   s   "r   r   r   F   s`   € ÷ <ñ <Ü—(‘(ð<ä—8‘8ð<ô œ•9ð<ô ”#�Yð	<ô
 Œ4”�9”dœ4¤�9•oÐ%Õ&ñ<r   c                 ór  € \        V4      p. p\        V4       F‡  w  rgWV^,           1RV13,          pWV^,           1WD^,           13,          p	\        P                  ! W‰.^R7      RRV13,          P	                  R4      P                  4       p
VP                  V
4       K‰  	  V.\        V4      ,          V3# )r   Nr;   r:   r   )Úminr    r   Úconcatenater   r   r0   Úlen)r   r   r4   r,   r%   r5   r&   r   Úaccepted_prefixÚbonusÚnews   &&&&       r   Ú%_eagle3_walk_batch_uniform_acceptancerQ   F   sµ   € ô �=Ó!€HØ')€OÜ˜wÖ'‰	ˆØ&¨1¨q­5 y°)°8°)Ð';Õ<ˆØ ! a¥%˜i¨¸qµLÐ)@Ð@ÕAˆä�NŠN˜OÐ3¸!Ô<¸QÀÀÀ¸ZÕHß‰W�R‹[ß‰V‹Xð 	ð
 	×Ñ˜sÖ#ñ (ð ˆ:œ˜MÓ*Õ*¨OÐ;Ð;r   c                óP   € V ^8„  d   QhR\         P                  R\        R\        /# )r   Údraft_modelÚgreedy_samplingr   )ÚnnÚModuleÚboolÚdict)r   s   "r   r   r   Z   s%   € ÷ ñ ¤b§i¡ið Ä$ð Ì4ñ r   c                 óD   € V'       d   \        V R R4      '       d   RR/# / # )Úsupports_greedy_draft_argmaxFÚgreedyT)Úgetattr)rS   rT   s   &&r   Ú_eagle3_draft_kwargsr]   Z   s%   € ßœ7 ;Ð0NÐPU×VÒVØ˜$ÐÐØ€Ir   c          	      ó’   € V ^8„  d   QhR\         P                  R\        \        ,          R\        \        \        \
        3,          /# )r   rS   Údraft_block_sizer   )rU   rV   r   r   r   rW   )r   s   "r   r   r   `   s9   € ÷ 7ñ 7Ü—‘ð7äœs•mð7ô Œ3””Tˆ>Õñ7r   c                 óØ   € \        V P                  P                  4      pVe   \        V4      VR3# \        V P                  RR 4      pVf   W"R3# \	        V\        V4      4      pW2W28„  3# )NFÚadaptive_max_block_size)r   ÚconfigÚ
block_sizer\   Úmax)rS   r_   Ú
configuredÚauto_maxs   &&  r   Ú_eagle3_block_settingsrg   `   sw   € ô �[×'Ñ'×2Ñ2Ó3€JØÒ#ÜÐ#Ó$ j°%Ð7Ð7ä�{×)Ñ)Ð+DÀdÓK€HØÒØ uÐ,Ð,Ü�:œs 8›}Ó-€HØ Ñ!6Ð6Ð6r   c                óR   € V ^8„  d   QhR\         R\         R\        \         ,          /# )r   Úconfigured_block_totalÚmax_block_totalr   )r   r   )r   s   "r   r   r   o   s&   € ÷ <ñ <´ð <Äcð <ÌdÔSVÍiñ <r   c                 óÈ   € V .pR F(  pYu;8  d
   V8  g   K  M K  VP                  V4       K*  	  W9  d   VP                  V4       \        \        R V 4       4      4      # )é   c              3   óL   "  € T F  p\        ^\        V4      4      x € K  	  R# 5i)r   N)rd   r   )Ú.0Útiers   & r   Ú	<genexpr>Ú&_eagle3_block_tiers.<locals>.<genexpr>v   s   é € Ð:±E¨D”c˜!œS ›Y×'Ð'³Eùs   ‚"$)rl   é   é   )r0   ÚsortedÚset)ri   rj   Útiersro   s   &&  r   Ú_eagle3_block_tiersrw   o   sX   € Ø#Ð$€EÛˆØ!Ö:¨?×:Ò:Ø�L‰L˜Öñ ð Ô#Ø�‰�_Ô%Ü”#Ñ:±EÓ:Ó:Ó;Ð;r   c                ót   € V ^8„  d   QhR\         P                  R\        R\        R\        R\        R\        /# )r   rS   rj   ri   Úremaining_budgetÚadaptiver   )rU   rV   r   rW   )r   s   "r   r   r   y   sH   € ÷ %ñ %Ü—‘ð%äð%ô  ð%ô ð	%ô ð%ô 	ñ%r   c          	      óú  a€ V'       g   \        W4      # \        W!4      p\        V R R4      oSf
   V^ ,          o\        V RR4      ;'       g    . p\        V RR4      ;'       g    . p\        V4      ^8¼  Edf   \        V4      ^8¼  EdU   VRR  Uu. uF  p\	        V4      NK  	  p	pVRR  U
u. uF  p
\        ^\	        V
4      4      NK  	  pp
\        R V	 4       4      \        V	4      ,          p\        R \        W›4       4       4      \        V	4      ,          pVP                  \        VV3R lR7      4      p\        V4      ^8X  d+   SV8X  d$   \        V4      ^8”  d   \        V4      ^,
          pMaVR8  g   VR	8  d   V^ 8X  d   \        ^ V^,
          4      pM8VR
8¼  g   VSR,          8¼  d$   \        \        V4      ^,
          V^,           4      pW^,          o\        SW4      oSV n	        S# u upi u up
i )Ú_adaptive_block_sizeNÚaccept_lensÚ
draft_lensc              3   ó2   "  € T F  q^,           x € K  	  R# 5i©r   N© )rn   Úas   & r   rp   Ú*_eagle3_next_block_size.<locals>.<genexpr>Ž   s   é € Ð7© A˜aŸ%š%«ùs   ‚c              3   ó:   "  € T F  w  rW8¼  g   K  ^x € K  	  R# 5ir€   r�   )rn   r‚   Úds   &  r   rp   rƒ   �   s   é € ÐQÑ&F™d˜aÈ!É&ŸšÓ&Fùs   ‚
‘
c                 ó(   <€ \        V S,
          4      # ©N)Úabs)ro   Úcurrents   &€r   Ú<lambda>Ú)_eagle3_next_block_size.<locals>.<lambda>“   s   ø€ ¼3¸tÀg½~Ô;Nr   )Úkeyg       @g      @g…ëQ¸Õ?g      è?iúÿÿÿ)
rK   rw   r\   rM   r   rd   Úsumr!   Úindexr|   )rS   rj   ri   ry   rz   rv   r%   Údraftedr‚   Úrecent_acceptr…   Úrecent_draftÚmean_outputÚ	full_rateÚtier_idxr‰   s   &&&&$          @r   Ú_eagle3_next_block_sizer•   y   s¼  ø€ ÷ Ü�?Ó5Ð5äÐ 6ÓH€EÜ�kÐ#9¸4Ó@€GØ‚Ø˜•(ˆä�{ M°4Ó8×>Ð>¸B€HÜ�k <°Ó6×<Ð<¸"€GÜ
ˆ8ƒ}˜Õœc '›l¨aÕ/Ø)1°"°#©Ó7© Aœ˜Qž©ˆÐ7Ø07¸¸±Ó=±¨1œ˜Aœs 1›vž±ˆÐ=ÜÑ7©Ó7Ó7¼#¸mÓ:LÕLˆÜÑQ¤c¨-Ô&FÓQÓQÔTWØóU
õ 
ˆ	ð —;‘;œs 5Ô.NÔOÓPˆÜˆx‹=˜AÔ 'Ð-CÔ"CÌÈEË
ÐUVÌÜ˜5“z A•~‰HØ˜3Ô ;°Ô#4¸Àa¼Ü˜1˜h¨�lÓ+‰HØ˜$Ô +°¸4µÔ"?Üœ3˜u›:¨�>¨8°a­<Ó8ˆHØ•/ˆä�'˜?Ó=€GØ'.€KÔ$Ø€Nùò% 8ùÚ=s   ÂG3Â.G8c                óô   € V ^8„  d   QhR\         P                  R\        P                  R\        \
        ,          R\        \        P                  .\        P                  3,          R\        \        ,          /# )r   ÚlmÚverify_inputÚprompt_cacheÚsamplerÚtarget_layer_ids)rU   rV   r   r   r   r   r   r   )r   s   "r   r   r   ¡   s^   € ÷ '8ñ '8Ü
�	‰	ð'8ä—(‘(ð'8ô ”s•)ð'8ô ”r—x‘x�j¤"§(¡(Ð*Õ+ð	'8ô
 œ3•iñ'8r   c                 ó  € VP                   ^,          ^8”  Ed)   R\        V 4      P                  9   Ed   V ! VR,          VVR7      p\        P                  ! VP
                  RR7      .pV! VP                  4      .pVP                  pVP                   ^,          ^8”  dq   V ! VR,          VVR7      p	VP                  \        P                  ! V	P
                  RR7      4       VP                  V! V	P                  4      4       V	P                  p\        P                  ! V^R7      \        P                  ! V^R7      V3# V ! VVVR7      p
\        P                  ! V
P
                  RR7      pV! V
P                  4      pW¼V
P                  3# )r   Úgemma4)ÚcacheÚcapture_layer_idsr;   )r:   :Nr   Nr   ©r:   :r   NN)	r   ÚtypeÚ
__module__r   rL   Úhidden_statesÚlogitsÚ
gdn_statesr0   )r—   r˜   r™   rš   r›   Ú	first_outÚhidden_chunksÚtarget_chunksr¥   Útail_outÚ
verify_outÚhiddenr   s   &&&&&        r   Ú_eagle3_verify_targetr¬   ¡   s^  € ð ×Ñ˜!Õ˜qÕ  X´°b³×1DÑ1DÕ%DÙØ˜ÕØØ.ô
ˆ	ô
 Ÿš¨	×(?Ñ(?ÀbÔIÐJˆÙ  ×!1Ñ!1Ó2Ð3ˆØ×)Ñ)ˆ
à×Ñ˜aÕ  1Ô$ÙØ˜UÕ#Ø"Ø"2ôˆHð
 × Ñ ¤§¢°×0FÑ0FÈRÔ!PÔQØ× Ñ ¡¨¯©Ó!9Ô:Ø!×,Ñ,ˆJä�NŠN˜=¨qÔ1Ü�NŠN˜=¨qÔ1Øð
ð 	
ñ ØØØ*ô€Jô
 �^Š^˜J×4Ñ4¸2Ô>€FÙ˜J×-Ñ-Ó.€MØ *×"7Ñ"7Ð7Ð7r   c                óZ   € V ^8„  d   QhR\         P                  R\        \        ,          /# )r   r—   r   ©rU   rV   r   r   )r   s   "r   r   r   Ë   s"   € ÷ ñ œbŸi™ið ¬D´­Iñ r   c                 óæ   € \        V R R4      p\        VRR4      pVf   . # \        V\        \        \        34      '       d!   V Uu. uF  q3f   K  \        V4      NK  	  up# \        V4      .# u upi )rb   NÚeos_token_id)r\   Ú
isinstanceÚlistÚtupleru   r   )r—   rb   ÚeosÚtoks   &   r   Ú_eagle3_eos_token_idsr¶   Ë   sh   € Ü�R˜ 4Ó(€FÜ
�&˜.¨$Ó
/€CØ
‚{Øˆ	Ü�#œœe¤SÐ)×*Ò*Ù$'Ó;¡C˜S””�C–¡CÑ;Ð;Ü�‹Hˆ:Ðùò <s   ÁA.ÁA.c          	      óÆ   € V ^8„  d   QhR\         P                  R\         P                  R\        \        \        ,          ,          R\        \
        P                  ,          /# )r   r—   rS   Úeos_token_idsr   )rU   rV   r   r   r   r   r   )r   s   "r   r   r   Õ   sJ   € ÷ ñ Ü
�	‰	ðä—‘ðô œD¤�IÕ&ðô Œb�h‰hÕñ	r   c                 óH  € \        VR R4      pVf   R# Ve   TM
\        V 4      p\        V4      p\        VRR4      p\        VRR4      pVe	   We8X  d   V# \        P                  ! VP
                  ^ ,          \        P                  R7      VP                  \        P                  4      ,           pV'       d?   \        P                  ! V\        P                  ! V\        P                  R7      .^ R7      pVP                  \        P                  4      pWQn
        W�n        V# )Úd2tNÚ_hot_token_ids_cache_keyÚ_hot_token_ids_cacher=   r;   )r\   r¶   r³   r   Úaranger   rA   r@   rL   r   r»   r¼   )	r—   rS   r¸   rº   Úeos_idsÚ	cache_keyÚ
cached_keyÚ
cached_idsÚhot_idss	   &&&      r   Ú_eagle3_hot_token_idsrÃ   Õ   sÞ   € ô
 �+˜u dÓ
+€CØ
‚{Ùà,Ò8‰mÔ>SÐTVÓ>W€GÜ�g“€IÜ˜Ð&@À$ÓG€JÜ˜Ð&<¸dÓC€JØÒ *Ô"9ØÐä�iŠi˜Ÿ	™	 !�¬B¯H©HÔ5¸¿
¹
Ä2Ç8Á8Ó8LÕL€GßÜ—.’. '¬2¯8ª8°GÄ2Ç8Á8Ô+LÐ!MÐTUÔVˆØ�n‰nœRŸX™XÓ&€GØ+4Ô(Ø'.Ô$Ø€Nr   c          
      ó¸   € V ^8„  d   QhR\         P                  R\         P                  R\        P                  R\        P                  R\        P                  /# )r   r—   rS   r«   rÂ   r   )rU   rV   r   r   )r   s   "r   r   r   î   sN   € ÷ .ñ .Ü
�	‰	ð.ä—‘ð.ô �H‰Hð.ô �X‰Xð	.ô
 ‡X�Xñ.r   c                 ó  € V P                   P                  p\        V4      \        VP                  ^ ,          4      3p\        VRR4      p\        VRR4      pVe   We8w  d©   \        VR4      '       dq   VP                  V,          VP                  V,          \        VR4      '       d   VP                  V,          MRVP                  VP                  \        VRR4      R3pMVP                  V,          RRRRRR	3pWQn        Wqn        Vw  r‰r«rÍpV'       d    \        P                  ! VVV	V
RVVVR
7      pMW(P                   ,          p\        V RR4      pVe%   \        P"                  ! VV,          4      V,          pV# )r   Ú_hot_lm_head_cache_keyNÚ_hot_lm_head_cacheÚscalesÚbiasesÚmodeÚaffineTF)rÈ   rÉ   Ú	transposeÚ
group_sizeÚbitsrÊ   Úfinal_logit_softcapping)ÚmodelÚembed_tokensÚidr   r   r\   ÚhasattrÚweightrÈ   rÉ   rÍ   rÎ   rÆ   rÇ   r   Úquantized_matmulÚTÚtanh)r—   rS   r«   rÂ   Úembedr¿   rÀ   Úcached_headrÔ   rÈ   rÉ   rÍ   rÎ   rÊ   Ú	quantizedr¤   Úsoftcaps   &&&&             r   Ú_eagle3_hot_logits_from_hiddenrÜ   î   sb  € ð �H‰H×!Ñ!€EÜ�E“œC §¡¨aÕ 0Ó1Ð2€IÜ˜Ð&>ÀÓE€JÜ˜+Ð';¸TÓB€KàÒ˜jÔ5Ü�5˜(×#Ò#à—‘˜WÕ%Ø—‘˜WÕ%Ü)0°¸×)AÒ)A�—‘˜WÖ%ÀtØ× Ñ Ø—
‘
Ü˜˜v xÓ0Øð‰Kð !Ÿ<™<¨Õ0°$¸¸dÀDÈ$ÐPUÐVˆKØ-6Ô*Ø)4Ô&à@KÑ=€F�F¨°IçÜ×$Ò$ØØØØØØ!ØØô	
‰ð Ÿ(™(Õ"ˆä�bÐ3°TÓ:€GØÒÜ—’˜ 'Õ)Ó*¨WÕ4ˆØ€Mr   c                óÂ  € V ^8„  d   QhR\         P                  R\         P                  R\        P                  R\        \
        ,          R\        \        P                  .\        P                  3,          R\        \        ,          R\        \        \        ,          ,          R\        \        \        P                  \        P                  \
        3,          ,          /# )	r   r—   rS   r˜   r™   rš   r›   r¸   r   )
rU   rV   r   r   r   r   r   r   r   r   )r   s   "r   r   r     s¢   € ÷ ,'ñ ,'Ü
�	‰	ð,'ä—‘ð,'ô —(‘(ð,'ô ”s•)ð	,'ô
 ”r—x‘x�j¤"§(¡(Ð*Õ+ð,'ô œ3•ið,'ô œD¤�IÕ&ð,'ô Œe”B—H‘HœbŸh™h¬Ð+Õ,Õ-ñ,'r   c                 ó  € \        V R 4      '       d   \        V R4      '       g   R# \        WV4      pVf   R# . pV P                  VVVVR7      p	V'       g   \        R4      h\        P
                  ! VRR7      p
\        WW—4      p\        P                  ! VRR7      P                  \        P                  4      pW|,          pVR,          p\        P                  ! \        Wí4      \        VP                  ^,          4      ^,
          4      p\        P                  ! WŸR,          ^R7      pV! V P                  V4      4      P                  VP                   4      p\        P"                  ! VVR	,          V^R7      pW­R3# )
rÐ   Úlogits_from_hiddenN)rž   rŸ   Úhidden_sinkz3EAGLE-3 verification did not produce hidden states.r;   r   r    )r:   NN©r:   N)rÓ   rÃ   rÐ   ÚRuntimeErrorr   rL   rÜ   r?   r@   rA   ÚminimumrH   r   r   Útake_along_axisrß   r>   Úput_along_axis)r—   rS   r˜   r™   rš   r›   r¸   rÂ   rà   Úfinal_hiddenr«   Ú
hot_logitsÚhot_idxr   r   Úfull_posÚfull_hiddenÚ
full_tokens   &&&&&&&           r   Ú_eagle3_verify_target_hotrì     sS  € ô �2�w×Ò¤w¨rÐ3G×'HÒ'HÙä# B°]ÓC€GØ‚Ùà"$€KØ—8‘8ØØØ*Øð	 ó €L÷ ÜÐPÓQÐQä�^Š^˜K¨bÔ1€FÜ/°ÀÓW€JÜ�iŠi˜
¨Ô,×3Ñ3´B·H±HÓ=€GØÕ$€Mà Õ&€LÜ�zŠzÜ˜lÓ:ÜˆM×Ñ Õ"Ó# aÕ'ó€Hô ×$Ò$ \¸MÕ3JÐQRÔS€KÙ˜×.Ñ.¨{Ó;Ó<×CÑCÀM×DWÑDWÓX€JÜ×%Ò%ØØ�ÕØØô	€Mð  $Ð&Ð&r   c                óZ   € V ^8„  d   QhR\         P                  R\        \        ,          /# )r   rS   r   r®   )r   s   "r   r   r   N  s"   € ÷ ñ ¬2¯9©9ð ¼¼c½ñ r   c                 ój   € \        \        V P                  R V P                  P                  4      4      # )rŸ   )r²   r\   rb   r›   )rS   s   &r   Ú_eagle3_capture_layer_idsrï   N  s2   € ÜÜØ×ÑØØ×Ñ×/Ñ/ó	
óð r   Úprompt_tokensr_   Útoken_dtyperT   Fc                óÐ  € V ^8„  d   QhR\         P                  R\         P                  R\        \        ,          R\        P
                  R\        \        P
                  ,          R\        R\        R\        \        P
                  .\        P
                  3,          R	\        \        ,          R
\        P                  R\        R\        \        \        R3,          RR3,          /# )r   rÐ   rS   r™   r«   rð   Úfirst_bonusÚ
max_tokensrš   r_   rñ   rT   r   N)rU   rV   r   r   r   r   r   r   r   ÚDtyperW   r   r   )r   s   "r   r   r   X  sØ   € ÷ Cñ CÜ�9‰9ðCä—‘ðCô ”s•)ðCô �H‰Hð	Cô œBŸH™HÕ%ðCô ðCô ðCô ”r—x‘x�j¤"§(¡(Ð*Õ+ðCô œs•mðCô —‘ðCô ðCô Œu”S˜$�YÕ  tÐ+Õ,ñCr   c             #   óæ  "  € \        V R 4      '       d   V P                  MT p\        VR4      '       g"   \        \        V4      P                   R24      h\        V4      p\        W4      w  rÞpVf   \        V^4      p\        Wí4      pRpVP                  V 4      p\        W'       * R7      p\        VRR4      p\        V4      '       d&   Ve"   VP                  ! VVVVVV	3/ \        W4      B  VP                  ^,          ^8”  d   VRRR1R3,          pTp^pVV8  Ed&   \        VVVVV,
          ^,           VR7      pV^8:  d   R# VP                   ! VP"                  VVVVVV	3/ \        W4      B p\$        P&                  ! \(        4      ;_uu_ 4        \$        P*                  ! \$        P,                  ! V..V	R	7      V.^R
7      p\/        VVVVVV\1        V 4      4      pVf   \3        VVVVV4      w  pppMVw  pppRRR4       VP5                  XX4       \7        VVVV,
          4      w  pp\9        VVV^,
          4       V F  pVR3x € V^,          pVV8¼  g   K   R# 	  \        VRR4      p\        V4      '       d#   VP                  ! VVVVVVV	3/ \        W4      B  VRVV^,           1R3,          pV'       d
   VR,          MTpVV^,
          8  d@   \$        P&                  ! \(        4      ;_uu_ 4        VP;                  VXVV4       RRR4       VR,          ^ 8X  g   EK  \$        P<                  ! 4        EK-  R#   + '       g   i     ELP; i  + '       g   i     LO; i5i)Úlanguage_modelÚrollback_speculative_cacheú/ does not implement rollback_speculative_cache.NF©ÚenabledÚprefill_from_target_hiddenr:   ©rz   r=   r;   Úaccept_verified_tokensé   r   )rÓ   r÷   râ   r¡   Ú__name__rï   rg   rK   Úresetr   r\   ÚcallableÚ
draft_callr]   r   r•   r   Údraft_blockr   Ústreamr   rL   r   rì   r¶   r¬   Útarget_evalr*   r
   rø   Úclear_cache) rÐ   rS   r™   r«   rð   ró   rô   rš   r_   rñ   rT   r—   r›   Úblock_totalri   Úadaptive_block_sizeÚdraft_cacheÚsampler_rngÚprefill_draftÚbÚemittedÚbsr   r˜   Ú
hot_verifyÚverify_hiddenr   r¥   r%   r)   rµ   Úaccept_verifieds    &&&&$$$$$$$                     r   Ú_eagle3_roundsr  X  s‡  é € ô ")¨Ð0@×!AÒ!Aˆ×	Ò	Àu€BÜ�2Ð3×4Ò4ÜÜ�B‹x× Ñ Ð!Ð!PÐQó
ð 	
ô 1°Ó=ÐÜ?UØó@Ñ<€KÐ)<ð ÒÜ˜+ qÓ)ˆÜ!$Ð%;Ó!IÐØ#ÐØ×#Ñ# EÓ*€KÜ(¨Ô>QÔR€Kä˜KÐ)EÀtÓL€MÜ�×Ò =Ò#<Ø×ÒØØØØØØñ	
ô # ;Ó@ò	
ð ‡|�|�A…˜ÔØ˜˜2™3 ˜	Õ"ˆà€AØ€Gà
�JÕ
Ü$ØØØ"Ø˜Õ  1Õ$Ø(ô
ˆð �Œ7Ùà"×/Ò/Ø×#Ñ#ØØØØØØñ	
ô # ;Ó@ñ	
ˆô �YŠYÔ(×)Õ)ÜŸ>š>Ü—’˜A˜3˜% {Ô3°\ÐBØôˆLô 3ØØØØØØ Ü% eÓ,óˆJð Ò!Ü;PØØ Ø ØØ$ó<Ñ8�˜}©jð <FÑ8�˜}¨j÷/ *ð0 	×Ñ ¨}Ô=ä+ØØØ˜Õ ó 
Ñˆ�*ô
 	" +¨x¸¸a½Ô@ãˆCØ�t�)ŠOØ�q�LˆGØ˜*Ö$Úñ	 ô " +Ð/GÈÓNˆÜ�O×$Ò$Ø×"Ò"ØØØØØØØñ	ô ' {ÓDò	ð ˜q (¨X¸­\Ð"9¸1Ð<Õ=ˆß(ˆJ�rŽN¨aˆà�b˜1•fÔÜ—’Ô,×-Õ-Ø×-Ñ-¨l¸JÈÐRTÔU÷ .ð �S�=˜A×Ü�NŠN×ñe ÷, *×)Ð)ú÷p .×-üsK   ‚FM1ÆA+M
ÈAM1É%A!M1Ë;M1ÌMÌM1Ì0M1Í
M	Í	M1ÍM.	Í)M1Ú
stop_checkr¸   c                 ó‚  € V ^8„  d   QhR\         P                  R\         P                  R\        \        ,          R\        P
                  R\        \        P
                  ,          R\        P
                  R\        R\        \        P
                  .\        P
                  3,          R	\        \        ,          R
\        P                  R\        \        \        \        .\        3,          ,          R\        \        ,          R\        R\        \        \        \        \        ,          ,          R3,          RR3,          /# )r   rÐ   rS   r™   r«   rð   ró   rô   rš   r_   rñ   r  r¸   rT   r   N)rU   rV   r   r   r   r   r   r   r   rõ   rW   ru   r   r   )r   s   "r   r   r   Þ  s  € ÷ Rñ RÜ�9‰9ðRä—‘ðRô ”s•)ðRô �H‰Hð	Rô œBŸH™HÕ%ðRô —‘ðRô ðRô ”r—x‘x�j¤"§(¡(Ð*Õ+ðRô œs•mðRô —‘ðRô œ¤3¬ *¬dÐ"2Õ3Õ4ðRô œC•=ðRô ðRô Œu”Tœ(¤3�-Õ(¨$Ð.Õ/°°tÐ;Õ<ñRr   c       	      #   ó†  a8a9a:"  € \        V R 4      '       d   V P                  MT p\        VR4      '       g"   \        \        V4      P                   R24      hVP
                  ^ ,          p\        V4      pVf   ^ .\        V4      ,          pM>\        VR4      '       d   VP                  4       MTpV Uu. uF  p\        V4      NK  	  pp\        V4      p\        W4      w  pppVf   \        V^4      p\        VV4      pRpVP                  V VR7      p\        W'       * R7      p\        VRR4      p\        V4      '       d&   Ve"   VP                   ! VVVVVV	3/ \#        W4      B  VP
                  ^,          ^8”  d   VR	RR1R	3,          pVP                  4       p^.V,          pR.V,          o:\%        \'        V4      4      o8\%        \'        V4      4      p\)        S84      ^ 8”  Ed†   \'        \)        S84      4       Uu. uF+  p\+        ^VVS8V,          ,          ,
          ^,           4      NK-  	  pp\-        VVV\        V4      VR
7      o9S9^8:  d   R# \)        S84      p\        VRR 4      ! 4       p V ef   V V8w  d_   \        VRR4      p!\        V!4      '       dA   V!! \.        P0                  ! V\.        P2                  R7      4       \%        \'        V4      4      p\'        V4       Uu. uF  pVS8V,          ,          NK  	  p"p\.        P0                  ! V"V	R7      p#VP4                  ! VP6                  V#VVS9VV	3/ \#        W4      B p$\.        P8                  ! \:        4      ;_uu_ 4        \.        P<                  ! V#R,          V$.^R7      p%\?        VVV%VVV\A        V 4      4      p&V&f   \C        VV%VVV4      w  p'p(p)MV&w  p'p(p)RRR4       VPE                  X(X'4       \'        V4       Uu. uF  pVVS8V,          ,          ,
          NK  	  p*p\G        V$V(V*4      w  p+p,V^8”  d>   \        VRR4      '       d+   \)        \I        V+4      4      ^8”  d   \K        V$V(V+V*4      w  p+p,\M        V\O        V+4      \)        V+4      ,          S9^,
          4       \        VRR4      p-\        V-4      '       d#   VP                   ! V-V'V$V+V,VV	3/ \#        W4      B  \.        PP                  ! V4      p.\.        P0                  ! V+4      p/V'V.V/R	3,          R,          pV,'       d   \+        R V, 4       4      M^ p0\'        V04       FÁ  p1R.V,          p2\'        V4       FŸ  pS8V,          p3V1\)        V,V,          4      8  g   K%  S:V3,          '       d   K6  V,V,          V1,          p4V4V2V3&   VV3;;,          ^,          uu&   VV3,          V8¼  d   RS:V3&   Ve   V4V9   d   RS:V3&   V
f   K‰  V
! V3V44      '       g   Kš  RS:V3&   K¡  	  V2R3x € KÃ  	  \'        V4       F0  pS8V,          p3V,V,          '       g   K  V,V,          R,          VV3&   K2  	  \R        ;QJ d    V93R lV+ 4       F  '       g   K   RM	  RM! V93R lV+ 4       4      '       dT   \.        P8                  ! \:        4      ;_uu_ 4        VPU                  VX)\.        P0                  ! V+4      S94       RRR4       \V        ;QJ d)    V8V:3R l\'        V4       4       F  '       d   K   RM	  RM! V8V:3R l\'        V4       4       4      '       d   R# \'        V4       Uu. uF  pS:S8V,          ,          '       d   K  VNK  	  p5p\)        V54      V8  dÛ   \.        P0                  ! V5\.        P2                  R7      p6\V        ;QJ d    R V 4       F  '       d   K   RM	  RM! R V 4       4      '       d   V F  p7V7PY                  V64       K  	  \        VRR4      p!\        V!4      '       d	   V!! V64       VV6,          pV5 Uu. uF  pS8V,          NK  	  upo8V5 Uu. uF  pVV,          NK  	  pp\O        V4      R,          ^ 8X  g   EK~  \.        PZ                  ! 4        EK–  R# u upi u upi u upi   + '       g   i     ELÅ; iu upi   + '       g   i     ELß; iu upi u upi u upi 5i)r÷   rø   rù   Nr   F)Úleft_paddingrú   rü   r:   rý   Ú
batch_sizec                  ó   € R # r‡   r�   r�   r   r   rŠ   Ú&_eagle3_rounds_batch.<locals>.<lambda>/  s   € Ádr   Úfilter_batchr=   r;   Ú!requires_uniform_batch_acceptanceÚaccept_verified_tokens_batchc              3   ó8   "  € T F  p\        V4      x € K  	  R # 5ir‡   )rM   )rn   Únts   & r   rp   Ú'_eagle3_rounds_batch.<locals>.<genexpr>�  s   é € Ð8© "”c˜"—g�g«ùs   ‚Tc              3   ó:   <"  € T F  qS^,
          8  x € K  	  R# 5ir€   r�   )rn   r%   r  s   & €r   rp   r   —  s   øé € Ð?± X˜"˜q�&Ö ³ùs   ƒc              3   óD   <"  € T F  pSSV,          ,          x € K  	  R # 5ir‡   r�   )rn   r6   Ú
active_idxÚfinisheds   & €€r   rp   r      s   øé € Ð@±¨1ˆx˜
 1�×&Ò&³ùs   ƒ c              3   ó:   "  € T F  p\        VR 4      x € K  	  R# 5i)ÚfilterN)rÓ   )rn   Úcs   & r   rp   r   ¥  s   é € Ð>±¨A”7˜1˜h×'Ð'³ùs   ‚rÿ   r   rá   )r:   Nr:   ).rÓ   r÷   râ   r¡   r   r   r	   r   r   rï   rg   rK   r  r   r\   r  r  r]   r²   r/   rM   rd   r•   r   r   rA   r   r  r  r   rL   rì   r¶   r¬   r  r7   ru   rQ   r
   r�   r½   rB   rø   Úallr&  r  );rÐ   rS   r™   r«   rð   ró   rô   rš   r_   rñ   r  r¸   rT   r—   r1   Úcache_left_paddingr  Úpadding_valuesÚpadr›   r  ri   r	  r
  r  r  r  r  Úcache_slotsr6   Ú	remainingÚn_activeÚdraft_batch_sizeÚfilter_drafterÚb_activeÚb_arrr   r˜   r  r  r   r¥   r,   r4   r5   r  Úrow_idxÚcol_idxÚmax_newÚposÚ
tokens_outÚorigrµ   Ú
keep_slotsÚkeep_mxr'  r#  r  r$  s;   &&&&$$$$$$$$$                                           @@@r   Ú_eagle3_rounds_batchr;  Þ  sÏ  úé € ô  ")¨Ð0@×!AÒ!Aˆ×	Ò	Àu€BÜ�2Ð3×4Ò4ÜÜ�B‹x× Ñ Ð!Ð!PÐQó
ð 	
ð 	×Ñ˜!Õ€AÜ2°<Ó@ÐØÒ!Ø�sœS ›V•|‰ô Ð)¨8×4Ò4ð ×%Ñ%Ô'à#ð 	ñ
 -;Ó;©N Sœ˜Cž©NˆÐ;Ü0°Ó=ÐÜ?UØó@Ñ<€KÐ'Ð)<ð ÒÜ˜+ qÓ)ˆÜ!$Ð%;¸[Ó!IÐØ#ÐØ×#Ñ# E¸Ð#ÓE€KÜ(¨Ô>QÔR€Kä˜KÐ)EÀtÓL€MÜ�×Ò =Ò#<Ø×ÒØØØØØØñ	
ô # ;Ó@ò	
ð ‡|�|�A…˜ÔØ˜˜2™3 ˜	Õ"ˆà×ÑÓ€AØˆc�A�g€GØˆw˜�{€HÜ”e˜A“h“€JÜ”u˜Q“x“.€Kä
ˆj‹/˜AÕ
ô œ3˜z›?Ô+ó
á+�ô ��: ¨
°1­Õ 6Õ6¸Õ:Ö;Ù+ð 	ð 
ô %ØØØ"Ü�	‹NØ(ô
ˆð �Œ7Ùä�z“?ˆÜ" ;°¹lÔKÓMÐØÒ'Ð,<ÀÔ,HÜ$ [°.À$ÓGˆNÜ˜×'Ò'ÙœrŸxšx¨¼2¿8¹8ÔDÔEÜ"¤5¨£?Ó3�ä.3°H¬oÓ>©o¨�A�j •m×$Ð$©oˆÐ>Ü—’˜¨Ô5ˆà"×/Ò/Ø×#Ñ#ØØØØØØñ	
ô # ;Ó@ñ	
ˆô �YŠYÔ(×)Õ)ÜŸ>š>¨5°­>¸<Ð*HÈqÔQˆLÜ2ØØØØØØ Ü% eÓ,óˆJð Ò!Ü;PØØ Ø ØØ$ó<Ñ8�˜}©jð <FÑ8�˜}¨j÷) *ð* 	×Ñ ¨}Ô=ä@EÀhÄÓPÁ¸1�: ¨
°1­Õ 6×6Ð6ÁˆÐPÜ);Ø˜-¨ó*
Ñ&ˆ�ð �qŒLÜ˜Ð%HÈ%×PÒPÜ”C˜Ó&Ó'¨!Ô+ä-RØØØØó	.Ñ*ˆM˜?ô 	"ØÜ�Ó¤ ]Ó!3Õ3Ø��Fô	
ô " +Ð/MÈtÓTˆÜ�O×$Ò$Ø×"Ò"ØØØØØØØñ	ô ' {ÓDò	ô —)’)˜HÓ%ˆÜ—(’(˜=Ó)ˆØ˜w¨°Ð2Õ3°JÕ?ˆç<K”#Ñ8©Ó8Ô8ÐQRˆÜ˜–>ˆCØ/3¨f°q­jˆJÜ˜8–_�Ø! !•}�Øœ˜_¨QÕ/Ó0Ö0¸À$¿¾Ø)¨!Õ,¨SÕ1�CØ'*�J˜tÑ$Ø˜D—M QÕ&“MØ˜t•}¨
Ô2Ø)-˜ ™Ø$Ò0°S¸MÔ5IØ)-˜ ™Ø!Ô-±*¸TÀ3×2GÔ2GØ)-˜ ›ñ %ð ˜dÐ"Ô"ñ "ô  �x–ˆAØ˜a•=ˆDØ˜q×!Ö!Ø)¨!Õ,¨RÕ0��$“ñ !÷
 ‹3Ô?±Ó?�3�3Š3Ô?±Ó?×?Ò?Ü—’Ô,×-Õ-Ø×-Ñ-Ø ØÜ—H’H˜]Ó+Øô	÷ .÷ ‹3Õ@´°h´Ó@�3�3Š3Õ@´°h´Ó@×@Ò@ÙÜ!& x¤ÓP¡˜A¸ÀÈAÅ×8OÑ8O—a�a¡ˆ
ÐPÜˆz‹?˜XÔ%Ü—h’h˜z´·±Ô:ˆGß‹sÑ>±Ó>�s�sŠsÑ>±Ó>×>Ò>Û%�AØ—H‘H˜WÖ%ñ &ä$ [°.À$ÓGˆNÜ˜×'Ò'Ù˜wÔ'Ø˜G•_ˆFÙ1;Ó<±¨A˜* QŸ-˜-±Ñ<ˆJÙ3=Ó>±:¨a˜; qŸ>˜>±:ˆKÐ>äˆw‹<˜#Õ ×"Ü�NŠN×ñc ùòC <ùòD
ùò, ?÷ *×)Ð)üò. Q÷z .×-Ð-üò Qùò =ùÚ>ùsò   …B2cÂ7a6ÃD:cÈ1a;È6B8cË.b ÌA(cÍ/AbÏ	'cÏ0 bÐE!cÕ6cÖAc×c×+6cØ&!cÙcÙcÙ="cÚ)bÛcÛcÛ<'cÜ$cÜ3b2Ýb2ÝAcÞcÞ-cßAcàb7à'cà.b<á cá)câb	âcâb/	â)cr‡   )"Útypingr   r   r   r   r   r   Úmlx.coreÚcorer   Úmlx.nnrU   Úcommonr	   r
   r   r   r*   r7   rH   rQ   r]   rg   rw   r•   r¬   r¶   rÃ   rÜ   rì   rï   rA   r  r;  r�   r   r   Ú<module>rA     sî   ðß B× Bå Ý ÷ó õ)õ&*õ2	õ<õ(õ7õ<õ%õP'8õT÷õ2.÷b,'õ^ðCð )-ðCð '+ðCð ŸH™HðCð "÷CðLRð )-ðRð '+ðRð ŸH™HðRð 8<ðRð $(ðRð "÷Rñ Rr   