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