{% if source is defined and source.mode is defined %}
{% set mode = source.mode %}
{% 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) %}
{% if usesF16 is defined and usesF16 %}enable f16;
{% endif %}{{ 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 a vec4 along d with coalesced stores.
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>,
  @builtin(num_workgroups) nwg: vec3<u32>
) {
  // The dispatch folds oversized one-dimensional grids into x/y. Rebuild the
  // flat invocation index; this reduces to gid.x when no fold is needed.
  let index = gid.x + gid.y * nwg.x * WG;
{% 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 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) %}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>,
        {% endif %}{% if cooperativeMode %}@builtin(workgroup_id) wid: vec3<u32>,
        @builtin(local_invocation_id) lid: vec3<u32>,
        {% endif %}
        @builtin(num_workgroups) nwg: vec3<u32>) {
{% 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.
  let i = gid.x + gid.y * nwg.x * WG;
  if (i >= params.count) { return; }
  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 * nwg.x;
  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.
  let i = gid.x + gid.y * nwg.x * WG;
  if (i >= params.count) { return; }
  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.
  let i = gid.x + gid.y * nwg.x * WG;
  if (i >= params.count) { return; }
  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.
  let i = gid.x + gid.y * nwg.x * WG;
  if (i >= params.count) { return; }
  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" %}
  // Windowed cache: `params.seq` is a fixed capacity C, and the invariant is that
  // the L = min(T, C) most recent tokens live contiguously at rows [0, L).
  //
  //   residentBefore = min(T - S, C)      E = max(0, residentBefore + S - C)
  //   rows [0, appendStart)          <- past rows shifted down by E
  //   rows [appendStart, append+S)   <- the new tokens
  //   rows beyond that               <- outside the window, cleared
  //
  // The shift reads past and writes present, which must be distinct buffers: an
  // in-place compaction would have one invocation overwrite row t while another
  // still needs it as the source for row t-E, with no ordering between them.
  //
  // Attention needs no change for this layout. It derives its key range from
  // `min(capacity, seqlens_k[b] + 1)`, which is exactly L, and both the causal
  // and local-window masks depend only on the query/key distance:
  //   q_abs - k_abs = (T - qSeq + s) - (origin + t) = (L - qSeq + s) - t
  // so scoring a windowed cache as if it were a full L-length one is the same
  // arithmetic. Only RoPE would need the true absolute position, which is why
  // `windowShiftOk` refuses a rotary request outright rather than silently
  // rotating at the cache row.
  let i = gid.x + gid.y * nwg.x * WG;
  if (i >= params.count) { return; }
  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];
    }
  } 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 {
    // 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);
    }
  }
{% else %}
  // The dispatch folds oversized one-dimensional grids into x/y. Rebuild the
  // flat invocation index; this reduces to gid.x when no fold is needed.
  let i = gid.x + gid.y * nwg.x * WG;
  if (i >= params.count) { return; }
  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 %}