| {{ env.wgsl.resourceDeclarations }} |
| // Rotary positional embedding. The dispatch pins one slice on z and folds that slice's own |
| // rotation groups across x and y, so every row of a slice shares one base offset and all the |
| // head geometry below is a compile-time constant: every division in the index math is by a |
| // literal rather than by a uniform. |
| // |
| // A slice is one batch at rank 3, where a token's heads are adjacent and a workgroup covers a |
| // whole token, so it reads each cos/sin row once. At rank 4 a token's heads are a sequence |
| // apart, so a slice is instead one batch and ITEMS consecutive heads: the invocation gathers |
| // its cos/sin row once and rotates the same rotation group in each of those heads. |
| const WG: u32 = {{ workgroupSize }}u; |
| const ITEMS: u32 = {{ itemsPerLane }}u; |
| // Rotation groups per head row, the row's stride, and the number of groups that rotate, all |
| // counted in the storage element the bindings use: a vec4 on the four-pair route and a scalar |
| // on the single-pair route. |
| const HEAD_UNITS: u32 = {{ headUnits }}u; |
| const HEAD_STRIDE: u32 = {{ headStride }}u; |
| const ROT_UNITS: u32 = {{ rotUnits }}u; |
| {% if hasTail %} |
| // First element past the rotary window. |
| const ROT_END: u32 = {{ rotStart }}u; |
| {% endif %} |
| const NUM_HEADS: u32 = {{ numHeads }}u; |
| {% if rank == 4 %} |
| const HEAD_BLOCKS: u32 = {{ headBlocks }}u; |
| {% endif %} |
| {% if hasTail %} |
| |
| {% macro copy_tail(base) %} |
| let tail = ROT_END + (group - ROT_UNITS) * 2u; |
| y[{{ base }} + tail] = x[{{ base }} + tail]; |
| {% if tailGuard %} |
| // An odd head size leaves the last group with a single element to copy. |
| if (tail + 1u < HEAD_STRIDE) { |
| y[{{ base }} + tail + 1u] = x[{{ base }} + tail + 1u]; |
| } |
| {% else %} |
| y[{{ base }} + tail + 1u] = x[{{ base }} + tail + 1u]; |
| {% endif %} |
| {% endmacro %} |
| {% endif %} |
| |
| {% macro rotate(base) %} |
| {% if interleaved and useVec4 %} |
| // Four interleaved pairs span two vec4s: (w0.x, w0.y), (w0.z, w0.w), (w1.x, w1.y), |
| // (w1.z, w1.w). Deinterleave, rotate, and reinterleave on the way back out. |
| let pairBase = {{ base }} + group * 2u; |
| let w0 = vec4<f32>(x[pairBase]); |
| let w1 = vec4<f32>(x[pairBase + 1u]); |
| let a = vec4<f32>(w0.x, w0.z, w1.x, w1.z); |
| let b = vec4<f32>(w0.y, w0.w, w1.y, w1.w); |
| let ra = a * cf - b * sf; |
| let rb = a * sf + b * cf; |
| y[pairBase] = {{ castOut }}(vec4<f32>(ra.x, rb.x, ra.y, rb.y)); |
| y[pairBase + 1u] = {{ castOut }}(vec4<f32>(ra.z, rb.z, ra.w, rb.w)); |
| {% elif interleaved %} |
| // Adjacent elements form the rotated pair. |
| let pairBase = {{ base }} + group * 2u; |
| let a = {{ castIn }}(x[pairBase]); |
| let b = {{ castIn }}(x[pairBase + 1u]); |
| y[pairBase] = {{ castOut }}(a * cf - b * sf); |
| y[pairBase + 1u] = {{ castOut }}(a * sf + b * cf); |
| {% else %} |
| // Split halves: ROT_UNITS is both the rotating group count and the offset from an |
| // element of the first half to its partner in the second. |
| let a = {{ castIn }}(x[{{ base }} + group]); |
| let b = {{ castIn }}(x[{{ base }} + ROT_UNITS + group]); |
| y[{{ base }} + group] = {{ castOut }}(a * cf - b * sf); |
| y[{{ base }} + ROT_UNITS + group] = {{ castOut }}(a * sf + b * cf); |
| {% endif %} |
| {% endmacro %} |
| |
| @compute @workgroup_size(WG) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) { |
| {% if rank == 3 %} |
| // wg.y carries the high bits of the group index past the per-axis dispatch fold width. |
| let slice = wg.z; |
| let sliceUnits = params.sliceRows * HEAD_UNITS; |
| let sliceBase = slice * params.sliceRows * HEAD_STRIDE; |
| let laneUnit = (wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u) * (WG * ITEMS) + lid.x; |
| {% for item in range(itemsPerLane) %} |
| { |
| let unit = laneUnit + {{ item }}u * WG; |
| if (unit < sliceUnits) { |
| let group = unit % HEAD_UNITS; |
| let row = unit / HEAD_UNITS; |
| let base = sliceBase + row * HEAD_STRIDE; |
| let token = row / NUM_HEADS; |
| {% if hasTail %} |
| if (group >= ROT_UNITS) { |
| {{ copy_tail("base") }} |
| } else { |
| {% endif %} |
| {% if posOffset %} |
| // One base offset for the whole request: token s reads cache row p[0] + s. |
| let position = position_ids[0] + token; |
| {% else %} |
| // The (batch_size, sequence_length) table. params.sliceRows / NUM_HEADS is the sequence |
| // length, so this slice's table row is slice * sequence_length + token. |
| let position = position_ids[slice * (params.sliceRows / NUM_HEADS) + token]; |
| {% endif %} |
| let cache = position * ROT_UNITS + group; |
| let cf = {{ castIn }}(cos_cache[cache]); |
| let sf = {{ castIn }}(sin_cache[cache]); |
| {{ rotate("base") }} |
| {% if hasTail %} |
| } |
| {% endif %} |
| } |
| } |
| {% endfor %} |
| {% else %} |
| // A rank-4 slice is one batch and ITEMS consecutive heads; wg.y carries the high bits of |
| // the group index past the per-axis dispatch fold width. |
| let slice = wg.z; |
| let headBase = slice % HEAD_BLOCKS * ITEMS; |
| let batch = slice / HEAD_BLOCKS; |
| let sliceUnits = params.sliceRows * HEAD_UNITS; |
| let unit = (wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u) * WG + lid.x; |
| if (unit >= sliceUnits) { |
| return; |
| } |
| let group = unit % HEAD_UNITS; |
| let token = unit / HEAD_UNITS; |
| let rowBase = (batch * NUM_HEADS + headBase) * params.sliceRows + token; |
| {% if hasTail %} |
| if (group >= ROT_UNITS) { |
| {% for item in range(itemsPerLane) %} |
| { |
| let base = (rowBase + {{ item }}u * params.sliceRows) * HEAD_STRIDE; |
| {{ copy_tail("base") }} |
| } |
| {% endfor %} |
| return; |
| } |
| {% endif %} |
| {% if posOffset %} |
| // One base offset for the whole request: token s reads cache row p[0] + s. |
| let position = position_ids[0] + token; |
| {% else %} |
| // The (batch_size, sequence_length) table; params.sliceRows is the sequence length. |
| let position = position_ids[batch * params.sliceRows + token]; |
| {% endif %} |
| // Gathered once for the whole head block: cos and sin depend on the token and the rotation |
| // group, never on the head. |
| let cache = position * ROT_UNITS + group; |
| let cf = {{ castIn }}(cos_cache[cache]); |
| let sf = {{ castIn }}(sin_cache[cache]); |
| {% for item in range(itemsPerLane) %} |
| { |
| let base = (rowBase + {{ item }}u * params.sliceRows) * HEAD_STRIDE; |
| {{ rotate("base") }} |
| } |
| {% endfor %} |
| {% endif %} |
| } |
| |
