a
    h                     @   sH  d Z ddlmZmZmZ ddlZddlZddlmZ ddlm	Z	 ddl
mZmZ ddlmZmZmZmZmZmZ ddlmZmZmZmZ d	d
lmZ d	dlmZmZ d	dlmZ d	dl 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* rddl-m.Z. ddl/m0Z0m1Z1 ndZ.e) r<ddl2m3Z3m4Z4 nd\Z4Z3e5e.e3e4fZ6e%7e8Z9G dd deddZ:G dd deZG dd deZ;d6dd Z<G d!d" d"eZ=G d#d$ d$eZ>d%d& Z?G d'd( d(ej@ZAG d)d* d*eZBG d+d, d,eZCG d-d. d.eZDe#G d/d0 d0eZEe#G d1d2 d2eEZFG d3d4 d4eZGg d5ZHdS )7zPyTorch Bamba model.    )Optional	TypedDictUnionN)nn)ACT2FN) HybridMambaAttentionDynamicCacheJambaAttentionDecoderLayer)LlamaAttentionLlamaForCausalLMLlamaMLPLlamaRMSNormLlamaRotaryEmbeddingrotate_half)MambaRMSNormGatedpad_tensor_by_sizereshape_into_chunkssegment_sum   )AttentionMaskConverter)BaseModelOutputWithPastCausalLMOutputWithPast)PreTrainedModel)Unpack)auto_docstringcan_return_tuplelogging)deprecate_kwarg)is_causal_conv1d_availableis_mamba_2_ssm_available   )BambaConfig)selective_state_update)mamba_chunk_scan_combined mamba_split_conv1d_scan_combined)causal_conv1d_fncausal_conv1d_update)NNc                   @   s@   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< dS )BambaFlashAttentionKwargsa  
    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.

    Attributes:
        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.
    Zcu_seq_lens_qZcu_seq_lens_kZmax_length_qZmax_length_kseq_idxN)	__name__
__module____qualname____doc__torch
LongTensor__annotations__int	IntTensor r1   r1   c/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/bamba/modular_bamba.pyr&   L   s   


r&   F)totalc                   @   s&   e Zd ZdZejdfedddZdS )r   a  
    A dynamic cache that can handle both the attention cache (which has a seq_len dimension) and the mamba cache
    (which has a constant shape regardless of seq_len).

    This cache has two sets of lists of tensors: `key_cache` and `value_cache` for attention cache and `conv_states`
    and `ssm_states` for mamba cache. Each of these lists has `num_layers` tensors. The expected shape for each tensor
    For attention layers, `key_cache` and `value_cache` have a shape of `(batch_size, num_heads, seq_len, head_dim)`,
    while `conv_states` and `ssm_states` have a shape of `(batch_size, 0)` (empty tensors).
    For mamba layers, `key_cache` and `value_cache` have a shape of `(batch_size, 0)` (empty tensors),
    while `conv_states` represents the convolution state and has a shape of `(batch_size, d_inner, d_conv)`,
    and `ssm_states` represents the ssm state and has a shape of `(batch_size, d_inner, d_state)`.
    Nconfigc                    s0  |j | _ d| _|j}|j}g | _g | _g | _t|jD ]}| j | dkr|  jt	j
 |j|j d|j |  ||dg7  _|  jt	j
 |j|j||dg7  _q6|  jt	jg g  dg7  _|  jt	jg g  dg7  _| j| q6 fddt|jD | _ fddt|jD | _d S )	NFmamba   devicedtyper9   c                    s    g | ]}t jg g  d qS r;   r,   tensor.0_
batch_sizer9   r1   r2   
<listcomp>       z=HybridMambaAttentionDynamicCache.__init__.<locals>.<listcomp>c                    s    g | ]}t jg g  d qS r<   r=   r?   rB   r1   r2   rD      rE   )layers_block_typehas_previous_statemamba_d_convmamba_d_stateconv_states
ssm_statesZtransformer_layersrangenum_hidden_layersr,   Zzerosmamba_expandhidden_sizemamba_n_groupsmamba_n_headsmamba_d_headr>   appendZ	key_cacheZvalue_cache)selfr5   rC   r:   r9   conv_kernel_sizessm_state_sizeir1   rB   r2   __init__t   sB    	
  z)HybridMambaAttentionDynamicCache.__init__)r(   r)   r*   r+   r,   Zfloat16r    rX   r1   r1   r1   r2   r   f   s   r   c                   @   s   e Zd ZdS )BambaRotaryEmbeddingNr(   r)   r*   r1   r1   r1   r2   rY      s   rY   c                 C   s   | |}| |}|jd }| dd|f | d|df  }}|dd|f |d|df  }	}
|| t||  }|	| t|	|  }tj||gdd}tj||
gdd}||fS )a  Applies Rotary Position Embedding to the query and key tensors.

    Removes the interleaving of cos and sin from GLM

    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.
        position_ids (`torch.Tensor`, *optional*):
            Deprecated and unused.
        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.
    .Ndim)	unsqueezeshaper   r,   cat)qkcossinposition_idsZunsqueeze_dimZ
rotary_dimZq_rotZq_passZk_rotZk_passZq_embedZk_embedr1   r1   r2   apply_rotary_pos_emb   s    


""rf   c                   @   s   e Zd ZdS )BambaAttentionNrZ   r1   r1   r1   r2   rg      s   rg   c                   @   s   e Zd ZdS )BambaRMSNormGatedNrZ   r1   r1   r1   r2   rh      s   rh   c                 C   sN   |durJ|j d dkrJ|j d dkrJ| j}| |dddddf  |} | S )zm
    Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
    Nr   r   )r_   r:   to)hidden_statesattention_maskr:   r1   r1   r2   apply_mask_to_padding_states   s    $ rl   c                       s   e Zd ZdZeed fddZdeje	e
 e	ej e	ej e	ej dddZde	e
 e	ej e	ej d	d
dZde	e
 e	ej e	ej e	ej dddZ  ZS )
BambaMixeruP  
    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
    r5   	layer_idxc                    s  t    |j| _|j| _|j| _|j| _t	|j
| j | _|| _|j| _|j| _t|j | _|j| _|j| _|j| _|j| _|j| _dtdf| _d| _d| _ | jd| j | j  | _!t"j#| j!| j!|j| j| j!| jd d| _$| j| j! | j }t"j%| j|| jd| _&t"'t()| j| _*t(+d| jd }t"'t(,|| _-d	| j-_.t/| j| jd
| _0t"'t()| j| _1d	| j1_.t"j%| j| j| jd| _2t3st45d n
t45d d S )N        infgMbP?g?r7   r   )Zin_channelsZout_channelsbiasZkernel_sizegroupspadding)rr   Tepsa  The fast path is not available because on 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-conv1dzDThe fast path for Bamba will be used when running the model on a GPU)6superrX   rQ   	num_headsrO   rI   rV   rH   rU   r/   rN   intermediate_sizero   Zmamba_conv_biasuse_conv_biasZ
hidden_act
activationr   actZmamba_proj_biasZuse_biasrms_norm_epsZlayer_norm_epsilonrP   n_groupsrR   head_dimZmamba_chunk_size
chunk_sizefloattime_step_limitZtime_step_minZtime_step_maxconv_dimr   ZConv1dconv1dZLinearin_proj	Parameterr,   Zonesdt_biasarangelogA_logZ_no_weight_decayrh   normDout_projis_fast_path_availableloggerwarning_once)rT   r5   ro   Zprojection_sizeA	__class__r1   r2   rX      s\    

	zBambaMixer.__init__N)rj   cache_paramscache_positionrk   r'   c                 C   s  t ||}| |}|j\}}}	| j| j }
|d uo|jo|dko|j| j jd |j| j jd   kop|kn  o|d uo|d dk}|rR|	dj
| j| j| jgdd\}}}t||j| j | jj	d| jj| j}tj
|| j|
|
gdd\}}}t| j  }|d d d df d d d d d f d| j| jjtjd}|d d d d d f dd| j}| jd d d df d| j}| jd d d df d| j}||| j|jd | j }||| j|jd | j }||| j| j}t|j| j ||||||d |dd
}||| j| j }| ||}|  |d d d df }nFt| j  }| j!d	td
fkr|i nd| j!i}| j"r|d u rt#|| jj	d| jj| j|f| j| j$|| j| jj| jj%| j j| j j| j| jddd|}n|j
| j| j| jgdd\}}}|d ur^|&dd}t'j()|| j*|jd  df}|j| j +| | jdvr| ,| |&dddd |f &dd}n2t-|&dd| jj	d| jj| j|d&dd}t ||}tj
|| j|
|
gdd\}}}t.|||d| j|||||| jd|||| jdf| j$| jd |d| jdd|\}}|d urt|d urt|j| j +| |||d}| ||}|  |}|S )Nr   r   r[   r\   .r:   T)zr   dt_softplusrp   rq   Zdt_limitF)r   r   r'   r{   Zrmsnorm_weightZrmsnorm_epsZoutproj_weightZoutproj_biasZheaddimZngroupsZnorm_before_gatereturn_final_statesr7   )ZsiluZswish)xweightrr   r{   r'   )r   r   r   r'   r   r   r   )/rl   r   r_   r~   rV   rG   rJ   ro   rK   squeezesplitry   r   rx   r%   r   r   rr   r{   r,   expr   r   expandr   ri   float32r   r   viewr!   r   r   r   trainingr#   r   Zvariance_epsilon	transposer   
functionalpadrU   copy_r|   r$   r"   )rT   rj   r   r   rk   r'   projected_statesrC   seq_lenrA   Zgroups_time_state_sizeuse_precomputed_statesgatehidden_states_B_CdtBCr   r   r   Zhidden_states_reshapedoutZdt_limit_kwargshidden_states_B_C_transposedrJ   scan_output	ssm_stater1   r1   r2   cuda_kernels_forward-  s   	





<"
"

$




zBambaMixer.cuda_kernels_forward)r   r   rk   c           3   
      s  |j \}}}|j}t||}|}	|	jjjjgdd\}
}}|d uo|jo|dko|j	j
 j d |jj
 j d   ko|kn  o|d uo|d dk}|r^|j	j
 jddd|j	j
< |d d dd d f |j	j
 j|j	j
 d d d d df< |j	j
 jjjjd}tj|jjd dd}jrR|jj }|}nr|d ur|dd}tj|j|j d  df}|j	j
 | |dddd |f dd}t||}tj|jjj jj gdd\}}}tj !  }|r|jj
 j}|d d dd d f d d d df }|dd"||j d j#}j$d	 "j$j d j#}tjj%|||j }t&|j'd j'd }|d
 "jj#jjtj(d}t|d	 | j|d}|)|jddd d d f }|"|jjj |j d * }|)|d|j d }|d	 |dd d d f  }|)|dj#}||d	  j|d}|jj
 |jj
 | |  |)|jddd d d f }|"|jjj |j d * }|)|d|j d }|jj
 j|j|jd}|+|j j#j}|+|j jd}t,||}|+|jj#}j-d	 "j-j d j#}|||  |j}|)|dd d d df }ntj%|j$ }t&|j'd j'd }|)||dj#! }|)||dj! }|)||dj! }|j.jj djd}|j.jj djd}j/|j/  j/  j-d	 t0|  }||d	  }||j| } fdd||||fD \}}}}|1dddd}tj2|dd}tt3|} |d d d d d d d d d d d f |d d d d d d d d d d d f  }!|!jdd}"|"d	 | 1dddddd	  }#|#jdd}$|$d	 |d d d d d f  jdd}%t|d d d d d d dd f | }&||&1ddddd	  }'|'dd d d f |d	  jdd}(|r0|jj
 d d d df j|(jd})nt4|(d d d df })tj5|)|(gdd}(tt3tj|d d d d d d df d}*|*dd}*|*d
 |(d d d d d df  jdd}+|+d d d df |+d d df  }(},t|}-|dd d d f |(d d d d d df  }.|-1dddd}/|.d|/d	  }0|%|0 }|)|djj#}|| } dkr|d d d |d d d d f }|)||d}|,d ur|d ur|jj
 |, 6||
}17|1|}2|2S )Nr[   r\   r   r   )Zshiftsdimsr;   r7   .).N).NNr   r8   )r]   Zoutput_sizec                    s   g | ]}t | jqS r1   )r   r   )r@   tZpad_sizerT   r1   r2   rD   `  rE   z,BambaMixer.torch_forward.<locals>.<listcomp>r      )r   r   )8r_   r:   rl   r   r   ry   r   rx   rG   rJ   ro   rK   Zrollri   r9   r   r   r,   sumr   rz   rr   r|   r   r   r   r   rU   r   r~   rV   r   r   r   r   r   r   Zsoftplusclampr   r   reshape
contiguousr   Zbmmr   Zrepeat_interleaver   r   Zpermutecumsumr   Z
zeros_liker`   r   r   )3rT   Zinput_statesr   r   rk   rC   r   rA   r:   r   r   r   r   r   rJ   r   rj   r   r   r   Zcache_devicer   ZdAZdBZdBxrK   Zssm_states_reshapedZ
C_reshapedyr   Z
D_residualZA_cumsumLZG_intermediateGZM_intermediateMZY_diagZdecay_statesZB_decayZstatesZprevious_statesZdecay_chunkZ
new_statesr   Zstate_decay_outZC_times_statesZstate_decay_out_permutedZY_offr   Zcontextualized_statesr1   r   r2   torch_forward  s    


@
,
$"$$$P&*"&0(&
*
 zBambaMixer.torch_forward)r   r   rk   r'   c                 K   s   t r&d| jjjjv r&| |||||S |d ur6td|j}|d ur|jd dkr|jd dkr||d d d d d f  	|}| 
||||S )Ncudaz\`seq_idx` support requires fast path support. Please install `mamba_ssm` and `causal_conv1d`r   r   )r   r   r   r9   typer   NotImplementedErrorr:   r_   ri   r   )rT   rj   r   r   rk   r'   kwargsr:   r1   r1   r2   forward  s    	$ zBambaMixer.forward)NNNN)NNN)NNNN)r(   r)   r*   r+   r    r/   rX   r,   Tensorr   r   r-   r0   r   r   r   __classcell__r1   r1   r   r2   rm      sB   F     .    S    rm   c                   @   s   e Zd ZdS )BambaMLPNrZ   r1   r1   r1   r2   r     s   r   c                   @   s   e Zd ZdS )BambaRMSNormNrZ   r1   r1   r1   r2   r     s   r   c                       s   e Zd Zdeeed fddZedddddej	e
ej	 e
ej e
e e
e e
e e
ej e
eej	ej	f  ee eeje
eejejf  f d
ddZ  ZS )BambaDecoderLayerr6   )r5   ro   
layer_typec                    sp   t  || | `d}|dkr"tnd }||| _|| _|dkrNt||d| _n|dkrdt||| _nt	dd S )Nr   r6   rn   	attentionzInvalid layer_type)
rw   rX   	self_attnr   feed_forwardr   rm   r6   rg   
ValueError)rT   r5   ro   r   Znum_expertsZffn_layer_classr   r1   r2   rX     s    
zBambaDecoderLayer.__init__Zpast_key_valuepast_key_valuesz4.58)new_nameversionNF)
rj   rk   re   r   output_attentions	use_cacher   position_embeddingsr   returnc	                 K   s   |}
|  |}| jdkr8| jf ||||d|	}d}n0| jdkrh| jf ||||||||d|	\}}|
| }|}
| |}| |}|
| }|f}|r||f7 }|S )a  
        Args:
            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
                `(batch, sequence_length)` where padding elements are indicated by 0.
            past_key_values (`HybridMambaAttentionDynamicCache`, *optional*): cached past key and value projection states
            output_attentions (`bool`, *optional*):
                Whether or not to return the attentions tensors of all attention layers. See `attentions` under
                returned tensors for more detail.
            use_cache (`bool`, *optional*):
                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
                (see `past_key_values`).
            cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
                Indices depicting the position of the input sequence tokens in the sequence.
            position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
                Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
                with `head_dim` being the embedding dimension of each attention head.
            kwargs (`dict`, *optional*):
                Arbitrary kwargs. Can be used to provide `BambaFlashAttentionKwargs` for
                padding-free training and/or improve torch.compile performance.
        r6   )rj   r   r   rk   Nr   )rj   rk   re   r   r   r   r   r   )Zinput_layernormr   r6   r   Zpre_ff_layernormr   )rT   rj   rk   re   r   r   r   r   r   r   ZresidualZself_attn_weightsoutputsr1   r1   r2   r     sD    #


	



zBambaDecoderLayer.forward)r6   )NNNFFNN)r(   r)   r*   r    r/   strrX   r   r,   r   r   r-   r   booltupler   r&   FloatTensorr   r   r1   r1   r   r2   r     s*          r   c                       sD   e Zd ZU eed< dZdZdgZdZdZ	dZ
dZ fddZ  ZS )BambaPreTrainedModelr5   modelTr   r   c                    sR   t  | t|trN|jjd tt	d|j
d |j_|jjd d S )Ng      ?r   )rw   _init_weights
isinstancerm   r   dataZfill_r,   r   r   rx   r   r   )rT   moduler   r1   r2   r   2  s
    
z"BambaPreTrainedModel._init_weights)r(   r)   r*   r    r.   Zbase_model_prefixZsupports_gradient_checkpointingZ_no_split_modulesZ_skip_keys_device_placementZ_supports_flash_attnZ_supports_sdpaZ_is_statefulr   r   r1   r1   r   r2   r   &  s   
r   c                       s   e Zd Zed fddZeedeej	 eej
 eej	 ee eej ee ee ee eej	 ee edddZej
ej
ej
eedd	d
Zeej
eeejej
edddZdd Z  ZS )
BambaModelr4   c                    s   t  | |j| _|j| _t|j|j| j| _g }t	|j
D ]}|t|||j| d q@t|| _|j| _t|j|jd| _t|d| _d| _|   d S )N)ro   r   ru   r4   F)rw   rX   Zpad_token_idZpadding_idx
vocab_sizer   Z	EmbeddingrO   embed_tokensrL   rM   rS   r   rF   Z
ModuleListlayers_attn_implementationr   r}   final_layernormrY   
rotary_embgradient_checkpointing	post_init)rT   r5   Zdecoder_layersrW   r   r1   r2   rX   <  s    zBambaModel.__init__N)	input_idsrk   re   r   inputs_embedsr   r   output_hidden_statesr   r   r   c
                 K   s  |d ur|n| j j}|d ur |n| j j}|d ur4|n| j j}|d u |d uA rTtd| jrr| jrr|rrtd d}|d u r| 	|}|}|r|d u rtd |	d u rt
j|jd |jd}	|d u r|	d}| |||	||}| ||	}| ||}|rdnd }|rdnd }| jD ]t}|jd	kr.|n|}|rB||f7 }||f||||||	|d
|
}|d }|r|d d ur||d f7 }q| |}|r||f7 }|r|jsd|_|sd n|}t||||dS )Nz:You must specify exactly one of input_ids or inputs_embedszX`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.FzBamba requires an initialized `HybridMambaAttentionDynamicCache` to return a cache. None was provided, so no cache will be returned.r   r;   r   r1   r6   )rk   re   r   r   r   r   r   T)last_hidden_stater   rj   
attentions)r5   r   r   r   r   r   r   r   r   r   r,   r   r_   r9   r^   _update_causal_mask_update_mamba_maskr   r   r   r   rG   r   )rT   r   rk   re   r   r   r   r   r   r   r   rj   causal_mask
mamba_maskr   Zall_hidden_statesZall_self_attnsZdecoder_layerZ
layer_maskZlayer_outputsZ
next_cacher1   r1   r2   r   O  s|    




	

zBambaModel.forward)rk   input_tensorr   r   r   c                 C   s   | j jdkr$|d ur d|v r |S d S |d ur4| nd}| j jdkr`|s`tj|||| jdr`d S |j}|jd }t|t	j
r|jd n
|| d }	| j|||	|||jd d}
| j jdkr|d ur|jjd	v r|st	|j}t|
|}
|
S )
NZflash_attention_2rp   r   Zsdpa)r   Zpast_key_values_lengthZis_trainingr   r[   )sequence_lengthtarget_lengthr:   r   rC   )r   ZxpuZnpu)r5   r   Zget_seq_lengthr   Z_ignore_causal_mask_sdpar   r:   r_   r   r,   r   5_prepare_4d_causal_attention_mask_with_cache_positionr9   r   finfominZ_unmask_unattended)rT   rk   r   r   r   r   Zpast_seen_tokensr:   r   r   r   	min_dtyper1   r1   r2   r     sL    





	zBambaModel._update_causal_mask)rk   r   r   r:   r   rC   c                 K   s~  | dur|   dkr| }n^t|j}tj||f|||jd}|dkrVtj|dd}|tj||jd|ddk9 }|ddddddf 	|ddd}| durz|
 }| jd }	| ddddddf | ddddddf kdddd| dddf |}
|ddddddd|	f |
 }|dk}|ddddddd|	f |||ddddddd|	f< |S )	aM  
        Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape
        `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.

        Args:
            attention_mask (`torch.Tensor`):
                A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape
                `(batch_size, 1, query_length, key_value_length)`.
            sequence_length (`int`):
                The sequence length being processed.
            target_length (`int`):
                The target length: when generating with static cache, the mask should be as long as the static cache,
                to account for the 0 padding, the part of the cache that is not filled yet.
            dtype (`torch.dtype`):
                The dtype to use for the 4D attention mask.
            cache_position (`torch.Tensor`):
                Indices depicting the position of the input sequence tokens in the sequence.
            batch_size (`torch.Tensor`):
                Batch size.
        Nr   )Z
fill_valuer:   r9   r   )Zdiagonalr;   r[   r   )r]   r,   r   r   fullr9   Ztriur   r   r   cloner_   ri   Zmasked_fill)rk   r   r   r:   r   rC   r   r   r   Zmask_lengthZpadding_attention_maskZpadding_maskr1   r1   r2   r     s0     $

.$  z@BambaModel._prepare_4d_causal_attention_mask_with_cache_positionc                 C   s.   |}|d dks&|dur*t |dkr*d}|S )zv
        No need for zeroing states when
            1. Cached forward
            2. Attending to all inputs
        r   Nr   )r,   all)rT   rk   r   r   r1   r1   r2   r   '  s    "zBambaModel._update_mamba_mask)	NNNNNNNNN)r(   r)   r*   r    rX   r   r   r   r,   r-   r   r   r   r   r   r&   r   r   r   staticmethodr/   r:   r   r   r   r1   r1   r   r2   r   :  sP            d<7r   c                       s   e Zd Z fddZdeej eej eej ee eej	 eej ee
 ee
 ee
 eej eeejf edddZdd	d
Z  ZS )BambaForCausalLMc                    s    t  | |j| _|   d S )N)rw   rX   z_loss_coefficientr   )rT   r5   r   r1   r2   rX   4  s    zBambaForCausalLM.__init__Nr   )r   rk   re   r   r   labelsr   r   r   r   logits_to_keepr   c                 K   s   |dur|n| j j}|	dur |	n| j j}	| jf ||||||||	|
d	|}|j}t|trht| dn|}| |dd|ddf }d}|dur| j	f ||| j j
d|}| jdkr|jddj|jdd }|| j|  }t|||j|j|jd	S )
aJ  
        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, BambaForCausalLM

        >>> model = BambaForCausalLM.from_pretrained("...")
        >>> tokenizer = AutoTokenizer.from_pretrained("...")

        >>> 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   rk   re   r   r   r   r   r   r   )logitsr  r   r   r[   r\   r   r7   )lossr  r   rj   r   )r5   r   r   r   r   r   r/   sliceZlm_headZloss_functionr   r  Z	logsumexpri   r:   powmeanr   r   rj   r   )rT   r   rk   re   r   r   r  r   r   r   r   r  r   r   rj   Zslice_indicesr  r  Zz_lossr1   r1   r2   r   ;  s@    %

 zBambaForCausalLM.forwardTc              	   K   s  |d u }	|	sj|d us&|d |j d krD|d d |j d  d f }q|j d |j d kr|d d |f }nt| j|j d | j| jd}|d ur|d u r| dd }||dkd |	s|d d |j d  d f }|d ur|	rd|i}
nd| i}
|
	||||| jj
|d |
S )Nr[   r   r   r;   r   r   )re   r   r   rk   r  r   )r_   r   r5   r:   r9   longr   Zmasked_fill_r   updateZnum_logits_to_keep)rT   r   r   rk   r   r   re   r   r   Zempty_past_kvZmodel_inputsr1   r1   r2   prepare_inputs_for_generation  s<    

z.BambaForCausalLM.prepare_inputs_for_generation)NNNNNNNNNNr   )NNNNNT)r(   r)   r*   rX   r   r,   r-   r   r   r   r   r   r/   r   r   r  r   r1   r1   r   r2   r  3  sB   	           P      r  )r   r  r   )Nr   )Ir+   typingr   r   r   r,   Ztorch.utils.checkpointr   Ztransformers.activationsr   Z(transformers.models.jamba.modeling_jambar   r   Z(transformers.models.llama.modeling_llamar	   r
   r   r   r   r   Z*transformers.models.mamba2.modeling_mamba2r   r   r   r   Zmodeling_attn_mask_utilsr   Zmodeling_outputsr   r   Zmodeling_utilsr   Zprocessing_utilsr   utilsr   r   r   Zutils.deprecationr   Zutils.import_utilsr   r   Zconfiguration_bambar    Z+mamba_ssm.ops.triton.selective_state_updater!   Z!mamba_ssm.ops.triton.ssd_combinedr"   r#   Zcausal_conv1dr$   r%   r   r   Z
get_loggerr(   r   r&   rY   rf   rg   rh   rl   Modulerm   r   r   r   r   r   r  __all__r1   r1   r1   r2   <module>   s^    
5
(   da y 