+
    ÒÜ³jN  ã                   ó¸   € ^ RI t ^ RIt^ RIHtHt ^RIHt ^ RIHt ^ RI	H
t
 ^ RIHt ^ RIHtHt ^ RIHtHt ^ R	IHt ^ R
IHt  ! R R]4      t ! R R4      tR# )é    N)ÚListÚTuple)ÚBaseAWQForCausalLM)ÚMixtralBlock)ÚMixtralModel)ÚFusedSparseMoeBlock)Úfuse_qkvÚfuse_linears)ÚMixtralDecoderLayerÚMixtralForCausalLM)ÚWQLinear_GEMM)ÚFasterTransformerRMSNormc                   ó®   a € ] tR t^t o RtRtR.t]V 3R lR l4       t]V 3R lR l4       t	]R 4       t
]V 3R	 lR
 l4       t]V 3R lR l4       tRtV tR# )ÚMixtralAWQForCausalLMr   Úmax_position_embeddingsÚgatec                ó    <€ V ^8„  d   QhRS[ /# ©é   Úmodel©ÚOldMixtralForCausalLM)ÚformatÚ__classdict__s   "€ÚS/home/luca/screen-checking/.venv/lib/python3.14/site-packages/awq/models/mixtral.pyÚ__annotate__Ú"MixtralAWQForCausalLM.__annotate__   s   ø€ ÷ !ñ !Ñ0ñ !ó    c                ó<   € \        V 4      pVP                  4        R # ©N)ÚMixtralFuserÚfuse_transformer)r   Úfusers   & r   Úfuse_layersÚ!MixtralAWQForCausalLM.fuse_layers   s   € ä˜UÓ#ˆØ×ÑÖ r   c                ó    <€ V ^8„  d   QhRS[ /# r   r   )r   r   s   "€r   r   r      s   ø€ ÷ "ñ "Ñ 5ñ "r   c                ó.   € V P                   P                  # r    )r   Úlayers)r   s   &r   Úget_model_layersÚ&MixtralAWQForCausalLM.get_model_layers   s   € à�{‰{×!Ñ!Ð!r   c                ó   € \        R R7      # )F)Úis_scalable)Údict)Úmodules   &r   Úget_act_for_scalingÚ)MixtralAWQForCausalLM.get_act_for_scaling   s   € ä Ô&Ð&r   c                ó&   <€ V ^8„  d   QhRS[ RS[/# )r   r   Údevice)r   Ústr)r   r   s   "€r   r   r   $   s    ø€ ÷ Gñ GÑ/ð G¹ñ Gr   c                ón   € V P                   P                  P                  V4      V P                   n        R # r    )r   Úembed_tokensÚto)r   r2   s   &&r   Ú
move_embedÚ MixtralAWQForCausalLM.move_embed#   s#   € à#(§;¡;×#;Ñ#;×#>Ñ#>¸vÓ#Fˆ�‰Ö r   c                ó    <€ V ^8„  d   QhRS[ /# )r   r.   )ÚOldMixtralDecoderLayer)r   r   s   "€r   r   r   (   s   ø€ ÷ 6ñ 6Ù&ñ6r   c                ó  € . pVP                  \        V P                  V P                  P                  V P                  P
                  V P                  P                  .VR ,          V P                  VR7      4       V P                  P                  P                  P                  V P                  P                  P                  P                  8X  dN   VP                  \        V P                  P                  V P                  P                  .VR,          R7      4       TP                  \        V P                  V P                  P                   UUu. uF$  pVP                  VP                  3 F  pVNK  	  K&  	  uppVR,          V P                  R7      4       \        V P                  P                  4       FB  w  rdVP                  \        VP                  VP                   .VRV R2,          R7      4       KD  	  V# u uppi )zself_attn.q_proj)Úprev_opr(   ÚinpÚmodule2inspectÚkwargszself_attn.o_proj)r<   r(   r=   Úblock_sparse_moe)r<   r(   r=   r>   zblock_sparse_moe.experts.z.w2)Úappendr-   Úinput_layernormÚ	self_attnÚq_projÚk_projÚv_projÚweightÚshapeÚo_projÚpost_attention_layernormr@   ÚexpertsÚw1Úw3Ú	enumerateÚw2)r.   Ú
input_featÚmodule_kwargsr(   ÚexpertÚwÚis   &&&    r   Úget_layers_for_scalingÚ,MixtralAWQForCausalLM.get_layers_for_scaling'   sÁ  € ð ˆð 	�‰ÜØ×.Ñ.à×$Ñ$×+Ñ+Ø×$Ñ$×+Ñ+Ø×$Ñ$×+Ñ+ðð
 Ð1Õ2Ø%×/Ñ/Ø$ô
ô	
ð ×Ñ×"Ñ"×)Ñ)×/Ñ/°6×3CÑ3C×3JÑ3J×3QÑ3Q×3WÑ3WÔWØ�M‰MÜØ"×,Ñ,×3Ñ3Ø"×,Ñ,×3Ñ3Ð4Ø"Ð#5Õ6ôôð 	�‰ÜØ×7Ñ7ð #)×"9Ñ"9×"AÒ"Aôá"A˜Ø$Ÿi™i¨¯©Ó3˜ó á3ñ Ù"Aòð
 Ð1Õ2Ø%×6Ñ6ô	ô	
ô # 6×#:Ñ#:×#BÑ#BÖC‰IˆAØ�M‰MÜØ"ŸI™IØ"ŸI™I˜;Ø"Ð%>¸q¸cÀÐ#EÕFôöñ Dð ˆùó)s   Å*H© N)Ú__name__Ú
__module__Ú__qualname__Ú__firstlineno__Ú
layer_typeÚmax_seq_len_keyÚmodules_to_not_convertÚstaticmethodr$   r)   r/   r7   rU   Ú__static_attributes__Ú__classdictcell__©r   s   @r   r   r      s{   ø‡ € Ø&€JØ/€OØ$˜XÐà÷!ó ð!ð ÷"ó ð"ð ñ'ó ð'ð ÷Gó ðGð ÷6ó ö6r   r   c                   ó8   a € ] tR t^at o V 3R lR ltR tRtV tR# )r!   c                ó    <€ V ^8„  d   QhRS[ /# r   r   )r   r   s   "€r   r   ÚMixtralFuser.__annotate__b   s   ø€ ÷ 
ñ 
Ñ3ñ 
r   c                óø   € Wn         V P                   P                  4        UUu. uFA  w  r#R P                  4       VP                  P                  P                  4       9   g   K>  W#3NKC  	  uppV n        R# u uppi )r   N)r   Únamed_modulesÚlowerÚ	__class__rX   Úmixtral_blocks)Úselfr   Únamer.   s   &&  r   Ú__init__ÚMixtralFuser.__init__b   sj   € ØŒ
ð !%§
¡
× 8Ñ 8Ô :ôI
á :‘�Ø$×*Ñ*Ó,°×0@Ñ0@×0IÑ0I×0OÑ0OÓ0QÑQô ˆT‹NÙ :òI
ˆÖùó I
s   ¤;A6Á$A6c                óP  € . p\         P                   ! V P                  P                  P                  R R7       EFê  p\        \	        VP                  4       P                  4       4      4      P                  p\        VVP                  P                  VP                  P                  VP                  P                  4      p\        VP                  P                  VP                  P                   4      p\        VP"                  P                  VP"                  P                   4      pVP$                  p\'        VP(                  ^ ,          P*                  \,        4      '       dï   \/        \1        VP(                  4      4       Uu. uFG  p\3        VP(                  V,          P*                  VP(                  V,          P4                  .V4      NKI  	  p	p\3        W“^ \6        P8                  R7      p
\3        VP(                   Uu. uF  q»P:                  NK  	  upV^ \6        P8                  R7      p\=        VP>                  VP@                  V
VR7      pVPC                  \E        V P                  PF                  PH                  V P                  PF                  PJ                  V P                  PF                  PL                  VVP                  PN                  VVVVV P                  PF                  PP                  V P                  PF                  PR                  R7      4       EKí  	  \        V P                  P                  PT                  P                  V P                  P                  PT                  P                   4      p\W        V P                  PF                  PX                  VV P                  P                  PZ                  V4      V P                  n        \]        V P                  P                  RV P                  P                  P^                  4       R# u upi u upi )zFusing layers...)Údesc)ÚdimÚ	operation)Útop_kr   ÚwsÚw2s)Úhidden_sizeÚn_headsÚ
n_kv_headsÚ	qkv_layerrI   ÚmoeÚnorm_1Únorm_2ÚdevÚmax_seq_lenÚ
rope_thetaÚblocksN)0Útqdmr   r(   ÚnextÚiterÚ
state_dictÚvaluesr2   r	   rC   rD   rE   rF   r   rB   rG   Úvariance_epsilonrJ   r@   Ú
isinstancerK   rL   r   ÚrangeÚlenr
   rM   ÚtorchÚstackrO   r   rs   r   rA   r   Úconfigrv   Únum_attention_headsÚnum_key_value_headsrI   r~   r   Únormr   Ú
vocab_sizer5   Úsetattrr€   )rk   r€   r.   r2   Úqkvr{   r|   Ú
sparse_moerT   Úfused_w1w3sÚstacked_w1w3srR   Ústacked_w2sÚ
model_norms   &             r   r"   ÚMixtralFuser.fuse_transformerk   s3  € Øˆô —i’i §
¡
× 0Ñ 0× 7Ñ 7Ð>P×QÐQˆFÜœ$˜v×0Ñ0Ó2×9Ñ9Ó;Ó<Ó=×DÑDˆFäØØ× Ñ ×'Ñ'Ø× Ñ ×'Ñ'Ø× Ñ ×'Ñ'ó	ˆCô .Ø×&Ñ&×-Ñ-¨v×/EÑ/E×/VÑ/VóˆFô .Ø×/Ñ/×6Ñ6Ø×/Ñ/×@Ñ@óˆFð
  ×0Ñ0ˆJÜ˜*×,Ñ,¨QÕ/×2Ñ2´M×BÒBô #¤3 z×'9Ñ'9Ó#:Ô;ó	ñ <˜ô !à&×.Ñ.¨qÕ1×4Ñ4Ø&×.Ñ.¨qÕ1×4Ñ4ðð öñ <ð ð 	ô !-Ø¨Q¼%¿+¹+ô!�ô +Ø-7×-?Ò-?Ó@Ñ-? 6—Y”YÑ-?Ñ@ØØÜ#Ÿk™kô	�ô 1Ø$×*Ñ*Ø#Ÿ™Ø$Ø#ô	�
ð �M‰MÜØ $§
¡
× 1Ñ 1× =Ñ =Ø ŸJ™J×-Ñ-×AÑAØ#Ÿz™z×0Ñ0×DÑDØ!Ø!×+Ñ+×2Ñ2Ø"Ø!Ø!ØØ $§
¡
× 1Ñ 1× =Ñ =Ø#Ÿz™z×0Ñ0×;Ñ;ô÷ñc RôB .Ø�J‰J×Ñ×!Ñ!×(Ñ(Ø�J‰J×Ñ×!Ñ!×2Ñ2ó
ˆ
ô
 (Ø�J‰J×Ñ×(Ñ(ØØ�J‰J×Ñ×)Ñ)Øó	
ˆ�
‰
Ôô 	�—
‘
× Ñ  (¨D¯J©J×,<Ñ,<×,CÑ,CÖDùòq	ùò  As   ÆAPÈP#
)rj   r   N)rX   rY   rZ   r[   rm   r"   r`   ra   rb   s   @r   r!   r!   a   s   ø‡ € ÷
ð 
÷PEð PEr   r!   )r�   rŠ   Útypingr   r   Úbaser   Úawq.modules.fused.blockr   Úawq.modules.fused.modelr   Úawq.modules.fused.moer   Úawq.utils.fused_utilsr	   r
   Ú,transformers.models.mixtral.modeling_mixtralr   r:   r   r   Úawq.modules.linearr   Úawq.modules.fused.normr   r   r!   rW   r   r   Ú<module>r¢      sH   ðÛ Û ß Ý $Ý 0Ý 0Ý 5ß 8÷õ -Ý ;ôMÐ.ô M÷`ZEó ZEr   