Function block_sparse_select 

pub fn block_sparse_select(
    index_q: &[Vec<f32>],
    index_k: &[Vec<f32>],
    query_pos: usize,
    cfg: &BlockSparseConfig,
) -> Vec<Vec<usize>>

The blocks each KV head may read for one query position, ascending.

index_q is one index head per KV head ([n_kv_heads] [index_head_dim]); index_k is one index key per position ([n_positions][index_head_dim]), shared across heads. Returns one selection per KV head, in the same order.

Never returns an empty selection for a query at a valid position – see the module docs on why that matters.