| { |
| "domain": "com.microsoft", |
| "name": "MoE", |
| "sinceVersion": 1, |
| "description": "Mixture of Experts: applies softmax to `router_probs`, routes each token to the top-`k` experts, applies FC1 and `activation_type`, projects through FC2, then sums the selected outputs using their routing probabilities. SwiGLU takes its operands from a separate FC3 (`swiglu_fusion` 0) or a fused FC1 in interleaved (1) or concatenated (2) order; SiLU may also use FC3 as its multiplicative linear projection. This inference package supports float32 and dense routing (`use_sparse_mixer = 0`); float16, bfloat16, and sparse mixing are not implemented. Quantized weights use `com.microsoft.QMoE`.", |
| "inputs": [ |
| { |
| "role": "input", |
| "dtype": "T", |
| "description": "Token activations, either 2D `(num_tokens, hidden_size)` or 3D `(batch_size, sequence_length, hidden_size)`." |
| }, |
| { |
| "role": "router_probs", |
| "dtype": "T", |
| "rank": 2, |
| "description": "2D router logits of shape `(num_tokens, num_experts)`, where `num_tokens` is the product of every leading dimension of `input`. Despite the historical port name, the operator applies a full softmax before top-k selection." |
| }, |
| { |
| "role": "fc1_experts_weights", |
| "dtype": "T", |
| "rank": 3, |
| "description": "3D first-layer expert weights of shape `(num_experts, fusion_size * inter_size, hidden_size)`, where `fusion_size` is 2 for fused SwiGLU (`swiglu_fusion` 1 or 2) and 1 otherwise." |
| }, |
| { |
| "role": "fc1_experts_bias", |
| "dtype": "T", |
| "rank": 2, |
| "optional": true, |
| "description": "Optional 2D FC1 bias of shape `(num_experts, fusion_size * inter_size)`." |
| }, |
| { |
| "role": "fc2_experts_weights", |
| "dtype": "T", |
| "rank": 3, |
| "description": "3D second-layer expert weights of shape `(num_experts, hidden_size, inter_size)`." |
| }, |
| { |
| "role": "fc2_experts_bias", |
| "dtype": "T", |
| "rank": 2, |
| "optional": true, |
| "description": "Optional 2D FC2 bias of shape `(num_experts, hidden_size)`, added per expert before that expert's routing weight is applied." |
| }, |
| { |
| "role": "fc3_experts_weights", |
| "dtype": "T", |
| "rank": 3, |
| "optional": true, |
| "description": "Optional 3D third-layer expert weights of shape `(num_experts, inter_size, hidden_size)`. It supplies the separate linear operand for SwiGLU when `swiglu_fusion` is 0, or the multiplicative linear projection for SiLU gating. Other activations do not consume FC3." |
| }, |
| { |
| "role": "fc3_experts_bias", |
| "dtype": "T", |
| "rank": 2, |
| "optional": true, |
| "description": "Optional 2D FC3 bias of shape `(num_experts, inter_size)`." |
| } |
| ], |
| "outputs": [ |
| { |
| "role": "output", |
| "dtype": "T", |
| "rank": "ranks.input", |
| "shape": "shapes.input", |
| "description": "Routed expert output with the same shape as `input`." |
| } |
| ], |
| "attributes": { |
| "activation_alpha": 1, |
| "activation_beta": 0, |
| "activation_type": "relu", |
| "k": 1, |
| "normalize_routing_weights": 0, |
| "swiglu_fusion": 0, |
| "use_sparse_mixer": 0 |
| }, |
| "attributeConstraints": { |
| "activation_type": { "values": ["relu", "gelu", "silu", "swiglu", "identity"] }, |
| "normalize_routing_weights": { "values": [0, 1] }, |
| "swiglu_fusion": { "values": [0, 1, 2] }, |
| "use_sparse_mixer": { "values": [0] } |
| }, |
| "attributeDescriptions": { |
| "activation_alpha": "Alpha parameter used by the activation; the schema default is 1.", |
| "activation_beta": "Beta parameter used by the activation; the schema default is 0.", |
| "activation_type": "Activation applied to the FC1 projection: `relu`, `gelu`, `silu`, `swiglu`, or `identity`. The schema default is `relu`.", |
| "k": "Number of experts selected per token; the schema default is 1.", |
| "normalize_routing_weights": "Whether to normalize the selected routing weights; the schema default is 0.", |
| "swiglu_fusion": "0 keeps the SwiGLU operands in separate FC1/FC3 GEMMs, 1 interleaves them in one FC1 row, and 2 concatenates them. The schema default is 0.", |
| "swiglu_limit": "Optional SwiGLU clamp limit; omission means no clamp.", |
| "use_sparse_mixer": "Whether to use sparse-mixer routing. The standard default and only supported value is 0." |
| }, |
| "typeConstraints": { "T": ["float32"] }, |
| "args": { |
| "inputT": { "kind": "tensor", "semantic": "input", "role": "input" }, |
| "routerT": { "kind": "tensor", "semantic": "router_probs", "role": "input" }, |
| "fc1T": { "kind": "tensor", "semantic": "fc1_experts_weights", "role": "input" }, |
| "fc1BiasT": { "kind": "tensor", "semantic": "fc1_experts_bias", "role": "input", "required": false }, |
| "fc2T": { "kind": "tensor", "semantic": "fc2_experts_weights", "role": "input" }, |
| "fc2BiasT": { "kind": "tensor", "semantic": "fc2_experts_bias", "role": "input", "required": false }, |
| "fc3T": { "kind": "tensor", "semantic": "fc3_experts_weights", "role": "input", "required": false }, |
| "fc3BiasT": { "kind": "tensor", "semantic": "fc3_experts_bias", "role": "input", "required": false }, |
| "outputT": { "kind": "tensor", "semantic": "output", "role": "output" } |
| }, |
| "tunables": { |
| "workgroupSize": 64, |
| "decodeLanes": 32, |
| "decodeBlockTarget": 1024, |
| "decodeMinLaneTrips": 4, |
| "groupThreads": 8, |
| "groupRegM": 4, |
| "groupRegN": 4, |
| "groupTileK": 16, |
| "groupRouteWorkgroup": 256 |
| }, |
| "derive": { |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "topK": "attrs.k", |
| "hiddenDim": "dim(shapes.inputT, ranks.inputT - 1)", |
| "tokenCount": "numel(shapes.inputT) / max(1, hiddenDim)", |
| "expertCount": "dim(shapes.routerT, 1)", |
| "activationType": "attrs.activation_type", |
| "isSwiglu": "activationType == \"swiglu\"", |
| "fusionMode": "attrs.swiglu_fusion", |
| "fusionSize": "2 if isSwiglu and fusionMode != 0 else 1", |
| "interSize": "dim(shapes.fc1T, 1) / max(1, fusionSize)", |
| "fusionOk": "(fusionMode == 0 or fusionMode == 1 or fusionMode == 2) if isSwiglu else fusionMode == 0", |
| "fc3ShapeOk": "not present.fc3T or (ranks.fc3T == 3 and dim(shapes.fc3T, 0) == expertCount and dim(shapes.fc3T, 1) == interSize and dim(shapes.fc3T, 2) == hiddenDim)", |
| "fc3PresenceOk": "(present.fc3T if fusionMode == 0 else not present.fc3T) if isSwiglu else (not present.fc3T or activationType == \"silu\")", |
| "fc1BiasOk": "not present.fc1BiasT or (ranks.fc1BiasT == 2 and dim(shapes.fc1BiasT, 0) == expertCount and dim(shapes.fc1BiasT, 1) == dim(shapes.fc1T, 1))", |
| "fc2BiasOk": "not present.fc2BiasT or (ranks.fc2BiasT == 2 and dim(shapes.fc2BiasT, 0) == expertCount and dim(shapes.fc2BiasT, 1) == hiddenDim)", |
| "fc3BiasOk": "not present.fc3BiasT or (present.fc3T and ranks.fc3BiasT == 2 and dim(shapes.fc3BiasT, 0) == expertCount and dim(shapes.fc3BiasT, 1) == interSize)", |
| "shapeOk": "(ranks.inputT == 2 or ranks.inputT == 3) and ranks.outputT == ranks.inputT and sameShape(shapes.inputT, shapes.outputT) and ranks.routerT == 2 and ranks.fc1T == 3 and ranks.fc2T == 3 and hiddenDim > 0 and interSize > 0 and dim(shapes.routerT, 0) == tokenCount and dim(shapes.fc1T, 0) == expertCount and dim(shapes.fc2T, 0) == expertCount and dim(shapes.fc1T, 2) == hiddenDim and dim(shapes.fc1T, 1) % max(1, fusionSize) == 0 and dim(shapes.fc2T, 1) == hiddenDim and dim(shapes.fc2T, 2) == interSize and topK >= 1 and topK <= expertCount", |
| "contractOk": "shapeOk and fusionOk and fc3ShapeOk and fc3PresenceOk and fc1BiasOk and fc2BiasOk and fc3BiasOk", |
| "workgroupSizeOk": "tunables.workgroupSize >= 1 and tunables.workgroupSize <= deviceWorkgroupCap", |
| "hiddenChunkBytes": "topK * interSize * 4", |
| "hiddenChunkFits": "hiddenChunkBytes <= device.limits.maxStorageBufferBindingSize and hiddenChunkBytes <= device.limits.maxBufferSize", |
| "hiddenChunkTokens": "min(tokenCount, max(1, floor(min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize) / max(1, hiddenChunkBytes))))", |
| "hiddenChunkCount": "max(1, ceilDiv(tokenCount, max(1, hiddenChunkTokens)))", |
| "routeScratchBytes": "tokenCount * topK * 4", |
| "routeScratchFits": "routeScratchBytes <= device.limits.maxStorageBufferBindingSize and routeScratchBytes <= device.limits.maxBufferSize", |
| "splitContract": "contractOk and workgroupSizeOk and hiddenChunkFits and routeScratchFits", |
| "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)) + expertCount", |
| "groupSharedBytes": "(groupTileM * tunables.groupTileK + 2 * groupTileN * tunables.groupTileK + groupTileM) * 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 expertCount * 8 <= device.limits.maxComputeWorkgroupStorageSize and tunables.groupTileK % 4 == 0", |
| "groupedDispatchOk": "groupMaxTiles <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(interSize, max(1, groupTileN)) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(hiddenDim, max(1, groupTileN)) <= device.limits.maxComputeWorkgroupsPerDimension", |
| "groupedShapeOk": "tokenCount * topK * 4 >= groupTileM * expertCount", |
| "groupSlotOutBytes": "hiddenChunkTokens * topK * hiddenDim * 4", |
| "groupSlotOutFits": "groupSlotOutBytes <= device.limits.maxStorageBufferBindingSize and groupSlotOutBytes <= device.limits.maxBufferSize", |
| "groupedContract": "splitContract and groupedDeviceOk and groupedDispatchOk and groupSlotOutFits and groupedShapeOk", |
| "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) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(hiddenDim, decodeRows) <= device.limits.maxComputeWorkgroupsPerDimension and hiddenChunkTokens * topK <= device.limits.maxComputeWorkgroupsPerDimension", |
| "decodeLaneDepth": "min(hiddenDim, interSize) / max(1, tunables.decodeLanes)", |
| "decodeDepthOk": "decodeLaneDepth >= tunables.decodeMinLaneTrips", |
| "splitFfnBlocks": "ceilDiv(tokenCount * topK * interSize, max(1, tunables.workgroupSize))", |
| "decodeOccupancyOk": "splitFfnBlocks <= tunables.decodeBlockTarget", |
| "decodeContract": "splitContract and decodeDeviceOk and decodeDispatchOk and decodeDepthOk and (decodeOccupancyOk or not groupedContract)", |
| "groupedVec4Ok": "hiddenDim % 4 == 0 and interSize % 4 == 0" |
| }, |
| "constants": { |
| "hidden": "hiddenDim", |
| "inter": "interSize", |
| "fc1Rows": "dim(shapes.fc1T, 1)", |
| "experts": "expertCount", |
| "topK": "topK", |
| "tokens": "tokenCount", |
| "activation": "activationType", |
| "swigluFusion": "fusionMode", |
| "hasSwigluLimit": "has(attrs, \"swiglu_limit\")", |
| "swigluLimit": "attrs.swiglu_limit if has(attrs, \"swiglu_limit\") else 0", |
| "hasFc1Bias": "present.fc1BiasT", |
| "hasFc2Bias": "present.fc2BiasT", |
| "hasFc3": "present.fc3T", |
| "hasFc3Bias": "present.fc3BiasT", |
| "workgroupSize": "tunables.workgroupSize" |
| }, |
| "bindingSets": { |
| "route": [ |
| { |
| "name": "router_probs", |
| "arg": "routerT", |
| "semantic": "router_probs", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "route_expert", "semantic": "routeExpert", "buffer": { "type": "storage" }, "elementType": "u32" }, |
| { "name": "route_mix", "semantic": "routeMix", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [{ "name": "normalizeRoutingWeights", "type": "u32", "value": "attrs.normalize_routing_weights" }] |
| } |
| } |
| ], |
| "ffn_plain_none": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }, |
| { |
| "name": "tokenCount", |
| "type": "u32", |
| "value": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| ] |
| } |
| } |
| ], |
| "ffn_plain_plain": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }, |
| { |
| "name": "tokenCount", |
| "type": "u32", |
| "value": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| ] |
| } |
| } |
| ], |
| "ffn_plain_biased": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_bias", |
| "arg": "fc3BiasT", |
| "semantic": "fc3_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }, |
| { |
| "name": "tokenCount", |
| "type": "u32", |
| "value": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| ] |
| } |
| } |
| ], |
| "ffn_bias_none": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc1_experts_bias", |
| "arg": "fc1BiasT", |
| "semantic": "fc1_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }, |
| { |
| "name": "tokenCount", |
| "type": "u32", |
| "value": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| ] |
| } |
| } |
| ], |
| "ffn_bias_plain": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc1_experts_bias", |
| "arg": "fc1BiasT", |
| "semantic": "fc1_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }, |
| { |
| "name": "tokenCount", |
| "type": "u32", |
| "value": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| ] |
| } |
| } |
| ], |
| "ffn_bias_biased": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc1_experts_bias", |
| "arg": "fc1BiasT", |
| "semantic": "fc1_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_bias", |
| "arg": "fc3BiasT", |
| "semantic": "fc3_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }, |
| { |
| "name": "tokenCount", |
| "type": "u32", |
| "value": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| ] |
| } |
| } |
| ], |
| "out_plain": [ |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { "name": "route_mix", "semantic": "routeMix", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "fc2_experts_weights", |
| "arg": "fc2T", |
| "semantic": "fc2_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }, |
| { |
| "name": "tokenCount", |
| "type": "u32", |
| "value": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| ] |
| } |
| } |
| ], |
| "out_bias": [ |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { "name": "route_mix", "semantic": "routeMix", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "fc2_experts_weights", |
| "arg": "fc2T", |
| "semantic": "fc2_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc2_experts_bias", |
| "arg": "fc2BiasT", |
| "semantic": "fc2_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }, |
| { |
| "name": "tokenCount", |
| "type": "u32", |
| "value": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| ] |
| } |
| } |
| ], |
| "group": [ |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { "name": "slot_list", "semantic": "slotList", "buffer": { "type": "storage" }, "elementType": "u32" }, |
| { "name": "tile_meta", "semantic": "tileMeta", "buffer": { "type": "storage" }, "elementType": "u32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }, |
| { |
| "name": "tokenCount", |
| "type": "u32", |
| "value": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| ] |
| } |
| } |
| ], |
| "ffng_plain_none": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { "name": "slot_list", "semantic": "slotList", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { "name": "tile_meta", "semantic": "tileMeta", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" } |
| ] |
| } |
| } |
| ], |
| "ffng_plain_plain": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { "name": "slot_list", "semantic": "slotList", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { "name": "tile_meta", "semantic": "tileMeta", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" } |
| ] |
| } |
| } |
| ], |
| "ffng_plain_biased": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { "name": "slot_list", "semantic": "slotList", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { "name": "tile_meta", "semantic": "tileMeta", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { |
| "name": "fc3_experts_bias", |
| "arg": "fc3BiasT", |
| "semantic": "fc3_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" } |
| ] |
| } |
| } |
| ], |
| "ffng_bias_none": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { "name": "slot_list", "semantic": "slotList", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { "name": "tile_meta", "semantic": "tileMeta", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { |
| "name": "fc1_experts_bias", |
| "arg": "fc1BiasT", |
| "semantic": "fc1_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" } |
| ] |
| } |
| } |
| ], |
| "ffng_bias_plain": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { "name": "slot_list", "semantic": "slotList", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { "name": "tile_meta", "semantic": "tileMeta", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { |
| "name": "fc1_experts_bias", |
| "arg": "fc1BiasT", |
| "semantic": "fc1_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" } |
| ] |
| } |
| } |
| ], |
| "ffng_bias_biased": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { "name": "slot_list", "semantic": "slotList", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { "name": "tile_meta", "semantic": "tileMeta", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { |
| "name": "fc1_experts_bias", |
| "arg": "fc1BiasT", |
| "semantic": "fc1_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { |
| "name": "fc3_experts_bias", |
| "arg": "fc3BiasT", |
| "semantic": "fc3_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" } |
| ] |
| } |
| } |
| ], |
| "outg_plain": [ |
| { |
| "name": "hidden_act", |
| "semantic": "hiddenAct", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { "name": "slot_list", "semantic": "slotList", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { "name": "tile_meta", "semantic": "tileMeta", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { |
| "name": "fc2_experts_weights", |
| "arg": "fc2T", |
| "semantic": "fc2_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { "name": "slot_out", "semantic": "slotOut", "buffer": { "type": "storage" }, "elementType": "f32" } |
| ], |
| "outg_bias": [ |
| { |
| "name": "hidden_act", |
| "semantic": "hiddenAct", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { "name": "slot_list", "semantic": "slotList", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { "name": "tile_meta", "semantic": "tileMeta", "buffer": { "type": "read-only-storage" }, "elementType": "u32" }, |
| { |
| "name": "fc2_experts_weights", |
| "arg": "fc2T", |
| "semantic": "fc2_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$groupElem" |
| }, |
| { |
| "name": "fc2_experts_bias", |
| "arg": "fc2BiasT", |
| "semantic": "fc2_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "slot_out", "semantic": "slotOut", "buffer": { "type": "storage" }, "elementType": "f32" } |
| ], |
| "mix": [ |
| { "name": "slot_out", "semantic": "slotOut", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "name": "route_mix", "semantic": "routeMix", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }, |
| { |
| "name": "tokenCount", |
| "type": "u32", |
| "value": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| ] |
| } |
| } |
| ], |
| "ffn_plain_none_io": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" } |
| ], |
| "ffn_plain_none_no_count": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" } |
| ] |
| } |
| } |
| ], |
| "ffn_plain_plain_io": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" } |
| ], |
| "ffn_plain_plain_no_count": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" } |
| ] |
| } |
| } |
| ], |
| "ffn_plain_biased_io": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_bias", |
| "arg": "fc3BiasT", |
| "semantic": "fc3_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" } |
| ], |
| "ffn_plain_biased_no_count": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_bias", |
| "arg": "fc3BiasT", |
| "semantic": "fc3_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" } |
| ] |
| } |
| } |
| ], |
| "ffn_bias_none_io": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc1_experts_bias", |
| "arg": "fc1BiasT", |
| "semantic": "fc1_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" } |
| ], |
| "ffn_bias_none_no_count": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc1_experts_bias", |
| "arg": "fc1BiasT", |
| "semantic": "fc1_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" } |
| ] |
| } |
| } |
| ], |
| "ffn_bias_plain_io": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc1_experts_bias", |
| "arg": "fc1BiasT", |
| "semantic": "fc1_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" } |
| ], |
| "ffn_bias_plain_no_count": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc1_experts_bias", |
| "arg": "fc1BiasT", |
| "semantic": "fc1_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" } |
| ] |
| } |
| } |
| ], |
| "ffn_bias_biased_io": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc1_experts_bias", |
| "arg": "fc1BiasT", |
| "semantic": "fc1_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_bias", |
| "arg": "fc3BiasT", |
| "semantic": "fc3_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" } |
| ], |
| "ffn_bias_biased_no_count": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { |
| "name": "fc1_experts_weights", |
| "arg": "fc1T", |
| "semantic": "fc1_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc1_experts_bias", |
| "arg": "fc1BiasT", |
| "semantic": "fc1_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_weights", |
| "arg": "fc3T", |
| "semantic": "fc3_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc3_experts_bias", |
| "arg": "fc3BiasT", |
| "semantic": "fc3_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "activationAlpha", "type": "f32", "value": "attrs.activation_alpha" }, |
| { "name": "activationBeta", "type": "f32", "value": "attrs.activation_beta" }, |
| { "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" } |
| ] |
| } |
| } |
| ], |
| "out_plain_io": [ |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { "name": "route_mix", "semantic": "routeMix", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "fc2_experts_weights", |
| "arg": "fc2T", |
| "semantic": "fc2_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| } |
| ], |
| "out_plain_no_count": [ |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { "name": "route_mix", "semantic": "routeMix", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "fc2_experts_weights", |
| "arg": "fc2T", |
| "semantic": "fc2_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [{ "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }] |
| } |
| } |
| ], |
| "out_bias_io": [ |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { "name": "route_mix", "semantic": "routeMix", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "fc2_experts_weights", |
| "arg": "fc2T", |
| "semantic": "fc2_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc2_experts_bias", |
| "arg": "fc2BiasT", |
| "semantic": "fc2_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| } |
| ], |
| "out_bias_no_count": [ |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "route_expert", |
| "semantic": "routeExpert", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "u32" |
| }, |
| { "name": "route_mix", "semantic": "routeMix", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { |
| "name": "fc2_experts_weights", |
| "arg": "fc2T", |
| "semantic": "fc2_experts_weights", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "fc2_experts_bias", |
| "arg": "fc2BiasT", |
| "semantic": "fc2_experts_bias", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [{ "name": "tokenOffset", "type": "u32", "value": "repeat.chunk * hiddenChunkTokens" }] |
| } |
| } |
| ] |
| }, |
| "variants": [ |
| { |
| "id": "split_routed_fc1plain_fc3none_fc2plain", |
| "priority": 10, |
| "when": ["splitContract", "not present.fc1BiasT", "not present.fc3T", "not present.fc2BiasT"], |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-stage.wgsl.jinja", |
| "bindings": "ffn_plain_none", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK * interSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-stage.wgsl.jinja", |
| "bindings": "out_plain", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "gemv_routed_fc1plain_fc3none_fc2plain", |
| "priority": 20, |
| "when": ["decodeContract", "not present.fc1BiasT", "not present.fc3T", "not present.fc2BiasT"], |
| "constants": { "decodeLanes": "tunables.decodeLanes", "decodeRows": "decodeRows" }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-gemv.wgsl.jinja", |
| "bindings": "ffn_plain_none_no_count", |
| "dispatch": { |
| "x": "ceilDiv(interSize, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-gemv.wgsl.jinja", |
| "bindings": "out_plain_no_count", |
| "dispatch": { |
| "x": "ceilDiv(hiddenDim, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "split_routed_fc1plain_fc3none_fc2bias", |
| "priority": 10, |
| "when": ["splitContract", "not present.fc1BiasT", "not present.fc3T", "present.fc2BiasT"], |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-stage.wgsl.jinja", |
| "bindings": "ffn_plain_none", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK * interSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-stage.wgsl.jinja", |
| "bindings": "out_bias", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "gemv_routed_fc1plain_fc3none_fc2bias", |
| "priority": 20, |
| "when": ["decodeContract", "not present.fc1BiasT", "not present.fc3T", "present.fc2BiasT"], |
| "constants": { "decodeLanes": "tunables.decodeLanes", "decodeRows": "decodeRows" }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-gemv.wgsl.jinja", |
| "bindings": "ffn_plain_none_no_count", |
| "dispatch": { |
| "x": "ceilDiv(interSize, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-gemv.wgsl.jinja", |
| "bindings": "out_bias_no_count", |
| "dispatch": { |
| "x": "ceilDiv(hiddenDim, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "split_routed_fc1plain_fc3plain_fc2plain", |
| "priority": 10, |
| "when": ["splitContract", "not present.fc1BiasT", "present.fc3T and not present.fc3BiasT", "not present.fc2BiasT"], |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-stage.wgsl.jinja", |
| "bindings": "ffn_plain_plain", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK * interSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-stage.wgsl.jinja", |
| "bindings": "out_plain", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "gemv_routed_fc1plain_fc3plain_fc2plain", |
| "priority": 20, |
| "when": ["decodeContract", "not present.fc1BiasT", "present.fc3T and not present.fc3BiasT", "not present.fc2BiasT"], |
| "constants": { "decodeLanes": "tunables.decodeLanes", "decodeRows": "decodeRows" }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-gemv.wgsl.jinja", |
| "bindings": "ffn_plain_plain_no_count", |
| "dispatch": { |
| "x": "ceilDiv(interSize, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-gemv.wgsl.jinja", |
| "bindings": "out_plain_no_count", |
| "dispatch": { |
| "x": "ceilDiv(hiddenDim, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "split_routed_fc1plain_fc3plain_fc2bias", |
| "priority": 10, |
| "when": ["splitContract", "not present.fc1BiasT", "present.fc3T and not present.fc3BiasT", "present.fc2BiasT"], |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-stage.wgsl.jinja", |
| "bindings": "ffn_plain_plain", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK * interSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-stage.wgsl.jinja", |
| "bindings": "out_bias", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "gemv_routed_fc1plain_fc3plain_fc2bias", |
| "priority": 20, |
| "when": ["decodeContract", "not present.fc1BiasT", "present.fc3T and not present.fc3BiasT", "present.fc2BiasT"], |
| "constants": { "decodeLanes": "tunables.decodeLanes", "decodeRows": "decodeRows" }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-gemv.wgsl.jinja", |
| "bindings": "ffn_plain_plain_no_count", |
| "dispatch": { |
| "x": "ceilDiv(interSize, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-gemv.wgsl.jinja", |
| "bindings": "out_bias_no_count", |
| "dispatch": { |
| "x": "ceilDiv(hiddenDim, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "split_routed_fc1plain_fc3biased_fc2plain", |
| "priority": 10, |
| "when": ["splitContract", "not present.fc1BiasT", "present.fc3BiasT", "not present.fc2BiasT"], |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-stage.wgsl.jinja", |
| "bindings": "ffn_plain_biased", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK * interSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-stage.wgsl.jinja", |
| "bindings": "out_plain", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "gemv_routed_fc1plain_fc3biased_fc2plain", |
| "priority": 20, |
| "when": ["decodeContract", "not present.fc1BiasT", "present.fc3BiasT", "not present.fc2BiasT"], |
| "constants": { "decodeLanes": "tunables.decodeLanes", "decodeRows": "decodeRows" }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-gemv.wgsl.jinja", |
| "bindings": "ffn_plain_biased_no_count", |
| "dispatch": { |
| "x": "ceilDiv(interSize, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-gemv.wgsl.jinja", |
| "bindings": "out_plain_no_count", |
| "dispatch": { |
| "x": "ceilDiv(hiddenDim, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "split_routed_fc1plain_fc3biased_fc2bias", |
| "priority": 10, |
| "when": ["splitContract", "not present.fc1BiasT", "present.fc3BiasT", "present.fc2BiasT"], |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-stage.wgsl.jinja", |
| "bindings": "ffn_plain_biased", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK * interSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-stage.wgsl.jinja", |
| "bindings": "out_bias", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "gemv_routed_fc1plain_fc3biased_fc2bias", |
| "priority": 20, |
| "when": ["decodeContract", "not present.fc1BiasT", "present.fc3BiasT", "present.fc2BiasT"], |
| "constants": { "decodeLanes": "tunables.decodeLanes", "decodeRows": "decodeRows" }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-gemv.wgsl.jinja", |
| "bindings": "ffn_plain_biased_no_count", |
| "dispatch": { |
| "x": "ceilDiv(interSize, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-gemv.wgsl.jinja", |
| "bindings": "out_bias_no_count", |
| "dispatch": { |
| "x": "ceilDiv(hiddenDim, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "split_routed_fc1bias_fc3none_fc2plain", |
| "priority": 10, |
| "when": ["splitContract", "present.fc1BiasT", "not present.fc3T", "not present.fc2BiasT"], |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-stage.wgsl.jinja", |
| "bindings": "ffn_bias_none", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK * interSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-stage.wgsl.jinja", |
| "bindings": "out_plain", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "gemv_routed_fc1bias_fc3none_fc2plain", |
| "priority": 20, |
| "when": ["decodeContract", "present.fc1BiasT", "not present.fc3T", "not present.fc2BiasT"], |
| "constants": { "decodeLanes": "tunables.decodeLanes", "decodeRows": "decodeRows" }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-gemv.wgsl.jinja", |
| "bindings": "ffn_bias_none_no_count", |
| "dispatch": { |
| "x": "ceilDiv(interSize, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-gemv.wgsl.jinja", |
| "bindings": "out_plain_no_count", |
| "dispatch": { |
| "x": "ceilDiv(hiddenDim, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "split_routed_fc1bias_fc3none_fc2bias", |
| "priority": 10, |
| "when": ["splitContract", "present.fc1BiasT", "not present.fc3T", "present.fc2BiasT"], |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-stage.wgsl.jinja", |
| "bindings": "ffn_bias_none", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK * interSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-stage.wgsl.jinja", |
| "bindings": "out_bias", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "gemv_routed_fc1bias_fc3none_fc2bias", |
| "priority": 20, |
| "when": ["decodeContract", "present.fc1BiasT", "not present.fc3T", "present.fc2BiasT"], |
| "constants": { "decodeLanes": "tunables.decodeLanes", "decodeRows": "decodeRows" }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-gemv.wgsl.jinja", |
| "bindings": "ffn_bias_none_no_count", |
| "dispatch": { |
| "x": "ceilDiv(interSize, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-gemv.wgsl.jinja", |
| "bindings": "out_bias_no_count", |
| "dispatch": { |
| "x": "ceilDiv(hiddenDim, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "split_routed_fc1bias_fc3plain_fc2plain", |
| "priority": 10, |
| "when": ["splitContract", "present.fc1BiasT", "present.fc3T and not present.fc3BiasT", "not present.fc2BiasT"], |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-stage.wgsl.jinja", |
| "bindings": "ffn_bias_plain", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK * interSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-stage.wgsl.jinja", |
| "bindings": "out_plain", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "gemv_routed_fc1bias_fc3plain_fc2plain", |
| "priority": 20, |
| "when": ["decodeContract", "present.fc1BiasT", "present.fc3T and not present.fc3BiasT", "not present.fc2BiasT"], |
| "constants": { "decodeLanes": "tunables.decodeLanes", "decodeRows": "decodeRows" }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-gemv.wgsl.jinja", |
| "bindings": "ffn_bias_plain_no_count", |
| "dispatch": { |
| "x": "ceilDiv(interSize, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-gemv.wgsl.jinja", |
| "bindings": "out_plain_no_count", |
| "dispatch": { |
| "x": "ceilDiv(hiddenDim, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "split_routed_fc1bias_fc3plain_fc2bias", |
| "priority": 10, |
| "when": ["splitContract", "present.fc1BiasT", "present.fc3T and not present.fc3BiasT", "present.fc2BiasT"], |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-stage.wgsl.jinja", |
| "bindings": "ffn_bias_plain", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK * interSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-stage.wgsl.jinja", |
| "bindings": "out_bias", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "gemv_routed_fc1bias_fc3plain_fc2bias", |
| "priority": 20, |
| "when": ["decodeContract", "present.fc1BiasT", "present.fc3T and not present.fc3BiasT", "present.fc2BiasT"], |
| "constants": { "decodeLanes": "tunables.decodeLanes", "decodeRows": "decodeRows" }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-gemv.wgsl.jinja", |
| "bindings": "ffn_bias_plain_no_count", |
| "dispatch": { |
| "x": "ceilDiv(interSize, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-gemv.wgsl.jinja", |
| "bindings": "out_bias_no_count", |
| "dispatch": { |
| "x": "ceilDiv(hiddenDim, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "split_routed_fc1bias_fc3biased_fc2plain", |
| "priority": 10, |
| "when": ["splitContract", "present.fc1BiasT", "present.fc3BiasT", "not present.fc2BiasT"], |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-stage.wgsl.jinja", |
| "bindings": "ffn_bias_biased", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK * interSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-stage.wgsl.jinja", |
| "bindings": "out_plain", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "gemv_routed_fc1bias_fc3biased_fc2plain", |
| "priority": 20, |
| "when": ["decodeContract", "present.fc1BiasT", "present.fc3BiasT", "not present.fc2BiasT"], |
| "constants": { "decodeLanes": "tunables.decodeLanes", "decodeRows": "decodeRows" }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-gemv.wgsl.jinja", |
| "bindings": "ffn_bias_biased_no_count", |
| "dispatch": { |
| "x": "ceilDiv(interSize, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-gemv.wgsl.jinja", |
| "bindings": "out_plain_no_count", |
| "dispatch": { |
| "x": "ceilDiv(hiddenDim, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "split_routed_fc1bias_fc3biased_fc2bias", |
| "priority": 10, |
| "when": ["splitContract", "present.fc1BiasT", "present.fc3BiasT", "present.fc2BiasT"], |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-stage.wgsl.jinja", |
| "bindings": "ffn_bias_biased", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK * interSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-stage.wgsl.jinja", |
| "bindings": "out_bias", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "gemv_routed_fc1bias_fc3biased_fc2bias", |
| "priority": 20, |
| "when": ["decodeContract", "present.fc1BiasT", "present.fc3BiasT", "present.fc2BiasT"], |
| "constants": { "decodeLanes": "tunables.decodeLanes", "decodeRows": "decodeRows" }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStage", |
| "shader": "moe-ffn-gemv.wgsl.jinja", |
| "bindings": "ffn_bias_biased_no_count", |
| "dispatch": { |
| "x": "ceilDiv(interSize, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * topK" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStage", |
| "shader": "moe-output-gemv.wgsl.jinja", |
| "bindings": "out_bias_no_count", |
| "dispatch": { |
| "x": "ceilDiv(hiddenDim, decodeRows)", |
| "y": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens)" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "grouped_routed_fc1plain_fc3none_fc2plain", |
| "priority": 30, |
| "when": ["groupedContract", "not present.fc1BiasT", "not present.fc3T", "not present.fc2BiasT"], |
| "constants": { |
| "groupTileM": "groupTileM", |
| "groupTileN": "groupTileN", |
| "groupTileK": "tunables.groupTileK", |
| "groupTileKVec": "groupTileKVec", |
| "groupThreads": "tunables.groupThreads", |
| "regM": "tunables.groupRegM", |
| "regN": "tunables.groupRegN", |
| "groupRouteWorkgroup": "tunables.groupRouteWorkgroup", |
| "groupVec4": "groupedVec4Ok", |
| "groupElem": "\"vec4<f32>\" if groupedVec4Ok else \"f32\"" |
| }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" }, |
| { "id": "slotList", "dtype": "uint32", "shape": "[max(1, groupSlots)]" }, |
| { "id": "tileMeta", "dtype": "uint32", "shape": "[1 + 3 * groupMaxTiles]" }, |
| { "id": "slotOut", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * hiddenDim)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "group_stage", |
| "name": "MoE.GroupStage", |
| "shader": "expert-group-slots.wgsl.jinja", |
| "bindings": "group", |
| "dispatch": { |
| "threads": "constants.groupRouteWorkgroup", |
| "workgroupSize": "constants.groupRouteWorkgroup" |
| } |
| }, |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStageGrouped", |
| "shader": "moe-ffn-grouped.wgsl.jinja", |
| "bindings": "ffng_plain_none", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStageGrouped", |
| "shader": "moe-output-grouped.wgsl.jinja", |
| "bindings": "outg_plain", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenDim, groupTileN)" } |
| }, |
| { |
| "id": "mix_stage", |
| "name": "MoE.MixStage", |
| "shader": "moe-mix-stage.wgsl.jinja", |
| "bindings": "mix", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "grouped_routed_fc1plain_fc3none_fc2bias", |
| "priority": 30, |
| "when": ["groupedContract", "not present.fc1BiasT", "not present.fc3T", "present.fc2BiasT"], |
| "constants": { |
| "groupTileM": "groupTileM", |
| "groupTileN": "groupTileN", |
| "groupTileK": "tunables.groupTileK", |
| "groupTileKVec": "groupTileKVec", |
| "groupThreads": "tunables.groupThreads", |
| "regM": "tunables.groupRegM", |
| "regN": "tunables.groupRegN", |
| "groupRouteWorkgroup": "tunables.groupRouteWorkgroup", |
| "groupVec4": "groupedVec4Ok", |
| "groupElem": "\"vec4<f32>\" if groupedVec4Ok else \"f32\"" |
| }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" }, |
| { "id": "slotList", "dtype": "uint32", "shape": "[max(1, groupSlots)]" }, |
| { "id": "tileMeta", "dtype": "uint32", "shape": "[1 + 3 * groupMaxTiles]" }, |
| { "id": "slotOut", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * hiddenDim)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "group_stage", |
| "name": "MoE.GroupStage", |
| "shader": "expert-group-slots.wgsl.jinja", |
| "bindings": "group", |
| "dispatch": { |
| "threads": "constants.groupRouteWorkgroup", |
| "workgroupSize": "constants.groupRouteWorkgroup" |
| } |
| }, |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStageGrouped", |
| "shader": "moe-ffn-grouped.wgsl.jinja", |
| "bindings": "ffng_plain_none", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStageGrouped", |
| "shader": "moe-output-grouped.wgsl.jinja", |
| "bindings": "outg_bias", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenDim, groupTileN)" } |
| }, |
| { |
| "id": "mix_stage", |
| "name": "MoE.MixStage", |
| "shader": "moe-mix-stage.wgsl.jinja", |
| "bindings": "mix", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "grouped_routed_fc1plain_fc3plain_fc2plain", |
| "priority": 30, |
| "when": ["groupedContract", "not present.fc1BiasT", "present.fc3T and not present.fc3BiasT", "not present.fc2BiasT"], |
| "constants": { |
| "groupTileM": "groupTileM", |
| "groupTileN": "groupTileN", |
| "groupTileK": "tunables.groupTileK", |
| "groupTileKVec": "groupTileKVec", |
| "groupThreads": "tunables.groupThreads", |
| "regM": "tunables.groupRegM", |
| "regN": "tunables.groupRegN", |
| "groupRouteWorkgroup": "tunables.groupRouteWorkgroup", |
| "groupVec4": "groupedVec4Ok", |
| "groupElem": "\"vec4<f32>\" if groupedVec4Ok else \"f32\"" |
| }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" }, |
| { "id": "slotList", "dtype": "uint32", "shape": "[max(1, groupSlots)]" }, |
| { "id": "tileMeta", "dtype": "uint32", "shape": "[1 + 3 * groupMaxTiles]" }, |
| { "id": "slotOut", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * hiddenDim)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "group_stage", |
| "name": "MoE.GroupStage", |
| "shader": "expert-group-slots.wgsl.jinja", |
| "bindings": "group", |
| "dispatch": { |
| "threads": "constants.groupRouteWorkgroup", |
| "workgroupSize": "constants.groupRouteWorkgroup" |
| } |
| }, |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStageGrouped", |
| "shader": "moe-ffn-grouped.wgsl.jinja", |
| "bindings": "ffng_plain_plain", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStageGrouped", |
| "shader": "moe-output-grouped.wgsl.jinja", |
| "bindings": "outg_plain", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenDim, groupTileN)" } |
| }, |
| { |
| "id": "mix_stage", |
| "name": "MoE.MixStage", |
| "shader": "moe-mix-stage.wgsl.jinja", |
| "bindings": "mix", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "grouped_routed_fc1plain_fc3plain_fc2bias", |
| "priority": 30, |
| "when": ["groupedContract", "not present.fc1BiasT", "present.fc3T and not present.fc3BiasT", "present.fc2BiasT"], |
| "constants": { |
| "groupTileM": "groupTileM", |
| "groupTileN": "groupTileN", |
| "groupTileK": "tunables.groupTileK", |
| "groupTileKVec": "groupTileKVec", |
| "groupThreads": "tunables.groupThreads", |
| "regM": "tunables.groupRegM", |
| "regN": "tunables.groupRegN", |
| "groupRouteWorkgroup": "tunables.groupRouteWorkgroup", |
| "groupVec4": "groupedVec4Ok", |
| "groupElem": "\"vec4<f32>\" if groupedVec4Ok else \"f32\"" |
| }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" }, |
| { "id": "slotList", "dtype": "uint32", "shape": "[max(1, groupSlots)]" }, |
| { "id": "tileMeta", "dtype": "uint32", "shape": "[1 + 3 * groupMaxTiles]" }, |
| { "id": "slotOut", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * hiddenDim)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "group_stage", |
| "name": "MoE.GroupStage", |
| "shader": "expert-group-slots.wgsl.jinja", |
| "bindings": "group", |
| "dispatch": { |
| "threads": "constants.groupRouteWorkgroup", |
| "workgroupSize": "constants.groupRouteWorkgroup" |
| } |
| }, |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStageGrouped", |
| "shader": "moe-ffn-grouped.wgsl.jinja", |
| "bindings": "ffng_plain_plain", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStageGrouped", |
| "shader": "moe-output-grouped.wgsl.jinja", |
| "bindings": "outg_bias", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenDim, groupTileN)" } |
| }, |
| { |
| "id": "mix_stage", |
| "name": "MoE.MixStage", |
| "shader": "moe-mix-stage.wgsl.jinja", |
| "bindings": "mix", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "grouped_routed_fc1plain_fc3biased_fc2plain", |
| "priority": 30, |
| "when": ["groupedContract", "not present.fc1BiasT", "present.fc3BiasT", "not present.fc2BiasT"], |
| "constants": { |
| "groupTileM": "groupTileM", |
| "groupTileN": "groupTileN", |
| "groupTileK": "tunables.groupTileK", |
| "groupTileKVec": "groupTileKVec", |
| "groupThreads": "tunables.groupThreads", |
| "regM": "tunables.groupRegM", |
| "regN": "tunables.groupRegN", |
| "groupRouteWorkgroup": "tunables.groupRouteWorkgroup", |
| "groupVec4": "groupedVec4Ok", |
| "groupElem": "\"vec4<f32>\" if groupedVec4Ok else \"f32\"" |
| }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" }, |
| { "id": "slotList", "dtype": "uint32", "shape": "[max(1, groupSlots)]" }, |
| { "id": "tileMeta", "dtype": "uint32", "shape": "[1 + 3 * groupMaxTiles]" }, |
| { "id": "slotOut", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * hiddenDim)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "group_stage", |
| "name": "MoE.GroupStage", |
| "shader": "expert-group-slots.wgsl.jinja", |
| "bindings": "group", |
| "dispatch": { |
| "threads": "constants.groupRouteWorkgroup", |
| "workgroupSize": "constants.groupRouteWorkgroup" |
| } |
| }, |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStageGrouped", |
| "shader": "moe-ffn-grouped.wgsl.jinja", |
| "bindings": "ffng_plain_biased", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStageGrouped", |
| "shader": "moe-output-grouped.wgsl.jinja", |
| "bindings": "outg_plain", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenDim, groupTileN)" } |
| }, |
| { |
| "id": "mix_stage", |
| "name": "MoE.MixStage", |
| "shader": "moe-mix-stage.wgsl.jinja", |
| "bindings": "mix", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "grouped_routed_fc1plain_fc3biased_fc2bias", |
| "priority": 30, |
| "when": ["groupedContract", "not present.fc1BiasT", "present.fc3BiasT", "present.fc2BiasT"], |
| "constants": { |
| "groupTileM": "groupTileM", |
| "groupTileN": "groupTileN", |
| "groupTileK": "tunables.groupTileK", |
| "groupTileKVec": "groupTileKVec", |
| "groupThreads": "tunables.groupThreads", |
| "regM": "tunables.groupRegM", |
| "regN": "tunables.groupRegN", |
| "groupRouteWorkgroup": "tunables.groupRouteWorkgroup", |
| "groupVec4": "groupedVec4Ok", |
| "groupElem": "\"vec4<f32>\" if groupedVec4Ok else \"f32\"" |
| }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" }, |
| { "id": "slotList", "dtype": "uint32", "shape": "[max(1, groupSlots)]" }, |
| { "id": "tileMeta", "dtype": "uint32", "shape": "[1 + 3 * groupMaxTiles]" }, |
| { "id": "slotOut", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * hiddenDim)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "group_stage", |
| "name": "MoE.GroupStage", |
| "shader": "expert-group-slots.wgsl.jinja", |
| "bindings": "group", |
| "dispatch": { |
| "threads": "constants.groupRouteWorkgroup", |
| "workgroupSize": "constants.groupRouteWorkgroup" |
| } |
| }, |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStageGrouped", |
| "shader": "moe-ffn-grouped.wgsl.jinja", |
| "bindings": "ffng_plain_biased", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStageGrouped", |
| "shader": "moe-output-grouped.wgsl.jinja", |
| "bindings": "outg_bias", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenDim, groupTileN)" } |
| }, |
| { |
| "id": "mix_stage", |
| "name": "MoE.MixStage", |
| "shader": "moe-mix-stage.wgsl.jinja", |
| "bindings": "mix", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "grouped_routed_fc1bias_fc3none_fc2plain", |
| "priority": 30, |
| "when": ["groupedContract", "present.fc1BiasT", "not present.fc3T", "not present.fc2BiasT"], |
| "constants": { |
| "groupTileM": "groupTileM", |
| "groupTileN": "groupTileN", |
| "groupTileK": "tunables.groupTileK", |
| "groupTileKVec": "groupTileKVec", |
| "groupThreads": "tunables.groupThreads", |
| "regM": "tunables.groupRegM", |
| "regN": "tunables.groupRegN", |
| "groupRouteWorkgroup": "tunables.groupRouteWorkgroup", |
| "groupVec4": "groupedVec4Ok", |
| "groupElem": "\"vec4<f32>\" if groupedVec4Ok else \"f32\"" |
| }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" }, |
| { "id": "slotList", "dtype": "uint32", "shape": "[max(1, groupSlots)]" }, |
| { "id": "tileMeta", "dtype": "uint32", "shape": "[1 + 3 * groupMaxTiles]" }, |
| { "id": "slotOut", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * hiddenDim)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "group_stage", |
| "name": "MoE.GroupStage", |
| "shader": "expert-group-slots.wgsl.jinja", |
| "bindings": "group", |
| "dispatch": { |
| "threads": "constants.groupRouteWorkgroup", |
| "workgroupSize": "constants.groupRouteWorkgroup" |
| } |
| }, |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStageGrouped", |
| "shader": "moe-ffn-grouped.wgsl.jinja", |
| "bindings": "ffng_bias_none", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStageGrouped", |
| "shader": "moe-output-grouped.wgsl.jinja", |
| "bindings": "outg_plain", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenDim, groupTileN)" } |
| }, |
| { |
| "id": "mix_stage", |
| "name": "MoE.MixStage", |
| "shader": "moe-mix-stage.wgsl.jinja", |
| "bindings": "mix", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "grouped_routed_fc1bias_fc3none_fc2bias", |
| "priority": 30, |
| "when": ["groupedContract", "present.fc1BiasT", "not present.fc3T", "present.fc2BiasT"], |
| "constants": { |
| "groupTileM": "groupTileM", |
| "groupTileN": "groupTileN", |
| "groupTileK": "tunables.groupTileK", |
| "groupTileKVec": "groupTileKVec", |
| "groupThreads": "tunables.groupThreads", |
| "regM": "tunables.groupRegM", |
| "regN": "tunables.groupRegN", |
| "groupRouteWorkgroup": "tunables.groupRouteWorkgroup", |
| "groupVec4": "groupedVec4Ok", |
| "groupElem": "\"vec4<f32>\" if groupedVec4Ok else \"f32\"" |
| }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" }, |
| { "id": "slotList", "dtype": "uint32", "shape": "[max(1, groupSlots)]" }, |
| { "id": "tileMeta", "dtype": "uint32", "shape": "[1 + 3 * groupMaxTiles]" }, |
| { "id": "slotOut", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * hiddenDim)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "group_stage", |
| "name": "MoE.GroupStage", |
| "shader": "expert-group-slots.wgsl.jinja", |
| "bindings": "group", |
| "dispatch": { |
| "threads": "constants.groupRouteWorkgroup", |
| "workgroupSize": "constants.groupRouteWorkgroup" |
| } |
| }, |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStageGrouped", |
| "shader": "moe-ffn-grouped.wgsl.jinja", |
| "bindings": "ffng_bias_none", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStageGrouped", |
| "shader": "moe-output-grouped.wgsl.jinja", |
| "bindings": "outg_bias", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenDim, groupTileN)" } |
| }, |
| { |
| "id": "mix_stage", |
| "name": "MoE.MixStage", |
| "shader": "moe-mix-stage.wgsl.jinja", |
| "bindings": "mix", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "grouped_routed_fc1bias_fc3plain_fc2plain", |
| "priority": 30, |
| "when": ["groupedContract", "present.fc1BiasT", "present.fc3T and not present.fc3BiasT", "not present.fc2BiasT"], |
| "constants": { |
| "groupTileM": "groupTileM", |
| "groupTileN": "groupTileN", |
| "groupTileK": "tunables.groupTileK", |
| "groupTileKVec": "groupTileKVec", |
| "groupThreads": "tunables.groupThreads", |
| "regM": "tunables.groupRegM", |
| "regN": "tunables.groupRegN", |
| "groupRouteWorkgroup": "tunables.groupRouteWorkgroup", |
| "groupVec4": "groupedVec4Ok", |
| "groupElem": "\"vec4<f32>\" if groupedVec4Ok else \"f32\"" |
| }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" }, |
| { "id": "slotList", "dtype": "uint32", "shape": "[max(1, groupSlots)]" }, |
| { "id": "tileMeta", "dtype": "uint32", "shape": "[1 + 3 * groupMaxTiles]" }, |
| { "id": "slotOut", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * hiddenDim)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "group_stage", |
| "name": "MoE.GroupStage", |
| "shader": "expert-group-slots.wgsl.jinja", |
| "bindings": "group", |
| "dispatch": { |
| "threads": "constants.groupRouteWorkgroup", |
| "workgroupSize": "constants.groupRouteWorkgroup" |
| } |
| }, |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStageGrouped", |
| "shader": "moe-ffn-grouped.wgsl.jinja", |
| "bindings": "ffng_bias_plain", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStageGrouped", |
| "shader": "moe-output-grouped.wgsl.jinja", |
| "bindings": "outg_plain", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenDim, groupTileN)" } |
| }, |
| { |
| "id": "mix_stage", |
| "name": "MoE.MixStage", |
| "shader": "moe-mix-stage.wgsl.jinja", |
| "bindings": "mix", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "grouped_routed_fc1bias_fc3plain_fc2bias", |
| "priority": 30, |
| "when": ["groupedContract", "present.fc1BiasT", "present.fc3T and not present.fc3BiasT", "present.fc2BiasT"], |
| "constants": { |
| "groupTileM": "groupTileM", |
| "groupTileN": "groupTileN", |
| "groupTileK": "tunables.groupTileK", |
| "groupTileKVec": "groupTileKVec", |
| "groupThreads": "tunables.groupThreads", |
| "regM": "tunables.groupRegM", |
| "regN": "tunables.groupRegN", |
| "groupRouteWorkgroup": "tunables.groupRouteWorkgroup", |
| "groupVec4": "groupedVec4Ok", |
| "groupElem": "\"vec4<f32>\" if groupedVec4Ok else \"f32\"" |
| }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" }, |
| { "id": "slotList", "dtype": "uint32", "shape": "[max(1, groupSlots)]" }, |
| { "id": "tileMeta", "dtype": "uint32", "shape": "[1 + 3 * groupMaxTiles]" }, |
| { "id": "slotOut", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * hiddenDim)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "group_stage", |
| "name": "MoE.GroupStage", |
| "shader": "expert-group-slots.wgsl.jinja", |
| "bindings": "group", |
| "dispatch": { |
| "threads": "constants.groupRouteWorkgroup", |
| "workgroupSize": "constants.groupRouteWorkgroup" |
| } |
| }, |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStageGrouped", |
| "shader": "moe-ffn-grouped.wgsl.jinja", |
| "bindings": "ffng_bias_plain", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStageGrouped", |
| "shader": "moe-output-grouped.wgsl.jinja", |
| "bindings": "outg_bias", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenDim, groupTileN)" } |
| }, |
| { |
| "id": "mix_stage", |
| "name": "MoE.MixStage", |
| "shader": "moe-mix-stage.wgsl.jinja", |
| "bindings": "mix", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "grouped_routed_fc1bias_fc3biased_fc2plain", |
| "priority": 30, |
| "when": ["groupedContract", "present.fc1BiasT", "present.fc3BiasT", "not present.fc2BiasT"], |
| "constants": { |
| "groupTileM": "groupTileM", |
| "groupTileN": "groupTileN", |
| "groupTileK": "tunables.groupTileK", |
| "groupTileKVec": "groupTileKVec", |
| "groupThreads": "tunables.groupThreads", |
| "regM": "tunables.groupRegM", |
| "regN": "tunables.groupRegN", |
| "groupRouteWorkgroup": "tunables.groupRouteWorkgroup", |
| "groupVec4": "groupedVec4Ok", |
| "groupElem": "\"vec4<f32>\" if groupedVec4Ok else \"f32\"" |
| }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" }, |
| { "id": "slotList", "dtype": "uint32", "shape": "[max(1, groupSlots)]" }, |
| { "id": "tileMeta", "dtype": "uint32", "shape": "[1 + 3 * groupMaxTiles]" }, |
| { "id": "slotOut", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * hiddenDim)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "group_stage", |
| "name": "MoE.GroupStage", |
| "shader": "expert-group-slots.wgsl.jinja", |
| "bindings": "group", |
| "dispatch": { |
| "threads": "constants.groupRouteWorkgroup", |
| "workgroupSize": "constants.groupRouteWorkgroup" |
| } |
| }, |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStageGrouped", |
| "shader": "moe-ffn-grouped.wgsl.jinja", |
| "bindings": "ffng_bias_biased", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStageGrouped", |
| "shader": "moe-output-grouped.wgsl.jinja", |
| "bindings": "outg_plain", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenDim, groupTileN)" } |
| }, |
| { |
| "id": "mix_stage", |
| "name": "MoE.MixStage", |
| "shader": "moe-mix-stage.wgsl.jinja", |
| "bindings": "mix", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "grouped_routed_fc1bias_fc3biased_fc2bias", |
| "priority": 30, |
| "when": ["groupedContract", "present.fc1BiasT", "present.fc3BiasT", "present.fc2BiasT"], |
| "constants": { |
| "groupTileM": "groupTileM", |
| "groupTileN": "groupTileN", |
| "groupTileK": "tunables.groupTileK", |
| "groupTileKVec": "groupTileKVec", |
| "groupThreads": "tunables.groupThreads", |
| "regM": "tunables.groupRegM", |
| "regN": "tunables.groupRegN", |
| "groupRouteWorkgroup": "tunables.groupRouteWorkgroup", |
| "groupVec4": "groupedVec4Ok", |
| "groupElem": "\"vec4<f32>\" if groupedVec4Ok else \"f32\"" |
| }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[max(1, tokenCount * topK)]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * interSize)]" }, |
| { "id": "slotList", "dtype": "uint32", "shape": "[max(1, groupSlots)]" }, |
| { "id": "tileMeta", "dtype": "uint32", "shape": "[1 + 3 * groupMaxTiles]" }, |
| { "id": "slotOut", "dtype": "float32", "shape": "[max(1, hiddenChunkTokens * topK * hiddenDim)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "MoE.RouteStage", |
| "shader": "moe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "tokenCount", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "group_stage", |
| "name": "MoE.GroupStage", |
| "shader": "expert-group-slots.wgsl.jinja", |
| "bindings": "group", |
| "dispatch": { |
| "threads": "constants.groupRouteWorkgroup", |
| "workgroupSize": "constants.groupRouteWorkgroup" |
| } |
| }, |
| { |
| "id": "ffn_stage", |
| "name": "MoE.FfnStageGrouped", |
| "shader": "moe-ffn-grouped.wgsl.jinja", |
| "bindings": "ffng_bias_biased", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" } |
| }, |
| { |
| "id": "output_stage", |
| "name": "MoE.OutputStageGrouped", |
| "shader": "moe-output-grouped.wgsl.jinja", |
| "bindings": "outg_bias", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenDim, groupTileN)" } |
| }, |
| { |
| "id": "mix_stage", |
| "name": "MoE.MixStage", |
| "shader": "moe-mix-stage.wgsl.jinja", |
| "bindings": "mix", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, tokenCount - repeat.chunk * hiddenChunkTokens) * hiddenDim", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| } |
| ] |
| } |
|
|
