// Token + position embedding: Y[r,i] = f32(table[id,i])·EMBED_SCALE // + f32(posEmbed[pos,i]) (pos NOT scaled) // One workgroup per row: dispatchWorkgroups(nRows). Two modes, selected at // build time (exactly one of SRC_IDS / DECODE): // SRC_IDS (encoder): ids is src token ids [B*S]; id = ids[r], pos = r % s // (batch rows are consecutive: r = b*S + m). PAD rows // embed the pad token normally — masking happens in // attention, not here. With PACKED (encoder row-packing: // pad rows dropped, T = Σ lens rows) each word carries // its own position: ids[r] = (pos << 16) | id — id fits // (VOCAB 24000 < 2^16) and the dispatch enforces // S < 2^16. // DECODE: ids is the token ring [T_max*B]; row r = batch index b; // id = DECODER_START when t == 0, else ids[(t-1)*batch + r]; // pos = t. // // Template placeholders (buildShader in pipelines.js): // ENABLE_F16 the f16 enable directive when T is f16, else empty // T storage type of table/posEmbed/Y (f16|f32) // WG workgroup size (224 → 2 elements per thread at D=448) // D row width (d_model, 448) // EMBED_SCALE √d_model as a full-precision literal // IF_SRC_IDS / IF_DECODE mode blocks // DECODER_START decoder start token id (DECODE mode only) {{ENABLE_IMMEDIATE}} {{ENABLE_F16}} struct Params { nRows: u32, t: u32, batch: u32, s: u32 } {{PARAM_BINDING}}var<{{PARAM_ADDRESS}}> params: Params; @group(0) @binding(1) var ids: array; @group(0) @binding(2) var table: array<{{T}}>; @group(0) @binding(3) var posEmbed: array<{{T}}>; @group(0) @binding(4) var Y: array<{{T}}>; const D: u32 = {{D}}u; const WG: u32 = {{WG}}u; @compute @workgroup_size({{WG}}) fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { // Uniform per workgroup (one workgroup per row) — safe early return. if (wid.x >= params.nRows) { return; } let r = wid.x; {{IF_SRC_IDS}} {{IF_PACKED}} let id = ids[r] & 0xffffu; let pos = ids[r] >> 16u; {{/IF_PACKED}} {{IF_NOPACKED}} let id = ids[r]; let pos = r % params.s; {{/IF_NOPACKED}} {{/IF_SRC_IDS}} {{IF_DECODE}} var id: u32 = {{DECODER_START}}u; if (params.t != 0u) { id = ids[(params.t - 1u) * params.batch + r]; } let pos = params.t; {{/IF_DECODE}} let toff = id * D; let poff = pos * D; let base = r * D; for (var i = lid.x; i < D; i = i + WG) { Y[base + i] = {{T}}(f32(table[toff + i]) * {{EMBED_SCALE}} + f32(posEmbed[poff + i])); } }