Struct CausalMaskCache 

pub struct CausalMaskCache { /* private fields */ }

Precomputed causal attention mask

Causal masks ensure that position i can only attend to positions <= i. We precompute a mask for max_seq_len and slice as needed.

Implementations

impl CausalMaskCache

pub fn new(max_seq_len: usize, device: &Device) -> Result<Self>

Create a new causal mask cache

Arguments

• max_seq_len - Maximum sequence length to precompute
• device - Device to store the mask on

pub fn get(&self, seq_len: usize) -> Result<Tensor>

Get causal mask for a specific sequence length

Returns a [seq_len, seq_len] mask

pub fn get_broadcast(&self, seq_len: usize) -> Result<Tensor>

Get causal mask reshaped for attention broadcasting

Returns mask with shape [1, 1, seq_len, seq_len] for broadcasting over [batch, heads, seq, seq] attention scores

pub fn max_seq_len(&self) -> usize

Get the maximum sequence length