pub fn finalize_ref( state: &AttnState, num_seqs: usize, num_q_heads: usize, head_dim: usize, ) -> Vec<bf16>