pub fn argmax_bf16(
gpu: &dyn GpuBackend,
kernel: KernelHandle,
logits: DevicePtr,
out: DevicePtr,
vocab_size: u32,
stream: u64,
) -> Result<()>Expand description
GPU-side argmax over BF16 logits.
Finds the index of the maximum value, writes a single u32 to out.
Kernel: argmax_bf16(logits, out, n)
Grid: (1, 1, 1) Block: (1024, 1, 1)