+
    œ°³jª  ã                  óÂ   € ^ 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 ]! 4       '       d   ^ RIt]'       d   ^RIHt ^R	IHt R
 R lt ! R R]4      tR# )é    )Úannotations)ÚTYPE_CHECKING)Úis_kernels_availableÚis_torch_available)ÚKERNELS_MAX_VERSIONÚKERNELS_MIN_VERSION)ÚHfQuantizer)Úget_module_from_nameN)ÚPreTrainedModel)ÚNVFP4Configc               ó   € V ^8„  d   QhRR/# )é   Úreturnztorch.device | None© )Úformats   "Úh/home/luca/screen-checking/.venv/lib/python3.14/site-packages/transformers/quantizers/quantizer_nvfp4.pyÚ__annotate__r       s   € ÷ 5ñ 5Ð2ñ 5ó    c                óä   € \        V \        4      '       d   \        P                  ! R V 4      #  \        P                  ! V 4      p T P                  R 8X  d   T # R#   \        \
        3 d     R# i ; i)ÚcudaN)Ú
isinstanceÚintÚtorchÚdeviceÚ	TypeErrorÚRuntimeErrorÚtype)r   s   &r   Ú_as_cuda_devicer       sd   € Ü�&œ#×ÒÜ�|Š|˜F FÓ+Ð+ðÜ—’˜fÓ%ˆð —[‘[ FÔ*ˆ6Ð4°Ð4øô ”|Ð$ô Úðús   ¯A ÁA/Á.A/c                  ó®   a € ] tR t^*t$ RtRtR]R&   R tR tR R lt	R	 V 3R
 llt
R R ltR t]R R l4       t]R R l4       t]R R l4       tRtV ;t# )ÚNVFP4HfQuantizerzTQuantize eligible linear weights to NVFP4 while loading a full-precision checkpoint.Fr   Úquantization_configc           	     	óÊ  € V P                   '       d   \        R 4      h\        4       '       g"   \        R\         R\
         R\         R24      h\        P                  P                  4       '       g   \        R4      hVf   \        R4      h\        V\        4      '       d   \        VP                  4       4      MV0pV Uu0 uF  p\        V4      kK  	  ppRV9   d   \        R4      h\        V4      ^8w  d   \        R	4      hVP!                  4       V n        \        P                  P%                  V P"                  4      w  rgV^
8  d   \        R
V V R24      hR# u upi )z=Loading pre-quantized NVFP4 checkpoints is not supported yet.zQNVFP4 quantization requires the `kernels` package. Install a compatible version (z <= version < z*), for example with `pip install kernels==z`.z'NVFP4 quantization requires a CUDA GPU.Nz0NVFP4 quantization requires a CUDA `device_map`.z8NVFP4 quantization does not support CPU or disk offload.z�NVFP4 quantization currently supports one CUDA device. Tensor parallelism and multi-device `device_map` configurations are not supported yet.zENVFP4 requires Blackwell (sm100+) block-scaled tensor cores; found smÚ.)Úpre_quantizedÚ
ValueErrorr   ÚImportErrorr   r   r   r   Úis_availabler   r   ÚdictÚsetÚvaluesr   ÚlenÚpopÚquantization_deviceÚget_device_capability)ÚselfÚ
device_mapÚkwargsÚdevicesr   Úcuda_devicesÚmajorÚminors   &&,     r   Úvalidate_environmentÚ%NVFP4HfQuantizer.validate_environment0   s[  € Ø××ÐÜÐ\Ó]Ð]Ü#×%Ò%Üð1Ü1DÐ0EÀ^ÔTgÐShð i:Ü:MÐ9NÈbðRóð ô
 �z‰z×&Ñ&×(Ò(ÜÐHÓIÐIàÒÜÐOÓPÐPÜ.8¸ÄT×.JÒ.J”#�j×'Ñ'Ó)Ô*ÐQ[ÐP\ˆÙ>EÓF¹g°Fœ¨Ö/¹gˆÐFØ�<ÔÜÐWÓXÐXÜˆ|Ó Ô!ÜðEóð ð $0×#3Ñ#3Ó#5ˆÔ ä—z‘z×7Ñ7¸×8PÑ8PÓQ‰ˆØ�2Œ:ÜÐ!fÐglÐfmÐnsÐmtÐtuÐvÓwÐwñ ùò Gs   ÃE c                	óf   € \        VR R4      p\        VRR4      pVe   V^8”  d   \        R4      hV# )Údistributed_configNÚtp_sizeztNVFP4 quantization does not support tensor parallelism yet because its scale metadata does not have a sharding plan.)Úgetattrr%   )r/   Úconfigr9   r:   s   &&  r   Úupdate_tp_planÚNVFP4HfQuantizer.update_tp_planM   sF   € Ü$ VÐ-AÀ4ÓHÐÜÐ,¨i¸Ó>ˆØÒ 7¨Q¤;Üð#óð ð ˆr   c               ó$   € V ^8„  d   QhRRRRRR/# )r   Úmodelr   Ú
param_nameÚstrr   Úboolr   )r   s   "r   r   ÚNVFP4HfQuantizer.__annotate__W   s'   € ÷ Kñ K¨oð KÈ3ð KÐ]añ Kr   c                	óX   € ^RI Hp \        W4      w  rV\        WT4      ;'       d    VR8H  # )r   )ÚNVFP4LinearÚweight)Úintegrations.nvfp4rF   r
   r   )r/   r@   rA   r1   rF   ÚmoduleÚtensor_names   &&&,   r   Úparam_needs_quantizationÚ)NVFP4HfQuantizer.param_needs_quantizationW   s*   € Ý4ä2°5ÓEÑˆÜ˜&Ó.×JÐJ°;À(Ñ3JÐJr   c               ó(   € V ^8„  d   QhRRRRRRRR/# )	r   r@   r   rA   rB   Úparamztorch.Tensorr   Úfloatr   )r   s   "r   r   rD   ]   s0   € ÷ Dñ D¨ð DÀSð DÐQ]ð DÐbgñ Dr   c                	óV   <€ V P                  W4      '       d   R # \        SV `	  WV4      # )g      à?)rK   ÚsuperÚparam_element_size)r/   r@   rA   rN   Ú	__class__s   &&&&€r   rR   Ú#NVFP4HfQuantizer.param_element_size]   s)   ø€ Ø×(Ñ(¨×;Ò;áÜ‰wÑ)¨%¸UÓCÐCr   c               ó   € V ^8„  d   QhRR/# )r   r@   r   r   )r   s   "r   r   rD   c   s   € ÷ ]ñ ]¸/ñ ]r   c                	ó¢   € ^RI Hp V P                  WP                  P                  VP
                  4      V n        V! WP                  R7       R# )r   )Úreplace_with_nvfp4_linear)Úmodules_to_not_convertN)rH   rW   Úget_modules_to_not_convertr!   rX   Ú_keep_in_fp32_modules)r/   r@   r1   rW   s   &&, r   Ú$_process_model_before_weight_loadingÚ5NVFP4HfQuantizer._process_model_before_weight_loadingc   s?   € ÝBà&*×&EÑ&EØ×+Ñ+×BÑBÀE×D_ÑD_ó'
ˆÔ#ñ 	" %×@[Ñ@[×\r   c                	ó2   € ^RI Hp V! V P                  4      # )r   )ÚNVFP4Quantize)rH   r^   r-   )r/   r^   s   & r   Úget_quantize_opsÚ!NVFP4HfQuantizer.get_quantize_opsk   s   € Ý6á˜T×5Ñ5Ó6Ð6r   c               ó   € V ^8„  d   QhRR/# ©r   r   rC   r   )r   s   "r   r   rD   q   s   € ÷ ñ  ñ r   c                	ó   € R # ©Fr   ©r/   s   &r   Úis_serializableÚ NVFP4HfQuantizer.is_serializablep   ó   € ár   c               ó   € V ^8„  d   QhRR/# rb   r   )r   s   "r   r   rD   u   s   € ÷ ñ ˜dñ r   c                	ó   € R # rd   r   re   s   &r   Úis_trainableÚNVFP4HfQuantizer.is_trainablet   rh   r   c               ó   € V ^8„  d   QhRR/# rb   r   )r   s   "r   r   rD   y   s   € ÷ ñ  ñ r   c                	ó   € R # )Tr   re   s   &r   Úis_compileableÚNVFP4HfQuantizer.is_compileablex   s   € ár   )rX   r-   )Ú__name__Ú
__module__Ú__qualname__Ú__firstlineno__Ú__doc__Úrequires_calibrationÚ__annotations__r6   r=   rK   rR   r[   r_   Úpropertyrf   rk   ro   Ú__static_attributes__Ú__classcell__)rS   s   @r   r    r    *   sv   ø‡ Ù^à ÐØ$Ó$òxò:õK÷Dð Dõ]ò7ð
 ôó ðð ôó ðð ôó ör   r    )Ú
__future__r   Útypingr   Úutilsr   r   Úutils.import_utilsr   r   Úbaser	   Úquantizers_utilsr
   r   Úmodeling_utilsr   Úutils.quantization_configr   r   r    r   r   r   Ú<module>rƒ      sC   ðõ #å  ç <ß IÝ Ý 2ñ ×ÒÛçÝ0Ý7õ5ôP�{ö Pr   