| {{ env.wgsl.resourceDeclarations }} |
| |
| const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u; |
| const TILE: u32 = {{ sharedTileSlots }}u; |
| const STATE_LENGTH: u32 = {{ stateLength }}u; |
| const HEADS: u32 = {{ numHeads }}u; |
| const HEADS_PER_ORDER: u32 = {{ headsPerNgram }}u; |
| {% if not hasPast %} |
| const PAD_ID: i32 = {{ padId }}; |
| {% endif %} |
| // Order-major storage coalesces prefix writes and shares each value across heads. |
| var<workgroup> prefixes: array<i32, {{ sharedTileSlots * stateLength }}>; |
| |
| fn positive_mod(value: i32, divisor: i32) -> i32 { |
| let r = value % max(divisor, 1); |
| return select(r, r + divisor, r < 0); |
| } |
| @compute @workgroup_size(WG) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) { |
| let b = wg.y; |
| let tile = wg.x + wg.z * {{ DISPATCH_FOLD_WIDTH }}u; |
| let seq = params.seqLength; |
| let tileBase = tile * TILE; |
| let inputBase = b * seq; |
| {% if hasPast %} |
| let pastBase = b * STATE_LENGTH; |
| {% endif %} |
| if (lid.x < TILE) { |
| let t = tileBase + lid.x; |
| if (t < seq) { |
| // Unsigned products preserve the specified two's-complement wrapping. |
| var mix = bitcast<u32>(input_ids[inputBase + t]) * bitcast<u32>(multipliers[0u]); |
| {% for k in range(1, maxNgramSize) %} |
| {% if hasPast %} |
| var token{{ k }}: i32; |
| if (t >= {{ k }}u) { token{{ k }} = input_ids[inputBase + t - {{ k }}u]; } |
| else { token{{ k }} = past_ids[pastBase + STATE_LENGTH + t - {{ k }}u]; } |
| {% else %} |
| var token{{ k }} = PAD_ID; |
| if (t >= {{ k }}u) { token{{ k }} = input_ids[inputBase + t - {{ k }}u]; } |
| {% endif %} |
| mix = mix ^ (bitcast<u32>(token{{ k }}) * bitcast<u32>(multipliers[{{ k }}u])); |
| prefixes[{{ k - 1 }}u * TILE + lid.x] = bitcast<i32>(mix); |
| {% endfor %} |
| } else if (t < seq + STATE_LENGTH) { |
| {% if hasPast %} |
| var carry: i32; |
| if (t >= STATE_LENGTH) { carry = input_ids[inputBase + t - STATE_LENGTH]; } |
| else { carry = past_ids[pastBase + t]; } |
| {% else %} |
| var carry = PAD_ID; |
| if (t >= STATE_LENGTH) { carry = input_ids[inputBase + t - STATE_LENGTH]; } |
| {% endif %} |
| present_ids[b * STATE_LENGTH + t - seq] = carry; |
| } |
| } |
| // Every lane reaches the barrier, including carry-only and partial tiles. |
| workgroupBarrier(); |
| var tokens = 0u; |
| if (tileBase < seq) { tokens = min(TILE, seq - tileBase); } |
| for (var index = lid.x; index < tokens * HEADS; index = index + WG) { |
| let token = index / HEADS; |
| let head = index % HEADS; |
| let mix = prefixes[(head / HEADS_PER_ORDER) * TILE + token]; |
| hash_ids[(inputBase + tileBase) * HEADS + index] = positive_mod(mix, vocab_sizes[head]); |
| } |
| } |
| |
