com.microsoft.DecoderMaskedMultiHeadAttention / build /webgpu /attn-flash-decode-splitk-merge.wgsl.jinja
| {% if splitQueries is not defined %}{% set splitQueries = false %}{% endif %} | |
| {% set qSeq = 0 %} | |
| {{ env.wgsl.resourceDeclarations }} | |
| // Split-K flash decode, pass 2 of 2. Combines the per-split un-normalized online | |
| // states (max, denom, weighted V sum) written by the split pass into the final | |
| // attention output for each query token, using the online softmax merge: | |
| // globalMax = max_s partialMax[s] | |
| // globalDen = sum_s partialDenom[s] * exp(partialMax[s] - globalMax) | |
| // out[d] = (sum_s partialOut[s, d] * exp(partialMax[s] - globalMax)) / globalDen | |
| // One workgroup per (batch, head); one thread per output vec4 channel. The | |
| // per-split (max, denom) scan is tiny (NUM_SPLITS is small), so every thread | |
| // recomputes it locally and keeps the merge state in registers. | |
| // Output layout: rank-4 [batch, heads, seq, headDim], matching the split pass. | |
| const HEAD_DIM_V4: u32 = {{ headDimV4 }}u; | |
| const Q_HEADS: u32 = {{ qNumHeads }}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)); | |
| } | |
| @compute @workgroup_size(HEAD_DIM_V4, 1, 1) | |
| fn main( | |
| @builtin(workgroup_id) wg: vec3<u32>, | |
| @builtin(local_invocation_id) lid: vec3<u32> | |
| ) { | |
| {% if splitQueries %} | |
| let queryToken = wg.x; | |
| {% endif %} | |
| let h = wg.y; | |
| let b = wg.z; | |
| let d4 = lid.x; | |
| if (h >= Q_HEADS || d4 >= HEAD_DIM_V4{% if splitQueries %} || queryToken >= Q_SEQ{% endif %}) { | |
| return; | |
| } | |
| {% if splitQueries %} | |
| let mdBase = ((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS; | |
| {% else %} | |
| let mdBase = (b * Q_HEADS + h) * NUM_SPLITS; | |
| {% endif %} | |
| var globalMax = -FLT_MAX; | |
| for (var s: u32 = 0u; s < NUM_SPLITS; s = s + 1u) { | |
| globalMax = max(globalMax, partial_stats[mdBase + s].x); | |
| } | |
| var globalDenom = 0.0; | |
| var outV = vec4<f32>(0.0); | |
| for (var s: u32 = 0u; s < NUM_SPLITS; s = s + 1u) { | |
| let stats = partial_stats[mdBase + s]; | |
| let w = exp_shift(stats.x, globalMax); | |
| globalDenom = globalDenom + stats.y * w; | |
| {% if splitQueries %} | |
| let pBase = (((b * Q_SEQ + queryToken) * Q_HEADS + h) * NUM_SPLITS + s) * HEAD_DIM_V4; | |
| {% else %} | |
| let pBase = ((b * Q_HEADS + h) * NUM_SPLITS + s) * HEAD_DIM_V4; | |
| {% endif %} | |
| outV = outV + partial_out[pBase + d4] * w; | |
| } | |
| {% 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 %} | |
| // A split that covers no keys writes vec2(-FLT_MAX, 0), so when every split is empty — | |
| // a query whose window admits nothing — globalMax stays -FLT_MAX, exp_shift(x, x) is 1, | |
| // and globalDenom lands on exactly 0. Emit 0 rather than dividing so an all-empty | |
| // row has a defined zero result. | |
| var outValue = select(vec4<f32>(0.0), outV / globalDenom, globalDenom > 0.0); | |
| output[qBaseV4 + d4] = vec4<{{ scalar }}>(outValue); | |
| } | |