Struct CausalMasker 

pub struct CausalMasker;

Implementations

impl CausalMasker

pub fn expand_mask( &self, mask: &Tensor, dtype: DType, tgt_len: Option<usize>, ) -> Result<Tensor>

Expands a mask from (bs, seq_len) to (bs, 1, tgt_len, seq_len) If tgt_len is None, use seq_len

pub fn calculate_past_kv_len( &self, cache: &[Option<(Tensor, Tensor)>], ) -> Result<usize>

pub fn make_causal_mask( &self, input_ids: &Tensor, cache: &dyn PastKvLenCache, dtype: DType, cfg: &CausalMaskConfig, ) -> Result<AttentionMask>

Build a causal attention mask.

• Returns AttentionMask::None for single-token decode.
• Returns AttentionMask::CausalFlash on CUDA+flash when force_custom is false (flash handles causality internally).
• Returns AttentionMask::Custom with a real mask tensor otherwise.

pub fn make_chunked_mask_matrix( &self, input_ids: &Tensor, chunk_size: usize, cache: &dyn PastKvLenCache, dtype: DType, _n_attn_heads: usize, ) -> Result<Option<Tensor>>

pub fn apply_mask_one_and_zero( &self, mask: &Option<Tensor>, att: Tensor, neg_inf: &Tensor, ) -> Result<Tensor>