{% if modeSpec is defined %}
{% set mode = modeSpec %}
{% elif mode is not defined %}
{% set mode = "transpose" %}
{% endif %}
{% set presentScalar = presentScalar | default(inputScalar | default("f32")) %}
{% set presentElement = presentElement | default(presentScalar) %}
{% set packed = packed | default(0) %}
{% set bits = bits | default(0) %}
{% set qmax = qmax | default(0) %}
{% set qmin = qmin | default(0) %}
{% 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.
{% if bound == "" %}
  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% else %}
  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
  if ({{ name }} >= {{ bound }}) { return; }{% endif %}{% endmacro %}
{{ env.wgsl.resourceDeclarations }}
{% if mode == "transpose" %}

const KV_HEADS: u32 = {{ kvNumHeads }}u;
const WG: u32 = {{ copyWorkgroupSize }}u;
{% if presentVec4 %}
// Transpose copy [batch, token, head*headDim] -> present [batch, head, token,
// headDim]. The headDim row is contiguous on both sides and divisible by four,
// so each invocation copies one vec4 along d and adjacent invocations write
// adjacent destination vectors.
const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
{% else %}
const HEAD_DIM: u32 = {{ headDim }}u;
const KV_HIDDEN: u32 = {{ kvHidden }}u;
{% endif %}

@compute @workgroup_size(WG, 1, 1)
fn main(
  @builtin(global_invocation_id) gid: vec3<u32>
) {
{{ flat_index_2d("WG", "index", "") }}
{% if presentVec4 %}
  let total = params.batchSize * KV_HEADS * params.kvSeq * HEAD_DIM_V4;
  if (index >= total) {
    return;
  }
  let d4 = index % HEAD_DIM_V4;
  let token = (index / HEAD_DIM_V4) % params.kvSeq;
  let head = (index / (HEAD_DIM_V4 * params.kvSeq)) % KV_HEADS;
  let batch = index / (HEAD_DIM_V4 * params.kvSeq * KV_HEADS);
  let packed = (batch * params.kvSeq + token) * KV_HIDDEN_V4 + head * HEAD_DIM_V4 + d4;
  present_key[index] = {{ presentElement }}(key[packed]);
  present_value[index] = {{ presentElement }}(value[packed]);
{% else %}
  let total = params.batchSize * KV_HEADS * params.kvSeq * HEAD_DIM;
  if (index >= total) {
    return;
  }
  let d = index % HEAD_DIM;
  let token = (index / HEAD_DIM) % params.kvSeq;
  let head = (index / (HEAD_DIM * params.kvSeq)) % KV_HEADS;
  let batch = index / (HEAD_DIM * params.kvSeq * KV_HEADS);
  let packed_index = (batch * params.kvSeq + token) * KV_HIDDEN + head * HEAD_DIM + d;
  present_key[index] = {{ presentScalar }}(key[packed_index]);
  present_value[index] = {{ presentScalar }}(value[packed_index]);
{% endif %}
}
{% else %}

{% set shareAppend = mode == "merge_share" and shareRegion == "append" %}
{% set shiftRegion = shiftRegion | default("all") %}
{% set shiftRetainOnly = mode == "window_shift" and shiftRegion == "shift" %}
{% set cooperativeCopy = device.adapterInfo.vendor != "apple" %}
{% set cooperativeMerge = cooperativeCopy or shareAppend %}
{% set cooperativeMode = cooperativeMerge and (mode == "copy" or mode == "merge" or mode == "merge_share") %}
{% set needsGid = not cooperativeMode or (shareAppend and not cooperativeCopy) %}
{% if mode == "copy" %}
// Copy the existing BNSH cache unchanged. Compression changes only the number
// of stored words per row.
{% elif mode == "merge" %}
// Copy the past BNSH cache, then append new BSH K/V tokens into the present cache.
{% elif mode == "merge_share" %}
// Buffer-sharing append: past and present share the full-capacity BNSH stride.
// Rows outside each batch's append window [seqlens_k[b]+1-keySeq, seqlens_k[b]+1)
// are an identity copy of the past cache; rows inside it come from the new BSH
// K/V (K optionally rotated at its absolute position).
// The append and retain passes own disjoint ranges, so their order is immaterial.
{% elif mode == "build" %}
// Build a BNSH cache from new BSH K/V, optionally applying K RMSNorm and rotary.
{% elif mode == "window_shift" %}
// Windowed cache: compact the surviving past rows down by the eviction count and
// append the new BSH K/V, keeping the most recent min(T, capacity) tokens
// contiguous at rows [0, L). Reads past, writes present (distinct buffers).
{% elif mode == "append_quant" %}
// Quantize new BSH K/V into an existing buffer-sharing INT8 BNSH cache. The
// active sequence end comes from seqlens_k; the physical cache stride remains
// the full present-cache capacity.
{% else %}
// Quantize new BSH K/V into an INT8 or packed INT4 BNSH cache.
{% endif %}
const HEAD_DIM: u32 = {{ headDim }}u;
{% if mode != "copy" %}const KV_HEADS: u32 = {{ kvHeads }}u;
{% endif %}
const WG: u32 = {{ tunables.COPY_WORKGROUP_SIZE }}u;
{% if mode != "copy" and (mode != "merge_share" or shareAppend) and not shiftRetainOnly %}const KV_HIDDEN: u32 = {{ kvHeads }}u * {{ headDim }}u;
{% endif %}
{% if mode == "build_quant" or mode == "append_quant" %}const PACKED: u32 = {{ packed }}u;
{% endif %}
{% if hasRotary and mode != "copy" and (mode != "merge_share" or shareAppend) %}const HALF: u32 = {{ half }}u;
{% endif %}
{% if hasKNorm and mode != "copy" and (mode != "merge_share" or shareAppend) %}const QK_EPS: f32 = {{ qkEps }};
{% endif %}
{% if mode == "build_quant" or mode == "append_quant" %}
const QMAX: f32 = {{ qmax }};
const QMIN: f32 = {{ qmin }};
fn kscale(d: u32, hk: u32) -> f32 { return k_scale[select(0u, hk * HEAD_DIM + d, params.perChannel != 0u)]; }
fn vscale(d: u32, hk: u32) -> f32 { return v_scale[select(0u, hk * HEAD_DIM + d, params.perChannel != 0u)]; }
{% endif %}

@compute @workgroup_size(WG)
fn main({% if needsGid %}@builtin(global_invocation_id) gid: vec3<u32>{% if cooperativeMode %},
        {% endif %}{% endif %}{% if cooperativeMode %}@builtin(workgroup_id) wid: vec3<u32>,
        @builtin(local_invocation_id) lid: vec3<u32>{% endif %}) {
{% macro append_cooperative_walk() %}
  // Compact cooperative append walk: the thread range spans only the appended
  // window (batch x kvHead x keySeq rows), not the full cache capacity; the
  // destination row is the dynamic append offset plus the local row.
  // The loop is bounded by the elements this workgroup owns rather than a fixed
  // HEAD_DIM trip count.
  let wgEnd = min(elemBase + WG * HEAD_DIM, totalElems);
  for (var e = elemBase + lid.x; e < wgEnd; e = e + WG) {
    {
      let d = e % HEAD_DIM;
      let et = e / HEAD_DIM;
      let j = et % params.keySeq;
      let etmp = et / params.keySeq;
      let hk = etmp % KV_HEADS;
      let b = etmp / KV_HEADS;
      // Clamping below by keySeq maps a right-padded first prompt
      // (seqlens_k[b]+1 < keySeq) to append offset 0.
      let activeEnd = min(params.seq, max(params.keySeq, u32(seqlens_k[b]) + 1u));
      let t = activeEnd - params.keySeq + j;
      let dst = ((b * KV_HEADS + hk) * params.seq + t) * HEAD_DIM + d;
      let nSrc = (b * params.keySeq + j) * KV_HIDDEN + hk * HEAD_DIM;
{% if hasRotary %}
      // Element-wise NeoX rotation: the pair (dr, dr+HALF) shares one
      // cos/sin entry, and row t is its own absolute position.
      let dr = select(d - HALF, d, d < HALF);
      let cs = cos_cache[t * HALF + dr];
      let sn = sin_cache[t * HALF + dr];
      let x0 = new_k[nSrc + dr];
      let x1 = new_k[nSrc + dr + HALF];
      present_key[dst] = select(x1 * cs + x0 * sn, x0 * cs - x1 * sn, d < HALF);
{% else %}
      present_key[dst] = new_k[nSrc + d];
{% endif %}
      present_value[dst] = new_v[nSrc + d];
    }
  }
{% endmacro %}
{% macro append_per_row_walk() %}
  // Per-row append walk: one thread streams its row's contiguous bytes.
{{ flat_index_2d("WG", guardInline=true) }}
  let j = i % params.keySeq;
  let tmp = i / params.keySeq;
  let hk = tmp % KV_HEADS;
  let b = tmp / KV_HEADS;
  // Clamping below by keySeq maps a right-padded first prompt
  // (seqlens_k[b]+1 < keySeq) to append offset 0.
  let activeEnd = min(params.seq, max(params.keySeq, u32(seqlens_k[b]) + 1u));
  let t = activeEnd - params.keySeq + j;
  let dstBase = ((b * KV_HEADS + hk) * params.seq + t) * HEAD_DIM;
  let nSrc = (b * params.keySeq + j) * KV_HIDDEN + hk * HEAD_DIM;
{% if hasRotary %}
  var k: array<f32, HEAD_DIM>;
  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { k[d] = f32(new_k[nSrc + d]); }
  // Row t is its own absolute position in the shared cache.
  for (var d = 0u; d < HALF; d = d + 1u) {
    let cs = cos_cache[t * HALF + d];
    let sn = sin_cache[t * HALF + d];
    let x0 = k[d];
    let x1 = k[d + HALF];
    k[d] = x0 * cs - x1 * sn;
    k[d + HALF] = x1 * cs + x0 * sn;
  }
  for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
    present_key[dstBase + d] = k[d];
    present_value[dstBase + d] = new_v[nSrc + d];
  }
{% else %}
  for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
    present_key[dstBase + d] = new_k[nSrc + d];
    present_value[dstBase + d] = new_v[nSrc + d];
  }
{% endif %}
{% endmacro %}
{% if cooperativeMerge and (mode == "copy" or mode == "merge" or mode == "merge_share") %}
  // Each workgroup copies the contiguous element span of its WG rows
  // cooperatively: at step s lane l touches element s*WG + l, so adjacent
  // lanes hit adjacent addresses on both source and destination. Same
  // dispatch geometry (ceil(rows/WG) workgroups, folded into x/y).
  let wgFlat = wid.x + wid.y * {{ DISPATCH_FOLD_WIDTH }}u;
  let elemBase = wgFlat * (WG * HEAD_DIM);
  let totalElems = params.count * HEAD_DIM;
{% endif %}
{% if mode == "copy" and cooperativeCopy %}
  // Copy mode is an identity copy over the whole BNSH range: dst index == src index.
  for (var s = 0u; s < HEAD_DIM; s = s + 1u) {
    let e = elemBase + s * WG + lid.x;
    if (e < totalElems) {
      present_key[e] = src_k[e];
      present_value[e] = src_v[e];
    }
  }
{% elif mode == "merge" and cooperativeCopy %}
  // present[0:pastSeq] = past (BNSH copy); present[pastSeq:totalSeq] = new K/V (BSH appended).
  for (var s = 0u; s < HEAD_DIM; s = s + 1u) {
    let e = elemBase + s * WG + lid.x;
    if (e < totalElems) {
      let d = e % HEAD_DIM;
      let et = e / HEAD_DIM;
      let t = et % params.seq;
      let etmp = et / params.seq;
      let hk = etmp % KV_HEADS;
      let b = etmp / KV_HEADS;
      if (t < params.pastSeq) {
        let src = ((b * KV_HEADS + hk) * params.pastSeq + t) * HEAD_DIM + d;
        present_key[e] = past_k[src];
        present_value[e] = past_v[src];
      } else {
        let nSrc = (b * params.keySeq + (t - params.pastSeq)) * KV_HIDDEN + hk * HEAD_DIM + d;
        present_key[e] = new_k[nSrc];
        present_value[e] = new_v[nSrc];
      }
    }
  }
{% elif mode == "merge_share" and cooperativeMerge %}
{% if shareAppend %}
{% if cooperativeCopy %}
{{ append_cooperative_walk() }}
{% else %}
  if (params.count < WG) {
{{ append_cooperative_walk() }}
    return;
  }
{{ append_per_row_walk() }}
{% endif %}
{% else %}
  // Cooperative element walk (see the copy branch); the same-stride layout
  // makes the outside-window rows an identity copy (dst index == src index).
  for (var s = 0u; s < HEAD_DIM; s = s + 1u) {
    let e = elemBase + s * WG + lid.x;
    if (e < totalElems) {
      let et = e / HEAD_DIM;
      let t = et % params.seq;
      let b = et / (params.seq * KV_HEADS);
      // Clamping below by keySeq maps a right-padded first prompt
      // (seqlens_k[b]+1 < keySeq) to append offset 0.
      let activeEnd = min(params.seq, max(params.keySeq, u32(seqlens_k[b]) + 1u));
      let appendStart = activeEnd - params.keySeq;
      // Everything outside the append window. The append pass owns the window itself,
      // and the shared full-capacity stride makes this an identity copy (dst == src).
      if (t < appendStart || t >= activeEnd) {
        present_key[e] = past_k[e];
        present_value[e] = past_v[e];
      }
    }
  }
{% endif %}
{% elif mode == "copy" %}
  // Thread i streams one contiguous (batch, kvHead, token) row.
{{ flat_index_2d("WG", guardInline=true) }}
  let base = i * HEAD_DIM;
  for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
    present_key[base + d] = src_k[base + d];
    present_value[base + d] = src_v[base + d];
  }
{% elif mode == "merge" %}
  // Per-row merge walk.
{{ flat_index_2d("WG", guardInline=true) }}
  let t = i % params.seq;
  let tmp = i / params.seq;
  let hk = tmp % KV_HEADS;
  let b = tmp / KV_HEADS;
  // present[0:pastSeq] = past (BNSH copy); present[pastSeq:totalSeq] = new K/V (BSH appended).
  let dstBase = ((b * KV_HEADS + hk) * params.seq + t) * HEAD_DIM;
  if (t < params.pastSeq) {
    let srcBase = ((b * KV_HEADS + hk) * params.pastSeq + t) * HEAD_DIM;
    for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
      present_key[dstBase + d] = past_k[srcBase + d];
      present_value[dstBase + d] = past_v[srcBase + d];
    }
  } else {
    let nSrc = (b * params.keySeq + (t - params.pastSeq)) * KV_HIDDEN + hk * HEAD_DIM;
    for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
      present_key[dstBase + d] = new_k[nSrc + d];
      present_value[dstBase + d] = new_v[nSrc + d];
    }
  }
{% elif mode == "merge_share" %}
  // Past and present share the full-capacity stride, so outside-window rows
  // copy at the same index.
{{ flat_index_2d("WG", guardInline=true) }}
  let t = i % params.seq;
  let tmp = i / params.seq;
  let hk = tmp % KV_HEADS;
  let b = tmp / KV_HEADS;
  // Clamping below by keySeq maps a right-padded first prompt
  // (seqlens_k[b]+1 < keySeq) to append offset 0.
  let activeEnd = min(params.seq, max(params.keySeq, u32(seqlens_k[b]) + 1u));
  let appendStart = activeEnd - params.keySeq;
  let dstBase = ((b * KV_HEADS + hk) * params.seq + t) * HEAD_DIM;
  // Everything outside the append window. The append pass owns the window itself,
  // and the shared full-capacity stride makes this an identity copy (dst == src).
  if (t < appendStart || t >= activeEnd) {
    for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
      present_key[dstBase + d] = past_k[dstBase + d];
      present_value[dstBase + d] = past_v[dstBase + d];
  }
  }
{% elif mode == "window_shift" %}
  // Keep the most recent min(total, capacity) tokens in contiguous cache rows.
  // Read past and write a distinct present buffer so compaction cannot race.
  // Rotary append uses the absolute token position; retained keys are already rotated.
{% if shiftRegion == "append" %}
  // Append-only pass: the thread range spans just the appended window (batch x
  // kvHead x keySeq rows), while `params.seq` stays the cache capacity. The shift
  // pass owns every other row and the two ranges are disjoint, so their order is
  // immaterial.
{{ flat_index_2d("WG", guardInline=true) }}
  let j = i % params.keySeq;
  let tmp = i / params.keySeq;
  let hk = tmp % KV_HEADS;
  let b = tmp / KV_HEADS;
  let absTotal = max(params.keySeq, u32(seqlens_k[b]) + 1u);
  let residentBefore = min(absTotal - params.keySeq, params.seq);
  // `max(0u, a + b - c)` does not clamp in u32: the subtraction wraps first, so a
  // step that evicts nothing reads back ~2^32 instead of 0. Compare before
  // subtracting.
  let filled = residentBefore + params.keySeq;
  let evicted = select(0u, filled - params.seq, filled > params.seq);
  let appendStart = residentBefore - evicted;
  let t = appendStart + j;
  let dstBase = ((b * KV_HEADS + hk) * params.seq + t) * HEAD_DIM;
  let nSrc = (b * params.keySeq + j) * KV_HIDDEN + hk * HEAD_DIM;
  // Row appendStart + j carries absolute position absTotal - keySeq + j: the cache
  // holds only the last L positions, so the row index is short of the absolute one
  // by the window origin. The reference rotates the same rows at
  // `appendStart + cacheOrigin`, which is the same number.
  let pos = absTotal - params.keySeq + j;
  var k: array<f32, HEAD_DIM>;
  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { k[d] = new_k[nSrc + d]; }
  for (var d = 0u; d < HALF; d = d + 1u) {
    let cs = cos_cache[pos * HALF + d];
    let sn = sin_cache[pos * HALF + d];
    let x0 = k[d];
    let x1 = k[d + HALF];
    k[d] = x0 * cs - x1 * sn;
    k[d + HALF] = x1 * cs + x0 * sn;
  }
  for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
    present_key[dstBase + d] = k[d];
    present_value[dstBase + d] = new_v[nSrc + d];
  }
{% else %}
{{ flat_index_2d("WG", guardInline=true) }}
  let t = i % params.seq;
  let tmp = i / params.seq;
  let hk = tmp % KV_HEADS;
  let b = tmp / KV_HEADS;
  let absTotal = max(params.keySeq, u32(seqlens_k[b]) + 1u);
  let residentBefore = min(absTotal - params.keySeq, params.seq);
  // `max(0u, a + b - c)` does not clamp in u32: the subtraction wraps first, so a
  // step that evicts nothing reads back ~2^32 instead of 0. Compare before
  // subtracting.
  let filled = residentBefore + params.keySeq;
  let evicted = select(0u, filled - params.seq, filled > params.seq);
  let appendStart = residentBefore - evicted;
  let dstBase = ((b * KV_HEADS + hk) * params.seq + t) * HEAD_DIM;
  if (t < appendStart) {
    // Surviving past rows slide down by the eviction count.
    let srcBase = ((b * KV_HEADS + hk) * params.seq + (t + evicted)) * HEAD_DIM;
    for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
      present_key[dstBase + d] = past_k[srcBase + d];
      present_value[dstBase + d] = past_v[srcBase + d];
    }
{% if shiftRegion == "shift" %}
  } else if (t >= appendStart + params.keySeq) {
{% else %}
  } else if (t < appendStart + params.keySeq) {
    let nSrc = (b * params.keySeq + (t - appendStart)) * KV_HIDDEN + hk * HEAD_DIM;
    for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
      present_key[dstBase + d] = new_k[nSrc + d];
      present_value[dstBase + d] = new_v[nSrc + d];
    }
  } else {
{% endif %}
    // Outside the resident window. Cleared rather than left stale so the present
    // buffer does not expose stale contents from its distinct allocation.
    for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
      present_key[dstBase + d] = {{ zeroScalar }}(0.0);
      present_value[dstBase + d] = {{ zeroScalar }}(0.0);
    }
  }
{% endif %}
{% else %}
{{ flat_index_2d("WG", guardInline=true) }}
  let t = i % params.seq;
  let tmp = i / params.seq;
  let hk = tmp % KV_HEADS;
  let b = tmp / KV_HEADS;
  let srcBase = (b * params.seq + t) * KV_HIDDEN + hk * HEAD_DIM; // BSH new K/V
  var k: array<f32, HEAD_DIM>;
  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { k[d] = src_k[srcBase + d]; }
{% if hasKNorm %}
  var ms = 0.0;
  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { ms = ms + k[d] * k[d]; }
  let invRms = inverseSqrt(ms / f32(HEAD_DIM) + QK_EPS);
  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { k[d] = k[d] * invRms * k_norm_weight[d]; }
{% endif %}
{% if hasRotary %}
  let pos = params.pastSeq + t;
  for (var d = 0u; d < HALF; d = d + 1u) {
    let cs = cos_cache[pos * HALF + d];
    let sn = sin_cache[pos * HALF + d];
    let x0 = k[d];
    let x1 = k[d + HALF];
    k[d] = x0 * cs - x1 * sn;
    k[d + HALF] = x1 * cs + x0 * sn;
  }
{% endif %}
{% if mode == "build" %}
  let dstBase = ((b * KV_HEADS + hk) * params.seq + t) * HEAD_DIM;
  for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
    present_key[dstBase + d] = k[d];
    present_value[dstBase + d] = src_v[srcBase + d];
  }
{% else %}
{% if mode == "append_quant" %}
  let activeEnd = u32(seqlens_k[b]) + 1u;
  let appendStart = activeEnd - params.seq;
  let dstBase = ((b * KV_HEADS + hk) * params.totalSeq + appendStart + t) * PACKED;
{% else %}
  let dstBase = ((b * KV_HEADS + hk) * params.seq + t) * PACKED;
{% endif %}
{% if bits == 8 %}
  for (var d = 0u; d < HEAD_DIM; d = d + 1u) {
    present_key[dstBase + d] = i32(clamp(round(k[d] / kscale(d, hk)), QMIN, QMAX));
    present_value[dstBase + d] = i32(clamp(round(src_v[srcBase + d] / vscale(d, hk)), QMIN, QMAX));
  }
{% else %}
  for (var dd = 0u; dd < PACKED; dd = dd + 1u) {
    let d0 = dd * 2u;
    let d1 = dd * 2u + 1u;
    let qk0 = i32(clamp(round(k[d0] / kscale(d0, hk)), QMIN, QMAX));
    let qk1 = i32(clamp(round(k[d1] / kscale(d1, hk)), QMIN, QMAX));
    let qv0 = i32(clamp(round(src_v[srcBase + d0] / vscale(d0, hk)), QMIN, QMAX));
    let qv1 = i32(clamp(round(src_v[srcBase + d1] / vscale(d1, hk)), QMIN, QMAX));
    present_key[dstBase + dd] = (u32(qk0 + 8) & 0xFu) | ((u32(qk1 + 8) & 0xFu) << 4u);
    present_value[dstBase + dd] = (u32(qv0 + 8) & 0xFu) | ((u32(qv1 + 8) & 0xFu) << 4u);
  }
{% endif %}
{% endif %}
{% endif %}
}
{% endif %}