Function embedding 

pub fn embedding(
    ids: &WgpuStorage,
    chunks: &[(&WgpuStorage, usize)],
    chunk_rows: usize,
    wpe: Option<&WgpuStorage>,
    t: usize,
    c: usize,
    n_ctx: usize,
    pos: usize,
) -> WgpuStorage

Gather from row-chunked embedding tables. chunks[i] holds rows [i * chunk_rows, …) of the full table; each token is written by exactly the dispatch whose chunk owns its id.