sparse_attn

Function sparse_attn 

Source
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].