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.