Kernels:

---

webgpu-kernels

/

com.microsoft.SparseAttention

Like 0

Follow

WebGPU Kernels Community 42

{% 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]);
}