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