{% 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 %}