sync 6fdf6301e2bb

Browse files

Files changed (14) hide show

README.md +3 -3
build/webgpu/{qmoe-mix-stage.wgsl.jinja → expert-slot-mix.wgsl.jinja} +6 -3
build/webgpu/manifest.json +69 -81
build/webgpu/metadata.json +17 -17
build/webgpu/qmoe-fc1-activation-gemv.wgsl.jinja +22 -22
build/webgpu/qmoe-fc1-activation-grouped-sgmat.wgsl.jinja +25 -21
build/webgpu/qmoe-fc1-activation-grouped.wgsl.jinja +25 -26
build/webgpu/qmoe-fc1-activation-stage.wgsl.jinja +48 -21
build/webgpu/qmoe-output-gemv.wgsl.jinja +22 -21
build/webgpu/qmoe-output-grouped-sgmat.wgsl.jinja +25 -20
build/webgpu/qmoe-output-grouped.wgsl.jinja +25 -34
build/webgpu/qmoe-output-stage.wgsl.jinja +47 -19
build/webgpu/qmoe-route-stage.wgsl.jinja +9 -9
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

@@ -75,13 +75,13 @@ 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
 - [`expert-group-slots.wgsl.jinja`](build/webgpu/expert-group-slots.wgsl.jinja)
 - [`qmoe-fc1-activation-gemv.wgsl.jinja`](build/webgpu/qmoe-fc1-activation-gemv.wgsl.jinja)
 - [`qmoe-fc1-activation-grouped-sgmat.wgsl.jinja`](build/webgpu/qmoe-fc1-activation-grouped-sgmat.wgsl.jinja)
 - [`qmoe-fc1-activation-grouped.wgsl.jinja`](build/webgpu/qmoe-fc1-activation-grouped.wgsl.jinja)
 - [`qmoe-fc1-activation-stage.wgsl.jinja`](build/webgpu/qmoe-fc1-activation-stage.wgsl.jinja)
-- [`qmoe-mix-stage.wgsl.jinja`](build/webgpu/qmoe-mix-stage.wgsl.jinja)
 - [`qmoe-output-gemv.wgsl.jinja`](build/webgpu/qmoe-output-gemv.wgsl.jinja)
 - [`qmoe-output-grouped-sgmat.wgsl.jinja`](build/webgpu/qmoe-output-grouped-sgmat.wgsl.jinja)
 - [`qmoe-output-grouped.wgsl.jinja`](build/webgpu/qmoe-output-grouped.wgsl.jinja)
@@ -92,7 +92,7 @@ Some implementation variants require `subgroup-matrix` and `subgroups`. These ar
 ## 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
 - [`expert-group-slots.wgsl.jinja`](build/webgpu/expert-group-slots.wgsl.jinja)
+- [`expert-slot-mix.wgsl.jinja`](build/webgpu/expert-slot-mix.wgsl.jinja)
 - [`qmoe-fc1-activation-gemv.wgsl.jinja`](build/webgpu/qmoe-fc1-activation-gemv.wgsl.jinja)
 - [`qmoe-fc1-activation-grouped-sgmat.wgsl.jinja`](build/webgpu/qmoe-fc1-activation-grouped-sgmat.wgsl.jinja)
 - [`qmoe-fc1-activation-grouped.wgsl.jinja`](build/webgpu/qmoe-fc1-activation-grouped.wgsl.jinja)
 - [`qmoe-fc1-activation-stage.wgsl.jinja`](build/webgpu/qmoe-fc1-activation-stage.wgsl.jinja)
 - [`qmoe-output-gemv.wgsl.jinja`](build/webgpu/qmoe-output-gemv.wgsl.jinja)
 - [`qmoe-output-grouped-sgmat.wgsl.jinja`](build/webgpu/qmoe-output-grouped-sgmat.wgsl.jinja)
 - [`qmoe-output-grouped.wgsl.jinja`](build/webgpu/qmoe-output-grouped.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/{qmoe-mix-stage.wgsl.jinja → expert-slot-mix.wgsl.jinja} RENAMED Viewed

@@ -1,3 +1,8 @@
 {{ env.wgsl.resourceDeclarations }}
 // Routed sum for the grouped schedule. The grouped FC2 stage produces one projected row per
@@ -11,9 +16,7 @@ const WG: u32 = {{ workgroupSize }}u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the per-axis dispatch fold width.
-  // Reduces to gid.x when the dispatch does not fold.
-  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   let total = params.tokenCount * HIDDEN;
   if (index >= total) {
     return;

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // Routed sum for the grouped schedule. The grouped FC2 stage produces one projected row per
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WG", "index", "") }}
   let total = params.tokenCount * HIDDEN;
   if (index >= total) {
     return;

build/webgpu/manifest.json CHANGED Viewed

@@ -52,9 +52,28 @@
     "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",
     "hiddenSize": "dim(shapes.inputT, ranks.inputT - 1)",
     "numTokens": "numel(shapes.inputT) / max(1, hiddenSize)",
-    "topK": "attrs.k",
     "weightBits": "attrs.expert_weight_bits",
     "quantBlockSize": "attrs.block_size if has(attrs, \"block_size\") else 0",
     "packSize": "2 if weightBits == 4 else 1",
@@ -74,19 +93,13 @@
     "quantScalesOk": "tensorDtypes.fc1ScalesT == \"float32\" and tensorDtypes.fc2ScalesT == \"float32\" and dim(shapes.fc1ScalesT, 0) == dim(shapes.routerT, 1) and dim(shapes.fc2ScalesT, 0) == dim(shapes.routerT, 1) and dim(shapes.fc1ScalesT, 1) == dim(shapes.fc1T, 1) and dim(shapes.fc2ScalesT, 1) == hiddenSize and ((ranks.fc1ScalesT == 2 and ranks.fc2ScalesT == 2) if colWiseScales else (ranks.fc1ScalesT == 3 and ranks.fc2ScalesT == 3 and dim(shapes.fc1ScalesT, 2) == fc1ScaleBlocks and dim(shapes.fc2ScalesT, 2) == fc2ScaleBlocks))",
     "quantShapeOk": "inputOutputShapeOk and ranks.routerT == 2 and ranks.fc1T == 3 and ranks.fc2T == 3 and dim(shapes.routerT, 0) == numTokens and dim(shapes.fc1T, 0) == dim(shapes.routerT, 1) and dim(shapes.fc2T, 0) == dim(shapes.routerT, 1) and dim(shapes.fc1T, 1) % fusionSize == 0 and dim(shapes.fc2T, 1) == hiddenSize and dim(shapes.fc1T, 2) * packSize == hiddenSize and dim(shapes.fc2T, 2) * packSize == interSize and quantBlockSizeOk and quantScalesOk",
     "quantContract": "activationSupported and routingModeSupported and rawWeightLayout and quantShapeOk and topK >= 1 and topK <= dim(shapes.routerT, 1)",
-    "workgroupSizeOk": "tunables.workgroupSize >= 1 and tunables.workgroupSize <= deviceWorkgroupCap",
     "hiddenChunkFits": "topK * interSize * 4 <= device.limits.maxStorageBufferBindingSize and topK * interSize * 4 <= device.limits.maxBufferSize",
     "hiddenChunkTokens": "numTokens if interSize == 0 else min(numTokens, max(1, floor(min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize) / (topK * interSize * 4))))",
     "hiddenChunkCount": "max(1, ceilDiv(numTokens, max(1, hiddenChunkTokens)))",
     "routeScratchBytes": "numTokens * topK * 4",
     "routedScratchFits": "routeScratchBytes <= device.limits.maxStorageBufferBindingSize and routeScratchBytes <= device.limits.maxBufferSize",
-    "groupTileM": "tunables.groupThreads * tunables.groupRegM",
-    "groupTileN": "tunables.groupThreads * tunables.groupRegN",
-    "groupTileKVec": "ceilDiv(tunables.groupTileK, 4)",
-    "groupThreadCount": "tunables.groupThreads * tunables.groupThreads",
     "groupSlots": "hiddenChunkTokens * topK",
     "groupMaxTiles": "ceilDiv(groupSlots, max(1, groupTileM)) + dim(shapes.routerT, 1)",
-    "groupSharedBytes": "(groupTileM * tunables.groupTileK + 2 * groupTileN * tunables.groupTileK + groupTileM) * 4",
     "groupSlotOutBytes": "hiddenChunkTokens * topK * hiddenSize * 4",
     "groupedDeviceOk": "groupThreadCount <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.groupThreads <= device.limits.maxComputeWorkgroupSizeX and tunables.groupThreads <= device.limits.maxComputeWorkgroupSizeY and tunables.groupRouteWorkgroup <= deviceWorkgroupCap and groupSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and dim(shapes.routerT, 1) * 8 <= device.limits.maxComputeWorkgroupStorageSize and tunables.groupTileK % 4 == 0",
     "groupedDispatchOk": "groupMaxTiles <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(interSize, max(1, groupTileN)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(hiddenSize, max(1, groupTileN)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
@@ -96,43 +109,35 @@
     "sgmatWorkgroup": "128",
     "sgmatSubgroups": "4",
     "sgmatRowSubtiles": "2",
     "groupedSgmatStageElements": "max(groupTileM * 32, sgmatSubgroups * 4 * 64)",
-    "groupedSgmatSharedBytes": "(groupedSgmatStageElements + 64 * 32) * 4",
-    "groupedSgmatOk": "groupedContract and groupTileM == 32 and interSize % 32 == 0 and groupedSgmatSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(hiddenSize, 64) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and hiddenSize % 32 == 0",
-    "decodeLanesOk": "tunables.decodeLanes >= 1 and tunables.decodeLanes <= tunables.workgroupSize and tunables.workgroupSize % tunables.decodeLanes == 0",
-    "decodeRows": "max(1, tunables.workgroupSize / max(1, tunables.decodeLanes))",
-    "decodeDeviceOk": "workgroupSizeOk and decodeLanesOk and tunables.decodeLanes <= device.limits.maxComputeWorkgroupSizeX and decodeRows <= device.limits.maxComputeWorkgroupSizeY and tunables.workgroupSize * 8 <= device.limits.maxComputeWorkgroupStorageSize",
     "decodeDispatchOk": "ceilDiv(interSize, decodeRows) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(hiddenSize, decodeRows) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and hiddenChunkTokens * topK <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "decodeLaneDepth": "min(hiddenSize, interSize) / max(1, tunables.decodeLanes)",
     "decodeDepthOk": "decodeLaneDepth >= tunables.decodeMinLaneTrips",
     "splitFc1Blocks": "ceilDiv(numTokens * topK * interSize, max(1, tunables.workgroupSize))",
     "decodeOccupancyOk": "splitFc1Blocks <= tunables.decodeBlockTarget",
     "decodeContract": "quantContract and decodeDeviceOk and decodeDispatchOk and hiddenChunkFits and routedScratchFits and interSize > 0 and decodeDepthOk and (decodeOccupancyOk or not groupedContract)",
-    "tokens": "numTokens",
     "hidden": "hiddenSize",
     "experts": "dim(shapes.routerT, 1)",
-    "fc1Rows": "dim(shapes.fc1T, 1)",
-    "inter": "interSize",
-    "workgroupSize": "tunables.workgroupSize",
-    "activationType": "attrs.activation_type",
-    "hasSwigluLimit": "has(attrs, \"swiglu_limit\")",
-    "swigluLimit": "attrs.swiglu_limit if has(attrs, \"swiglu_limit\") else 0"
   },
   "bindings": {
-    "output": { "arg": "outputT", "buffer": "storage", "elementType": "f32" },
-    "router_probs": { "arg": "routerT", "buffer": "read-only-storage", "elementType": "f32" },
-    "route_expert": { "scratch": "routeExpert", "buffer": "storage", "elementType": "u32" },
-    "route_mix": { "scratch": "routeMix", "buffer": "storage", "elementType": "f32" },
-    "route_expert_2": {
       "scratch": "routeExpert",
       "name": "route_expert",
       "buffer": "read-only-storage",
       "elementType": "u32"
     },
-    "slot_list": { "scratch": "slotList", "buffer": "storage", "elementType": "u32" },
-    "tile_meta": { "scratch": "tileMeta", "buffer": "storage", "elementType": "u32" },
     "params": {
-      "buffer": "uniform",
       "struct": [
         { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" },
         {
@@ -142,32 +147,31 @@
         }
       ]
     },
-    "input": { "arg": "inputT", "buffer": "read-only-storage", "elementType": "f32" },
-    "slot_list_2": { "scratch": "slotList", "name": "slot_list", "buffer": "read-only-storage", "elementType": "u32" },
-    "tile_meta_2": { "scratch": "tileMeta", "name": "tile_meta", "buffer": "read-only-storage", "elementType": "u32" },
-    "fc1_experts_weights": { "arg": "fc1T", "buffer": "read-only-storage", "elementType": "u32" },
-    "fc1_scales": { "arg": "fc1ScalesT", "buffer": "read-only-storage", "elementType": "f32" },
-    "hidden_act": { "scratch": "hiddenAct", "buffer": "storage", "elementType": "f32" },
-    "params_2": {
       "name": "params",
-      "buffer": "uniform",
       "struct": [
         { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" },
         { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" },
         { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }
       ]
     },
-    "hidden_act_2": {
       "scratch": "hiddenAct",
       "name": "hidden_act",
       "buffer": "read-only-storage",
       "elementType": "f32"
     },
-    "fc2_experts_weights": { "arg": "fc2T", "buffer": "read-only-storage", "elementType": "u32" },
-    "fc2_scales": { "arg": "fc2ScalesT", "buffer": "read-only-storage", "elementType": "f32" },
-    "slot_out": { "scratch": "slotOut", "buffer": "storage", "elementType": "f32" },
-    "slot_out_2": { "scratch": "slotOut", "name": "slot_out", "buffer": "read-only-storage", "elementType": "f32" },
-    "route_mix_2": { "scratch": "routeMix", "name": "route_mix", "buffer": "read-only-storage", "elementType": "f32" }
   },
   "variants": [
     {
@@ -197,7 +201,6 @@
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
-      "derive": { "groupRouteWorkgroup": "tunables.groupRouteWorkgroup" },
       "intermediates": [
         { "id": "routeExpert", "dtype": "uint32", "shape": "[numTokens * topK]" },
         { "id": "routeMix", "dtype": "float32", "shape": "[numTokens * topK]" },
@@ -211,7 +214,7 @@
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
-          "bindings": ["router_probs", "route_expert", "route_mix"],
           "dispatch": {
             "x": "min(ceilDiv((numTokens), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numTokens), (workgroupSize)), 65535)",
@@ -226,32 +229,28 @@
               "id": "group_stage",
               "name": "QMoE.GroupStage",
               "shader": "expert-group-slots.wgsl.jinja",
-              "bindings": ["route_expert_2", "slot_list", "tile_meta", "params"],
-              "dispatch": {
-                "x": "min(ceilDiv((groupRouteWorkgroup), (groupRouteWorkgroup)), 65535)",
-                "y": "ceilDiv(ceilDiv((groupRouteWorkgroup), (groupRouteWorkgroup)), 65535)",
-                "z": 1
-              }
             },
             {
               "id": "fc1_activation_stage",
               "name": "QMoE.FC1ActivationStageGroupedSgmat",
               "shader": "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja",
-              "bindings": ["input", "slot_list_2", "tile_meta_2", "fc1_experts_weights", "fc1_scales", "hidden_act", "params_2"],
-              "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize * fusionSize, 64)" }
             },
             {
               "id": "output_stage",
               "name": "QMoE.OutputStageGroupedSgmat",
               "shader": "qmoe-output-grouped-sgmat.wgsl.jinja",
-              "bindings": ["hidden_act_2", "slot_list_2", "tile_meta_2", "fc2_experts_weights", "fc2_scales", "slot_out"],
-              "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenSize, 64)" }
             },
             {
               "id": "mix_stage",
               "name": "QMoE.MixStage",
-              "shader": "qmoe-mix-stage.wgsl.jinja",
-              "bindings": ["slot_out_2", "route_mix_2", "output", "params"],
               "dispatch": {
                 "x": "min(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
                 "y": "ceilDiv(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
@@ -266,13 +265,6 @@
       "id": "quant_grouped_routed",
       "priority": 30,
       "when": ["groupedContract"],
-      "derive": {
-        "groupTileK": "tunables.groupTileK",
-        "groupThreads": "tunables.groupThreads",
-        "regM": "tunables.groupRegM",
-        "regN": "tunables.groupRegN",
-        "groupRouteWorkgroup": "tunables.groupRouteWorkgroup"
-      },
       "intermediates": [
         { "id": "routeExpert", "dtype": "uint32", "shape": "[numTokens * topK]" },
         { "id": "routeMix", "dtype": "float32", "shape": "[numTokens * topK]" },
@@ -286,7 +278,7 @@
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
-          "bindings": ["router_probs", "route_expert", "route_mix"],
           "dispatch": {
             "x": "min(ceilDiv((numTokens), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numTokens), (workgroupSize)), 65535)",
@@ -301,32 +293,28 @@
               "id": "group_stage",
               "name": "QMoE.GroupStage",
               "shader": "expert-group-slots.wgsl.jinja",
-              "bindings": ["route_expert_2", "slot_list", "tile_meta", "params"],
-              "dispatch": {
-                "x": "min(ceilDiv((groupRouteWorkgroup), (groupRouteWorkgroup)), 65535)",
-                "y": "ceilDiv(ceilDiv((groupRouteWorkgroup), (groupRouteWorkgroup)), 65535)",
-                "z": 1
-              }
             },
             {
               "id": "fc1_activation_stage",
               "name": "QMoE.FC1ActivationStageGrouped",
               "shader": "qmoe-fc1-activation-grouped.wgsl.jinja",
-              "bindings": ["input", "slot_list_2", "tile_meta_2", "fc1_experts_weights", "fc1_scales", "hidden_act", "params_2"],
               "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" }
             },
             {
               "id": "output_stage",
               "name": "QMoE.OutputStageGrouped",
               "shader": "qmoe-output-grouped.wgsl.jinja",
-              "bindings": ["hidden_act_2", "slot_list_2", "tile_meta_2", "fc2_experts_weights", "fc2_scales", "slot_out"],
               "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenSize, groupTileN)" }
             },
             {
               "id": "mix_stage",
               "name": "QMoE.MixStage",
-              "shader": "qmoe-mix-stage.wgsl.jinja",
-              "bindings": ["slot_out_2", "route_mix_2", "output", "params"],
               "dispatch": {
                 "x": "min(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
                 "y": "ceilDiv(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
@@ -352,7 +340,7 @@
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
-          "bindings": ["router_probs", "route_expert", "route_mix"],
           "dispatch": {
             "x": "min(ceilDiv((numTokens), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numTokens), (workgroupSize)), 65535)",
@@ -367,7 +355,7 @@
               "id": "fc1_activation_stage",
               "name": "QMoE.FC1ActivationStageGemv",
               "shader": "qmoe-fc1-activation-gemv.wgsl.jinja",
-              "bindings": ["input", "route_expert_2", "fc1_experts_weights", "fc1_scales", "hidden_act", "params_2"],
               "dispatch": {
                 "x": "ceilDiv(interSize, decodeRows)",
                 "y": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * topK"
@@ -378,9 +366,9 @@
               "name": "QMoE.OutputStageGemv",
               "shader": "qmoe-output-gemv.wgsl.jinja",
               "bindings": [
-                "hidden_act_2",
-                "route_expert_2",
-                "route_mix_2",
                 "fc2_experts_weights",
                 "fc2_scales",
                 "output",
@@ -412,7 +400,7 @@
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
-          "bindings": ["router_probs", "route_expert", "route_mix"],
           "dispatch": {
             "x": "min(ceilDiv((numTokens), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numTokens), (workgroupSize)), 65535)",
@@ -429,7 +417,7 @@
               "shader": "qmoe-fc1-activation-stage.wgsl.jinja",
               "bindings": [
                 "input",
-                "route_expert_2",
                 "fc1_experts_weights",
                 "fc1_scales",
                 "hidden_act",
@@ -457,7 +445,7 @@
               "id": "output_stage",
               "name": "QMoE.OutputStage",
               "shader": "qmoe-output-stage.wgsl.jinja",
-              "bindings": ["hidden_act_2", "route_expert_2", "route_mix_2", "fc2_experts_weights", "fc2_scales", "output", "params"],
               "dispatch": {
                 "x": "min(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
                 "y": "ceilDiv(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",

     "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",
+    "groupTileK": "tunables.groupTileK",
+    "groupThreads": "tunables.groupThreads",
+    "regM": "tunables.groupRegM",
+    "regN": "tunables.groupRegN",
+    "groupRouteWorkgroup": "tunables.groupRouteWorkgroup",
+    "groupTileM": "tunables.groupThreads * tunables.groupRegM",
+    "groupTileN": "tunables.groupThreads * tunables.groupRegN",
+    "groupTileKVec": "ceilDiv(tunables.groupTileK, 4)",
+    "groupThreadCount": "tunables.groupThreads * tunables.groupThreads",
+    "groupSharedBytes": "(groupTileM * tunables.groupTileK + 2 * groupTileN * tunables.groupTileK + groupTileM) * 4",
+    "topK": "attrs.k",
+    "activationType": "attrs.activation_type",
+    "workgroupSizeOk": "tunables.workgroupSize >= 1 and tunables.workgroupSize <= deviceWorkgroupCap",
+    "decodeLanesOk": "tunables.decodeLanes >= 1 and tunables.decodeLanes <= tunables.workgroupSize and tunables.workgroupSize % tunables.decodeLanes == 0",
+    "decodeRows": "max(1, tunables.workgroupSize / max(1, tunables.decodeLanes))",
+    "decodeDeviceOk": "workgroupSizeOk and decodeLanesOk and tunables.decodeLanes <= device.limits.maxComputeWorkgroupSizeX and decodeRows <= device.limits.maxComputeWorkgroupSizeY and tunables.workgroupSize * 8 <= device.limits.maxComputeWorkgroupStorageSize",
+    "fc1Rows": "dim(shapes.fc1T, 1)",
+    "hasSwigluLimit": "has(attrs, \"swiglu_limit\")",
+    "swigluLimit": "attrs.swiglu_limit if has(attrs, \"swiglu_limit\") else 0",
+    "workgroupSize": "tunables.workgroupSize",
     "hiddenSize": "dim(shapes.inputT, ranks.inputT - 1)",
     "numTokens": "numel(shapes.inputT) / max(1, hiddenSize)",
     "weightBits": "attrs.expert_weight_bits",
     "quantBlockSize": "attrs.block_size if has(attrs, \"block_size\") else 0",
     "packSize": "2 if weightBits == 4 else 1",
     "quantScalesOk": "tensorDtypes.fc1ScalesT == \"float32\" and tensorDtypes.fc2ScalesT == \"float32\" and dim(shapes.fc1ScalesT, 0) == dim(shapes.routerT, 1) and dim(shapes.fc2ScalesT, 0) == dim(shapes.routerT, 1) and dim(shapes.fc1ScalesT, 1) == dim(shapes.fc1T, 1) and dim(shapes.fc2ScalesT, 1) == hiddenSize and ((ranks.fc1ScalesT == 2 and ranks.fc2ScalesT == 2) if colWiseScales else (ranks.fc1ScalesT == 3 and ranks.fc2ScalesT == 3 and dim(shapes.fc1ScalesT, 2) == fc1ScaleBlocks and dim(shapes.fc2ScalesT, 2) == fc2ScaleBlocks))",
     "quantShapeOk": "inputOutputShapeOk and ranks.routerT == 2 and ranks.fc1T == 3 and ranks.fc2T == 3 and dim(shapes.routerT, 0) == numTokens and dim(shapes.fc1T, 0) == dim(shapes.routerT, 1) and dim(shapes.fc2T, 0) == dim(shapes.routerT, 1) and dim(shapes.fc1T, 1) % fusionSize == 0 and dim(shapes.fc2T, 1) == hiddenSize and dim(shapes.fc1T, 2) * packSize == hiddenSize and dim(shapes.fc2T, 2) * packSize == interSize and quantBlockSizeOk and quantScalesOk",
     "quantContract": "activationSupported and routingModeSupported and rawWeightLayout and quantShapeOk and topK >= 1 and topK <= dim(shapes.routerT, 1)",
     "hiddenChunkFits": "topK * interSize * 4 <= device.limits.maxStorageBufferBindingSize and topK * interSize * 4 <= device.limits.maxBufferSize",
     "hiddenChunkTokens": "numTokens if interSize == 0 else min(numTokens, max(1, floor(min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize) / (topK * interSize * 4))))",
     "hiddenChunkCount": "max(1, ceilDiv(numTokens, max(1, hiddenChunkTokens)))",
     "routeScratchBytes": "numTokens * topK * 4",
     "routedScratchFits": "routeScratchBytes <= device.limits.maxStorageBufferBindingSize and routeScratchBytes <= device.limits.maxBufferSize",
     "groupSlots": "hiddenChunkTokens * topK",
     "groupMaxTiles": "ceilDiv(groupSlots, max(1, groupTileM)) + dim(shapes.routerT, 1)",
     "groupSlotOutBytes": "hiddenChunkTokens * topK * hiddenSize * 4",
     "groupedDeviceOk": "groupThreadCount <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.groupThreads <= device.limits.maxComputeWorkgroupSizeX and tunables.groupThreads <= device.limits.maxComputeWorkgroupSizeY and tunables.groupRouteWorkgroup <= deviceWorkgroupCap and groupSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and dim(shapes.routerT, 1) * 8 <= device.limits.maxComputeWorkgroupStorageSize and tunables.groupTileK % 4 == 0",
     "groupedDispatchOk": "groupMaxTiles <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(interSize, max(1, groupTileN)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(hiddenSize, max(1, groupTileN)) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "sgmatWorkgroup": "128",
     "sgmatSubgroups": "4",
     "sgmatRowSubtiles": "2",
+    "sgmatTileCols": "64",
     "groupedSgmatStageElements": "max(groupTileM * 32, sgmatSubgroups * 4 * 64)",
+    "groupedSgmatSharedBytes": "(groupedSgmatStageElements + sgmatTileCols * 32) * 4",
+    "groupedSgmatOk": "groupedContract and groupTileM == 32 and interSize % 32 == 0 and groupedSgmatSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(hiddenSize, sgmatTileCols) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and hiddenSize % 32 == 0",
     "decodeDispatchOk": "ceilDiv(interSize, decodeRows) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and ceilDiv(hiddenSize, decodeRows) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and hiddenChunkTokens * topK <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "decodeLaneDepth": "min(hiddenSize, interSize) / max(1, tunables.decodeLanes)",
     "decodeDepthOk": "decodeLaneDepth >= tunables.decodeMinLaneTrips",
     "splitFc1Blocks": "ceilDiv(numTokens * topK * interSize, max(1, tunables.workgroupSize))",
     "decodeOccupancyOk": "splitFc1Blocks <= tunables.decodeBlockTarget",
     "decodeContract": "quantContract and decodeDeviceOk and decodeDispatchOk and hiddenChunkFits and routedScratchFits and interSize > 0 and decodeDepthOk and (decodeOccupancyOk or not groupedContract)",
     "hidden": "hiddenSize",
     "experts": "dim(shapes.routerT, 1)",
+    "inter": "interSize"
   },
   "bindings": {
+    "output": { "arg": "outputT", "elementType": "f32" },
+    "router_probs": { "arg": "routerT", "elementType": "f32" },
+    "route_expert": { "scratch": "routeExpert", "elementType": "u32" },
+    "route_mix": { "scratch": "routeMix", "elementType": "f32" },
+    "route_expert_u32": {
       "scratch": "routeExpert",
       "name": "route_expert",
       "buffer": "read-only-storage",
       "elementType": "u32"
     },
+    "slot_list": { "scratch": "slotList", "elementType": "u32" },
+    "tile_meta": { "scratch": "tileMeta", "elementType": "u32" },
+    "params_route": { "name": "params", "struct": [{ "name": "tokenCount", "type": "u32", "value": "numTokens" }] },
     "params": {
       "struct": [
         { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" },
         {
         }
       ]
     },
+    "input": { "arg": "inputT", "elementType": "f32" },
+    "slot_list_u32": { "scratch": "slotList", "name": "slot_list", "buffer": "read-only-storage", "elementType": "u32" },
+    "tile_meta_u32": { "scratch": "tileMeta", "name": "tile_meta", "buffer": "read-only-storage", "elementType": "u32" },
+    "fc1_experts_weights": { "arg": "fc1T", "elementType": "u32" },
+    "fc1_scales": { "arg": "fc1ScalesT", "elementType": "f32" },
+    "hidden_act": { "scratch": "hiddenAct", "elementType": "f32" },
+    "params__uniform": {
       "name": "params",
       "struct": [
         { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" },
         { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" },
         { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }
       ]
     },
+    "hidden_act_f32": {
       "scratch": "hiddenAct",
       "name": "hidden_act",
       "buffer": "read-only-storage",
       "elementType": "f32"
     },
+    "fc2_experts_weights": { "arg": "fc2T", "elementType": "u32" },
+    "fc2_scales": { "arg": "fc2ScalesT", "elementType": "f32" },
+    "slot_out": { "scratch": "slotOut", "elementType": "f32" },
+    "slot_out_f32": { "scratch": "slotOut", "name": "slot_out", "buffer": "read-only-storage", "elementType": "f32" },
+    "route_mix_f32": { "scratch": "routeMix", "name": "route_mix", "buffer": "read-only-storage", "elementType": "f32" }
   },
   "variants": [
     {
         "features": ["subgroups", "chromium-experimental-subgroup-matrix"],
         "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }]
       },
       "intermediates": [
         { "id": "routeExpert", "dtype": "uint32", "shape": "[numTokens * topK]" },
         { "id": "routeMix", "dtype": "float32", "shape": "[numTokens * topK]" },
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
+          "bindings": ["router_probs", "route_expert", "route_mix", "params_route"],
           "dispatch": {
             "x": "min(ceilDiv((numTokens), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numTokens), (workgroupSize)), 65535)",
               "id": "group_stage",
               "name": "QMoE.GroupStage",
               "shader": "expert-group-slots.wgsl.jinja",
+              "bindings": ["route_expert_u32", "slot_list", "tile_meta", "params"],
+              "dispatch": { "x": 1 }
             },
             {
               "id": "fc1_activation_stage",
               "name": "QMoE.FC1ActivationStageGroupedSgmat",
               "shader": "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja",
+              "bindings": ["input", "slot_list_u32", "tile_meta_u32", "fc1_experts_weights", "fc1_scales", "hidden_act", "params__uniform"],
+              "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize * fusionSize, sgmatTileCols)" }
             },
             {
               "id": "output_stage",
               "name": "QMoE.OutputStageGroupedSgmat",
               "shader": "qmoe-output-grouped-sgmat.wgsl.jinja",
+              "bindings": ["hidden_act_f32", "slot_list_u32", "tile_meta_u32", "fc2_experts_weights", "fc2_scales", "slot_out"],
+              "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenSize, sgmatTileCols)" }
             },
             {
               "id": "mix_stage",
               "name": "QMoE.MixStage",
+              "shader": "expert-slot-mix.wgsl.jinja",
+              "bindings": ["slot_out_f32", "route_mix_f32", "output", "params"],
               "dispatch": {
                 "x": "min(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
                 "y": "ceilDiv(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
       "id": "quant_grouped_routed",
       "priority": 30,
       "when": ["groupedContract"],
       "intermediates": [
         { "id": "routeExpert", "dtype": "uint32", "shape": "[numTokens * topK]" },
         { "id": "routeMix", "dtype": "float32", "shape": "[numTokens * topK]" },
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
+          "bindings": ["router_probs", "route_expert", "route_mix", "params_route"],
           "dispatch": {
             "x": "min(ceilDiv((numTokens), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numTokens), (workgroupSize)), 65535)",
               "id": "group_stage",
               "name": "QMoE.GroupStage",
               "shader": "expert-group-slots.wgsl.jinja",
+              "bindings": ["route_expert_u32", "slot_list", "tile_meta", "params"],
+              "dispatch": { "x": 1 }
             },
             {
               "id": "fc1_activation_stage",
               "name": "QMoE.FC1ActivationStageGrouped",
               "shader": "qmoe-fc1-activation-grouped.wgsl.jinja",
+              "bindings": ["input", "slot_list_u32", "tile_meta_u32", "fc1_experts_weights", "fc1_scales", "hidden_act", "params__uniform"],
               "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" }
             },
             {
               "id": "output_stage",
               "name": "QMoE.OutputStageGrouped",
               "shader": "qmoe-output-grouped.wgsl.jinja",
+              "bindings": ["hidden_act_f32", "slot_list_u32", "tile_meta_u32", "fc2_experts_weights", "fc2_scales", "slot_out"],
               "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenSize, groupTileN)" }
             },
             {
               "id": "mix_stage",
               "name": "QMoE.MixStage",
+              "shader": "expert-slot-mix.wgsl.jinja",
+              "bindings": ["slot_out_f32", "route_mix_f32", "output", "params"],
               "dispatch": {
                 "x": "min(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
                 "y": "ceilDiv(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
+          "bindings": ["router_probs", "route_expert", "route_mix", "params_route"],
           "dispatch": {
             "x": "min(ceilDiv((numTokens), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numTokens), (workgroupSize)), 65535)",
               "id": "fc1_activation_stage",
               "name": "QMoE.FC1ActivationStageGemv",
               "shader": "qmoe-fc1-activation-gemv.wgsl.jinja",
+              "bindings": ["input", "route_expert_u32", "fc1_experts_weights", "fc1_scales", "hidden_act", "params__uniform"],
               "dispatch": {
                 "x": "ceilDiv(interSize, decodeRows)",
                 "y": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * topK"
               "name": "QMoE.OutputStageGemv",
               "shader": "qmoe-output-gemv.wgsl.jinja",
               "bindings": [
+                "hidden_act_f32",
+                "route_expert_u32",
+                "route_mix_f32",
                 "fc2_experts_weights",
                 "fc2_scales",
                 "output",
           "id": "route_stage",
           "name": "QMoE.RouteStage",
           "shader": "qmoe-route-stage.wgsl.jinja",
+          "bindings": ["router_probs", "route_expert", "route_mix", "params_route"],
           "dispatch": {
             "x": "min(ceilDiv((numTokens), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((numTokens), (workgroupSize)), 65535)",
               "shader": "qmoe-fc1-activation-stage.wgsl.jinja",
               "bindings": [
                 "input",
+                "route_expert_u32",
                 "fc1_experts_weights",
                 "fc1_scales",
                 "hidden_act",
               "id": "output_stage",
               "name": "QMoE.OutputStage",
               "shader": "qmoe-output-stage.wgsl.jinja",
+              "bindings": ["hidden_act_f32", "route_expert_u32", "route_mix_f32", "fc2_experts_weights", "fc2_scales", "output", "params"],
               "dispatch": {
                 "x": "min(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",
                 "y": "ceilDiv(ceilDiv((min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize), (workgroupSize)), 65535)",

build/webgpu/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "com.microsoft.QMoE",
-  "id": "_com_microsoft_qmoe_webgpu_b3768d9",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
@@ -9,28 +9,28 @@
     "files": {
       "bench.json": "YuXSEWa4U+6Ek0E0ZkIeDrzqQ672IelokTilKJuKHf0=",
       "expert-group-slots.wgsl.jinja": "Ta+3H2FA1qRRzksdkwgKLZiO+AokmuM+B9JVN8cv1EQ=",
-      "manifest.json": "r3eaOyzncaumYYrgh5XHHfvd9mHdvVPR3cl5xbiwDDA=",
-      "qmoe-fc1-activation-gemv.wgsl.jinja": "UvP8Vyx1E9+/+I09w/z9iB6h6MHrSrOmjKfIelyIS0M=",
-      "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja": "uWNSXPm0htntxHKCNBzDIo5qFmROP9wDF3mjBGQo+Wo=",
-      "qmoe-fc1-activation-grouped.wgsl.jinja": "6NEYMHvSNgCp6KHJveI4LELrTqr8u/lxIJjSkRX68V4=",
-      "qmoe-fc1-activation-stage.wgsl.jinja": "C9DvNGQ2u4eBAAceyEdbGX13goQFJbxSunkXXhYN604=",
-      "qmoe-mix-stage.wgsl.jinja": "+xEc6boXboFbx+4AJ2/KWElM04LRuufmiIKfEK51/oo=",
-      "qmoe-output-gemv.wgsl.jinja": "OJLoDDWMYE5yaTeDGY7qUeuf6AepJL7ZZLx8ko6OTEQ=",
-      "qmoe-output-grouped-sgmat.wgsl.jinja": "jopTCpUPoIw9MBRhrmjQuZioTbJJIygwBvbwRyhHvUY=",
-      "qmoe-output-grouped.wgsl.jinja": "dOoFq8c0+yBCwOae+QVKBQyodJUK4vwam1RQb46oBNU=",
-      "qmoe-output-stage.wgsl.jinja": "yuSg3G05eKQOunyR5lcCEaJ5y5lXoT73RzT9rq6ahJQ=",
       "qmoe-output-zero-inter.wgsl.jinja": "6jHYpVxU5+9WMflhwMste4v17hq7WWA1BDDS4q4DLS4=",
-      "qmoe-route-stage.wgsl.jinja": "SF5VN6d0zp/ssCGnTwCWccgbnOkZ0mptZiwIbMX3gXA=",
-      "test.json": "/ffqKBtjk2jseUZPswe8gOyYZfhNJu9aOZg8GNlB+9A="
     }
   },
-  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
   "webgpu": {
-    "manifestSpec": "2.0",
     "variants": {
       "quant_zero_inter": ["qmoe-output-zero-inter.wgsl.jinja"],
-      "quant_grouped_sgmat_routed": ["expert-group-slots.wgsl.jinja", "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja", "qmoe-mix-stage.wgsl.jinja", "qmoe-output-grouped-sgmat.wgsl.jinja", "qmoe-route-stage.wgsl.jinja"],
-      "quant_grouped_routed": ["expert-group-slots.wgsl.jinja", "qmoe-fc1-activation-grouped.wgsl.jinja", "qmoe-mix-stage.wgsl.jinja", "qmoe-output-grouped.wgsl.jinja", "qmoe-route-stage.wgsl.jinja"],
       "quant_gemv_routed": ["qmoe-fc1-activation-gemv.wgsl.jinja", "qmoe-output-gemv.wgsl.jinja", "qmoe-route-stage.wgsl.jinja"],
       "quant_split_routed": ["qmoe-fc1-activation-stage.wgsl.jinja", "qmoe-output-stage.wgsl.jinja", "qmoe-route-stage.wgsl.jinja"]
     }

 {
   "name": "com.microsoft.QMoE",
+  "id": "_com_microsoft_qmoe_webgpu_58337b0",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
     "files": {
       "bench.json": "YuXSEWa4U+6Ek0E0ZkIeDrzqQ672IelokTilKJuKHf0=",
       "expert-group-slots.wgsl.jinja": "Ta+3H2FA1qRRzksdkwgKLZiO+AokmuM+B9JVN8cv1EQ=",
+      "expert-slot-mix.wgsl.jinja": "VwX7A1f2r+lxMzG/I5r0KUpGr0Y5gW/CrI/wdwx0Ah4=",
+      "manifest.json": "9eXT25hIM0qWJhCsJjnAodd/OgG9+nO+AvuscEpzk5U=",
+      "qmoe-fc1-activation-gemv.wgsl.jinja": "ZYXWDSbc10LldI8RICBHK6Z/9qrrXbuYfLpVnsa0ON4=",
+      "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja": "emJWdN2BOxm610l2YstemK7OkegF96pKvhnpRSNQ5LY=",
+      "qmoe-fc1-activation-grouped.wgsl.jinja": "TXmpZXCGJhq/T26VJpnp9OcBu14jg4oJCnAvnqY7uF4=",
+      "qmoe-fc1-activation-stage.wgsl.jinja": "zpkM+RhDGOwHtlQBDa6k30BgDLJl2MzkCQDH1+fEFGk=",
+      "qmoe-output-gemv.wgsl.jinja": "aB5licSktZYGWKyHBnL+K4a7ogTZCG5b0IewORHZTVs=",
+      "qmoe-output-grouped-sgmat.wgsl.jinja": "NZocGQlY7dInufxtEwMxylKXDqspt3MyJqUftZ/M6Uk=",
+      "qmoe-output-grouped.wgsl.jinja": "P7AAth2FTgq4p67WQpt6rWMdWs5PsAEn0HURW44wIRQ=",
+      "qmoe-output-stage.wgsl.jinja": "x1VLEwM3N45620tcpvbRcD94PSrt/gxDUjXsYBDQJAw=",
       "qmoe-output-zero-inter.wgsl.jinja": "6jHYpVxU5+9WMflhwMste4v17hq7WWA1BDDS4q4DLS4=",
+      "qmoe-route-stage.wgsl.jinja": "wi5Q1NwN1ZugT/otZryovayihC1gP6EdLO9yw1skhtc=",
+      "test.json": "6Rt3peYHm7WQraXPO0uYfIdb8VlBXwFwgy9pAvQwiac="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "quant_zero_inter": ["qmoe-output-zero-inter.wgsl.jinja"],
+      "quant_grouped_sgmat_routed": ["expert-group-slots.wgsl.jinja", "expert-slot-mix.wgsl.jinja", "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja", "qmoe-output-grouped-sgmat.wgsl.jinja", "qmoe-route-stage.wgsl.jinja"],
+      "quant_grouped_routed": ["expert-group-slots.wgsl.jinja", "expert-slot-mix.wgsl.jinja", "qmoe-fc1-activation-grouped.wgsl.jinja", "qmoe-output-grouped.wgsl.jinja", "qmoe-route-stage.wgsl.jinja"],
       "quant_gemv_routed": ["qmoe-fc1-activation-gemv.wgsl.jinja", "qmoe-output-gemv.wgsl.jinja", "qmoe-route-stage.wgsl.jinja"],
       "quant_split_routed": ["qmoe-fc1-activation-stage.wgsl.jinja", "qmoe-output-stage.wgsl.jinja", "qmoe-route-stage.wgsl.jinja"]
     }

build/webgpu/qmoe-fc1-activation-gemv.wgsl.jinja CHANGED Viewed

@@ -10,54 +10,57 @@ const TOP_K: u32 = {{ topK }}u;
 // reduction, which is LANES times the invocations for the same arithmetic.
 const LANES: u32 = {{ decodeLanes }}u;
 const ROWS: u32 = {{ decodeRows }}u;
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
-const FC1_PACKED_COLS: u32 = {{ fc1PackedCols }}u;
 {% if not colWiseScales %}
-// Block-wise scales (fc1_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
-const FC1_BLOCKS: u32 = {{ fc1ScaleBlocks }}u;
 {% endif %}
-fn fc1_weight(expert: u32, row: u32, col: u32) -> f32 {
-  let base = expert * FC1_ROWS + row;
 {% if packSize == 2 %}
-  let word = fc1_experts_weights[base * FC1_PACKED_COLS + col / 2u];
   let code = f32((word >> ((col & 1u) * 4u)) & 15u);
 {% else %}
-  let code = f32(fc1_experts_weights[base * FC1_PACKED_COLS + col] & 255u);
 {% endif %}
 {% if not colWiseScales %}
-  return (code - MIDPOINT) * fc1_scales[base * FC1_BLOCKS + col / QUANT_BLOCK];
 {% else %}
-  return (code - MIDPOINT) * fc1_scales[base];
 {% endif %}
 }
 // Decode two consecutive reduction-axis codes together. Four-bit codes share a
 // stored word; eight-bit codes occupy adjacent words.
-fn fc1_weight_pair(expert: u32, row: u32, col: u32) -> vec2<f32> {
-  let base = expert * FC1_ROWS + row;
 {% if packSize == 2 %}
-  let word = fc1_experts_weights[base * FC1_PACKED_COLS + col / 2u];
   let lo = f32(word & 15u) - MIDPOINT;
   let hi = f32((word >> 4u) & 15u) - MIDPOINT;
 {% else %}
-  let lo = f32(fc1_experts_weights[base * FC1_PACKED_COLS + col] & 255u) - MIDPOINT;
-  let hi = f32(fc1_experts_weights[base * FC1_PACKED_COLS + col + 1u] & 255u) - MIDPOINT;
 {% endif %}
 {% if not colWiseScales %}
-  return vec2<f32>(lo * fc1_scales[base * FC1_BLOCKS + col / QUANT_BLOCK],
-                   hi * fc1_scales[base * FC1_BLOCKS + (col + 1u) / QUANT_BLOCK]);
 {% else %}
-  let scale = fc1_scales[base];
   return vec2<f32>(lo * scale, hi * scale);
 {% endif %}
 }
 {% if activationType == "swiglu" %}
 fn activate(gate_in: f32, up_in: f32) -> f32 {
 {% if hasSwigluLimit %}
@@ -77,7 +80,6 @@ fn activate(value: f32) -> f32 {
 }
 {% endif %}
 {% macro rowlane_fold(arrays, lanes="LANES", lane="lane", slot="slot") %}
   var n = {{ lanes }};
   while (n > 1u) {
@@ -89,9 +91,7 @@ fn activate(value: f32) -> f32 {
     }
     workgroupBarrier();
     n = half;
-  }
-{%- endmacro %}
 var<workgroup> gate_partial: array<f32, {{ decodeLanes * decodeRows }}>;
 {% if activationType == "swiglu" %}
 // The gate and up rows of an interleaved SwiGLU pair share the same activation

 // reduction, which is LANES times the invocations for the same arithmetic.
 const LANES: u32 = {{ decodeLanes }}u;
 const ROWS: u32 = {{ decodeRows }}u;
+{% set qw = "fc1" %}
+{% set qwRows = "FC1_ROWS" %}
+{% set qwPackedCols = fc1PackedCols %}
+{% set qwBlocks = fc1ScaleBlocks %}
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const {{ qw | upper }}_PACKED_COLS: u32 = {{ qwPackedCols }}u;
 {% if not colWiseScales %}
+// Block-wise scales ({{ qw }}_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const {{ qw | upper }}_BLOCKS: u32 = {{ qwBlocks }}u;
 {% endif %}
+fn {{ qw }}_weight(expert: u32, row: u32, col: u32) -> f32 {
+  let base = expert * {{ qwRows }} + row;
 {% if packSize == 2 %}
+  let word = {{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col / 2u];
   let code = f32((word >> ((col & 1u) * 4u)) & 15u);
 {% else %}
+  let code = f32({{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col] & 255u);
 {% endif %}
 {% if not colWiseScales %}
+  return (code - MIDPOINT) * {{ qw }}_scales[base * {{ qw | upper }}_BLOCKS + col / QUANT_BLOCK];
 {% else %}
+  return (code - MIDPOINT) * {{ qw }}_scales[base];
 {% endif %}
 }
 // Decode two consecutive reduction-axis codes together. Four-bit codes share a
 // stored word; eight-bit codes occupy adjacent words.
+fn {{ qw }}_weight_pair(expert: u32, row: u32, col: u32) -> vec2<f32> {
+  let base = expert * {{ qwRows }} + row;
 {% if packSize == 2 %}
+  let word = {{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col / 2u];
   let lo = f32(word & 15u) - MIDPOINT;
   let hi = f32((word >> 4u) & 15u) - MIDPOINT;
 {% else %}
+  let lo = f32({{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col] & 255u) - MIDPOINT;
+  let hi = f32({{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col + 1u] & 255u) - MIDPOINT;
 {% endif %}
 {% if not colWiseScales %}
+  return vec2<f32>(lo * {{ qw }}_scales[base * {{ qw | upper }}_BLOCKS + col / QUANT_BLOCK],
+                   hi * {{ qw }}_scales[base * {{ qw | upper }}_BLOCKS + (col + 1u) / QUANT_BLOCK]);
 {% else %}
+  let scale = {{ qw }}_scales[base];
   return vec2<f32>(lo * scale, hi * scale);
 {% endif %}
 }
 {% if activationType == "swiglu" %}
 fn activate(gate_in: f32, up_in: f32) -> f32 {
 {% if hasSwigluLimit %}
 }
 {% endif %}
 {% macro rowlane_fold(arrays, lanes="LANES", lane="lane", slot="slot") %}
   var n = {{ lanes }};
   while (n > 1u) {
     }
     workgroupBarrier();
     n = half;
+  }{% endmacro %}
 var<workgroup> gate_partial: array<f32, {{ decodeLanes * decodeRows }}>;
 {% if activationType == "swiglu" %}
 // The gate and up rows of an interleaved SwiGLU pair share the same activation

build/webgpu/qmoe-fc1-activation-grouped-sgmat.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 }}
 // Expert-grouped FC1 and activation stage using subgroup-matrix operations.
@@ -20,50 +19,53 @@ const HIDDEN: u32 = {{ hidden }}u;
 const FC1_ROWS: u32 = {{ fc1Rows }}u;
 const INTER: u32 = {{ inter }}u;
 const TOP_K: u32 = {{ topK }}u;
-const TILE_COLS: u32 = 64u;
 const TILE_K: u32 = 32u;
 const SUB_COLS: u32 = 32u;
 const SUB_ROWS: u32 = 16u;
 {% set isSwiglu = activationType == "swiglu" %}
 // Weight rows per output column: the fused gate/up pair, or one.
 const FUSION: u32 = {{ 2 if isSwiglu else 1 }}u;
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
-const FC1_PACKED_COLS: u32 = {{ fc1PackedCols }}u;
 {% if not colWiseScales %}
-// Block-wise scales (fc1_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
-const FC1_BLOCKS: u32 = {{ fc1ScaleBlocks }}u;
 {% endif %}
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
 // this helper requires `col` to be a multiple of four.
-fn fc1_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
-  let base = expert * FC1_ROWS + row;
 {% if packSize == 2 %}
-  let pair = base * FC1_PACKED_COLS + col / 2u;
-  let w0 = fc1_experts_weights[pair];
-  let w1 = fc1_experts_weights[pair + 1u];
   let codes = vec4<f32>(f32(w0 & 15u), f32((w0 >> 4u) & 15u),
                         f32(w1 & 15u), f32((w1 >> 4u) & 15u));
 {% else %}
-  let word = base * FC1_PACKED_COLS + col;
-  let codes = vec4<f32>(f32(fc1_experts_weights[word] & 255u), f32(fc1_experts_weights[word + 1u] & 255u),
-                        f32(fc1_experts_weights[word + 2u] & 255u), f32(fc1_experts_weights[word + 3u] & 255u));
 {% endif %}
 {% if not colWiseScales %}
-  let scale = fc1_scales[base * FC1_BLOCKS + col / QUANT_BLOCK];
 {% else %}
-  let scale = fc1_scales[base];
 {% endif %}
   return (codes - vec4<f32>(MIDPOINT)) * scale;
 }
 {% if activationType == "swiglu" %}
 fn activate(gate_in: f32, up_in: f32) -> f32 {
 {% if hasSwigluLimit %}
@@ -83,11 +85,10 @@ fn activate(value: f32) -> f32 {
 }
 {% endif %}
 // The final K-loop barrier ends the A operand's lifetime. Its storage then
 // holds four 8x8 output banks per subgroup, fenced by the epilogue barriers.
 var<workgroup> tile_A: array<f32, {{ groupedSgmatStageElements }}>;
-var<workgroup> tile_B: array<f32, 64 * 32>;
 // Rows past the tile's end stage zero and are dropped at the store, so a short final
 // tile needs no separate arm.
@@ -149,9 +150,11 @@ fn main(
   let base_A = subtile_idy * SUB_ROWS;
   let base_B = subtile_idx * SUB_COLS;
-{% for row in [0, 1] %}{% for col in [0, 1, 2, 3] %}
   var matC{{ row }}{{ col }}: subgroup_matrix_result<f32, 8, 8>;
-{% endfor %}{% endfor %}
   for (var kidx = 0u; kidx < HIDDEN; kidx = kidx + TILE_K) {
     loadSHMA(kidx, rows, local_idx / 4u, local_idx % 4u, a_slot);
@@ -159,7 +162,8 @@ fn main(
     workgroupBarrier();
     for (var step = 0u; step < TILE_K; step = step + 8u) {
-{% set operandScalar = "f32" %}{% set directInputs = directMatrixInputs is defined and directMatrixInputs %}
       let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
 {% for r in range(2) %}
       var matA{{ r }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ "w" if directInputs else "tile_A" }}, matrix_a_offset{% if r > 0 %} + 8u * {{ "K" if directInputs else "TILE_K" }}{% endif %}, {{ "K" if directInputs else "TILE_K" }});

 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 // Expert-grouped FC1 and activation stage using subgroup-matrix operations.
 const FC1_ROWS: u32 = {{ fc1Rows }}u;
 const INTER: u32 = {{ inter }}u;
 const TOP_K: u32 = {{ topK }}u;
+const TILE_COLS: u32 = {{ sgmatTileCols }}u;
 const TILE_K: u32 = 32u;
 const SUB_COLS: u32 = 32u;
 const SUB_ROWS: u32 = 16u;
 {% set isSwiglu = activationType == "swiglu" %}
 // Weight rows per output column: the fused gate/up pair, or one.
 const FUSION: u32 = {{ 2 if isSwiglu else 1 }}u;
+{% set qw = "fc1" %}
+{% set qwRows = "FC1_ROWS" %}
+{% set qwPackedCols = fc1PackedCols %}
+{% set qwBlocks = fc1ScaleBlocks %}
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const {{ qw | upper }}_PACKED_COLS: u32 = {{ qwPackedCols }}u;
 {% if not colWiseScales %}
+// Block-wise scales ({{ qw }}_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const {{ qw | upper }}_BLOCKS: u32 = {{ qwBlocks }}u;
 {% endif %}
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
 // this helper requires `col` to be a multiple of four.
+fn {{ qw }}_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
+  let base = expert * {{ qwRows }} + row;
 {% if packSize == 2 %}
+  let pair = base * {{ qw | upper }}_PACKED_COLS + col / 2u;
+  let w0 = {{ qw }}_experts_weights[pair];
+  let w1 = {{ qw }}_experts_weights[pair + 1u];
   let codes = vec4<f32>(f32(w0 & 15u), f32((w0 >> 4u) & 15u),
                         f32(w1 & 15u), f32((w1 >> 4u) & 15u));
 {% else %}
+  let word = base * {{ qw | upper }}_PACKED_COLS + col;
+  let codes = vec4<f32>(f32({{ qw }}_experts_weights[word] & 255u), f32({{ qw }}_experts_weights[word + 1u] & 255u),
+                        f32({{ qw }}_experts_weights[word + 2u] & 255u), f32({{ qw }}_experts_weights[word + 3u] & 255u));
 {% endif %}
 {% if not colWiseScales %}
+  let scale = {{ qw }}_scales[base * {{ qw | upper }}_BLOCKS + col / QUANT_BLOCK];
 {% else %}
+  let scale = {{ qw }}_scales[base];
 {% endif %}
   return (codes - vec4<f32>(MIDPOINT)) * scale;
 }
 {% if activationType == "swiglu" %}
 fn activate(gate_in: f32, up_in: f32) -> f32 {
 {% if hasSwigluLimit %}
 }
 {% endif %}
 // The final K-loop barrier ends the A operand's lifetime. Its storage then
 // holds four 8x8 output banks per subgroup, fenced by the epilogue barriers.
 var<workgroup> tile_A: array<f32, {{ groupedSgmatStageElements }}>;
+var<workgroup> tile_B: array<f32, {{ sgmatTileCols }} * 32>;
 // Rows past the tile's end stage zero and are dropped at the store, so a short final
 // tile needs no separate arm.
   let base_A = subtile_idy * SUB_ROWS;
   let base_B = subtile_idx * SUB_COLS;
+{% for row in [0, 1] %}
+{% for col in [0, 1, 2, 3] %}
   var matC{{ row }}{{ col }}: subgroup_matrix_result<f32, 8, 8>;
+{% endfor %}
+{% endfor %}
   for (var kidx = 0u; kidx < HIDDEN; kidx = kidx + TILE_K) {
     loadSHMA(kidx, rows, local_idx / 4u, local_idx % 4u, a_slot);
     workgroupBarrier();
     for (var step = 0u; step < TILE_K; step = step + 8u) {
+{% set operandScalar = "f32" %}
+{% set directInputs = false %}
       let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
 {% for r in range(2) %}
       var matA{{ r }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ "w" if directInputs else "tile_A" }}, matrix_a_offset{% if r > 0 %} + 8u * {{ "K" if directInputs else "TILE_K" }}{% endif %}, {{ "K" if directInputs else "TILE_K" }});

build/webgpu/qmoe-fc1-activation-grouped.wgsl.jinja CHANGED Viewed

@@ -15,57 +15,60 @@ const THREADS_SIDE: u32 = {{ groupThreads }}u;
 const THREADS: u32 = THREADS_SIDE * THREADS_SIDE;
 {% set isSwiglu = activationType == "swiglu" %}
 {% set primaryRow = "col * 2u" if isSwiglu else "col" %}
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
-const FC1_PACKED_COLS: u32 = {{ fc1PackedCols }}u;
 {% if not colWiseScales %}
-// Block-wise scales (fc1_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
-const FC1_BLOCKS: u32 = {{ fc1ScaleBlocks }}u;
 {% endif %}
-fn fc1_weight(expert: u32, row: u32, col: u32) -> f32 {
-  let base = expert * FC1_ROWS + row;
 {% if packSize == 2 %}
-  let word = fc1_experts_weights[base * FC1_PACKED_COLS + col / 2u];
   let code = f32((word >> ((col & 1u) * 4u)) & 15u);
 {% else %}
-  let code = f32(fc1_experts_weights[base * FC1_PACKED_COLS + col] & 255u);
 {% endif %}
 {% if not colWiseScales %}
-  return (code - MIDPOINT) * fc1_scales[base * FC1_BLOCKS + col / QUANT_BLOCK];
 {% else %}
-  return (code - MIDPOINT) * fc1_scales[base];
 {% endif %}
 }
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
 // this helper requires `col` to be a multiple of four.
-fn fc1_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
-  let base = expert * FC1_ROWS + row;
 {% if packSize == 2 %}
-  let pair = base * FC1_PACKED_COLS + col / 2u;
-  let w0 = fc1_experts_weights[pair];
-  let w1 = fc1_experts_weights[pair + 1u];
   let codes = vec4<f32>(f32(w0 & 15u), f32((w0 >> 4u) & 15u),
                         f32(w1 & 15u), f32((w1 >> 4u) & 15u));
 {% else %}
-  let word = base * FC1_PACKED_COLS + col;
-  let codes = vec4<f32>(f32(fc1_experts_weights[word] & 255u), f32(fc1_experts_weights[word + 1u] & 255u),
-                        f32(fc1_experts_weights[word + 2u] & 255u), f32(fc1_experts_weights[word + 3u] & 255u));
 {% endif %}
 {% if not colWiseScales %}
-  let scale = fc1_scales[base * FC1_BLOCKS + col / QUANT_BLOCK];
 {% else %}
-  let scale = fc1_scales[base];
 {% endif %}
   return (codes - vec4<f32>(MIDPOINT)) * scale;
 }
 {% if activationType == "swiglu" %}
 fn activate(gate_in: f32, up_in: f32) -> f32 {
 {% if hasSwigluLimit %}
@@ -85,7 +88,6 @@ fn activate(value: f32) -> f32 {
 }
 {% endif %}
 const KVEC: u32 = {{ groupTileKVec }}u;
 {% macro stage_group_tiles(aLoad, bLoad, b2Load, kExtent, nExtent, guarded, bLoadVec4="", b2LoadVec4="", aLoadVec4="") %}
       for (var idx = tid; idx < MTILE * KVEC; idx = idx + THREADS) {
@@ -138,8 +140,7 @@ const KVEC: u32 = {{ groupTileKVec }}u;
 {% if b2Load %}
         b2_tile[idx] = b2_vec;
 {% endif %}
-      }
-{%- endmacro %}
 {% macro group_tile_loop(aLoad, bLoad, b2Load, kExtent, nExtent, regM, regN, bLoadVec4="", b2LoadVec4="", aLoadVec4="") %}
 {% for r in range(regM) %}
@@ -192,9 +193,7 @@ const KVEC: u32 = {{ groupTileKVec }}u;
     // Orders this step's tile reads before the next step overwrites them.
     workgroupBarrier();
     k_base = k_base + KTILE;
-  }
-{%- endmacro %}
 var<workgroup> row_slot: array<u32, {{ groupTileM }}>;
 var<workgroup> a_tile: array<vec4<f32>, {{ groupTileM * groupTileKVec }}>;
 var<workgroup> b_tile: array<vec4<f32>, {{ groupTileN * groupTileKVec }}>;

 const THREADS: u32 = THREADS_SIDE * THREADS_SIDE;
 {% set isSwiglu = activationType == "swiglu" %}
 {% set primaryRow = "col * 2u" if isSwiglu else "col" %}
+{% set qw = "fc1" %}
+{% set qwRows = "FC1_ROWS" %}
+{% set qwPackedCols = fc1PackedCols %}
+{% set qwBlocks = fc1ScaleBlocks %}
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const {{ qw | upper }}_PACKED_COLS: u32 = {{ qwPackedCols }}u;
 {% if not colWiseScales %}
+// Block-wise scales ({{ qw }}_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const {{ qw | upper }}_BLOCKS: u32 = {{ qwBlocks }}u;
 {% endif %}
+fn {{ qw }}_weight(expert: u32, row: u32, col: u32) -> f32 {
+  let base = expert * {{ qwRows }} + row;
 {% if packSize == 2 %}
+  let word = {{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col / 2u];
   let code = f32((word >> ((col & 1u) * 4u)) & 15u);
 {% else %}
+  let code = f32({{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col] & 255u);
 {% endif %}
 {% if not colWiseScales %}
+  return (code - MIDPOINT) * {{ qw }}_scales[base * {{ qw | upper }}_BLOCKS + col / QUANT_BLOCK];
 {% else %}
+  return (code - MIDPOINT) * {{ qw }}_scales[base];
 {% endif %}
 }
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
 // this helper requires `col` to be a multiple of four.
+fn {{ qw }}_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
+  let base = expert * {{ qwRows }} + row;
 {% if packSize == 2 %}
+  let pair = base * {{ qw | upper }}_PACKED_COLS + col / 2u;
+  let w0 = {{ qw }}_experts_weights[pair];
+  let w1 = {{ qw }}_experts_weights[pair + 1u];
   let codes = vec4<f32>(f32(w0 & 15u), f32((w0 >> 4u) & 15u),
                         f32(w1 & 15u), f32((w1 >> 4u) & 15u));
 {% else %}
+  let word = base * {{ qw | upper }}_PACKED_COLS + col;
+  let codes = vec4<f32>(f32({{ qw }}_experts_weights[word] & 255u), f32({{ qw }}_experts_weights[word + 1u] & 255u),
+                        f32({{ qw }}_experts_weights[word + 2u] & 255u), f32({{ qw }}_experts_weights[word + 3u] & 255u));
 {% endif %}
 {% if not colWiseScales %}
+  let scale = {{ qw }}_scales[base * {{ qw | upper }}_BLOCKS + col / QUANT_BLOCK];
 {% else %}
+  let scale = {{ qw }}_scales[base];
 {% endif %}
   return (codes - vec4<f32>(MIDPOINT)) * scale;
 }
 {% if activationType == "swiglu" %}
 fn activate(gate_in: f32, up_in: f32) -> f32 {
 {% if hasSwigluLimit %}
 }
 {% endif %}
 const KVEC: u32 = {{ groupTileKVec }}u;
 {% macro stage_group_tiles(aLoad, bLoad, b2Load, kExtent, nExtent, guarded, bLoadVec4="", b2LoadVec4="", aLoadVec4="") %}
       for (var idx = tid; idx < MTILE * KVEC; idx = idx + THREADS) {
 {% if b2Load %}
         b2_tile[idx] = b2_vec;
 {% endif %}
+      }{% endmacro %}
 {% macro group_tile_loop(aLoad, bLoad, b2Load, kExtent, nExtent, regM, regN, bLoadVec4="", b2LoadVec4="", aLoadVec4="") %}
 {% for r in range(regM) %}
     // Orders this step's tile reads before the next step overwrites them.
     workgroupBarrier();
     k_base = k_base + KTILE;
+  }{% endmacro %}
 var<workgroup> row_slot: array<u32, {{ groupTileM }}>;
 var<workgroup> a_tile: array<vec4<f32>, {{ groupTileM * groupTileKVec }}>;
 var<workgroup> b_tile: array<vec4<f32>, {{ groupTileN * groupTileKVec }}>;

build/webgpu/qmoe-fc1-activation-stage.wgsl.jinja CHANGED Viewed

@@ -1,4 +1,31 @@
 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hidden }}u;
 const FC1_ROWS: u32 = {{ fc1Rows }}u;
@@ -11,34 +38,37 @@ const WG: u32 = {{ workgroupSize }}u;
 {% if topK > 1 %}
 const TOP_K: u32 = {{ topK }}u;
 {% endif %}
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
-const FC1_PACKED_COLS: u32 = {{ fc1PackedCols }}u;
 {% if not colWiseScales %}
-// Block-wise scales (fc1_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
-const FC1_BLOCKS: u32 = {{ fc1ScaleBlocks }}u;
 {% endif %}
-fn fc1_weight(expert: u32, row: u32, col: u32) -> f32 {
-  let base = expert * FC1_ROWS + row;
 {% if packSize == 2 %}
-  let word = fc1_experts_weights[base * FC1_PACKED_COLS + col / 2u];
   let code = f32((word >> ((col & 1u) * 4u)) & 15u);
 {% else %}
-  let code = f32(fc1_experts_weights[base * FC1_PACKED_COLS + col] & 255u);
 {% endif %}
 {% if not colWiseScales %}
-  return (code - MIDPOINT) * fc1_scales[base * FC1_BLOCKS + col / QUANT_BLOCK];
 {% else %}
-  return (code - MIDPOINT) * fc1_scales[base];
 {% endif %}
 }
 {% if activationType == "swiglu" %}
 fn activate(gate_in: f32, up_in: f32) -> f32 {
 {% if hasSwigluLimit %}
@@ -58,12 +88,9 @@ fn activate(value: f32) -> f32 {
 }
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the per-axis dispatch fold width.
-  // Reduces to gid.x when the dispatch does not fold.
-  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
 {% if topK > 1 %}
   // One activation row per routed slot rather than per token: the grid is
   // (token, slot, inter) and hidden_act is addressed by the same flat index.
@@ -91,25 +118,25 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   let expert = route_expert[token];
 {% endif %}
-  var gate = 0.0;
 {% if activationType == "swiglu" %}
-  var up = 0.0;
 {% endif %}
   for (var h: u32 = 0u; h < HIDDEN; h = h + 1u) {
     let x = f32(input[token * HIDDEN + h]);
 {% if activationType == "swiglu" %}
-    gate = gate + x * fc1_weight(expert, i * 2u, h);
 {% else %}
-    gate = gate + x * fc1_weight(expert, i, h);
 {% endif %}
 {% if activationType == "swiglu" %}
-    up = up + x * fc1_weight(expert, i * 2u + 1u, h);
 {% endif %}
   }
 {% if activationType == "swiglu" %}
-  hidden_act[index] = activate(gate, up);
 {% else %}
-  hidden_act[index] = activate(gate);
 {% endif %}
 }

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
+{% set dotType = "f32" %}
+// Retain product and addition residuals across a dot product. Explicit fma
+// boundaries preserve the addition error transform under reassociation.
+struct DotAccumulator {
+  hi: {{ dotType }},
+  lo: {{ dotType }},
+}
+fn dot_accumulate(acc: DotAccumulator, a: {{ dotType }}, b: {{ dotType }}) -> DotAccumulator {
+  let product = fma(a, b, {{ dotType }}(0.0));
+  let productError = fma(a, b, -product);
+  let sum = fma(acc.hi, {{ dotType }}(1.0), product);
+  let bv = fma({{ dotType }}(-1.0), acc.hi, sum);
+  let av = fma({{ dotType }}(-1.0), bv, sum);
+  let ae = fma({{ dotType }}(-1.0), av, acc.hi);
+  let be = fma({{ dotType }}(-1.0), bv, product);
+  let error = fma(fma(fma(ae, {{ dotType }}(1.0), be), {{ dotType }}(1.0), acc.lo), {{ dotType }}(1.0), productError);
+  let hi = fma(sum, {{ dotType }}(1.0), error);
+  return DotAccumulator(hi, fma({{ dotType }}(-1.0), fma({{ dotType }}(-1.0), sum, hi), error));
+}
+fn dot_value(acc: DotAccumulator) -> {{ dotType }} {
+  return fma(acc.hi, {{ dotType }}(1.0), acc.lo);
+}
 const HIDDEN: u32 = {{ hidden }}u;
 const FC1_ROWS: u32 = {{ fc1Rows }}u;
 {% if topK > 1 %}
 const TOP_K: u32 = {{ topK }}u;
 {% endif %}
+{% set qw = "fc1" %}
+{% set qwRows = "FC1_ROWS" %}
+{% set qwPackedCols = fc1PackedCols %}
+{% set qwBlocks = fc1ScaleBlocks %}
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const {{ qw | upper }}_PACKED_COLS: u32 = {{ qwPackedCols }}u;
 {% if not colWiseScales %}
+// Block-wise scales ({{ qw }}_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const {{ qw | upper }}_BLOCKS: u32 = {{ qwBlocks }}u;
 {% endif %}
+fn {{ qw }}_weight(expert: u32, row: u32, col: u32) -> f32 {
+  let base = expert * {{ qwRows }} + row;
 {% if packSize == 2 %}
+  let word = {{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col / 2u];
   let code = f32((word >> ((col & 1u) * 4u)) & 15u);
 {% else %}
+  let code = f32({{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col] & 255u);
 {% endif %}
 {% if not colWiseScales %}
+  return (code - MIDPOINT) * {{ qw }}_scales[base * {{ qw | upper }}_BLOCKS + col / QUANT_BLOCK];
 {% else %}
+  return (code - MIDPOINT) * {{ qw }}_scales[base];
 {% endif %}
 }
 {% if activationType == "swiglu" %}
 fn activate(gate_in: f32, up_in: f32) -> f32 {
 {% if hasSwigluLimit %}
 }
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WG", "index", "") }}
 {% if topK > 1 %}
   // One activation row per routed slot rather than per token: the grid is
   // (token, slot, inter) and hidden_act is addressed by the same flat index.
   let expert = route_expert[token];
 {% endif %}
+  var gate = DotAccumulator(0.0, 0.0);
 {% if activationType == "swiglu" %}
+  var up = DotAccumulator(0.0, 0.0);
 {% endif %}
   for (var h: u32 = 0u; h < HIDDEN; h = h + 1u) {
     let x = f32(input[token * HIDDEN + h]);
 {% if activationType == "swiglu" %}
+    gate = dot_accumulate(gate, x, fc1_weight(expert, i * 2u, h));
 {% else %}
+    gate = dot_accumulate(gate, x, fc1_weight(expert, i, h));
 {% endif %}
 {% if activationType == "swiglu" %}
+    up = dot_accumulate(up, x, fc1_weight(expert, i * 2u + 1u, h));
 {% endif %}
   }
 {% if activationType == "swiglu" %}
+  hidden_act[index] = activate(dot_value(gate), dot_value(up));
 {% else %}
+  hidden_act[index] = activate(dot_value(gate));
 {% endif %}
 }

build/webgpu/qmoe-output-gemv.wgsl.jinja CHANGED Viewed

@@ -8,54 +8,57 @@ const TOP_K: u32 = {{ topK }}u;
 // LANES-wide cooperative reduction.
 const LANES: u32 = {{ decodeLanes }}u;
 const ROWS: u32 = {{ decodeRows }}u;
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
-const FC2_PACKED_COLS: u32 = {{ fc2PackedCols }}u;
 {% if not colWiseScales %}
-// Block-wise scales (fc2_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
-const FC2_BLOCKS: u32 = {{ fc2ScaleBlocks }}u;
 {% endif %}
-fn fc2_weight(expert: u32, row: u32, col: u32) -> f32 {
-  let base = expert * HIDDEN + row;
 {% if packSize == 2 %}
-  let word = fc2_experts_weights[base * FC2_PACKED_COLS + col / 2u];
   let code = f32((word >> ((col & 1u) * 4u)) & 15u);
 {% else %}
-  let code = f32(fc2_experts_weights[base * FC2_PACKED_COLS + col] & 255u);
 {% endif %}
 {% if not colWiseScales %}
-  return (code - MIDPOINT) * fc2_scales[base * FC2_BLOCKS + col / QUANT_BLOCK];
 {% else %}
-  return (code - MIDPOINT) * fc2_scales[base];
 {% endif %}
 }
 // Decode two consecutive reduction-axis codes together. Four-bit codes share a
 // stored word; eight-bit codes occupy adjacent words.
-fn fc2_weight_pair(expert: u32, row: u32, col: u32) -> vec2<f32> {
-  let base = expert * HIDDEN + row;
 {% if packSize == 2 %}
-  let word = fc2_experts_weights[base * FC2_PACKED_COLS + col / 2u];
   let lo = f32(word & 15u) - MIDPOINT;
   let hi = f32((word >> 4u) & 15u) - MIDPOINT;
 {% else %}
-  let lo = f32(fc2_experts_weights[base * FC2_PACKED_COLS + col] & 255u) - MIDPOINT;
-  let hi = f32(fc2_experts_weights[base * FC2_PACKED_COLS + col + 1u] & 255u) - MIDPOINT;
 {% endif %}
 {% if not colWiseScales %}
-  return vec2<f32>(lo * fc2_scales[base * FC2_BLOCKS + col / QUANT_BLOCK],
-                   hi * fc2_scales[base * FC2_BLOCKS + (col + 1u) / QUANT_BLOCK]);
 {% else %}
-  let scale = fc2_scales[base];
   return vec2<f32>(lo * scale, hi * scale);
 {% endif %}
 }
 {% macro rowlane_fold(arrays, lanes="LANES", lane="lane", slot="slot") %}
   var n = {{ lanes }};
   while (n > 1u) {
@@ -67,9 +70,7 @@ fn fc2_weight_pair(expert: u32, row: u32, col: u32) -> vec2<f32> {
     }
     workgroupBarrier();
     n = half;
-  }
-{%- endmacro %}
 var<workgroup> partial: array<f32, {{ decodeLanes * decodeRows }}>;
 @compute @workgroup_size(LANES, ROWS, 1)

 // LANES-wide cooperative reduction.
 const LANES: u32 = {{ decodeLanes }}u;
 const ROWS: u32 = {{ decodeRows }}u;
+{% set qw = "fc2" %}
+{% set qwRows = "HIDDEN" %}
+{% set qwPackedCols = fc2PackedCols %}
+{% set qwBlocks = fc2ScaleBlocks %}
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const {{ qw | upper }}_PACKED_COLS: u32 = {{ qwPackedCols }}u;
 {% if not colWiseScales %}
+// Block-wise scales ({{ qw }}_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const {{ qw | upper }}_BLOCKS: u32 = {{ qwBlocks }}u;
 {% endif %}
+fn {{ qw }}_weight(expert: u32, row: u32, col: u32) -> f32 {
+  let base = expert * {{ qwRows }} + row;
 {% if packSize == 2 %}
+  let word = {{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col / 2u];
   let code = f32((word >> ((col & 1u) * 4u)) & 15u);
 {% else %}
+  let code = f32({{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col] & 255u);
 {% endif %}
 {% if not colWiseScales %}
+  return (code - MIDPOINT) * {{ qw }}_scales[base * {{ qw | upper }}_BLOCKS + col / QUANT_BLOCK];
 {% else %}
+  return (code - MIDPOINT) * {{ qw }}_scales[base];
 {% endif %}
 }
 // Decode two consecutive reduction-axis codes together. Four-bit codes share a
 // stored word; eight-bit codes occupy adjacent words.
+fn {{ qw }}_weight_pair(expert: u32, row: u32, col: u32) -> vec2<f32> {
+  let base = expert * {{ qwRows }} + row;
 {% if packSize == 2 %}
+  let word = {{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col / 2u];
   let lo = f32(word & 15u) - MIDPOINT;
   let hi = f32((word >> 4u) & 15u) - MIDPOINT;
 {% else %}
+  let lo = f32({{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col] & 255u) - MIDPOINT;
+  let hi = f32({{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col + 1u] & 255u) - MIDPOINT;
 {% endif %}
 {% if not colWiseScales %}
+  return vec2<f32>(lo * {{ qw }}_scales[base * {{ qw | upper }}_BLOCKS + col / QUANT_BLOCK],
+                   hi * {{ qw }}_scales[base * {{ qw | upper }}_BLOCKS + (col + 1u) / QUANT_BLOCK]);
 {% else %}
+  let scale = {{ qw }}_scales[base];
   return vec2<f32>(lo * scale, hi * scale);
 {% endif %}
 }
 {% macro rowlane_fold(arrays, lanes="LANES", lane="lane", slot="slot") %}
   var n = {{ lanes }};
   while (n > 1u) {
     }
     workgroupBarrier();
     n = half;
+  }{% endmacro %}
 var<workgroup> partial: array<f32, {{ decodeLanes * decodeRows }}>;
 @compute @workgroup_size(LANES, ROWS, 1)

build/webgpu/qmoe-output-grouped-sgmat.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 }}
 // Expert-grouped FC2 stage using subgroup-matrix operations. Routed slots are
@@ -17,51 +16,54 @@ diagnostic(off, chromium.subgroup_matrix_uniformity);
 // A loader gathers through `slot_list` and the store scatters back through it.
 const HIDDEN: u32 = {{ hidden }}u;
 const INTER: u32 = {{ inter }}u;
-const TILE_COLS: u32 = 64u;
 const TILE_K: u32 = 32u;
 const SUB_COLS: u32 = 32u;
 const SUB_ROWS: u32 = 16u;
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
-const FC2_PACKED_COLS: u32 = {{ fc2PackedCols }}u;
 {% if not colWiseScales %}
-// Block-wise scales (fc2_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
-const FC2_BLOCKS: u32 = {{ fc2ScaleBlocks }}u;
 {% endif %}
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
 // this helper requires `col` to be a multiple of four.
-fn fc2_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
-  let base = expert * HIDDEN + row;
 {% if packSize == 2 %}
-  let pair = base * FC2_PACKED_COLS + col / 2u;
-  let w0 = fc2_experts_weights[pair];
-  let w1 = fc2_experts_weights[pair + 1u];
   let codes = vec4<f32>(f32(w0 & 15u), f32((w0 >> 4u) & 15u),
                         f32(w1 & 15u), f32((w1 >> 4u) & 15u));
 {% else %}
-  let word = base * FC2_PACKED_COLS + col;
-  let codes = vec4<f32>(f32(fc2_experts_weights[word] & 255u), f32(fc2_experts_weights[word + 1u] & 255u),
-                        f32(fc2_experts_weights[word + 2u] & 255u), f32(fc2_experts_weights[word + 3u] & 255u));
 {% endif %}
 {% if not colWiseScales %}
-  let scale = fc2_scales[base * FC2_BLOCKS + col / QUANT_BLOCK];
 {% else %}
-  let scale = fc2_scales[base];
 {% endif %}
   return (codes - vec4<f32>(MIDPOINT)) * scale;
 }
 // The final K-loop barrier ends the A operand's lifetime. Its storage then
 // holds four 8x8 output banks per subgroup, fenced by the epilogue barriers.
 var<workgroup> tile_A: array<f32, {{ groupedSgmatStageElements }}>;
-var<workgroup> tile_B: array<f32, 64 * 32>;
 // Rows past the tile's end stage zero and are dropped at the store, so a short final
 // tile needs no separate arm.
@@ -141,9 +143,11 @@ fn main(
   let base_A = subtile_idy * SUB_ROWS;
   let base_B = subtile_idx * SUB_COLS;
-{% for row in [0, 1] %}{% for col in [0, 1, 2, 3] %}
   var matC{{ row }}{{ col }}: subgroup_matrix_result<f32, 8, 8>;
-{% endfor %}{% endfor %}
   for (var kidx = 0u; kidx < INTER; kidx = kidx + TILE_K) {
     loadSHMA(kidx, rows, local_idx / 4u, local_idx % 4u, a_slot);
@@ -151,7 +155,8 @@ fn main(
     workgroupBarrier();
     for (var step = 0u; step < TILE_K; step = step + 8u) {
-{% set operandScalar = "f32" %}{% set directInputs = directMatrixInputs is defined and directMatrixInputs %}
       let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
 {% for r in range(2) %}
       var matA{{ r }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ "w" if directInputs else "tile_A" }}, matrix_a_offset{% if r > 0 %} + 8u * {{ "K" if directInputs else "TILE_K" }}{% endif %}, {{ "K" if directInputs else "TILE_K" }});

 enable chromium_experimental_subgroup_matrix;
 diagnostic(off, chromium.subgroup_matrix_uniformity);
 {{ env.wgsl.resourceDeclarations }}
 // Expert-grouped FC2 stage using subgroup-matrix operations. Routed slots are
 // A loader gathers through `slot_list` and the store scatters back through it.
 const HIDDEN: u32 = {{ hidden }}u;
 const INTER: u32 = {{ inter }}u;
+const TILE_COLS: u32 = {{ sgmatTileCols }}u;
 const TILE_K: u32 = 32u;
 const SUB_COLS: u32 = 32u;
 const SUB_ROWS: u32 = 16u;
+{% set qw = "fc2" %}
+{% set qwRows = "HIDDEN" %}
+{% set qwPackedCols = fc2PackedCols %}
+{% set qwBlocks = fc2ScaleBlocks %}
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const {{ qw | upper }}_PACKED_COLS: u32 = {{ qwPackedCols }}u;
 {% if not colWiseScales %}
+// Block-wise scales ({{ qw }}_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const {{ qw | upper }}_BLOCKS: u32 = {{ qwBlocks }}u;
 {% endif %}
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
 // this helper requires `col` to be a multiple of four.
+fn {{ qw }}_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
+  let base = expert * {{ qwRows }} + row;
 {% if packSize == 2 %}
+  let pair = base * {{ qw | upper }}_PACKED_COLS + col / 2u;
+  let w0 = {{ qw }}_experts_weights[pair];
+  let w1 = {{ qw }}_experts_weights[pair + 1u];
   let codes = vec4<f32>(f32(w0 & 15u), f32((w0 >> 4u) & 15u),
                         f32(w1 & 15u), f32((w1 >> 4u) & 15u));
 {% else %}
+  let word = base * {{ qw | upper }}_PACKED_COLS + col;
+  let codes = vec4<f32>(f32({{ qw }}_experts_weights[word] & 255u), f32({{ qw }}_experts_weights[word + 1u] & 255u),
+                        f32({{ qw }}_experts_weights[word + 2u] & 255u), f32({{ qw }}_experts_weights[word + 3u] & 255u));
 {% endif %}
 {% if not colWiseScales %}
+  let scale = {{ qw }}_scales[base * {{ qw | upper }}_BLOCKS + col / QUANT_BLOCK];
 {% else %}
+  let scale = {{ qw }}_scales[base];
 {% endif %}
   return (codes - vec4<f32>(MIDPOINT)) * scale;
 }
 // The final K-loop barrier ends the A operand's lifetime. Its storage then
 // holds four 8x8 output banks per subgroup, fenced by the epilogue barriers.
 var<workgroup> tile_A: array<f32, {{ groupedSgmatStageElements }}>;
+var<workgroup> tile_B: array<f32, {{ sgmatTileCols }} * 32>;
 // Rows past the tile's end stage zero and are dropped at the store, so a short final
 // tile needs no separate arm.
   let base_A = subtile_idy * SUB_ROWS;
   let base_B = subtile_idx * SUB_COLS;
+{% for row in [0, 1] %}
+{% for col in [0, 1, 2, 3] %}
   var matC{{ row }}{{ col }}: subgroup_matrix_result<f32, 8, 8>;
+{% endfor %}
+{% endfor %}
   for (var kidx = 0u; kidx < INTER; kidx = kidx + TILE_K) {
     loadSHMA(kidx, rows, local_idx / 4u, local_idx % 4u, a_slot);
     workgroupBarrier();
     for (var step = 0u; step < TILE_K; step = step + 8u) {
+{% set operandScalar = "f32" %}
+{% set directInputs = false %}
       let matrix_a_offset = subtile_idy * SUB_ROWS * TILE_K + step;
 {% for r in range(2) %}
       var matA{{ r }}: subgroup_matrix_left<{{ operandScalar }}, 8, 8> = subgroupMatrixLoad<subgroup_matrix_left<{{ operandScalar }}, 8, 8>, row_major>(&{{ "w" if directInputs else "tile_A" }}, matrix_a_offset{% if r > 0 %} + 8u * {{ "K" if directInputs else "TILE_K" }}{% endif %}, {{ "K" if directInputs else "TILE_K" }});

build/webgpu/qmoe-output-grouped.wgsl.jinja CHANGED Viewed

@@ -13,57 +13,60 @@ const NTILE: u32 = {{ groupTileN }}u;
 const KTILE: u32 = {{ groupTileK }}u;
 const THREADS_SIDE: u32 = {{ groupThreads }}u;
 const THREADS: u32 = THREADS_SIDE * THREADS_SIDE;
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
-const FC2_PACKED_COLS: u32 = {{ fc2PackedCols }}u;
 {% if not colWiseScales %}
-// Block-wise scales (fc2_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
-const FC2_BLOCKS: u32 = {{ fc2ScaleBlocks }}u;
 {% endif %}
-fn fc2_weight(expert: u32, row: u32, col: u32) -> f32 {
-  let base = expert * HIDDEN + row;
 {% if packSize == 2 %}
-  let word = fc2_experts_weights[base * FC2_PACKED_COLS + col / 2u];
   let code = f32((word >> ((col & 1u) * 4u)) & 15u);
 {% else %}
-  let code = f32(fc2_experts_weights[base * FC2_PACKED_COLS + col] & 255u);
 {% endif %}
 {% if not colWiseScales %}
-  return (code - MIDPOINT) * fc2_scales[base * FC2_BLOCKS + col / QUANT_BLOCK];
 {% else %}
-  return (code - MIDPOINT) * fc2_scales[base];
 {% endif %}
 }
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
 // this helper requires `col` to be a multiple of four.
-fn fc2_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
-  let base = expert * HIDDEN + row;
 {% if packSize == 2 %}
-  let pair = base * FC2_PACKED_COLS + col / 2u;
-  let w0 = fc2_experts_weights[pair];
-  let w1 = fc2_experts_weights[pair + 1u];
   let codes = vec4<f32>(f32(w0 & 15u), f32((w0 >> 4u) & 15u),
                         f32(w1 & 15u), f32((w1 >> 4u) & 15u));
 {% else %}
-  let word = base * FC2_PACKED_COLS + col;
-  let codes = vec4<f32>(f32(fc2_experts_weights[word] & 255u), f32(fc2_experts_weights[word + 1u] & 255u),
-                        f32(fc2_experts_weights[word + 2u] & 255u), f32(fc2_experts_weights[word + 3u] & 255u));
 {% endif %}
 {% if not colWiseScales %}
-  let scale = fc2_scales[base * FC2_BLOCKS + col / QUANT_BLOCK];
 {% else %}
-  let scale = fc2_scales[base];
 {% endif %}
   return (codes - vec4<f32>(MIDPOINT)) * scale;
 }
 const KVEC: u32 = {{ groupTileKVec }}u;
 {% macro stage_group_tiles(aLoad, bLoad, b2Load, kExtent, nExtent, guarded, bLoadVec4="", b2LoadVec4="", aLoadVec4="") %}
       for (var idx = tid; idx < MTILE * KVEC; idx = idx + THREADS) {
@@ -116,16 +119,12 @@ const KVEC: u32 = {{ groupTileKVec }}u;
 {% if b2Load %}
         b2_tile[idx] = b2_vec;
 {% endif %}
-      }
-{%- endmacro %}
 {% macro group_tile_loop(aLoad, bLoad, b2Load, kExtent, nExtent, regM, regN, bLoadVec4="", b2LoadVec4="", aLoadVec4="") %}
 {% for r in range(regM) %}
 {% for c in range(regN) %}
   var acc_{{ r }}_{{ c }} = 0.0;
-{% if b2Load %}
-  var acc2_{{ r }}_{{ c }} = 0.0;
-{% endif %}
 {% endfor %}
 {% endfor %}
@@ -154,25 +153,17 @@ const KVEC: u32 = {{ groupTileKVec }}u;
 {% endfor %}
 {% for c in range(regN) %}
       let b{{ c }} = b_tile[(lid.x * {{ regN }}u + {{ c }}u) * KVEC + kv];
-{% if b2Load %}
-      let s{{ c }} = b2_tile[(lid.x * {{ regN }}u + {{ c }}u) * KVEC + kv];
-{% endif %}
 {% endfor %}
 {% for r in range(regM) %}
 {% for c in range(regN) %}
       acc_{{ r }}_{{ c }} = acc_{{ r }}_{{ c }} + dot(a{{ r }}, b{{ c }});
-{% if b2Load %}
-      acc2_{{ r }}_{{ c }} = acc2_{{ r }}_{{ c }} + dot(a{{ r }}, s{{ c }});
-{% endif %}
 {% endfor %}
 {% endfor %}
     }
     // Orders this step's tile reads before the next step overwrites them.
     workgroupBarrier();
     k_base = k_base + KTILE;
-  }
-{%- endmacro %}
 var<workgroup> row_slot: array<u32, {{ groupTileM }}>;
 var<workgroup> a_tile: array<vec4<f32>, {{ groupTileM * groupTileKVec }}>;
 var<workgroup> b_tile: array<vec4<f32>, {{ groupTileN * groupTileKVec }}>;

 const KTILE: u32 = {{ groupTileK }}u;
 const THREADS_SIDE: u32 = {{ groupThreads }}u;
 const THREADS: u32 = THREADS_SIDE * THREADS_SIDE;
+{% set qw = "fc2" %}
+{% set qwRows = "HIDDEN" %}
+{% set qwPackedCols = fc2PackedCols %}
+{% set qwBlocks = fc2ScaleBlocks %}
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const {{ qw | upper }}_PACKED_COLS: u32 = {{ qwPackedCols }}u;
 {% if not colWiseScales %}
+// Block-wise scales ({{ qw }}_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const {{ qw | upper }}_BLOCKS: u32 = {{ qwBlocks }}u;
 {% endif %}
+fn {{ qw }}_weight(expert: u32, row: u32, col: u32) -> f32 {
+  let base = expert * {{ qwRows }} + row;
 {% if packSize == 2 %}
+  let word = {{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col / 2u];
   let code = f32((word >> ((col & 1u) * 4u)) & 15u);
 {% else %}
+  let code = f32({{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col] & 255u);
 {% endif %}
 {% if not colWiseScales %}
+  return (code - MIDPOINT) * {{ qw }}_scales[base * {{ qw | upper }}_BLOCKS + col / QUANT_BLOCK];
 {% else %}
+  return (code - MIDPOINT) * {{ qw }}_scales[base];
 {% endif %}
 }
 // Decode the four-code group consumed by the staging loop. The aligned codes
 // share one quantization scale: blocks are powers of two of at least 16, and
 // this helper requires `col` to be a multiple of four.
+fn {{ qw }}_weight_vec4(expert: u32, row: u32, col: u32) -> vec4<f32> {
+  let base = expert * {{ qwRows }} + row;
 {% if packSize == 2 %}
+  let pair = base * {{ qw | upper }}_PACKED_COLS + col / 2u;
+  let w0 = {{ qw }}_experts_weights[pair];
+  let w1 = {{ qw }}_experts_weights[pair + 1u];
   let codes = vec4<f32>(f32(w0 & 15u), f32((w0 >> 4u) & 15u),
                         f32(w1 & 15u), f32((w1 >> 4u) & 15u));
 {% else %}
+  let word = base * {{ qw | upper }}_PACKED_COLS + col;
+  let codes = vec4<f32>(f32({{ qw }}_experts_weights[word] & 255u), f32({{ qw }}_experts_weights[word + 1u] & 255u),
+                        f32({{ qw }}_experts_weights[word + 2u] & 255u), f32({{ qw }}_experts_weights[word + 3u] & 255u));
 {% endif %}
 {% if not colWiseScales %}
+  let scale = {{ qw }}_scales[base * {{ qw | upper }}_BLOCKS + col / QUANT_BLOCK];
 {% else %}
+  let scale = {{ qw }}_scales[base];
 {% endif %}
   return (codes - vec4<f32>(MIDPOINT)) * scale;
 }
 const KVEC: u32 = {{ groupTileKVec }}u;
 {% macro stage_group_tiles(aLoad, bLoad, b2Load, kExtent, nExtent, guarded, bLoadVec4="", b2LoadVec4="", aLoadVec4="") %}
       for (var idx = tid; idx < MTILE * KVEC; idx = idx + THREADS) {
 {% if b2Load %}
         b2_tile[idx] = b2_vec;
 {% endif %}
+      }{% endmacro %}
 {% macro group_tile_loop(aLoad, bLoad, b2Load, kExtent, nExtent, regM, regN, bLoadVec4="", b2LoadVec4="", aLoadVec4="") %}
 {% for r in range(regM) %}
 {% for c in range(regN) %}
   var acc_{{ r }}_{{ c }} = 0.0;
 {% endfor %}
 {% endfor %}
 {% endfor %}
 {% for c in range(regN) %}
       let b{{ c }} = b_tile[(lid.x * {{ regN }}u + {{ c }}u) * KVEC + kv];
 {% endfor %}
 {% for r in range(regM) %}
 {% for c in range(regN) %}
       acc_{{ r }}_{{ c }} = acc_{{ r }}_{{ c }} + dot(a{{ r }}, b{{ c }});
 {% endfor %}
 {% endfor %}
     }
     // Orders this step's tile reads before the next step overwrites them.
     workgroupBarrier();
     k_base = k_base + KTILE;
+  }{% endmacro %}
 var<workgroup> row_slot: array<u32, {{ groupTileM }}>;
 var<workgroup> a_tile: array<vec4<f32>, {{ groupTileM * groupTileKVec }}>;
 var<workgroup> b_tile: array<vec4<f32>, {{ groupTileN * groupTileKVec }}>;

build/webgpu/qmoe-output-stage.wgsl.jinja CHANGED Viewed

@@ -1,4 +1,31 @@
 {{ env.wgsl.resourceDeclarations }}
 const HIDDEN: u32 = {{ hidden }}u;
 const INTER: u32 = {{ inter }}u;
@@ -6,39 +33,40 @@ const WG: u32 = {{ workgroupSize }}u;
 {% if topK > 1 %}
 const TOP_K: u32 = {{ topK }}u;
 {% endif %}
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
-const FC2_PACKED_COLS: u32 = {{ fc2PackedCols }}u;
 {% if not colWiseScales %}
-// Block-wise scales (fc2_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
-const FC2_BLOCKS: u32 = {{ fc2ScaleBlocks }}u;
 {% endif %}
-fn fc2_weight(expert: u32, row: u32, col: u32) -> f32 {
-  let base = expert * HIDDEN + row;
 {% if packSize == 2 %}
-  let word = fc2_experts_weights[base * FC2_PACKED_COLS + col / 2u];
   let code = f32((word >> ((col & 1u) * 4u)) & 15u);
 {% else %}
-  let code = f32(fc2_experts_weights[base * FC2_PACKED_COLS + col] & 255u);
 {% endif %}
 {% if not colWiseScales %}
-  return (code - MIDPOINT) * fc2_scales[base * FC2_BLOCKS + col / QUANT_BLOCK];
 {% else %}
-  return (code - MIDPOINT) * fc2_scales[base];
 {% endif %}
 }
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the per-axis dispatch fold width.
-  // Reduces to gid.x when the dispatch does not fold.
-  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   let total = params.tokenCount * HIDDEN;
   if (index >= total) {
     return;
@@ -55,11 +83,11 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   var acc = 0.0;
   for (var j = 0u; j < TOP_K; j = j + 1u) {
     let expert = route_expert[token * TOP_K + j];
-    var part = 0.0;
     for (var i: u32 = 0u; i < INTER; i = i + 1u) {
-      part = part + hidden_act[(local_token * TOP_K + j) * INTER + i] * fc2_weight(expert, out_col, i);
     }
-    acc = acc + part * route_mix[token * TOP_K + j];
   }
   output[token * HIDDEN + out_col] = acc;
@@ -67,11 +95,11 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   let expert = route_expert[token];
   let mix = route_mix[token];
-  var acc = 0.0;
   for (var i: u32 = 0u; i < INTER; i = i + 1u) {
-    acc = acc + hidden_act[local_token * INTER + i] * fc2_weight(expert, out_col, i);
   }
-  output[token * HIDDEN + out_col] = acc * mix;
 {% endif %}
 }

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
+{% set dotType = "f32" %}
+// Retain product and addition residuals across a dot product. Explicit fma
+// boundaries preserve the addition error transform under reassociation.
+struct DotAccumulator {
+  hi: {{ dotType }},
+  lo: {{ dotType }},
+}
+fn dot_accumulate(acc: DotAccumulator, a: {{ dotType }}, b: {{ dotType }}) -> DotAccumulator {
+  let product = fma(a, b, {{ dotType }}(0.0));
+  let productError = fma(a, b, -product);
+  let sum = fma(acc.hi, {{ dotType }}(1.0), product);
+  let bv = fma({{ dotType }}(-1.0), acc.hi, sum);
+  let av = fma({{ dotType }}(-1.0), bv, sum);
+  let ae = fma({{ dotType }}(-1.0), av, acc.hi);
+  let be = fma({{ dotType }}(-1.0), bv, product);
+  let error = fma(fma(fma(ae, {{ dotType }}(1.0), be), {{ dotType }}(1.0), acc.lo), {{ dotType }}(1.0), productError);
+  let hi = fma(sum, {{ dotType }}(1.0), error);
+  return DotAccumulator(hi, fma({{ dotType }}(-1.0), fma({{ dotType }}(-1.0), sum, hi), error));
+}
+fn dot_value(acc: DotAccumulator) -> {{ dotType }} {
+  return fma(acc.hi, {{ dotType }}(1.0), acc.lo);
+}
 const HIDDEN: u32 = {{ hidden }}u;
 const INTER: u32 = {{ inter }}u;
 {% if topK > 1 %}
 const TOP_K: u32 = {{ topK }}u;
 {% endif %}
+{% set qw = "fc2" %}
+{% set qwRows = "HIDDEN" %}
+{% set qwPackedCols = fc2PackedCols %}
+{% set qwBlocks = fc2ScaleBlocks %}
 // Quantized expert weights use unsigned codes offset by 2^(bits-1), so a
 // weight is (code - MIDPOINT) * scale. Four-bit codes are packed low-nibble
 // first along the reduction axis, two codes per stored word.
 const MIDPOINT: f32 = f32({{ quantMidpoint }});
+const {{ qw | upper }}_PACKED_COLS: u32 = {{ qwPackedCols }}u;
 {% if not colWiseScales %}
+// Block-wise scales ({{ qw }}_scales rank 3): one scale per QUANT_BLOCK columns of the reduction
 // axis. Column-wise scales (block_size 0) are the rank-2 single-scale-per-row case below.
 const QUANT_BLOCK: u32 = {{ quantBlockSize }}u;
+const {{ qw | upper }}_BLOCKS: u32 = {{ qwBlocks }}u;
 {% endif %}
+fn {{ qw }}_weight(expert: u32, row: u32, col: u32) -> f32 {
+  let base = expert * {{ qwRows }} + row;
 {% if packSize == 2 %}
+  let word = {{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col / 2u];
   let code = f32((word >> ((col & 1u) * 4u)) & 15u);
 {% else %}
+  let code = f32({{ qw }}_experts_weights[base * {{ qw | upper }}_PACKED_COLS + col] & 255u);
 {% endif %}
 {% if not colWiseScales %}
+  return (code - MIDPOINT) * {{ qw }}_scales[base * {{ qw | upper }}_BLOCKS + col / QUANT_BLOCK];
 {% else %}
+  return (code - MIDPOINT) * {{ qw }}_scales[base];
 {% endif %}
 }
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WG", "index", "") }}
   let total = params.tokenCount * HIDDEN;
   if (index >= total) {
     return;
   var acc = 0.0;
   for (var j = 0u; j < TOP_K; j = j + 1u) {
     let expert = route_expert[token * TOP_K + j];
+    var part = DotAccumulator(0.0, 0.0);
     for (var i: u32 = 0u; i < INTER; i = i + 1u) {
+      part = dot_accumulate(part, hidden_act[(local_token * TOP_K + j) * INTER + i], fc2_weight(expert, out_col, i));
     }
+    acc = acc + dot_value(part) * route_mix[token * TOP_K + j];
   }
   output[token * HIDDEN + out_col] = acc;
   let expert = route_expert[token];
   let mix = route_mix[token];
+  var acc = DotAccumulator(0.0, 0.0);
   for (var i: u32 = 0u; i < INTER; i = i + 1u) {
+    acc = dot_accumulate(acc, hidden_act[local_token * INTER + i], fc2_weight(expert, out_col, i));
   }
+  output[token * HIDDEN + out_col] = dot_value(acc) * mix;
 {% endif %}
 }

build/webgpu/qmoe-route-stage.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,13 @@
 {{ env.wgsl.resourceDeclarations }}
-const TOKENS: u32 = {{ tokens }}u;
 const EXPERTS: u32 = {{ experts }}u;
 const WG: u32 = {{ workgroupSize }}u;
 {% if topK > 1 %}
@@ -54,7 +61,6 @@ fn topk_route(token: u32) -> TopKRoute {
   }
   return result;
 }
 {% else %}
 // Serial top-1 selection over `(logit, expert_index)` pairs ordered descending,
 // so equal logits select the higher expert index.
@@ -75,17 +81,11 @@ fn top1_route(token: u32) -> Top1Route {
   }
   return result;
 }
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
-  // 2D-folded flat index: gid.y carries the high bits past the per-axis dispatch fold width.
-  // Reduces to gid.x when the dispatch does not fold.
-  let token = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
-  if (token >= TOKENS) {
-    return;
-  }
 {% if topK > 1 %}
   // Slot j of each token's TOP_K-wide run holds its j-th ranked expert and

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 const EXPERTS: u32 = {{ experts }}u;
 const WG: u32 = {{ workgroupSize }}u;
 {% if topK > 1 %}
   }
   return result;
 }
 {% else %}
 // Serial top-1 selection over `(logit, expert_index)` pairs ordered descending,
 // so equal logits select the higher expert index.
   }
   return result;
 }
 {% endif %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WG", "token", "params.tokenCount") }}
 {% if topK > 1 %}
   // Slot j of each token's TOP_K-wide run holds its j-th ranked expert and

build/webgpu/test.json CHANGED Viewed

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