| {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %} |
| {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %} |
| // 2D-folded flat index: gid.y carries the high bits past the dispatch's |
| // per-axis workgroup fold width. |
| let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }}; |
| if ({{ name }} >= {{ bound }}) { |
| return; |
| }{% endmacro %} |
| {{ env.wgsl.resourceDeclarations }} |
| |
| // Cache maintenance for decode self attention, in two roles from one template. |
| // |
| // retain — copy the whole past buffer into the present buffer at capacity. |
| // The rows beyond the live length are copied too: the operator's |
| // contract is that present equals past everywhere except the one |
| // appended row, and upstream's own kernel memcpys the full buffer. |
| // The pass is a pure copy, so it declares `viewAlias`: it is elided |
| // when `past_key` and `present_key` resolve to one buffer. |
| // append — write this token's key and value rows at index past_sequence_length, |
| // which arrives as tensor data and so is read here rather than baked. |
| // |
| // Key and value share one dispatch: they have identical geometry, and one |
| // invocation touching both halves keeps the index arithmetic in registers. |
| const WG: u32 = {{ workgroupSize }}u; |
| {% if appendRow %} |
| const HEAD_SIZE_V4: u32 = {{ headSizeV4 }}u; |
| const NUM_HEADS: u32 = {{ numHeads }}u; |
| const HIDDEN_V4: u32 = NUM_HEADS * HEAD_SIZE_V4; |
| |
| @compute @workgroup_size(WG) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| // One invocation owns one vec4 of one (batch, head) row. |
| {{ flat_index_2d("WG", "index", "params.rowCountV4") }} |
| let d4 = index % HEAD_SIZE_V4; |
| let head = (index / HEAD_SIZE_V4) % NUM_HEADS; |
| let batch = index / HIDDEN_V4; |
| // past_sequence_length is a single value shared by every batch row. |
| let slot = u32(past_sequence_length[0]); |
| if (slot >= params.capacity) { |
| return; |
| } |
| let cacheIndex = ((batch * NUM_HEADS + head) * params.capacity + slot) * HEAD_SIZE_V4 + d4; |
| present_key[cacheIndex] = key[index]; |
| present_value[cacheIndex] = value[index]; |
| } |
| {% else %} |
| // Each lane stages ITEMS words of each cache before storing any of them, so that |
| // many loads are in flight before it has to wait. One word per lane leaves the |
| // copy latency-bound wherever the caches stay resident in cache; where the copy is |
| // DRAM-bound the depth is neutral. Consecutive lanes still read consecutive words; |
| // the stride between a lane's own items is a whole workgroup. ITEMS is bounded so |
| // the dispatch keeps a floor number of workgroups, since raising it divides them. |
| const ITEMS: u32 = {{ itemsPerLane }}u; |
| |
| @compute @workgroup_size(WG) |
| fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) { |
| let block = (wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u) * WG * ITEMS + lid.x; |
| var stagedKey: array<vec4<{{ cacheScalar }}>, ITEMS>; |
| var stagedValue: array<vec4<{{ cacheScalar }}>, ITEMS>; |
| {% for i in range(itemsPerLane) %} |
| { |
| let index = block + {{ i }}u * WG; |
| if (index < params.countV4) { |
| stagedKey[{{ i }}] = past_key[index]; |
| stagedValue[{{ i }}] = past_value[index]; |
| } |
| } |
| {% endfor %} |
| {% for i in range(itemsPerLane) %} |
| { |
| let index = block + {{ i }}u * WG; |
| if (index < params.countV4) { |
| present_key[index] = stagedKey[{{ i }}]; |
| present_value[index] = stagedValue[{{ i }}]; |
| } |
| } |
| {% endfor %} |
| } |
| {% endif %} |
| |
