Skip to main content

mla_forward_token

Function mla_forward_token 

Source
pub fn mla_forward_token(
    weights: &MlaAttnWeights,
    cfg: &MlaConfig,
    hidden: &[f32],
    rms_norm_eps: f32,
    k_cache: &mut Vec<f32>,
    v_cache: &mut Vec<f32>,
) -> Vec<f32>
Expand description

One decode step. k_cache/v_cache are growable, caller-owned buffers in [seq_len_so_far, n_heads, head_dim] layout (head_dim = qk_nope_head_dim + qk_rope_head_dim for k, v_head_dim for v) — plain Vec<f32>, not yet ferrox_core::cache::KvCache (see module doc comment). This function appends the current position’s K/V to both before running attention over every position pushed so far.