Skip to main content

attention

Function attention 

Source
pub fn attention(
    q: &[f32],
    k_cache: &[f32],
    v_cache: &[f32],
    n_heads: usize,
    n_kv_heads: usize,
    head_dim: usize,
) -> Result<Vec<f32>, EngineError>
Expand description

Causal attention for single query step against KV cache. q: [n_heads * head_dim], k_cache/v_cache: [seq, n_kv_heads * head_dim]