sync c928d21e6cc1

Browse files

Files changed (4) hide show

build/webgpu/attn-materialized-sgmat-f32.wgsl.jinja +51 -38
build/webgpu/manifest.json +163 -11
build/webgpu/metadata.json +5 -5
build/webgpu/test.json +50 -0

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

@@ -1,4 +1,11 @@
 enable subgroups;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
@@ -6,10 +13,11 @@ diagnostic(off, chromium.subgroup_matrix_uniformity);
 {% set layout = source.layout | default("bsh") %}
 {% set headMajor = layout == "bhsd" %}
 {% macro q_index(row, d) %}(b * params.qSeq + {{ row }}) * HIDDEN + h * HEAD_DIM + {{ d }}{% endmacro %}
 {% macro kv_index(seq, d) %}(b * params.kvSeq + {{ seq }}) * KV_HIDDEN + h_kv * HEAD_DIM + {{ d }}{% endmacro %}
 {% set OUT_ROW_STRIDE = "HEAD_DIM" if headMajor else "HIDDEN" %}
-{% set KV_ROW_STRIDE = "HEAD_DIM" if headMajor else "KV_HIDDEN" %}
 {% set scorePhase = source.phase == "score" %}
 {% set FUSED_SOFTMAX = source.fusedSoftmax is defined and source.fusedSoftmax %}
 {% set PRIVATE_ROW_STATS = FUSED_SOFTMAX and (materializedSgmatPrivateRowStats is defined and materializedSgmatPrivateRowStats) %}
@@ -38,11 +46,11 @@ select(0.0, scores[{{ index }}], {{ guard[1] }})
 {% 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 %}
 {% 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 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) %}
@@ -80,8 +88,14 @@ const SUBGROUP_ROWS: u32 = {{ SUBGROUP_ROWS }}u;
 const WORKGROUP_THREADS: u32 = {{ WORKGROUP_THREADS }}u;
 {% endif %}
-var<workgroup> tile_A: array<f32, {{ TILE_M_VALUE * TILE_K_VALUE }}>;
-var<workgroup> tile_B: array<f32, {{ TILE_N_VALUE * TILE_K_VALUE }}>;
 {% 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.
@@ -120,7 +134,7 @@ fn exp_shift(value: f32, maxValue: f32) -> f32 {
 {% endif %}
-@compute @workgroup_size({{ WORKGROUP_THREADS }}, 1, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_index) li: u32{% if not DIRECT_OUTPUT_STORE %},
@@ -177,7 +191,7 @@ fn main(
 {% endif %}
 {% endif %}
   for (var k_base = 0u; k_base < inner; k_base += TILE_K) {
-{% if source.phase == "apply" and not FUSED_SOFTMAX %}
     // Full interior PV tiles can be loaded directly from storage. Query,
     // reduction, and output-dimension tails use the guarded shared path below.
     if (
@@ -190,16 +204,16 @@ fn main(
         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>>(
-            &scores, score_offset{{ row_block }}, false, 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>>(
-            &value, value_offset{{ col_block }}, false, {{ KV_ROW_STRIDE }}
           );
 {% endfor %}
 {% for row_block in range(ROW_BLOCKS) %}
@@ -227,12 +241,12 @@ fn main(
 {% if scorePhase %}
 {% if headDim % 32 == 0 %}
       tile_A[a_row * TILE_K + a_col + i] = select(
-        0.0,
         {{ q_tile_value(q_index("row", "k")) }},
         row < params.qSeq
       );
 {% else %}
-      var loaded = 0.0;
       if (row < params.qSeq && k < HEAD_DIM) {
         loaded = {{ q_tile_value(q_index("row", "k")) }};
       }
@@ -241,7 +255,7 @@ fn main(
 {% else %}
       let score_base = (b * HEADS + h) * params.qSeq * params.kvSeq;
       tile_A[a_row * TILE_K + a_col + i] =
-        {{ score_value("score_base + row * params.kvSeq + k", ["a_row", "row < params.qSeq && k < params.kvSeq"]) }};
 {% endif %}
     }
@@ -253,12 +267,12 @@ fn main(
 {% if scorePhase %}
 {% if headDim % 32 == 0 %}
       tile_B[b_row * TILE_K + b_col + i] = select(
-        0.0,
         key[{{ kv_index("col", "k") }}],
         col < params.kvSeq
       );
 {% else %}
-      var loaded = 0.0;
       if (col < params.kvSeq && k < HEAD_DIM) {
         loaded = key[{{ kv_index("col", "k") }}];
       }
@@ -266,7 +280,7 @@ fn main(
 {% endif %}
 {% else %}
       tile_B[b_row * TILE_K + b_col + i] = select(
-        0.0,
         value[{{ kv_index("k", "col") }}],
         k < params.kvSeq && col < HEAD_DIM
       );
@@ -281,19 +295,19 @@ fn main(
       let row = m_base + tile_row;
       let k = k_base + tile_k;
 {% if scorePhase %}
-      var loaded = 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 =
-        {{ score_value("score_base + row * params.kvSeq + k", ["tile_row", "row < params.qSeq && k < params.kvSeq"]) }};
 {% else %}
-      var loaded = 0.0;
       if (row < params.qSeq && k < params.kvSeq) {
         let score_base = (b * HEADS + h) * params.qSeq * params.kvSeq;
-        loaded = scores[score_base + row * params.kvSeq + k];
       }
 {% endif %}
       tile_A[idx] = loaded;
@@ -305,12 +319,12 @@ fn main(
       let col = n_base + tile_col;
       let k = k_base + tile_k;
 {% if scorePhase %}
-      var loaded = 0.0;
       if (col < params.kvSeq && k < HEAD_DIM) {
         loaded = key[{{ kv_index("col", "k") }}];
       }
 {% else %}
-      var loaded = 0.0;
       if (k < params.kvSeq && col < HEAD_DIM) {
         loaded = value[{{ kv_index("k", "col") }}];
       }
@@ -324,17 +338,17 @@ fn main(
 {% 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<f32, 8, 8> =
-        subgroupMatrixLoad<subgroup_matrix_left<f32, 8, 8>>(
-          &tile_A, matrix_a_offset{{ row_block }}, false, 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<f32, 8, 8> =
-        subgroupMatrixLoad<subgroup_matrix_right<f32, 8, 8>>(
-          &tile_B, matrix_b_offset{{ col_block }}, true, TILE_K
         );
 {% endfor %}
 {% for row_block in range(ROW_BLOCKS) %}
@@ -368,16 +382,16 @@ fn main(
     (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(
     &scores, output_offset{{ row_block }}{{ col_block }},
-    matC{{ row_block }}{{ col_block }}, false, 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(
     &output, output_offset{{ row_block }}{{ col_block }},
-    matC{{ row_block }}{{ col_block }}, false, {{ OUT_ROW_STRIDE }}
   );
 {% endif %}
 {% endfor %}
@@ -404,11 +418,10 @@ fn main(
   var stat_d{{ row_block }} = 0.0;
 {% endif %}
 {% for col_block in range(COL_BLOCKS) %}
-  subgroupMatrixStore(
-    &tile_A,
     (subgroup * {{ COL_BLOCKS }}u + {{ col_block }}u) * 64u,
     matC{{ row_block }}{{ col_block }},
-    false,
     8u
   );
 {% endfor %}
@@ -422,7 +435,7 @@ fn main(
       row < params.qSeq &&
       col < {% if scorePhase %}params.kvSeq{% else %}HEAD_DIM{% endif %}
     ) {
-      let result = tile_A[
         (subgroup * {{ COL_BLOCKS }}u + {{ col_block }}u) * 64u
           + row_in_block * 8u + col_in_block + pair
       ];
@@ -446,7 +459,7 @@ fn main(
 {% if hasBias %}
       // V bias row base: skip the packed Q and K blocks, then index this head.
 {% endif %}
-      output[{{ q_index("row", "col") }}] = result{% if hasBias %} + bias[2u * HIDDEN + h * HEAD_DIM + col]{% endif %};
 {% endif %}
     }
   }

+{% 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);
 {% set layout = source.layout | default("bsh") %}
 {% set headMajor = layout == "bhsd" %}
+{% set kvHeadMajor = (source.kvLayout | default(layout)) == "bhsd" %}
 {% macro q_index(row, d) %}(b * params.qSeq + {{ row }}) * HIDDEN + h * HEAD_DIM + {{ d }}{% endmacro %}
 {% macro kv_index(seq, d) %}(b * params.kvSeq + {{ seq }}) * KV_HIDDEN + h_kv * HEAD_DIM + {{ d }}{% 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 FUSED_SOFTMAX = source.fusedSoftmax is defined and source.fusedSoftmax %}
 {% set PRIVATE_ROW_STATS = FUSED_SOFTMAX and (materializedSgmatPrivateRowStats is defined and materializedSgmatPrivateRowStats) %}
 {% 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) %}
 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.
 {% 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 %},
 {% endif %}
 {% endif %}
   for (var k_base = 0u; k_base < inner; 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 (
         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) %}
 {% 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")) }};
       }
 {% 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 %}
     }
 {% 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") }}];
       }
 {% 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
       );
       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;
       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") }}];
       }
 {% 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) %}
     (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 %}
   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 %}
       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 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 %}
     }
   }

build/webgpu/manifest.json CHANGED Viewed

@@ -103,6 +103,9 @@
     "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",
     "headDim": "dim(shapes.query, 2) / attrs.num_heads if (ranks.query == 3 and attrs.num_heads > 0) else 0",
     "qkvDtypesOk": "tensorDtypes.key == tensorDtypes.query and tensorDtypes.value == tensorDtypes.query and tensorDtypes.output == tensorDtypes.query",
     "floatDtypeOk": "(tensorDtypes.query == \"float32\" or tensorDtypes.query == \"float16\") and f16Ok(tensorDtypes.query)",
@@ -171,7 +174,7 @@
     "materializedSgmatDirectScoreStore": "dim(shapes.query, 1) % materializedSgmatQueryTile == 0 and dim(shapes.key, 1) % materializedSgmatKeyTile == 0",
     "materializedSgmatDirectApplyStore": "dim(shapes.query, 1) % materializedSgmatQueryTile == 0 and headDim % materializedSgmatKeyTile == 0",
     "materializedSgmatRuntimeDirectStore": "dim(shapes.query, 1) >= 2 * materializedSgmatQueryTile and dim(shapes.key, 1) >= 2 * materializedSgmatKeyTile",
-    "materializedSgmatCoreOk": "qkvContractOk and tensorDtypes.query == \"float32\" and attrs.unidirectional == 0 and headDim >= 32 and headDim <= 256 and dim(shapes.query, 1) >= 512 and dim(shapes.key, 1) >= 512 and subgroupsWave32 and device.features.has(\"chromium-experimental-subgroup-matrix\") and materializedSgmatBuffersFit and materializedSgmatResourcesFit and materializedSgmatDispatchFits",
     "materializedCachedSoftmaxWg": "min(tunables.MATERIALIZED_CACHED_SOFTMAX_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "materializedCachedSoftmaxVecsPerLane": "ceilDiv(ceilDiv(dim(shapes.key, 1), 4), max(1, materializedCachedSoftmaxWg))",
     "materializedCachedSoftmaxStorageBytes": "materializedCachedSoftmaxWg * 8 + 8",
@@ -196,7 +199,9 @@
     "smallSeqWorkgroupSize": "max(32, pow2ceil(dim(shapes.query, 1)))",
     "smallSeqSharedBytes": "dim(shapes.key, 1) * headDim * 8",
     "smallSeqResourcesFit": "smallSeqPrivateFloats <= tunables.SMALL_SEQ_MAX_PRIVATE_FLOATS and smallSeqWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup and smallSeqWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX and smallSeqSharedBytes <= device.limits.maxComputeWorkgroupStorageSize",
-    "smallSeqDispatchFits": "attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension and dim(shapes.query, 0) <= device.limits.maxComputeWorkgroupsPerDimension"
   },
   "constants": {
     "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\""
@@ -1193,6 +1198,82 @@
           ]
         }
       }
     ]
   },
   "variants": [
@@ -1343,7 +1424,6 @@
       "priority": 52,
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
-        "subgroupMinSize": 32,
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "when": ["materializedSgmatFusedOk", "not present.biasT", "materializedFusedSoftmaxWorthIt"],
@@ -1416,7 +1496,6 @@
       "priority": 52,
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
-        "subgroupMinSize": 32,
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "when": ["materializedSgmatFusedOk", "biasOk", "materializedFusedSoftmaxWorthIt"],
@@ -1483,6 +1562,81 @@
         }
       ]
     },
     {
       "id": "qkv_no_bias_flash_cluster_lpq4_nosg",
       "description": "Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.",
@@ -1626,7 +1780,7 @@
         "headDim": "dim(shapes.query, 2) / attrs.num_heads",
         "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4",
         "TILE_Q": 16,
-        "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 128) else 8",
         "LPQ": 4
       },
       "passes": [
@@ -1664,7 +1818,7 @@
         "headDim": "dim(shapes.query, 2) / attrs.num_heads",
         "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4",
         "TILE_Q": 16,
-        "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 128) else 8",
         "LPQ": 8
       },
       "passes": [
@@ -1702,7 +1856,7 @@
         "headDim": "dim(shapes.query, 2) / attrs.num_heads",
         "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4",
         "TILE_Q": 16,
-        "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 128) else 8",
         "qHidden": "dim(shapes.query, 2)",
         "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"",
         "LPQ": 8
@@ -1745,7 +1899,7 @@
         "headDim": "dim(shapes.query, 2) / attrs.num_heads",
         "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4",
         "TILE_Q": 16,
-        "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 128) else 8",
         "LPQ": 8
       },
       "passes": [
@@ -1786,7 +1940,7 @@
         "headDim": "dim(shapes.query, 2) / attrs.num_heads",
         "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4",
         "TILE_Q": 16,
-        "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 128) else 8",
         "qHidden": "dim(shapes.query, 2)",
         "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"",
         "LPQ": 8
@@ -2792,7 +2946,6 @@
       "priority": 52,
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
-        "subgroupMinSize": 32,
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "when": ["materializedSgmatOk", "not present.biasT", "not materializedFusedSoftmaxWorthIt"],
@@ -2861,7 +3014,6 @@
       "priority": 52,
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
-        "subgroupMinSize": 32,
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "when": ["materializedSgmatOk", "biasOk", "not materializedFusedSoftmaxWorthIt"],

     "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",
     "headDim": "dim(shapes.query, 2) / attrs.num_heads if (ranks.query == 3 and attrs.num_heads > 0) else 0",
     "qkvDtypesOk": "tensorDtypes.key == tensorDtypes.query and tensorDtypes.value == tensorDtypes.query and tensorDtypes.output == tensorDtypes.query",
     "floatDtypeOk": "(tensorDtypes.query == \"float32\" or tensorDtypes.query == \"float16\") and f16Ok(tensorDtypes.query)",
     "materializedSgmatDirectScoreStore": "dim(shapes.query, 1) % materializedSgmatQueryTile == 0 and dim(shapes.key, 1) % materializedSgmatKeyTile == 0",
     "materializedSgmatDirectApplyStore": "dim(shapes.query, 1) % materializedSgmatQueryTile == 0 and headDim % materializedSgmatKeyTile == 0",
     "materializedSgmatRuntimeDirectStore": "dim(shapes.query, 1) >= 2 * materializedSgmatQueryTile and dim(shapes.key, 1) >= 2 * materializedSgmatKeyTile",
+    "materializedSgmatCoreOk": "qkvContractOk and tensorDtypes.query == \"float32\" and attrs.unidirectional == 0 and headDim >= 32 and headDim <= 256 and dim(shapes.query, 1) >= 512 and dim(shapes.key, 1) >= 512 and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and materializedSgmatBuffersFit and materializedSgmatResourcesFit and materializedSgmatDispatchFits",
     "materializedCachedSoftmaxWg": "min(tunables.MATERIALIZED_CACHED_SOFTMAX_WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
     "materializedCachedSoftmaxVecsPerLane": "ceilDiv(ceilDiv(dim(shapes.key, 1), 4), max(1, materializedCachedSoftmaxWg))",
     "materializedCachedSoftmaxStorageBytes": "materializedCachedSoftmaxWg * 8 + 8",
     "smallSeqWorkgroupSize": "max(32, pow2ceil(dim(shapes.query, 1)))",
     "smallSeqSharedBytes": "dim(shapes.key, 1) * headDim * 8",
     "smallSeqResourcesFit": "smallSeqPrivateFloats <= tunables.SMALL_SEQ_MAX_PRIVATE_FLOATS and smallSeqWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup and smallSeqWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX and smallSeqSharedBytes <= device.limits.maxComputeWorkgroupStorageSize",
+    "smallSeqDispatchFits": "attrs.num_heads <= device.limits.maxComputeWorkgroupsPerDimension and dim(shapes.query, 0) <= device.limits.maxComputeWorkgroupsPerDimension",
+    "materializedSgmatCoreF16Ok": "qkvContractOk and tensorDtypes.query == \"float16\" and tensorDtypes.key == \"float16\" and tensorDtypes.value == \"float16\" and device.features.has(\"shader-f16\") and attrs.unidirectional == 0 and headDim >= 32 and headDim <= 256 and dim(shapes.query, 1) >= 512 and dim(shapes.key, 1) >= 512 and device.features.has(\"subgroups\") and wave32Effective and device.features.has(\"chromium-experimental-subgroup-matrix\") and materializedSgmatBuffersFit and materializedSgmatResourcesFit and materializedSgmatDispatchFits",
+    "materializedSgmatFusedF16Ok": "materializedSgmatCoreF16Ok"
   },
   "constants": {
     "attentionScaleExpression": "\"select(inverseSqrt(f32(HEAD_DIM)), params.scale, params.scale != 0.0)\""
           ]
         }
       }
+    ],
+    "materializedScoreStatsDyn": [
+      {
+        "name": "query",
+        "arg": "queryT",
+        "semantic": "query",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$inputElement"
+      },
+      {
+        "name": "key",
+        "arg": "keyT",
+        "semantic": "key",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$inputElement"
+      },
+      { "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" }
+          ]
+        }
+      }
+    ],
+    "materializedApplyFusedDyn": [
+      {
+        "name": "scores",
+        "semantic": "materializedScores",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      {
+        "name": "value",
+        "arg": "valueT",
+        "semantic": "value",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$inputElement"
+      },
+      {
+        "name": "output",
+        "arg": "outputT",
+        "semantic": "output",
+        "buffer": { "type": "storage" },
+        "elementType": "$outputElement"
+      },
+      {
+        "name": "rowStats",
+        "semantic": "materializedRowStats",
+        "buffer": { "type": "read-only-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)" }
+          ]
+        }
+      }
     ]
   },
   "variants": [
       "priority": 52,
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "when": ["materializedSgmatFusedOk", "not present.biasT", "materializedFusedSoftmaxWorthIt"],
       "priority": 52,
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "when": ["materializedSgmatFusedOk", "biasOk", "materializedFusedSoftmaxWorthIt"],
         }
       ]
     },
+    {
+      "id": "qkv_no_bias_materialized_sgmat_fused_f16",
+      "description": "Materializes FP32 scores for subgroup-matrix score and value products, and folds the softmax into the apply pass using row statistics the score pass emits. Selected over the separate-softmax route when the score matrix is large enough that reading it back costs more than the apply-side staging fusing forces.",
+      "priority": 52,
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "when": ["materializedSgmatFusedF16Ok", "not present.biasT", "materializedFusedSoftmaxWorthIt"],
+      "constants": {
+        "qNumHeads": "attrs.num_heads",
+        "headDim": "headDim",
+        "qHidden": "dim(shapes.query, 2)",
+        "materializedSgmatQueryTile": "materializedSgmatQueryTile",
+        "materializedSgmatKeyTile": "materializedSgmatKeyTile",
+        "materializedSgmatInnerTile": "materializedSgmatInnerTile",
+        "hasBias": false,
+        "useSubgroups": true,
+        "materializedRowStatsWg": "materializedRowStatsWg",
+        "statSlots": "materializedSgmatStatSlots",
+        "statQuerySeq": "dim(shapes.query, 1)",
+        "operandF16": true,
+        "inputElement": "\"f16\"",
+        "outputElement": "\"f16\""
+      },
+      "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": "[materializedRowStatsElements]" },
+        { "id": "materializedScorePartials", "dtype": "float32", "shape": "[materializedScorePartialElements]" }
+      ],
+      "passes": [
+        {
+          "id": "scores",
+          "name": "MultiHeadAttention.MaterializedScoresSgmatF16",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": { "phase": "\"score\"", "emitRowStats": true }
+          },
+          "bindings": "materializedScoreStatsDyn",
+          "dispatch": {
+            "x": "ceilDiv(dim(shapes.key, 1), materializedSgmatKeyTile)",
+            "y": "ceilDiv(dim(shapes.query, 1), materializedSgmatQueryTile)",
+            "z": "dim(shapes.query, 0) * attrs.num_heads"
+          }
+        },
+        {
+          "id": "rowstats",
+          "name": "MultiHeadAttention.MaterializedRowStatsCombine",
+          "source": { "shader": "attn-materialized-rowstats-combine-f32.wgsl.jinja" },
+          "bindings": "materializedRowStatsCombine",
+          "dispatch": {
+            "gridStride": "dim(shapes.query, 0) * attrs.num_heads * dim(shapes.query, 1)",
+            "workgroupSize": "materializedRowStatsWg"
+          }
+        },
+        {
+          "id": "apply",
+          "name": "MultiHeadAttention.MaterializedApplySgmatF16",
+          "source": {
+            "shader": "attn-materialized-sgmat-f32.wgsl.jinja",
+            "inputs": { "phase": "\"apply\"", "fusedSoftmax": true }
+          },
+          "bindings": "materializedApplyFusedDyn",
+          "dispatch": {
+            "x": "ceilDiv(headDim, materializedSgmatKeyTile)",
+            "y": "ceilDiv(dim(shapes.query, 1), materializedSgmatQueryTile)",
+            "z": "dim(shapes.query, 0) * attrs.num_heads"
+          }
+        }
+      ]
+    },
     {
       "id": "qkv_no_bias_flash_cluster_lpq4_nosg",
       "description": "Keeps query/key tiles and online-softmax state inside one workgroup instead of materializing all scores. The family supplies subgroup and portable forms for compatible prefill shapes.",
         "headDim": "dim(shapes.query, 2) / attrs.num_heads",
         "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4",
         "TILE_Q": 16,
+        "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 64) else 8",
         "LPQ": 4
       },
       "passes": [
         "headDim": "dim(shapes.query, 2) / attrs.num_heads",
         "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4",
         "TILE_Q": 16,
+        "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 64) else 8",
         "LPQ": 8
       },
       "passes": [
         "headDim": "dim(shapes.query, 2) / attrs.num_heads",
         "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4",
         "TILE_Q": 16,
+        "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 64) else 8",
         "qHidden": "dim(shapes.query, 2)",
         "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"",
         "LPQ": 8
         "headDim": "dim(shapes.query, 2) / attrs.num_heads",
         "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4",
         "TILE_Q": 16,
+        "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 64) else 8",
         "LPQ": 8
       },
       "passes": [
         "headDim": "dim(shapes.query, 2) / attrs.num_heads",
         "headDimV4": "(dim(shapes.query, 2) / attrs.num_heads) / 4",
         "TILE_Q": 16,
+        "TILE_K": "32 if (tensorDtypes.query != \"float16\" and (dim(shapes.query, 2) / attrs.num_heads) <= 64) else 8",
         "qHidden": "dim(shapes.query, 2)",
         "inputScalar": "\"f16\" if tensorDtypes.query == \"float16\" else \"f32\"",
         "LPQ": 8
       "priority": 52,
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "when": ["materializedSgmatOk", "not present.biasT", "not materializedFusedSoftmaxWorthIt"],
       "priority": 52,
       "requires": {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "when": ["materializedSgmatOk", "biasOk", "not materializedFusedSoftmaxWorthIt"],

build/webgpu/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "com.microsoft.MultiHeadAttention",
-  "id": "_com_microsoft_multiheadattention_webgpu_9e2cafc",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
@@ -16,16 +16,16 @@
       "attn-materialized-apply-f32.wgsl.jinja": "qbrsh9S14QYuAM+xF49dgcooO6f19yNg9lC8eoBeWFQ=",
       "attn-materialized-rowstats-combine-f32.wgsl.jinja": "rdpdnh6y3m0tBwFMVCc903j2cvGsnGTcT8EAWRr5TRk=",
       "attn-materialized-score-f32.wgsl.jinja": "EAhU/wHSgmJZ/YHCnwG/takpgfMAtGqKXo3Cf8TfvxU=",
-      "attn-materialized-sgmat-f32.wgsl.jinja": "EpEZp/6PjkWmTzB2njeciGho5JQbCMadXVxBhUqfVyE=",
       "attn-materialized-softmax-f32.wgsl.jinja": "7z288syeGhUUyBpaZzVVqdtFVMQALuN88sNfytvClWw=",
       "attn-online-scalar.wgsl.jinja": "CW0ru5yg8YFpIv31XnnD5IuKBXppsKc+2BU7u2uAQv4=",
       "attn-small-head-parallel.wgsl.jinja": "ejotx9NjdqElO3EIQWj8E3hJk2iW8Exq1vUjxK4FUk0=",
       "bench.json": "VV4kVBFJ112r3U2RGQlQd1Z5eAzgm+y30+oGG6TpK3g=",
-      "manifest.json": "KlaFGLtELDJFyP/xCymid5FGW+u1XJRXukAwblz62b4=",
       "mha-small-seq.wgsl.jinja": "F825nrB6RLu55ziDnDI7BTJ3JSdT1GwsuMYmtrHQSic=",
-      "test.json": "8uMdVLj7AWq1cMy7/Iyl35jJPDjE5z/VztKieOgg20U="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
   "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.MultiHeadAttention" }
 }

 {
   "name": "com.microsoft.MultiHeadAttention",
+  "id": "_com_microsoft_multiheadattention_webgpu_03bfaa6",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
       "attn-materialized-apply-f32.wgsl.jinja": "qbrsh9S14QYuAM+xF49dgcooO6f19yNg9lC8eoBeWFQ=",
       "attn-materialized-rowstats-combine-f32.wgsl.jinja": "rdpdnh6y3m0tBwFMVCc903j2cvGsnGTcT8EAWRr5TRk=",
       "attn-materialized-score-f32.wgsl.jinja": "EAhU/wHSgmJZ/YHCnwG/takpgfMAtGqKXo3Cf8TfvxU=",
+      "attn-materialized-sgmat-f32.wgsl.jinja": "7fJFDK+kWEhN4rqFs3fE9uPkZdHh5C9xl7PXe9vZze4=",
       "attn-materialized-softmax-f32.wgsl.jinja": "7z288syeGhUUyBpaZzVVqdtFVMQALuN88sNfytvClWw=",
       "attn-online-scalar.wgsl.jinja": "CW0ru5yg8YFpIv31XnnD5IuKBXppsKc+2BU7u2uAQv4=",
       "attn-small-head-parallel.wgsl.jinja": "ejotx9NjdqElO3EIQWj8E3hJk2iW8Exq1vUjxK4FUk0=",
       "bench.json": "VV4kVBFJ112r3U2RGQlQd1Z5eAzgm+y30+oGG6TpK3g=",
+      "manifest.json": "kaCZGh3/PB8EmCgBc5QCc0aVGSrjf55xB6OM4s0yhgg=",
       "mha-small-seq.wgsl.jinja": "F825nrB6RLu55ziDnDI7BTJ3JSdT1GwsuMYmtrHQSic=",
+      "test.json": "6PCZYazz1hXA9+nbVJmd0/btWsIkkp29/xYqWZ4OSrs="
     }
   },
+  "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
   "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.MultiHeadAttention" }
 }

build/webgpu/test.json CHANGED Viewed

@@ -1051,6 +1051,56 @@
         "outputT": { "dtype": "float32", "shape": [1, 512, 128], "tolerance": 0.0004, "relTolerance": 0.0004 }
       }
     },
     {
       "name": "materialized_sgmat_q512_kv512_h1_d128_bias",
       "provenance": {

         "outputT": { "dtype": "float32", "shape": [1, 512, 128], "tolerance": 0.0004, "relTolerance": 0.0004 }
       }
     },
+    {
+      "name": "materialized_sgmat_fused_f16_q1024_h4_d64_no_bias",
+      "provenance": {
+        "notes": "Pins the f16-operand fused materialized subgroup-matrix route just over its 16MB score-bytes admission floor (4 heads x 1024 x 1024 x 4B): f16 QK operands with f32 accumulation, raw f32 scores plus row stats, and the f16-probability PV apply with the dedicated f32 store scratch."
+      },
+      "attrs": { "num_heads": 4, "scale": 0.125 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 1024, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.013, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 1024, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.019, "cosStep": 0.037 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 1024, 256],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.029, "cosStep": 0.017, "offset": 1.0 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 1024, 256], "tolerance": 0.02, "relTolerance": 0.02 } }
+    },
+    {
+      "name": "materialized_sgmat_fused_f16_q1024_h4_d80_no_bias",
+      "provenance": {
+        "notes": "The d=80 twin of the f16 fused materialized fixture: a head dim that is not a multiple of 32 takes the guarded staging forms instead of the unguarded selects, on the same 16MB score-bytes admission floor."
+      },
+      "attrs": { "num_heads": 4, "scale": 0.125 },
+      "inputs": {
+        "queryT": {
+          "dtype": "float16",
+          "shape": [1, 1024, 320],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.013, "cosStep": 0.029 }
+        },
+        "keyT": {
+          "dtype": "float16",
+          "shape": [1, 1024, 320],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.019, "cosStep": 0.037 }
+        },
+        "valueT": {
+          "dtype": "float16",
+          "shape": [1, 1024, 320],
+          "data": { "kind": "fillFloat32", "scale": 0.08, "sinStep": 0.029, "cosStep": 0.017, "offset": 1.0 }
+        }
+      },
+      "outputs": { "outputT": { "dtype": "float16", "shape": [1, 1024, 320], "tolerance": 0.02, "relTolerance": 0.02 } }
+    },
     {
       "name": "materialized_sgmat_q512_kv512_h1_d128_bias",
       "provenance": {