File size: 2,797 Bytes
c971a45 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | // 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<storage, read> ids: array<u32>;
@group(0) @binding(2) var<storage, read> table: array<{{T}}>;
@group(0) @binding(3) var<storage, read> posEmbed: array<{{T}}>;
@group(0) @binding(4) var<storage, read_write> Y: array<{{T}}>;
const D: u32 = {{D}}u;
const WG: u32 = {{WG}}u;
@compute @workgroup_size({{WG}})
fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
// 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]));
}
}
|