ground_truth_block_weights

Function ground_truth_block_weights 

Source
pub fn ground_truth_block_weights(
    q: &[bf16],
    k: &[bf16],
    num_q_heads: usize,
    num_kv_heads: usize,
    head_dim: usize,
    block_size: usize,
    num_blocks: usize,
) -> Vec<f32>
Expand description

Ground-truth attention-weight per block: max over q_heads of the softmax-normalized attention to the block’s tokens. Used by the recall test as the “oracle” the predictor’s scores are compared against.