engine/kernels/embed.wgsl

// 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]));
  }
}