sync 6fdf6301e2bb

Browse files

Files changed (7) hide show

README.md +2 -2
build/webgpu/bench.json +18 -0
build/webgpu/manifest.json +114 -148
build/webgpu/matmul-nbits-fused-rms-norm.wgsl.jinja +5 -41
build/webgpu/metadata.json +8 -8
build/webgpu/mlp-gate-up.wgsl.jinja +47 -94
build/webgpu/test.json +49 -1

README.md CHANGED Viewed

@@ -62,14 +62,14 @@ Attributes and default values (overridable per request):
 - [`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
 - [`matmul-nbits-fused-rms-norm.wgsl.jinja`](build/webgpu/matmul-nbits-fused-rms-norm.wgsl.jinja)
 - [`mlp-gate-up.wgsl.jinja`](build/webgpu/mlp-gate-up.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.

 - [`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
 - [`matmul-nbits-fused-rms-norm.wgsl.jinja`](build/webgpu/matmul-nbits-fused-rms-norm.wgsl.jinja)
 - [`mlp-gate-up.wgsl.jinja`](build/webgpu/mlp-gate-up.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.

build/webgpu/bench.json CHANGED Viewed

@@ -18,6 +18,24 @@
       "outputs": { "yT": { "shape": [1, 5632], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 5632 * 64 * 16" }] }
     },
     {
       "name": "mlp-q4-prefill-m64-k2048-n5632",
       "preset": "model",

       "outputs": { "yT": { "shape": [1, 5632], "dtype": "float32" } },
       "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 5632 * 64 * 16" }] }
     },
+    {
+      "name": "mlp-q4-b16-skip-decode-k2048-n5632",
+      "tunableSpace": {},
+      "preset": "smoke",
+      "vars": { "dtype": "float32" },
+      "attrs": { "K": 2048, "N": 5632, "bits": 4, "block_size": 16, "activation": "silu" },
+      "inputs": {
+        "aT": { "shape": [1, 2048], "dtype": "float32", "dist": "normal", "seed": 8121, "scale": 1 },
+        "skipT": { "shape": [1, 2048], "dtype": "float32", "dist": "normal", "seed": 8122, "scale": 1 },
+        "normScaleT": { "shape": [2048], "dtype": "float32", "dist": "normal", "seed": 8123, "scale": 1 },
+        "gateBT": { "shape": [5632, 128, 8], "dtype": "uint8", "dist": "uniform", "seed": 8124, "min": 0, "max": 256 },
+        "gateScalesT": { "shape": [5632, 128], "dtype": "float32", "dist": "normal", "seed": 8125, "scale": 0.05 },
+        "upBT": { "shape": [5632, 128, 8], "dtype": "uint8", "dist": "uniform", "seed": 8126, "min": 0, "max": 256 },
+        "upScalesT": { "shape": [5632, 128], "dtype": "float32", "dist": "normal", "seed": 8127, "scale": 0.05 }
+      },
+      "outputs": { "yT": { "shape": [1, 5632], "dtype": "float32" } },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 5632 * 128 * 8" }] }
+    },
     {
       "name": "mlp-q4-prefill-m64-k2048-n5632",
       "preset": "model",

build/webgpu/manifest.json CHANGED Viewed

@@ -50,92 +50,93 @@
   },
   "derive": {
     "aRows": "numel(shapes.aT) / max(1, attrs.K)",
-    "rowTilePlan": "1 if aRows <= 1 else min(aRows, tunables.ROW_TILE)",
-    "rowGroups": "ceilDiv(aRows, rowTilePlan)",
-    "kBlocks": "dim(shapes.gateBT, 1)",
-    "blobSize": "dim(shapes.gateBT, 2)",
     "codesPerByte": "8 / attrs.bits",
     "codeMask": "3 if attrs.bits == 2 else (15 if attrs.bits == 4 else 255)",
     "epsilonValue": "attrs.epsilon",
     "bitsSupported": "attrs.bits == 2 or attrs.bits == 4 or attrs.bits == 8",
     "weightShapeOk": "ranks.gateBT == 3 and ranks.upBT == 3 and dim(shapes.gateBT, 0) == attrs.N and dim(shapes.upBT, 0) == attrs.N and dim(shapes.upBT, 1) == kBlocks and dim(shapes.upBT, 2) == blobSize and kBlocks == ceilDiv(attrs.K, attrs.block_size) and blobSize * 8 == attrs.block_size * attrs.bits",
     "scaleShapeOk": "ranks.gateScalesT == 2 and ranks.upScalesT == 2 and dim(shapes.gateScalesT, 0) == attrs.N and dim(shapes.gateScalesT, 1) == kBlocks and dim(shapes.upScalesT, 0) == attrs.N and dim(shapes.upScalesT, 1) == kBlocks",
     "ioShapeOk": "(ranks.aT == 2 or ranks.aT == 3) and dim(shapes.aT, ranks.aT - 1) == attrs.K and ranks.yT == ranks.aT and dim(shapes.yT, ranks.yT - 1) == attrs.N and sameShape(prefix(shapes.yT, ranks.yT - 1), prefix(shapes.aT, ranks.aT - 1))",
     "biasShapeOk": "(ranks.gateBiasT == 1 and dim(shapes.gateBiasT, 0) == attrs.N if present.gateBiasT else true) and (ranks.upBiasT == 1 and dim(shapes.upBiasT, 0) == attrs.N if present.upBiasT else true)",
     "dtypeOk": "tensorDtypes.gateScalesT == tensorDtypes.aT and tensorDtypes.upScalesT == tensorDtypes.aT and tensorDtypes.yT == tensorDtypes.aT and f16Ok(tensorDtypes.aT)",
-    "lanesPow2": "tunables.LANES == pow2ceil(tunables.LANES)",
     "mlpShapeOk": "bitsSupported and weightShapeOk and scaleShapeOk and ioShapeOk and biasShapeOk and dtypeOk and lanesPow2 and attrs.K > 0 and attrs.N > 0 and attrs.block_size > 0",
     "normContractOk": "present.normScaleT and ranks.normScaleT == 1 and dim(shapes.normScaleT, 0) == attrs.K and tensorDtypes.normScaleT == tensorDtypes.aT and (sameShape(shapes.skipT, shapes.aT) and tensorDtypes.skipT == tensorDtypes.aT if present.skipT else true) and (sameShape(shapes.residualT, shapes.aT) and tensorDtypes.residualT == tensorDtypes.aT and present.skipT if present.residualT else true)",
-    "decodeWalk": "aRows <= 1",
     "decodeVec": "decodeWalk and blobSize % 16 == 0",
     "decodeActVec4": "decodeVec and attrs.K % attrs.block_size == 0",
     "decodeLaneSplit": "decodeVec and kBlocks * blobSize <= tunables.DECODE_WORKGROUP_SIZE * 16",
     "decodeCols": "8 if decodeLaneSplit else 4",
     "decodeLanes": "tunables.DECODE_WORKGROUP_SIZE * (2 if decodeLaneSplit else 1)",
-    "tileCols": "decodeCols if decodeWalk else tunables.TILE_N",
     "decodeWorkgroupOk": "tunables.DECODE_WORKGROUP_SIZE >= 4 and pow2ceil(tunables.DECODE_WORKGROUP_SIZE) == tunables.DECODE_WORKGROUP_SIZE and decodeLanes <= device.limits.maxComputeInvocationsPerWorkgroup and decodeLanes <= device.limits.maxComputeWorkgroupSizeX",
-    "gateUpDispatchFits": "decodeWorkgroupOk and ceilDiv(attrs.N, tileCols) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and aRows <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and tunables.TILE_N * tunables.LANES <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.TILE_N * tunables.LANES <= device.limits.maxComputeWorkgroupSizeX",
     "normDispatchFits": "tunables.NORM_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.NORM_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX",
     "biasPresence_nogb_noub": "not present.gateBiasT and not present.upBiasT",
     "biasPresence_nogb_ub": "not present.gateBiasT and present.upBiasT",
     "biasPresence_gb_noub": "present.gateBiasT and not present.upBiasT",
     "biasPresence_gb_ub": "present.gateBiasT and present.upBiasT",
-    "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-    "scalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
-    "K": "attrs.K",
     "N": "attrs.N",
-    "blockSize": "attrs.block_size",
-    "bits": "attrs.bits",
     "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
-    "tileN": "tunables.TILE_N",
-    "lanes": "tunables.LANES",
-    "rowTile": "rowTilePlan",
-    "rowCount": "aRows",
     "decodeNCols": "decodeCols",
     "decodeWorkgroupSize": "decodeLanes",
     "laneGroups": "2 if decodeLaneSplit else 1",
-    "useSubgroups": "device.features.has(\"subgroups\")",
     "weightElement": "\"vec4<u32>\" if decodeVec else \"u32\"",
     "actVec4": "decodeActVec4",
     "normedElement": "\"vec4<f32>\" if decodeActVec4 else \"f32\"",
-    "hidden": "attrs.K",
-    "workgroupSize": "tunables.NORM_WORKGROUP_SIZE",
-    "epsilon": "epsilonValue",
     "hasGateBias": "present.gateBiasT",
     "hasUpBias": "present.upBiasT",
-    "hasSkip": "present.skipT",
-    "writeResidual": "present.residualT",
-    "K_LEN": "attrs.K",
     "N_LEN": "attrs.N"
   },
   "when": ["mlpShapeOk", "gateUpDispatchFits"],
   "bindings": {
-    "a": { "arg": "aT", "buffer": "read-only-storage", "elementType": "$aScalar" },
-    "gate_b": { "arg": "gateBT", "buffer": "read-only-storage", "elementType": "$weightElement" },
-    "gate_scales": { "arg": "gateScalesT", "buffer": "read-only-storage", "elementType": "$aScalar" },
-    "up_b": { "arg": "upBT", "buffer": "read-only-storage", "elementType": "$weightElement" },
-    "up_scales": { "arg": "upScalesT", "buffer": "read-only-storage", "elementType": "$aScalar" },
-    "y": { "arg": "yT", "buffer": "storage", "elementType": "$aScalar" },
-    "up_bias": { "arg": "upBiasT", "buffer": "read-only-storage", "elementType": "$aScalar", "length": "$N_LEN" },
-    "gate_bias": { "arg": "gateBiasT", "buffer": "read-only-storage", "elementType": "$aScalar", "length": "$N_LEN" },
-    "norm_scale": { "arg": "normScaleT", "buffer": "read-only-storage", "elementType": "$aScalar", "length": "$K_LEN" },
-    "normed": { "scratch": "normedA", "buffer": "storage", "elementType": "f32" },
-    "params": { "buffer": "uniform", "struct": [{ "name": "rows", "type": "u32", "value": "aRows" }] },
-    "normed_2": {
       "scratch": "normedA",
       "name": "normed",
       "buffer": "read-only-storage",
       "elementType": "$normedElement"
     },
-    "skip": { "arg": "skipT", "buffer": "read-only-storage", "elementType": "$aScalar" },
-    "residual": { "arg": "residualT", "buffer": "storage", "elementType": "$aScalar" }
   },
   "variants": [
     {
       "id": "plain_nogb_noub",
       "priority": 10,
       "when": ["not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_nogb_noub"],
-      "derive": { "inlineNorm": "0", "fromNormed": "0", "rowTile": "rowTilePlan" },
       "passes": [
         {
           "id": "main",
@@ -150,7 +151,7 @@
       "id": "plain_nogb_ub",
       "priority": 10,
       "when": ["not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_nogb_ub"],
-      "derive": { "inlineNorm": "0", "fromNormed": "0", "rowTile": "rowTilePlan" },
       "passes": [
         {
           "id": "main",
@@ -165,7 +166,7 @@
       "id": "plain_gb_noub",
       "priority": 10,
       "when": ["not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_gb_noub"],
-      "derive": { "inlineNorm": "0", "fromNormed": "0", "rowTile": "rowTilePlan" },
       "passes": [
         {
           "id": "main",
@@ -180,23 +181,22 @@
       "id": "plain_gb_ub",
       "priority": 10,
       "when": ["not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_gb_ub"],
-      "derive": { "inlineNorm": "0", "fromNormed": "0", "rowTile": "rowTilePlan" },
       "passes": [
         {
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "staged_norm_nogb_noub",
       "priority": 10,
-      "when": ["normContractOk", "biasPresence_nogb_noub", "normDispatchFits", "not present.skipT", "not present.residualT"],
-      "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
@@ -210,17 +210,16 @@
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
-          "bindings": ["normed_2", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "staged_skip_nogb_noub",
       "priority": 10,
-      "when": ["normContractOk", "biasPresence_nogb_noub", "normDispatchFits", "present.skipT", "not present.residualT"],
-      "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
@@ -234,17 +233,16 @@
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
-          "bindings": ["normed_2", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "staged_skipsum_nogb_noub",
       "priority": 10,
-      "when": ["normContractOk", "biasPresence_nogb_noub", "normDispatchFits", "present.skipT", "present.residualT"],
-      "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
@@ -258,17 +256,16 @@
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
-          "bindings": ["normed_2", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "staged_norm_nogb_ub",
       "priority": 10,
-      "when": ["normContractOk", "biasPresence_nogb_ub", "normDispatchFits", "not present.skipT", "not present.residualT"],
-      "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
@@ -282,17 +279,16 @@
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
-          "bindings": ["normed_2", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "staged_skip_nogb_ub",
       "priority": 10,
-      "when": ["normContractOk", "biasPresence_nogb_ub", "normDispatchFits", "present.skipT", "not present.residualT"],
-      "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
@@ -306,17 +302,16 @@
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
-          "bindings": ["normed_2", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "staged_skipsum_nogb_ub",
       "priority": 10,
-      "when": ["normContractOk", "biasPresence_nogb_ub", "normDispatchFits", "present.skipT", "present.residualT"],
-      "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
@@ -330,17 +325,16 @@
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
-          "bindings": ["normed_2", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "staged_norm_gb_noub",
       "priority": 10,
-      "when": ["normContractOk", "biasPresence_gb_noub", "normDispatchFits", "not present.skipT", "not present.residualT"],
-      "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
@@ -354,17 +348,16 @@
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
-          "bindings": ["normed_2", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "staged_skip_gb_noub",
       "priority": 10,
-      "when": ["normContractOk", "biasPresence_gb_noub", "normDispatchFits", "present.skipT", "not present.residualT"],
-      "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
@@ -378,17 +371,16 @@
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
-          "bindings": ["normed_2", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "staged_skipsum_gb_noub",
       "priority": 10,
-      "when": ["normContractOk", "biasPresence_gb_noub", "normDispatchFits", "present.skipT", "present.residualT"],
-      "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
@@ -402,17 +394,16 @@
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
-          "bindings": ["normed_2", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "staged_norm_gb_ub",
       "priority": 10,
-      "when": ["normContractOk", "biasPresence_gb_ub", "normDispatchFits", "not present.skipT", "not present.residualT"],
-      "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
@@ -426,17 +417,16 @@
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
-          "bindings": ["normed_2", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "staged_skip_gb_ub",
       "priority": 10,
-      "when": ["normContractOk", "biasPresence_gb_ub", "normDispatchFits", "present.skipT", "not present.residualT"],
-      "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
@@ -450,17 +440,16 @@
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
-          "bindings": ["normed_2", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "staged_skipsum_gb_ub",
       "priority": 10,
-      "when": ["normContractOk", "biasPresence_gb_ub", "normDispatchFits", "present.skipT", "present.residualT"],
-      "derive": { "inlineNorm": "0", "fromNormed": "1", "rowTile": "rowTilePlan" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
@@ -474,17 +463,15 @@
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
-          "bindings": ["normed_2", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "fused_norm_nogb_noub",
       "priority": 30,
-      "when": ["normContractOk", "biasPresence_nogb_noub", "aRows == 1", "not decodeVec", "not present.skipT", "not present.residualT"],
-      "requires": { "limits": { "maxStorageBuffersPerShaderStage": 7 } },
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
@@ -492,16 +479,14 @@
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "fused_skip_nogb_noub",
       "priority": 30,
-      "when": ["normContractOk", "biasPresence_nogb_noub", "aRows == 1", "not decodeVec", "present.skipT", "not present.residualT"],
-      "requires": { "limits": { "maxStorageBuffersPerShaderStage": 8 } },
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
@@ -509,16 +494,14 @@
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "fused_skipsum_nogb_noub",
       "priority": 30,
-      "when": ["normContractOk", "biasPresence_nogb_noub", "aRows == 1", "not decodeVec", "present.skipT", "present.residualT"],
-      "requires": { "limits": { "maxStorageBuffersPerShaderStage": 9 } },
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
@@ -526,16 +509,14 @@
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "y", "residual"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "fused_norm_nogb_ub",
       "priority": 30,
-      "when": ["normContractOk", "biasPresence_nogb_ub", "aRows == 1", "not decodeVec", "not present.skipT", "not present.residualT"],
-      "requires": { "limits": { "maxStorageBuffersPerShaderStage": 8 } },
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
@@ -543,16 +524,14 @@
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "fused_skip_nogb_ub",
       "priority": 30,
-      "when": ["normContractOk", "biasPresence_nogb_ub", "aRows == 1", "not decodeVec", "present.skipT", "not present.residualT"],
-      "requires": { "limits": { "maxStorageBuffersPerShaderStage": 9 } },
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
@@ -560,16 +539,14 @@
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "fused_skipsum_nogb_ub",
       "priority": 30,
-      "when": ["normContractOk", "biasPresence_nogb_ub", "aRows == 1", "not decodeVec", "present.skipT", "present.residualT"],
-      "requires": { "limits": { "maxStorageBuffersPerShaderStage": 10 } },
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
@@ -577,16 +554,14 @@
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y", "residual"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "fused_norm_gb_noub",
       "priority": 30,
-      "when": ["normContractOk", "biasPresence_gb_noub", "aRows == 1", "not decodeVec", "not present.skipT", "not present.residualT"],
-      "requires": { "limits": { "maxStorageBuffersPerShaderStage": 8 } },
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
@@ -594,16 +569,14 @@
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "fused_skip_gb_noub",
       "priority": 30,
-      "when": ["normContractOk", "biasPresence_gb_noub", "aRows == 1", "not decodeVec", "present.skipT", "not present.residualT"],
-      "requires": { "limits": { "maxStorageBuffersPerShaderStage": 9 } },
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
@@ -611,16 +584,14 @@
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "fused_skipsum_gb_noub",
       "priority": 30,
-      "when": ["normContractOk", "biasPresence_gb_noub", "aRows == 1", "not decodeVec", "present.skipT", "present.residualT"],
-      "requires": { "limits": { "maxStorageBuffersPerShaderStage": 10 } },
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
@@ -628,16 +599,14 @@
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y", "residual"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "fused_norm_gb_ub",
       "priority": 30,
-      "when": ["normContractOk", "biasPresence_gb_ub", "aRows == 1", "not decodeVec", "not present.skipT", "not present.residualT"],
-      "requires": { "limits": { "maxStorageBuffersPerShaderStage": 9 } },
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
@@ -645,16 +614,14 @@
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     },
     {
       "id": "fused_skip_gb_ub",
       "priority": 30,
-      "when": ["normContractOk", "biasPresence_gb_ub", "aRows == 1", "not decodeVec", "present.skipT", "not present.residualT"],
-      "requires": { "limits": { "maxStorageBuffersPerShaderStage": 10 } },
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
@@ -662,8 +629,7 @@
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
-          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" },
-          "subgroupCollectivesWidth": "portable"
         }
       ]
     }

   },
   "derive": {
     "aRows": "numel(shapes.aT) / max(1, attrs.K)",
     "codesPerByte": "8 / attrs.bits",
     "codeMask": "3 if attrs.bits == 2 else (15 if attrs.bits == 4 else 255)",
     "epsilonValue": "attrs.epsilon",
+    "lanesPow2": "tunables.LANES == pow2ceil(tunables.LANES)",
+    "decodeWalk": "aRows <= 1",
+    "aScalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
+    "scalar": "\"f16\" if tensorDtypes.aT == \"float16\" else \"f32\"",
+    "K": "attrs.K",
+    "blockSize": "attrs.block_size",
+    "bits": "attrs.bits",
+    "lanes": "tunables.LANES",
+    "rowCount": "aRows",
+    "useSubgroups": "device.features.has(\"subgroups\")",
+    "hidden": "attrs.K",
+    "workgroupSize": "tunables.NORM_WORKGROUP_SIZE",
+    "epsilon": "epsilonValue",
+    "hasSkip": "present.skipT",
+    "writeResidual": "present.residualT",
+    "K_LEN": "attrs.K",
+    "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",
+    "portableSmallTile": "variableSubgroup16To32 and tensorDtypes.aT == \"float32\" and attrs.bits == 4 and aRows >= 2 * tunables.ROW_TILE and tunables.TILE_N >= 4 and tunables.LANES >= 1 and floor(tunables.TILE_N / 2) * tunables.LANES <= device.limits.maxComputeInvocationsPerWorkgroup and floor(tunables.TILE_N / 2) * tunables.LANES <= device.limits.maxComputeWorkgroupSizeX and ceilDiv(attrs.N, floor(tunables.TILE_N / 2)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and device.limits.maxComputeWorkgroupStorageSize >= 8 * floor(tunables.TILE_N / 2) * tunables.LANES * 2 * tunables.ROW_TILE",
+    "tileN": "floor(tunables.TILE_N / 2) if portableSmallTile else tunables.TILE_N",
+    "rowTile": "min(aRows, 2 * tunables.ROW_TILE, floor(device.limits.maxComputeWorkgroupStorageSize / (8 * tileN * tunables.LANES))) if portableSmallTile else (1 if aRows <= 1 else min(aRows, tunables.ROW_TILE))",
+    "rowGroups": "ceilDiv(aRows, rowTile)",
+    "kBlocks": "dim(shapes.gateBT, 1)",
+    "blobSize": "dim(shapes.gateBT, 2)",
     "bitsSupported": "attrs.bits == 2 or attrs.bits == 4 or attrs.bits == 8",
     "weightShapeOk": "ranks.gateBT == 3 and ranks.upBT == 3 and dim(shapes.gateBT, 0) == attrs.N and dim(shapes.upBT, 0) == attrs.N and dim(shapes.upBT, 1) == kBlocks and dim(shapes.upBT, 2) == blobSize and kBlocks == ceilDiv(attrs.K, attrs.block_size) and blobSize * 8 == attrs.block_size * attrs.bits",
     "scaleShapeOk": "ranks.gateScalesT == 2 and ranks.upScalesT == 2 and dim(shapes.gateScalesT, 0) == attrs.N and dim(shapes.gateScalesT, 1) == kBlocks and dim(shapes.upScalesT, 0) == attrs.N and dim(shapes.upScalesT, 1) == kBlocks",
     "ioShapeOk": "(ranks.aT == 2 or ranks.aT == 3) and dim(shapes.aT, ranks.aT - 1) == attrs.K and ranks.yT == ranks.aT and dim(shapes.yT, ranks.yT - 1) == attrs.N and sameShape(prefix(shapes.yT, ranks.yT - 1), prefix(shapes.aT, ranks.aT - 1))",
     "biasShapeOk": "(ranks.gateBiasT == 1 and dim(shapes.gateBiasT, 0) == attrs.N if present.gateBiasT else true) and (ranks.upBiasT == 1 and dim(shapes.upBiasT, 0) == attrs.N if present.upBiasT else true)",
     "dtypeOk": "tensorDtypes.gateScalesT == tensorDtypes.aT and tensorDtypes.upScalesT == tensorDtypes.aT and tensorDtypes.yT == tensorDtypes.aT and f16Ok(tensorDtypes.aT)",
     "mlpShapeOk": "bitsSupported and weightShapeOk and scaleShapeOk and ioShapeOk and biasShapeOk and dtypeOk and lanesPow2 and attrs.K > 0 and attrs.N > 0 and attrs.block_size > 0",
     "normContractOk": "present.normScaleT and ranks.normScaleT == 1 and dim(shapes.normScaleT, 0) == attrs.K and tensorDtypes.normScaleT == tensorDtypes.aT and (sameShape(shapes.skipT, shapes.aT) and tensorDtypes.skipT == tensorDtypes.aT if present.skipT else true) and (sameShape(shapes.residualT, shapes.aT) and tensorDtypes.residualT == tensorDtypes.aT and present.skipT if present.residualT else true)",
     "decodeVec": "decodeWalk and blobSize % 16 == 0",
     "decodeActVec4": "decodeVec and attrs.K % attrs.block_size == 0",
     "decodeLaneSplit": "decodeVec and kBlocks * blobSize <= tunables.DECODE_WORKGROUP_SIZE * 16",
     "decodeCols": "8 if decodeLaneSplit else 4",
     "decodeLanes": "tunables.DECODE_WORKGROUP_SIZE * (2 if decodeLaneSplit else 1)",
+    "tileCols": "decodeCols if decodeWalk else tileN",
     "decodeWorkgroupOk": "tunables.DECODE_WORKGROUP_SIZE >= 4 and pow2ceil(tunables.DECODE_WORKGROUP_SIZE) == tunables.DECODE_WORKGROUP_SIZE and decodeLanes <= device.limits.maxComputeInvocationsPerWorkgroup and decodeLanes <= device.limits.maxComputeWorkgroupSizeX",
+    "gateUpDispatchFits": "decodeWorkgroupOk and ceilDiv(attrs.N, tileCols) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and aRows <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and tileN * tunables.LANES <= device.limits.maxComputeInvocationsPerWorkgroup and tileN * tunables.LANES <= device.limits.maxComputeWorkgroupSizeX",
     "normDispatchFits": "tunables.NORM_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.NORM_WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX",
     "biasPresence_nogb_noub": "not present.gateBiasT and not present.upBiasT",
     "biasPresence_nogb_ub": "not present.gateBiasT and present.upBiasT",
     "biasPresence_gb_noub": "present.gateBiasT and not present.upBiasT",
     "biasPresence_gb_ub": "present.gateBiasT and present.upBiasT",
     "N": "attrs.N",
     "defaultZero": "\"2.0\" if attrs.bits == 2 else (\"8.0\" if attrs.bits == 4 else \"128.0\")",
     "decodeNCols": "decodeCols",
     "decodeWorkgroupSize": "decodeLanes",
     "laneGroups": "2 if decodeLaneSplit else 1",
     "weightElement": "\"vec4<u32>\" if decodeVec else \"u32\"",
     "actVec4": "decodeActVec4",
     "normedElement": "\"vec4<f32>\" if decodeActVec4 else \"f32\"",
     "hasGateBias": "present.gateBiasT",
     "hasUpBias": "present.upBiasT",
     "N_LEN": "attrs.N"
   },
   "when": ["mlpShapeOk", "gateUpDispatchFits"],
   "bindings": {
+    "a": { "arg": "aT", "elementType": "$aScalar" },
+    "gate_b": { "arg": "gateBT", "elementType": "$weightElement" },
+    "gate_scales": { "arg": "gateScalesT", "elementType": "$aScalar" },
+    "up_b": { "arg": "upBT", "elementType": "$weightElement" },
+    "up_scales": { "arg": "upScalesT", "elementType": "$aScalar" },
+    "y": { "arg": "yT", "elementType": "$aScalar" },
+    "up_bias": { "arg": "upBiasT", "elementType": "$aScalar", "length": "$N_LEN" },
+    "gate_bias": { "arg": "gateBiasT", "elementType": "$aScalar", "length": "$N_LEN" },
+    "norm_scale": { "arg": "normScaleT", "elementType": "$aScalar", "length": "$K_LEN" },
+    "normed": { "scratch": "normedA", "elementType": "f32" },
+    "params": { "struct": [{ "name": "rows", "type": "u32", "value": "aRows" }] },
+    "normed_main": {
       "scratch": "normedA",
       "name": "normed",
       "buffer": "read-only-storage",
       "elementType": "$normedElement"
     },
+    "skip": { "arg": "skipT", "elementType": "$aScalar" },
+    "residual": { "arg": "residualT", "elementType": "$aScalar" }
   },
   "variants": [
     {
       "id": "plain_nogb_noub",
       "priority": 10,
       "when": ["not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_nogb_noub"],
+      "derive": { "inlineNorm": "0", "fromNormed": "0" },
       "passes": [
         {
           "id": "main",
       "id": "plain_nogb_ub",
       "priority": 10,
       "when": ["not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_nogb_ub"],
+      "derive": { "inlineNorm": "0", "fromNormed": "0" },
       "passes": [
         {
           "id": "main",
       "id": "plain_gb_noub",
       "priority": 10,
       "when": ["not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_gb_noub"],
+      "derive": { "inlineNorm": "0", "fromNormed": "0" },
       "passes": [
         {
           "id": "main",
       "id": "plain_gb_ub",
       "priority": 10,
       "when": ["not present.normScaleT", "not present.skipT", "not present.residualT", "biasPresence_gb_ub"],
+      "derive": { "inlineNorm": "0", "fromNormed": "0" },
       "passes": [
         {
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
         }
       ]
     },
     {
       "id": "staged_norm_nogb_noub",
       "priority": 10,
+      "when": ["normContractOk", "biasPresence_nogb_noub", "normDispatchFits", "not hasSkip and not writeResidual"],
+      "derive": { "inlineNorm": "0", "fromNormed": "1" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
+          "bindings": ["normed_main", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
         }
       ]
     },
     {
       "id": "staged_skip_nogb_noub",
       "priority": 10,
+      "when": ["normContractOk", "biasPresence_nogb_noub", "normDispatchFits", "hasSkip and not writeResidual"],
+      "derive": { "inlineNorm": "0", "fromNormed": "1" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
+          "bindings": ["normed_main", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
         }
       ]
     },
     {
       "id": "staged_skipsum_nogb_noub",
       "priority": 10,
+      "when": ["normContractOk", "biasPresence_nogb_noub", "normDispatchFits", "hasSkip and writeResidual"],
+      "derive": { "inlineNorm": "0", "fromNormed": "1" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
+          "bindings": ["normed_main", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
         }
       ]
     },
     {
       "id": "staged_norm_nogb_ub",
       "priority": 10,
+      "when": ["normContractOk", "biasPresence_nogb_ub", "normDispatchFits", "not hasSkip and not writeResidual"],
+      "derive": { "inlineNorm": "0", "fromNormed": "1" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
+          "bindings": ["normed_main", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
         }
       ]
     },
     {
       "id": "staged_skip_nogb_ub",
       "priority": 10,
+      "when": ["normContractOk", "biasPresence_nogb_ub", "normDispatchFits", "hasSkip and not writeResidual"],
+      "derive": { "inlineNorm": "0", "fromNormed": "1" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
+          "bindings": ["normed_main", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
         }
       ]
     },
     {
       "id": "staged_skipsum_nogb_ub",
       "priority": 10,
+      "when": ["normContractOk", "biasPresence_nogb_ub", "normDispatchFits", "hasSkip and writeResidual"],
+      "derive": { "inlineNorm": "0", "fromNormed": "1" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
+          "bindings": ["normed_main", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
         }
       ]
     },
     {
       "id": "staged_norm_gb_noub",
       "priority": 10,
+      "when": ["normContractOk", "biasPresence_gb_noub", "normDispatchFits", "not hasSkip and not writeResidual"],
+      "derive": { "inlineNorm": "0", "fromNormed": "1" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
+          "bindings": ["normed_main", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
         }
       ]
     },
     {
       "id": "staged_skip_gb_noub",
       "priority": 10,
+      "when": ["normContractOk", "biasPresence_gb_noub", "normDispatchFits", "hasSkip and not writeResidual"],
+      "derive": { "inlineNorm": "0", "fromNormed": "1" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
+          "bindings": ["normed_main", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
         }
       ]
     },
     {
       "id": "staged_skipsum_gb_noub",
       "priority": 10,
+      "when": ["normContractOk", "biasPresence_gb_noub", "normDispatchFits", "hasSkip and writeResidual"],
+      "derive": { "inlineNorm": "0", "fromNormed": "1" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
+          "bindings": ["normed_main", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
         }
       ]
     },
     {
       "id": "staged_norm_gb_ub",
       "priority": 10,
+      "when": ["normContractOk", "biasPresence_gb_ub", "normDispatchFits", "not hasSkip and not writeResidual"],
+      "derive": { "inlineNorm": "0", "fromNormed": "1" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
+          "bindings": ["normed_main", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
         }
       ]
     },
     {
       "id": "staged_skip_gb_ub",
       "priority": 10,
+      "when": ["normContractOk", "biasPresence_gb_ub", "normDispatchFits", "hasSkip and not writeResidual"],
+      "derive": { "inlineNorm": "0", "fromNormed": "1" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
+          "bindings": ["normed_main", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
         }
       ]
     },
     {
       "id": "staged_skipsum_gb_ub",
       "priority": 10,
+      "when": ["normContractOk", "biasPresence_gb_ub", "normDispatchFits", "hasSkip and writeResidual"],
+      "derive": { "inlineNorm": "0", "fromNormed": "1" },
       "intermediates": [{ "id": "normedA", "dtype": "float32", "shape": "[numel(shapes.aT)]" }],
       "passes": [
         {
           "id": "main",
           "name": "MatMulNBitsMlp.GateUp",
           "shader": "mlp-gate-up.wgsl.jinja",
+          "bindings": ["normed_main", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "rowGroups" }
         }
       ]
     },
     {
       "id": "fused_norm_nogb_noub",
       "priority": 30,
+      "when": ["normContractOk", "biasPresence_nogb_noub", "aRows == 1", "not decodeVec", "not hasSkip and not writeResidual"],
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" }
         }
       ]
     },
     {
       "id": "fused_skip_nogb_noub",
       "priority": 30,
+      "when": ["normContractOk", "biasPresence_nogb_noub", "aRows == 1", "not decodeVec", "hasSkip and not writeResidual"],
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" }
         }
       ]
     },
     {
       "id": "fused_skipsum_nogb_noub",
       "priority": 30,
+      "when": ["normContractOk", "biasPresence_nogb_noub", "aRows == 1", "not decodeVec", "hasSkip and writeResidual"],
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "y", "residual"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" }
         }
       ]
     },
     {
       "id": "fused_norm_nogb_ub",
       "priority": 30,
+      "when": ["normContractOk", "biasPresence_nogb_ub", "aRows == 1", "not decodeVec", "not hasSkip and not writeResidual"],
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" }
         }
       ]
     },
     {
       "id": "fused_skip_nogb_ub",
       "priority": 30,
+      "when": ["normContractOk", "biasPresence_nogb_ub", "aRows == 1", "not decodeVec", "hasSkip and not writeResidual"],
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" }
         }
       ]
     },
     {
       "id": "fused_skipsum_nogb_ub",
       "priority": 30,
+      "when": ["normContractOk", "biasPresence_nogb_ub", "aRows == 1", "not decodeVec", "hasSkip and writeResidual"],
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "up_b", "up_scales", "up_bias", "y", "residual"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" }
         }
       ]
     },
     {
       "id": "fused_norm_gb_noub",
       "priority": 30,
+      "when": ["normContractOk", "biasPresence_gb_noub", "aRows == 1", "not decodeVec", "not hasSkip and not writeResidual"],
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" }
         }
       ]
     },
     {
       "id": "fused_skip_gb_noub",
       "priority": 30,
+      "when": ["normContractOk", "biasPresence_gb_noub", "aRows == 1", "not decodeVec", "hasSkip and not writeResidual"],
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" }
         }
       ]
     },
     {
       "id": "fused_skipsum_gb_noub",
       "priority": 30,
+      "when": ["normContractOk", "biasPresence_gb_noub", "aRows == 1", "not decodeVec", "hasSkip and writeResidual"],
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "y", "residual"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" }
         }
       ]
     },
     {
       "id": "fused_norm_gb_ub",
       "priority": 30,
+      "when": ["normContractOk", "biasPresence_gb_ub", "aRows == 1", "not decodeVec", "not hasSkip and not writeResidual"],
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" }
         }
       ]
     },
     {
       "id": "fused_skip_gb_ub",
       "priority": 30,
+      "when": ["normContractOk", "biasPresence_gb_ub", "aRows == 1", "not decodeVec", "hasSkip and not writeResidual"],
       "derive": { "inlineNorm": "1", "fromNormed": "0", "rowTile": "1" },
       "passes": [
         {
           "name": "MatMulNBitsMlp.FusedDecode",
           "shader": "mlp-gate-up.wgsl.jinja",
           "bindings": ["a", "skip", "norm_scale", "gate_b", "gate_scales", "gate_bias", "up_b", "up_scales", "up_bias", "y"],
+          "dispatch": { "x": "ceilDiv(attrs.N, tileCols)", "y": "aRows" }
         }
       ]
     }

build/webgpu/matmul-nbits-fused-rms-norm.wgsl.jinja CHANGED Viewed

@@ -12,68 +12,32 @@ const EPSILON: f32 = {{ epsilon }};
 var<workgroup> partial: array<f32, WG>;
 {% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
-{% if op == "max" %}
-{{ a }}[{{ idx }}] = max({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);
-{%- else %}
-{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] + {{ a }}[{{ idx }} + {{ svar }}];
-{%- endif %}
-{% endmacro %}
-{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false) %}
   var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
   loop {
-{% if form == "head" %}
-{% if breakInline %}
-    if ({{ svar }} == 0u) { break; }
-{% else %}
     if ({{ svar }} == 0u) {
       break;
     }
-{% endif %}
-{% endif %}
-{% if bodyInline %}
-    if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} }
-{% else %}
     if ({{ idx }} < {{ svar }}) {
 {% for a in arrays %}
       {{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
 {% endfor %}
     }
-{% endif %}
-{% if form == "head" %}
-{% if barrierFirst %}
-    workgroupBarrier();
     {{ svar }} = {{ svar }} / 2u;
-{% else %}
-    {{ svar }} = {{ svar }} / 2u;
-    workgroupBarrier();
-{% endif %}
-{% else %}
     workgroupBarrier();
-    if ({{ svar }} == 1u) {
-      break;
-    }
-    {{ svar }} = {{ svar }} / 2u;
-{% endif %}
-  }
-{%- endmacro %}
 // Reusing partial after this reduction requires a barrier between the read of
 // partial[0] and the next write, or the next round can race the prior readers.
-{% set trailingBarrier = trailingBarrier is defined and trailingBarrier %}
 fn reduce_sum(value: f32, tid: u32) -> f32 {
   partial[tid] = value;
   workgroupBarrier();
 {{ wgsl_tree_fold(["partial"], idx="tid", wg="WG", form="head") }}
-{% if trailingBarrier %}
-  let total = partial[0];
-  workgroupBarrier();
-  return total;
-{% else %}
   return partial[0];
-{% endif %}
 }
 fn row_value(index: u32) -> f32 {
 {% if hasSkip %}
   return f32(a[index]) + f32(skip[index]);

 var<workgroup> partial: array<f32, WG>;
 {% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
+{% if op == "max" or op == "min" %}
+{{ a }}[{{ idx }}] = {{ op }}({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);{% else %}
+{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] {{ "*" if op == "prod" else "+" }} {{ a }}[{{ idx }} + {{ svar }}];{% endif %}{% endmacro %}
+{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false, reuse=false) %}
   var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
   loop {
     if ({{ svar }} == 0u) {
       break;
     }
     if ({{ idx }} < {{ svar }}) {
 {% for a in arrays %}
       {{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
 {% endfor %}
     }
     {{ svar }} = {{ svar }} / 2u;
     workgroupBarrier();
+  }{% endmacro %}
 // Reusing partial after this reduction requires a barrier between the read of
 // partial[0] and the next write, or the next round can race the prior readers.
 fn reduce_sum(value: f32, tid: u32) -> f32 {
   partial[tid] = value;
   workgroupBarrier();
 {{ wgsl_tree_fold(["partial"], idx="tid", wg="WG", form="head") }}
   return partial[0];
 }
 fn row_value(index: u32) -> f32 {
 {% if hasSkip %}
   return f32(a[index]) + f32(skip[index]);

build/webgpu/metadata.json CHANGED Viewed

@@ -1,22 +1,22 @@
 {
   "name": "com.microsoft.MatMulNBitsMlp",
-  "id": "_com_microsoft_matmulnbitsmlp_webgpu_d711f87",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "teEo4C2oaUqxuIKr6wLZW6vFlmI9gTl/akkWJhjWxqg=",
-      "manifest.json": "ceiKDb7Itix3E2Ab3NY1WqE4vb5P3kKtrPxTrImqQQk=",
-      "matmul-nbits-fused-rms-norm.wgsl.jinja": "4lOdB+RprQh3iv29i6RV8UWkn8S5y1aK6Te5nJpxuEk=",
-      "mlp-gate-up.wgsl.jinja": "VuwljvMtS5vV09dkhewYJLe6s6Syr6zEtOTiGQMMukI=",
-      "test.json": "x1T4v2UMpwHHuiJcAq6Y8p4TkcyvZwz5pbpbV38m5D4="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "plain_nogb_noub": ["mlp-gate-up.wgsl.jinja"],
       "plain_nogb_ub": ["mlp-gate-up.wgsl.jinja"],

 {
   "name": "com.microsoft.MatMulNBitsMlp",
+  "id": "_com_microsoft_matmulnbitsmlp_webgpu_9ceab4a",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "ehJ4b7fUHWFgrEfT9QMCQJ+Lmm7CNvhClWNKMbmSgo0=",
+      "manifest.json": "wcISD9DIQ/KU8a1i5qDiwKVx1Tnk79JURhMQLb9w62s=",
+      "matmul-nbits-fused-rms-norm.wgsl.jinja": "lbzl6cUXSdORYxwAlSR40vtv2Qq+oNUsoJ7DAzgk0zg=",
+      "mlp-gate-up.wgsl.jinja": "AEO81iEkyXsDSnjuJdFYKcSmCnQXKFLWbtrOn1pWVlc=",
+      "test.json": "SKeEd86oERPUrGq2L1eafvM53Ft4NkoMTAwNuBb3BgQ="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "plain_nogb_noub": ["mlp-gate-up.wgsl.jinja"],
       "plain_nogb_ub": ["mlp-gate-up.wgsl.jinja"],

build/webgpu/mlp-gate-up.wgsl.jinja CHANGED Viewed

@@ -12,9 +12,7 @@ fn {{ fn }}(n: u32, block: u32, offset: u32) -> u32 {
   let byte_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + offset;
   return ({{ buffer }}[byte_index >> 2u] >> ((byte_index & 3u) * 8u)) & 255u;
 {% endif %}
-}
-{%- endmacro %}
 {% if rowTile == 1 and useSubgroups %}
 enable subgroups;
 {% endif %}
@@ -76,12 +74,13 @@ const CODES_PER_VEC: u32 = 16u * CODES_PER_BYTE;
 {% if inlineNorm %}
 const EPSILON: f32 = {{ epsilon }};
 {% endif %}
 {% for stream in ["gate", "up"] %}
 {% if not gemvWalk %}
 {{ matmul_nbits_packed_code(fn=stream ~ "_code", buffer=stream ~ "_b", kBlocks="KBLOCKS", blobSize="BLOB_SIZE", bits=bits) }}
 {% endif %}
 {% if not decodeVec %}
 // Decode two consecutive reduction-axis codes. Below 8 bits an even offset and
 // its successor share one stored byte; at 8 bits they occupy adjacent bytes of
 // one word, because an even byte index in a 4-byte-aligned blob never ends a
@@ -113,71 +112,7 @@ var<workgroup> red_gate: array<f32, WG * ROW_TILE>;
 var<workgroup> red_up: array<f32, WG * ROW_TILE>;
 {% endif %}
 {% if inlineNorm %}
-var<workgroup> partial: array<f32, WG>;
-var<workgroup> row_inv: f32;
-{% macro wgsl_tree_fold_stmt(a, op, idx, svar) %}
-{% if op == "max" %}
-{{ a }}[{{ idx }}] = max({{ a }}[{{ idx }}], {{ a }}[{{ idx }} + {{ svar }}]);
-{%- else %}
-{{ a }}[{{ idx }}] = {{ a }}[{{ idx }}] + {{ a }}[{{ idx }} + {{ svar }}];
-{%- endif %}
-{% endmacro %}
-{% macro wgsl_tree_fold(arrays, op="add", idx="lid", wg="WORKGROUP_SIZE", svar="stride", typed=false, form="tail", breakInline=false, bodyInline=false, barrierFirst=false) %}
-  var {{ svar }}{{ ": u32 " if typed else " " }}= {{ wg }} / 2u;
-  loop {
-{% if form == "head" %}
-{% if breakInline %}
-    if ({{ svar }} == 0u) { break; }
-{% else %}
-    if ({{ svar }} == 0u) {
-      break;
-    }
-{% endif %}
-{% endif %}
-{% if bodyInline %}
-    if ({{ idx }} < {{ svar }}) { {{ wgsl_tree_fold_stmt(arrays[0], op, idx, svar) }} }
-{% else %}
-    if ({{ idx }} < {{ svar }}) {
-{% for a in arrays %}
-      {{ wgsl_tree_fold_stmt(a, op, idx, svar) }}
-{% endfor %}
-    }
-{% endif %}
-{% if form == "head" %}
-{% if barrierFirst %}
-    workgroupBarrier();
-    {{ svar }} = {{ svar }} / 2u;
-{% else %}
-    {{ svar }} = {{ svar }} / 2u;
-    workgroupBarrier();
-{% endif %}
-{% else %}
-    workgroupBarrier();
-    if ({{ svar }} == 1u) {
-      break;
-    }
-    {{ svar }} = {{ svar }} / 2u;
-{% endif %}
-  }
-{%- endmacro %}
-// Reusing partial after this reduction requires a barrier between the read of
-// partial[0] and the next write, or the next round can race the prior readers.
-{% set trailingBarrier = trailingBarrier is defined and trailingBarrier %}
-fn reduce_sum(value: f32, tid: u32) -> f32 {
-  partial[tid] = value;
-  workgroupBarrier();
-{{ wgsl_tree_fold(["partial"], idx="tid", wg="WG", form="head") }}
-{% if trailingBarrier %}
-  let total = partial[0];
-  workgroupBarrier();
-  return total;
-{% else %}
-  return partial[0];
-{% endif %}
-}
 fn row_value(index: u32) -> f32 {
 {% if hasSkip %}
@@ -188,8 +123,7 @@ fn row_value(index: u32) -> f32 {
 }
 {% endif %}
-{% macro act(b, k) %}{% if inlineNorm %}row_value({{ b }} + {{ k }}) * row_inv * f32(norm_scale[{{ k }}]){% elif fromNormed %}normed[{{ b }} + {{ k }}]{% else %}f32(a[{{ b }} + {{ k }}]){% endif %}{%- endmacro %}
 {% if gemvWalk %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
@@ -200,20 +134,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   let row = min(wg.y, ROWS - 1u);
   let base_0 = row * K;
   let col_base = wg.x * N_COLS;
-{% if inlineNorm %}
-  var local_sq = 0.0;
-  for (var d = tid; d < K; d = d + WG) {
-    let value = row_value(base_0 + d);
-    local_sq = local_sq + value * value;
-  }
-  let inv = inverseSqrt(reduce_sum(local_sq, tid) / f32(K) + EPSILON);
-  if (tid == 0u) {
-    row_inv = inv;
-  }
-  // Separates the reduction's readers of partial[0] from the projection's
-  // reuse of the same workgroup array below.
-  workgroupBarrier();
-{% if writeResidual %}
   // Every N tile computes the same residual row; only the first one stores it,
   // so the tiles never write the same location.
   if (wg.x == 0u) {
@@ -222,9 +143,8 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
     }
   }
 {% endif %}
-{% endif %}
-  // Whole workgroups past the last column return here, after the normalization
-  // barriers above and before the reduction's.
   if (col_base >= N) {
     return;
   }
@@ -233,6 +153,9 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   var acc_gate{{ sfx }} = vec4<f32>(0.0);
   var acc_up{{ sfx }} = vec4<f32>(0.0);
 {% endfor %}
 {% if decodeVec %}
 {% set codesPerWord = 4 * codesPerByte %}
 {% if laneGroups == 2 %}
@@ -325,9 +248,18 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   for (var k = tid * 2u; k < K; k = k + WG * 2u) {
     let block = k / BLOCK_SIZE;
     let offset = k - block * BLOCK_SIZE;
     let v0 = {{ act("base_0", "k") }};
     // K need not be even; a code past the end contributes zero.
     let v1 = select(0.0, {{ act("base_0", "min(k + 1u, K - 1u)") }}, k + 1u < K);
 {% for c in range(4) %}
 {% set comp = ["x", "y", "z", "w"][c] %}
 {% if c == 0 %}
@@ -360,7 +292,13 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   let sgGate{{ sfx }} = subgroupAdd(acc_gate{{ sfx }});
   let sgUp{{ sfx }} = subgroupAdd(acc_up{{ sfx }});
 {% endfor %}
   if (sgLane == 0u) {
 {% for gi in range(colGroups) %}
 {% set sfx = "" if gi == 0 else gi %}
     red_gate{{ sfx }}[tid / sgSize] = sgGate{{ sfx }};
@@ -376,7 +314,13 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
     var gate_total{{ sfx }} = red_gate{{ sfx }}[0];
     var up_total{{ sfx }} = red_up{{ sfx }}[0];
 {% endfor %}
     for (var i = 1u; i < subgroupCount; i = i + 1u) {
 {% for gi in range(colGroups) %}
 {% set sfx = "" if gi == 0 else gi %}
       gate_total{{ sfx }} = gate_total{{ sfx }} + red_gate{{ sfx }}[i];
@@ -389,6 +333,9 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   red_gate{{ sfx }}[tid] = acc_gate{{ sfx }};
   red_up{{ sfx }}[tid] = acc_up{{ sfx }};
 {% endfor %}
   workgroupBarrier();
   var stride = WG / 2u;
   loop {
@@ -396,6 +343,9 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
       break;
     }
     if (tid < stride) {
 {% for gi in range(colGroups) %}
 {% set sfx = "" if gi == 0 else gi %}
       red_gate{{ sfx }}[tid] = red_gate{{ sfx }}[tid] + red_gate{{ sfx }}[tid + stride];
@@ -412,6 +362,12 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
     let gate_total{{ sfx }} = red_gate{{ sfx }}[0];
     let up_total{{ sfx }} = red_up{{ sfx }}[0];
 {% endfor %}
 {% endif %}
 {% for gi in range(colGroups) %}
 {% set sfx = "" if gi == 0 else gi %}
@@ -424,8 +380,8 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
     if (col_base + {{ i }}u < N) {
 {% endif %}
       let n = col_base + {{ i }}u;
-      var gate_value = gate_total{{ sfx }}.{{ comp }};
-      var up_value = up_total{{ sfx }}.{{ comp }};
 {% if hasGateBias %}
       gate_value = gate_value + f32(gate_bias[n]);
 {% endif %}
@@ -472,9 +428,7 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
 {% endfor %}
           }
 {% endif %}
-        }
-{%- endmacro %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
   let row0 = wg.y * ROW_TILE;
@@ -488,7 +442,6 @@ fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid:
   let base_{{ r }} = min(row0 + {{ r }}u, ROWS - 1u) * K;
 {% endfor %}
 {% for r in range(rowTile) %}
   var acc_gate_{{ r }} = 0.0;
   var acc_up_{{ r }} = 0.0;

   let byte_index = (n * {{ kBlocks }} + block) * {{ blobSize }} + offset;
   return ({{ buffer }}[byte_index >> 2u] >> ((byte_index & 3u) * 8u)) & 255u;
 {% endif %}
+}{% endmacro %}
 {% if rowTile == 1 and useSubgroups %}
 enable subgroups;
 {% endif %}
 {% if inlineNorm %}
 const EPSILON: f32 = {{ epsilon }};
 {% endif %}
 {% for stream in ["gate", "up"] %}
 {% if not gemvWalk %}
 {{ matmul_nbits_packed_code(fn=stream ~ "_code", buffer=stream ~ "_b", kBlocks="KBLOCKS", blobSize="BLOB_SIZE", bits=bits) }}
 {% endif %}
 {% if not decodeVec %}
 // Decode two consecutive reduction-axis codes. Below 8 bits an even offset and
 // its successor share one stored byte; at 8 bits they occupy adjacent bytes of
 // one word, because an even byte index in a 4-byte-aligned blob never ends a
 var<workgroup> red_up: array<f32, WG * ROW_TILE>;
 {% endif %}
 {% if inlineNorm %}
+var<workgroup> red_sq: array<f32, WG>;
 fn row_value(index: u32) -> f32 {
 {% if hasSkip %}
 }
 {% endif %}
+{% macro act(b, k) %}{% if fromNormed %}normed[{{ b }} + {{ k }}]{% else %}f32(a[{{ b }} + {{ k }}]){% endif %}{% endmacro %}
 {% if gemvWalk %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>{% if useSubgroups %},
   let row = min(wg.y, ROWS - 1u);
   let base_0 = row * K;
   let col_base = wg.x * N_COLS;
+{% if inlineNorm and writeResidual %}
   // Every N tile computes the same residual row; only the first one stores it,
   // so the tiles never write the same location.
   if (wg.x == 0u) {
     }
   }
 {% endif %}
+  // Whole workgroups past the last column return here, before the reduction's
+  // barriers.
   if (col_base >= N) {
     return;
   }
   var acc_gate{{ sfx }} = vec4<f32>(0.0);
   var acc_up{{ sfx }} = vec4<f32>(0.0);
 {% endfor %}
+{% if inlineNorm %}
+  var acc_sq = 0.0;
+{% endif %}
 {% if decodeVec %}
 {% set codesPerWord = 4 * codesPerByte %}
 {% if laneGroups == 2 %}
   for (var k = tid * 2u; k < K; k = k + WG * 2u) {
     let block = k / BLOCK_SIZE;
     let offset = k - block * BLOCK_SIZE;
+{% if inlineNorm %}
+    let r0 = row_value(base_0 + k);
+    // K need not be even; a code past the end contributes zero.
+    let r1 = select(0.0, row_value(base_0 + min(k + 1u, K - 1u)), k + 1u < K);
+    acc_sq = acc_sq + r0 * r0 + r1 * r1;
+    let v0 = r0 * f32(norm_scale[k]);
+    let v1 = r1 * f32(norm_scale[min(k + 1u, K - 1u)]);
+{% else %}
     let v0 = {{ act("base_0", "k") }};
     // K need not be even; a code past the end contributes zero.
     let v1 = select(0.0, {{ act("base_0", "min(k + 1u, K - 1u)") }}, k + 1u < K);
+{% endif %}
 {% for c in range(4) %}
 {% set comp = ["x", "y", "z", "w"][c] %}
 {% if c == 0 %}
   let sgGate{{ sfx }} = subgroupAdd(acc_gate{{ sfx }});
   let sgUp{{ sfx }} = subgroupAdd(acc_up{{ sfx }});
 {% endfor %}
+{% if inlineNorm %}
+  let sgSq = subgroupAdd(acc_sq);
+{% endif %}
   if (sgLane == 0u) {
+{% if inlineNorm %}
+    red_sq[tid / sgSize] = sgSq;
+{% endif %}
 {% for gi in range(colGroups) %}
 {% set sfx = "" if gi == 0 else gi %}
     red_gate{{ sfx }}[tid / sgSize] = sgGate{{ sfx }};
     var gate_total{{ sfx }} = red_gate{{ sfx }}[0];
     var up_total{{ sfx }} = red_up{{ sfx }}[0];
 {% endfor %}
+{% if inlineNorm %}
+    var sq_total = red_sq[0];
+{% endif %}
     for (var i = 1u; i < subgroupCount; i = i + 1u) {
+{% if inlineNorm %}
+      sq_total = sq_total + red_sq[i];
+{% endif %}
 {% for gi in range(colGroups) %}
 {% set sfx = "" if gi == 0 else gi %}
       gate_total{{ sfx }} = gate_total{{ sfx }} + red_gate{{ sfx }}[i];
   red_gate{{ sfx }}[tid] = acc_gate{{ sfx }};
   red_up{{ sfx }}[tid] = acc_up{{ sfx }};
 {% endfor %}
+{% if inlineNorm %}
+  red_sq[tid] = acc_sq;
+{% endif %}
   workgroupBarrier();
   var stride = WG / 2u;
   loop {
       break;
     }
     if (tid < stride) {
+{% if inlineNorm %}
+      red_sq[tid] = red_sq[tid] + red_sq[tid + stride];
+{% endif %}
 {% for gi in range(colGroups) %}
 {% set sfx = "" if gi == 0 else gi %}
       red_gate{{ sfx }}[tid] = red_gate{{ sfx }}[tid] + red_gate{{ sfx }}[tid + stride];
     let gate_total{{ sfx }} = red_gate{{ sfx }}[0];
     let up_total{{ sfx }} = red_up{{ sfx }}[0];
 {% endfor %}
+{% if inlineNorm %}
+    let sq_total = red_sq[0];
+{% endif %}
+{% endif %}
+{% if inlineNorm %}
+    let inv = inverseSqrt(sq_total / f32(K) + EPSILON);
 {% endif %}
 {% for gi in range(colGroups) %}
 {% set sfx = "" if gi == 0 else gi %}
     if (col_base + {{ i }}u < N) {
 {% endif %}
       let n = col_base + {{ i }}u;
+      var gate_value = gate_total{{ sfx }}.{{ comp }}{% if inlineNorm %} * inv{% endif %};
+      var up_value = up_total{{ sfx }}.{{ comp }}{% if inlineNorm %} * inv{% endif %};
 {% if hasGateBias %}
       gate_value = gate_value + f32(gate_bias[n]);
 {% endif %}
 {% endfor %}
           }
 {% endif %}
+        }{% endmacro %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
   let row0 = wg.y * ROW_TILE;
   let base_{{ r }} = min(row0 + {{ r }}u, ROWS - 1u) * K;
 {% endfor %}
 {% for r in range(rowTile) %}
   var acc_gate_{{ r }} = 0.0;
   var acc_up_{{ r }} = 0.0;

build/webgpu/test.json CHANGED Viewed

@@ -1618,7 +1618,7 @@
     {
       "name": "prefill_skip_nogb_noub",
       "provenance": {
-        "notes": "Five activation rows force the two-pass prefill schedule. Supplying skip and norm_scale while omitting both projection biases exercises staged SkipSimplifiedLayerNormalization without the residual-sum output."
       },
       "attrs": { "K": 32, "N": 8, "bits": 4, "block_size": 16, "activation": "silu" },
       "inputs": {
@@ -1798,6 +1798,54 @@
         "notes": "Single row with an odd K, so the last reduction pair has one live code; skip input and residual output."
       }
     },
     {
       "name": "bits8_decode_two_trips",
       "attrs": { "K": 160, "N": 8, "bits": 8, "block_size": 32, "activation": "silu" },

     {
       "name": "prefill_skip_nogb_noub",
       "provenance": {
+        "notes": "Five activation rows (M=5) with skip and norm_scale supplied but both projection biases omitted checks staged SkipSimplifiedLayerNormalization without the residual-sum output, at 4-bit block_size=16 quantization."
       },
       "attrs": { "K": 32, "N": 8, "bits": 4, "block_size": 16, "activation": "silu" },
       "inputs": {
         "notes": "Single row with an odd K, so the last reduction pair has one live code; skip input and residual output."
       }
     },
+    {
+      "name": "fused_skipsum_multi_trip_odd_k",
+      "attrs": { "K": 301, "N": 5, "bits": 4, "block_size": 16, "activation": "silu", "epsilon": 0.001 },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 301],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 0.8 }
+        },
+        "skipT": {
+          "dtype": "float32",
+          "shape": [1, 301],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.11, "scale": 0.5 }
+        },
+        "normScaleT": {
+          "dtype": "float32",
+          "shape": [301],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.4, "offset": 1.0 }
+        },
+        "gateBT": {
+          "dtype": "uint8",
+          "shape": [5, 19, 8],
+          "data": { "kind": "cycle", "values": [27, 180, 75, 226, 33, 150, 201, 108, 57, 246, 129, 66, 195] }
+        },
+        "gateScalesT": {
+          "dtype": "float32",
+          "shape": [5, 19],
+          "data": { "kind": "cycle", "values": [0.04, 0.055, 0.05, 0.065, 0.06, 0.075, 0.07] }
+        },
+        "upBT": {
+          "dtype": "uint8",
+          "shape": [5, 19, 8],
+          "data": { "kind": "cycle", "values": [211, 44, 137, 98, 165, 20, 233, 121, 78, 190, 15, 252, 87, 143, 61] }
+        },
+        "upScalesT": {
+          "dtype": "float32",
+          "shape": [5, 19],
+          "data": { "kind": "cycle", "values": [0.045, 0.03, 0.07, 0.05, 0.08, 0.035] }
+        }
+      },
+      "outputs": {
+        "yT": { "dtype": "float32", "shape": [1, 5], "tolerance": 0.0001, "relTolerance": 0.0001 },
+        "residualT": { "dtype": "float32", "shape": [1, 301], "tolerance": 0.000001, "relTolerance": 0.000001 }
+      },
+      "provenance": {
+        "notes": "Inline-norm fused route (blob 8, so not the vector walk): the sum of squares accumulates across three 128-code trips of the pair walk, the odd K leaves one live code in the last pair, and N is not a whole column group."
+      }
+    },
     {
       "name": "bits8_decode_two_trips",
       "attrs": { "K": 160, "N": 8, "bits": 8, "block_size": 32, "activation": "silu" },