
    $HJj                       d dl mZ d dlZd dlZd dlmZ d dlmZ d dlZd dl	m
Z
 d dl	mZ ddlmZ ddlmZ dd	lmZmZ dd
lmZ ddlmZmZmZmZ ddlmZmZmZ ddlm Z  ddl!m"Z"m#Z# ddl$m%Z%  ejL                  e'      Z(ejR                  Z* ejV                  e*      jX                  Z- ejV                  e*      j\                  Z/ej`                  d-d       Z1d Z2 ed       G d d             Z3ej`                  d.d       Z4ejj                  jl                  d/d       Z7d.dZ8d0dZ9	 	 	 	 d1	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 d2dZ:	 	 	 	 d3	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 d4dZ;	 	 	 	 d3	 	 	 	 	 	 	 	 	 	 	 	 	 	 	 d4dZ< G d de
jz                        Z> G d  d!e>      Z?	 	 	 	 	 	 	 	 	 	 d5d"Z@	 	 	 	 	 	 	 	 	 	 d5d#ZA G d$ d%e
j                        ZC G d& d'e"      ZD eD       ZE	 d6	 d7d(ZF G d) d*e      ZG G d+ d,e      ZHy)8    )annotationsN)Callable)	dataclass)
functional   )ACT2FN)ConversionOps)get_module_from_nameshould_convert_module)logging)KERNELS_MAX_VERSIONKERNELS_MIN_VERSIONis_kernels_availableis_torchdynamo_compiling   ) deepgemm_fp8_fp4_experts_forwarddeepgemm_fp8_fp4_linear(deepgemm_fp8_fp4_megamoe_experts_forward)lazy_load_kernel)ExpertsInterfaceuse_experts_implementation)to_localc                 |    t        t        d      st        dt        j                   d      t        j                  S )uS  Return ``torch.float8_e8m0fnu`` or raise a clear error on torch without FP8 support.

    UE8M0 scales are always stored/consumed as this single dtype — the kernels (Triton
    finegrained + DeepGEMM) read it natively, and supporting the same scales in mixed
    container dtypes would be a mess — so fail loudly rather than fall back.float8_e8m0fnuzbscale_fmt='ue8m0' requires torch.float8_e8m0fnu, which is only available in PyTorch >= 2.7 (found z.). Upgrade torch to use UE8M0 FP8 checkpoints.)hasattrtorchRuntimeError__version__r        s/Users/ahmed/devFolder/claude-voice/.venv/lib/python3.12/site-packages/transformers/integrations/finegrained_fp8.py_get_ue8m0_dtyper"   5   sD     5*+%%*%6%6$77eg
 	
 r    c                    |D ]  }t        | |      st        | |      c S  t        t        |       j                   d|       )Nz has none of: )r   getattrAttributeErrortype__name__)objnamesnames      r!   _first_attrr+   D   sF    33%%  DI../~eWE
FFr    T)frozenc                  0    e Zd ZU dZded<   ded<   ded<   y)FineGrainedFP8zNEntry points exposed by the `kernels-community/finegrained-fp8` Triton kernel.r   matmulbatched_matmulgrouped_matmulN)r'   
__module____qualname____doc____annotations__r   r    r!   r.   r.   K   s    Xr    r.   c                    t               s+t               s!t        dt         dt         dt         d      t        d      } | t        d      t        | dd      }t        | d	d      }t        | d
d      }d|fd	|fd
|ffD cg c]	  \  }}|| }}}|r3t        ddj                  |       dt         dt         dt         d	      t        |||      S c c}}w )z
    Load the finegrained-fp8 Triton kernel once and return its entry points.

    Raises `ImportError` if the `kernels` package is missing, or the kernel or required
    symbols cannot be found.
    z\finegrained-fp8 kernel requires the `kernels` package. Please install a compatible version (z <= version < z), e.g. `pip install kernels==`zfinegrained-fp8Nu   Failed to load the finegrained-fp8 kernel — check that `kernels-community/finegrained-fp8` has a build matching the current torch/CUDA.	matmul_2dmatmul_batchedmatmul_groupedz4finegrained-fp8 kernel is missing required symbols: , z'. Please install a compatible version ()r/   r0   r1   )	r   r   ImportErrorr   r   r   r$   joinr.   )kernelr/   r0   r1   r*   attrmissings          r!   _load_finegrained_fp8_kernelrA   T   sI    $%#%88K7LN[nZo p..A-B!E  /0F~;
 	

 V[$/FV%5t<NV%5t<N
 &!~.~.

JD$
 < 	
   B499WCUBV W44G3HWjVk l**=)>aA
 	
 %% !s   Cc                     t               } y N)rA   )_s    r!    _populate_finegrained_fp8_kernelrE      s    $&Ar    c                 >    t               r
t                t               S rC   )r   rE   rA   r   r    r!   load_finegrained_fp8_kernelrG      s    !(*'))r    c                    | |z   dz
  |z  S )zCeiling division.r   r   )abs     r!   _cdivrK      s    EAI!r    c	                \   t        j                  | |||z  |      }	t        j                  |	|	j	                               }
t        |t        ||      nd|      }|t        ||      nd}t        j                  | |||      }t        j                  ||j	                               }|
|fS )u  Allocate `(weight, weight_scale_inv)` parameters for one expert projection.

    `weight_k_div` halves the K dim for FP4-packed storage (2 e2m1 values per byte).
    `sf_gran_n` / `sf_gran_k` set per-block (None → per-row/per-tensor) SF granularity.
    `min_sf_out` floors the SF tensor's output dim — used by the fused gate_up
    projection to keep room for both halves (pass `2`) even when `proj_out < sf_gran_n`
    would otherwise collapse the SF dim to 1.
    dtyperequires_gradr   )r   emptynn	Parameteris_floating_pointmaxrK   )num_expertsproj_outproj_inweight_dtypesf_dtypeweight_k_div	sf_gran_n	sf_gran_k
min_sf_outweight_tweightsf_outsf_insf_tsfs                  r!   _alloc_expert_projre      s    & {{;'\2IQ]^H\\((2L2L2NOFy/Dx+!ZXF)2)>E'9%AE;;{FEBD	d$*@*@*B	CB2:r    c                n    t               }|j                  | |||||      }||j                  |       |S )u0  Triton FP8/FP4 linear: fused act-quant + matmul, then optional bias add.

    ``activation_scale=None`` → dynamic per-K-block scales (inline); set it for
    static per-tensor quant. ``weight_scale_inv`` accepts fp32 or UE8M0; the
    dispatcher routes FP4 (``int8``-packed) weights automatically.
    activation_scale)rG   r/   add_)	inputr`   weight_scale_inv
block_sizebiasrh   output_dtypefinegrained_fp8outputs	            r!   finegrained_fp8_linearrq      sL     23O##) $ F DMr    c           	         |du xr |j                   j                  dk(  xr t        j                  j	                         j
                  dk\  xrm |j                  t        j                  k(  xs |duxr |d   |d   cxk(  xr dk(  nc xr0 t        j                  j                  dd      d	k7  xr t                }|r	 t        | ||||||
      S t!        | ||||||      S # t        $ r"}t        j                  d|        Y d}~7d}~ww xY w)u  End-to-end FP8/FP4 linear used by `FP8Linear` and the eager `FP8Experts` loop.

    Dispatch order — both backends handle FP8 and FP4 weights with fp32 or UE8M0 scales:
      1. DeepGEMM (`deepgemm_fp8_fp4_linear`) — 3-6× faster on the shapes it supports.
         Preferred for FP4, UE8M0 SFs, and 128×128 block FP8.
      2. Triton finegrained-fp8 fallback — used when DeepGEMM is unavailable, when the
         caller passes ``activation_scale`` (DeepGEMM is dynamic-only), or for any
         shape DeepGEMM declined.

    Args:
        input: (..., K) bf16/fp16 activations.
        weight: (N, K) `float8_e4m3fn` or (N, K // 2) `int8` (FP4-packed).
        weight_scale_inv: per-block weight scales — `float32` (V3-style) or `float8_e8m0fnu`
            (V4-style; reinterpreted as int32 at the DeepGEMM kernel boundary).
        block_size: [block_n, block_k] for FP8 block-wise quant, or None/[N, K] for per-tensor.
            Ignored for FP4 weights (the kernel infers SF granularity from the dtype).
        bias: optional bias added to the matmul output.
        activation_scale: pass a per-tensor scalar to use static activation quant; leave `None`
            for dynamic (per-token) quant.
        output_dtype: desired output dtype.
    Ncuda	   r   r      $TRANSFORMERS_DISABLE_DEEPGEMM_LINEAR01)rl   rn   rh   rm   zDDeepGEMM unavailable for this call, falling back to Triton. Reason: )devicer&   r   rs   get_device_propertiesmajorrN   int8osenvirongetr   r   r<   loggerwarning_oncerq   )	rj   r`   rk   rl   rm   rh   rn   deepgemm_preferredes	            r!   
fp8_linearr      s/   T 	D  	+MM&(	+JJ,,.449	+ \\UZZ'mJd,B,lzRS}XbcdXeGlilGl	+ JJNNA3G3N		+
 )**  	l* %)!1  "%1A:tUegstt  	l "fghfi jkk	ls   ?C" "	D+DDc                  H     e Zd Z	 	 	 	 d	 	 	 	 	 	 	 	 	 	 	 d fdZddZ xZS )	FP8Linearc                   t         
|   ||       || _        || _        || _        t
        j                  j                  t        j                  ||t                    | _
        | j                  >t        j                  t        j                  dt
        j                              | _        n|dk(  r
t               nt
        j                  }|| j                  d   z   dz
  | j                  d   z  }|| j                  d   z   dz
  | j                  d   z  }	t        j                  t        j                  ||	|      |j                        | _        | j                  dk(  r>t        j                  t        j                  dt
        j                              | _        n| j#                  dd        | j                  r8t        j                  t        j                  | j$                              | _        y | j#                  d	d        y )
NrM         ?ue8m0r   r   rO   staticrh   rm   )super__init__has_biasrl   activation_schemer   rR   rS   rQ   
_FP8_DTYPEr`   tensorfloat32rk   r"   rT   rh   register_parameterout_featuresrm   )selfin_featuresr   rl   r   	scale_fmtr   rZ   scale_out_featuresscale_in_features	__class__s             r!   r   zFP8Linear.__init__  s    	l3 $!2hh((\;V`)ab??"$&LLc1W$XD!-6'-A')u}}H".1C"Ca"GDOO\]L^!^!,tq/A!AA!E$//Z[J\ \$&LL.0AR&88%D!
 !!X-$&LLc1W$XD!##$6===U[[1B1B%CDDI##FD1r    c           	     Z   | j                   j                         dkD  r+t        j                  || j                   | j                        S t        | j                         }t        | j                        }t        |||| j                  | j                  |j                  | j                        S )Nr   )rl   rh   rn   rm   )r`   element_sizeFlinearrm   r   rk   r   rl   rh   rN   )r   rj   r`   	scale_invs       r!   forwardzFP8Linear.forward<  s    ;;##%)88E4;;		::$++&T223	!22
 	
r    NdynamicfloatF)r   intr   r   rl   tuple[int, int] | Noner   strr   r   r   bool)rj   torch.Tensorreturnr   )r'   r2   r3   r   r   __classcell__r   s   @r!   r   r     sV    
 .2!* $2$2 $2 +	$2
 $2 $2 $2L
r    r   c                  P     e Zd ZdZ	 	 	 	 d	 	 	 	 	 	 	 	 	 	 	 	 	 d fdZddZ xZS )FP8GroupedLinearu  FP8 drop-in for block-diagonal grouped linears.

    The underlying nn.Linear stores a single `(n_groups * out_per_group, in_per_group)`
    weight; logically that's `n_groups` independent `(out_per_group, in_per_group)`
    sub-matrices, each consuming a disjoint slice of the input's last-but-one dim.
    Forward expects input of shape `(..., n_groups, in_per_group)` and returns
    `(..., n_groups, out_per_group)` — same contract as the vanilla bf16 grouped
    linear it replaces.

    c                >    t         |   ||||||       || _        y )Nr   r   rl   r   r   r   )r   r   n_groups)	r   in_features_per_groupr   r   rl   r   r   r   r   s	           r!   r   zFP8GroupedLinear.__init__Z  s2     	-%!/ 	 	
 !r    c                $   |j                   d d }|j                   d   }| j                  j                         dkD  r| j                  j                  | j                  d|      j                  dd      }|j                  d| j                  |      j                  dd      }t        j                  ||      j                  dd      } |j                  g || j                  d }| j                  r5|j                  | j                  j                  | j                  d             |S t        | j                        }t        | j                        }|j                  | j                  d|      }|j                  dd      j                  d|      }|j                  | j                  |j                  d      | j                  z  |j                  d            }|j                  d      | j                  z  }t        j                   | j                  f||j"                  t        j$                        }t        j&                  d| j                  dz   |j"                  t        j$                        |z  }	t)               }
|
j+                  ||||	|| j,                        } |j                  | j                  g|d j                  dd      }| j                  r5|j                  | j                  j                  | j                  d             |S )Nr   r   r   )ry   rN   offsetstokens_per_expertrl   )shaper`   r   viewr   	transposereshaper   bmmr   ri   rm   r   rk   movedimsizefullry   int32arangerG   r1   rl   )r   xinput_shape
hidden_dimwyr   tokens_per_groupr   r   ro   s              r!   r   zFP8GroupedLinear.forwardn  se   ggcrlWWR[
;;##%)  J?II!QOA		"dmmZ8BB1aHA		!Q))!Q/A		:;::r:A}}tyy~~dmmR89HT[[!T223	FF4=="j1IIb!$$R4NN4==)..2Ct}}2TV_VdVdefVgh	66!95!JJ'79IRSRZRZbgbmbmn,,q$--!"3AHHEKKX[kk57**/ + 
 AIIdmm6k626>>q"E==FF499>>$--45r    r   )r   r   r   r   r   r   rl   r   r   r   r   r   r   r   )r   r   r   r   )r'   r2   r3   r4   r   r   r   r   s   @r!   r   r   N  sd    	  .2!* !"! ! 	!
 +! ! ! !($r    r   c                6   | j                   dk(  rt        d      t               }|j                  d      }|j                  d      }|j                  d      }|j	                  |d      }|j                  d      }	|j                  d      }
|
| j                  k\  j                  d      }t        | j                  r| j                  n| j                        }t        | j                  r| j                  n| j                        }t        | j                        }t        | j                        }|j!                  |||| j"                  |
      }| j                  r| j%                  |      }n| j'                  |      }|j!                  |||| j"                  |
      }||	j)                  |j*                        j                  d      z  }|j-                  |d       |j/                  |||      j1                  d      }|j)                  |j*                        S )	Nr   zbatched_mm experts dispatch does not support activation_scheme='static'. Use the default eager dispatch or switch to activation_scheme='dynamic'.r   r   dim)rl   
expert_ids        r   )r   NotImplementedErrorrG   r   repeat_interleaver   rV   	unsqueezer   has_gategate_up_projup_projgate_up_proj_scale_invup_proj_scale_inv	down_projdown_proj_scale_invr0   rl   _apply_gateact_fntorN   masked_fill_r   sum)r   hidden_statestop_k_indextop_k_weightsro   	num_top_k
num_tokensr   selected_hidden_statessample_weightsr   sentinel_mask	weight_upweight_scale_upweight_downweight_scale_downrW   weighted_outfinal_hidden_statess                      r!   fp8_batched_mm_experts_forwardr     s     )!W
 	

 23O  $I##A&J##B'J +<<YA<N"**2.N$$R(J
  4#3#33>>rBMdmm**NIdmmt::QUQgQghO4>>*K !9!9: --?? . H }}##H- ;;x( --?? . H n//?II"MML mS1 '++J	:NRRWXRY!!-"5"566r    c                   | j                   dk(  rt        d      t               }|j                  }|j	                  d      }|j	                  d      }|j	                  d      }|j                  d      }	|j                  d      }
t        j                  |
      \  }}|||z     }|	|   }|j                  dk(  r|j                         n|j                         }t        j                  || j                  d| j                  dz
        }t        j                  |dt        j                        }|| j                  k\  j                  d      }t!        | j"                  r| j$                  n| j&                        }t!        | j"                  r| j(                  n| j*                        }t!        | j,                        }t!        | j.                        }|j1                  |||||| j2                  	      }| j"                  r| j5                  |      }n| j7                  |      }|j1                  |||||| j2                  	      }||j9                  |j:                        j                  d      z  }|j=                  |d
       t        j>                  |      }t        j@                  |j	                  d      |      ||<   ||   }|jC                  |||      jE                  d      }|j9                  |j:                        S )Nr   zgrouped_mm experts dispatch does not support activation_scheme='static'. Use the default eager dispatch or switch to activation_scheme='dynamic'.r   r   cpur   )binsminrU   )r   rN   r   r   )ry   r   )#r   r   rG   ry   r   r   r   sortr&   r   r   histcrV   cumsumr   r   r   r   r   r   r   r   r   r   r1   rl   r   r   r   rN   r   
empty_liker   r   r   )r   r   r   r   ro   ry   r   r   r   r   r   expert_ids_gpermselected_hidden_states_gsample_weights_ghistc_inputr   r   r   r   r   r   r   rW   r   inv_permr   s                              r!   fp8_grouped_mm_experts_forwardr     s    )!W
 	

 23O!!F  $I##A&J##B'J #**2.N$$R(J J/L$,TY->?%d+
 +1++*>,$$&LDTDTDVKKd6F6FASWScScfgSghll,!5;;GG "T%5%55@@DMdmm**NIdmmt::QUQgQghO4>>*K !9!9: -- +?? . H }}##H- ;;x( --+?? . H .11(..AKKBOOL mS1 %H\\$))A,v>HTN)L '++J	:NRRWXRY!!-"5"566r    c                       e Zd ZU ddddiZded<   	 	 	 	 	 d	 	 	 	 	 	 	 	 	 d fdZddZ	 	 	 	 	 	 	 	 dd	Z	 d	 	 	 	 	 	 	 	 	 dd
Z xZ	S )
FP8Expertsdeepgemm_megamoemegamoe_expertsmegamoe_router)moe_tp_experts	ep_routerzdict[str, dict[str, str]]_impl_tp_layer_overridesc                j   t         
|           |du sJ d       || _        || _        || _        || _        |j                  | _        || _        t        |dd      | _
        t        |dd      | _        t        |dd       | _        t        |dd       | _        t        t        |d	d
         | _        t        |dd       | _        t        |dd      dk(  }|dk(  r
t%               nt&        j(                  }|rt&        j*                  |dddd}	nt,        |||d   nd ||d   nd d}	| j                  rQt/        | j                  d| j                  z  | j                  fddi|	\  | _        | _        | j5                  dd        nKt/        | j                  | j                  | j                  fi |	\  | _        | _        | j5                  dd        t/        | j                  | j                  | j                  fi |	\  | _        | _        | j5                  dd        | j                  dk(  rt?        j@                  t'        jB                  | j                  t&        j(                              | _"        t?        j@                  t'        jB                  | j                  t&        j(                              | _#        y y )NFzWFP8Experts does not support bias for now, please open an issue if you want this featurenum_local_expertsrV   moe_intermediate_sizeintermediate_sizeswiglu_alphaswiglu_limithidden_activation
hidden_actexpert_dtypefp8fp4r   r   r       )rY   rZ   r[   r\   r]   r   )rY   rZ   r\   r]   r^   gate_up_proj_biasup_proj_biasdown_proj_biasr   rM   )$r   r   configr   r   rl   hidden_sizer   r   r+   rV   intermediate_dimr$   r  r  r   r   limitr"   r   r   r|   r   re   r   r   r   r   r   r   r   rR   rS   onesgate_up_proj_activation_scaledown_proj_activation_scale)r   r  rl   r   r   r   r   is_fp4rZ   alloc_kwargsr   s             r!   r   zFP8Experts.__init__L  s    	5  	
e	
    $ ,,!2&v/BMR +F4KM` a#FNDA#FNDA[1DlSTV^T:
 75@)2g)=#%5== %

$ !L !+$.8.DZ]$.8.DZ]$	L ===O  !d&;&;";T__>YZ>^j>:Dt: ##$7>3E  $"7"74LX40DL$0 ##ND93Edoot/D/D4
HT4
00 	 0$7!!X-13ejjIYIYafanan>o1pD..0ll5::dFVFV^c^k^k;l.mD+ .r    c                   |j                  dd      \  }}| j                  q|j                  | j                        }|j                  | j                   | j                        }|t	        j
                  || j                  z        z  }|dz   |z  S | j                  D|j                  | j                        }|j                  | j                   | j                        }| j                  |      |z  S )Nr   r   r   )rU   r   rU   r   )chunkr  clampr  r   sigmoidr  r   )r   gate_upgateupglus        r!   r   zFP8Experts._apply_gate  s    ===+b(::$"3"3:4Dt000d6G6GHBtd.?.?'?@@CH##ZZ#::$**:-Dtzzktzz:B{{4 2%%r    c                   t        j                  |t         j                        }t        j                         5  t         j                  j
                  j                  || j                  dz         }|j                  ddd      }t        j                  |j                  d      d      j                  d	      j                  d
      }d d d        D ]  }|| j                  k(  rt        j                  |         \  }}	||	   }
| j                  dk(  r| j                  |   nd }| j!                  |
| j"                  r| j$                  |   n| j&                  |   | j"                  r| j(                  |   n| j*                  |   |      }| j"                  r| j-                  |      n| j/                  |      }| j                  dk(  r| j0                  |   nd }| j!                  || j2                  |   | j4                  |   |      }||	|d f   }||j7                  |j8                        z  }|j;                  d|	|j7                  |j8                                |j7                  |j8                        S # 1 sw Y   xY w)NrM   r   )num_classesr   r   )r   r   r   F)as_tupler   r   rg   )r   
zeros_liker   no_gradrR   r   one_hotrV   permutegreaterr   nonzeror   wherer   r  r   r   r   r   r   r   r   r   r  r   r   r   rN   
index_add_)r   r   r   r   r   expert_mask
expert_hit
expert_idx	top_k_pos	token_idxcurrent_stategate_up_act_scalerW   down_act_scalerouting_weightsr   s                   r!   r   zFP8Experts.forward  sK   
 $..}EMMR]]_((--55ktO_O_bcOc5dK%--aA6K{8'DaHPPZ_P`eefhiJ 
 %JT---#(;;{:/F#G Iy))4MBFBXBX\dBd22:>jn  {{15!!*-DLLQ[D\;?==++J7dNdNdeoNp!2	 # H 6:]]t''1T\H]H?C?U?UYa?a//
;gk  {{z*((4!/	 # H ,Iy$,FGO#o&8&8&HHL**1iI\IbIb9cd7 %8 #%%m&9&9::C _s   BI;;Jc                    |j                         dkD  rt        j                  ||d       S t        |||| j                  ||j
                        S )Nr   )rh   rn   )r   r   r   r   rl   rN   )r   rj   r`   rk   rh   s        r!   r   zFP8Experts.linear  sO      1$88E6400OO-
 	
r    )Nr   r   FT)
rl   r   r   r   r   r   r   r   r   r   )r  r   r   r   )r   r   r   r   r   r   r   r   rC   )
rj   r   r`   r   rk   r   rh   torch.Tensor | Noner   r   )
r'   r2   r3   r   r5   r   r   r   r   r   r   s   @r!   r   r   =  s     	/)
;7  .2!* Dn +Dn 	Dn
 Dn Dn DnL&(;)(;8D(;Ua(;	(;^ 15

 
 '	

 .
 

r    r   c                      e Zd ZdZeeeedZy)FP8ExpertsInterfacez?Interface for registering custom FP8 experts forward functions.)
batched_mm
grouped_mmdeepgemmr   N)	r'   r2   r3   r4   r   r   r   r   _global_mappingr   r    r!   r9  r9    s    I 544D	Or    r9  c                <   |j                   r| S d}| j                         D ]  \  }}t        ||      sd}t        j                  d      5  |j                  d      rt        |dd      }t        |dd      }	t        |d| j                  j                               }
t        t        t        |	|	      } ||
|j                  |j                  |j                  |	|
      }nt        |      t         j"                  u rPt%        |j&                  |j(                  |j                  |j                  |j                  |j*                  du      }nt-        |t         j"                        rqdt        |      j.                  v rZt1        |j&                  |j(                  |j2                  |j                  |j                  |j                  |j*                  du      }|| j5                  ||       d}ddd        |st6        j9                  d       | S # 1 sw Y   xY w)a  
    A helper function to replace all `torch.nn.Linear` modules by `FP8Linear` modules.

    Parameters:
        model (`torch.nn.Module`):
            Input model or `torch.nn.Module` as the function is run recursively.
        modules_to_not_convert (`list[`str`]`, *optional*, defaults to `None`):
            Names of the modules to not convert. In practice we keep the `lm_head` in full precision for numerical stability reasons.
        quantization_config (`FineGrainedFP8Config`):
            The quantization config object that contains the quantization parameters.
        pre_quantized (`book`, defaults to `False`):
            Whether the model is pre-quantized or not
    FNmetaz.expertsr   Tr   r  )experts_classexperts_interfacer   r   )r  rl   r   r   r   r   r   GroupedLinear)r   r   r   rl   r   r   r   zYou are loading your model using fp8 but no linear modules were found in your model. Please double check your model architecture.)
dequantizenamed_modulesr   r   ry   endswithr$   r  get_text_configr   r   ALL_FP8_EXPERTS_FUNCTIONSweight_block_sizer   r   r&   rR   Linearr   r   r   rm   
isinstancer'   r   r   set_submoduler   warning)modelmodules_to_not_convertquantization_configpre_quantizedhas_been_replacedmodule_namemodule
new_moduler   r   r  	new_classs               r!   replace_with_fp8_linearrV    s   " %%$224V$[2HI
\\&!##J/"6:t<"6:u= 5<<3O3O3QR6",&?%%		 '!2DD&9&K&K1;;%%
 f*& & 2 2!'!4!42DD&9&K&K1;;#[[4
 FBII.?d6lF[F[3[ .*0*<*<!'!4!4#__2DD&9&K&K1;;#[[4
 %##K<$(!_ "!  5l <	
 Lm "!s   F!HH	c                  @    e Zd ZdZd ZddZd	dZd
dZedd       Z	y)Fp8Quantizez^
    A quantization operation that creates two tensors, weight and scale out of a weight.
    c                    || _         y rC   hf_quantizerr   r[  s     r!   r   zFp8Quantize.__init__C  
    (r    c                ^   d }| j                   j                  kt        | j                   j                  t              r&| j                   j                  j	                  d      }n!t        | j                   j                  dd       }||j                  d   |j                  d   f}t        |      S )NrH  r   r   )r[  rO  rJ  dictr   r$   r   tuple)r   valuerl   s      r!   _resolve_block_sizezFp8Quantize._resolve_block_sizeF  s    
00<$++??F!..BBFFGZ[
$T%6%6%J%JL_aef
++b/5;;r?;JZ  r    c                   |j                   dk  r||iS | j                  |      \  }}|j                  d   |j                  d   }}||z  dk7  s||z  dk7  r||iS |j                  d d }||z  }||z  }	|j                  }
|j                  t        j
                        } |j                  g ||||	| }|j                         j                  d      }t	        j                  |dkD  |t	        j                  |            }t        |z  }t	        j                  |dkD  |t	        j                  |            }d|z  j                  t        j
                        }| j                  j                  j                  dk(  rt	        j                  d	t	        j                   t	        j"                  |j%                  t	        j&                  t        j
                        j(                  
                        }|j                  t+                     }d|j                  t        j
                        z  }|j-                  d      j-                  d      }||z  }t	        j$                  |t.        t              j                  t0              }|j                  |
      }|j3                  d      r|j5                  dd      d   dz   n|dz   }||||iS )Nr   r   r   r   )r   r   r   r          @)r   rd  r  .weight.r   .weight_scale_inv
_scale_inv)ndimrb  r   r   r   r   r   absamaxr+  	ones_like_FP8_MAXr[  rO  r   powceillog2r  finfotinyr"   r   _FP8_MINr   rE  rsplit)r   keyra  block_mblock_nrowscolsleading_shape
rows_tiles
cols_tilesoriginal_shape
value_fp32reshapedmax_abssafe_max_absscales
inv_scalesscales_broadcastscaled	quantized	scale_keys                        r!   _quantize_onezFp8Quantize._quantize_oneQ  sk    ::><33E:[[_ekk"od'>Q$.A"5< CR(W_
W_
XXemm,
%:%%_}_j_'_:_W^_,,.%%(%3{{7Q;9QRL(Wq[&%//&2IJFl&&u}}5
 00::gE3

5::j>N>NSXS^S^_d_l_lSmSrSr>N>s3t(uvJ#'7'9:J:==77F!++B/99"=,,KKH(CFFzR	%%n5	CF<<PYCZCJJsA&q),??`cfr`r	Y	:66r    c                    i }|j                         D ]=  \  }}t        |t              r|d   n|}|j                  | j	                  ||             ? |S )Nr   )itemsrJ  listupdater  )r   
input_dictkwargsresultrv  ra  r   s          r!   convertzFp8Quantize.converty  sT     +-$**,JC!+E4!8U1XeFMM$,,S&9: - r    c                ,    t        | j                        S rC   )Fp8Dequantizer[  r   s    r!   
reverse_opzFp8Quantize.reverse_op  s    T..//r    N)ra  r   r   ztuple[int, int])rv  r   ra  r   r   dict[str, torch.Tensor])r  r   r   r  r   r	   )
r'   r2   r3   r4   r   rb  r  r  propertyr  r   r    r!   rX  rX  >  s0    )	!&7P 0 0r    rX  c                  z    e Zd ZdZd ZddZdZddZ	 d	 	 	 	 	 	 	 ddZddZ		 	 d	 	 	 	 	 	 	 dd	Z
edd
       Zy)r  u  Dequantize FP8 weights using their per-block ``weight_scale_inv``.

    Designed to run as the *first* op in any :class:`WeightConverter` chain when
    loading with ``dequantize=True`` — :meth:`update_weight_conversions` on the
    FP8 quantizer attaches it to each existing model-specific converter so that
    per-expert (weight, scale) pairs are folded into full-precision tensors before
    the chain's merge / concat ops collapse the per-expert structure.

    Pattern semantics
        Input ``input_dict`` carries one entry per source pattern; each value is a
        list of tensors (one per ``*`` match). For every weight pattern that has a
        sibling ``*.weight_scale_inv`` pattern in the dict, this op pairs them up by
        index, dequantizes per-pair, and emits the dequantized list under the
        original *weight* key. Scale entries are dropped from the output so the
        remaining ops only see weights.
    c                    || _         y rC   rZ  r\  s     r!   r   zFp8Dequantize.__init__  r]  r    c                    |j                  d      }|r|d d n|}|j                  d      r|d t        d        dz   }n|dk(  rd}n|dz   }|r|dz   S |S )N$r   rf  rh  r`   rk   ri  )rE  len)r   weight_patternanchoredbasescales        r!   _scale_pattern_forz Fp8Dequantize._scale_pattern_for  so    !**3/&.~cr"N==#*C	N?+.AAEX&E<'E&us{1E1r    )r   g      ?r   g      ?re  g      @g      @g      @g       g      g      g      g       g      g      g      c                   t        j                  | j                  t         j                  |j                        }|j                         j                  t         j                        }|dz  j                         }|dz	  dz  j                         }t        j                  ||   ||   gd      } |j                  g |j                  dd d|j                  d   z   S )uR   Two ``e2m1`` FP4 values per byte → float32 tensor twice as wide on the last dim.)rN   ry         r   r   Nr   )r   r   _FP4_E2M1_LUTr   ry   
contiguousr   uint8longstackr   r   )r   packedlutu8lowhighunpackeds          r!   _unpack_fp4zFp8Dequantize._unpack_fp4  s    ll4--U]]6==Y %%ekk2CxooqC%%';;C#d)4"=xIcr!2IAR8H4HIIr    Nc                    t        t        dd       }|j                  t        j                  k(  s|!|j                  |k(  r| j	                  |      }n|j                  t        j                        }|j                  dd  \  }}	 |j                  dd  \  }}	||z  s||	z  rt        d| d| d| d|	 d	      ||z  }
||	z  }|E|j                  j                  r|j                         dk\  r|j                  nt        j                  }|j                  t        j                  k(  r1|j                  t        j                        d	z
  j                         }n|j                  t        j                        }|j                  }|j                  d
||
|	|      }|j                  d
||	      j!                  d
      j!                  d      }||z  j                  |      j                  |      S # t        $ r	 d\  }}	Y fw xY w)Nfloat4_e2m1fn_x2r   )r   r   zWeight shape (r;   z) not divisible by scale grid (z).r   g     _@r   )r$   r   rN   r|   r  r   r   r   	Exception
ValueErrorrT   r   bfloat16r  exp2r   r   )r   r  r  rn   	fp4_dtypequantized_fp32ry  rz  
scale_rows
scale_colsrw  rx  s_fp32r~  qss                   r!   _dequantize_onezFp8Dequantize._dequantize_one  s   
 E#5t<	??ejj(Y-ByZcGc!--i8N&\\%--8N#))"#.
d	*%+\\"#%6"J
 *z 1 b.Mj\Y[\f[ggij  *$*$
  & > >6CVCVCX\]C]chcqcq  <<5;;&ii.6<<>FYYu}}-F'--""2z7JPNN2z:6@@DNNqQAzz,'//??7  	*%)"J
	*s   G+ +G=<G=c                ^    ||y t        ||      \  }}t        ||d       }t        |dd       S )NrN   )r
   r$   )r   rM  full_layer_namerS  tensor_nameparams         r!   _get_target_dtypezFp8Dequantize._get_target_dtype  s=    =O325/JT2ugt,,r    c                   | j                  ||      }d|v r\||nd}|d   }t        |t              r|d   n|}d|v r2|d   }t        |t              r|d   n|}|| j                  |||      iS ||iS i }	|j	                         D ]  \  }
}d|
v sd|
v r| j                  |
      }||vr||	|
<   *t        |t              r|n|g}||   }t        |t              r|n|g}t        |      t        |      k7  r't        d|
 dt        |       d	t        |       d
      t        ||      D cg c]  \  }}| j                  |||       c}}|	|
<    |	S c c}}w )Nzweight$r`   r   rk   )rn   rh   z/Fp8Dequantize: weight/scale count mismatch for z (z weights vs z	 scales).)	r  rJ  r  r  r  r  r  r  zip)r   r  r  rM  r  rn   
target_keyr  r  r  rv  ra  r  weightsr   r  s                   r!   r  zFp8Dequantize.convert  s    --e_E
 
" -<,GXJ"9-I(29d(C	!I!Z/#$67&0&>F"D$8$8FYe$8$fgg	** @B$**,JC!S(,>#,E//4I
*#s)%6eUGG	*F)&$7VfXF7|s6{* EcU KG~\#f+iI  ^aahjp]qr]qUYUVXY4//1</P]qrF3K! -"  ss   +Ec                ,    t        | j                        S rC   )rX  r[  r  s    r!   r  zFp8Dequantize.reverse_op  s    
 4,,--r    )r  r   r   r   )r  r   r   r   rC   )r  r   r  r   rn   torch.dtype | Noner   r   )rM  torch.nn.Module | Noner  
str | Noner   r  )NN)r  ,dict[str, list[torch.Tensor] | torch.Tensor]r  r  rM  r  r   r  r  )r'   r2   r3   r4   r   r  r  r  r  r  r  r  r  r   r    r!   r  r    s    ")
2 mMJ ae+@%+@/;+@K]+@	+@Z- '+(,	,@, $, &	, 
6,\ . .r    r  )r   torch.dtype)r   r.   )r   None)rI   r   rJ   r   r   r   )r   NNr   )rV   r   rW   r   rX   r   rY   r  rZ   r  r[   r   r\   
int | Noner]   r  r^   r   r   z!tuple[nn.Parameter, nn.Parameter])NNNN)rj   r   r`   r   rk   r   rl   zlist[int] | Nonerm   r7  rh   r7  rn   r  r   r   )
r   ztorch.nn.Moduler   r   r   r   r   r   r   r   )NNF)rN  zlist[str] | None)I
__future__r   	functoolsr}   collections.abcr   dataclassesr   r   torch.nnrR   r   r   activationsr   core_model_loadingr	   quantizers.quantizers_utilsr
   r   utilsr   utils.import_utilsr   r   r   r   r<  r   r   r   hub_kernelsr   moer   r   tensor_parallelr   
get_loggerr'   r   float8_e4m3fnr   rr  r   rt  rU   rn  cacher"   r+   r.   rA   _dynamoallow_in_graphrE   rG   rK   re   rq   r   rI  r   r   r   r   Moduler   r9  rG  rV  rX  r  r   r    r!   <module>r     sM   #  	 $ !   $   . U   
 * = % 
		H	%   
5;;z"&&5;;z"&&    G $   . .b  
*     	
      '@ $( $,0'+ # !	
  * % B $( $,0'+BuBuBu #Bu !	Bu
 Bu *Bu %Bu BuJ6
		 6
rDy DNG7
G7G7 G7  	G7
 G7T[7
[7[7 [7  	[7
 [7|]
 ]
@*  01  ejP#3PfG0- G0TV.M V.r    