| {{ env.wgsl.resourceDeclarations }} |
| |
| // Engram n-gram hash ids. For every order n in [2, MAX_NGRAM] the kernel mixes the n causal |
| // shifts of the compressed tokenizer ids, |
| // mix_n = ids[t] * m[0] ^ ids[t-1] * m[1] ^ ... ^ ids[t-n+1] * m[n-1], |
| // and emits mix_n modulo each of that order's head vocabulary sizes. mix_n extends mix_{n-1} |
| // by a single term, so one walk over the window produces every order. |
| // |
| // Each product is formed in u32 so that overflow wraps in two's complement, which is the |
| // arithmetic the operator specifies. A signed multiply that overflows has an indeterminate |
| // result in WGSL, so the bitcasts here are load-bearing rather than cosmetic. |
| // |
| // One dispatch covers both outputs. Workgroup row wg.y owns batch row b, and the tile index |
| // folded across wg.x/wg.z walks a span of SEQ + STATE_LENGTH slots: the first SEQ slots hash a |
| // token position, and the trailing STATE_LENGTH slots emit the carry state that the next call |
| // feeds back as past_ids. A slot past the end of the span exits, so a short sequence costs one |
| // workgroup rather than a second dispatch. |
| const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u; |
| const ITEMS: u32 = {{ tunables.ITEMS_PER_LANE }}u; |
| const STATE_LENGTH: u32 = {{ stateLength }}u; |
| const PAD_ID: i32 = {{ padId }}; |
| |
| // Euclidean remainder for a positive divisor: WGSL `%` truncates toward zero exactly as C |
| // does, so a negative dividend needs the one correction the operator specifies. A head whose |
| // vocabulary size is not strictly positive has no meaningful modulo, and clamping the divisor |
| // to 1 makes its remainder 0 -- the value every GPU implementation of this operator emits for |
| // that head, and the reason it stays a device-side division rather than a rejected input. |
| fn positive_mod(value: i32, divisor: i32) -> i32 { |
| let r = value % max(divisor, 1); |
| return select(r, r + divisor, r < 0); |
| } |
| {% macro hash_position(interior) %} |
| let idBase = inputBase + t; |
| var mix = bitcast<u32>(input_ids[idBase]) * mult0; |
| {% for k in range(1, maxNgramSize) %} |
| {% if interior %} |
| mix = mix ^ (bitcast<u32>(input_ids[idBase - {{ k }}u]) * mult{{ k }}); |
| {% else %} |
| var token{{ k }} = PAD_ID; |
| if (t >= {{ k }}u) { |
| token{{ k }} = input_ids[idBase - {{ k }}u]; |
| } |
| {% if hasPast %} |
| else { |
| // past_ids is right-aligned, so position -1 is its last slot. t < k bounds the index |
| // inside the window, which is why no second range test is needed here. |
| token{{ k }} = past_ids[pastBase + STATE_LENGTH + t - {{ k }}u]; |
| } |
| {% endif %} |
| mix = mix ^ (bitcast<u32>(token{{ k }}) * mult{{ k }}); |
| {% endif %} |
| let mixOrder{{ k }} = bitcast<i32>(mix); |
| {% endfor %} |
| {% for n in range(1, maxNgramSize) %} |
| {% for h in range(headsPerNgram) %} |
| let head{{ (n - 1) * headsPerNgram + h }} = positive_mod(mixOrder{{ n }}, vocab{{ (n - 1) * headsPerNgram + h }}); |
| {% endfor %} |
| {% endfor %} |
| let outBase = idBase * {{ numHeads }}u; |
| {% for h in range(numHeads) %} |
| hash_ids[outBase + {{ h }}u] = head{{ h }}; |
| {% endfor %} |
| {% endmacro %} |
| |
| @compute @workgroup_size(WG) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) { |
| // wg.y is a real axis -- one workgroup row per batch row -- so only the span tiles are |
| // folded, across wg.x and wg.z at the per-axis workgroup fold width. |
| let b = wg.y; |
| let tile = wg.x + wg.z * {{ DISPATCH_FOLD_WIDTH }}u; |
| let seq = params.seqLength; |
| let span = seq + STATE_LENGTH; |
| let inputBase = b * seq; |
| let tileBase = tile * (WG * ITEMS); |
| {% if hasPast %} |
| let pastBase = b * STATE_LENGTH; |
| {% endif %} |
| |
| // The multiplier and vocabulary tables are attribute-sized and every position reads all of |
| // them, so an invocation stages them once instead of reloading them inside the mix walk. |
| {% for k in range(maxNgramSize) %} |
| let mult{{ k }} = bitcast<u32>(multipliers[{{ k }}u]); |
| {% endfor %} |
| {% for h in range(numHeads) %} |
| let vocab{{ h }} = vocab_sizes[{{ h }}u]; |
| {% endfor %} |
| |
| // An interior tile is one whose every slot hashes a token position and whose every causal |
| // shift lands inside input_ids. Both of the per-slot tests below, and the per-shift history |
| // fallback, are then statically false, which is worth one workgroup-uniform branch. |
| if (tileBase >= STATE_LENGTH && tileBase + WG * ITEMS <= seq) { |
| for (var item = 0u; item < ITEMS; item = item + 1u) { |
| let t = tileBase + lid.x + item * WG; |
| {{ hash_position(true) }} |
| } |
| return; |
| } |
| |
| for (var item = 0u; item < ITEMS; item = item + 1u) { |
| let slot = tileBase + lid.x + item * WG; |
| if (slot >= span) { |
| break; |
| } |
| if (slot < seq) { |
| let t = slot; |
| {{ hash_position(false) }} |
| } else { |
| // Carry slot j holds the id at position seq - STATE_LENGTH + j of this call, and slot j |
| // of past_ids when the call is shorter than the window. Writing it from this dispatch is |
| // safe because present_ids never shares its buffer with past_ids here. |
| var carry = PAD_ID; |
| if (slot >= STATE_LENGTH) { |
| carry = input_ids[inputBase + slot - STATE_LENGTH]; |
| } |
| {% if hasPast %} |
| else { |
| carry = past_ids[pastBase + slot]; |
| } |
| {% endif %} |
| present_ids[b * STATE_LENGTH + slot - seq] = carry; |
| } |
| } |
| } |
| |
