
    ^j                        d dl mZ d dlmZmZ d dlZd dlmc 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 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% ddl&m'Z'm(Z( ddl)m*Z*m+Z+ ddl,m-Z- ddl.m/Z/m0Z0m1Z1m2Z2m3Z3 ddl4m5Z5m6Z6 ddl7m8Z8 ddl9m:Z: ddl;m<Z<  e3jz                  e>      Z?d Z@ ed      dRd       ZAdej                  deCdej                  fdZD	 dSd ej                  d!ej                  d"ej                  d#ej                  d$ej                  dz  d%eFd&eFd'e-e/   fd(ZG eeA       G d) d*ej                               ZHd+ej                  d,eCfd-ZId. ZJd/ ZKd0 ZL G d1 d2ej                        ZM G d3 d4ej                  j                        ZN G d5 d6ej                        ZO G d7 d8ej                        ZP G d9 d:ej                        ZQe G d; d<ej                               ZR G d= d>ej                        ZS G d? d@edAB      ZT edC       G dD dEej                               ZU G dF dGe!      ZVe0 G dH dIe+             ZWe0 G dJ dKeW             ZX	 	 	 dTdLej                  eYej                     z  dz  dMeCdz  d$ej                  dz  dej                  eCz  fdNZZe0 G dO dPeWe             Z[g dQZ\y)U    )Callable)Optional	TypedDictN)nn   )initialization)ACT2FN)CacheDynamicCache)GenerationMixin)use_experts_implementationuse_kernel_forward_from_hubuse_kernel_func_from_hubuse_kernelized_func)force_accelerate_hooks)lazy_load_kernel)create_causal_maskcreate_recurrent_attention_mask)GradientCheckpointingLayer)BaseModelOutputWithPastMoeCausalLMOutputWithPastMoeModelOutputWithPast)ROPE_INIT_FUNCTIONSdynamic_rope_update)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)TransformersKwargsauto_docstringcan_return_tupleis_torchdynamo_compilinglogging)maybe_autocastmerge_with_config_defaults)resolve_internal_import)capture_outputs   )GraniteMoeHybridConfigc                     | dd| j                   d   dz  f   }| d| j                   d   dz  df   }t        j                  | |fd      S )z*Rotates half the hidden dims of the input..N   dim)shapetorchcat)xx1x2s      /var/www/ramen.bs-engineer-server.com/venv/lib/python3.12/site-packages/transformers/models/granitemoehybrid/modeling_granitemoehybrid.pyrotate_halfr5   8   sZ    	
3"!''"+"""	#B	
3q ""	#B99rc2YB''    rotary_pos_embc                     |j                  |      }|j                  |      }| |z  t        |       |z  z   }||z  t        |      |z  z   }||fS )a  Applies Rotary Position Embedding to the query and key tensors.

    Args:
        q (`torch.Tensor`): The query tensor.
        k (`torch.Tensor`): The key tensor.
        cos (`torch.Tensor`): The cosine part of the rotary embedding.
        sin (`torch.Tensor`): The sine part of the rotary embedding.
        unsqueeze_dim (`int`, *optional*, defaults to 1):
            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
    Returns:
        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
    )	unsqueezer5   )qkcossinunsqueeze_dimq_embedk_embeds          r4   apply_rotary_pos_embrA   ?   sY    & --
&C
--
&C3w;q>C/0G3w;q>C/0GGr6   hidden_statesn_repreturnc                     | j                   \  }}}}|dk(  r| S | dddddddddf   j                  |||||      } | j                  |||z  ||      S )z
    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
    r'   N)r.   expandreshape)rB   rC   batchnum_key_value_headsslenhead_dims         r4   	repeat_kvrL   Y   so    
 2?1D1D.Ehz!!Qa"23::5BUW\^bdlmM  (;e(CT8TTr6   modulequerykeyvalueattention_maskscalingdropoutkwargsc                    t        || j                        }t        || j                        }	t        j                  ||j	                  dd            |z  }
||
|z   }
t
        j                  j                  |
dt        j                        j                  |j                        }
t
        j                  j                  |
|| j                        }
t        j                  |
|	      }|j	                  dd      j                         }||
fS )Nr+   r   r*   )r-   dtype)ptrainingr'   )rL   num_key_value_groupsr/   matmul	transposer   
functionalsoftmaxfloat32torV   rS   rX   
contiguous)rM   rN   rO   rP   rQ   rR   rS   rT   
key_statesvalue_statesattn_weightsattn_outputs               r4   eager_attention_forwardre   e   s     3 ; ;<JUF$?$?@L<<z';';Aq'ABWLL!#n4==((2U]](SVVW\WbWbcL==((6??([L,,|\:K''1-88:K$$r6   c                        e Zd ZdZdedef fdZ	 	 ddej                  dej                  dz  de	dz  d	e
ej                  ej                  f   dz  d
ee   de
ej                  ej                  f   fdZ xZS )GraniteMoeHybridAttentionu   Hybrid variant that handles ``position_embeddings is None`` — granitemoe-hybrid configs can
    opt out of RoPE via ``position_embedding_type=None``, in which case the model passes ``None``
    instead of a ``(cos, sin)`` tuple.config	layer_idxc                 ^   t         |           || _        || _        t	        |d|j
                  |j                  z        | _        |j                  |j                  z  | _	        |j                  | _        |j                  | _        d| _        t        j                  |j
                  |j                  | j                  z  |j                         | _        t        j                  |j
                  |j                  | j                  z  |j                         | _        t        j                  |j
                  |j                  | j                  z  |j                         | _        t        j                  |j                  | j                  z  |j
                  |j                         | _        y )NrK   Tbias)super__init__rh   ri   getattrhidden_sizenum_attention_headsrK   rI   rY   attention_multiplierrR   attention_dropout	is_causalr   Linearattention_biasq_projk_projv_projo_projselfrh   ri   	__class__s      r4   rn   z"GraniteMoeHybridAttention.__init__   sJ   "
F4F4F&JdJd4de$*$>$>&B\B\$\!22!'!9!9ii : :T]] JQWQfQf
 ii : :T]] JQWQfQf
 ii : :T]] JQWQfQf
 ii&&68J8JQWQfQf
r6   NrB   rQ   past_key_valuesposition_embeddingsrT   rD   c                    |j                   d d }g |d| j                  }| j                  |      j                  |      j	                  dd      }| j                  |      j                  |      j	                  dd      }	| j                  |      j                  |      j	                  dd      }
||\  }}t        ||	||      \  }}	| |j                  |	|
| j                        \  }	}
t        j                  | j                  j                  t              } || ||	|
|f| j                  sdn| j                   | j"                  d|\  }} |j$                  g |d j'                         }| j)                  |      }||fS )Nr*   r'   r+           )rS   rR   )r.   rK   rw   viewr[   rx   ry   rA   updateri   r   get_interfacerh   _attn_implementationre   rX   rs   rR   rG   r`   rz   )r|   rB   rQ   r~   r   rT   input_shapehidden_shapequery_statesra   rb   r<   r=   attention_interfacerd   rc   s                   r4   forwardz!GraniteMoeHybridAttention.forward   s    $))#2.88b8$--8{{=166|DNNqRST[[/44\BLLQPQR
{{=166|DNNqRST**HC';L*VY[^'_$L*&'6'='=j,X\XfXf'g$J(?(M(MKK,,.E)
 %8	%
  $}}C$2H2HLL	%
 	%
!\ *k));;;;FFHkk+.L((r6   NN)__name__
__module____qualname____doc__r(   intrn   r/   Tensorr
   tupler   r   r   __classcell__r}   s   @r4   rg   rg   ~   s    *
5 
# 
6 )-HL%)||%) t+%) 	%)
 #5<<#=>E%) +,%) 
u||U\\)	*%)r6   rg   input_tensorpad_sizec                     t        | j                        dk(  r
ddddd|ddfnddd|ddf}t        j                  j                  j                  | |dd      S )z
    Padding x tensor with `pad_size` on the seq_len dim (dim=1)

    Assumes that we only have tensors of either size 4 or 3
       r   constant)moderP   )lenr.   r/   r   r\   pad)r   r   	pad_shapes      r4   pad_tensor_by_sizer      sf     47|7I7I3Ja3OAq!Q!Q/VWYZ\]_gijlmUnI88""<ST"UUr6   c                    t        | |      } t        | j                        dk(  r.| j                  | j                  d   d|| j                  d         S | j                  | j                  d   d|| j                  d   | j                  d         S )z
    Padding input_tensor with `pad_size` on the seq_len dim (dim=1) and
    simultaneously splitting it into chunk sequences.

    Assumes that we only have tensors of either size 4 or 3
    r   r   r*   r+   )r   r   r.   rG   )r   r   
chunk_sizes      r4   reshape_into_chunksr      s     &lH=L
<!###L$6$6q$92z<K]K]^_K`aa ##q!2z<3E3Ea3H,J\J\]^J_
 	
r6   c                 "   | j                  d      } | d   j                  g | j                         | } t        j                  t        j                  ||| j
                  t        j                        d      }| j                  | d      } t        j                  | d      }t        j                  t        j                  ||| j
                  t        j                        d      }|j                  | t        j                         }|S )zo
    More stable segment sum calculation. Uses cumulative sums and masking instead of direct subtractions.
    r*   .NdevicerV   )diagonalr   r,   )
sizerF   r/   trilonesr   boolmasked_fillcumsuminf)r   r   masktensor_segsums       r4   segment_sumr      s     ""2&J 2<	*11S<3D3D3FS
SL::ejjZ@S@S[`[e[efqstD++TE15LLL26M ::ejjZ@S@S[`[e[efqrsD!--teeiiZ@Mr6   c                     |N|j                   d   dkD  r<|j                   d   dkD  r*| j                  }| |dddddf   z  j                  |      } | S )zm
    Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
    Nr'   r   )r.   rV   r_   )rB   rQ   rV   s      r4   apply_mask_to_padding_statesr      sa    
 !n&:&:1&=&AnFZFZ[\F]`aFa##&1d
)CCGGNr6   c            
       <    e Zd ZdZdedef fdZ	 	 	 ddej                  de	dz  dej                  dz  d	ej                  dz  fd
Z	 	 dde	dz  dej                  dz  fdZ ed      	 	 	 dde	dz  dej                  dz  d	ej                  dz  fd       Z xZS )GraniteMoeHybridMambaLayeruP  
    Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`.
    A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective)
    ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4,
    and is why Mamba is called **selective** state spaces)

    The are a few differences between this and Mamba2Mixer:
    - The variable use_precomputed_states is slightly different due to the hybrid cache structure
    - There's a few non-obvious bugs fixed with batching in the slow path that exist in main
    - Some extra variables that our layer doesn't need have been removed
    - We ported most of the refactors in https://github.com/huggingface/transformers/pull/35154, which is (as of Dec 18, 2024) unmerged
    rh   ri   c           	         t         |           |j                  | _        |j                  | _        |j
                  | _        |j                  | _        t        |j                  | j                  z        | _        || _        |j                  | _        |j                  | _        t"        |j                     | _        |j&                  | _        |j*                  | _        |j.                  | _        |j2                  | _        |j6                  | _        |j:                  | _        |j<                  | _        |j>                  | _        | j                  d| j0                  z  | j                  z  z   | _         tC        jD                  | j@                  | j@                  |j                  | j                  | j@                  | j                  dz
        | _#        | j                  | j@                  z   | j                  z   }tC        jH                  | j                  || j(                        | _%        tC        jL                  tO        jP                  | j                              | _)        tO        jT                  d| j                  dz         }tC        jL                  tO        jV                  |            | _,        t[        | j                  | j,                        | _.        tC        jL                  tO        jP                  | j                              | _/        tC        jH                  | j                  | j                  | j(                        | _0        tc        d      }te        |dd       a3te        |dd       a4tc        d	      }tk        |d
      a6tk        |d      a7tk        |d      a8ts        tl        tn        tp        th        tf        f      a:tt        stv        jy                  d       ntv        jy                  d       |jz                  |   | _>        y )Nr+   r'   )in_channelsout_channelsrl   kernel_sizegroupspaddingrk   epszcausal-conv1dcausal_conv1d_updatecausal_conv1d_fnz	mamba-ssmz8ops.triton.selective_state_update.selective_state_update)chained_pathz1ops.triton.ssd_combined.mamba_chunk_scan_combinedz8ops.triton.ssd_combined.mamba_split_conv1d_scan_combineda  The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)` is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and https://github.com/Dao-AILab/causal-conv1dzOThe fast path for GraniteMoeHybrid will be used when running the model on a GPU)?rm   rn   mamba_n_heads	num_headsrp   mamba_d_statessm_state_sizemamba_d_convconv_kernel_sizer   mamba_expandintermediate_sizeri   mamba_conv_biasuse_conv_bias
hidden_act
activationr	   actmamba_proj_biasuse_biasrms_norm_epslayer_norm_epsilonmamba_n_groupsn_groupsmamba_d_headrK   mamba_chunk_sizer   time_step_limittime_step_mintime_step_maxconv_dimr   Conv1dconv1dru   in_proj	Parameterr/   r   dt_biasarangelogA_logGraniteMoeHybridRMSNormGatednormDout_projr   ro   r   r   r%   selective_state_updatemamba_chunk_scan_combined mamba_split_conv1d_scan_combinedallis_fast_path_availableloggerwarning_oncelayer_types
layer_type)r|   rh   ri   projection_sizeAcausal_conv1d	mamba_ssmr}   s          r4   rn   z#GraniteMoeHybridMambaLayer.__init__  s,   --!--$22 & 3 3!$V%8%84;K;K%K!L"#33 ++&++,.."("5"5--++ 11%55#11#11..T]]1BTEXEX1XXii''--==))A-
 004==@4>>Qyy
 ||EJJt~~$>? LLDNNQ./\\%))A,/
01G1GTMdMde	ejj89		$"8"8$:J:JQUQ^Q^_ )9&}6LdS"=2DdK %[1	!8$^"
 %<$W%
! ,C$^,
(
 "%&)0 $"
 &>  qr ,,Y7r6   NrB   cache_paramsrQ   seq_idxc                    t        ||      }| j                  |      }|j                  \  }}}| j                  | j                  z  }	|d uxr |j                  | j                        }
|
rL|j                  | j                     j                  d   }|j                  | j                     j                  d   }|
r|dk(  r|j                  d      j                  | j                  | j                  | j                  gd      \  }}}t        || j                   j"                  j                  d      | j                   j$                  | j&                        }t)        j                  || j                  |	|	gd      \  }}}t)        j*                  | j,                  j/                                }|d d d df   d d d d d f   j1                  d| j2                  | j                        j5                  t(        j6                        }|d d d d d f   j1                  dd| j2                        }| j8                  d d d df   j1                  d| j2                        }| j:                  d d d df   j1                  d| j2                        }|j=                  || j                  |j                  d   | j                  z        }|j=                  || j                  |j                  d   | j                  z        }|j=                  || j                  | j2                        }t?        ||||||d |d
      }|j=                  || j                  | j2                  z        }| jA                  ||      }| jC                  |      d d d df   }|S t)        j*                  | j,                  j/                                }| jD                  d	t/        d
      fk(  ri nd| jD                  i}| jF                  r|tI        || j                   j"                  j                  d      | j                   j$                  | j8                  |f| j:                  | jJ                  || j&                  | j@                  j"                  | j@                  jL                  | jB                  j"                  | jB                  j$                  | j2                  | j                  ddd|}|S |j                  | j                  | j                  | j                  gd      \  }}}|jO                  dd      }|
rt)        jP                  |gd      }|XtR        jT                  jW                  || jX                  |j                  d   z
  df      }|j[                  || j                         | j&                  dvr5| j]                  | j!                  |      dd |j                  d   f         }nQt_        || j                   j"                  j                  d      | j                   j$                  | j&                  |      }|
r|d d d d | d f   }|jO                  dd      }t        ||      }t)        j                  || j                  |	|	gd      \  }}}ta        |j=                  ||d| j2                        |||j=                  ||| j                  d      |j=                  ||| j                  d      f| jJ                  | j:                  d |d| j8                  d|
rnd d|\  }}|||jc                  || j                        }|j=                  ||d      }| jA                  ||      }| jC                  |      }|S )Nr   r'   r*   r,   .rV   T)zr   dt_softplusr   r   dt_limitF)r   r   r   r   rmsnorm_weightrmsnorm_epsoutproj_weightoutproj_biasheaddimngroupsnorm_before_gatereturn_final_statesr+   )siluswish)r1   weightrl   r   r   )r   r   r   r   r  r   r   initial_states)2r   r   r.   r   r   has_previous_stateri   layersconv_statesrecurrent_statessqueezesplitr   r   r   r   r   r  rl   r   r/   expr   floatrF   rK   r_   r^   r   r   r   r   r   r   r   rX   r   r   variance_epsilonr[   r0   r   r\   r   r   update_conv_stater   r   r   update_recurrent_state)r|   rB   r   rQ   r   projected_states
batch_sizeseq_len_groups_time_state_sizeuse_precomputed_states
conv_staterecurrent_stategatehidden_states_B_CdtBCr   r   r   hidden_states_reshapedoutdt_limit_kwargsr  scan_output	ssm_states                              r4   cuda_kernels_forwardz/GraniteMoeHybridMambaLayer.cuda_kernels_forwardr  s    5]NS<<6 "/!4!4
GQ!%1D1D!D!-T!9!ml>]>]^b^l^l>m!%,,T^^<HHKJ*11$..ARRSTUO "gl*:*B*B1*E*K*K''GR +L +'D#R
 !5!""**1-  ! #(++!'')?AWX#M1a 4::++-..A!T3,1d
+222t}}dFYFYZ]]didqdq]rAAq$J&&r2t}}=Bll1dC<077DMMJGq$|$++B>Az4==!''!*2MNAz4==!''!*2MNA%2%7%7
DNNTXTaTa%b"2& M *..z4>>DMM;YZM IImT:M --.q$|<C@ 
{ 4::++-..A$($8$8S%,<O$ObV`bfbvbvUwO }}!56$KK&&..q1KK$$LL ff####'99#3#3 $		 : :#'==#7#7!%!3!3 MM MM%*(-#$ &%p 
E /?.D.D++T]]DNNKQS /E /+'
 %6$?$?1$E!) ).		:?P2QWY(Z%+"$--"3"3)..1B1H1H1LLaP#K !22;O??*;;(,=N1OPSUrWhWnWnoqWrUrPr1s(t%(8+#{{1199!<![[--#'?? ')% *(9!Q	/(J%$5$?$?1$E!$@ARTb$c!&+kk%++-CE[\'#q! *C!&&z7BNFF:wrBFF:wrB*  $ff#(, LL $6L?RV* &*&Y$ (\-E , C CIt~~ ^I)..z7BG"iiT: mmK0
r6   c                    |j                   \  }}}|j                  }t        ||      }| j                  |      }|j	                  | j
                  | j                  | j                  gd      \  }	}
}|
j                  dd      }
|d uxr |j                  | j                        }|r&|j                  | j                     j                  d   }|r|dk(  r|j                  |
| j                        d| j                   d f   }t        j                   || j"                  j$                  j'                  d      z  d      }
| j(                  r|
| j"                  j*                  z   }
| j-                  |
      }
n|rt        j.                  |
gd      }
|Xt0        j2                  j5                  |
| j                  |
j                   d   z
  df      }|j                  || j                         | j-                  | j#                  |
      dd |
j                   d   f         }
|r
|
d| d f   }
|
j                  dd      }
t        |
|      }
t        j                  |
| j
                  | j6                  | j8                  z  | j6                  | j8                  z  gd      \  }}}t        j:                  | j<                  j?                                }|r|dk(  r|j                  | j                     j@                  d   jB                  }|d d dd d f   d d d df   }|j                  dd      jE                  ||j                   d   | jF                        }| jH                  d   jE                  | jH                  j                   d   | jF                        }t        j0                  j2                  jK                  ||jM                  |j                        z         }t        jN                  || jP                  d   | jP                  d         }|d   jE                  | j                  | jF                  | j8                        jM                  t        jR                  	      }t        j:                  |d   |z        jM                  |
      }|jU                  || j6                  d      dd d d f   }|jE                  || j6                  | j                  | j6                  z  |j                   d         jW                         }|jU                  |d|j                   d         }|d   |dd d d f   z  }|jU                  |d| jF                        }||d   z  jM                  |
      }|j                  | j                     j@                  d   |z  |z   }|jY                  || j                        }|jU                  || j6                  d      dd d d f   }|jE                  || j6                  | j                  | j6                  z  |j                   d         jW                         }|jU                  |d|j                   d         }|jM                  |jB                  |j                        }|j[                  || j                  z  | jF                  | j8                        }|j[                  || j                  z  | j8                  d      }t        j\                  ||      }|j[                  || j                  | jF                        }| j^                  d   jE                  | j^                  j                   d   | jF                        }|||z  z   jM                  |j                        }|jU                  |d      d d d df   }nt0        j2                  jK                  || jH                  z         }t        jN                  || jP                  d   | jP                  d         }|jU                  ||d| jF                        j?                         }|jU                  ||d| j8                        j?                         }|jU                  ||d| j8                        j?                         }|ja                  | j                  | j6                  z  d| j                        }|ja                  | j                  | j6                  z  d| j                        }| jb                  || jb                  z  z
  | jb                  z  }| j^                  d   te        ||      z  }||d   z  }|jM                  |j                        |z  }||||fD cg c]  }tg        ||| jb                         c}\  }}}}|ji                  dddd      }t        jj                  |d      } t        j:                  tm        |            }!|d d d d d d d d d d d f   |d d d d d d d d d d d f   z  }"|"j!                  d      }#|#d   |!ji                  ddddd      d   z  }$|$j!                  d      }%|%d   |d d d d d f   z  j!                  d      }&t        j:                  | d d d d d d dd f   | z
        }'||'ji                  dddd      d   z  }(|(dd d d f   |d   z  j!                  d      })|rR|j                  | j                     j@                  d   d d d f   jM                  |)j                  |)jB                        nt        jn                  |)d d d df         }*t        j.                  |*|)gd      })t        j:                  tm        t0        j2                  j5                  | d d d d d d df   d                  }+|+j                  dd      }+|+d   |)d d d d d df   z  j!                  d      },|,d d d df   |,d d df   }-})t        j:                  |       }.|dd d d f   |)d d d d d df   z  }/|.ji                  dddd      }0|/j!                  d      |0d   z  }1|&|1z   }|jU                  |d| j                  | jF                        }||z   }|dkD  r|d d d |d d d d f   }|jU                  ||d      }|-||jY                  |-| j                        }-| jq                  ||	      }2| js                  |2jM                  |            }3|3S c c}w )Nr*   r,   r'   r+   r   .r   ).NNr   r   r   )r-   output_sizer   r   r   )rV   r   )r'   r   ):r.   rV   r   r   r  r   r   r   r[   r  ri   r  r  r  r   r/   sumr   r  r
  r   rl   r   r0   r   r\   r   r   r   r  r   r  r	  r   rF   rK   r   softplusr_   clampr   r^   rG   r`   r  r   bmmr   repeat_interleaver   r   r   permuter   r   
zeros_liker   r   )4r|   input_statesr   rQ   r  r  r  rV   r  r  r  r  r  r  r  rB   r  r  r   cache_devicer   dAdBdBx
ssm_statesssm_states_reshaped
C_reshapedyr   r   
D_residualtA_cumsumLG_intermediateGM_intermediateMY_diagdecay_statesB_decaystatesprevious_statesdecay_chunk
new_statesr"  state_decay_outC_times_statesstate_decay_out_permutedY_offr!  contextualized_statess4                                                       r4   torch_forwardz(GraniteMoeHybridMambaLayer.torch_forward  s    ".!3!3
GQ"" 4L.Q<<5&6&<&<''GR '= '
# .77!<!-T!9!ml>]>]^b^l^l>m!%,,T^^<HHKJ "gl&889JDNN[\_bfbwbwawax\xyK %		dkk0088;;! !!$58H8H$H! $): ;%$)IIz;L.MSU$V!' mm//%(=(=@Q@W@WXZ@[([]^'_ ..{DNNK $5F)GMjO`OfOfgiOjMjHj)k l%$5cG89n$E! 1 ; ;Aq A89JN[#kk##T]]T5H5H%H$--Z^ZmZmJmn
q! YYtzz'')**!gl'..t~~>OOPQRYYL Aq!GQc\*Ba#**:rxx|T]]SBll9-44T\\5G5G5JDMMZG$$--b7::bhh3G.GHBR!5!5a!8$:N:Nq:QRB/"))$..$--I\I\]``glgtgt`uA))ByMA-.22,2GB
 		*dmmR8dAFAT]]DNNdmm4SUVU\U\]_U`allnA		*b!''"+6AI3a<0B *11*b$--PMi0044L4IC &,,T^^<MMaPSUUX[[J%<<ZXJ 		*dmmR8dAFAT]]DNNdmm4SUVU\U\]_U`allnA		*b!''"+6A $ahhaggFJ",//*t~~2Mt}}^b^q^q"r
T^^ ;T=P=PRSTJ		-z:Az4>>4==AA y!((a$--HA]Q&&**1773A 		*b)!T3,7A ''T\\(9:BR!5!5a!8$:N:Nq:QRB)11*gr4==Y__aM		*gr43F3FGMMOA		*gr43F3FGMMOA##DNNdmm$CX\XfXf#gA##DNNdmm$CX\XfXf#gA'DOO*CCtVH	*-?x-XXJ *ByM9M](()B.A cpqrtuwxay%z\]&9!Xt&W%z"M1a 		!Q1%A||A2.H 		+a.)A q!Qa23a1dAq!8K6LLN""r"*A y\AIIaAq!,DY,OON""r"*A 	l]1a:%>>CCCJF !99XaArsl%;h%FGL,..q"b!<YGGGc4l+mI.FFKKPQKRF * ##DNN3DDQG4PSSZ`ZfZfouo|o|S}%%fQUm4 
 YY8a@F))K0A0A(1aQRTV;BWY_0`$abK%//15K%o61dC9PPUUZ[U\J *1crc6 2Jq"u4EIF $ii1OT1oq!T30GGN'6'>'>q!Q'J$#''+.Fy.QQE A		*b$..$--HAJA!|a'1a'(		*gr2A $)A(??	4>>Z	ii4(
 !%knnU.C D$$I &{s   <r?r   c                    t         rJd| j                  j                  j                  j                  v rt               s| j                  ||||      S |t        d      |j                  }|B|j                  d   dkD  r0|j                  d   dkD  r||d d d d d f   z  j                  |      }| j                  |||      S )Ncudaz\`seq_idx` support requires fast path support. Please install `mamba_ssm` and `causal_conv1d`r'   r   )r   r   r  r   typer!   r#  NotImplementedErrorrV   r.   r_   rK  )r|   rB   r   rQ   r   rT   rV   s          r4   r   z"GraniteMoeHybridMambaLayer.forward  s     "f0C0C0J0J0O0O&OXpXr,,]L.Zabb%n  ##%.*>*>q*AA*E.J^J^_`JadeJe*^Aq$J-GGKKERM!!-~NNr6   NNNr   )r   r   r   r   r(   r   rn   r/   r   r
   	IntTensorr#  rK  r   r   r   r   s   @r4   r   r     s   \85 \8# \8B &*.2*.b||b dlb t+	b
 4'bP &*.2	D% dlD% t+	D%N H% &*.2*.O dlO t+	O
 4'O &Or6   r   c                   (     e Zd Zd fd	ZddZ xZS )r   c                     t         |           t        j                  t	        j
                  |            | _        || _        y Nrm   rn   r   r   r/   r   r  r  r|   rp   r   r}   s      r4   rn   z%GraniteMoeHybridRMSNormGated.__init__  s/    ll5::k#:; #r6   c                    |j                   }|j                  t        j                        }|?|t        j
                  j                  |j                  t        j                              z  }|j                  d      j                  dd      }|t        j                  || j                  z         z  }| j                  |j                  |      z  S Nr+   r*   T)keepdim)rV   r_   r/   r^   r   r\   r  powmeanrsqrtr  r  )r|   rB   r  input_dtypevariances        r4   r   z$GraniteMoeHybridRMSNormGated.forward  s    #))%((7)BMM,>,>twwu}}?U,VVM $$Q',,R,>%Ht?T?T4T(UU{{]--k:::r6   gư>rT  )r   r   r   rn   r   r   r   s   @r4   r   r     s    $
	;r6   r   c                   `     e Zd ZdZdef fdZdej                  dej                  fdZ xZ	S )GraniteMoeHybridMLPz~
    MLP layer for shared experts

    Args:
        config:
            Configuration object with model hyperparameters.
    rh   c                 `   t         |           |j                  | _        |j                  | _        t
        |j                     | _        t        j                  | j                  | j                  dz  d      | _
        t        j                  | j                  | j                  d      | _        y )Nr+   Frk   )rm   rn   rp   
input_sizeshared_intermediate_sizer	   r   r   r   ru   input_linearoutput_linearr|   rh   r}   s     r4   rn   zGraniteMoeHybridMLP.__init__  s     ,,!:: !2!23IIdoot7G7G!7KRWXYYt'7'7uUr6   rB   rD   c                     | j                  |      }|j                  dd      }| j                  |d         |d   z  }| j                  |      }|S )Nr+   r*   r,   r   r'   )re  chunkr   rf  )r|   rB   chunked_hidden_statess      r4   r   zGraniteMoeHybridMLP.forward  s^    ))-8 - 3 3A2 3 >(=a(@ADYZ[D\\**=9r6   
r   r   r   r   r(   rn   r/   r   r   r   r   s   @r4   ra  ra    s2    V5 VU\\ ell r6   ra  c                        e Zd ZU ej                  ed<   ddef fdZe	 	 	 ddedz  de	d   de
dz  ded	ef   fd
       Z ej                         ed               Z xZS )GraniteMoeHybridRotaryEmbeddinginv_freqNrh   c                    t         |           |j                  | _        |j                  | _        || _        | j
                  j                  d   | _        | j                  }| j                  dk7  rt        | j                     } || j
                  |      \  }| _
        | j                  d|d       | j                  d|j                         d       y )N	rope_typedefaultrn  F)
persistentoriginal_inv_freq)rm   rn   max_position_embeddingsmax_seq_len_cachedoriginal_max_seq_lenrh   rope_parametersrp  compute_default_rope_parametersr   attention_scalingregister_bufferclone)r|   rh   r   rope_init_fnrn  r}   s        r4   rn   z(GraniteMoeHybridRotaryEmbedding.__init__$  s    "("@"@$*$B$B!44[A!%!E!E>>Y&.t~~>L+7V+L($(ZeD0(..2BuUr6   r   ztorch.devicer  rD   ztorch.Tensorc                    | j                   d   }t        | dd      xs | j                  | j                  z  }d}d|t	        j
                  d|dt        j                        j                  |t        j                        |z  z  z  }||fS )	a  
        Computes the inverse frequencies according to the original RoPE implementation
        Args:
            config ([`~transformers.PreTrainedConfig`]):
                The model configuration.
            device (`torch.device`):
                The device to use for initialization of the inverse frequencies.
            seq_len (`int`, *optional*):
                The current sequence length. Unused for this type of RoPE.
        Returns:
            Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
            post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
        
rope_thetarK   Ng      ?r   r+   r   r   )	rw  ro   rp   rq   r/   r   int64r_   r  )rh   r   r  baser-   attention_factorrn  s          r4   rx  z?GraniteMoeHybridRotaryEmbedding.compute_default_rope_parameters4  s    & %%l3fj$/c63E3EIcIc3c U\\!S!5;;?BB&X]XcXcBdgjjk
 )))r6   c                 N   | j                   d d d d f   j                         j                  |j                  d   dd      j	                  |j
                        }|d d d d d f   j                         }t        |j
                  j                  t              r/|j
                  j                  dk7  r|j
                  j                  nd}t        |d      5  |j                         |j                         z  j                  dd      }t        j                  ||fd	      }|j                         | j                  z  }|j                         | j                  z  }	d d d        j	                  |j                   
      	j	                  |j                   
      fS # 1 sw Y   AxY w)Nr   r*   r'   mpscpuF)device_typeenabledr+   r,   r   )rn  r  rF   r.   r_   r   
isinstancerN  strr#   r[   r/   r0   r<   ry  r=   rV   )
r|   r1   position_idsinv_freq_expandedposition_ids_expandedr  freqsembr<   r=   s
             r4   r   z'GraniteMoeHybridRotaryEmbedding.forwardR  sR    !MM$4-8>>@GGHZHZ[\H]_acdehhijiqiqr ,QaZ 8 > > @'1!((--'E!((--[`J`ahhmmfkUC 	5&,,.1F1L1L1NNYYZ[]^_E))UEN3C'')d444C'')d444C		5 vvAGGv$cff177f&;;;	5 	5s   BFF$rT  rP  )r   r   r   r/   r   __annotations__r(   rn   staticmethodr   r   r   r  rx  no_gradr   r   r   r   s   @r4   rm  rm  !  s    llV5 V  04+/"*&-*(* t* 
~u$	%	* *: U]]_<  <r6   rm  c                        e Zd ZdZdef fdZdej                  deej                  ej                  ej                  f   fdZ	 xZ
S )GraniteMoeHybridTopKRoutera  Top-k gating that returns the routing decisions without grouping tokens by expert.

    Returns ``(top_k_index, top_k_weights, router_logits)``; the grouping/scattering used to live
    here (via ``expert_size.tolist()``, which broke fullgraph compile) and now happens inside the
    experts forward via ``use_experts_implementation`` so the default ``grouped_mm`` / ``batched_mm``
    paths can compile cleanly.
    rh   c                     t         |           |j                  | _        |j                  | _        t        j                  t        j                  | j                  |j                              | _        y rT  )rm   rn   num_local_expertsnum_expertsnum_experts_per_toktop_kr   r   r/   emptyrp   r  rg  s     r4   rn   z#GraniteMoeHybridTopKRouter.__init__k  sP    !33//
ll5;;t/?/?ASAS#TUr6   rB   rD   c                     t        j                  || j                        j                         }|j	                  | j
                  d      \  }}t        j                  |d      j                  |      }|||fS )Nr*   r,   )	Flinearr  r  topkr  r/   r]   type_as)r|   rB   router_logitstop_k_logitstop_k_indextop_k_weightss         r4   r   z"GraniteMoeHybridTopKRouter.forwardq  se    <BBD$1$6$6tzzr$6$J!kl;CCMRM=88r6   )r   r   r   r   r(   rn   r/   r   r   r   r   r   s   @r4   r  r  b  sI    V5 V9U\\ 9eELL%,,X]XdXd<d6e 9r6   r  c                        e Zd ZdZdef fdZdej                  dej                  dej                  dej                  fdZ xZ	S )	GraniteMoeHybridExpertsz2Collection of expert weights stored as 3D tensors.rh   c                    t         |           |j                  | _        |j                  | _        |j                  | _        t        j                  t        j                  | j                  d| j                  z  | j
                              | _        t        j                  t        j                  | j                  | j
                  | j                              | _        t        |j                     | _        y )Nr+   )rm   rn   r  r  rp   
hidden_dimr   intermediate_dimr   r   r/   r  gate_up_proj	down_projr	   r   act_fnrg  s     r4   rn   z GraniteMoeHybridExperts.__init__|  s    !33 ,, & 8 8LLT5E5Eq4K`K`G`bfbqbq)rsekk$2B2BDOOUYUjUj&klV../r6   rB   r  r  rD   c                 f   t        j                  |      }t        j                         5  t         j                  j                  j                  || j                        }|j                  ddd      }t        j                  |j                  d      d      j                         }d d d        D ]  }|d   }|| j                  k(  rt        j                  |         \  }}	||	   }
t        j                  j                  |
| j                  |         j                  dd      \  }}| j                  |      |z  }t        j                  j                  || j                   |         }|||	|d f   z  }|j#                  d|	|j%                  |j&                                |S # 1 sw Y   xY w)N)num_classesr+   r'   r   )r*   r   r,   r*   )r/   r-  r  r   r\   one_hotr  r,  greaterr'  nonzerowherer  r  ri  r  r  
index_add_r_   rV   )r|   rB   r  r  final_hidden_statesexpert_mask
expert_hit
expert_idx	top_k_pos	token_idxcurrent_stater  upcurrent_hidden_statess                 r4   r   zGraniteMoeHybridExperts.forward  s    $..}=]]_ 	S((--55ktO_O_5`K%--aA6K{8'DaHPPRJ	S
 % 
	nJ#AJT---#(;;{:/F#G Iy))4M}}++M4;L;LZ;XY__`agi_jHD"$(KK$5$:!$&MM$8$89NPTP^P^_iPj$k!$9M)U^`dJd<e$e!**1i9N9Q9QReRkRk9lm
	n #"#	S 	Ss   A=F&&F0rk  r   s   @r4   r  r  x  sN    <05 0#||# \\# ||	#
 
#r6   r  c                   `     e Zd ZdZdef fdZdej                  dej                  fdZ xZ	S )GraniteMoeHybridMoEzISparsely-gated mixture-of-experts block: router decides, experts compute.rh   c                     t         |           |j                  | _        t	        |      | _        t        |      | _        y rT  )rm   rn   rp   rc  r  routerr  expertsrg  s     r4   rn   zGraniteMoeHybridMoE.__init__  s3     ,,08.v6r6   layer_inputrD   c                     |j                         \  }}}|j                  d|      }| j                  |      \  }}}| j                  |||      }	|	j	                  ||| j
                        S )Nr*   )r   rG   r  r  r   rc  )
r|   r  bszlengthemb_sizerB   r  r  r  layer_outputs
             r4   r   zGraniteMoeHybridMoE.forward  si     + 0 0 2VX#++B9(,M(B%]A||M;N  fdoo>>r6   rk  r   s   @r4   r  r    s.    S75 7?5<< ?ELL ?r6   r  c                       e Zd ZU dZej
                  ed<   ej
                  ed<   eed<   eed<   ej                  ed<   y)GraniteFlashAttentionKwargsaT  
    Keyword arguments for advanced Flash Attention, causal-conv1d, and mamba_ssm kernel usage.
    Use cases include padding-free training and fewer `torch.compile` graph breaks.

    cu_seq_lens_q (`torch.LongTensor`):
        Gets cumulative sequence length for query state.
    cu_seq_lens_k (`torch.LongTensor`):
        Gets cumulative sequence length for key state.
    max_length_q (`int`):
        Maximum sequence length for query state.
    max_length_k (`int`):
        Maximum sequence length for key state.
    seq_idx (`torch.IntTensor):
        Index of each packed sequence.
    cu_seq_lens_qcu_seq_lens_kmax_length_qmax_length_kr   N)	r   r   r   r   r/   
LongTensorr  r   rQ   r6   r4   r  r    s7      ######__r6   r  F)totalRMSNormc                   h     e Zd Zddeddf fdZdej                  dej                  fdZd Z xZ	S )	GraniteMoeHybridRMSNormr   rD   Nc                     t         |           t        j                  t	        j
                  |            | _        || _        y)zF
        GraniteMoeHybridRMSNorm is equivalent to T5LayerNorm
        NrU  rV  s      r4   rn   z GraniteMoeHybridRMSNorm.__init__  s1     	ll5::k#:; #r6   rB   c                 "   |j                   }|j                  t        j                        }|j	                  d      j                  dd      }|t        j                  || j                  z         z  }| j                  |j                  |      z  S rX  )	rV   r_   r/   r^   rZ  r[  r\  r  r  )r|   rB   r]  r^  s       r4   r   zGraniteMoeHybridRMSNorm.forward  sy    #))%((7 $$Q',,R,>%Ht?T?T4T(UU{{]--k:::r6   c                 ^    t        | j                  j                         d| j                   S )Nz, eps=)r   r  r.   r  )r|   s    r4   
extra_reprz"GraniteMoeHybridRMSNorm.extra_repr  s*    ))*+6$2G2G1HIIr6   r_  )
r   r   r   r  rn   r/   r   r   r  r   r   s   @r4   r  r    s7    $ $$ $;U\\ ;ell ;Jr6   r  c                   .    e Zd Zdedef fdZe	 	 	 	 ddej                  dej                  dz  de	dz  de
dz  d	eej                  ej                  f   dz  d
ee   deej                  eej                  ej                  f   dz  f   fd       Z xZS )GraniteMoeHybridDecoderLayerrh   ri   c                 8   t         |           |j                  | _        d | _        t	        |j                  |j
                        | _        t	        |j                  |j
                        | _        |j                  dkD  rt        |      nd | _
        |j                  | _        t        |      | _        d | _        |j                  |   dk(  rt!        ||      | _        nt#        ||      | _        |j                  |   | _        t'        |dd      dkD  | _        y )Nr   r   linear_attentionr  )rm   rn   rp   	self_attnr  r   input_layernormpost_attention_layernormr  r  block_sparse_moeresidual_multiplierra  
shared_mlpmambalayers_block_typer   rg   
block_typero   has_expertsr{   s      r4   rn   z%GraniteMoeHybridDecoderLayer.__init__  s    !--6v7I7IvObObc(?@R@RX^XkXk(l% @F?W?WZ[?[ 3F ;ae#)#=#= -f5
##I.2DD3FIFDJ6vyIDN 229= #6+>BQFr6   NrB   rQ   r~   	use_cacher   rT   rD   c           	         |}| j                  |      }| j                   | j                  d|||d|}n | j                  d|||||d|\  }}||| j                  z  z   }|}| j	                  |      }| j
                  r&| j                  |      }	|	| j                  |      z   }n| j                  |      }||| j                  z  z   }|S )N)rB   r   rQ   )rB   rQ   r~   r  r   r  )r  r  r  r  r  r  r  r  )
r|   rB   rQ   r~   r  r   rT   residualr  moe_hidden_statess
             r4   r   z$GraniteMoeHybridDecoderLayer.forward  s    !,,];::!&DJJ +,- 	M  .t~~  +- /#$7   M1 !=43K3K#KK 55mD $ 5 5m D-0NNM OOM:M =43K3K#KKr6   )NNFN)r   r   r   r(   r   rn   r   r/   r   r
   r   r   r   r  FloatTensorr   r   r   s   @r4   r  r    s    G5 G# G.  /3(,!&HL(||( t+( 	(
 $;( #5<<#=>E( 45( 
u  %(9(95;L;L(L"MPT"TT	U( (r6   r  c                        e Zd ZU eed<   dZdZdgZdgZdZ	dZ
dZdZdZeedZdZ ej&                          fd       Z xZS )GraniteMoeHybridPreTrainedModelrh   modelTr  r~   )rB   
attentionsc           
      P   t         |   |       t        |t              rmt	        j
                  |j                  d| j                  j                         t	        j
                  |j                  d| j                  j                         nFt        |t              r6t	        j
                  |j                  d| j                  j                         t        |t              rt	        j                  |j                         t	        j                  |j                   t#        j$                  t#        j&                  d|j(                  dz                      t	        j                  |j*                         y t        |t,              r t	        j                  |j                         y y )Nr   )r[  stdr'   )rm   _init_weightsr  r  initnormal_r  rh   initializer_ranger  r  r  r   ones_r   copy_r   r/   r   r   r   r   r   )r|   rM   r}   s     r4   r  z-GraniteMoeHybridPreTrainedModel._init_weights4  s   f%f56LL,,3DKK<Y<YZLL))9V9VW :;LLSdkk6S6STf89JJv~~&JJv||UYYu||Av?O?ORS?S/T%UVJJvxx  <=JJv}}% >r6   )r   r   r   r(   r  base_model_prefixsupports_gradient_checkpointing_no_split_modules_skip_keys_device_placement_supports_flash_attn_supports_sdpa_supports_flex_attn_can_compile_fullgraph_supports_attention_backendr  rg   _can_record_outputs_is_statefulr/   r  r  r   r   s   @r4   r  r  "  ss    ""&*#78#4"5N!"&5/ LU]]_& &r6   r  c                        e Zd Zdef fdZeee	 	 	 	 	 	 ddej                  dz  dej                  dz  dej                  dz  dedz  dej                  dz  d	edz  d
ee   deez  fd                     Z xZS )GraniteMoeHybridModelrh   c           	      N   t         |   |       |j                  | _        |j                  | _        t        j                  |j                  |j                  | j                        | _        t        j                  t        |j                        D cg c]  }t        ||       c}      | _        t        |j                  |j                        | _        |j"                  dk(  rt%        |      nd | _        d| _        |j*                  | _        | j-                          y c c}w )Nr   ropeF)rm   rn   pad_token_idpadding_idx
vocab_sizer   	Embeddingrp   embed_tokens
ModuleListrangenum_hidden_layersr  r  r  r   r   position_embedding_typerm  
rotary_embgradient_checkpointingembedding_multiplier	post_initr{   s      r4   rn   zGraniteMoeHybridModel.__init__F  s     !.. ++LL):):F<N<NPTP`P`ammNSTZTlTlNmn)&)<n
 ,F,>,>FDWDWX	EKEcEcgmEm9&Asw&+#$*$?$?! 	 os   D"N	input_idsrQ   r  r~   inputs_embedsr  rT   rD   c           	         |d u |d uz  rt        d      || j                  |      }|| j                  z  }|r|t        | j                        }|V||j                         nd}t        j                  |j                  d   |j                        |z   }|j                  d      }t        |x}	t              s(| j                  |||d}
t        d
i |
t        d
i |
d}	|}d }| j                  | j                  ||      }t!        | j"                        D ]-  \  }} ||f|	| j                  j$                  |      |||d|}/ | j'                  |      }t)        ||	      S )Nz:You must specify exactly one of input_ids or inputs_embeds)rh   r   r'   r%  )rh   r  rQ   r~   )full_attentionr  )rQ   r~   r  r   )last_hidden_stater~   r  )
ValueErrorr  r  r   rh   get_seq_lengthr/   r   r.   r   r9   r  dictr   r   r  	enumerater  r  r   r   )r|   r  rQ   r  r~   r  r  rT   past_seen_tokenscausal_mask_mappingmask_kwargsrB   r   idecoder_layers                  r4   r   zGraniteMoeHybridModel.forwardW  s    -t";<YZZ  --i8M%(A(AA0*$++>OCRC^==?de <<(;(;A(>}G[G[\_ooL'11!4L?-F ++!."0#2	K #5"C{"C$C$Rk$R# &"??&"&//-"N )$++ 6 	A})24;;3P3PQR3ST /#$7 M	 		-0%++
 	
r6   )NNNNNN)r   r   r   r(   rn   r   r$   r&   r/   r  r   r
   r  r   r   r  r   r   r   r   r   s   @r4   r   r   D  s    5 "  .2.204(,26!%<
##d*<
 t+<
 &&-	<

 <
 ((4/<
 $;<
 45<
 
(	(<
    <
r6   r   gate_logitsr  c                    | t        | t              syt        | t              rC| d   j                  }t        j                  | D cg c]  }|j                  |       c}d      }t        j                  j                  j                  d      }t        j                  ||d      \  }}	t        j                  j                  j                  |	|      }
|>t        j                  |
j                         d      }t        j                  |d      }n|j                  \  }}|j                  d   ||z  z  }|dddddddf   j                  |||||f      j                  d||      j                        }t        j                   |
j                         |z  d      t        j                   |d      z  }|ddddddf   j                  ||||f      j                  d|      j                  |      }t        j                   ||z  d      t        j                   |d      z  }t        j                   ||j#                  d      z        }||z  S c c}w )a  
    Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.

    See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
    function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
    experts is too unbalanced.

    Args:
        gate_logits:
            Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
            shape [batch_size X sequence_length, num_experts].
        num_experts:
            Number of experts
        top_k:
            The number of experts to route per-token, can be also interpreted as the `top-k` routing
            parameter.
        attention_mask (`torch.Tensor`, *optional*):
            The attention_mask used in forward function
            shape [batch_size X sequence_length] if not None.

    Returns:
        The auxiliary loss.
    Nr   r,   r*   )r  r   r   r/   r0   r_   r   r\   r]   r  r  r[  r  r.   rF   rG   r'  r9   )r  r  r  rQ   compute_device
layer_gateconcatenated_gate_logitsrouting_weightsr  selected_expertsr  tokens_per_expertrouter_prob_per_expertr  sequence_lengthr
  expert_attention_mask router_per_expert_attention_maskoverall_losss                      r4   load_balancing_loss_funcr+    s9   : *[%"@+u%$Q..#(99^i-jPZjmmN.K-jpq#r hh))112JPR1SO**_eDA((%%--.>LK!JJ{'8'8':B "'O!C&4&:&:#
O4::1=*B^_ 4AtT12V&
OUKXYWR,R	 	 "IIk&7&7&9<Q&QWXY\a\e\e!q]
 
 4At+,V&
O[QRWR%R	 	) "'?=]+]cd!ehmhqhq,!i
 "
 99.1G1Q1QRS1TTUL+%%[ .ks   Ic                   D    e Zd ZddiZddiZddgdgfiZdef fdZee		 	 	 	 	 	 	 	 dd
e
j                  d	z  de
j                  d	z  de
j                  d	z  ded	z  de
j                  d	z  de
j                  d	z  ded	z  dee
j                  z  deez  fd              Z xZS )GraniteMoeHybridForCausalLMzlm_head.weightzmodel.embed_tokens.weightlm_headcolwise_gather_outputrB   logitsrh   c                 p   t         |   |       t        |      | _        |j                  | _        t        j                  |j                  |j                  d      | _        |j                  | _	        |j                  | _        |j                  | _        |j                  | _        | j                          y )NFrk   )rm   rn   r   r  r  r   ru   rp   r.  router_aux_loss_coefr  r  r  logits_scalingr  rg  s     r4   rn   z$GraniteMoeHybridForCausalLM.__init__  s     *62
 ++yy!3!3V5F5FUS$*$?$?!!33#)#=#= $33 	r6   Nr  rQ   r  r~   r  labelsoutput_router_logitslogits_to_keeprD   c	           	         ||n| j                   j                  } | j                  d|||||d|	}
|
j                  }t	        |t
              rt        | d      n|}| j                  |dd|ddf         }|| j                   j                  z  }d}|* | j                  ||fd| j                   j                  i|	}d}|rYt        |
j                  | j                  | j                  |      }|+|| j                  |j!                  |j"                        z  z  }t%        ||||
j&                  |
j(                  |
j*                  |
j                        S )a  
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.

        Example:

        ```python
        >>> from transformers import AutoTokenizer, GraniteMoeHybridForCausalLM

        >>> model = GraniteMoeHybridForCausalLM.from_pretrained("ibm-granite/granite-4.0-h-tiny")
        >>> tokenizer = AutoTokenizer.from_pretrained("ibm-granite/granite-4.0-h-tiny")

        >>> prompt = "Hey, are you conscious? Can you talk to me?"
        >>> inputs = tokenizer(prompt, return_tensors="pt")

        >>> # Generate
        >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
        "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
        ```N)r  rQ   r  r~   r  r  )lossaux_lossr0  r~   rB   r  r  r  )rh   r5  r  r  r  r   slicer.  r3  loss_functionr  r+  r  r  r  r2  r_   r   r   r~   rB   r  )r|   r  rQ   r  r~   r  r4  r5  r6  rT   outputsrB   slice_indicesr0  r8  r9  s                   r4   r   z#GraniteMoeHybridForCausalLM.forward  s|   J %9$D $++JjJj 	 $** 
)%+'
 
  118B>SV8W~ot4]kmA}a,?@A$++444%4%%  ;;11 	D /%%  ((	H !11HKK4LLL(#33!//))!//
 	
r6   )NNNNNNNr   )r   r   r   _tied_weights_keys_tp_plan_pp_planr(   rn   r   r    r/   r  r   r
   r  r   r   r   r   r   r   r   s   @r4   r-  r-    s    *,GH23H_-z:;H5   .2.204(,26*.,0-.Q
##d*Q
 t+Q
 &&-	Q

 Q
 ((4/Q
   4'Q
 #TkQ
 ell*Q
 
*	*Q
  Q
r6   r-  )r-  r   r  )r'   )r   )Nr+   N)]collections.abcr   typingr   r   r/   torch.nn.functionalr   r\   r   r   r  activationsr	   cache_utilsr
   r   
generationr   integrationsr   r   r   r   integrations.accelerater   integrations.hub_kernelsr   masking_utilsr   r   modeling_layersr   modeling_outputsr   r   r   modeling_rope_utilsr   r   modeling_utilsr   r   processing_utilsr   utilsr   r   r    r!   r"   utils.genericr#   r$   utils.import_utilsr%   utils.output_capturingr&   configuration_granitemoehybridr(   
get_loggerr   r   r5   rA   r   r   rL   Moduler  re   rg   r   r   r   r   r   r   ra  rm  r  r  r  r  r  r  r  r   r   r+  r-  __all__r  r6   r4   <module>rY     s  * % &     & ! . )  > 8 P 9 j j K F & l l G 9 5 B 
		H	%( *+ ,2	UU\\ 	U# 	U%,, 	U& %II%<<% 
% <<	%
 LL4'% % % '(%2 )*A)		 A) +A)NVU\\ VS V
((	lO lO^;588?? ;$")) 4><bii ><B9 9, $#bii $# $#N?")) ?")5 0 Y'Jbii J (J(A#= AH &o & &B Q
; Q
 Q
l #
*.	O&ell 33d:O&tO& LL4'	O&
 \\CO&d e
"A? e
 e
P fr6   