File size: 18,490 Bytes

{% 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 %}
{{ env.wgsl.resourceDeclarations }}

// com.microsoft.SparseAttention, attention pass.
// One workgroup per (batch, query-token tile, query head). The workgroup sweeps only the key
// blocks this query's layout row selects, in tiles of WG keys with one key per thread,
// merging each tile into a running online softmax so the score row is never
// materialized. Repeated CSR columns are suppressed because the CSR data represents a
// boolean block mask.
//
// The block mask is CSR: row q_abs / SPARSE_BLOCK of layout (head % NUM_LAYOUT) spans
// block_col_indices[start .. end), and each of those columns names SPARSE_BLOCK
// contiguous keys. Sweeping (end - start) * SPARSE_BLOCK virtual slots and resolving
// each through the column list is what makes the kernel cost the SELECTED blocks rather
// than the whole history -- the reason the operator exists.
//
// Causality bounds every visited row at this query's absolute position. A layout
// with the triangular dense count and a sparse row with row + 1 entries take the
// dense shortcut without inspecting their column values.
const Q_HEADS: u32 = {{ numHeads }}u;
const KV_HEADS: u32 = {{ kvNumHeads }}u;
const HEAD_DIM: u32 = {{ headSize }}u;
// The cache is read four elements at a time. `head_size` is a multiple of eight,
// so every row divides into whole vec4 values and requires no scalar tail.
const HEAD_VEC: u32 = {{ headVec }}u;
const MAX_CACHE_SEQ: u32 = {{ maxCacheSeq }}u;
const SPARSE_BLOCK: u32 = {{ sparseBlockSize }}u;
const NUM_LAYOUT: u32 = {{ numLayout }}u;
const ROW_STRIDE: u32 = {{ maxBlocks + 1 }}u;
const COL_STRIDE: u32 = {{ maxNnz }}u;
const Q_HIDDEN: u32 = {{ numHeads * headSize }}u;
{% if not usesRotary %}
const Q_STRIDE: u32 = {{ packedStride if packedQkv else numHeads * headSize }}u;
{% endif %}
const WG: u32 = {{ attnWorkgroup }}u;
const Q_TILE: u32 = {{ qTile }}u;

// FLT_MAX, not -inf, as the online (m, d) accumulator init: merges must keep
// `m - m` finite so an empty lane / all--inf row contributes the exact
// accumulator identity (m, d) = (-FLT_MAX, 0). Operator epilogues interpret
// a zero final denominator according to their public semantics. Using -inf
// here changes +inf-row behavior.
const FLT_MAX: f32 = 3.4028234663852886e38;

fn is_finite_f32(value: f32) -> bool {
  return select(false, value <= FLT_MAX, value >= -FLT_MAX);
}

// x - m that is exactly 0 when x equals a finite m, so exp(shifted) == 1
// exactly at the row max. `x - x` on an infinite max is a legal fast-math
// fold to 0, which would silently turn +inf rows finite — the explicit
// equality test keeps the NaN propagation of the serial kernels.
fn shifted_value(value: f32, maxValue: f32) -> f32 {
  let equalFiniteMax = select(false, value == maxValue, is_finite_f32(maxValue));
  return select(value - maxValue, 0.0, equalFiniteMax);
}

fn exp_shift(value: f32, maxValue: f32) -> f32 {
  return exp(shifted_value(value, maxValue));
}

var<workgroup> q_shared: array<vec4<f32>, Q_TILE * HEAD_VEC>;
var<workgroup> running_out: array<vec4<f32>, Q_TILE * HEAD_VEC>;
var<workgroup> probs: array<f32, Q_TILE * WG>;
{% if vStageWorthIt %}
// The workgroup stages each chunk of value rows, then HEAD_VEC lanes consume
// those rows from shared memory during the running-output update.
const V_STAGE_KEYS: u32 = 16u;
var<workgroup> v_stage: array<vec4<f32>, V_STAGE_KEYS * HEAD_VEC>;
{% endif %}
{% if valueParts > 1 %}
// Each lane owns one (key partition, output vec4); query streams share V loads.
var<workgroup> value_partials: array<vec4<f32>, {{ valueParts * headVec * qTile | int }}>;
{% endif %}
// One resolved cache row base per key of the current tile, so the value accumulation
// re-reads a base instead of re-walking the column list per head dimension.
var<workgroup> key_rows: array<u32, WG>;

{% set mdStreams = qTile %}
// Workgroup-cooperative merge of per-thread online-softmax (m, d) partials:
// mNew = max(m1, m2)
// dNew = d1 * exp(m1 - mNew) + d2 * exp(m2 - mNew)
// Both the subgroup and portable barrier-tree engines return the same merged
// pair to every invocation. Repeated merges require a workgroup barrier between
// calls before their shared partial storage is reused.
{% set mdStreamed = mdStreams is defined %}
{% set mdStreams = mdStreams if mdStreams is defined else 1 %}
{% set mdExtent = "WG" if mdStreams == 1 else "WG * " ~ mdStreams ~ "u" %}
var<workgroup> partialM: array<f32, {{ mdExtent }}>;
var<workgroup> partialD: array<f32, {{ mdExtent }}>;
{% if mdStreamed %}

// In-place fold of {{ mdStreams }} streams. Input partials occupy
// partialM/partialD; stream s returns its merged pair in slot s * WG.
fn combine_partials_streams(lidx: u32) {
  workgroupBarrier();
  var stride = WG / 2u;
  loop {
    if (stride == 0u) {
      break;
    }
    if (lidx < stride) {
{% for s in range(mdStreams) %}
      {
        let slot = {{ s }}u * WG + lidx;
        let m1 = partialM[slot];
        let d1 = partialD[slot];
        let m2 = partialM[slot + stride];
        let d2 = partialD[slot + stride];
        let mNew = max(m1, m2);
        partialD[slot] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
        partialM[slot] = mNew;
      }
{% endfor %}
    }
    workgroupBarrier();
    stride = stride / 2u;
  }
}
{% else %}

fn combine_partials(m: f32, d: f32, lidx: u32) -> vec2<f32> {
  partialM[lidx] = m;
  partialD[lidx] = d;
  workgroupBarrier();
  var stride = WG / 2u;
  loop {
    if (stride == 0u) {
      break;
    }
    if (lidx < stride) {
      let m1 = partialM[lidx];
      let d1 = partialD[lidx];
      let m2 = partialM[lidx + stride];
      let d2 = partialD[lidx + stride];
      let mNew = max(m1, m2);
      partialD[lidx] = d1 * exp_shift(m1, mNew) + d2 * exp_shift(m2, mNew);
      partialM[lidx] = mNew;
    }
    workgroupBarrier();
    stride = stride / 2u;
  }
  let merged = vec2<f32>(partialM[0], partialD[0]);
  // Trailing barrier so back-to-back calls cannot race a next call's partial
  // stores against this call's reads of slot 0.
  workgroupBarrier();
  return merged;
}
{% endif %}

{% set ATTN_SCALE_DIM = "HEAD_DIM" %}fn scale_value() -> f32 {
  if (params.scale != 0.0) { return params.scale; }
  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
}

{{ sparse_schedule() }}
@compute @workgroup_size(WG, 1, 1)
fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
{% if promptTail %}
  // The matrix pass retains every query when this call includes past history.
  if (u32(total_sequence_length[0]) != params.seqLen) { return; }
{% endif %}
  let tile0 = wg.x * Q_TILE{% if queryOffset > 0 %} + {{ queryOffset }}u{% endif %};
  let head = wg.y % Q_HEADS;
  let batch = wg.y / Q_HEADS;
  let tid = lid.x;

  let past = past_sequence_length(batch);
  let kv_head = head / (Q_HEADS / KV_HEADS);
  let scale = scale_value();

  // Layouts cycle over the heads: four layouts and eight heads give S0 S1 S2 S3 S0 ...
  let layout_id = head % NUM_LAYOUT;
  let row_base = layout_id * ROW_STRIDE;
  // A layout at or above the triangular dense-block count is treated as dense.
  let dense_nnz = ROW_STRIDE * (ROW_STRIDE - 1u) / 2u;
  let layout_is_dense = u32(block_row_indices[row_base + ROW_STRIDE - 1u]) >= dense_nnz;

{% for j in range(qTile) %}
  // Keys 0 .. q_abs are visible; the +1 makes the bound exclusive. A tile that runs past
  // the last token carries inactive queries; they take no mask row and store nothing.
  let live_{{ j }} = tile0 + {{ j }}u < params.seqLen;
  let q_abs_{{ j }} = past + tile0 + {{ j }}u;
  let key_bound_{{ j }} = q_abs_{{ j }} + 1u;
  let mask_row_{{ j }} = q_abs_{{ j }} / SPARSE_BLOCK;
{% endfor %}

{% for j in range(qTile) %}
  {
{% if usesRotary %}
    let q_base = ((batch * Q_HEADS + head) * params.seqLen + min(tile0 + {{ j }}u, params.seqLen - 1u)) * HEAD_DIM;
{% else %}
    let q_base = (batch * params.seqLen + min(tile0 + {{ j }}u, params.seqLen - 1u)) * Q_STRIDE + head * HEAD_DIM;
{% endif %}
    for (var dv: u32 = tid; dv < HEAD_VEC; dv = dv + WG) {
      let qb = q_base + dv * 4u;
{% if usesRotary %}
      q_shared[{{ j }}u * HEAD_VEC + dv] = vec4<f32>(q_rotary[qb], q_rotary[qb + 1u], q_rotary[qb + 2u], q_rotary[qb + 3u]);
{% else %}
      q_shared[{{ j }}u * HEAD_VEC + dv] = vec4<f32>(f32(query[qb]), f32(query[qb + 1u]), f32(query[qb + 2u]), f32(query[qb + 3u]));
{% endif %}
      running_out[{{ j }}u * HEAD_VEC + dv] = vec4<f32>(0.0);
    }
  }
{% endfor %}
  workgroupBarrier();

{% for j in range(qTile) %}
  var runningMax_{{ j }} = -FLT_MAX;
  var runningDenom_{{ j }} = 0.0;
{% endfor %}

  // Consecutive tokens span at most two mask rows, and every query of a row selects the
  // same blocks, so one sweep per row covers the tile. A query contributes only to the
  // sweep of its own row, which is why its online state is never merged across rows.
{% if qTile > 1 %}
  let row_first = q_abs_0 / SPARSE_BLOCK;
{% if (seqLen - queryOffset) % qTile != 0 %}
  // Inactive query lanes must not extend traversal past the final CSR row.
  let row_last = (past + min(tile0 + Q_TILE, params.seqLen) - 1u) / SPARSE_BLOCK;
{% else %}
  let row_last = mask_row_{{ qTile - 1 }};
{% endif %}
  for (var mask_row = row_first; mask_row <= row_last; mask_row = mask_row + 1u) {
{% else %}
  {
    let mask_row = mask_row_0;
{% endif %}
    let start = u32(block_row_indices[row_base + mask_row]);
    let end = u32(block_row_indices[row_base + mask_row + 1u]);
    // A row listing every causal block is also dense, so its column values are ignored.
    let dense_row = layout_is_dense || (end - start == mask_row + 1u);
    // The dense sweep runs to the furthest causal bound of the queries on this row; each
    // query still drops the keys past its own.
    var bound_max = 0u;
{% for j in range(qTile) %}
    if (live_{{ j }} && mask_row_{{ j }} == mask_row) { bound_max = max(bound_max, key_bound_{{ j }}); }
{% endfor %}
    let slot_count = select((end - start) * SPARSE_BLOCK, bound_max, dense_row);

    var tileBase: u32 = 0u;
    loop {
      if (tileBase >= slot_count) {
        break;
      }
      let slot = tileBase + tid;

      // Resolving the slot is query-independent: the block list belongs to the row, and
      // only the causal bound below differs between the queries sharing it.
      var key = 0u;
      var keyResolved = false;
      if (slot < slot_count) {
        if (!dense_row) {
          // CSR is a boolean mask. Repeating a column therefore cannot repeat its
          // probability mass, even though malformed-but-valid CSR may contain duplicates.
          let entry = slot / SPARSE_BLOCK;
          let block = u32(block_col_indices[layout_id * COL_STRIDE + start + entry]);
          var duplicate = false;
          for (var previous = 0u; previous < entry; previous = previous + 1u) {
            duplicate = duplicate ||
              u32(block_col_indices[layout_id * COL_STRIDE + start + previous]) == block;
          }
          key = block * SPARSE_BLOCK + slot % SPARSE_BLOCK;
          keyResolved = !duplicate;
        } else {
          key = slot;
          keyResolved = true;
        }
      }
      // Always a row this cache holds, so an unresolved slot leaves a readable base rather
      // than a stale or uninitialized one and the accumulation below needs no guard.
      let base = ((batch * KV_HEADS + kv_head) * MAX_CACHE_SEQ + select(0u, key, keyResolved)) * HEAD_DIM;
      key_rows[tid] = base;

{% for j in range(qTile) %}
      let allowed_{{ j }} = keyResolved && live_{{ j }} && mask_row_{{ j }} == mask_row && key < key_bound_{{ j }};
      var acc_{{ j }} = 0.0;
{% endfor %}
      if (keyResolved) {
        // One key vector, every query's score: this is the reuse the tile exists for.
        let row_v = base / 4u;
        for (var dv: u32 = 0u; dv < HEAD_VEC; dv = dv + 1u) {
          let kv = vec4<f32>(present_key[row_v + dv]);
{% for j in range(qTile) %}
          acc_{{ j }} = acc_{{ j }} + dot(q_shared[{{ j }}u * HEAD_VEC + dv], kv);
{% endfor %}
        }
      }

      // One barrier tree folds all Q_TILE online-softmax partials.
{% for j in range(qTile) %}
      let score_{{ j }} = acc_{{ j }} * scale;
      partialM[{{ j }}u * WG + tid] = select(-FLT_MAX, score_{{ j }}, allowed_{{ j }});
      partialD[{{ j }}u * WG + tid] = select(0.0, 1.0, allowed_{{ j }});
{% endfor %}
      combine_partials_streams(tid);
{% for j in range(qTile) %}
      let tile_{{ j }} = vec2<f32>(partialM[{{ j }}u * WG], partialD[{{ j }}u * WG]);
{% endfor %}

{% for j in range(qTile) %}
      let newMax_{{ j }} = max(runningMax_{{ j }}, tile_{{ j }}.x);
      let correction_{{ j }} = exp_shift(runningMax_{{ j }}, newMax_{{ j }});
      runningDenom_{{ j }} = runningDenom_{{ j }} * correction_{{ j }}
        + tile_{{ j }}.y * exp_shift(tile_{{ j }}.x, newMax_{{ j }});
      runningMax_{{ j }} = newMax_{{ j }};
      probs[{{ j }}u * WG + tid] = select(0.0, exp_shift(score_{{ j }}, newMax_{{ j }}), allowed_{{ j }});
{% endfor %}
      workgroupBarrier();

      // One value vector serves every query. When head columns underfill the
      // workgroup, spare lanes sum disjoint contiguous key ranges and publish
      // partials; each output owner then folds them in key-range order. The
      // manifest limits this storage and keeps a single-owner fallback.
{% if valueParts > 1 %}
      let tileCount = min(WG, slot_count - tileBase);
      let valueDim = tid % HEAD_VEC;
      let valuePart = tid / HEAD_VEC;
      if (valuePart < {{ valueParts }}u) {
        let first = tileCount * valuePart / {{ valueParts }}u;
        let last = tileCount * (valuePart + 1u) / {{ valueParts }}u;
{% for j in range(qTile) %}
        var valueSum_{{ j }} = vec4<f32>(0.0);
{% endfor %}
        for (var i = first; i < last; i++) {
          let vv = vec4<f32>(present_value[key_rows[i] / 4u + valueDim]);
{% for j in range(qTile) %}
          valueSum_{{ j }} += probs[{{ j }}u * WG + i] * vv;
{% endfor %}
        }
{% for j in range(qTile) %}
        value_partials[{{ j * valueParts * headVec | int }}u + tid] = valueSum_{{ j }};
{% endfor %}
      }
      workgroupBarrier();
      if (tid < HEAD_VEC) {
{% for j in range(qTile) %}
        var valueSum_{{ j }} = value_partials[{{ j * valueParts * headVec | int }}u + tid];
{% for p in range(1, valueParts) %}
        valueSum_{{ j }} += value_partials[{{ (j * valueParts + p) * headVec | int }}u + tid];
{% endfor %}
        running_out[{{ j }}u * HEAD_VEC + tid] = running_out[{{ j }}u * HEAD_VEC + tid] * correction_{{ j }} + valueSum_{{ j }};
{% endfor %}
      }
      workgroupBarrier();
{% elif vStageWorthIt %}
      let tileCount = min(WG, slot_count - tileBase);
{% for j in range(qTile) %}
      var vSum_{{ j }} = vec4<f32>(0.0);
{% endfor %}
      for (var chunkBase: u32 = 0u; chunkBase < tileCount; chunkBase = chunkBase + V_STAGE_KEYS) {
        let chunkCount = min(V_STAGE_KEYS, tileCount - chunkBase);
        for (var slot = tid; slot < chunkCount * HEAD_VEC; slot = slot + WG) {
          let ki = slot / HEAD_VEC;
          let dvs = slot % HEAD_VEC;
          v_stage[slot] = vec4<f32>(present_value[key_rows[chunkBase + ki] / 4u + dvs]);
        }
        workgroupBarrier();
        for (var dv: u32 = tid; dv < HEAD_VEC; dv = dv + WG) {
          for (var i: u32 = 0u; i < chunkCount; i = i + 1u) {
            let vv = v_stage[i * HEAD_VEC + dv];
{% for j in range(qTile) %}
            vSum_{{ j }} = vSum_{{ j }} + probs[{{ j }}u * WG + chunkBase + i] * vv;
{% endfor %}
          }
        }
        // Orders this chunk's v_stage reads before the next chunk overwrites it.
        workgroupBarrier();
      }
      for (var dv: u32 = tid; dv < HEAD_VEC; dv = dv + WG) {
{% for j in range(qTile) %}
        running_out[{{ j }}u * HEAD_VEC + dv] = running_out[{{ j }}u * HEAD_VEC + dv] * correction_{{ j }} + vSum_{{ j }};
{% endfor %}
      }
      // Orders this tile's probs and key_rows reads before the next tile writes them.
      workgroupBarrier();
{% else %}
      let tileCount = min(WG, slot_count - tileBase);
      for (var dv: u32 = tid; dv < HEAD_VEC; dv = dv + WG) {
{% for j in range(qTile) %}
        var vSum_{{ j }} = vec4<f32>(0.0);
{% endfor %}
        for (var i: u32 = 0u; i < tileCount; i = i + 1u) {
          let vv = vec4<f32>(present_value[key_rows[i] / 4u + dv]);
{% for j in range(qTile) %}
          vSum_{{ j }} = vSum_{{ j }} + probs[{{ j }}u * WG + i] * vv;
{% endfor %}
        }
{% for j in range(qTile) %}
        running_out[{{ j }}u * HEAD_VEC + dv] = running_out[{{ j }}u * HEAD_VEC + dv] * correction_{{ j }} + vSum_{{ j }};
{% endfor %}
      }
      // Orders this tile's probs and key_rows reads before the next tile writes them.
      workgroupBarrier();
{% endif %}

      tileBase = tileBase + WG;
    }
  }

  // An all-masked row becomes a uniform average over causal history rather than zero.
{% for j in range(qTile) %}
  if (live_{{ j }}) {
    let hasMass_{{ j }} = runningDenom_{{ j }} > 0.0;
    let invDenom_{{ j }} = select(0.0, 1.0 / runningDenom_{{ j }}, hasMass_{{ j }});
    let out_base_{{ j }} = (batch * params.seqLen + tile0 + {{ j }}u) * Q_HIDDEN + head * HEAD_DIM;
    for (var dv: u32 = tid; dv < HEAD_VEC; dv = dv + WG) {
      var out_v: vec4<f32>;
      if (hasMass_{{ j }}) {
        out_v = running_out[{{ j }}u * HEAD_VEC + dv] * invDenom_{{ j }};
      } else {
        var total = vec4<f32>(0.0);
        for (var key = 0u; key < key_bound_{{ j }}; key = key + 1u) {
          let row_v = (((batch * KV_HEADS + kv_head) * MAX_CACHE_SEQ + key) * HEAD_DIM) / 4u;
          total = total + vec4<f32>(present_value[row_v + dv]);
        }
        out_v = total / f32(key_bound_{{ j }});
      }
      let ob = out_base_{{ j }} + dv * 4u;
      output[ob] = {{ scalar }}(out_v.x);
      output[ob + 1u] = {{ scalar }}(out_v.y);
      output[ob + 2u] = {{ scalar }}(out_v.z);
      output[ob + 3u] = {{ scalar }}(out_v.w);
    }
  }
{% endfor %}
}