com.microsoft.SparseAttention / build /webgpu /sparse-attention.wgsl.jinja

Xenova HF Staff

sync 91d990483a17

f03331e verified 10 days ago

19 kB

	{% 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 %}
	// 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 combineSubgroups = combineSubgroups is defined and combineSubgroups %}
	{% if combineSubgroups %}
	// Cross-subgroup merge that assumes nothing about which invocations share a
	// subgroup or how many subgroups there are: each subgroup's elected lane
	// publishes the subgroup pair in the slot at its OWN invocation index and sets
	// that index's bit in a workgroup bitmask; thread 0 then folds exactly the
	// published slots, in ascending index order (the online (m, d) merge is not
	// float-associative, so the order is fixed), and clears the mask for the next
	// call as it reads it. Workgroup memory starts zeroed, so the mask needs no
	// setup. Same three collectives as a single-subgroup reduce, two barriers.
	var<workgroup> partialM: array<f32, WG>;
	var<workgroup> partialD: array<f32, WG>;
	var<workgroup> leaderMask: array<atomic<u32>, (WG + 31u) / 32u>;
	var<workgroup> combinedMD: vec2<f32>;

	// When the whole workgroup is one subgroup the subgroup reduce already covers
	// it (no barriers, no shared state). `subgroup_size` is the size of the current
	// subgroup and uniform, so the test is exact and may guard the barriers below.
	fn combine_partials(m: f32, d: f32, lidx: u32, sgSize: u32) -> vec2<f32> {
	let sgM = subgroupMax(m);
	// A lane with no elements contributes d == 0 (exact identity). A +inf
	// element made exp(inf - inf) = NaN stick in that lane's d; a NaN element
	// landed in d via exp(NaN); both survive the merge and are detected by the
	// code after the reduction.
	let sgD = subgroupAdd(d * exp_shift(m, sgM));
	if (sgSize == WG) {
	return vec2<f32>(sgM, sgD);
	}
	if (subgroupElect()) {
	partialM[lidx] = sgM;
	partialD[lidx] = sgD;
	atomicOr(&leaderMask[lidx / 32u], 1u << (lidx % 32u));
	}
	workgroupBarrier();
	if (lidx == 0u) {
	var accM = -FLT_MAX;
	var accD = 0.0;
	for (var w = 0u; w < (WG + 31u) / 32u; w = w + 1u) {
	var bits = atomicExchange(&leaderMask[w], 0u);
	while (bits != 0u) {
	let slot = w * 32u + firstTrailingBit(bits);
	bits = bits & (bits - 1u);
	let mNew = max(accM, partialM[slot]);
	accD = accD * exp_shift(accM, mNew) + partialD[slot] * exp_shift(partialM[slot], mNew);
	accM = mNew;
	}
	}
	combinedMD = vec2<f32>(accM, accD);
	}
	workgroupBarrier();
	return combinedMD;
	}
	{% else %}
	{% 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 %}
	{% endif %}


	{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
	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>) {
	let tile0 = wg.x * Q_TILE;
	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;
	let row_last = mask_row_{{ qTile - 1 }};
	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();

	// running_out[j][d] is owned by the same thread across every tile (tid = d mod WG),
	// so this rescale-and-accumulate needs no further synchronization. One value vector
	// serves every query, which is the other half of the tile's reuse; a key outside a
	// query's causal bound carries prob 0 and is multiplied away.
	{% if 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 %}
	}