| {% macro sparse_schedule() %} |
| // How much history precedes this call's tokens. Prompt mode is determined from |
| // scalar total_sequence_length rather than per-batch values and forces the past |
| // length to zero; a padded prompt row must not be read as if it had history. |
| fn past_sequence_length(batch: u32) -> u32 { |
| if (u32(total_sequence_length[0]) == params.seqLen) { |
| return 0u; |
| } |
| let total = u32(key_total_sequence_lengths[batch]); |
| return select(0u, total - params.seqLen, total >= params.seqLen); |
| } |
| {% endmacro %} |
| {% macro sparse_rotary(interleaved) %} |
| // Which cos/sin entry a component uses, and which member of its rotation pair it is. |
| // The two layouts differ only here: the NeoX split pairs d with d + ROTARY_HALF, and the |
| // interleaved layout pairs the even lane with the odd one beside it. |
| fn rotary_slot(d: u32) -> u32 { |
| {% if interleaved %} |
| return d / 2u; |
| {% else %} |
| return select(d - ROTARY_HALF, d, d < ROTARY_HALF); |
| {% endif %} |
| } |
| |
| fn rotary_partner(d: u32) -> u32 { |
| {% if interleaved %} |
| return select(d - 1u, d + 1u, (d % 2u) == 0u); |
| {% else %} |
| return select(d - ROTARY_HALF, d + ROTARY_HALF, d < ROTARY_HALF); |
| {% endif %} |
| } |
| |
| fn rotary_is_first(d: u32) -> bool { |
| {% if interleaved %} |
| return (d % 2u) == 0u; |
| {% else %} |
| return d < ROTARY_HALF; |
| {% endif %} |
| } |
| |
| // One component of the rotation, written from the component's own point of view so a |
| // single expression covers both members of the pair: the leading one subtracts its |
| // partner's sine term, the trailing one adds it. |
| fn rotary_value(own: f32, partner: f32, cs: f32, sn: f32, first: bool) -> f32 { |
| return select(own * cs + partner * sn, own * cs - partner * sn, first); |
| } |
| {% endmacro %} |
| {% 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 }};{% endmacro %} |
| {{ env.wgsl.resourceDeclarations }} |
| |
| // com.microsoft.SparseAttention, KV append pass. |
| // past_key and present_key are the same allocation, so this writes only the new rows: |
| // one invocation per (batch, kv head, new token, component) of the BNSH cache, landing |
| // at absolute position past + t. Nothing rewrites the retained history. |
| const KV_HEADS: u32 = {{ kvNumHeads }}u; |
| const HEAD_DIM: u32 = {{ headSize }}u; |
| const MAX_CACHE_SEQ: u32 = {{ maxCacheSeq }}u; |
| {% if packedQkv %} |
| // Packed layout: query carries [Q | K | V] on one row, so K starts after the query |
| // heads and V after the key heads. |
| const Q_STRIDE: u32 = {{ packedStride }}u; |
| const PACKED_K_OFFSET: u32 = {{ numHeads * headSize }}u; |
| const PACKED_V_OFFSET: u32 = {{ (numHeads + kvNumHeads) * headSize }}u; |
| {% else %} |
| const KV_HIDDEN: u32 = {{ kvNumHeads * headSize }}u; |
| {% endif %} |
| const WG: u32 = {{ appendWorkgroupSize }}u; |
| {% if usesRotary %} |
| const ROTARY_HALF: u32 = {{ rotaryHalf }}u; |
| const ROTARY_DIM: u32 = {{ rotaryDim }}u; |
| {% endif %} |
| |
| {{ sparse_schedule() }} |
| {% if usesRotary %} |
| {{ sparse_rotary(rotaryInterleaved) }} |
| {% endif %} |
| @compute @workgroup_size(WG, 1, 1) |
| fn main(@builtin(global_invocation_id) gid: vec3<u32>) { |
| {{ flat_index_2d("WG", "index", "") }} |
| let count = params.batchSize * KV_HEADS * params.seqLen * HEAD_DIM; |
| if (index >= count) { |
| return; |
| } |
| let d = index % HEAD_DIM; |
| let t = (index / HEAD_DIM) % params.seqLen; |
| let kv_head = (index / (HEAD_DIM * params.seqLen)) % KV_HEADS; |
| let batch = index / (HEAD_DIM * params.seqLen * KV_HEADS); |
| |
| let position = past_sequence_length(batch) + t; |
| let dst = ((batch * KV_HEADS + kv_head) * MAX_CACHE_SEQ + position) * HEAD_DIM + d; |
| {% if packedQkv %} |
| let row = (batch * params.seqLen + t) * Q_STRIDE; |
| let k_base = row + PACKED_K_OFFSET + kv_head * HEAD_DIM; |
| let v_base = row + PACKED_V_OFFSET + kv_head * HEAD_DIM; |
| {% else %} |
| let k_base = ((batch * params.seqLen + t) * KV_HIDDEN) + kv_head * HEAD_DIM; |
| let v_base = k_base; |
| {% endif %} |
| |
| {% if usesRotary %} |
| // The key is rotated at its own absolute position before it enters the cache, so the |
| // cache holds post-rotary keys and the attention pass never rotates again. |
| if (d < ROTARY_DIM) { |
| let slot = rotary_slot(d); |
| let cs = f32(cos_cache[position * ROTARY_HALF + slot]); |
| let sn = f32(sin_cache[position * ROTARY_HALF + slot]); |
| let own = f32({{ kvSource }}[k_base + d]); |
| let partner = f32({{ kvSource }}[k_base + rotary_partner(d)]); |
| present_key[dst] = {{ scalar }}(rotary_value(own, partner, cs, sn, rotary_is_first(d))); |
| } else { |
| present_key[dst] = {{ scalar }}({{ kvSource }}[k_base + d]); |
| } |
| {% else %} |
| present_key[dst] = {{ scalar }}({{ kvSource }}[k_base + d]); |
| {% endif %} |
| present_value[dst] = {{ scalar }}({{ vSource }}[v_base + d]); |
| } |
| |
