sync c928d21e6cc1

Browse files

Files changed (8) hide show

README.md +3 -1
build/webgpu/attn-materialized-rowstats-combine-f32.wgsl.jinja +72 -0
build/webgpu/attn-materialized-sgmat-f32.wgsl.jinja +562 -0
build/webgpu/bench.json +23 -0
build/webgpu/gqa-qprep.wgsl.jinja +11 -4
build/webgpu/manifest.json +1343 -2
build/webgpu/metadata.json +8 -6
build/webgpu/test.json +719 -0

README.md CHANGED Viewed

@@ -75,7 +75,7 @@ Attributes and default values (overridable per request):
 ## Device requirements
-Some implementation variants require `shader-f16` and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
@@ -89,6 +89,8 @@ Some implementation variants require `shader-f16` and `subgroups`. These are rou
 - [`attn-flash-online.wgsl.jinja`](build/webgpu/attn-flash-online.wgsl.jinja)
 - [`attn-flash-prefill-cluster.wgsl.jinja`](build/webgpu/attn-flash-prefill-cluster.wgsl.jinja)
 - [`attn-flash-q32-broadcast.wgsl.jinja`](build/webgpu/attn-flash-q32-broadcast.wgsl.jinja)
 - [`attn-online-scalar.wgsl.jinja`](build/webgpu/attn-online-scalar.wgsl.jinja)
 - [`gqa-attention.wgsl.jinja`](build/webgpu/gqa-attention.wgsl.jinja)
 - [`gqa-present.wgsl.jinja`](build/webgpu/gqa-present.wgsl.jinja)

 ## Device requirements
+Some implementation variants require `subgroup-matrix`, `shader-f16`, and `subgroups`. These are route-specific capabilities, not package-wide requirements; availability also depends on the request shape and dtype.
 ## Files
 - [`attn-flash-online.wgsl.jinja`](build/webgpu/attn-flash-online.wgsl.jinja)
 - [`attn-flash-prefill-cluster.wgsl.jinja`](build/webgpu/attn-flash-prefill-cluster.wgsl.jinja)
 - [`attn-flash-q32-broadcast.wgsl.jinja`](build/webgpu/attn-flash-q32-broadcast.wgsl.jinja)
+- [`attn-materialized-rowstats-combine-f32.wgsl.jinja`](build/webgpu/attn-materialized-rowstats-combine-f32.wgsl.jinja)
+- [`attn-materialized-sgmat-f32.wgsl.jinja`](build/webgpu/attn-materialized-sgmat-f32.wgsl.jinja)
 - [`attn-online-scalar.wgsl.jinja`](build/webgpu/attn-online-scalar.wgsl.jinja)
 - [`gqa-attention.wgsl.jinja`](build/webgpu/gqa-attention.wgsl.jinja)
 - [`gqa-present.wgsl.jinja`](build/webgpu/gqa-present.wgsl.jinja)

build/webgpu/attn-materialized-rowstats-combine-f32.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,72 @@

+// Folds the per-tile softmax constants the score pass emitted into one (m, d)
+// pair per query row, for the apply pass to normalize with.
+//
+// The score pass publishes one partial per row, key tile, and subgroup column,
+// so this pass can combine the constants without rescanning the materialized
+// score matrix. It uses the same online merge rule as the row-local reduction.
+//
+// One thread per row. Each row's partials are contiguous along the row axis, so
+// consecutive threads read consecutive pairs; the fold walks slots in index
+// order, which is fixed for a given shape, so the result does not depend on how
+// the GPU schedules anything.
+// `maxOnly` means the producer published only a row max per slot because
+// computing the denominator there would double the exp count. Fold maxima and
+// leave the denominator to the apply pass, which sees every row element anyway.
+{{ env.wgsl.resourceDeclarations }}
+const SLOTS: u32 = {{ statSlots }}u;
+const Q_SEQ: u32 = {{ statQuerySeq }}u;
+const WG: u32 = {{ materializedRowStatsWg }}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));
+}
+@compute @workgroup_size(WG, 1, 1)
+fn main(
+  @builtin(global_invocation_id) gid: vec3<u32>,
+  @builtin(num_workgroups) nwg: vec3<u32>
+) {
+  let row = gid.x + gid.y * nwg.x * WG;
+  if (row >= params.rows) { return; }
+  // `row` already runs over (batch, head, query) together, and the partial
+  // layout puts that same product one axis out from the slot, so the stride
+  // between a row's slots is the number of rows in its (batch, head) plane.
+  let plane = row / Q_SEQ;
+  let inPlane = row % Q_SEQ;
+  let base = ((plane * SLOTS) * Q_SEQ + inPlane) * 2u;
+  var m = -FLT_MAX;
+  var d = 0.0;
+  for (var slot = 0u; slot < SLOTS; slot = slot + 1u) {
+    let index = base + slot * Q_SEQ * 2u;
+    let slotM = scorePartials[index];
+    let slotD = scorePartials[index + 1u];
+    let merged = max(m, slotM);
+    d = d * exp_shift(m, merged) + slotD * exp_shift(slotM, merged);
+    m = merged;
+  }
+  rowStats[row * 2u] = m;
+  rowStats[row * 2u + 1u] = d;
+}

build/webgpu/attn-materialized-sgmat-f32.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,562 @@

+{% set MT = "f16" if (operandF16 is defined and operandF16) else "f32" %}
+{% if MT == "f16" %}
+enable f16;
+{% endif %}
+enable subgroups;
+{% if pinSubgroupSize32 %}
+enable subgroup_size_control;
+{% endif %}
+enable chromium_experimental_subgroup_matrix;
+diagnostic(off, chromium.subgroup_matrix_uniformity);
+{{ env.wgsl.resourceDeclarations }}
+{% set layout = source.layout | default("bsh") %}
+{% set headMajor = layout == "bhsd" %}
+{% set kvHeadMajor = (source.kvLayout | default(layout)) == "bhsd" %}
+{% set CAUSAL_UPPER_LEFT = source.causalUpperLeft is defined and source.causalUpperLeft %}
+{% set CAUSAL = (source.causalRightAlign is defined and source.causalRightAlign) or CAUSAL_UPPER_LEFT %}
+{% macro q_index(row, d) %}(b * params.qSeq + {{ row }}) * HIDDEN + h * HEAD_DIM + {{ d }}{% endmacro %}
+{% macro kv_index(seq, d) %}{% if kvHeadMajor %}((b * KV_HEADS + h_kv) * params.kvSeq + {{ seq }}) * HEAD_DIM + {{ d }}{% else %}(b * params.kvSeq + {{ seq }}) * KV_HIDDEN + h_kv * HEAD_DIM + {{ d }}{% endif %}{% endmacro %}
+{% set OUT_ROW_STRIDE = "HEAD_DIM" if headMajor else "HIDDEN" %}
+{% set KV_ROW_STRIDE = "HEAD_DIM" if kvHeadMajor else "KV_HIDDEN" %}
+{% set scorePhase = source.phase == "score" %}
+{% set SCORE_BIAS = scorePhase and source.scoreBias is defined and source.scoreBias %}
+{% set SCORE_WINDOW = CAUSAL and source.scoreWindow is defined and source.scoreWindow %}
+{% set USE_SEQLENS = source.useSeqlens is defined and source.useSeqlens %}
+{% set FUSED_SOFTMAX = source.fusedSoftmax is defined and source.fusedSoftmax %}
+{% set PRIVATE_ROW_STATS = FUSED_SOFTMAX and (materializedSgmatPrivateRowStats is defined and materializedSgmatPrivateRowStats) %}
+{% macro score_value(index, guard) %}
+{% if FUSED_SOFTMAX %}
+{% if PRIVATE_ROW_STATS %}
+select(0.0, exp_shift(scores[{{ index }}], private_softmax_m) / private_softmax_d, {{ guard[1] }})
+{%- else %}
+select(0.0, exp_shift(scores[{{ index }}], softmax_m[{{ guard[0] }}]) / softmax_d[{{ guard[0] }}], {{ guard[1] }})
+{%- endif %}
+{% else %}
+select(0.0, scores[{{ index }}], {{ guard[1] }})
+{%- endif %}
+{% endmacro %}
+{% set TILE_M_VALUE = materializedSgmatQueryTile %}
+{% set TILE_N_VALUE = materializedSgmatKeyTile %}
+{% set TILE_K_VALUE = materializedSgmatInnerTile %}
+{% set SUB_ROWS_VALUE = materializedSgmatSubgroupTileRows if materializedSgmatSubgroupTileRows is defined else 16 %}
+{% set SUB_COLS_VALUE = materializedSgmatSubgroupTileCols if materializedSgmatSubgroupTileCols is defined else 32 %}
+{% set ROW_BLOCKS = (SUB_ROWS_VALUE / 8)|int %}
+{% set COL_BLOCKS = (SUB_COLS_VALUE / 8)|int %}
+{% set SUBGROUP_ROWS = (TILE_M_VALUE / SUB_ROWS_VALUE)|int %}
+{% set SUBGROUP_COLS = (TILE_N_VALUE / SUB_COLS_VALUE)|int %}
+{% set SUBGROUP_COUNT = SUBGROUP_ROWS * SUBGROUP_COLS %}
+{% set WORKGROUP_THREADS = SUBGROUP_COUNT * 32 %}
+{% if hasBias is not defined %}{% set hasBias = false %}{% endif %}
+{% set EMIT_ROW_STATS = source.emitRowStats is defined and source.emitRowStats %}
+{% set DIRECT_SCORE_STORE = materializedSgmatDirectScoreStore and not EMIT_ROW_STATS %}
+{% set DIRECT_APPLY_STORE = materializedSgmatDirectApplyStore and not hasBias and MT == "f32" %}
+{% set DIRECT_OUTPUT_STORE = DIRECT_SCORE_STORE if scorePhase else DIRECT_APPLY_STORE %}
+{% set RUNTIME_DIRECT_STORE = (not DIRECT_OUTPUT_STORE)
+   and materializedSgmatRuntimeDirectStore
+   and (scorePhase or (not hasBias and MT == "f32"))
+   and not EMIT_ROW_STATS %}
+{% set ANY_DIRECT_STORE = DIRECT_OUTPUT_STORE or RUNTIME_DIRECT_STORE %}
+{% set SCALE_IN_Q = DIRECT_SCORE_STORE or (RUNTIME_DIRECT_STORE and scorePhase) %}
+{% macro q_tile_value(index) %}{% if hasBias %}(query[{{ index }}] + bias[h * HEAD_DIM + k]){% else %}query[{{ index }}]{% endif %}{% if SCALE_IN_Q %} * score_scale{% endif %}{% endmacro %}
+const HEADS: u32 = {{ qNumHeads }}u;
+{% if kvNumHeads is defined %}
+const KV_HEADS: u32 = {{ kvNumHeads }}u;
+{% else %}
+const KV_HEADS: u32 = HEADS;
+{% endif %}
+const HEAD_DIM: u32 = {{ headDim }}u;
+const HIDDEN: u32 = {{ qHidden }}u;
+{% if hasBias %}
+/* Packed [Q; K; V] bias. Q folds into the query tile before the GEMM. K is
+ * omitted: expanding (q + bq).(k + bk) leaves a (q + bq).bk term that is
+ * constant across every key in the row, and softmax is invariant under that
+ * constant. V is token-independent, so sum_k p[k] * bv = bv after
+ * normalization; it is added after the apply GEMM. */
+{% endif %}
+{% if not kvHeadMajor %}
+// Packed K/V rows span KV_HEADS heads, so they are narrower than the query
+// row whenever queries are grouped. Q and the output keep HIDDEN.
+const KV_HIDDEN: u32 = KV_HEADS * HEAD_DIM;
+{% endif %}
+{% if EMIT_ROW_STATS %}
+const STAT_SLOTS: u32 = {{ statSlots }}u;
+{% endif %}
+const TILE_M: u32 = {{ TILE_M_VALUE }}u;
+const TILE_N: u32 = {{ TILE_N_VALUE }}u;
+const TILE_K: u32 = {{ TILE_K_VALUE }}u;
+const SUB_ROWS: u32 = {{ SUB_ROWS_VALUE }}u;
+const SUB_COLS: u32 = {{ SUB_COLS_VALUE }}u;
+const SUBGROUP_ROWS: u32 = {{ SUBGROUP_ROWS }}u;
+{% if (FUSED_SOFTMAX and not PRIVATE_ROW_STATS)
+   or not (TILE_M_VALUE == 64 and TILE_N_VALUE == 64 and TILE_K_VALUE == 32 and WORKGROUP_THREADS == 256) %}
+const WORKGROUP_THREADS: u32 = {{ WORKGROUP_THREADS }}u;
+{% endif %}
+var<workgroup> tile_A: array<{{ MT }}, {{ TILE_M_VALUE * TILE_K_VALUE }}>;
+var<workgroup> tile_B: array<{{ MT }}, {{ TILE_N_VALUE * TILE_K_VALUE }}>;
+{% set SCRATCH = "tile_A" if MT == "f32" else "store_scratch" %}
+{% if MT == "f16" %}
+// The compact epilogue banks f32 result fragments; an f16 operand tile cannot
+// alias them, so the f16 build carries a dedicated store scratch.
+var<workgroup> store_scratch: array<f32, {{ TILE_M_VALUE * TILE_K_VALUE }}>;
+{% endif %}
+{% if FUSED_SOFTMAX and not PRIVATE_ROW_STATS %}
+// The tile's rows own their softmax constants for the whole k loop, so they are
+// read once per workgroup rather than once per staged element.
+var<workgroup> softmax_m: array<f32, {{ TILE_M_VALUE }}>;
+var<workgroup> softmax_d: array<f32, {{ TILE_M_VALUE }}>;
+{% endif %}
+{% if FUSED_SOFTMAX or EMIT_ROW_STATS %}
+{% set stableUsage = stableHelperUsage if stableHelperUsage is defined else "all" %}
+// 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;
+{% if stableUsage != "constant" %}
+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);
+}
+{%- endif %}
+{% if stableUsage == "all" %}
+fn exp_shift(value: f32, maxValue: f32) -> f32 {
+  return exp(shifted_value(value, maxValue));
+}
+{%- endif %}
+{% endif %}
+@compute @workgroup_size({{ WORKGROUP_THREADS }}, 1, 1){{ " @subgroup_size(32)" if pinSubgroupSize32 else "" }}
+fn main(
+  @builtin(workgroup_id) wg: vec3<u32>,
+  @builtin(local_invocation_index) li: u32{% if not DIRECT_OUTPUT_STORE %},
+  @builtin(subgroup_invocation_id) lane: u32{% endif %}
+) {
+  let h = wg.z % HEADS;
+  let b = wg.z / HEADS;
+  let h_kv = h / (HEADS / KV_HEADS);
+  let m_base = wg.y * TILE_M;
+  let n_base = wg.x * TILE_N;
+  let subgroup = li / 32u;
+  let subtile_idy = subgroup % SUBGROUP_ROWS;
+  let subtile_idx = subgroup / SUBGROUP_ROWS;
+  let base_A = subtile_idy * SUB_ROWS;
+  let base_B = subtile_idx * SUB_COLS;
+{% if USE_SEQLENS %}
+  // Rows the cache-update pass left resident: the survivors live in
+  // [0, kv_active), and every causal/window bound below uses this live length
+  // while params.kvSeq keeps the allocated capacity for strides.
+  let kv_active = min(params.kvSeq, u32(seqlens_k[b]) + 1u);
+{% endif %}
+{% if CAUSAL and scorePhase %}
+  // Workgroup-uniform causal skip: a key tile starting past the last row's
+  // bound holds no valid column. Publish the stats identity for its slots so
+  // the combine pass reads initialized pairs, then leave before any loads.
+{% if SCORE_WINDOW %}
+  // A tile whose last column sits behind the FIRST row's window floor is dead
+  // the same way: later rows only move the floor further right.
+{% endif %}
+  let kv_causal_off = i32({% if USE_SEQLENS %}kv_active{% else %}params.kvSeq{% endif %}) - i32(params.qSeq);
+  if (i32(n_base) > kv_causal_off + i32(m_base + TILE_M) - 1{% if SCORE_WINDOW %}
+      || i32(n_base + TILE_N) <= kv_causal_off + i32(m_base) + 1 - i32(params.windowSize){% endif %}) {
+    for (var idx = li; idx < TILE_M * {{ SUBGROUP_COLS }}u; idx += {{ WORKGROUP_THREADS }}u) {
+      let stat_row = m_base + idx / {{ SUBGROUP_COLS }}u;
+      if (stat_row < params.qSeq) {
+        let slot = wg.x * {{ SUBGROUP_COLS }}u + idx % {{ SUBGROUP_COLS }}u;
+        let out_index = (((b * HEADS + h) * STAT_SLOTS + slot) * params.qSeq + stat_row) * 2u;
+        scorePartials[out_index] = -FLT_MAX;
+        scorePartials[out_index + 1u] = 0.0;
+      }
+    }
+    return;
+  }
+{% endif %}
+{% for row_block in range(ROW_BLOCKS) %}
+{% for col_block in range(COL_BLOCKS) %}
+  var matC{{ row_block }}{{ col_block }}: subgroup_matrix_result<f32, 8, 8>;
+{% endfor %}
+{% endfor %}
+{% if scorePhase %}
+  let inner = HEAD_DIM;
+{% if SCALE_IN_Q %}
+  // The direct store has no epilogue, so apply the score scale to Q. Direct and
+  // guarded store paths then share one pre-scaled query tile.
+  let score_scale = {{ attentionScaleExpression }};
+{% endif %}
+{% else %}
+  let inner = {% if USE_SEQLENS %}kv_active{% else %}params.kvSeq{% endif %};
+{% if CAUSAL %}
+  // Score tiles past this query tile's causal bound were never written; stop
+  // the reduction at the last written key column for the tile's rows.
+  let kv_causal_off = i32({% if USE_SEQLENS %}inner{% else %}params.kvSeq{% endif %}) - i32(params.qSeq);
+  let inner_bound = u32(clamp(kv_causal_off + i32(m_base + TILE_M), 0, i32(inner)));
+{% if SCORE_WINDOW %}
+  // Score tiles behind every row's window floor were never written either. The
+  // floor of this tile's FIRST row is the leftmost any of its rows can reach,
+  // and rounding it down to a tile boundary only re-reads columns the score
+  // phase did write (its skip test uses whole key tiles).
+  let inner_start = (u32(max(kv_causal_off + i32(m_base) + 1 - i32(params.windowSize), 0)) / TILE_K) * TILE_K;
+{% endif %}
+{% endif %}
+{% endif %}
+{% if FUSED_SOFTMAX %}
+{% if PRIVATE_ROW_STATS %}
+  // In the admitted BM64/BN64/BK32/WG256 loader, four adjacent lanes own the
+  // same query row for every reduction tile. Keep that row's constants private:
+  // this removes both 512 bytes of workgroup storage and the initialization
+  // barrier while preserving the exact exp/divide sequence of the shared-memory
+  // row-stats arm.
+  let private_stat_row =
+    (b * HEADS + h) * params.qSeq + min(m_base + li / 4u, params.qSeq - 1u);
+  let private_softmax_m = rowStats[private_stat_row * 2u];
+  let private_softmax_d = rowStats[private_stat_row * 2u + 1u];
+{% else %}
+  // One row-stats pair per query row of the tile. A query tail clamps to the last
+  // real row rather than reading past the buffer; those lanes are discarded by the
+  // staging guard anyway, and the clamp keeps the denominator non-zero.
+  for (var r = li; r < TILE_M; r += WORKGROUP_THREADS) {
+    let stat_row = (b * HEADS + h) * params.qSeq + min(m_base + r, params.qSeq - 1u);
+    softmax_m[r] = rowStats[stat_row * 2u];
+    softmax_d[r] = rowStats[stat_row * 2u + 1u];
+  }
+  workgroupBarrier();
+{% endif %}
+{% endif %}
+  for (var k_base = {% if SCORE_WINDOW and not scorePhase %}inner_start{% else %}0u{% endif %}; k_base < {% if CAUSAL and not scorePhase %}inner_bound{% else %}inner{% endif %}; k_base += TILE_K) {
+{% if source.phase == "apply" and not FUSED_SOFTMAX and MT == "f32" %}
+    // Full interior PV tiles can be loaded directly from storage. Query,
+    // reduction, and output-dimension tails use the guarded shared path below.
+    if (
+      k_base + TILE_K <= inner &&
+      m_base + TILE_M <= params.qSeq &&
+      n_base + TILE_N <= HEAD_DIM
+    ) {
+      for (var step = 0u; step < TILE_K; step += 8u) {
+{% for row_block in range(ROW_BLOCKS) %}
+        let score_offset{{ row_block }} = (b * HEADS + h) * params.qSeq * params.kvSeq
+          + (m_base + base_A + {{ row_block * 8 }}u) * params.kvSeq + k_base + step;
+        var matA{{ row_block }}: subgroup_matrix_left<f32, 8, 8> =
+          subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>, row_major>(
+            &scores, score_offset{{ row_block }}, params.kvSeq
+          );
+{% endfor %}
+{% for col_block in range(COL_BLOCKS) %}
+        let value_offset{{ col_block }} =
+          {{ kv_index("k_base + step", "n_base + base_B + " ~ (col_block * 8) ~ "u") }};
+        var matB{{ col_block }}: subgroup_matrix_right<f32, 8, 8> =
+          subgroupMatrixLoad<subgroup_matrix_right<f32, 8, 8>, row_major>(
+            &value, value_offset{{ col_block }}, {{ KV_ROW_STRIDE }}
+          );
+{% endfor %}
+{% for row_block in range(ROW_BLOCKS) %}
+{% for col_block in range(COL_BLOCKS) %}
+        matC{{ row_block }}{{ col_block }} = subgroupMatrixMultiplyAccumulate(
+          matA{{ row_block }}, matB{{ col_block }}, matC{{ row_block }}{{ col_block }}
+        );
+{% endfor %}
+{% endfor %}
+      }
+      continue;
+    }
+{% endif %}
+{% if TILE_M_VALUE == 64 and TILE_N_VALUE == 64 and TILE_K_VALUE == 32 and WORKGROUP_THREADS == 256 %}
+    // BM64/BN64/WG256 is the high-throughput geometry on wide devices. Four
+    // adjacent threads each own one contiguous eight-element segment of a
+    // row, preserving the original coalesced load schedule while retaining
+    // the generalized K-tail guards.
+    let a_row = li / 4u;
+    let a_col = (li % 4u) * 8u;
+    for (var i = 0u; i < 8u; i++) {
+      let row = m_base + a_row;
+      let k = k_base + a_col + i;
+{% if scorePhase %}
+{% if headDim % 32 == 0 %}
+      tile_A[a_row * TILE_K + a_col + i] = select(
+        {{ "0.0h" if MT == "f16" else "0.0" }},
+        {{ q_tile_value(q_index("row", "k")) }},
+        row < params.qSeq
+      );
+{% else %}
+      var loaded = {{ "0.0h" if MT == "f16" else "0.0" }};
+      if (row < params.qSeq && k < HEAD_DIM) {
+        loaded = {{ q_tile_value(q_index("row", "k")) }};
+      }
+      tile_A[a_row * TILE_K + a_col + i] = loaded;
+{% endif %}
+{% else %}
+      let score_base = (b * HEADS + h) * params.qSeq * params.kvSeq;
+      tile_A[a_row * TILE_K + a_col + i] =
+        {{ "f16(" if MT == "f16" else "" }}{{ score_value("score_base + row * params.kvSeq + k", ["a_row", "row < params.qSeq && k < params.kvSeq"]) }}{{ ")" if MT == "f16" else "" }};
+{% endif %}
+    }
+    let b_row = li / 4u;
+    let b_col = (li % 4u) * 8u;
+    for (var i = 0u; i < 8u; i++) {
+      let col = n_base + b_row;
+      let k = k_base + b_col + i;
+{% if scorePhase %}
+{% if headDim % 32 == 0 %}
+      tile_B[b_row * TILE_K + b_col + i] = select(
+        {{ "0.0h" if MT == "f16" else "0.0" }},
+        key[{{ kv_index("col", "k") }}],
+        col < params.kvSeq
+      );
+{% else %}
+      var loaded = {{ "0.0h" if MT == "f16" else "0.0" }};
+      if (col < params.kvSeq && k < HEAD_DIM) {
+        loaded = key[{{ kv_index("col", "k") }}];
+      }
+      tile_B[b_row * TILE_K + b_col + i] = loaded;
+{% endif %}
+{% else %}
+      tile_B[b_row * TILE_K + b_col + i] = select(
+        {{ "0.0h" if MT == "f16" else "0.0" }},
+        value[{{ kv_index("k", "col") }}],
+        k < params.kvSeq && col < HEAD_DIM
+      );
+{% endif %}
+    }
+{% else %}
+    // Cooperative linear loads make smaller tile dimensions independent of
+    // the selected workgroup size.
+    for (var idx = li; idx < TILE_M * TILE_K; idx += WORKGROUP_THREADS) {
+      let tile_row = idx / TILE_K;
+      let tile_k = idx % TILE_K;
+      let row = m_base + tile_row;
+      let k = k_base + tile_k;
+{% if scorePhase %}
+      var loaded = {{ "0.0h" if MT == "f16" else "0.0" }};
+      if (row < params.qSeq && k < HEAD_DIM) {
+        loaded = {{ q_tile_value(q_index("row", "k")) }};
+      }
+{% elif FUSED_SOFTMAX %}
+      let score_base = (b * HEADS + h) * params.qSeq * params.kvSeq;
+      let loaded =
+        {{ "f16(" if MT == "f16" else "" }}{{ score_value("score_base + row * params.kvSeq + k", ["tile_row", "row < params.qSeq && k < params.kvSeq"]) }}{{ ")" if MT == "f16" else "" }};
+{% else %}
+      var loaded = {{ "0.0h" if MT == "f16" else "0.0" }};
+      if (row < params.qSeq && k < params.kvSeq) {
+        let score_base = (b * HEADS + h) * params.qSeq * params.kvSeq;
+        loaded = {{ "f16(" if MT == "f16" else "" }}scores[score_base + row * params.kvSeq + k]{{ ")" if MT == "f16" else "" }};
+      }
+{% endif %}
+      tile_A[idx] = loaded;
+    }
+    for (var idx = li; idx < TILE_N * TILE_K; idx += WORKGROUP_THREADS) {
+      let tile_col = idx / TILE_K;
+      let tile_k = idx % TILE_K;
+      let col = n_base + tile_col;
+      let k = k_base + tile_k;
+{% if scorePhase %}
+      var loaded = {{ "0.0h" if MT == "f16" else "0.0" }};
+      if (col < params.kvSeq && k < HEAD_DIM) {
+        loaded = key[{{ kv_index("col", "k") }}];
+      }
+{% else %}
+      var loaded = {{ "0.0h" if MT == "f16" else "0.0" }};
+      if (k < params.kvSeq && col < HEAD_DIM) {
+        loaded = value[{{ kv_index("k", "col") }}];
+      }
+{% endif %}
+      tile_B[idx] = loaded;
+    }
+{% endif %}
+    workgroupBarrier();
+    for (var step = 0u; step < TILE_K; step += 8u) {
+{% for row_block in range(ROW_BLOCKS) %}
+      let matrix_a_offset{{ row_block }} =
+        (subtile_idy * SUB_ROWS + {{ row_block * 8 }}u) * TILE_K + step;
+      var matA{{ row_block }}: subgroup_matrix_left<{{ MT }}, 8, 8> =
+        subgroupMatrixLoad<subgroup_matrix_left<{{ MT }}, 8, 8>, row_major>(
+          &tile_A, matrix_a_offset{{ row_block }}, TILE_K
+        );
+{% endfor %}
+{% for col_block in range(COL_BLOCKS) %}
+      let matrix_b_offset{{ col_block }} =
+        (subtile_idx * SUB_COLS + {{ col_block * 8 }}u) * TILE_K + step;
+      var matB{{ col_block }}: subgroup_matrix_right<{{ MT }}, 8, 8> =
+        subgroupMatrixLoad<subgroup_matrix_right<{{ MT }}, 8, 8>, col_major>(
+          &tile_B, matrix_b_offset{{ col_block }}, TILE_K
+        );
+{% endfor %}
+{% for row_block in range(ROW_BLOCKS) %}
+{% for col_block in range(COL_BLOCKS) %}
+      matC{{ row_block }}{{ col_block }} = subgroupMatrixMultiplyAccumulate(
+        matA{{ row_block }}, matB{{ col_block }}, matC{{ row_block }}{{ col_block }}
+      );
+{% endfor %}
+{% endfor %}
+    }
+    workgroupBarrier();
+  }
+{% if ANY_DIRECT_STORE %}
+{% if RUNTIME_DIRECT_STORE %}
+  // Workgroup-uniform (both bases come from workgroup_id alone): an interior tile
+  // publishes through the subgroup-matrix collectives and is done; only an edge
+  // tile falls through to the scratch round trip and its barriers.
+  if (
+    m_base + TILE_M <= params.qSeq &&
+    n_base + TILE_N <= {% if scorePhase %}params.kvSeq{% else %}HEAD_DIM{% endif %}
+  ) {
+{% else %}
+  // Every dispatched tile is interior, so the matrices can go straight to
+  // storage without the compact scratch/readback epilogue.
+{% endif %}
+{% for row_block in range(ROW_BLOCKS) %}
+{% for col_block in range(COL_BLOCKS) %}
+{% if scorePhase %}
+  let output_offset{{ row_block }}{{ col_block }} =
+    (b * HEADS + h) * params.qSeq * params.kvSeq
+      + (m_base + base_A + {{ row_block * 8 }}u) * params.kvSeq
+      + n_base + base_B + {{ col_block * 8 }}u;
+  subgroupMatrixStore<row_major>(
+    &scores, output_offset{{ row_block }}{{ col_block }},
+    matC{{ row_block }}{{ col_block }}, params.kvSeq
+  );
+{% else %}
+  let output_offset{{ row_block }}{{ col_block }} =
+    {{ q_index("m_base + base_A + " ~ (row_block * 8) ~ "u", "n_base + base_B + " ~ (col_block * 8) ~ "u") }};
+  subgroupMatrixStore<row_major>(
+    &output, output_offset{{ row_block }}{{ col_block }},
+    matC{{ row_block }}{{ col_block }}, {{ OUT_ROW_STRIDE }}
+  );
+{% endif %}
+{% endfor %}
+{% endfor %}
+{% if RUNTIME_DIRECT_STORE %}
+    return;
+  }
+{% endif %}
+{% endif %}
+{% if not DIRECT_OUTPUT_STORE %}
+  let row_in_block = lane / 4u;
+  let col_in_block = (lane % 4u) * 2u;
+{% for row_block in range(ROW_BLOCKS) %}
+{% if row_block > 0 %}
+  // The compact path aliases the input tile as matrix-store scratch, so every
+  // lane must finish reading the preceding row block before it is overwritten.
+  workgroupBarrier();
+{% endif %}
+{% if EMIT_ROW_STATS %}
+    // All four lanes of this quad carry the same row (row_in_block is lane / 4),
+  // so the accumulator below is a partial over one row and the butterfly that
+  // merges it is quad-uniform — no divergent shuffle even on a query tail.
+  var stat_m{{ row_block }} = -FLT_MAX;
+  var stat_d{{ row_block }} = 0.0;
+{% endif %}
+{% for col_block in range(COL_BLOCKS) %}
+  subgroupMatrixStore<row_major>(
+    &{{ SCRATCH }},
+    (subgroup * {{ COL_BLOCKS }}u + {{ col_block }}u) * 64u,
+    matC{{ row_block }}{{ col_block }},
+    8u
+  );
+{% endfor %}
+  workgroupBarrier();
+{% for col_block in range(COL_BLOCKS) %}
+  for (var pair = 0u; pair < 2u; pair++) {
+    let row = m_base + base_A + {{ row_block * 8 }}u + row_in_block;
+    let col = n_base + base_B + {{ col_block * 8 }}u + col_in_block + pair;
+    if (
+      row < params.qSeq &&
+      col < {% if scorePhase %}params.kvSeq{% else %}HEAD_DIM{% endif %}
+    ) {
+      let result = {{ SCRATCH }}[
+        (subgroup * {{ COL_BLOCKS }}u + {{ col_block }}u) * 64u
+          + row_in_block * 8u + col_in_block + pair
+      ];
+{% if scorePhase %}
+{% if not SCALE_IN_Q %}
+      let scale = {{ attentionScaleExpression }};
+{% endif %}
+{% if CAUSAL %}
+      var scored = result * scale;
+{% if SCORE_BIAS %}
+      let bias_b = select(b, 0u, params.biasBatch == 1u);
+      let bias_h = select(h, 0u, params.biasHeads == 1u);
+      scored += attention_bias[
+        ((bias_b * params.biasHeads + bias_h) * params.qSeq + row) * params.kvSeq + col
+      ];
+{% endif %}
+      if (i32(col) > kv_causal_off + i32(row)) { scored = -FLT_MAX; }
+{% if SCORE_WINDOW %}
+      if (i32(col) + i32(params.windowSize) <= kv_causal_off + i32(row)) { scored = -FLT_MAX; }
+{% endif %}
+{% else %}
+      let scored = result{% if not SCALE_IN_Q %} * scale{% endif %};
+{% endif %}
+      scores[
+        (b * HEADS + h) * params.qSeq * params.kvSeq + row * params.kvSeq + col
+      ] = scored;
+{% if EMIT_ROW_STATS %}
+      // Softmax sees the STORED value, so the statistics have to be taken on it
+      // and not on the raw accumulator.
+      let stat_prev = stat_m{{ row_block }};
+      stat_m{{ row_block }} = max(stat_m{{ row_block }}, scored);
+      stat_d{{ row_block }} = stat_d{{ row_block }} * exp_shift(stat_prev, stat_m{{ row_block }})
+        + exp_shift(scored, stat_m{{ row_block }});
+{% endif %}
+{% else %}
+{% if hasBias %}
+      // V bias row base: skip the packed Q and K blocks, then index this head.
+{% endif %}
+      output[{{ q_index("row", "col") }}] = {{ "f16(" if MT == "f16" else "" }}result{{ ")" if MT == "f16" else "" }}{% if hasBias %} + bias[2u * HIDDEN + h * HEAD_DIM + col]{% endif %};
+{% endif %}
+    }
+  }
+{% endfor %}
+{% if EMIT_ROW_STATS %}
+  // Butterfly the quad UNCONDITIONALLY. A lane whose row ran past the query tail
+  // never entered the guard above and still carries the exact identity
+  // (-FLT_MAX, 0), which merges to a no-op — that is cheaper than making the
+  // shuffle conditional, and a subgroup operation under a guard that only some
+  // quads of the subgroup satisfy would not be uniform.
+{% for xor in [1, 2] %}
+  {
+    let other_m = subgroupShuffleXor(stat_m{{ row_block }}, {{ xor }}u);
+    let other_d = subgroupShuffleXor(stat_d{{ row_block }}, {{ xor }}u);
+    let merged_m = max(stat_m{{ row_block }}, other_m);
+    stat_d{{ row_block }} = stat_d{{ row_block }} * exp_shift(stat_m{{ row_block }}, merged_m)
+      + other_d * exp_shift(other_m, merged_m);
+    stat_m{{ row_block }} = merged_m;
+  }
+{% endfor %}
+  {
+    let stat_row = m_base + base_A + {{ row_block * 8 }}u + row_in_block;
+    if (lane % 4u == 0u && stat_row < params.qSeq) {
+      // Row is the FASTEST axis so the eight writing lanes of a subgroup land on
+      // eight consecutive pairs, and the combine pass reads a slot's whole column
+      // of rows contiguously.
+      let slot = wg.x * {{ SUBGROUP_COLS }}u + subtile_idx;
+      let out_index = (((b * HEADS + h) * STAT_SLOTS + slot) * params.qSeq + stat_row) * 2u;
+      scorePartials[out_index] = stat_m{{ row_block }};
+      scorePartials[out_index + 1u] = stat_d{{ row_block }};
+    }
+  }
+{% endif %}
+{% endfor %}
+{% endif %}
+}

build/webgpu/bench.json CHANGED Viewed

@@ -1560,6 +1560,29 @@
         "source": "register-geometry gate asymmetry",
         "notes": "Cached f32 prefill at headDim 256, above the shared-memory cluster's register-geometry boundary. This guards consistent route admission between no-past and cached prefill families."
       }
     }
   ]
 }

         "source": "register-geometry gate asymmetry",
         "notes": "Cached f32 prefill at headDim 256, above the shared-memory cluster's register-geometry boundary. This guards consistent route admission between no-past and cached prefill families."
       }
+    },
+    {
+      "name": "qwen3-prefill-32h8kv-d128-s256",
+      "preset": "smoke",
+      "vars": { "batch": 1, "qSeq": 256, "kvSeq": 256, "heads": 32, "kvHeads": 8, "headDim": 128 },
+      "attrs": { "num_heads": 32, "kv_num_heads": 8, "scale": 0.08838834764831845, "causal": 0 },
+      "inputs": {
+        "queryT": { "shape": [1, 256, 4096], "dtype": "float32", "dist": "normal", "seed": 960, "scale": 0.2 },
+        "keyT": { "shape": [1, 256, 1024], "dtype": "float32", "dist": "normal", "seed": 961, "scale": 0.2 },
+        "valueT": { "shape": [1, 256, 1024], "dtype": "float32", "dist": "normal", "seed": 962, "scale": 0.2 },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [255] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 256, 4096], "dtype": "float32" },
+        "presentKeyT": { "shape": [1, 8, 256, 128], "dtype": "float32" },
+        "presentValueT": { "shape": [1, 8, 256, 128], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [
+          { "type": "gflops", "value": "4 * args.batch * args.qSeq * args.kvSeq * args.heads * args.headDim" }
+        ]
+      }
     }
   ]
 }

build/webgpu/gqa-qprep.wgsl.jinja CHANGED Viewed

@@ -8,17 +8,22 @@
 // exactly (the present-cache K is already norm/rotary-applied, so only Q is touched),
 // so the flash result is bit-identical to the cooperative/threaded path.
 //
-// f16 queries are widened before cos/sin and norm/rotary arithmetic;
-// the qPrep output also stays f32 — RoPE precision matters at large positions, so
-// the rotated Q enters attention in full f32 (the cluster reads f32 Q + f16 K/V).
 const HEAD_DIM: u32 = {{ headDim }}u;
 const HEAD_DIM_V4: u32 = {{ headDim }}u / 4u;
 const Q_HEADS: u32 = {{ qHeads }}u;
 const Q_HIDDEN: u32 = {{ qHidden }}u;
 const Q_HIDDEN_V4: u32 = {{ qHidden }}u / 4u;
 const WG: u32 = {{ copyWorkgroupSize }}u;
 const HALF: u32 = {{ half }}u;
 const QK_EPS: f32 = {{ qkEps }};
 @compute @workgroup_size(WG)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
@@ -37,10 +42,12 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
   var q: array<f32, HEAD_DIM>;
   for (var d = 0u; d < HEAD_DIM; d = d + 1u) { q[d] = f32(query[base + d]); }
   var ss = 0.0;
   for (var d = 0u; d < HEAD_DIM; d = d + 1u) { ss = ss + q[d] * q[d]; }
   let invRms = inverseSqrt(ss / f32(HEAD_DIM) + QK_EPS);
   for (var d = 0u; d < HEAD_DIM; d = d + 1u) { q[d] = q[d] * invRms * f32(q_norm_weight[d]); }
   for (var d = 0u; d < HALF; d = d + 1u) {
     let cs = f32(cos_cache[absPos * HALF + d]);
     let sn = f32(sin_cache[absPos * HALF + d]);
@@ -52,6 +59,6 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>,
   let base4 = (b * params.qSeq + s) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
   for (var c = 0u; c < HEAD_DIM_V4; c = c + 1u) {
-    qout[base4 + c] = vec4<f32>(q[c * 4u], q[c * 4u + 1u], q[c * 4u + 2u], q[c * 4u + 3u]);
   }
 }

 // exactly (the present-cache K is already norm/rotary-applied, so only Q is touched),
 // so the flash result is bit-identical to the cooperative/threaded path.
 //
+// f16 queries are widened before cos/sin and norm/rotary arithmetic, and the
+// rotation itself always runs in f32 — RoPE precision matters at large
+// positions. The store narrows only for a consumer whose operand tiles are
+// f16 anyway (the subgroup-matrix route), where a wider intermediate would be
+// narrowed at staging regardless; every other consumer keeps the f32 output.
 const HEAD_DIM: u32 = {{ headDim }}u;
 const HEAD_DIM_V4: u32 = {{ headDim }}u / 4u;
 const Q_HEADS: u32 = {{ qHeads }}u;
 const Q_HIDDEN: u32 = {{ qHidden }}u;
 const Q_HIDDEN_V4: u32 = {{ qHidden }}u / 4u;
 const WG: u32 = {{ copyWorkgroupSize }}u;
+{% set HAS_QNORM = hasQNorm is not defined or hasQNorm %}
 const HALF: u32 = {{ half }}u;
+{% if HAS_QNORM %}
 const QK_EPS: f32 = {{ qkEps }};
+{% endif %}
 @compute @workgroup_size(WG)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>,
   var q: array<f32, HEAD_DIM>;
   for (var d = 0u; d < HEAD_DIM; d = d + 1u) { q[d] = f32(query[base + d]); }
+{% if HAS_QNORM %}
   var ss = 0.0;
   for (var d = 0u; d < HEAD_DIM; d = d + 1u) { ss = ss + q[d] * q[d]; }
   let invRms = inverseSqrt(ss / f32(HEAD_DIM) + QK_EPS);
   for (var d = 0u; d < HEAD_DIM; d = d + 1u) { q[d] = q[d] * invRms * f32(q_norm_weight[d]); }
+{% endif %}
   for (var d = 0u; d < HALF; d = d + 1u) {
     let cs = f32(cos_cache[absPos * HALF + d]);
     let sn = f32(sin_cache[absPos * HALF + d]);
   let base4 = (b * params.qSeq + s) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
   for (var c = 0u; c < HEAD_DIM_V4; c = c + 1u) {
+    qout[base4 + c] = {{ "vec4<f16>(" if (qPrepF16 is defined and qPrepF16) else "" }}vec4<f32>(q[c * 4u], q[c * 4u + 1u], q[c * 4u + 2u], q[c * 4u + 3u]){{ ")" if (qPrepF16 is defined and qPrepF16) else "" }};
   }
 }

build/webgpu/manifest.json CHANGED Viewed

@@ -210,13 +210,23 @@
     "CLUSTER_TILE_K_F16": 8,
     "NO_SG_TILE_K_MAX": 16,
     "COOPERATIVE_WORKGROUP_SIZE": 32,
-    "CLUSTER_MAX_SLICE": 8
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "subgroupsWave32": "device.features.has(\"subgroups\") and wave32Adapter",
     "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
     "windowCacheRequested": "attrs.sliding_window_cache == 1",
     "headDim": "dim(shapes.query, 2) / attrs.num_heads if (ranks.query == 3 and attrs.num_heads > 0) else 0",
     "copyWorkgroupSize": "min(tunables.COPY_WORKGROUP_SIZE, min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX))",
@@ -310,7 +320,41 @@
     "cachedNoSgTileBytes": "gqaHeadDim * (8 if tensorDtypes.queryT == \"float32\" else 4) + cachedNoSgReductionBytesPerKey",
     "cachedNoSgTileK": "min(tunables.NO_SG_TILE_K_MAX, max(1, floor(device.limits.maxComputeWorkgroupStorageSize / cachedNoSgTileBytes)))",
     "cachedNoSgWorkgroupBytes": "cachedNoSgTileK * cachedNoSgTileBytes",
-    "cachedNoSgClusterStorageOk": "cachedNoSgWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize"
   },
   "bindingSets": {
     "splitAttention": [
@@ -1895,9 +1939,1306 @@
           ]
         }
       }
     ]
   },
   "variants": [
     {
       "id": "new_kv_share_append_split",
       "description": "Retains the existing cache and appends new key/value rows in separate passes before portable attention. It avoids rebuilding unchanged cache positions when the past and present allocations share capacity.",

     "CLUSTER_TILE_K_F16": 8,
     "NO_SG_TILE_K_MAX": 16,
     "COOPERATIVE_WORKGROUP_SIZE": 32,
+    "CLUSTER_MAX_SLICE": 8,
+    "MATERIALIZED_SGMAT_QUERY_TILE": 64,
+    "MATERIALIZED_SGMAT_KEY_TILE": 64,
+    "MATERIALIZED_SGMAT_INNER_TILE": 32,
+    "MATERIALIZED_SOFTMAX_WORKGROUP_SIZE": 256,
+    "MATERIALIZED_SGMAT_MIN_SEQ": 256,
+    "MATERIALIZED_SGMAT_PAST_MIN_SEQ": 256,
+    "MATERIALIZED_SGMAT_WINDOW_MIN_QSEQ": 128
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32",
     "subgroupsWave32": "device.features.has(\"subgroups\") and wave32Adapter",
     "narrowSubgroupRange": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize < device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= 16",
+    "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32",
+    "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter",
+    "wave32Effective": "wave32Adapter or pinSubgroupSize32",
     "windowCacheRequested": "attrs.sliding_window_cache == 1",
     "headDim": "dim(shapes.query, 2) / attrs.num_heads if (ranks.query == 3 and attrs.num_heads > 0) else 0",
     "copyWorkgroupSize": "min(tunables.COPY_WORKGROUP_SIZE, min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX))",
     "cachedNoSgTileBytes": "gqaHeadDim * (8 if tensorDtypes.queryT == \"float32\" else 4) + cachedNoSgReductionBytesPerKey",
     "cachedNoSgTileK": "min(tunables.NO_SG_TILE_K_MAX, max(1, floor(device.limits.maxComputeWorkgroupStorageSize / cachedNoSgTileBytes)))",
     "cachedNoSgWorkgroupBytes": "cachedNoSgTileK * cachedNoSgTileBytes",
+    "cachedNoSgClusterStorageOk": "cachedNoSgWorkgroupBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "gqaMatQueryTile": "tunables.MATERIALIZED_SGMAT_QUERY_TILE",
+    "gqaMatKeyTile": "tunables.MATERIALIZED_SGMAT_KEY_TILE",
+    "gqaMatInnerTile": "tunables.MATERIALIZED_SGMAT_INNER_TILE",
+    "gqaMatSubgroupRows": "floor(gqaMatQueryTile / 16)",
+    "gqaMatSubgroupCols": "floor(gqaMatKeyTile / 32)",
+    "gqaMatWorkgroupSize": "gqaMatSubgroupRows * gqaMatSubgroupCols * 32",
+    "gqaMatCompactStorageBytes": "(gqaMatQueryTile + gqaMatKeyTile) * gqaMatInnerTile * 4",
+    "gqaMatGeometryOk": "gqaMatQueryTile >= 16 and gqaMatQueryTile % 16 == 0 and gqaMatKeyTile >= 32 and gqaMatKeyTile <= 64 and gqaMatKeyTile % 32 == 0 and gqaMatInnerTile == 32 and gqaMatQueryTile >= 64 and gqaMatKeyTile >= 64",
+    "gqaMatResourcesFit": "gqaMatGeometryOk and gqaMatWorkgroupSize <= deviceWorkgroupCap and gqaMatCompactStorageBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "gqaMatScoreBytes": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1) * 4",
+    "gqaMatScoreFits": "gqaMatScoreBytes <= device.limits.maxStorageBufferBindingSize and gqaMatScoreBytes <= device.limits.maxBufferSize",
+    "gqaMatBuffersFit": "numel(shapes.query) * 4 <= device.limits.maxStorageBufferBindingSize and numel(shapes.key) * 4 <= device.limits.maxStorageBufferBindingSize and numel(shapes.value) * 4 <= device.limits.maxStorageBufferBindingSize",
+    "gqaMatDispatchFits": "ceilDiv(dim(shapes.key, 1), gqaMatKeyTile) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(dim(shapes.query, 1), gqaMatQueryTile) <= device.limits.maxComputeWorkgroupsPerDimension and dim(shapes.query, 0) * attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension",
+    "gqaMatStatSlots": "ceilDiv(dim(shapes.key, 1), gqaMatKeyTile) * gqaMatSubgroupCols",
+    "gqaMatRowStatsWg": "min(tunables.MATERIALIZED_SOFTMAX_WORKGROUP_SIZE, deviceWorkgroupCap)",
+    "gqaMatRowStatsElements": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * 2",
+    "gqaMatScorePartialElements": "dim(shapes.query, 0) * attrs.num_heads * gqaMatStatSlots * dim(shapes.query, 1) * 2",
+    "gqaMatCoreOk": "qkvContractOk and attrs.causal == 0 and tensorDtypes.query == \"float32\" and attrs.local_window_size == -1 and headDim >= 64 and headDim <= 256 and headDim % 16 == 0 and dim(shapes.query, 1) >= tunables.MATERIALIZED_SGMAT_MIN_SEQ and dim(shapes.key, 1) >= tunables.MATERIALIZED_SGMAT_MIN_SEQ and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatScoreFits and gqaMatBuffersFit and gqaMatResourcesFit and gqaMatDispatchFits",
+    "gqaMatPastQSeq": "dim(shapes.queryT, 1)",
+    "gqaMatPastKvSeq": "dim(shapes.presentKeyT, 2)",
+    "gqaMatPastScoreBytes": "dim(shapes.queryT, 0) * attrs.num_heads * gqaMatPastQSeq * gqaMatPastKvSeq * 4",
+    "gqaMatPastScoreFits": "gqaMatPastScoreBytes <= device.limits.maxStorageBufferBindingSize and gqaMatPastScoreBytes <= device.limits.maxBufferSize",
+    "gqaMatPastBuffersFit": "numel(shapes.queryT) * 4 <= device.limits.maxStorageBufferBindingSize and numel(shapes.presentKeyT) * 4 <= device.limits.maxStorageBufferBindingSize and numel(shapes.presentValueT) * 4 <= device.limits.maxStorageBufferBindingSize",
+    "gqaMatPastStatSlots": "ceilDiv(gqaMatPastKvSeq, gqaMatKeyTile) * gqaMatSubgroupCols",
+    "gqaMatPastRowStatsElements": "dim(shapes.queryT, 0) * attrs.num_heads * gqaMatPastQSeq * 2",
+    "gqaMatPastScorePartialElements": "dim(shapes.queryT, 0) * attrs.num_heads * gqaMatPastStatSlots * gqaMatPastQSeq * 2",
+    "gqaMatPastDispatchFits": "ceilDiv(gqaMatPastKvSeq, gqaMatKeyTile) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(gqaMatPastQSeq, gqaMatQueryTile) <= device.limits.maxComputeWorkgroupsPerDimension and dim(shapes.queryT, 0) * attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension",
+    "gqaMatWindowOk": "windowShiftOk and plainAttentionOptions and tensorDtypes.queryT == \"float32\" and attrs.local_window_size > 0 and dim(shapes.queryT, 1) == dim(shapes.keyT, 1) and gqaHeadDim >= 64 and gqaHeadDim <= 256 and gqaHeadDim % 16 == 0 and dim(shapes.queryT, 1) >= tunables.MATERIALIZED_SGMAT_WINDOW_MIN_QSEQ and windowCapacity >= dim(shapes.queryT, 1) and attrs.local_window_size + dim(shapes.queryT, 1) >= windowCapacity and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits",
+    "gqaMatPastCoreOk": "sharedKvFloatOk and plainAttentionOptions and tensorDtypes.queryT == \"float32\" and attrs.local_window_size == -1 and gqaHeadDim >= 64 and gqaHeadDim <= 256 and gqaHeadDim % 16 == 0 and gqaMatPastQSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= gqaMatPastQSeq and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits",
+    "gqaMatPastRotaryOk": "sharedKvFloatOk and standardSoftmax and present.cosCacheT and rotaryRequested and not present.qNormWeightT and not present.attentionBiasT and not present.headSinkT and tensorDtypes.queryT == \"float32\" and attrs.local_window_size == -1 and gqaHeadDim >= 64 and gqaHeadDim <= 256 and gqaHeadDim % 16 == 0 and gqaHeadDim % 8 == 0 and gqaMatPastQSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= gqaMatPastQSeq and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits",
+    "gqaMatNewPastOk": "newKvPastOk and plainAttentionOptions and tensorDtypes.queryT == \"float32\" and attrs.local_window_size == -1 and gqaHeadDim >= 64 and gqaHeadDim <= 256 and gqaHeadDim % 16 == 0 and gqaMatPastQSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= gqaMatPastQSeq and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits",
+    "gqaMatPastCoreF16Ok": "sharedKvFloatOk and plainAttentionOptions and tensorDtypes.queryT == \"float16\" and device.features.has(\"shader-f16\") and attrs.local_window_size == -1 and gqaHeadDim >= 64 and gqaHeadDim <= 256 and gqaHeadDim % 16 == 0 and gqaMatPastQSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= gqaMatPastQSeq and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits",
+    "gqaMatPastRotaryF16Ok": "sharedKvFloatOk and standardSoftmax and present.cosCacheT and rotaryRequested and not present.qNormWeightT and not present.attentionBiasT and not present.headSinkT and tensorDtypes.queryT == \"float16\" and device.features.has(\"shader-f16\") and attrs.local_window_size == -1 and gqaHeadDim >= 64 and gqaHeadDim <= 256 and gqaHeadDim % 16 == 0 and gqaHeadDim % 8 == 0 and gqaMatPastQSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= gqaMatPastQSeq and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits",
+    "gqaMatPastBiasOk": "sharedKvFloatOk and standardSoftmax and not present.cosCacheT and not rotaryRequested and not present.qNormWeightT and not present.headSinkT and present.attentionBiasT and ranks.attentionBiasT == 4 and tensorDtypes.attentionBiasT == \"float32\" and (dim(shapes.attentionBiasT, 0) == dim(shapes.queryT, 0) or dim(shapes.attentionBiasT, 0) == 1) and (dim(shapes.attentionBiasT, 1) == attrs.num_heads or dim(shapes.attentionBiasT, 1) == 1) and dim(shapes.attentionBiasT, 2) == dim(shapes.queryT, 1) and dim(shapes.attentionBiasT, 3) == dim(shapes.presentKeyT, 2) and tensorDtypes.queryT == \"float32\" and attrs.local_window_size == -1 and gqaHeadDim >= 64 and gqaHeadDim <= 256 and gqaHeadDim % 16 == 0 and gqaMatPastQSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= tunables.MATERIALIZED_SGMAT_PAST_MIN_SEQ and gqaMatPastKvSeq >= gqaMatPastQSeq and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and gqaMatPastScoreFits and gqaMatPastBuffersFit and gqaMatResourcesFit and gqaMatPastDispatchFits"
   },
   "bindingSets": {
     "splitAttention": [
           ]
         }
       }
+    ],
+    "gqaMatScoreStats": [
+      {
+        "name": "query",
+        "arg": "queryT",
+        "semantic": "query",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "key",
+        "arg": "keyT",
+        "semantic": "key",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "scorePartials",
+        "semantic": "materializedScorePartials",
+        "buffer": { "type": "storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" },
+            { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" },
+            { "name": "scale", "type": "f32", "value": "attrs.scale if has(attrs, \"scale\") else 0" }
+          ]
+        }
+      }
+    ],
+    "gqaMatRowStatsCombine": [
+      {
+        "name": "scorePartials",
+        "semantic": "materializedScorePartials",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      { "name": "rowStats", "semantic": "materializedRowStats", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "rows", "type": "u32", "value": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)" }
+          ]
+        }
+      }
+    ],
+    "gqaMatApplyFused": [
+      {
+        "name": "scores",
+        "semantic": "materializedScores",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "value",
+        "arg": "valueT",
+        "semantic": "value",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "rowStats",
+        "semantic": "materializedRowStats",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "output",
+        "arg": "outputT",
+        "semantic": "output",
+        "buffer": { "type": "storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "qSeq", "type": "u32", "value": "dim(shapes.query, 1)" },
+            { "name": "kvSeq", "type": "u32", "value": "dim(shapes.key, 1)" }
+          ]
+        }
+      }
+    ],
+    "gqaMatPastScoreStats": [
+      { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "scorePartials",
+        "semantic": "materializedScorePartials",
+        "buffer": { "type": "storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
+            { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" },
+            { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }
+          ]
+        }
+      }
+    ],
+    "gqaMatWindowScoreStats": [
+      { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "scorePartials",
+        "semantic": "materializedScorePartials",
+        "buffer": { "type": "storage" },
+        "elementType": "f32"
+      },
+      { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
+            { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" },
+            { "name": "windowSize", "type": "u32", "value": "attrs.local_window_size" },
+            { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }
+          ]
+        }
+      }
+    ],
+    "gqaMatWindowApplyFused": [
+      {
+        "name": "scores",
+        "semantic": "materializedScores",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      { "name": "value", "arg": "presentValueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      {
+        "name": "rowStats",
+        "semantic": "materializedRowStats",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "f32" },
+      { "name": "seqlens_k", "arg": "seqlensKT", "buffer": { "type": "read-only-storage" }, "elementType": "i32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
+            { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" },
+            { "name": "windowSize", "type": "u32", "value": "attrs.local_window_size" }
+          ]
+        }
+      }
+    ],
+    "gqaMatPastRotaryScoreStats": [
+      { "name": "query", "semantic": "qPrep", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "scorePartials",
+        "semantic": "materializedScorePartials",
+        "buffer": { "type": "storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
+            { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" },
+            { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }
+          ]
+        }
+      }
+    ],
+    "rotaryQprep": [
+      { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" },
+      {
+        "name": "cos_cache",
+        "arg": "cosCacheT",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$cosScalar"
+      },
+      {
+        "name": "sin_cache",
+        "arg": "sinCacheT",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$cosScalar"
+      },
+      { "name": "qout", "semantic": "qPrep", "buffer": { "type": "storage" }, "elementType": "vec4<f32>" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" },
+            { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
+            { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }
+          ]
+        }
+      }
+    ],
+    "gqaMatPastRowStatsCombine": [
+      {
+        "name": "scorePartials",
+        "semantic": "materializedScorePartials",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      { "name": "rowStats", "semantic": "materializedRowStats", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            {
+              "name": "rows",
+              "type": "u32",
+              "value": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)"
+            }
+          ]
+        }
+      }
+    ],
+    "gqaMatPastApplyFused": [
+      {
+        "name": "scores",
+        "semantic": "materializedScores",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      { "name": "value", "arg": "presentValueT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      {
+        "name": "rowStats",
+        "semantic": "materializedRowStats",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
+            { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }
+          ]
+        }
+      }
+    ],
+    "gqaMatPastScoreStatsF16": [
+      { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f16" },
+      { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f16" },
+      { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "scorePartials",
+        "semantic": "materializedScorePartials",
+        "buffer": { "type": "storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
+            { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" },
+            { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }
+          ]
+        }
+      }
+    ],
+    "gqaMatPastApplyFusedF16": [
+      {
+        "name": "scores",
+        "semantic": "materializedScores",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      { "name": "value", "arg": "presentValueT", "buffer": { "type": "read-only-storage" }, "elementType": "f16" },
+      {
+        "name": "rowStats",
+        "semantic": "materializedRowStats",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      { "name": "output", "arg": "outputT", "buffer": { "type": "storage" }, "elementType": "f16" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
+            { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" }
+          ]
+        }
+      }
+    ],
+    "gqaMatPastRotaryScoreStatsF16": [
+      { "name": "query", "semantic": "qPrep", "buffer": { "type": "read-only-storage" }, "elementType": "f16" },
+      { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f16" },
+      { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "scorePartials",
+        "semantic": "materializedScorePartials",
+        "buffer": { "type": "storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
+            { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" },
+            { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" }
+          ]
+        }
+      }
+    ],
+    "rotaryQprepF16": [
+      { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "$inputScalar" },
+      {
+        "name": "cos_cache",
+        "arg": "cosCacheT",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$cosScalar"
+      },
+      {
+        "name": "sin_cache",
+        "arg": "sinCacheT",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$cosScalar"
+      },
+      { "name": "qout", "semantic": "qPrep", "buffer": { "type": "storage" }, "elementType": "vec4<f16>" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "batch", "type": "u32", "value": "dim(shapes.queryT, 0)" },
+            { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
+            { "name": "totalSeq", "type": "u32", "value": "dim(shapes.pastKeyT, 2)" }
+          ]
+        }
+      }
+    ],
+    "gqaMatPastBiasScoreStats": [
+      { "name": "query", "arg": "queryT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      { "name": "key", "arg": "presentKeyT", "buffer": { "type": "read-only-storage" }, "elementType": "f32" },
+      {
+        "name": "attention_bias",
+        "arg": "attentionBiasT",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      { "name": "scores", "semantic": "materializedScores", "buffer": { "type": "storage" }, "elementType": "f32" },
+      {
+        "name": "scorePartials",
+        "semantic": "materializedScorePartials",
+        "buffer": { "type": "storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "qSeq", "type": "u32", "value": "dim(shapes.queryT, 1)" },
+            { "name": "kvSeq", "type": "u32", "value": "dim(shapes.presentKeyT, 2)" },
+            { "name": "scale", "type": "f32", "value": "attrs.scale if attrs.scale else 0" },
+            { "name": "biasBatch", "type": "u32", "value": "dim(shapes.attentionBiasT, 0)" },
+            { "name": "biasHeads", "type": "u32", "value": "dim(shapes.attentionBiasT, 1)" }
+          ]
+        }
+      }
     ]
   },
   "variants": [
+    {
+      "id": "qkv_present_materialized_sgmat_f32",
+      "description": "Materialized float32 subgroup-matrix prefill for the bidirectional no-past qkv route, where `seqlens_k` is metadata-only like the flash routes: the score pass emits per-row softmax statistics, the apply pass folds the softmax, and the present copy is unchanged. It serves square-ish f32 prefill where the score and value GEMMs dominate.",
+      "priority": 23,
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "when": ["gqaMatCoreOk"],
+      "constants": {
+        "qNumHeads": "attrs.num_heads",
+        "kvNumHeads": "attrs.kv_num_heads",
+        "headDim": "dim(shapes.query, 2) / attrs.num_heads",
+        "qHidden": "dim(shapes.query, 2)",
+        "kvHidden": "dim(shapes.key, 2)",
+        "hasBias": false,
+        "useSubgroups": true,
+        "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"",
+        "materializedSgmatQueryTile": "gqaMatQueryTile",
+        "materializedSgmatKeyTile": "gqaMatKeyTile",
+        "materializedSgmatInnerTile": "gqaMatInnerTile",
+        "materializedSgmatDirectScoreStore": false,
+        "materializedSgmatDirectApplyStore": false,
+        "materializedSgmatRuntimeDirectStore": false,
+        "materializedRowStatsWg": "gqaMatRowStatsWg",
+        "statSlots": "gqaMatStatSlots",
+        "statQuerySeq": "dim(shapes.query, 1)",
+        "presentScalar": "qkvCacheScalar",
+        "presentElement": "qkvCacheVec4",
+        "presentVec4": true,
+        "copyWorkgroupSize": "copyWorkgroupSize",
+        "inputElement": "qkvInputVec4",
+        "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4",
+        "kvHiddenV4": "dim(shapes.key, 2) / 4"
+      },
+      "intermediates": [
+        {
+          "id": "materializedScores",
+          "dtype": "float32",
+          "shape": "[dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1) * dim(shapes.key, 1)]"
+        },
+        { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatRowStatsElements]" },
+        { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatScorePartialElements]" }
+      ],
+      "passes": [
+        {
+          "id": "scores",
+          "name": "GroupQueryAttention.MaterializedScoresSgmat",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": { "phase": "\"score\"", "emitRowStats": true, "layout": "\"bsh\"" }
+          },
+          "bindings": "gqaMatScoreStats",
+          "dispatch": {
+            "x": "ceilDiv(dim(shapes.key, 1), gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.query, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.query, 0) * attrs.num_heads"
+          }
+        },
+        {
+          "id": "rowstats",
+          "name": "GroupQueryAttention.MaterializedRowStatsCombine",
+          "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja" },
+          "bindings": "gqaMatRowStatsCombine",
+          "dispatch": {
+            "gridStride": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)",
+            "workgroupSize": "gqaMatRowStatsWg"
+          }
+        },
+        {
+          "id": "apply",
+          "name": "GroupQueryAttention.MaterializedApplySgmat",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": { "phase": "\"apply\"", "fusedSoftmax": true, "layout": "\"bsh\"" }
+          },
+          "bindings": "gqaMatApplyFused",
+          "dispatch": {
+            "x": "ceilDiv(headDim, gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.query, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.query, 0) * attrs.num_heads"
+          }
+        },
+        {
+          "id": "present",
+          "name": "GroupQueryAttention.MaterializedPresent",
+          "shader": "gqa-present.wgsl.jinja",
+          "bindings": "qkvPresent",
+          "dispatch": { "threads": "numel(shapes.presentKeyT) / 4", "workgroupSize": "constants.copyWorkgroupSize" }
+        }
+      ]
+    },
+    {
+      "id": "past_kv_materialized_sgmat_f32",
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "description": "Materializes the causal score matrix with float32 subgroup-matrix tiles over the shared float cache and applies the softmax-normalized weights with the same tiles. The score pass skips key tiles past each query tile's causal bound and the apply pass stops its reduction there, matching flash's triangular work at matrix-unit throughput.",
+      "priority": 35,
+      "when": ["gqaMatPastCoreOk"],
+      "constants": {
+        "qNumHeads": "attrs.num_heads",
+        "kvNumHeads": "attrs.kv_num_heads",
+        "headDim": "gqaHeadDim",
+        "qHidden": "dim(shapes.queryT, 2)",
+        "hasBias": false,
+        "useSubgroups": true,
+        "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"",
+        "materializedSgmatQueryTile": "gqaMatQueryTile",
+        "materializedSgmatKeyTile": "gqaMatKeyTile",
+        "materializedSgmatInnerTile": "gqaMatInnerTile",
+        "materializedSgmatDirectScoreStore": false,
+        "materializedSgmatDirectApplyStore": false,
+        "materializedSgmatRuntimeDirectStore": false,
+        "materializedRowStatsWg": "gqaMatRowStatsWg",
+        "statSlots": "gqaMatPastStatSlots",
+        "statQuerySeq": "dim(shapes.queryT, 1)",
+        "mode": "\"copy\"",
+        "packed": "gqaHeadDim",
+        "kvHeads": "attrs.kv_num_heads",
+        "inputScalar": "gqaScalar",
+        "usesF16": false,
+        "copyWorkgroupSize": "copyWorkgroupSize"
+      },
+      "intermediates": [
+        {
+          "id": "materializedScores",
+          "dtype": "float32",
+          "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2)]"
+        },
+        { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatPastRowStatsElements]" },
+        { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatPastScorePartialElements]" }
+      ],
+      "passes": [
+        {
+          "id": "present",
+          "name": "GroupQueryAttention.Present",
+          "shader": "gqa-present.wgsl.jinja",
+          "bindings": "presentCopy",
+          "dispatch": {
+            "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)",
+            "workgroupSize": "copyWorkgroupSize"
+          },
+          "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }]
+        },
+        {
+          "id": "scores",
+          "name": "GroupQueryAttention.PastMaterializedScores",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": {
+              "phase": "\"score\"",
+              "emitRowStats": true,
+              "layout": "\"bsh\"",
+              "kvLayout": "\"bhsd\"",
+              "causalRightAlign": true
+            }
+          },
+          "bindings": "gqaMatPastScoreStats",
+          "dispatch": {
+            "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
+        },
+        {
+          "id": "rowstats",
+          "name": "GroupQueryAttention.PastMaterializedRowStats",
+          "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} },
+          "bindings": "gqaMatPastRowStatsCombine",
+          "dispatch": {
+            "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)",
+            "workgroupSize": "gqaMatRowStatsWg"
+          }
+        },
+        {
+          "id": "apply",
+          "name": "GroupQueryAttention.PastMaterializedApply",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": {
+              "phase": "\"apply\"",
+              "fusedSoftmax": true,
+              "layout": "\"bsh\"",
+              "kvLayout": "\"bhsd\"",
+              "causalRightAlign": true
+            }
+          },
+          "bindings": "gqaMatPastApplyFused",
+          "dispatch": {
+            "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
+        }
+      ]
+    },
+    {
+      "id": "past_kv_bias_materialized_sgmat_f32",
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "description": "Adds the additive attention bias (batch/head-broadcast at runtime) to the materialized causal scores as they are stored, so the emitted row statistics fold it in for free; everything else matches `past_kv_materialized_sgmat_f32`.",
+      "priority": 36,
+      "when": ["gqaMatPastBiasOk"],
+      "constants": {
+        "qNumHeads": "attrs.num_heads",
+        "kvNumHeads": "attrs.kv_num_heads",
+        "headDim": "gqaHeadDim",
+        "qHidden": "dim(shapes.queryT, 2)",
+        "hasBias": false,
+        "useSubgroups": true,
+        "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"",
+        "materializedSgmatQueryTile": "gqaMatQueryTile",
+        "materializedSgmatKeyTile": "gqaMatKeyTile",
+        "materializedSgmatInnerTile": "gqaMatInnerTile",
+        "materializedSgmatDirectScoreStore": false,
+        "materializedSgmatDirectApplyStore": false,
+        "materializedSgmatRuntimeDirectStore": false,
+        "materializedRowStatsWg": "gqaMatRowStatsWg",
+        "statSlots": "gqaMatPastStatSlots",
+        "statQuerySeq": "dim(shapes.queryT, 1)",
+        "mode": "\"copy\"",
+        "packed": "gqaHeadDim",
+        "kvHeads": "attrs.kv_num_heads",
+        "inputScalar": "gqaScalar",
+        "usesF16": false,
+        "copyWorkgroupSize": "copyWorkgroupSize"
+      },
+      "intermediates": [
+        {
+          "id": "materializedScores",
+          "dtype": "float32",
+          "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2)]"
+        },
+        { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatPastRowStatsElements]" },
+        { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatPastScorePartialElements]" }
+      ],
+      "passes": [
+        {
+          "id": "present",
+          "name": "GroupQueryAttention.Present",
+          "shader": "gqa-present.wgsl.jinja",
+          "bindings": "presentCopy",
+          "dispatch": {
+            "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)",
+            "workgroupSize": "copyWorkgroupSize"
+          },
+          "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }]
+        },
+        {
+          "id": "scores",
+          "name": "GroupQueryAttention.PastMaterializedScores",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": {
+              "phase": "\"score\"",
+              "emitRowStats": true,
+              "layout": "\"bsh\"",
+              "kvLayout": "\"bhsd\"",
+              "causalRightAlign": true,
+              "scoreBias": true
+            }
+          },
+          "bindings": "gqaMatPastBiasScoreStats",
+          "dispatch": {
+            "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
+        },
+        {
+          "id": "rowstats",
+          "name": "GroupQueryAttention.PastMaterializedRowStats",
+          "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} },
+          "bindings": "gqaMatPastRowStatsCombine",
+          "dispatch": {
+            "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)",
+            "workgroupSize": "gqaMatRowStatsWg"
+          }
+        },
+        {
+          "id": "apply",
+          "name": "GroupQueryAttention.PastMaterializedApply",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": {
+              "phase": "\"apply\"",
+              "fusedSoftmax": true,
+              "layout": "\"bsh\"",
+              "kvLayout": "\"bhsd\"",
+              "causalRightAlign": true
+            }
+          },
+          "bindings": "gqaMatPastApplyFused",
+          "dispatch": {
+            "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
+        }
+      ]
+    },
+    {
+      "id": "past_kv_materialized_sgmat_f16",
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "description": "Materializes the causal score matrix with `f16` operand tiles feeding `f32`-accumulating subgroup matrices over the shared cache and applies the softmax-normalized weights the same way; scores and row statistics stay `f32`. The score pass skips key tiles past each query tile's causal bound and the apply pass stops its reduction there.",
+      "priority": 35,
+      "when": ["gqaMatPastCoreF16Ok"],
+      "constants": {
+        "qNumHeads": "attrs.num_heads",
+        "kvNumHeads": "attrs.kv_num_heads",
+        "headDim": "gqaHeadDim",
+        "qHidden": "dim(shapes.queryT, 2)",
+        "hasBias": false,
+        "useSubgroups": true,
+        "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"",
+        "materializedSgmatQueryTile": "gqaMatQueryTile",
+        "materializedSgmatKeyTile": "gqaMatKeyTile",
+        "materializedSgmatInnerTile": "gqaMatInnerTile",
+        "materializedSgmatDirectScoreStore": false,
+        "materializedSgmatDirectApplyStore": false,
+        "materializedSgmatRuntimeDirectStore": false,
+        "materializedRowStatsWg": "gqaMatRowStatsWg",
+        "statSlots": "gqaMatPastStatSlots",
+        "statQuerySeq": "dim(shapes.queryT, 1)",
+        "mode": "\"copy\"",
+        "packed": "gqaHeadDim",
+        "kvHeads": "attrs.kv_num_heads",
+        "inputScalar": "gqaScalar",
+        "usesF16": true,
+        "copyWorkgroupSize": "copyWorkgroupSize",
+        "operandF16": true
+      },
+      "intermediates": [
+        {
+          "id": "materializedScores",
+          "dtype": "float32",
+          "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2)]"
+        },
+        { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatPastRowStatsElements]" },
+        { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatPastScorePartialElements]" }
+      ],
+      "passes": [
+        {
+          "id": "present",
+          "name": "GroupQueryAttention.Present",
+          "shader": "gqa-present.wgsl.jinja",
+          "bindings": "presentCopy",
+          "dispatch": {
+            "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)",
+            "workgroupSize": "copyWorkgroupSize"
+          },
+          "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }]
+        },
+        {
+          "id": "scores",
+          "name": "GroupQueryAttention.PastMaterializedScores",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": {
+              "phase": "\"score\"",
+              "emitRowStats": true,
+              "layout": "\"bsh\"",
+              "kvLayout": "\"bhsd\"",
+              "causalRightAlign": true
+            }
+          },
+          "bindings": "gqaMatPastScoreStatsF16",
+          "dispatch": {
+            "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
+        },
+        {
+          "id": "rowstats",
+          "name": "GroupQueryAttention.PastMaterializedRowStats",
+          "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} },
+          "bindings": "gqaMatPastRowStatsCombine",
+          "dispatch": {
+            "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)",
+            "workgroupSize": "gqaMatRowStatsWg"
+          }
+        },
+        {
+          "id": "apply",
+          "name": "GroupQueryAttention.PastMaterializedApply",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": {
+              "phase": "\"apply\"",
+              "fusedSoftmax": true,
+              "layout": "\"bsh\"",
+              "kvLayout": "\"bhsd\"",
+              "causalRightAlign": true
+            }
+          },
+          "bindings": "gqaMatPastApplyFusedF16",
+          "dispatch": {
+            "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
+        }
+      ]
+    },
+    {
+      "id": "past_kv_rotary_materialized_sgmat_f32",
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "description": "Rotary twin of the causal materialized route: the shared cache already holds rotary-transformed keys, so one preparation pass rotates the query block at its absolute positions and the tile-skipping score/apply passes run unchanged on the prepared buffer.",
+      "priority": 37,
+      "when": ["gqaMatPastRotaryOk"],
+      "constants": {
+        "qNumHeads": "attrs.num_heads",
+        "kvNumHeads": "attrs.kv_num_heads",
+        "headDim": "gqaHeadDim",
+        "qHidden": "dim(shapes.queryT, 2)",
+        "hasBias": false,
+        "useSubgroups": true,
+        "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"",
+        "materializedSgmatQueryTile": "gqaMatQueryTile",
+        "materializedSgmatKeyTile": "gqaMatKeyTile",
+        "materializedSgmatInnerTile": "gqaMatInnerTile",
+        "materializedSgmatDirectScoreStore": false,
+        "materializedSgmatDirectApplyStore": false,
+        "materializedSgmatRuntimeDirectStore": false,
+        "materializedRowStatsWg": "gqaMatRowStatsWg",
+        "statSlots": "gqaMatPastStatSlots",
+        "statQuerySeq": "dim(shapes.queryT, 1)",
+        "mode": "\"copy\"",
+        "packed": "gqaHeadDim",
+        "kvHeads": "attrs.kv_num_heads",
+        "inputScalar": "gqaScalar",
+        "usesF16": false,
+        "copyWorkgroupSize": "copyWorkgroupSize",
+        "half": "gqaHeadDim / 2",
+        "qHeads": "attrs.num_heads",
+        "cosScalar": "\"f16\" if tensorDtypes.cosCacheT == \"float16\" else \"f32\"",
+        "hasQNorm": false
+      },
+      "intermediates": [
+        {
+          "id": "qPrep",
+          "dtype": "float32",
+          "shape": "[dim(shapes.queryT, 0), dim(shapes.queryT, 1), dim(shapes.queryT, 2)]"
+        },
+        {
+          "id": "materializedScores",
+          "dtype": "float32",
+          "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2)]"
+        },
+        { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatPastRowStatsElements]" },
+        { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatPastScorePartialElements]" }
+      ],
+      "passes": [
+        {
+          "id": "present",
+          "name": "GroupQueryAttention.Present",
+          "shader": "gqa-present.wgsl.jinja",
+          "bindings": "presentCopy",
+          "dispatch": {
+            "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)",
+            "workgroupSize": "copyWorkgroupSize"
+          },
+          "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }]
+        },
+        {
+          "id": "qprep",
+          "name": "GroupQueryAttention.RotaryQPrep",
+          "shader": "gqa-qprep.wgsl.jinja",
+          "bindings": "rotaryQprep",
+          "dispatch": {
+            "threads": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)",
+            "workgroupSize": "copyWorkgroupSize"
+          }
+        },
+        {
+          "id": "scores",
+          "name": "GroupQueryAttention.PastRotaryMaterializedScores",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": {
+              "phase": "\"score\"",
+              "emitRowStats": true,
+              "layout": "\"bsh\"",
+              "kvLayout": "\"bhsd\"",
+              "causalRightAlign": true
+            }
+          },
+          "bindings": "gqaMatPastRotaryScoreStats",
+          "dispatch": {
+            "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
+        },
+        {
+          "id": "rowstats",
+          "name": "GroupQueryAttention.PastRotaryMaterializedRowStats",
+          "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} },
+          "bindings": "gqaMatPastRowStatsCombine",
+          "dispatch": {
+            "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)",
+            "workgroupSize": "gqaMatRowStatsWg"
+          }
+        },
+        {
+          "id": "apply",
+          "name": "GroupQueryAttention.PastRotaryMaterializedApply",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": {
+              "phase": "\"apply\"",
+              "fusedSoftmax": true,
+              "layout": "\"bsh\"",
+              "kvLayout": "\"bhsd\"",
+              "causalRightAlign": true
+            }
+          },
+          "bindings": "gqaMatPastApplyFused",
+          "dispatch": {
+            "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
+        }
+      ]
+    },
+    {
+      "id": "past_kv_rotary_materialized_sgmat_f16",
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "description": "Materializes the causal score matrix with `f16` operand tiles feeding `f32`-accumulating subgroup matrices over the shared cache and applies the softmax-normalized weights the same way; scores and row statistics stay `f32`. The score pass skips key tiles past each query tile's causal bound and the apply pass stops its reduction there.",
+      "priority": 37,
+      "when": ["gqaMatPastRotaryF16Ok"],
+      "constants": {
+        "qNumHeads": "attrs.num_heads",
+        "kvNumHeads": "attrs.kv_num_heads",
+        "headDim": "gqaHeadDim",
+        "qHidden": "dim(shapes.queryT, 2)",
+        "hasBias": false,
+        "useSubgroups": true,
+        "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"",
+        "materializedSgmatQueryTile": "gqaMatQueryTile",
+        "materializedSgmatKeyTile": "gqaMatKeyTile",
+        "materializedSgmatInnerTile": "gqaMatInnerTile",
+        "materializedSgmatDirectScoreStore": false,
+        "materializedSgmatDirectApplyStore": false,
+        "materializedSgmatRuntimeDirectStore": false,
+        "materializedRowStatsWg": "gqaMatRowStatsWg",
+        "statSlots": "gqaMatPastStatSlots",
+        "statQuerySeq": "dim(shapes.queryT, 1)",
+        "mode": "\"copy\"",
+        "packed": "gqaHeadDim",
+        "kvHeads": "attrs.kv_num_heads",
+        "inputScalar": "gqaScalar",
+        "usesF16": true,
+        "copyWorkgroupSize": "copyWorkgroupSize",
+        "half": "gqaHeadDim / 2",
+        "qHeads": "attrs.num_heads",
+        "cosScalar": "\"f16\" if tensorDtypes.cosCacheT == \"float16\" else \"f32\"",
+        "hasQNorm": false,
+        "operandF16": true
+      },
+      "intermediates": [
+        {
+          "id": "qPrep",
+          "dtype": "float16",
+          "shape": "[dim(shapes.queryT, 0), dim(shapes.queryT, 1), dim(shapes.queryT, 2)]"
+        },
+        {
+          "id": "materializedScores",
+          "dtype": "float32",
+          "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2)]"
+        },
+        { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatPastRowStatsElements]" },
+        { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatPastScorePartialElements]" }
+      ],
+      "passes": [
+        {
+          "id": "present",
+          "name": "GroupQueryAttention.Present",
+          "shader": "gqa-present.wgsl.jinja",
+          "bindings": "presentCopy",
+          "dispatch": {
+            "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.pastKeyT, 2)",
+            "workgroupSize": "copyWorkgroupSize"
+          },
+          "viewAlias": [{ "input": "src_k", "output": "present_key" }, { "input": "src_v", "output": "present_value" }]
+        },
+        {
+          "id": "qprep",
+          "name": "GroupQueryAttention.RotaryQPrep",
+          "shader": "gqa-qprep.wgsl.jinja",
+          "bindings": "rotaryQprepF16",
+          "dispatch": {
+            "threads": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)",
+            "workgroupSize": "copyWorkgroupSize"
+          },
+          "constants": { "qPrepF16": true }
+        },
+        {
+          "id": "scores",
+          "name": "GroupQueryAttention.PastRotaryMaterializedScores",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": {
+              "phase": "\"score\"",
+              "emitRowStats": true,
+              "layout": "\"bsh\"",
+              "kvLayout": "\"bhsd\"",
+              "causalRightAlign": true
+            }
+          },
+          "bindings": "gqaMatPastRotaryScoreStatsF16",
+          "dispatch": {
+            "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
+        },
+        {
+          "id": "rowstats",
+          "name": "GroupQueryAttention.PastRotaryMaterializedRowStats",
+          "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} },
+          "bindings": "gqaMatPastRowStatsCombine",
+          "dispatch": {
+            "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)",
+            "workgroupSize": "gqaMatRowStatsWg"
+          }
+        },
+        {
+          "id": "apply",
+          "name": "GroupQueryAttention.PastRotaryMaterializedApply",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": {
+              "phase": "\"apply\"",
+              "fusedSoftmax": true,
+              "layout": "\"bsh\"",
+              "kvLayout": "\"bhsd\"",
+              "causalRightAlign": true
+            }
+          },
+          "bindings": "gqaMatPastApplyFusedF16",
+          "dispatch": {
+            "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
+        }
+      ]
+    },
+    {
+      "id": "new_kv_past_materialized_sgmat_f32",
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "description": "Append-route twin of the causal materialized route for chunked prefill: one merge pass concatenates the past cache with the new key/value rows, and the tile-skipping score/apply passes run over the merged cache with the same right-aligned causal bound.",
+      "priority": 35,
+      "when": ["gqaMatNewPastOk"],
+      "constants": {
+        "qNumHeads": "attrs.num_heads",
+        "kvNumHeads": "attrs.kv_num_heads",
+        "headDim": "gqaHeadDim",
+        "qHidden": "dim(shapes.queryT, 2)",
+        "hasBias": false,
+        "useSubgroups": true,
+        "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"",
+        "materializedSgmatQueryTile": "gqaMatQueryTile",
+        "materializedSgmatKeyTile": "gqaMatKeyTile",
+        "materializedSgmatInnerTile": "gqaMatInnerTile",
+        "materializedSgmatDirectScoreStore": false,
+        "materializedSgmatDirectApplyStore": false,
+        "materializedSgmatRuntimeDirectStore": false,
+        "materializedRowStatsWg": "gqaMatRowStatsWg",
+        "statSlots": "gqaMatPastStatSlots",
+        "statQuerySeq": "dim(shapes.queryT, 1)",
+        "mode": "\"merge\"",
+        "packed": "gqaHeadDim",
+        "kvHeads": "attrs.kv_num_heads",
+        "inputScalar": "gqaScalar",
+        "usesF16": false,
+        "copyWorkgroupSize": "copyWorkgroupSize"
+      },
+      "intermediates": [
+        {
+          "id": "materializedScores",
+          "dtype": "float32",
+          "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2)]"
+        },
+        { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatPastRowStatsElements]" },
+        { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatPastScorePartialElements]" }
+      ],
+      "passes": [
+        {
+          "id": "present",
+          "name": "GroupQueryAttention.Merge",
+          "shader": "gqa-present.wgsl.jinja",
+          "bindings": "presentMerge",
+          "dispatch": {
+            "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)",
+            "workgroupSize": "copyWorkgroupSize"
+          }
+        },
+        {
+          "id": "scores",
+          "name": "GroupQueryAttention.NewPastMaterializedScores",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": {
+              "phase": "\"score\"",
+              "emitRowStats": true,
+              "layout": "\"bsh\"",
+              "kvLayout": "\"bhsd\"",
+              "causalRightAlign": true
+            }
+          },
+          "bindings": "gqaMatPastScoreStats",
+          "dispatch": {
+            "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
+        },
+        {
+          "id": "rowstats",
+          "name": "GroupQueryAttention.NewPastMaterializedRowStats",
+          "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} },
+          "bindings": "gqaMatPastRowStatsCombine",
+          "dispatch": {
+            "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)",
+            "workgroupSize": "gqaMatRowStatsWg"
+          }
+        },
+        {
+          "id": "apply",
+          "name": "GroupQueryAttention.NewPastMaterializedApply",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": {
+              "phase": "\"apply\"",
+              "fusedSoftmax": true,
+              "layout": "\"bsh\"",
+              "kvLayout": "\"bhsd\"",
+              "causalRightAlign": true
+            }
+          },
+          "bindings": "gqaMatPastApplyFused",
+          "dispatch": {
+            "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
+        }
+      ]
+    },
+    {
+      "id": "window_shift_materialized_sgmat_f32",
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
+      },
+      "description": "Windowed-cache twin of the causal materialized route for chunked prefill: the shift pass compacts surviving cache rows and appends the chunk, the score pass masks each row's sliding-window floor alongside the causal bound so the row statistics fold both out, and the apply pass starts each reduction at the first key tile the window can reach.",
+      "priority": 36,
+      "when": ["gqaMatWindowOk"],
+      "constants": {
+        "qNumHeads": "attrs.num_heads",
+        "kvNumHeads": "attrs.kv_num_heads",
+        "headDim": "gqaHeadDim",
+        "qHidden": "dim(shapes.queryT, 2)",
+        "hasBias": false,
+        "useSubgroups": true,
+        "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\"",
+        "materializedSgmatQueryTile": "gqaMatQueryTile",
+        "materializedSgmatKeyTile": "gqaMatKeyTile",
+        "materializedSgmatInnerTile": "gqaMatInnerTile",
+        "materializedSgmatDirectScoreStore": false,
+        "materializedSgmatDirectApplyStore": false,
+        "materializedSgmatRuntimeDirectStore": false,
+        "materializedRowStatsWg": "gqaMatRowStatsWg",
+        "statSlots": "gqaMatPastStatSlots",
+        "statQuerySeq": "dim(shapes.queryT, 1)",
+        "mode": "\"window_shift\"",
+        "packed": "gqaHeadDim",
+        "kvHeads": "attrs.kv_num_heads",
+        "inputScalar": "gqaScalar",
+        "zeroScalar": "gqaScalar",
+        "usesF16": false,
+        "copyWorkgroupSize": "copyWorkgroupSize"
+      },
+      "intermediates": [
+        {
+          "id": "materializedScores",
+          "dtype": "float32",
+          "shape": "[dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1) * dim(shapes.presentKeyT, 2)]"
+        },
+        { "id": "materializedRowStats", "dtype": "float32", "shape": "[gqaMatPastRowStatsElements]" },
+        { "id": "materializedScorePartials", "dtype": "float32", "shape": "[gqaMatPastScorePartialElements]" }
+      ],
+      "passes": [
+        {
+          "id": "present",
+          "name": "GroupQueryAttention.WindowShift",
+          "shader": "gqa-present.wgsl.jinja",
+          "bindings": "presentMergeShare",
+          "dispatch": {
+            "threads": "dim(shapes.queryT, 0) * attrs.kv_num_heads * dim(shapes.presentKeyT, 2)",
+            "workgroupSize": "copyWorkgroupSize"
+          }
+        },
+        {
+          "id": "scores",
+          "name": "GroupQueryAttention.WindowMaterializedScores",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": {
+              "phase": "\"score\"",
+              "emitRowStats": true,
+              "layout": "\"bsh\"",
+              "kvLayout": "\"bhsd\"",
+              "causalRightAlign": true,
+              "scoreWindow": true,
+              "useSeqlens": true
+            }
+          },
+          "bindings": "gqaMatWindowScoreStats",
+          "dispatch": {
+            "x": "ceilDiv(dim(shapes.presentKeyT, 2), gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
+        },
+        {
+          "id": "rowstats",
+          "name": "GroupQueryAttention.WindowMaterializedRowStats",
+          "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja", "inputs": {} },
+          "bindings": "gqaMatPastRowStatsCombine",
+          "dispatch": {
+            "gridStride": "dim(shapes.queryT, 0) * attrs.num_heads * dim(shapes.queryT, 1)",
+            "workgroupSize": "gqaMatRowStatsWg"
+          }
+        },
+        {
+          "id": "apply",
+          "name": "GroupQueryAttention.WindowMaterializedApply",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": {
+              "phase": "\"apply\"",
+              "fusedSoftmax": true,
+              "layout": "\"bsh\"",
+              "kvLayout": "\"bhsd\"",
+              "causalRightAlign": true,
+              "scoreWindow": true,
+              "useSeqlens": true
+            }
+          },
+          "bindings": "gqaMatWindowApplyFused",
+          "dispatch": {
+            "x": "ceilDiv(gqaHeadDim, gqaMatKeyTile)",
+            "y": "ceilDiv(dim(shapes.queryT, 1), gqaMatQueryTile)",
+            "z": "dim(shapes.queryT, 0) * attrs.num_heads"
+          }
+        }
+      ]
+    },
     {
       "id": "new_kv_share_append_split",
       "description": "Retains the existing cache and appends new key/value rows in separate passes before portable attention. It avoids rebuilding unchanged cache positions when the past and present allocations share capacity.",

build/webgpu/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "com.microsoft.GroupQueryAttention",
-  "id": "_com_microsoft_groupqueryattention_webgpu_75fbc9b",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
@@ -13,15 +13,17 @@
       "attn-flash-online.wgsl.jinja": "MStTMF1/l5L9Rzx1+rPz5T+70/46n/vO/+hZQ0lbFws=",
       "attn-flash-prefill-cluster.wgsl.jinja": "yk1WymJIbxCy0slbiho5DS223uZH9PQ15629yvQ9RFc=",
       "attn-flash-q32-broadcast.wgsl.jinja": "evSdaiYgvrikYZ8hpb5m4gxf+rIryKkHOTiKgyF7ioI=",
       "attn-online-scalar.wgsl.jinja": "X4yGNCYu9+uJX6Q3EKkxV6v28CPlIYWme0Nmf8F16L8=",
-      "bench.json": "rGPC3cQInjJZjL+u+MzVXCOGu4L+G0v+/0UGDI9DEuQ=",
       "gqa-attention.wgsl.jinja": "0EbwmdstNZMuskS170zYEo8whbQOFtupXLnNP2PaFL0=",
       "gqa-present.wgsl.jinja": "J9g6oAeY2BtbmWBMlmbIXbxzo4sKHdxYqcE4EAd7BYM=",
-      "gqa-qprep.wgsl.jinja": "SvA19/AQoN2FXduwBtMI5QzBeeth2y9g8EnaIkV3V/0=",
-      "manifest.json": "qODDtl2Snem6J/K6QJwIL5i6CnGQ898t0a/EEka4JnY=",
-      "test.json": "iZAuQQINr2i9Q2hwXzNhKO+UT+ByfRwPsNc+paVhR20="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
   "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.GroupQueryAttention" }
 }

 {
   "name": "com.microsoft.GroupQueryAttention",
+  "id": "_com_microsoft_groupqueryattention_webgpu_34ecc1b",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
       "attn-flash-online.wgsl.jinja": "MStTMF1/l5L9Rzx1+rPz5T+70/46n/vO/+hZQ0lbFws=",
       "attn-flash-prefill-cluster.wgsl.jinja": "yk1WymJIbxCy0slbiho5DS223uZH9PQ15629yvQ9RFc=",
       "attn-flash-q32-broadcast.wgsl.jinja": "evSdaiYgvrikYZ8hpb5m4gxf+rIryKkHOTiKgyF7ioI=",
+      "attn-materialized-rowstats-combine-f32.wgsl.jinja": "TtTBeeeP/ufhzeNHoeLWGLxKTfFT9fiEAstnXB79j48=",
+      "attn-materialized-sgmat-f32.wgsl.jinja": "iQeacl5PUCl75lx78qlSBaQ/syBDNxgjZg4UPGki9U0=",
       "attn-online-scalar.wgsl.jinja": "X4yGNCYu9+uJX6Q3EKkxV6v28CPlIYWme0Nmf8F16L8=",
+      "bench.json": "b4ZLhmCqTm6AgiNfIiRHYVeNp4q1mx12GMcDAextGqI=",
       "gqa-attention.wgsl.jinja": "0EbwmdstNZMuskS170zYEo8whbQOFtupXLnNP2PaFL0=",
       "gqa-present.wgsl.jinja": "J9g6oAeY2BtbmWBMlmbIXbxzo4sKHdxYqcE4EAd7BYM=",
+      "gqa-qprep.wgsl.jinja": "RCBDy7LA/yy4ayRAsnyq+7/IaXjsDGmPWidj51c+t8k=",
+      "manifest.json": "oxSz/yE5vrd9vU5lHkrRN2kG8sdBzWQmGP9/jjRoqNU=",
+      "test.json": "zfjYb8H4THOE5qN8K90UPdGq0lL/B2LWarbvokAz8Cg="
     }
   },
+  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
   "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.GroupQueryAttention" }
 }

build/webgpu/test.json CHANGED Viewed

@@ -4271,6 +4271,205 @@
         "notes": "Chunked flash prefill into a windowed cache: the present pass compacts entries, the attention bound comes from seqlens_k (kvActive), and the batch stride remains the cache capacity. Batch 0 is unfilled at T=40 while batch 1 evicts at T=80, requiring kvActive to be tracked per batch."
       }
     },
     {
       "name": "share_append_flash_chunk_slack_b1q32cap64_h2kv1d64",
       "attrs": { "num_heads": 2, "kv_num_heads": 1, "local_window_size": 16 },
@@ -4655,6 +4854,526 @@
         "presentKeyT": { "dtype": "uint8", "shape": [1, 1, 4, 4], "tolerance": 0 },
         "presentValueT": { "dtype": "uint8", "shape": [1, 1, 4, 4], "tolerance": 0 }
       }
     }
   ]
 }

         "notes": "Chunked flash prefill into a windowed cache: the present pass compacts entries, the attention bound comes from seqlens_k (kvActive), and the batch stride remains the cache capacity. Batch 0 is unfilled at T=40 while batch 1 evicts at T=80, requiring kvActive to be tracked per batch."
       }
     },
+    {
+      "name": "window_cache_sgmat_chunk_evicts_b1q128cap256_h2kv1d128_w256",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "sliding_window_cache": 1, "local_window_size": 256 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 128, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.27, "scale": 0.5 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 128, 128],
+          "data": { "kind": "linspace", "start": -1.0, "end": 1.0 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 128, 128],
+          "data": { "kind": "linspace", "start": 0.0, "end": 2.0 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 256, 128],
+          "data": { "kind": "linspace", "start": 0.0, "end": 4.0 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [383] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [384] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 128, 256], "tolerance": 0.005 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 256, 128], "tolerance": 0.0001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 256, 128], "tolerance": 0.0001 }
+      },
+      "provenance": {
+        "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
+        "test": "GQA sliding_window_cache chunk prefill on the materialized subgroup-matrix route",
+        "notes": "T=384 > C=256 evicts 128 rows; w=256=C keeps the floor inactive, so this pins the shift + right-aligned causal path of the materialized window route at d=128 (two apply column tiles)."
+      }
+    },
+    {
+      "name": "window_cache_sgmat_chunk_floor_b1q128cap384_h2kv1d64_w256",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "sliding_window_cache": 1, "local_window_size": 256 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.27, "scale": 0.5 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 128, 64], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 128, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 2.0 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 384, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 384, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 4.0 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 128, 128], "tolerance": 0.005 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 384, 64], "tolerance": 0.0001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 384, 64], "tolerance": 0.0001 }
+      },
+      "provenance": {
+        "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
+        "test": "GQA sliding_window_cache chunk prefill on the materialized subgroup-matrix route",
+        "notes": "w+q == C exactly, so every row's window floor is row+1: the second query tile's first key tile is dead (uniform window skip) and the apply pass starts its reduction one tile in. The two-sided band mask is fully exercised."
+      }
+    },
+    {
+      "name": "window_cache_sgmat_chunk_unfilled_b1q128cap256_h2kv1d64_w256",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "sliding_window_cache": 1, "local_window_size": 256 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.27, "scale": 0.5 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 128, 64], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 128, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 2.0 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 256, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 4.0 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [191] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [192] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 128, 128], "tolerance": 0.005 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 256, 64], "tolerance": 0.0001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 256, 64], "tolerance": 0.0001 }
+      },
+      "provenance": {
+        "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
+        "test": "GQA sliding_window_cache chunk prefill on the materialized subgroup-matrix route",
+        "notes": "Partially filled cache: kvActive=192 < C=256, so the live length comes from seqlens_k and the key tiles beyond it publish the stats identity instead of reading cleared rows."
+      }
+    },
+    {
+      "name": "window_cache_sgmat_chunk_mixed_batch_b2q128cap256_h2kv1d64_w192",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "sliding_window_cache": 1, "local_window_size": 192 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 128, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.27, "scale": 0.5 }
+        },
+        "keyT": { "dtype": "float32", "shape": [2, 128, 64], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [2, 128, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 2.0 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 1, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 1, 256, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 4.0 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [319, 191] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [320] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 128, 128], "tolerance": 0.005 },
+        "presentKeyT": { "dtype": "float32", "shape": [2, 1, 256, 64], "tolerance": 0.0001 },
+        "presentValueT": { "dtype": "float32", "shape": [2, 1, 256, 64], "tolerance": 0.0001 }
+      },
+      "provenance": {
+        "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
+        "test": "GQA sliding_window_cache chunk prefill on the materialized subgroup-matrix route",
+        "notes": "Per-batch live lengths: batch 0 evicts (kvActive=256) with an active window floor, batch 1 is unfilled (kvActive=192) with the floor inactive — both from the same dispatch."
+      }
+    },
+    {
+      "name": "window_cache_sgmat_chunk_qtail_b1q160cap320_h2kv1d64_w160",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "sliding_window_cache": 1, "local_window_size": 160 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 160, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.27, "scale": 0.5 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 160, 64], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 160, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 2.0 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 320, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.31, "scale": 0.5 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 320, 64],
+          "data": { "kind": "linspace", "start": 0.0, "end": 4.0 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [399] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [400] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 160, 128], "tolerance": 0.005 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 320, 64], "tolerance": 0.0001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 320, 64], "tolerance": 0.0001 }
+      },
+      "provenance": {
+        "source": "onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc",
+        "test": "GQA sliding_window_cache chunk prefill on the materialized subgroup-matrix route",
+        "notes": "q=160 is not a query-tile multiple and w+q == C makes every floor active, so the guarded tail rows and the tile-aligned reduction start are exercised together."
+      }
+    },
     {
       "name": "share_append_flash_chunk_slack_b1q32cap64_h2kv1d64",
       "attrs": { "num_heads": 2, "kv_num_heads": 1, "local_window_size": 16 },
         "presentKeyT": { "dtype": "uint8", "shape": [1, 1, 4, 4], "tolerance": 0 },
         "presentValueT": { "dtype": "uint8", "shape": [1, 1, 4, 4], "tolerance": 0 }
       }
+    },
+    {
+      "name": "qkv_materialized_sgmat_f32_q512_kv512_h2kv1_d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "causal": 0 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 128], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 512, 64], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 512, 64], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "qkv_materialized_sgmat_f32_padded_seqlens_q512_kv512_h2kv1_d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "causal": 0 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [299] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 128], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 512, 64], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 512, 64], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "qkv_materialized_sgmat_f32_rightalign_q512_kv576_h2kv1_d128",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "causal": 0 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 576, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 576, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [575] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [576] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 256], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 576, 128], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 576, 128], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "past_kv_materialized_sgmat_f32_q512_p512_h2kv1_d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 128], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 512, 64], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 512, 64], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "past_kv_materialized_sgmat_f32_qtail_q520_p576_h2kv1_d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 520, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 576, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 576, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [575] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [576] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 520, 128], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 576, 64], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 576, 64], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "past_kv_materialized_sgmat_f32_group_q512_p512_h4kv2_d128",
+      "attrs": { "num_heads": 4, "kv_num_heads": 2 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 256], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 256], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 2, 512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 2, 512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 512], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 2, 512, 128], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 2, 512, 128], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "past_kv_materialized_sgmat_f32_floor_q256_p256_h2kv1_d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [255] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 256, 128], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 256, 64], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 256, 64], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "past_kv_rotary_materialized_sgmat_f32_q512_p512_h2kv1_d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "do_rotary": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } },
+        "cosCacheT": {
+          "dtype": "float32",
+          "shape": [512, 32],
+          "data": { "kind": "rotaryCos", "thetaStart": 0.07, "thetaStep": 0.013 }
+        },
+        "sinCacheT": {
+          "dtype": "float32",
+          "shape": [512, 32],
+          "data": { "kind": "rotarySin", "thetaStart": 0.07, "thetaStep": 0.013 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 128], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 512, 64], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 512, 64], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "past_kv_rotary_materialized_sgmat_f32_offset_q512_p576_h2kv1_d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "do_rotary": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 576, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 576, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [575] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [576] } },
+        "cosCacheT": {
+          "dtype": "float32",
+          "shape": [576, 32],
+          "data": { "kind": "rotaryCos", "thetaStart": 0.07, "thetaStep": 0.013 }
+        },
+        "sinCacheT": {
+          "dtype": "float32",
+          "shape": [576, 32],
+          "data": { "kind": "rotarySin", "thetaStart": 0.07, "thetaStep": 0.013 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 128], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 576, 64], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 576, 64], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "new_kv_past_materialized_sgmat_f32_chunk_q512_p512_h2kv1_d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.019, "cosStep": 0.037 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.029, "cosStep": 0.013 }
+        },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1023] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [1024] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 128], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 1024, 64], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 1024, 64], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "qkv_materialized_sgmat_f32_floor_q256_kv256_h2kv1_d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "causal": 0 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": {
+          "dtype": "float32",
+          "shape": [1, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "valueT": {
+          "dtype": "float32",
+          "shape": [1, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [255] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 256, 128], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 256, 64], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 256, 64], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "past_kv_materialized_sgmat_f16_q512_p512_h2kv1_d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": { "dtype": "float16", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float16", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float16",
+          "shape": [1, 1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "pastValueT": {
+          "dtype": "float16",
+          "shape": [1, 1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 512, 128], "tolerance": 0.03 },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 1, 512, 64], "tolerance": 0.001 },
+        "presentValueT": { "dtype": "float16", "shape": [1, 1, 512, 64], "tolerance": 0.001 }
+      }
+    },
+    {
+      "name": "past_kv_materialized_sgmat_f16_floor_q256_p256_h2kv1_d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": { "dtype": "float16", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float16", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float16",
+          "shape": [1, 1, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "pastValueT": {
+          "dtype": "float16",
+          "shape": [1, 1, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [255] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 256, 128], "tolerance": 0.03 },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 1, 256, 64], "tolerance": 0.001 },
+        "presentValueT": { "dtype": "float16", "shape": [1, 1, 256, 64], "tolerance": 0.001 }
+      }
+    },
+    {
+      "name": "past_kv_rotary_materialized_sgmat_f16_q512_p512_h2kv1_d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1, "do_rotary": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": { "dtype": "float16", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float16", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float16",
+          "shape": [1, 1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "pastValueT": {
+          "dtype": "float16",
+          "shape": [1, 1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } },
+        "cosCacheT": {
+          "dtype": "float16",
+          "shape": [512, 32],
+          "data": { "kind": "rotaryCos", "thetaStart": 0.07, "thetaStep": 0.013 }
+        },
+        "sinCacheT": {
+          "dtype": "float16",
+          "shape": [512, 32],
+          "data": { "kind": "rotarySin", "thetaStart": 0.07, "thetaStep": 0.013 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float16", "shape": [1, 512, 128], "tolerance": 0.03 },
+        "presentKeyT": { "dtype": "float16", "shape": [1, 1, 512, 64], "tolerance": 0.001 },
+        "presentValueT": { "dtype": "float16", "shape": [1, 1, 512, 64], "tolerance": 0.001 }
+      }
+    },
+    {
+      "name": "past_kv_bias_materialized_sgmat_f32_q512_p512_h2kv1_d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [1, 512, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": { "dtype": "float32", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [1, 0, 64], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [1, 1, 512, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [511] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [512] } },
+        "attentionBiasT": {
+          "dtype": "float32",
+          "shape": [1, 2, 512, 512],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.13, "scale": 0.5 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [1, 512, 128], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [1, 1, 512, 64], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [1, 1, 512, 64], "tolerance": 0.000001 }
+      }
+    },
+    {
+      "name": "past_kv_bias_materialized_sgmat_f32_broadcast_b2_q256_p256_h2kv1_d64",
+      "attrs": { "num_heads": 2, "kv_num_heads": 1 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float32",
+          "shape": [2, 256, 128],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "keyT": { "dtype": "float32", "shape": [2, 0, 64], "data": { "kind": "values", "values": [] } },
+        "valueT": { "dtype": "float32", "shape": [2, 0, 64], "data": { "kind": "values", "values": [] } },
+        "pastKeyT": {
+          "dtype": "float32",
+          "shape": [2, 1, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.11, "cosStep": 0.23 }
+        },
+        "pastValueT": {
+          "dtype": "float32",
+          "shape": [2, 1, 256, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.07, "cosStep": 0.41 }
+        },
+        "seqlensKT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [255, 255] } },
+        "totalSequenceLengthT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [256] } },
+        "attentionBiasT": {
+          "dtype": "float32",
+          "shape": [2, 1, 256, 256],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.17, "scale": 0.5 }
+        }
+      },
+      "outputs": {
+        "outputT": { "dtype": "float32", "shape": [2, 256, 128], "tolerance": 0.00002 },
+        "presentKeyT": { "dtype": "float32", "shape": [2, 1, 256, 64], "tolerance": 0.000001 },
+        "presentValueT": { "dtype": "float32", "shape": [2, 1, 256, 64], "tolerance": 0.000001 }
+      }
     }
   ]
 }