pub fn attention_causal(
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 a batch of queries [seq_q, n_heads*head_dim].
When seq_q == seq_kv, query t attends to keys 0..=t (prefill).
When seq_q == 1, equivalent to attention (decode).