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