Function extract_bf16 

pub fn extract_bf16(
    bf16: &[u8],
    target: u32,
    k: usize,
    vocab: usize,
) -> PromptTokenLogprob

Extract one position’s PromptTokenLogprob from a BF16 [vocab] logits slice (the layout lm_head_batched writes per row).