sync 6fdf6301e2bb

Browse files

Files changed (8) hide show

README.md +17 -4
build/webgpu/bench.json +0 -0
build/webgpu/{gemm-fast-gelu.wgsl.jinja → gemm-epilogue-tiled-reg.wgsl.jinja} +63 -20
build/webgpu/gemm-subgroup-matrix.wgsl.jinja +24 -65
build/webgpu/manifest.json +453 -53
build/webgpu/matmul-vector-matrix-vec4.wgsl.jinja +70 -0
build/webgpu/metadata.json +21 -10
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

@@ -40,6 +40,16 @@ See the [ONNX Runtime `GemmFastGelu` contrib-operator spec](https://github.com/m
 One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
 - `sgmat_direct_bias` — Uses supported subgroup matrices with direct activation loads and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
 - `sgmat_direct` — Uses supported subgroup matrices with direct activation loads and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
 - `sgmat_direct_bias_f16` — Uses supported subgroup matrices with direct activation loads and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
@@ -48,6 +58,8 @@ One implementation is selected per call from the device capabilities, the reques
 - `sgmat` — Uses supported subgroup matrices with staged operands. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
 - `sgmat_bias_f16` — Uses supported subgroup matrices with staged operands. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
 - `sgmat_f16` — Uses supported subgroup matrices with staged operands. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
 ## Device requirements
@@ -58,14 +70,15 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
-- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
-- [`gemm-fast-gelu.wgsl.jinja`](build/webgpu/gemm-fast-gelu.wgsl.jinja)
 - [`gemm-subgroup-matrix.wgsl.jinja`](build/webgpu/gemm-subgroup-matrix.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
@@ -79,5 +92,5 @@ Replace each `*Data` placeholder with a typed array containing the corresponding
 import { getKernel } from "@huggingface/kernels";
 const kernel = await getKernel("webgpu-kernels/com.microsoft.GemmFastGelu", { version: 1 });
-const { Y } = await kernel({ X: { data: XData, shape: [5, 6] }, W: { data: WData, shape: [6, 4] } });
 ```

 One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `vector_matrix_bias` — Partitions a single-row vector-matrix product across device-bounded workgroup slices. Optional bias and FastGelu are fused into the f32 reduction before the single output cast; no subgroup feature is required.
+- `vector_matrix` — Partitions a single-row vector-matrix product across device-bounded workgroup slices. Optional bias and FastGelu are fused into the f32 reduction before the single output cast; no subgroup feature is required.
+- `sgmat_medium_bias` — Uses supported subgroup matrices with 32x32 direct activation tiles and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `sgmat_medium` — Uses supported subgroup matrices with 32x32 direct activation tiles and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `sgmat_medium_bias_f16` — Uses supported subgroup matrices with 32x32 direct activation tiles and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `sgmat_medium_f16` — Uses supported subgroup matrices with 32x32 direct activation tiles and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `sgmat_small_bias` — Uses supported subgroup matrices with 32x16 direct activation tiles and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `sgmat_small` — Uses supported subgroup matrices with 32x16 direct activation tiles and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `sgmat_small_bias_f16` — Uses supported subgroup matrices with 32x16 direct activation tiles and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `sgmat_small_f16` — Uses supported subgroup matrices with 32x16 direct activation tiles and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
 - `sgmat_direct_bias` — Uses supported subgroup matrices with direct activation loads and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
 - `sgmat_direct` — Uses supported subgroup matrices with direct activation loads and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
 - `sgmat_direct_bias_f16` — Uses supported subgroup matrices with direct activation loads and double-buffered weights. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
 - `sgmat` — Uses supported subgroup matrices with staged operands. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
 - `sgmat_bias_f16` — Uses supported subgroup matrices with staged operands. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
 - `sgmat_f16` — Uses supported subgroup matrices with staged operands. Bias and FastGelu run on the f32 accumulator before the single output cast. Resource, alignment and matrix-type guards retain the portable path on other devices.
+- `tiled_bias` — Portable register tiles use the output-grid budget and a device-specific f32 tile choice to improve occupancy, with device-bounded shared padding and compensated f32 tile sums before the fused activation.
+- `tiled` — Portable register tiles use the output-grid budget and a device-specific f32 tile choice to improve occupancy, with device-bounded shared padding and compensated f32 tile sums before the fused activation.
 ## Device requirements
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
+- [`gemm-epilogue-tiled-reg.wgsl.jinja`](build/webgpu/gemm-epilogue-tiled-reg.wgsl.jinja)
 - [`gemm-subgroup-matrix.wgsl.jinja`](build/webgpu/gemm-subgroup-matrix.wgsl.jinja)
+- [`matmul-vector-matrix-vec4.wgsl.jinja`](build/webgpu/matmul-vector-matrix-vec4.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 import { getKernel } from "@huggingface/kernels";
 const kernel = await getKernel("webgpu-kernels/com.microsoft.GemmFastGelu", { version: 1 });
+const { Y } = await kernel({ X: { data: XData, shape: [1, 1] }, W: { data: WData, shape: [1, 4] } });
 ```

build/webgpu/bench.json CHANGED Viewed

The diff for this file is too large to render. See raw diff

build/webgpu/{gemm-fast-gelu.wgsl.jinja → gemm-epilogue-tiled-reg.wgsl.jinja} RENAMED Viewed

@@ -1,47 +1,56 @@
-{% if usesF16 is defined and usesF16 %}enable f16;
-{% endif %}{% set OUT = "f16" if (usesF16 is defined and usesF16) else "f32" %}
 {{ env.wgsl.resourceDeclarations }}
 // Shape-driven, register-blocked f32 GEMM for all four transpose combinations.
-// Runtime M/N/K share one compiled pipeline. A 4x4 thread micro-tile forms a
-// 64x64 output tile with BK=16; transpose modes change only load strides and
-// keep the stored-contiguous axis coalesced. M/N/K tails are bounds-checked.
 // Both tiles are indexed by their own output axis and group four K values per
 // vector word, so the micro-tile accumulates through dot() and one step reads
 // TM + TN words instead of 4 * (TM + TN) scalars.
-{% set gemmEpi = gemmEpilogue if gemmEpilogue is defined else "none" %}
 fn tanh_safe(x: f32) -> f32 {
   if (x > 10.0) { return 1.0; }
   if (x < -10.0) { return -1.0; }
   return tanh(x);
 }
 fn gelu_tanh(v: f32) -> f32 {
   return 0.5 * v * (1.0 + tanh_safe(0.7978845608028654 * (v + 0.044715 * v * v * v)));
 }
-{% set epiOpen = "gelu_tanh(" if gemmEpi == "fastgelu" else ("fused_act(" if gemmEpi == "activation" else "") %}
-{% set epiClose = ")" if gemmEpi != "none" else "" %}
-const BK: u32 = 16u;
-const BM: u32 = 64u;
-const BN: u32 = 64u;
-const TM: u32 = 4u;   // per-thread micro-tile rows
-const TN: u32 = 4u;   // per-thread micro-tile cols
 // Stage tiles at the operands' storage width. For f16 inputs this is
 // bit-identical to widening on shared load, keeps f32 FMA registers, and halves
 // workgroup storage.
 {% set tileT = OUT %}
 const K_VECS: u32 = BK / 4u;
-var<workgroup> tileA: array<array<vec4<{{ tileT }}>, K_VECS>, BM>;   // [64][4]  logical A'[m][k/4]
-var<workgroup> tileB: array<array<vec4<{{ tileT }}>, K_VECS>, BN>;   // [64][4]  logical B'[n][k/4]
-@compute @workgroup_size(16, 16, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>,
   @builtin(num_workgroups) nwg: vec3<u32>
 ) {
-  let li = lid.y * 16u + lid.x;
   let M = params.M;
   let N = params.N;
   let K = params.K;
@@ -55,46 +64,70 @@ fn main(
       var acc: array<f32, TM * TN>;
       for (var i: u32 = 0u; i < TM * TN; i = i + 1u) { acc[i] = 0.0; }
       let numTiles = (K + BK - 1u) / BK;
       for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
         let kBase = kt * BK;
         // Cooperative load: one vector word per lane per pass. The lane grid is
         // laid out so the axis the operand stores contiguously stays coalesced.
-        for (var idx: u32 = li; idx < BM * K_VECS; idx = idx + 256u) {
           // A stored [M, K]: logical A[m,k] = a[m*K + k]. Coalesce on k (contiguous).
           let ar = idx / K_VECS;
           let ac4 = idx % K_VECS;
           let am = mBase + ar;
           let ak = kBase + ac4 * 4u;
           var aWord = vec4<{{ tileT }}>({{ tileT }}(0.0));
           if (am < M) {
 {% for component in range(4) %}
             if (ak + {{ component }}u < K) {
               aWord[{{ component }}u] = {{ tileT }}(a[am * K + ak + {{ component }}u]);
             }
 {% endfor %}
           }
           tileA[ar][ac4] = aWord;
         }
-        for (var idx: u32 = li; idx < BN * K_VECS; idx = idx + 256u) {
           // B stored [K, N]: logical B[k,n] = b[k*N + n]. Coalesce on n (contiguous).
           let bc = idx % BN;
           let br4 = idx / BN;
           let bn = nBase + bc;
           let bk = kBase + br4 * 4u;
           var bWord = vec4<{{ tileT }}>({{ tileT }}(0.0));
           if (bn < N) {
 {% for component in range(4) %}
             if (bk + {{ component }}u < K) {
               bWord[{{ component }}u] = {{ tileT }}(b[(bk + {{ component }}u) * N + bn]);
             }
 {% endfor %}
           }
           tileB[bc][br4] = bWord;
         }
         workgroupBarrier();
 {% set regT = "f32" %}{% set regCast = 1 if tileT != "f32" else 0 %}
 {% filter indent(8, true) %}
 let aRow = lid.y * TM;
 let bCol = lid.x * TN;
 for (var kv: u32 = 0u; kv < BK / 4u; kv = kv + 1u) {
@@ -104,11 +137,21 @@ for (var kv: u32 = 0u; kv < BK / 4u; kv = kv + 1u) {
   for (var j: u32 = 0u; j < TN; j = j + 1u) { bv[j] = {% if regCast %}vec4<{{ regT }}>(tileB[bCol + j][kv]){% else %}tileB[bCol + j][kv]{% endif %}; }
   for (var i: u32 = 0u; i < TM; i = i + 1u) {
     for (var j: u32 = 0u; j < TN; j = j + 1u) {
-      acc[i * TN + j] = acc[i * TN + j] + dot(av[i], bv[j]);
     }
   }
 }
 {% endfilter %}
         workgroupBarrier();
       }

+{% set OUT = "f16" if (usesF16 is defined and usesF16) else "f32" %}
 {{ env.wgsl.resourceDeclarations }}
 // Shape-driven, register-blocked f32 GEMM for all four transpose combinations.
+// Runtime M/N/K share one compiled pipeline for each tile configuration.
+// Register micro-tiles reuse staged K slices; transpose modes change load strides and
+// keep the stored-contiguous axis coalesced. M/N/K tails are bounds-checked
+// unless an authored dimension proves that its tile has no tail.
 // Both tiles are indexed by their own output axis and group four K values per
 // vector word, so the micro-tile accumulates through dot() and one step reads
 // TM + TN words instead of 4 * (TM + TN) scalars.
 fn tanh_safe(x: f32) -> f32 {
+  // tanh rounds to its saturated value for these tails in f32. Return that
+  // value directly, including for infinite input, before invoking the builtin.
   if (x > 10.0) { return 1.0; }
   if (x < -10.0) { return -1.0; }
+  // For tiny |x|, return x directly to preserve its sign and magnitude without
+  // relying on backend-specific builtin behavior near zero.
+  if (x > -1.0e-4 && x < 1.0e-4) { return x; }
   return tanh(x);
 }
 fn gelu_tanh(v: f32) -> f32 {
   return 0.5 * v * (1.0 + tanh_safe(0.7978845608028654 * (v + 0.044715 * v * v * v)));
 }
+{% set epiOpen = "gelu_tanh(" %}
+{% set epiClose = ")" %}
+{% set gemmWgX = ((gemmTileCols / gemmMicroCols) | int) %}
+{% set gemmWgY = ((gemmTileRows / gemmMicroRows) | int) %}
+{% set gemmRowTail = gemmStaticM is not defined or gemmStaticM % gemmTileRows != 0 %}
+{% set gemmColTail = gemmStaticN is not defined or gemmStaticN % gemmTileCols != 0 %}
+{% set gemmKTileTail = gemmStaticK is not defined or gemmStaticK % gemmTileK != 0 %}
+const BK: u32 = {{ gemmTileK }}u;
+const BM: u32 = {{ gemmTileRows }}u;
+const BN: u32 = {{ gemmTileCols }}u;
+const TM: u32 = {{ gemmMicroRows }}u;   // per-thread micro-tile rows
+const TN: u32 = {{ gemmMicroCols }}u;   // per-thread micro-tile cols
 // Stage tiles at the operands' storage width. For f16 inputs this is
 // bit-identical to widening on shared load, keeps f32 FMA registers, and halves
 // workgroup storage.
 {% set tileT = OUT %}
 const K_VECS: u32 = BK / 4u;
+var<workgroup> tileA: array<array<vec4<{{ tileT }}>, K_VECS{% if gemmTilePadding > 0 %} + {{ gemmTilePadding }}u{% endif %}>, BM>;   // logical A'[m][k/4]
+var<workgroup> tileB: array<array<vec4<{{ tileT }}>, K_VECS{% if gemmTilePadding > 0 %} + {{ gemmTilePadding }}u{% endif %}>, BN>;   // logical B'[n][k/4]
+@compute @workgroup_size({{ gemmWgX }}, {{ gemmWgY }}, 1)
 fn main(
   @builtin(workgroup_id) wg: vec3<u32>,
   @builtin(local_invocation_id) lid: vec3<u32>,
   @builtin(num_workgroups) nwg: vec3<u32>
 ) {
+  let li = lid.y * {{ gemmWgX }}u + lid.x;
   let M = params.M;
   let N = params.N;
   let K = params.K;
       var acc: array<f32, TM * TN>;
       for (var i: u32 = 0u; i < TM * TN; i = i + 1u) { acc[i] = 0.0; }
+{% if gemmCompensatedTiles %}
+      var accError: array<f32, TM * TN>;
+{% endif %}
       let numTiles = (K + BK - 1u) / BK;
       for (var kt: u32 = 0u; kt < numTiles; kt = kt + 1u) {
         let kBase = kt * BK;
         // Cooperative load: one vector word per lane per pass. The lane grid is
         // laid out so the axis the operand stores contiguously stays coalesced.
+        for (var idx: u32 = li; idx < BM * K_VECS; idx = idx + {{ gemmWgX * gemmWgY }}u) {
           // A stored [M, K]: logical A[m,k] = a[m*K + k]. Coalesce on k (contiguous).
           let ar = idx / K_VECS;
           let ac4 = idx % K_VECS;
           let am = mBase + ar;
           let ak = kBase + ac4 * 4u;
           var aWord = vec4<{{ tileT }}>({{ tileT }}(0.0));
+{% if gemmRowTail %}
           if (am < M) {
+{% endif %}
 {% for component in range(4) %}
+{% if gemmKTileTail %}
             if (ak + {{ component }}u < K) {
+{% endif %}
               aWord[{{ component }}u] = {{ tileT }}(a[am * K + ak + {{ component }}u]);
+{% if gemmKTileTail %}
             }
+{% endif %}
 {% endfor %}
+{% if gemmRowTail %}
           }
+{% endif %}
           tileA[ar][ac4] = aWord;
         }
+        for (var idx: u32 = li; idx < BN * K_VECS; idx = idx + {{ gemmWgX * gemmWgY }}u) {
           // B stored [K, N]: logical B[k,n] = b[k*N + n]. Coalesce on n (contiguous).
           let bc = idx % BN;
           let br4 = idx / BN;
           let bn = nBase + bc;
           let bk = kBase + br4 * 4u;
           var bWord = vec4<{{ tileT }}>({{ tileT }}(0.0));
+{% if gemmColTail %}
           if (bn < N) {
+{% endif %}
 {% for component in range(4) %}
+{% if gemmKTileTail %}
             if (bk + {{ component }}u < K) {
+{% endif %}
               bWord[{{ component }}u] = {{ tileT }}(b[(bk + {{ component }}u) * N + bn]);
+{% if gemmKTileTail %}
             }
+{% endif %}
 {% endfor %}
+{% if gemmColTail %}
           }
+{% endif %}
           tileB[bc][br4] = bWord;
         }
         workgroupBarrier();
+{% if gemmCompensatedTiles %}
+        var tileAcc: array<f32, TM * TN>;
+{% endif %}
+{% set regAccumulator = "tileAcc" if gemmCompensatedTiles else "acc" %}
 {% set regT = "f32" %}{% set regCast = 1 if tileT != "f32" else 0 %}
 {% filter indent(8, true) %}
+{% set regAccumulator = regAccumulator | default("acc") %}
 let aRow = lid.y * TM;
 let bCol = lid.x * TN;
 for (var kv: u32 = 0u; kv < BK / 4u; kv = kv + 1u) {
   for (var j: u32 = 0u; j < TN; j = j + 1u) { bv[j] = {% if regCast %}vec4<{{ regT }}>(tileB[bCol + j][kv]){% else %}tileB[bCol + j][kv]{% endif %}; }
   for (var i: u32 = 0u; i < TM; i = i + 1u) {
     for (var j: u32 = 0u; j < TN; j = j + 1u) {
+      {{ regAccumulator }}[i * TN + j] = {{ regAccumulator }}[i * TN + j] + dot(av[i], bv[j]);
     }
   }
 }
 {% endfilter %}
+{% if gemmCompensatedTiles %}
+        for (var i: u32 = 0u; i < TM * TN; i = i + 1u) {
+          // Explicit fma preserves the rounded differences in this correction.
+          // Plain subtraction can be reassociated away by fast-math backends.
+          let corrected = fma(-1.0, accError[i], tileAcc[i]);
+          let next = acc[i] + corrected;
+          accError[i] = select(0.0, fma(-1.0, corrected, fma(-1.0, acc[i], next)), abs(next) <= 3.4028234663852886e38);
+          acc[i] = next;
+        }
+{% endif %}
         workgroupBarrier();
       }

build/webgpu/gemm-subgroup-matrix.wgsl.jinja CHANGED Viewed

@@ -5,7 +5,6 @@ enable subgroup_size_control;
 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set operandScalar = mmaScalar %}
@@ -14,18 +13,15 @@ diagnostic(off, chromium.subgroup_matrix_uniformity);
 {% set OUT_SCALAR = outScalar if outScalar is defined else yScalar %}
 {% set tileMValue = tileM if tileM is defined else 32 %}
 {% set tileNValue = tileN if tileN is defined else 64 %}
-{% set tileKValue = tileK if tileK is defined else 32 %}
-{% set gemmEpi = gemmEpilogue if gemmEpilogue is defined else "none" %}
-{% set epiOpen = "gelu_tanh(" if gemmEpi == "fastgelu" else ("fused_act(" if gemmEpi == "activation" else "") %}
-{% set epiClose = ")" if gemmEpi != "none" else "" %}
-{% set useDirectStore = directStore is defined and directStore and gemmEpi == "none" %}
 {% set useDirectAStorage = directAStorage is defined and directAStorage %}
 {% set useDoubleBufferedB = doubleBufferedB is defined and doubleBufferedB %}
 {% set useBalancedAccumulation = balancedMatrixAccumulation is defined and balancedMatrixAccumulation and useDoubleBufferedB and not useDirectStore and operandScalar == "f32" %}
-{% set aTransposed = aTransposedStorage if aTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
-{% set bTransposed = bTransposedStorage if bTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
-{% set columnTileBaseValue = columnTileBase if columnTileBase is defined else 0 %}
-{% set rowTileBaseValue = rowTileBase if rowTileBase is defined else 0 %}
 // Four subgroups cover 32xN as a 2x2 grid, or taller tiles as four row bands.
 {% if tileMValue > 32 %}
 {% set subtileCols = tileNValue %}
@@ -38,7 +34,6 @@ diagnostic(off, chromium.subgroup_matrix_uniformity);
 {% set bTilesPerSubgroup = (subtileCols / 8)|int %}
 {% set subgroupCount = 4 %}
 {% set scratchBanks = subgroupCount * aTilesPerSubgroup * bTilesPerSubgroup %}
 {% macro matrixAccumulators() %}
 {% for m in range(aTilesPerSubgroup) %}
 {% for n in range(bTilesPerSubgroup) %}
@@ -50,9 +45,7 @@ diagnostic(off, chromium.subgroup_matrix_uniformity);
 {% endif %}
 {% endfor %}
 {% endfor %}
 {% endmacro %}
 {% macro storeMatrixAccumulators() %}
 {% for m in range(aTilesPerSubgroup) %}
 {% for n in range(bTilesPerSubgroup) %}
@@ -65,8 +58,13 @@ diagnostic(off, chromium.subgroup_matrix_uniformity);
 {% endmacro %}
 fn tanh_safe(x: f32) -> f32 {
   if (x > 10.0) { return 1.0; }
   if (x < -10.0) { return -1.0; }
   return tanh(x);
 }
@@ -75,7 +73,7 @@ fn gelu_tanh(v: f32) -> f32 {
 }
 const IN_F:       u32 = {{ inFeatures }}u;
 const OUT_F:      u32 = {{ outFeatures }}u;
-{% set kLoopEndValue = kLoopEnd if kLoopEnd is defined else ("K_LOOP" if tailSafe is defined and tailSafe else "IN_F") %}
 const TILE_COLS:  u32 = {{ tileNValue }}u;
 const TILE_ROWS:  u32 = {{ tileMValue }}u;
 const TILE_K:     u32 = {{ tileKValue }}u;
@@ -191,40 +189,28 @@ fn main(
 {% set matrixAccumulator = matrixAccumulator | default("matC") %}
 {% set matrixStepOffset = matrixStepOffset | default(0) %}
 {% set matrixStep = "step" if matrixStepOffset == 0 else "(step + " ~ matrixStepOffset ~ "u)" %}
-{% set aTransposed = aTransposedStorage if aTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
-{% set bTransposed = bTransposedStorage if bTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
 {% set directA = directAStorage is defined and directAStorage %}
 {% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
 {% if directA %}
       // Direct storage guarantees a full aligned M/K tile, so subgroup-matrix
       // loads need no guarded workgroup staging.
-      let matrix_a_offset = (a_global_base + base_A) * IN_F + kidx + {{ matrixStep }};
-{% elif aTransposed %}
-      let matrix_a_offset = {{ matrixStep }} * TILE_ROWS + subtile_idy * SUB_ROWS;
 {% else %}
       let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + {{ matrixStep }};
 {% endif %}
 {% for m in range(dynamicATiles) %}
 {% if directA %}
-      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&x, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * IN_F{% endif %}, IN_F);
-{% elif aTransposed %}
-      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, col_major>(&tile_A, matrix_a_offset + {{ m * 8 }}u, TILE_ROWS);
 {% else %}
       var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, TILE_K);
 {% endif %}
 {% endfor %}
-{% if bTransposed %}
       let matrix_b_offset = {{ B_BANK }}{{ matrixStep }} * TILE_COLS + subtile_idx * SUB_COLS;
-{% else %}
-      let matrix_b_offset = {{ B_BANK }}subtile_idx * SUB_COLS * TILE_K + {{ matrixStep }};
-{% endif %}
 {% for n in range(bTilesPerSubgroup) %}
-{% if bTransposed %}
       var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, row_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u, TILE_COLS);
-{% else %}
-      var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u * TILE_K, TILE_K);
-{% endif %}
 {% endfor %}
 {% for m in range(dynamicATiles) %}
@@ -232,7 +218,6 @@ fn main(
       {{ matrixAccumulator }}{{ m }}{{ n }} = subgroupMatrixMultiplyAccumulate(matA{{ m }}, matB{{ n }}, {{ matrixAccumulator }}{{ m }}{{ n }});
 {% endfor %}
 {% endfor %}
     }
 {% endfor %}
 {% else %}
@@ -240,40 +225,28 @@ fn main(
 {% set matrixAccumulator = matrixAccumulator | default("matC") %}
 {% set matrixStepOffset = matrixStepOffset | default(0) %}
 {% set matrixStep = "step" if matrixStepOffset == 0 else "(step + " ~ matrixStepOffset ~ "u)" %}
-{% set aTransposed = aTransposedStorage if aTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
-{% set bTransposed = bTransposedStorage if bTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
 {% set directA = directAStorage is defined and directAStorage %}
 {% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
 {% if directA %}
       // Direct storage guarantees a full aligned M/K tile, so subgroup-matrix
       // loads need no guarded workgroup staging.
-      let matrix_a_offset = (a_global_base + base_A) * IN_F + kidx + {{ matrixStep }};
-{% elif aTransposed %}
-      let matrix_a_offset = {{ matrixStep }} * TILE_ROWS + subtile_idy * SUB_ROWS;
 {% else %}
       let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + {{ matrixStep }};
 {% endif %}
 {% for m in range(dynamicATiles) %}
 {% if directA %}
-      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&x, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * IN_F{% endif %}, IN_F);
-{% elif aTransposed %}
-      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, col_major>(&tile_A, matrix_a_offset + {{ m * 8 }}u, TILE_ROWS);
 {% else %}
       var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, TILE_K);
 {% endif %}
 {% endfor %}
-{% if bTransposed %}
       let matrix_b_offset = {{ B_BANK }}{{ matrixStep }} * TILE_COLS + subtile_idx * SUB_COLS;
-{% else %}
-      let matrix_b_offset = {{ B_BANK }}subtile_idx * SUB_COLS * TILE_K + {{ matrixStep }};
-{% endif %}
 {% for n in range(bTilesPerSubgroup) %}
-{% if bTransposed %}
       var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, row_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u, TILE_COLS);
-{% else %}
-      var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u * TILE_K, TILE_K);
-{% endif %}
 {% endfor %}
 {% for m in range(dynamicATiles) %}
@@ -281,7 +254,6 @@ fn main(
       {{ matrixAccumulator }}{{ m }}{{ n }} = subgroupMatrixMultiplyAccumulate(matA{{ m }}, matB{{ n }}, {{ matrixAccumulator }}{{ m }}{{ n }});
 {% endfor %}
 {% endfor %}
 {% endif %}
     }
     let next_kidx = kidx + TILE_K;
@@ -305,40 +277,28 @@ fn main(
 {% set matrixAccumulator = matrixAccumulator | default("matC") %}
 {% set matrixStepOffset = matrixStepOffset | default(0) %}
 {% set matrixStep = "step" if matrixStepOffset == 0 else "(step + " ~ matrixStepOffset ~ "u)" %}
-{% set aTransposed = aTransposedStorage if aTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
-{% set bTransposed = bTransposedStorage if bTransposedStorage is defined else (matrixInputsTransposed is defined and matrixInputsTransposed) %}
 {% set directA = directAStorage is defined and directAStorage %}
 {% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
 {% if directA %}
       // Direct storage guarantees a full aligned M/K tile, so subgroup-matrix
       // loads need no guarded workgroup staging.
-      let matrix_a_offset = (a_global_base + base_A) * IN_F + kidx + {{ matrixStep }};
-{% elif aTransposed %}
-      let matrix_a_offset = {{ matrixStep }} * TILE_ROWS + subtile_idy * SUB_ROWS;
 {% else %}
       let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + {{ matrixStep }};
 {% endif %}
 {% for m in range(dynamicATiles) %}
 {% if directA %}
-      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&x, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * IN_F{% endif %}, IN_F);
-{% elif aTransposed %}
-      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, col_major>(&tile_A, matrix_a_offset + {{ m * 8 }}u, TILE_ROWS);
 {% else %}
       var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, TILE_K);
 {% endif %}
 {% endfor %}
-{% if bTransposed %}
       let matrix_b_offset = {{ B_BANK }}{{ matrixStep }} * TILE_COLS + subtile_idx * SUB_COLS;
-{% else %}
-      let matrix_b_offset = {{ B_BANK }}subtile_idx * SUB_COLS * TILE_K + {{ matrixStep }};
-{% endif %}
 {% for n in range(bTilesPerSubgroup) %}
-{% if bTransposed %}
       var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, row_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u, TILE_COLS);
-{% else %}
-      var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, col_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u * TILE_K, TILE_K);
-{% endif %}
 {% endfor %}
 {% for m in range(dynamicATiles) %}
@@ -346,7 +306,6 @@ fn main(
       {{ matrixAccumulator }}{{ m }}{{ n }} = subgroupMatrixMultiplyAccumulate(matA{{ m }}, matB{{ n }}, {{ matrixAccumulator }}{{ m }}{{ n }});
 {% endfor %}
 {% endfor %}
     }
     workgroupBarrier();
   }

 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 {% set operandScalar = mmaScalar %}
 {% set OUT_SCALAR = outScalar if outScalar is defined else yScalar %}
 {% set tileMValue = tileM if tileM is defined else 32 %}
 {% set tileNValue = tileN if tileN is defined else 64 %}
+{% set tileKValue = 32 %}
+{% set epiOpen = "gelu_tanh(" %}
+{% set epiClose = ")" %}
+{% set useDirectStore = false %}
 {% set useDirectAStorage = directAStorage is defined and directAStorage %}
 {% set useDoubleBufferedB = doubleBufferedB is defined and doubleBufferedB %}
 {% set useBalancedAccumulation = balancedMatrixAccumulation is defined and balancedMatrixAccumulation and useDoubleBufferedB and not useDirectStore and operandScalar == "f32" %}
+{% set columnTileBaseValue = 0 %}
+{% set rowTileBaseValue = 0 %}
 // Four subgroups cover 32xN as a 2x2 grid, or taller tiles as four row bands.
 {% if tileMValue > 32 %}
 {% set subtileCols = tileNValue %}
 {% set bTilesPerSubgroup = (subtileCols / 8)|int %}
 {% set subgroupCount = 4 %}
 {% set scratchBanks = subgroupCount * aTilesPerSubgroup * bTilesPerSubgroup %}
 {% macro matrixAccumulators() %}
 {% for m in range(aTilesPerSubgroup) %}
 {% for n in range(bTilesPerSubgroup) %}
 {% endif %}
 {% endfor %}
 {% endfor %}
 {% endmacro %}
 {% macro storeMatrixAccumulators() %}
 {% for m in range(aTilesPerSubgroup) %}
 {% for n in range(bTilesPerSubgroup) %}
 {% endmacro %}
 fn tanh_safe(x: f32) -> f32 {
+  // tanh rounds to its saturated value for these tails in f32. Return that
+  // value directly, including for infinite input, before invoking the builtin.
   if (x > 10.0) { return 1.0; }
   if (x < -10.0) { return -1.0; }
+  // For tiny |x|, return x directly to preserve its sign and magnitude without
+  // relying on backend-specific builtin behavior near zero.
+  if (x > -1.0e-4 && x < 1.0e-4) { return x; }
   return tanh(x);
 }
 }
 const IN_F:       u32 = {{ inFeatures }}u;
 const OUT_F:      u32 = {{ outFeatures }}u;
+{% set kLoopEndValue = "IN_F" %}
 const TILE_COLS:  u32 = {{ tileNValue }}u;
 const TILE_ROWS:  u32 = {{ tileMValue }}u;
 const TILE_K:     u32 = {{ tileKValue }}u;
 {% set matrixAccumulator = matrixAccumulator | default("matC") %}
 {% set matrixStepOffset = matrixStepOffset | default(0) %}
 {% set matrixStep = "step" if matrixStepOffset == 0 else "(step + " ~ matrixStepOffset ~ "u)" %}
 {% set directA = directAStorage is defined and directAStorage %}
 {% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
 {% if directA %}
+{% set matrixAStorage = matrixAStorage | default("x") %}
+{% set matrixAStride = matrixAStride | default("IN_F") %}
       // Direct storage guarantees a full aligned M/K tile, so subgroup-matrix
       // loads need no guarded workgroup staging.
+      let matrix_a_offset = (a_global_base + base_A) * {{ matrixAStride }} + kidx + {{ matrixStep }};
 {% else %}
       let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + {{ matrixStep }};
 {% endif %}
 {% for m in range(dynamicATiles) %}
 {% if directA %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ matrixAStorage }}, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * {{ matrixAStride }}{% endif %}, {{ matrixAStride }});
 {% else %}
       var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, TILE_K);
 {% endif %}
 {% endfor %}
       let matrix_b_offset = {{ B_BANK }}{{ matrixStep }} * TILE_COLS + subtile_idx * SUB_COLS;
 {% for n in range(bTilesPerSubgroup) %}
       var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, row_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u, TILE_COLS);
 {% endfor %}
 {% for m in range(dynamicATiles) %}
       {{ matrixAccumulator }}{{ m }}{{ n }} = subgroupMatrixMultiplyAccumulate(matA{{ m }}, matB{{ n }}, {{ matrixAccumulator }}{{ m }}{{ n }});
 {% endfor %}
 {% endfor %}
     }
 {% endfor %}
 {% else %}
 {% set matrixAccumulator = matrixAccumulator | default("matC") %}
 {% set matrixStepOffset = matrixStepOffset | default(0) %}
 {% set matrixStep = "step" if matrixStepOffset == 0 else "(step + " ~ matrixStepOffset ~ "u)" %}
 {% set directA = directAStorage is defined and directAStorage %}
 {% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
 {% if directA %}
+{% set matrixAStorage = matrixAStorage | default("x") %}
+{% set matrixAStride = matrixAStride | default("IN_F") %}
       // Direct storage guarantees a full aligned M/K tile, so subgroup-matrix
       // loads need no guarded workgroup staging.
+      let matrix_a_offset = (a_global_base + base_A) * {{ matrixAStride }} + kidx + {{ matrixStep }};
 {% else %}
       let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + {{ matrixStep }};
 {% endif %}
 {% for m in range(dynamicATiles) %}
 {% if directA %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ matrixAStorage }}, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * {{ matrixAStride }}{% endif %}, {{ matrixAStride }});
 {% else %}
       var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, TILE_K);
 {% endif %}
 {% endfor %}
       let matrix_b_offset = {{ B_BANK }}{{ matrixStep }} * TILE_COLS + subtile_idx * SUB_COLS;
 {% for n in range(bTilesPerSubgroup) %}
       var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, row_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u, TILE_COLS);
 {% endfor %}
 {% for m in range(dynamicATiles) %}
       {{ matrixAccumulator }}{{ m }}{{ n }} = subgroupMatrixMultiplyAccumulate(matA{{ m }}, matB{{ n }}, {{ matrixAccumulator }}{{ m }}{{ n }});
 {% endfor %}
 {% endfor %}
 {% endif %}
     }
     let next_kidx = kidx + TILE_K;
 {% set matrixAccumulator = matrixAccumulator | default("matC") %}
 {% set matrixStepOffset = matrixStepOffset | default(0) %}
 {% set matrixStep = "step" if matrixStepOffset == 0 else "(step + " ~ matrixStepOffset ~ "u)" %}
 {% set directA = directAStorage is defined and directAStorage %}
 {% set B_BANK = "b_bank_offset + " if doubleBufferedB is defined and doubleBufferedB else "" %}
 {% if directA %}
+{% set matrixAStorage = matrixAStorage | default("x") %}
+{% set matrixAStride = matrixAStride | default("IN_F") %}
       // Direct storage guarantees a full aligned M/K tile, so subgroup-matrix
       // loads need no guarded workgroup staging.
+      let matrix_a_offset = (a_global_base + base_A) * {{ matrixAStride }} + kidx + {{ matrixStep }};
 {% else %}
       let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + {{ matrixStep }};
 {% endif %}
 {% for m in range(dynamicATiles) %}
 {% if directA %}
+      var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ matrixAStorage }}, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * {{ matrixAStride }}{% endif %}, {{ matrixAStride }});
 {% else %}
       var matA{{ m }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&tile_A, matrix_a_offset{% if m > 0 %} + {{ m * 8 }}u * TILE_K{% endif %}, TILE_K);
 {% endif %}
 {% endfor %}
       let matrix_b_offset = {{ B_BANK }}{{ matrixStep }} * TILE_COLS + subtile_idx * SUB_COLS;
 {% for n in range(bTilesPerSubgroup) %}
       var matB{{ n }}: subgroup_matrix_right<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_right<{{ operandScalar }}, 8, 8>, row_major>(&tile_B, matrix_b_offset + {{ n * 8 }}u, TILE_COLS);
 {% endfor %}
 {% for m in range(dynamicATiles) %}
       {{ matrixAccumulator }}{{ m }}{{ n }} = subgroupMatrixMultiplyAccumulate(matA{{ m }}, matB{{ n }}, {{ matrixAccumulator }}{{ m }}{{ n }});
 {% endfor %}
 {% endfor %}
     }
     workgroupBarrier();
   }

build/webgpu/manifest.json CHANGED Viewed

@@ -10,13 +10,21 @@
   "outputs": { "Y": { "dtype": "T", "rank": "ranks.X", "shape": "matmulShape(shapes.X, shapes.W)" } },
   "typeConstraints": { "T": ["float32", "float16"] },
   "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",
     "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",
-    "kDim": "dim(shapes.W, 0)",
-    "nDim": "dim(shapes.W, 1)",
     "gfgRows": "numel(shapes.X) / kDim if kDim > 0 else 0",
     "shapeOk": "ranks.W == 2 and ranks.X >= 2 and kDim > 0 and nDim > 0 and dim(shapes.X, ranks.X - 1) == kDim and ranks.Y == ranks.X and dim(shapes.Y, ranks.Y - 1) == nDim and numel(shapes.Y) == gfgRows * nDim",
     "dtypeOk": "(tensorDtypes.X == \"float32\" or tensorDtypes.X == \"float16\") and tensorDtypes.W == tensorDtypes.X and tensorDtypes.Y == tensorDtypes.X and f16Ok(dtypes.T)",
@@ -30,24 +38,416 @@
     "sgmatLayoutOk": "gfgRows > 0 and kDim % 32 == 0 and nDim % 64 == 0 and ((tensorDtypes.X == \"float16\" and device.features.has(\"shader-f16\") and gfgRows >= 2) or (tensorDtypes.X == \"float32\" and gfgRows >= 32)) and ceilDiv(gfgRows, 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(nDim, 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "sgmatContract": "wave32Effective and sgmatLayoutOk and sgmatResourcesFit",
     "usesF16": "tensorDtypes.X == \"float16\"",
-    "transA": false,
     "transB": false,
-    "gemmEpilogue": "\"fastgelu\"",
     "sgmatDirectRows": 64,
     "sgmatDirectCols": 32,
     "sgmatDirectStorageBytes": "2 * sgmatDirectCols * 32 * sgmatOperandBytes + sgmatDirectRows * sgmatDirectCols * 4",
-    "sgmatDirectContract": "wave32Effective and gfgRows > 0 and gfgRows % sgmatDirectRows == 0 and kDim % 32 == 0 and nDim % sgmatDirectCols == 0 and deviceWorkgroupCap >= 128 and sgmatDirectStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(gfgRows, sgmatDirectRows) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(nDim, sgmatDirectCols) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"
   },
   "bindings": {
-    "x": { "arg": "X", "buffer": "read-only-storage", "elementType": "$xScalar" },
-    "w": { "arg": "W", "buffer": "read-only-storage", "elementType": "$wScalar" },
-    "y": { "arg": "Y", "buffer": "storage", "elementType": "$yScalar" },
-    "params": { "buffer": "uniform", "struct": [{ "name": "M", "type": "u32", "value": "gfgRows" }] },
-    "a": { "arg": "X", "buffer": "read-only-storage", "elementType": "$T" },
-    "b": { "arg": "W", "buffer": "read-only-storage", "elementType": "$T" },
-    "y_2": { "arg": "Y", "name": "y", "buffer": "storage", "elementType": "$T" }
   },
   "variants": [
     {
       "id": "sgmat_direct_bias",
       "priority": 110,
@@ -58,7 +458,6 @@
       },
       "derive": {
         "hasBias": true,
-        "bTransposedStorage": true,
         "xScalar": "\"f32\"",
         "wScalar": "\"f32\"",
         "bScalar": "\"f32\"",
@@ -68,7 +467,6 @@
         "outFeatures": "nDim",
         "tileM": "sgmatDirectRows",
         "tileN": "sgmatDirectCols",
-        "tileK": 32,
         "directAStorage": true,
         "doubleBufferedB": true,
         "balancedMatrixAccumulation": true
@@ -76,7 +474,7 @@
       "passes": [
         {
           "id": "main",
-          "name": "GemmFastGelu.sgmat_direct_bias",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
           "dispatch": {
@@ -85,7 +483,8 @@
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
-      ]
     },
     {
       "id": "sgmat_direct",
@@ -97,7 +496,6 @@
       },
       "derive": {
         "hasBias": false,
-        "bTransposedStorage": true,
         "xScalar": "\"f32\"",
         "wScalar": "\"f32\"",
         "bScalar": "\"f32\"",
@@ -107,7 +505,6 @@
         "outFeatures": "nDim",
         "tileM": "sgmatDirectRows",
         "tileN": "sgmatDirectCols",
-        "tileK": 32,
         "directAStorage": true,
         "doubleBufferedB": true,
         "balancedMatrixAccumulation": true
@@ -115,7 +512,7 @@
       "passes": [
         {
           "id": "main",
-          "name": "GemmFastGelu.sgmat_direct",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", "y", "params"],
           "dispatch": {
@@ -124,7 +521,8 @@
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
-      ]
     },
     {
       "id": "sgmat_direct_bias_f16",
@@ -136,7 +534,6 @@
       },
       "derive": {
         "hasBias": true,
-        "bTransposedStorage": true,
         "xScalar": "\"f16\"",
         "wScalar": "\"f16\"",
         "bScalar": "\"f16\"",
@@ -146,7 +543,6 @@
         "outFeatures": "nDim",
         "tileM": "sgmatDirectRows",
         "tileN": "sgmatDirectCols",
-        "tileK": 32,
         "directAStorage": true,
         "doubleBufferedB": true,
         "balancedMatrixAccumulation": false
@@ -154,7 +550,7 @@
       "passes": [
         {
           "id": "main",
-          "name": "GemmFastGelu.sgmat_direct_bias_f16",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
           "dispatch": {
@@ -163,7 +559,8 @@
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
-      ]
     },
     {
       "id": "sgmat_direct_f16",
@@ -175,7 +572,6 @@
       },
       "derive": {
         "hasBias": false,
-        "bTransposedStorage": true,
         "xScalar": "\"f16\"",
         "wScalar": "\"f16\"",
         "bScalar": "\"f16\"",
@@ -185,7 +581,6 @@
         "outFeatures": "nDim",
         "tileM": "sgmatDirectRows",
         "tileN": "sgmatDirectCols",
-        "tileK": 32,
         "directAStorage": true,
         "doubleBufferedB": true,
         "balancedMatrixAccumulation": false
@@ -193,7 +588,7 @@
       "passes": [
         {
           "id": "main",
-          "name": "GemmFastGelu.sgmat_direct_f16",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", "y", "params"],
           "dispatch": {
@@ -202,7 +597,8 @@
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
-      ]
     },
     {
       "id": "sgmat_bias",
@@ -214,7 +610,6 @@
       },
       "derive": {
         "hasBias": true,
-        "bTransposedStorage": true,
         "xScalar": "\"f32\"",
         "wScalar": "\"f32\"",
         "bScalar": "\"f32\"",
@@ -224,7 +619,6 @@
         "outFeatures": "nDim",
         "tileM": 32,
         "tileN": 64,
-        "tileK": 32,
         "directAStorage": false,
         "doubleBufferedB": false,
         "balancedMatrixAccumulation": false
@@ -232,7 +626,7 @@
       "passes": [
         {
           "id": "main",
-          "name": "GemmFastGelu.sgmat_bias",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
           "dispatch": {
@@ -241,7 +635,8 @@
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
-      ]
     },
     {
       "id": "sgmat",
@@ -253,7 +648,6 @@
       },
       "derive": {
         "hasBias": false,
-        "bTransposedStorage": true,
         "xScalar": "\"f32\"",
         "wScalar": "\"f32\"",
         "bScalar": "\"f32\"",
@@ -263,7 +657,6 @@
         "outFeatures": "nDim",
         "tileM": 32,
         "tileN": 64,
-        "tileK": 32,
         "directAStorage": false,
         "doubleBufferedB": false,
         "balancedMatrixAccumulation": false
@@ -271,7 +664,7 @@
       "passes": [
         {
           "id": "main",
-          "name": "GemmFastGelu.sgmat",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", "y", "params"],
           "dispatch": {
@@ -280,7 +673,8 @@
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
-      ]
     },
     {
       "id": "sgmat_bias_f16",
@@ -292,7 +686,6 @@
       },
       "derive": {
         "hasBias": true,
-        "bTransposedStorage": true,
         "xScalar": "\"f16\"",
         "wScalar": "\"f16\"",
         "bScalar": "\"f16\"",
@@ -302,7 +695,6 @@
         "outFeatures": "nDim",
         "tileM": 32,
         "tileN": 64,
-        "tileK": 32,
         "directAStorage": false,
         "doubleBufferedB": false,
         "balancedMatrixAccumulation": false
@@ -310,7 +702,7 @@
       "passes": [
         {
           "id": "main",
-          "name": "GemmFastGelu.sgmat_bias_f16",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
           "dispatch": {
@@ -319,7 +711,8 @@
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
-      ]
     },
     {
       "id": "sgmat_f16",
@@ -331,7 +724,6 @@
       },
       "derive": {
         "hasBias": false,
-        "bTransposedStorage": true,
         "xScalar": "\"f16\"",
         "wScalar": "\"f16\"",
         "bScalar": "\"f16\"",
@@ -341,7 +733,6 @@
         "outFeatures": "nDim",
         "tileM": 32,
         "tileN": 64,
-        "tileK": 32,
         "directAStorage": false,
         "doubleBufferedB": false,
         "balancedMatrixAccumulation": false
@@ -349,7 +740,7 @@
       "passes": [
         {
           "id": "main",
-          "name": "GemmFastGelu.sgmat_f16",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", "y", "params"],
           "dispatch": {
@@ -358,7 +749,8 @@
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
-      ]
     },
     {
       "id": "tiled_bias",
@@ -369,12 +761,12 @@
         {
           "id": "main",
           "name": "GemmFastGelu.Bias",
-          "shader": "gemm-fast-gelu.wgsl.jinja",
           "bindings": [
             "a",
             "b",
             "bias",
-            "y_2",
             {
               "name": "params",
               "struct": [
@@ -387,8 +779,8 @@
             }
           ],
           "dispatch": {
-            "x": "min(ceilDiv(nDim, 64), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
-            "y": "min(ceilDiv(gfgRows, 64), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))"
           }
         }
       ]
@@ -402,11 +794,11 @@
         {
           "id": "main",
           "name": "GemmFastGelu",
-          "shader": "gemm-fast-gelu.wgsl.jinja",
           "bindings": [
             "a",
             "b",
-            "y_2",
             {
               "name": "params",
               "struct": [
@@ -418,11 +810,19 @@
             }
           ],
           "dispatch": {
-            "x": "min(ceilDiv(nDim, 64), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
-            "y": "min(ceilDiv(gfgRows, 64), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))"
           }
         }
       ]
     }
-  ]
 }

   "outputs": { "Y": { "dtype": "T", "rank": "ranks.X", "shape": "matmulShape(shapes.X, shapes.W)" } },
   "typeConstraints": { "T": ["float32", "float16"] },
   "derive": {
+    "kDim": "dim(shapes.W, 0)",
+    "nDim": "dim(shapes.W, 1)",
+    "gemvN": "nDim",
     "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",
     "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",
+    "variableSubgroup16To32": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 16 and device.adapterInfo.subgroupMaxSize == 32",
+    "gemvLanes": 32,
+    "vec4OutputTile": "4 * gemvLanes",
+    "gemvWorkgroups": "ceilDiv(gemvN, vec4OutputTile)",
+    "gemvSliceCap": "min(32, device.limits.maxComputeWorkgroupSizeY, floor(device.limits.maxComputeInvocationsPerWorkgroup / gemvLanes), floor(device.limits.maxComputeWorkgroupStorageSize / (16 * gemvLanes)))",
+    "gemvSlices": "max(1, min(gemvSliceCap, max(8, pow2ceil(ceilDiv(tunables.GEMV_TARGET_BLOCKS, gemvWorkgroups)))))",
+    "gemvResourcesFit": "gemvLanes <= device.limits.maxComputeWorkgroupSizeX and gemvSlices <= device.limits.maxComputeWorkgroupSizeY and gemvLanes * gemvSlices <= device.limits.maxComputeInvocationsPerWorkgroup and 16 * gemvLanes * gemvSlices <= device.limits.maxComputeWorkgroupStorageSize",
     "gfgRows": "numel(shapes.X) / kDim if kDim > 0 else 0",
     "shapeOk": "ranks.W == 2 and ranks.X >= 2 and kDim > 0 and nDim > 0 and dim(shapes.X, ranks.X - 1) == kDim and ranks.Y == ranks.X and dim(shapes.Y, ranks.Y - 1) == nDim and numel(shapes.Y) == gfgRows * nDim",
     "dtypeOk": "(tensorDtypes.X == \"float32\" or tensorDtypes.X == \"float16\") and tensorDtypes.W == tensorDtypes.X and tensorDtypes.Y == tensorDtypes.X and f16Ok(dtypes.T)",
     "sgmatLayoutOk": "gfgRows > 0 and kDim % 32 == 0 and nDim % 64 == 0 and ((tensorDtypes.X == \"float16\" and device.features.has(\"shader-f16\") and gfgRows >= 2) or (tensorDtypes.X == \"float32\" and gfgRows >= 32)) and ceilDiv(gfgRows, 32) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(nDim, 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "sgmatContract": "wave32Effective and sgmatLayoutOk and sgmatResourcesFit",
     "usesF16": "tensorDtypes.X == \"float16\"",
     "transB": false,
     "sgmatDirectRows": 64,
     "sgmatDirectCols": 32,
     "sgmatDirectStorageBytes": "2 * sgmatDirectCols * 32 * sgmatOperandBytes + sgmatDirectRows * sgmatDirectCols * 4",
+    "sgmatDirectContract": "wave32Effective and gfgRows > 0 and gfgRows % sgmatDirectRows == 0 and kDim % 32 == 0 and nDim % sgmatDirectCols == 0 and deviceWorkgroupCap >= 128 and sgmatDirectStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(gfgRows, sgmatDirectRows) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(nDim, sgmatDirectCols) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "sgmatOutputGroups": "ceilDiv(gfgRows, sgmatDirectRows) * ceilDiv(nDim, sgmatDirectCols)",
+    "sgmatSmallRows": 32,
+    "sgmatSmallCols": 16,
+    "sgmatSmallStorageBytes": "2 * sgmatSmallCols * 32 * sgmatOperandBytes + sgmatSmallRows * sgmatSmallCols * 4",
+    "sgmatSmallContract": "wave32Effective and gfgRows > 0 and gfgRows % sgmatSmallRows == 0 and kDim % 32 == 0 and nDim % sgmatSmallCols == 0 and deviceWorkgroupCap >= 128 and sgmatSmallStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(gfgRows, sgmatSmallRows) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(nDim, sgmatSmallCols) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "sgmatMediumRows": 32,
+    "sgmatMediumCols": 32,
+    "sgmatMediumStorageBytes": "2 * sgmatMediumCols * 32 * sgmatOperandBytes + sgmatMediumRows * sgmatMediumCols * 4",
+    "sgmatMediumContract": "wave32Effective and gfgRows > 0 and gfgRows % sgmatMediumRows == 0 and kDim % 32 == 0 and nDim % sgmatMediumCols == 0 and deviceWorkgroupCap >= 128 and sgmatMediumStorageBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(gfgRows, sgmatMediumRows) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(nDim, sgmatMediumCols) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "portableSmallGroups": "ceilDiv(gfgRows, 32) * ceilDiv(nDim, 32)",
+    "portableLargeGroups": "ceilDiv(gfgRows, 64) * ceilDiv(nDim, 64)",
+    "portableSmallBudget": "tunables.PORTABLE_SMALL_TILE_GROUPS_F16 if usesF16 else tunables.PORTABLE_SMALL_TILE_GROUPS_F32",
+    "portableAvoidsPadding": "portableLargeGroups * 3 >= portableSmallGroups",
+    "portableXe3LpgF32": "not usesF16 and has(device.adapterInfo, \"vendor\") and device.adapterInfo.vendor == \"intel\" and has(device.adapterInfo, \"architecture\") and device.adapterInfo.architecture == \"xe-3lpg\"",
+    "portableFixedWave32NoF16": "wave32Effective and not device.features.has(\"shader-f16\") and gfgRows >= device.adapterInfo.subgroupMaxSize and kDim >= 16 * device.adapterInfo.subgroupMaxSize and portableSmallGroups >= 8 * device.adapterInfo.subgroupMaxSize",
+    "portableLargeNoSubgroupGrid": "not usesF16 and not device.features.has(\"subgroups\") and portableSmallGroups >= 4 * deviceWorkgroupCap",
+    "gemmTileRows": "32 if portableXe3LpgF32 or portableFixedWave32NoF16 or portableLargeNoSubgroupGrid or portableSmallGroups <= portableSmallBudget * (2 if portableAvoidsPadding else 1) else 64",
+    "gemmTileCols": "gemmTileRows",
+    "gfgWideMicroPreferred": "(variableSubgroup16To32 or portableFixedWave32NoF16) and (8 if usesF16 else 4) * device.adapterInfo.subgroupMaxSize <= ceilDiv(gfgRows, gemmTileRows) * ceilDiv(nDim, gemmTileCols) and gemmTileRows * gemmTileCols / 8 <= deviceWorkgroupCap and gemmTileCols / 2 <= device.limits.maxComputeWorkgroupSizeX and gemmTileRows / 4 <= device.limits.maxComputeWorkgroupSizeY",
+    "gemmMicroRows": "4 if gfgWideMicroPreferred else gemmTileRows / 16",
+    "gemmMicroCols": "2 if gfgWideMicroPreferred else gemmTileCols / 16",
+    "gemmTileK": "(16 * 16 * 4) / gemmTileRows if usesF16 else (16 if portableLargeNoSubgroupGrid and has(device.adapterInfo, \"vendor\") and device.adapterInfo.vendor == \"intel\" else 32)",
+    "gemmTilePadding": "1 if not usesF16 and 2 * gemmTileRows * (gemmTileK + 4) * 4 <= device.limits.maxComputeWorkgroupStorageSize else 0",
+    "gemmCompensatedTiles": "not usesF16 and kDim > gemmTileK",
+    "gemmStaticM": "gfgRows",
+    "gemmStaticN": "nDim",
+    "gemmStaticK": "kDim"
   },
   "bindings": {
+    "x": { "arg": "X", "elementType": "$xScalar" },
+    "w": { "arg": "W", "elementType": "$wScalar" },
+    "y": { "arg": "Y", "elementType": "$yScalar" },
+    "params": { "struct": [{ "name": "M", "type": "u32", "value": "gfgRows" }] },
+    "a": { "arg": "X", "elementType": "$T" },
+    "b": { "arg": "W", "elementType": "$T" },
+    "y_t": { "arg": "Y", "name": "y", "elementType": "$T" }
   },
   "variants": [
+    {
+      "id": "vector_matrix_bias",
+      "priority": 130,
+      "when": ["biasContract", "gfgRows == 1", "nDim % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "gemvResourcesFit"],
+      "derive": {
+        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "outputBuffer": "\"y\"",
+        "gemvBias": true,
+        "unrollK2": "usesF16"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.VectorMatrixBias",
+          "shader": "matmul-vector-matrix-vec4.wgsl.jinja",
+          "bindings": [
+            "a",
+            { "name": "b", "arg": "W", "elementType": "$vectorScalar" },
+            { "arg": "bias", "elementType": "$vectorScalar" },
+            { "name": "y", "arg": "Y", "elementType": "$vectorScalar" },
+            {
+              "name": "params",
+              "struct": [
+                { "name": "K", "type": "u32", "value": "kDim" },
+                { "name": "N4", "type": "u32", "value": "nDim / 4" }
+              ]
+            }
+          ],
+          "dispatch": { "x": "gemvWorkgroups" }
+        }
+      ]
+    },
+    {
+      "id": "vector_matrix",
+      "priority": 130,
+      "when": ["noBiasContract", "gfgRows == 1", "nDim % 4 == 0", "gemvWorkgroups <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", "gemvResourcesFit"],
+      "derive": {
+        "vectorScalar": "\"vec4<\" ~ dtypes.T ~ \">\"",
+        "outputBuffer": "\"y\"",
+        "gemvBias": false,
+        "unrollK2": "usesF16"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.VectorMatrix",
+          "shader": "matmul-vector-matrix-vec4.wgsl.jinja",
+          "bindings": [
+            "a",
+            { "name": "b", "arg": "W", "elementType": "$vectorScalar" },
+            { "name": "y", "arg": "Y", "elementType": "$vectorScalar" },
+            {
+              "name": "params",
+              "struct": [
+                { "name": "K", "type": "u32", "value": "kDim" },
+                { "name": "N4", "type": "u32", "value": "nDim / 4" }
+              ]
+            }
+          ],
+          "dispatch": { "x": "gemvWorkgroups" }
+        }
+      ]
+    },
+    {
+      "id": "sgmat_medium_bias",
+      "priority": 120,
+      "when": ["biasContract", "sgmatMediumContract", "tensorDtypes.X == \"float32\""],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }]
+      },
+      "derive": {
+        "hasBias": true,
+        "xScalar": "\"f32\"",
+        "wScalar": "\"f32\"",
+        "bScalar": "\"f32\"",
+        "yScalar": "\"f32\"",
+        "mmaScalar": "\"f32\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": "sgmatMediumRows",
+        "tileN": "sgmatMediumCols",
+        "directAStorage": true,
+        "doubleBufferedB": true,
+        "balancedMatrixAccumulation": true
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.SgmatMediumBias",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["sgmatOutputGroups > (tunables.MEDIUM_MATRIX_MAX_GROUPS_F16 if usesF16 else tunables.MEDIUM_MATRIX_MAX_GROUPS_F32) and (sgmatDirectContract or sgmatContract)"]
+    },
+    {
+      "id": "sgmat_medium",
+      "priority": 120,
+      "when": ["noBiasContract", "sgmatMediumContract", "tensorDtypes.X == \"float32\""],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }]
+      },
+      "derive": {
+        "hasBias": false,
+        "xScalar": "\"f32\"",
+        "wScalar": "\"f32\"",
+        "bScalar": "\"f32\"",
+        "yScalar": "\"f32\"",
+        "mmaScalar": "\"f32\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": "sgmatMediumRows",
+        "tileN": "sgmatMediumCols",
+        "directAStorage": true,
+        "doubleBufferedB": true,
+        "balancedMatrixAccumulation": true
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.SgmatMedium",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["sgmatOutputGroups > (tunables.MEDIUM_MATRIX_MAX_GROUPS_F16 if usesF16 else tunables.MEDIUM_MATRIX_MAX_GROUPS_F32) and (sgmatDirectContract or sgmatContract)"]
+    },
+    {
+      "id": "sgmat_medium_bias_f16",
+      "priority": 120,
+      "when": ["biasContract", "sgmatMediumContract", "tensorDtypes.X == \"float16\""],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "derive": {
+        "hasBias": true,
+        "xScalar": "\"f16\"",
+        "wScalar": "\"f16\"",
+        "bScalar": "\"f16\"",
+        "yScalar": "\"f16\"",
+        "mmaScalar": "\"f16\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": "sgmatMediumRows",
+        "tileN": "sgmatMediumCols",
+        "directAStorage": true,
+        "doubleBufferedB": true,
+        "balancedMatrixAccumulation": false
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.SgmatMediumBiasF16",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["sgmatOutputGroups > (tunables.MEDIUM_MATRIX_MAX_GROUPS_F16 if usesF16 else tunables.MEDIUM_MATRIX_MAX_GROUPS_F32) and (sgmatDirectContract or sgmatContract)"]
+    },
+    {
+      "id": "sgmat_medium_f16",
+      "priority": 120,
+      "when": ["noBiasContract", "sgmatMediumContract", "tensorDtypes.X == \"float16\""],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "derive": {
+        "hasBias": false,
+        "xScalar": "\"f16\"",
+        "wScalar": "\"f16\"",
+        "bScalar": "\"f16\"",
+        "yScalar": "\"f16\"",
+        "mmaScalar": "\"f16\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": "sgmatMediumRows",
+        "tileN": "sgmatMediumCols",
+        "directAStorage": true,
+        "doubleBufferedB": true,
+        "balancedMatrixAccumulation": false
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.SgmatMediumF16",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["sgmatOutputGroups > (tunables.MEDIUM_MATRIX_MAX_GROUPS_F16 if usesF16 else tunables.MEDIUM_MATRIX_MAX_GROUPS_F32) and (sgmatDirectContract or sgmatContract)"]
+    },
+    {
+      "id": "sgmat_small_bias",
+      "priority": 125,
+      "when": ["biasContract", "sgmatSmallContract", "tensorDtypes.X == \"float32\""],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }]
+      },
+      "derive": {
+        "hasBias": true,
+        "xScalar": "\"f32\"",
+        "wScalar": "\"f32\"",
+        "bScalar": "\"f32\"",
+        "yScalar": "\"f32\"",
+        "mmaScalar": "\"f32\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": "sgmatSmallRows",
+        "tileN": "sgmatSmallCols",
+        "directAStorage": true,
+        "doubleBufferedB": true,
+        "balancedMatrixAccumulation": true
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.SgmatSmallBias",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["sgmatOutputGroups > tunables.SMALL_MATRIX_MAX_GROUPS and (sgmatMediumContract or sgmatDirectContract or sgmatContract)"]
+    },
+    {
+      "id": "sgmat_small",
+      "priority": 125,
+      "when": ["noBiasContract", "sgmatSmallContract", "tensorDtypes.X == \"float32\""],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f32", "M": 8, "N": 8, "K": 8, "resultComponentType": "f32" }]
+      },
+      "derive": {
+        "hasBias": false,
+        "xScalar": "\"f32\"",
+        "wScalar": "\"f32\"",
+        "bScalar": "\"f32\"",
+        "yScalar": "\"f32\"",
+        "mmaScalar": "\"f32\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": "sgmatSmallRows",
+        "tileN": "sgmatSmallCols",
+        "directAStorage": true,
+        "doubleBufferedB": true,
+        "balancedMatrixAccumulation": true
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.SgmatSmall",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["sgmatOutputGroups > tunables.SMALL_MATRIX_MAX_GROUPS and (sgmatMediumContract or sgmatDirectContract or sgmatContract)"]
+    },
+    {
+      "id": "sgmat_small_bias_f16",
+      "priority": 125,
+      "when": ["biasContract", "sgmatSmallContract", "tensorDtypes.X == \"float16\""],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "derive": {
+        "hasBias": true,
+        "xScalar": "\"f16\"",
+        "wScalar": "\"f16\"",
+        "bScalar": "\"f16\"",
+        "yScalar": "\"f16\"",
+        "mmaScalar": "\"f16\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": "sgmatSmallRows",
+        "tileN": "sgmatSmallCols",
+        "directAStorage": true,
+        "doubleBufferedB": true,
+        "balancedMatrixAccumulation": false
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.SgmatSmallBiasF16",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["sgmatOutputGroups > tunables.SMALL_MATRIX_MAX_GROUPS and (sgmatMediumContract or sgmatDirectContract or sgmatContract)"]
+    },
+    {
+      "id": "sgmat_small_f16",
+      "priority": 125,
+      "when": ["noBiasContract", "sgmatSmallContract", "tensorDtypes.X == \"float16\""],
+      "requires": {
+        "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
+        "subgroupMatrixConfigs": [{ "componentType": "f16", "M": 8, "N": 8, "K": 8 }]
+      },
+      "derive": {
+        "hasBias": false,
+        "xScalar": "\"f16\"",
+        "wScalar": "\"f16\"",
+        "bScalar": "\"f16\"",
+        "yScalar": "\"f16\"",
+        "mmaScalar": "\"f16\"",
+        "inFeatures": "kDim",
+        "outFeatures": "nDim",
+        "tileM": "sgmatSmallRows",
+        "tileN": "sgmatSmallCols",
+        "directAStorage": true,
+        "doubleBufferedB": true,
+        "balancedMatrixAccumulation": false
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmFastGelu.SgmatSmallF16",
+          "shader": "gemm-subgroup-matrix.wgsl.jinja",
+          "bindings": ["x", "w", "y", "params"],
+          "dispatch": {
+            "x": "ceilDiv(nDim, tileN)",
+            "y": "min(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)",
+            "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "demoteWhen": ["sgmatOutputGroups > tunables.SMALL_MATRIX_MAX_GROUPS and (sgmatMediumContract or sgmatDirectContract or sgmatContract)"]
+    },
     {
       "id": "sgmat_direct_bias",
       "priority": 110,
       },
       "derive": {
         "hasBias": true,
         "xScalar": "\"f32\"",
         "wScalar": "\"f32\"",
         "bScalar": "\"f32\"",
         "outFeatures": "nDim",
         "tileM": "sgmatDirectRows",
         "tileN": "sgmatDirectCols",
         "directAStorage": true,
         "doubleBufferedB": true,
         "balancedMatrixAccumulation": true
       "passes": [
         {
           "id": "main",
+          "name": "GemmFastGelu.SgmatDirectBias",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
           "dispatch": {
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
+      ],
+      "demoteWhen": ["false"]
     },
     {
       "id": "sgmat_direct",
       },
       "derive": {
         "hasBias": false,
         "xScalar": "\"f32\"",
         "wScalar": "\"f32\"",
         "bScalar": "\"f32\"",
         "outFeatures": "nDim",
         "tileM": "sgmatDirectRows",
         "tileN": "sgmatDirectCols",
         "directAStorage": true,
         "doubleBufferedB": true,
         "balancedMatrixAccumulation": true
       "passes": [
         {
           "id": "main",
+          "name": "GemmFastGelu.SgmatDirect",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", "y", "params"],
           "dispatch": {
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
+      ],
+      "demoteWhen": ["false"]
     },
     {
       "id": "sgmat_direct_bias_f16",
       },
       "derive": {
         "hasBias": true,
         "xScalar": "\"f16\"",
         "wScalar": "\"f16\"",
         "bScalar": "\"f16\"",
         "outFeatures": "nDim",
         "tileM": "sgmatDirectRows",
         "tileN": "sgmatDirectCols",
         "directAStorage": true,
         "doubleBufferedB": true,
         "balancedMatrixAccumulation": false
       "passes": [
         {
           "id": "main",
+          "name": "GemmFastGelu.SgmatDirectBiasF16",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
           "dispatch": {
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
+      ],
+      "demoteWhen": ["false"]
     },
     {
       "id": "sgmat_direct_f16",
       },
       "derive": {
         "hasBias": false,
         "xScalar": "\"f16\"",
         "wScalar": "\"f16\"",
         "bScalar": "\"f16\"",
         "outFeatures": "nDim",
         "tileM": "sgmatDirectRows",
         "tileN": "sgmatDirectCols",
         "directAStorage": true,
         "doubleBufferedB": true,
         "balancedMatrixAccumulation": false
       "passes": [
         {
           "id": "main",
+          "name": "GemmFastGelu.SgmatDirectF16",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", "y", "params"],
           "dispatch": {
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
+      ],
+      "demoteWhen": ["false"]
     },
     {
       "id": "sgmat_bias",
       },
       "derive": {
         "hasBias": true,
         "xScalar": "\"f32\"",
         "wScalar": "\"f32\"",
         "bScalar": "\"f32\"",
         "outFeatures": "nDim",
         "tileM": 32,
         "tileN": 64,
         "directAStorage": false,
         "doubleBufferedB": false,
         "balancedMatrixAccumulation": false
       "passes": [
         {
           "id": "main",
+          "name": "GemmFastGelu.SgmatBias",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
           "dispatch": {
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
+      ],
+      "demoteWhen": ["false"]
     },
     {
       "id": "sgmat",
       },
       "derive": {
         "hasBias": false,
         "xScalar": "\"f32\"",
         "wScalar": "\"f32\"",
         "bScalar": "\"f32\"",
         "outFeatures": "nDim",
         "tileM": 32,
         "tileN": 64,
         "directAStorage": false,
         "doubleBufferedB": false,
         "balancedMatrixAccumulation": false
       "passes": [
         {
           "id": "main",
+          "name": "GemmFastGelu.Sgmat",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", "y", "params"],
           "dispatch": {
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
+      ],
+      "demoteWhen": ["false"]
     },
     {
       "id": "sgmat_bias_f16",
       },
       "derive": {
         "hasBias": true,
         "xScalar": "\"f16\"",
         "wScalar": "\"f16\"",
         "bScalar": "\"f16\"",
         "outFeatures": "nDim",
         "tileM": 32,
         "tileN": 64,
         "directAStorage": false,
         "doubleBufferedB": false,
         "balancedMatrixAccumulation": false
       "passes": [
         {
           "id": "main",
+          "name": "GemmFastGelu.SgmatBiasF16",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", { "arg": "bias", "elementType": "$bScalar" }, "y", "params"],
           "dispatch": {
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
+      ],
+      "demoteWhen": ["false"]
     },
     {
       "id": "sgmat_f16",
       },
       "derive": {
         "hasBias": false,
         "xScalar": "\"f16\"",
         "wScalar": "\"f16\"",
         "bScalar": "\"f16\"",
         "outFeatures": "nDim",
         "tileM": 32,
         "tileN": 64,
         "directAStorage": false,
         "doubleBufferedB": false,
         "balancedMatrixAccumulation": false
       "passes": [
         {
           "id": "main",
+          "name": "GemmFastGelu.SgmatF16",
           "shader": "gemm-subgroup-matrix.wgsl.jinja",
           "bindings": ["x", "w", "y", "params"],
           "dispatch": {
             "z": "ceilDiv(ceilDiv(gfgRows, tileM), DISPATCH_FOLD_WIDTH)"
           }
         }
+      ],
+      "demoteWhen": ["false"]
     },
     {
       "id": "tiled_bias",
         {
           "id": "main",
           "name": "GemmFastGelu.Bias",
+          "shader": "gemm-epilogue-tiled-reg.wgsl.jinja",
           "bindings": [
             "a",
             "b",
             "bias",
+            "y_t",
             {
               "name": "params",
               "struct": [
             }
           ],
           "dispatch": {
+            "x": "min(ceilDiv(nDim, gemmTileCols), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": "min(ceilDiv(gfgRows, gemmTileRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))"
           }
         }
       ]
         {
           "id": "main",
           "name": "GemmFastGelu",
+          "shader": "gemm-epilogue-tiled-reg.wgsl.jinja",
           "bindings": [
             "a",
             "b",
+            "y_t",
             {
               "name": "params",
               "struct": [
             }
           ],
           "dispatch": {
+            "x": "min(ceilDiv(nDim, gemmTileCols), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))",
+            "y": "min(ceilDiv(gfgRows, gemmTileRows), min(device.limits.maxComputeWorkgroupsPerDimension, 65535))"
           }
         }
       ]
     }
+  ],
+  "tunables": {
+    "SMALL_MATRIX_MAX_GROUPS": { "default": 48 },
+    "MEDIUM_MATRIX_MAX_GROUPS_F32": { "default": 96 },
+    "MEDIUM_MATRIX_MAX_GROUPS_F16": { "default": 80 },
+    "GEMV_TARGET_BLOCKS": { "default": 512 },
+    "PORTABLE_SMALL_TILE_GROUPS_F32": { "default": 64 },
+    "PORTABLE_SMALL_TILE_GROUPS_F16": { "default": 128 }
+  }
 }

build/webgpu/matmul-vector-matrix-vec4.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,70 @@

+// GEMV specialization for y[N] = a[K] @ B[K, N]. Each workgroup owns 32
+// consecutive vec4 column groups, or 128 output columns.
+{{ env.wgsl.resourceDeclarations }}
+{% set OUT = "y" %}
+fn tanh_safe(x: f32) -> f32 {
+  // tanh rounds to its saturated value for these tails in f32. Return that
+  // value directly, including for infinite input, before invoking the builtin.
+  if (x > 10.0) { return 1.0; }
+  if (x < -10.0) { return -1.0; }
+  // For tiny |x|, return x directly to preserve its sign and magnitude without
+  // relying on backend-specific builtin behavior near zero.
+  if (x > -1.0e-4 && x < 1.0e-4) { return x; }
+  return tanh(x);
+}
+fn gelu_tanh(v: f32) -> f32 {
+  return 0.5 * v * (1.0 + tanh_safe(0.7978845608028654 * (v + 0.044715 * v * v * v)));
+}
+const LANES: u32 = {{ gemvLanes }}u;
+// SLICES partitions the K reduction across the workgroup's second dimension.
+// Thread zero of each column group combines the slice partials in index order.
+const SLICES: u32 = {{ gemvSlices }}u;
+var<workgroup> partials: array<vec4<f32>, LANES * SLICES>;
+@compute @workgroup_size({{ gemvLanes }}, {{ gemvSlices }}, 1)
+fn main(
+  @builtin(workgroup_id) workgroup_id: vec3<u32>,
+  @builtin(local_invocation_id) lid: vec3<u32>
+) {
+  let lane = lid.x;
+  let slice = lid.y;
+  let cg = workgroup_id.x * LANES + lane;
+  var acc = vec4<f32>(0.0);
+  if (cg < params.N4) {
+{% if unrollK2 is defined and unrollK2 %}
+    // Two K positions per iteration amortize loop/address arithmetic on long
+    // decode projections while preserving each slice's exact strided order.
+    var k = slice;
+    for (; k + SLICES < params.K; k = k + 2u * SLICES) {
+      acc = acc + f32(a[k]) * vec4<f32>(b[k * params.N4 + cg]);
+      let k1 = k + SLICES;
+      acc = acc + f32(a[k1]) * vec4<f32>(b[k1 * params.N4 + cg]);
+    }
+    if (k < params.K) {
+      acc = acc + f32(a[k]) * vec4<f32>(b[k * params.N4 + cg]);
+    }
+{% else %}
+    for (var k = slice; k < params.K; k = k + SLICES) {
+      acc = acc + f32(a[k]) * vec4<f32>(b[k * params.N4 + cg]);
+    }
+{% endif %}
+  }
+  partials[slice * LANES + lane] = acc;
+  workgroupBarrier();
+  if (slice == 0u && cg < params.N4) {
+    var total = partials[lane];
+    for (var s = 1u; s < SLICES; s = s + 1u) {
+      total = total + partials[s * LANES + lane];
+    }
+{% if gemvBias | default(false) %}
+    total = total + vec4<f32>(bias[cg]);
+{% endif %}
+{% for i in range(4) %}
+    total[{{ i }}] = gelu_tanh(total[{{ i }}]);
+{% endfor %}
+    {{ OUT }}[cg] = vec4<{{ T }}>(total);
+  }
+}

build/webgpu/metadata.json CHANGED Viewed

@@ -1,23 +1,34 @@
 {
   "name": "com.microsoft.GemmFastGelu",
-  "id": "_com_microsoft_gemmfastgelu_webgpu_a713abb",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "oqlnhXJEZ8273ZAZC7ET+BMVMM8J84BrfZAVpaboErU=",
-      "gemm-fast-gelu.wgsl.jinja": "qh+CurhpiRX3buxeBnrOX2J2mzPrlqMP+m++faB5SfI=",
-      "gemm-subgroup-matrix.wgsl.jinja": "aaVCZdPC5akSTqDdYvM/MP8NnncAcK9kwSbvrd4Bvj8=",
-      "manifest.json": "EespK6n0UbB6+9TwNvJ40Pp/lYYHsHCO0BCKXxLRZhM=",
-      "test.json": "LFwbPPwea77gPJaoCjDJtA/uWfmqDYMUtIoDxd8iphY="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "sgmat_direct_bias": ["gemm-subgroup-matrix.wgsl.jinja"],
       "sgmat_direct": ["gemm-subgroup-matrix.wgsl.jinja"],
       "sgmat_direct_bias_f16": ["gemm-subgroup-matrix.wgsl.jinja"],
@@ -26,8 +37,8 @@
       "sgmat": ["gemm-subgroup-matrix.wgsl.jinja"],
       "sgmat_bias_f16": ["gemm-subgroup-matrix.wgsl.jinja"],
       "sgmat_f16": ["gemm-subgroup-matrix.wgsl.jinja"],
-      "tiled_bias": ["gemm-fast-gelu.wgsl.jinja"],
-      "tiled": ["gemm-fast-gelu.wgsl.jinja"]
     }
   }
 }

 {
   "name": "com.microsoft.GemmFastGelu",
+  "id": "_com_microsoft_gemmfastgelu_webgpu_f5e93e1",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "FfMbeBQWBZxpDttjltTuzAV+j6kE32ZmQfmcYdMKHQk=",
+      "gemm-epilogue-tiled-reg.wgsl.jinja": "R5ytrlOP+b62VdPVNEBfcMFde6AXx522owT2aH53EsQ=",
+      "gemm-subgroup-matrix.wgsl.jinja": "iDyC5cRpAczYx/bZ4E9ye5mFljkK0mlPhB0ZMAaLjRg=",
+      "manifest.json": "1QpiQuqS1C0R7nLKk7zC6wA9zt9GPnOdYFZ1+xNlE/U=",
+      "matmul-vector-matrix-vec4.wgsl.jinja": "uU7qRZ2/JhR5TZ4kZEeaf2pg1ZGqTvpO5VG1r2ft00E=",
+      "test.json": "IB4HCRTq8pos9TIr+Oy8FeeM8GCWURJ+d85QJYDv+bc="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
+      "vector_matrix_bias": ["matmul-vector-matrix-vec4.wgsl.jinja"],
+      "vector_matrix": ["matmul-vector-matrix-vec4.wgsl.jinja"],
+      "sgmat_medium_bias": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "sgmat_medium": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "sgmat_medium_bias_f16": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "sgmat_medium_f16": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "sgmat_small_bias": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "sgmat_small": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "sgmat_small_bias_f16": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "sgmat_small_f16": ["gemm-subgroup-matrix.wgsl.jinja"],
       "sgmat_direct_bias": ["gemm-subgroup-matrix.wgsl.jinja"],
       "sgmat_direct": ["gemm-subgroup-matrix.wgsl.jinja"],
       "sgmat_direct_bias_f16": ["gemm-subgroup-matrix.wgsl.jinja"],
       "sgmat": ["gemm-subgroup-matrix.wgsl.jinja"],
       "sgmat_bias_f16": ["gemm-subgroup-matrix.wgsl.jinja"],
       "sgmat_f16": ["gemm-subgroup-matrix.wgsl.jinja"],
+      "tiled_bias": ["gemm-epilogue-tiled-reg.wgsl.jinja"],
+      "tiled": ["gemm-epilogue-tiled-reg.wgsl.jinja"]
     }
   }
 }

build/webgpu/test.json CHANGED Viewed

The diff for this file is too large to render. See raw diff