Struct MultiHeadAttention 

pub struct MultiHeadAttention { /* private fields */ }

Multi-head attention.

Projects input through Q, K, V linear layers, splits into n_heads heads, runs per-head scaled dot-product attention (via the fused Attention op), concatenates heads, and applies an output projection.

Implementations

impl MultiHeadAttention

pub fn new( wq: Linear, wk: Linear, wv: Linear, wo: Linear, n_heads: usize, ) -> Self

Create a new MultiHeadAttention from pre-built linear layers.

• wq, wk, wv: projection layers with weight [n_heads * head_dim, model_dim]
• wo: output projection [model_dim, n_heads * head_dim]
• n_heads: number of attention heads

pub fn forward_causal(&self, x: &Tensor) -> Result<Tensor>

Forward pass with causal masking (self-attention, auto-regressive).

x has shape [seq_len, model_dim]. Returns [seq_len, model_dim].