pub fn batched_embed(
gpu: &dyn GpuBackend,
kernel: KernelHandle,
token_ids_dev: DevicePtr,
embed_table: DevicePtr,
output: DevicePtr,
num_tokens: u32,
hidden_size: u32,
stream: u64,
) -> Result<()>Expand description
Batched embedding: gather N rows from embedding table in one launch.
Replaces N individual D2D copies with a single kernel.
token_ids_dev must point to [num_tokens] u32 on device.
Kernel: batched_embed(token_ids, embed_table, output, hidden_size)
Grid: (num_tokens, 1, 1) Block: (256, 1, 1)