Skip to main content

attention_with_scale

Function attention_with_scale 

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

Causal attention with an explicit softmax scale (Gemma-4 uses 1.0 after QK-norm).