Function splitk_reduce_f32 

pub fn splitk_reduce_f32(
    attn: &[f32],
    lse: &[f32],
    batch: usize,
    heads: usize,
    splits: usize,
    head_dim: usize,
) -> Vec<f32>