File size: 20,391 Bytes

f96162c

{% if useSubgroups is not defined %}{% set useSubgroups = true %}{% endif %}
{% if splitQueries is not defined %}{% set splitQueries = false %}{% endif %}
{% if quantizedCache is not defined %}{% set quantizedCache = false %}{% endif %}
{% if cacheSeqlens is not defined %}{% set cacheSeqlens = false %}{% endif %}
{% if hasMask is not defined %}{% set hasMask = false %}{% endif %}
{% if maskIsBool is not defined %}{% set maskIsBool = false %}{% endif %}
{% set splitKWorkgroupSize = source.workgroupSize if source.workgroupSize is defined else tunables.WORKGROUP_SIZE %}
{% if useSubgroups %}
enable subgroups;
{% endif %}
{% if usesF16 %}
enable f16;
{% endif %}
{{ env.wgsl.resourceDeclarations }}

// Split-K flash attention, pass 1 of 2; the merge pass follows. Shared by
// dense-attention decode and short-query/long-context prefill paths.
//
// The non-split flash decode launches only `batch * numHeads` workgroups, each
// sweeping the whole KV sequence serially in WG-key tiles. This pass splits the
// KV sequence into `NUM_SPLITS` contiguous ranges and gives each range its own
// workgroup, so `batch * numHeads * NUM_SPLITS` workgroups run the tiled online
// softmax in parallel. Each workgroup emits the *un-normalized* online state for
// its range — the running (max, denom) and the softmax-weighted V sum before the
// final divide — and the merge pass combines the per-split states with the online
// rule.
{% if source.layout == "bhsd" %}
// Layout: rank-4 [batch, heads, seq, headDim] for Q/K/V.
{% elif source.layout == "layer_cache" %}
// Layout: flat query [heads, headDim] plus a persistent KV cache laid out
// [layer, cacheLen, kvHeads, headDim]. This is the Qwen3.5 decode layout; the
// dispatch has a single implicit batch.
{% else %}
// Layout: token-major [batch, seq, heads * headDim]; Q and KV hidden strides
// are compiled constants.
{% endif %}
{% if (fusedQNormRope is defined and fusedQNormRope) or source.layout != "layer_cache" %}const HEAD_DIM: u32 = {{ headDim }}u;
{% endif %}
const HEAD_DIM_V4: u32 = {{ headDimV4 }}u;
const Q_HEADS: u32 = {{ qNumHeads }}u;
const KV_HEADS: u32 = {{ kvNumHeads }}u;
{% if source.layout == "bsh" %}
const Q_HIDDEN_V4: u32 = {{ qHiddenV4 }}u;
const KV_HIDDEN_V4: u32 = {{ kvHiddenV4 }}u;
{% elif source.layout == "layer_cache" %}
const LAYER: u32 = {{ layer }}u;
const CACHE_LEN: u32 = {{ cacheLen }}u;
const ATTN_SCALE: f32 = {{ scale }};
{% endif %}
const WG: u32 = {{ splitKWorkgroupSize }}u;
const NUM_SPLITS: u32 = {{ numSplits }}u;
{% if splitQueries %}
const Q_SEQ: u32 = {{ qSeq }}u;
{% endif %}
// 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>, HEAD_DIM_V4>;
var<workgroup> running_out: array<vec4<f32>, HEAD_DIM_V4>;
var<workgroup> probs: array<f32, WG>;
{% set coopQk = useSubgroups and headDimV4 >= 8 and not (usesF16 and headDimV4 <= 32) %}
{% set jGroups = (splitKWorkgroupSize / headDimV4)|int %}
{% set jSplitV = (splitKWorkgroupSize % headDimV4 == 0) and (jGroups >= 2) %}
{% if coopQk %}
var<workgroup> sval_sh: array<f32, WG>;
{% endif %}
{% if jSplitV %}
var<workgroup> vacc_sh: array<vec4<f32>, WG>;
{% endif %}
{% set combineSubgroups = useSubgroups %}
// 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 %}
// Per-subgroup partials are published into a deterministic slot: the subgroup's
// ordinal index within the workgroup (lidx / sgSize). The online (m, d) merge
// is not float-associative, so thread 0 must fold partials in a fixed order.
// Subgroups partition a workgroup into contiguous ordinal ranges on supported
// backends, so the ordinal slot is unique per subgroup and every slot in
// [0, subgroupCount) is written (each subgroup elects one leader).
// Sized for the worst case of one partial per invocation.
var<workgroup> partialM: array<f32, WG>;
var<workgroup> partialD: array<f32, WG>;
var<workgroup> combinedMD: vec2<f32>;

// When the whole workgroup is one subgroup the subgroup reduce already covers
// it (no barriers, no shared state); otherwise subgroup leaders publish
// partials through shared memory and thread 0 folds them in ordinal order.
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);
  }
  let subgroupCount = (WG + sgSize - 1u) / sgSize;
  // Pre-seed every fold slot with the (max, denom) identity. The fold below reads a
  // fixed subgroupCount slots in ordinal order (for determinism), but a slot whose
  // subgroup elects no leader this call — e.g. a fully out-of-window key tile in the
  // flash-attention loop that re-uses this shared memory each iteration — would
  // otherwise read stale shared memory. Identity makes such a slot a no-op.
  // (max identity = -FLT_MAX, denom identity = 0.)
  if (lidx < subgroupCount) {
    partialM[lidx] = -FLT_MAX;
    partialD[lidx] = 0.0;
  }
  workgroupBarrier();
  if (subgroupElect()) {
    let slot = lidx / sgSize;
    partialM[slot] = sgM;
    partialD[slot] = sgD;
  }
  workgroupBarrier();
  if (lidx == 0u) {
    var accM = -FLT_MAX;
    var accD = 0.0;
    for (var i = 0u; i < subgroupCount; i = i + 1u) {
      let mNew = max(accM, partialM[i]);
      accD = accD * exp_shift(accM, mNew) + partialD[i] * exp_shift(partialM[i], 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. The caller stores its per-thread
// partials into partialM/partialD first and reads the merged pair of stream s
// from slot s * WG afterwards.
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 source.layout == "layer_cache" %}{% set ATTN_SCALE_OVERRIDE = "ATTN_SCALE" %}{% endif %}
{% if ATTN_SCALE_DIM is not defined %}{% set ATTN_SCALE_DIM = "HEAD_DIM" %}{% endif %}
fn scale_value() -> f32 {
{% if ATTN_SCALE_OVERRIDE is defined %}
  return {{ ATTN_SCALE_OVERRIDE }};
{% else %}
  if (params.scale != 0.0) { return params.scale; }
  return inverseSqrt(f32({{ ATTN_SCALE_DIM }}));
{% endif %}
}


{% if quantizedCache %}
{% macro emit_quant_scale4(kind, scaleBuffer) %}
fn {{ kind }}scale4(d4: u32, hk: u32) -> vec4<f32> {
  if (params.perChannel == 0u) {
    return vec4<f32>({{ scaleBuffer }}[0]);
  }
  let base = hk * HEAD_DIM + d4 * 4u;
  return vec4<f32>(
    {{ scaleBuffer }}[base],
    {{ scaleBuffer }}[base + 1u],
    {{ scaleBuffer }}[base + 2u],
    {{ scaleBuffer }}[base + 3u]
  );
}
{%- endmacro %}
{%- macro emit_quant_load4(format, kind, buffer, scaleBuffer) %}
{{ emit_quant_scale4(kind, scaleBuffer) }}
fn load_{{ kind }}4(indexV4: u32, d4: u32, hk: u32) -> vec4<f32> {
{%- if format == "int8" %}
  return vec4<f32>({{ buffer }}[indexV4]) * {{ kind }}scale4(d4, hk);
{%- else %}
  // Two elements cover this vec4: each carries two +8-biased nibbles, low first.
  let rowBase = indexV4 - d4;
  let lo = {{ buffer }}[rowBase + d4 * 2u];
  let hi = {{ buffer }}[rowBase + d4 * 2u + 1u];
  let nibbles = vec4<i32>(
    i32(lo & 0xFu), i32((lo >> 4u) & 0xFu),
    i32(hi & 0xFu), i32((hi >> 4u) & 0xFu)
  );
  let signed = nibbles - vec4<i32>(8);
  return vec4<f32>(signed) * {{ kind }}scale4(d4, hk);
{%- endif %}
}
{%- endmacro %}

{{ emit_quant_load4("int8", "key", "key", "k_scale") }}
{{ emit_quant_load4("int8", "value", "value", "v_scale") }}
{% else %}
fn load_key4(indexV4: u32) -> vec4<f32> {
  return vec4<f32>(key[indexV4]);
}

fn load_value4(indexV4: u32) -> vec4<f32> {
  return vec4<f32>(value[indexV4]);
}
{% endif %}

{% if hasBias %}
// Packed [Q; K; V] bias rows (token-independent). The Q bias folds into the
// query row before the Q.K dots; the K bias adds a constant to every key score
// that softmax cancels, so it is skipped; the V bias is token-independent and
// is applied once in the merge pass after the final normalize.
fn load_bias4(base: u32, d4: u32) -> vec4<f32> {
  let offset = base + d4 * 4u;
  return vec4<f32>(bias[offset], bias[offset + 1u], bias[offset + 2u], bias[offset + 3u]);
}

{% endif %}
@compute @workgroup_size(WG, 1, 1)
fn main(
  @builtin(workgroup_id) wg: vec3<u32>,
  @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
  @builtin(subgroup_size) sgSize: u32{% endif %}
) {
{% if useSubgroups %}
  // Subgroup tiles partition the fixed workgroup exactly. The advertised range
  // is validated before dispatch; retain this uniform guard for implementations that
  // choose an intermediate width at pipeline execution time.
  if (sgSize == 0u || sgSize > WG || WG % sgSize != 0u) { return; }
{% endif %}
{% if splitQueries %}
  let queryToken = wg.x / NUM_SPLITS;
  let split = wg.x % NUM_SPLITS;
{% else %}
  let split = wg.x;
{% endif %}
  let h = wg.y;
  let b = wg.z;
  if (h >= Q_HEADS || split >= NUM_SPLITS{% if splitQueries %} || queryToken >= Q_SEQ{% endif %}{% if source.layout == "layer_cache" %} || params.past_len >= CACHE_LEN{% endif %}) {
    return;
  }
  let tid = lid.x;
  let hKv = h / (Q_HEADS / KV_HEADS);
{% if source.layout == "layer_cache" %}
  let kvSeq = params.past_len + 1u;
{% else %}
  let cacheSeq = params.kvSeq;
{% if cacheSeqlens %}
  // Buffer-sharing caches retain their capacity in the physical BNSH stride;
  // seqlens_k supplies the active end independently for each batch.
  let kvSeq = min(cacheSeq, u32(seqlens_k[b]) + 1u);
{% else %}
  let kvSeq = cacheSeq;
{% endif %}
{% endif %}

  // Query row (decode uses token zero; short-query prefill folds the token into wg.x).
{% if source.layout == "bsh" %}
{% if splitQueries %}
  let qBaseV4 = (b * Q_SEQ + queryToken) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
{% else %}
  let qBaseV4 = b * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
{% endif %}
  let kvBaseV4 = b * kvSeq * KV_HIDDEN_V4 + hKv * HEAD_DIM_V4;
  let kvTokenStrideV4 = KV_HIDDEN_V4;
{% elif source.layout == "layer_cache" %}
  let qBaseV4 = h * HEAD_DIM_V4;
  let kvBaseV4 = (LAYER * CACHE_LEN * KV_HEADS + hKv) * HEAD_DIM_V4;
  let kvTokenStrideV4 = KV_HEADS * HEAD_DIM_V4;
{% else %}
{% if splitQueries %}
  let qBaseV4 = ((b * Q_HEADS + h) * Q_SEQ + queryToken) * HEAD_DIM_V4;
{% else %}
  let qBaseV4 = (b * Q_HEADS + h) * HEAD_DIM_V4;
{% endif %}
  let kvBaseV4 = (b * KV_HEADS + hKv) * cacheSeq * HEAD_DIM_V4;
  let kvTokenStrideV4 = HEAD_DIM_V4;
{% endif %}

  // Contiguous KV range owned by this split. Ceil division lets the last split
  // absorb any remainder; empty ranges write identity partials and are ignored
  // by the merge pass.
{% if hasWindow %}
  // Sliding window on the single decode query (absolute position
  // kvSeq-1): it attends only the last `windowSize` keys, so split the
  // contiguous [windowStart, kvSeq) range instead of the whole cache.
  var windowStart: u32 = 0u;
  if (kvSeq > params.windowSize) {
    windowStart = kvSeq - params.windowSize;
  }
  let activeKeys = kvSeq - windowStart;
  let keysPerSplit = (activeKeys + NUM_SPLITS - 1u) / NUM_SPLITS;
  let splitStart = windowStart + split * keysPerSplit;
{% else %}
  let keysPerSplit = (kvSeq + NUM_SPLITS - 1u) / NUM_SPLITS;
  let splitStart = split * keysPerSplit;
{% endif %}
  var splitEnd = splitStart + keysPerSplit;
  if (splitEnd > kvSeq) {
    splitEnd = kvSeq;
  }

{% set hasBias = hasBias is defined and hasBias %}
  for (var d4 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
    var qv = vec4<f32>(query[qBaseV4 + d4]);
{% if hasBias %}
    qv = qv + load_bias4(h * HEAD_DIM, d4);
{% endif %}
    q_shared[d4] = qv;
    running_out[d4] = vec4<f32>(0.0);
  }
  workgroupBarrier();
  let scale = scale_value();
  var runningMax = -FLT_MAX;
  var runningDenom = 0.0;

  var kjBase = splitStart;
  loop {
    if (kjBase >= splitEnd) {
      break;
    }
    let kj = kjBase + tid;
    var keyAllowed = kj < splitEnd;
    let tileCount = min(WG, splitEnd - kjBase);

    var score = -FLT_MAX;
    var m = -FLT_MAX;
    var dPart = 0.0;
{% if coopQk %}
    // Cooperative Q.K: one subgroup per key, lanes splitting HEAD_DIM_V4, then a hardware
    // subgroupAdd — turns the per-thread HEAD_DIM_V4-long dependent dot chain into a few
    // strided vec4 dots + one reduce. Uniform trip count keeps subgroupAdd in uniform flow.
    let sgPerWg = WG / sgSize;
    let qkRounds = (tileCount + sgPerWg - 1u) / sgPerWg;
    let lane = tid % sgSize;
    let sgInWg = tid / sgSize;
    for (var rr: u32 = 0u; rr < qkRounds; rr = rr + 1u) {
      let j = rr * sgPerWg + sgInWg;
      var accS: f32 = 0.0;
      if (j < tileCount) {
        let kRowV4 = kvBaseV4 + (kjBase + j) * kvTokenStrideV4;
        for (var d4: u32 = lane; d4 < HEAD_DIM_V4; d4 = d4 + sgSize) {
          accS = accS + dot(q_shared[d4], load_key4(kRowV4 + d4{% if quantizedCache %}, d4, hKv{% endif %}));
        }
      }
      let sj = subgroupAdd(accS);
      if (lane == 0u && j < tileCount) {
        sval_sh[j] = sj;
      }
    }
    workgroupBarrier();
    if (keyAllowed) {
      score = sval_sh[tid] * scale;
      m = score;
      dPart = 1.0;
    }
{% else %}
    if (keyAllowed) {
      let kRowV4 = kvBaseV4 + kj * kvTokenStrideV4;
      var acc: f32 = 0.0;
      for (var d4: u32 = 0u; d4 < HEAD_DIM_V4; d4 = d4 + 1u) {
        acc = acc + dot(q_shared[d4], load_key4(kRowV4 + d4{% if quantizedCache %}, d4, hKv{% endif %}));
      }
      score = acc * scale;
      m = score;
      dPart = 1.0;
    }
{% endif %}
{% if hasMask %}
    if (keyAllowed) {
{% if splitQueries %}
      let maskQuery = queryToken;
{% else %}
      let maskQuery = 0u;
{% endif %}
      let maskIndex = b * params.maskBatchStride + h * params.maskHeadStride + maskQuery * params.maskSeqStride + kj;
{% if maskIsBool %}
      // A rejected bool-mask key contributes no probability mass. The merge
      // pass already maps a zero global denominator to an all-zero output row.
      if (attn_mask[maskIndex] == 0u) {
        keyAllowed = false;
        score = -FLT_MAX;
        dPart = 0.0;
      }
{% else %}
      score = score + f32(attn_mask[maskIndex]);
{% endif %}
      m = score;
    }
{% endif %}
    let tile = combine_partials(m, dPart, tid{% if useSubgroups %}, sgSize{% endif %});

// Merge one key tile's online-softmax (maximum, denominator) partial into the
// running state, then store the per-key probabilities consumed by V accumulation.
    let newMax = max(runningMax, tile.x);
    let correction = exp_shift(runningMax, newMax);
    runningDenom = runningDenom * correction + tile.y * exp_shift(tile.x, newMax);
    runningMax = newMax;

    var prob = 0.0;
    if (keyAllowed) {
      prob = exp_shift(score, newMax);
    }
    probs[tid] = prob;
    workgroupBarrier();


{% if jSplitV %}
    // j-split V accumulation: thread (jg, d4v) sums keys j == jg mod
    // J_GROUPS for dim block d4v into a register, then the groups combine
    // through shared memory so all lanes participate.
    const J_GROUPS: u32 = {{ jGroups }}u;
    let jg = tid / HEAD_DIM_V4;
    let d4v = tid % HEAD_DIM_V4;
    var vacc = vec4<f32>(0.0);
    var jj = jg;
    loop {
      if (jj >= tileCount) { break; }
      vacc = vacc + probs[jj] * load_value4(
        kvBaseV4 + (kjBase + jj) * kvTokenStrideV4 + d4v{% if quantizedCache %},
        d4v,
        hKv{% endif %}
      );
      jj = jj + J_GROUPS;
    }
    vacc_sh[tid] = vacc;
    workgroupBarrier();
    for (var d4: u32 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
      var a4 = running_out[d4] * correction;
      for (var g: u32 = 0u; g < J_GROUPS; g = g + 1u) {
        a4 = a4 + vacc_sh[g * HEAD_DIM_V4 + d4];
      }
      running_out[d4] = a4;
    }
    workgroupBarrier();
{% else %}
    for (var d4: u32 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
      var vSum = vec4<f32>(0.0);
      for (var i: u32 = 0u; i < tileCount; i = i + 1u) {
        vSum = vSum + probs[i] * load_value4(
          kvBaseV4 + (kjBase + i) * kvTokenStrideV4 + d4{% if quantizedCache %},
          d4,
          hKv{% endif %}
        );
      }
      running_out[d4] = running_out[d4] * correction + vSum;
    }
    workgroupBarrier();
{% endif %}

    kjBase = kjBase + WG;
  }

  // Emit un-normalized partials for (b, h, split): the merge pass divides.
{% if splitQueries %}
  let partialBase = (((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS + split) * HEAD_DIM_V4;
{% else %}
  let partialBase = ((b * Q_HEADS + h) * NUM_SPLITS + split) * HEAD_DIM_V4;
{% endif %}
  for (var d4: u32 = tid; d4 < HEAD_DIM_V4; d4 = d4 + WG) {
    partial_out[partialBase + d4] = running_out[d4];
  }
  if (tid == 0u) {
{% if splitQueries %}
    let mdBase = ((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS + split;
{% else %}
    let mdBase = (b * Q_HEADS + h) * NUM_SPLITS + split;
{% endif %}
    // (max, denom) travel together to the merge, so they share one buffer as an
    // interleaved vec2 rather than costing two bindings. Interleaved, not two
    // halves, so the index needs no region size — and the merge reads both
    // fields of a split in a single load.
    partial_stats[mdBase] = vec2<f32>(runningMax, runningDenom);
  }
}