Function embedding_chunked 

pub fn embedding_chunked(
    ids: &Tensor,
    chunks: &[Tensor],
    chunk_rows: usize,
    wpe: Option<&Tensor>,
    pos: usize,
) -> Result<Tensor>

Embedding gather from a row-chunked table. chunks[i] holds rows [i * chunk_rows, …) of the full [vocab, c] table — chunking keeps every GPU binding under max_storage_buffer_binding_size (GPT-2’s wte alone is ~147 MiB, above the WebGPU default and llvmpipe’s hard 128 MiB limit).