Function embed_from_argmax 

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<()>

GPU-side argmax + embedding lookup — eliminates D2H sync in MTP propose.

Reads the argmax result from argmax_out, looks up the embedding row from embed_table, and writes it to embed_out. Also copies the token ID to token_id_out for deferred CPU readback.