pub fn argmax_bf16_batch(
gpu: &dyn GpuBackend,
kernel: KernelHandle,
logits: DevicePtr,
out: DevicePtr,
vocab_size: u32,
n_rows: u32,
row_stride: u32,
stream: u64,
) -> Result<()>Expand description
Batched argmax: ONE launch, one block per row, instead of n serial launches of
the single-row argmax_bf16 (which is a one-CTA reduction and so uses 1 of 48
SMs). Byte-identical — each block runs the identical per-row body.