pub fn embed_from_argmax(
gpu: &dyn GpuBackend,
kernel: KernelHandle,
argmax_out: DevicePtr,
embed_table: DevicePtr,
embed_out: DevicePtr,
token_id_out: DevicePtr,
hidden_size: u32,
stream: u64,
) -> Result<()>Expand description
2026-09-25: Read the token id at argmax_out, copy its BF16 row of
embed_table to embed_out, and copy the id to token_id_out, all on the
device.