Function argmax_bf16_batch_lp 

pub fn argmax_bf16_batch_lp(
    gpu: &dyn GpuBackend,
    kernel: KernelHandle,
    logits: DevicePtr,
    out: DevicePtr,
    out_logprob: DevicePtr,
    vocab_size: u32,
    n_rows: u32,
    row_stride: u32,
    stream: u64,
) -> Result<()>

Batched argmax that ALSO writes each row’s top-1 log-probability (out_logprob[row] = log softmax(row)[argmax], FP32), computed by online softmax in the same pass — same bandwidth as argmax_bf16_batch, same index semantics.

Consumer: D-Cut verification-depth pruning, whose ranking key is the prefix SUM of these log-probabilities (= the log of the prefix product of survival probabilities). Separate kernel so every existing argmax_bf16_batch caller stays byte-identical and an unresolved handle is a silent 0 the caller gates on.