mla_masked_attention

Function mla_masked_attention 

Source
pub fn mla_masked_attention(
    q: &[f32],
    k: &[f32],
    v: &[f32],
    mask: &[u8],
    dims: DsaDims,
    q_rows: usize,
    kv_len: usize,
) -> Vec<f32>
Expand description

NoPE MLA over a per-query selected key set.

q: [q_rows, heads, qk_head_dim], k: [kv_len, heads, qk_head_dim], v: [kv_len, heads, v_head_dim], mask: [q_rows, kv_len].

🪤 NoPE means qk_head_dim == qk_nope_head_dim and there is no rope section to skip. The scale is qk_head_dim^-0.5 over the FULL head dim, which on GLM-5.3 equals qk_nope_head_dim^-0.5 only because the rope part is zero-width — do not hardcode either.