pub fn argmax_bf16_batch_feed(
gpu: &dyn GpuBackend,
kernel: KernelHandle,
logits: DevicePtr,
masks: DevicePtr,
cells: DevicePtr,
vocab_size: u32,
n_rows: u32,
row_stride: u32,
stream: u64,
) -> Result<()>Expand description
2026-09-25: Batched argmax into the feed cells (device u32 per row) with a
per-row pair of masked ids masks ([m0, m1] per row, u32::MAX = none):
when the plain argmax is a masked id, the row gets the argmax with both ids
excluded. argmax_feed.cu states the tie rules. One block per row.