pub fn sparse_attn(
q: &[f32],
kv: &[f32],
sink: &[f32],
idx: &[i32],
s: usize,
h: usize,
d: usize,
topk: usize,
scale: f32,
) -> Vec<f32>Expand description
sparse_attn (kernel.py:310-405, shim form): per query and head, softmax over the gathered
rows plus the sink (denominator only), all in f32; output bf16. q: [s][h][d],
kv: [rows][d], idx: [s][topk].