Struct Attention 

pub struct Attention { /* private fields */ }

Multi-head attention with optional grouped-query attention

Supports:

• Standard multi-head attention (num_heads == num_key_value_heads)
• Grouped-query attention (num_key_value_heads < num_heads)
• RoPE positional embeddings
• Optional causal masking

Implementations

impl Attention

pub fn new( hidden_size: usize, head_dim: usize, num_heads: usize, num_key_value_heads: usize, causal: bool, vb: VarBuilder<'_>, ) -> Result<Self>

Create new Attention layer

Arguments

• hidden_size - Input/output dimension
• head_dim - Dimension per attention head
• num_heads - Number of query heads
• num_key_value_heads - Number of key/value heads (for GQA)
• causal - Whether to use causal masking
• vb - VarBuilder for parameter initialization

pub fn forward( &self, hidden_states: &Tensor, cos_sin: Option<(&Tensor, &Tensor)>, ) -> Result<Tensor>

Forward pass

Arguments

• hidden_states - Input tensor [batch, seq_len, hidden_size]
• cos_sin - Optional RoPE embeddings (cos, sin) each [seq_len, head_dim]

Returns

Output tensor [batch, seq_len, hidden_size]