pub fn mla_prefill_f32(
query: &[f32],
query_pe: &[f32],
key: &[f32],
key_pe: &[f32],
value: &[f32],
batch: usize,
query_len: usize,
key_len: usize,
heads: usize,
kv_heads: usize,
head_dim: usize,
pe_dim: usize,
scale: f32,
) -> Vec<f32>