| { |
| "domain": "com.microsoft", |
| "name": "QMoE", |
| "sinceVersion": 1, |
| "description": "Quantized mixture of experts over float32 activations and raw uint8-packed integer expert weights. The package implements 4- and 8-bit symmetric dequantization, ReLU, and interleaved SwiGLU for rank-2 or rank-3 input. Routing uses the softmax over the selected top-k logits, matching ONNX Runtime QMoE when the omitted `router_weights` input is absent. Biases, explicit zero points, FC3, sparse mixing, provider-prepacked layouts, and FP4/FP8 modes are not implemented.", |
| "inputs": [ |
| { |
| "role": "input", |
| "dtype": "T", |
| "description": "Token activations of shape `(num_tokens, hidden_size)` or `(batch_size, sequence_length, hidden_size)`." |
| }, |
| { |
| "role": "router_probs", |
| "dtype": "T", |
| "rank": 2, |
| "description": "Routing logits of shape `(num_tokens, num_experts)`, where `num_tokens` is the product of the leading input dimensions." |
| }, |
| { |
| "role": "fc1_experts_weights", |
| "dtype": "T1", |
| "rank": 3, |
| "description": "Raw packed FC1 weights of shape `(num_experts, fusion_size * inter_size, hidden_size / pack_size)`, where `fusion_size` is 2 only for interleaved SwiGLU and `pack_size` is 2 for 4-bit or 1 for 8-bit weights." |
| }, |
| { |
| "role": "fc1_scales", |
| "dtype": "T2", |
| "description": "Required FC1 dequantization scales: rank 2 `(num_experts, fusion_size * inter_size)` for column-wise quantization, or rank 3 with a trailing `hidden_size / block_size` axis." |
| }, |
| { |
| "role": "fc2_experts_weights", |
| "dtype": "T1", |
| "rank": 3, |
| "description": "Raw packed FC2 weights of shape `(num_experts, hidden_size, inter_size / pack_size)`." |
| }, |
| { |
| "role": "fc2_scales", |
| "dtype": "T2", |
| "description": "Required FC2 dequantization scales: rank 2 `(num_experts, hidden_size)` for column-wise quantization, or rank 3 with a trailing `inter_size / block_size` axis." |
| } |
| ], |
| "outputs": [ |
| { |
| "role": "output", |
| "dtype": "T", |
| "shape": "shapes.input", |
| "description": "Routed expert output with the same shape as `input`." |
| } |
| ], |
| "attributes": { |
| "activation_alpha": 1, |
| "activation_beta": 0, |
| "activation_type": "relu", |
| "expert_weight_bits": 4, |
| "k": 1, |
| "normalize_routing_weights": 0, |
| "quant_type": "int", |
| "swiglu_fusion": 0, |
| "use_sparse_mixer": 0, |
| "weights_prepacked": -1 |
| }, |
| "attributeDescriptions": { |
| "activation_alpha": "Alpha used by SwiGLU; the exact standard default is 1.", |
| "activation_beta": "Beta added to the SwiGLU linear branch; the exact standard default is 0.", |
| "activation_type": "Activation applied after FC1. This package supports `relu` and `swiglu`; the exact standard default is `relu`.", |
| "block_size": "Optional quantization block size along the reduction dimension. Omission selects column-wise scaling; a supplied value must be a power of two of at least 16 and divide both `hidden_size` and `inter_size`.", |
| "expert_weight_bits": "Integer expert-weight bit width. This package supports 4 and 8; the exact standard default is 4.", |
| "k": "Number of experts selected per token; the exact standard default is 1.", |
| "normalize_routing_weights": "Accepted values are 0 and 1. With the separate `router_weights` input omitted, ONNX Runtime QMoE applies a softmax over the selected top-k logits for either value; the attribute only distinguishes the unsupported separate-weight path.", |
| "quant_type": "Quantization family. This package supports only the exact standard default `int`.", |
| "swiglu_fusion": "SwiGLU packing mode. ReLU uses the exact standard default 0; supported SwiGLU interleaves gate/up FC1 rows with value 1.", |
| "swiglu_limit": "Optional SwiGLU clamp limit. Omission means no finite clamp.", |
| "use_sparse_mixer": "Whether to use sparse-mixer routing. The exact standard default and only supported value is 0.", |
| "weights_prepacked": "Provider weight-layout selector. Values -1 and 0 both consume the public raw packed tensor layout; provider-specific prepacked value 1 is not portable and is rejected." |
| }, |
| "attributeConstraints": { |
| "activation_type": { "values": ["relu", "swiglu"] }, |
| "expert_weight_bits": { "values": [4, 8] }, |
| "normalize_routing_weights": { "values": [0, 1] }, |
| "quant_type": { "values": ["int"] }, |
| "swiglu_fusion": { "values": [0, 1] }, |
| "use_sparse_mixer": { "values": [0] }, |
| "weights_prepacked": { "values": [-1, 0] } |
| }, |
| "typeConstraints": { "T": ["float32"], "T1": ["uint8"], "T2": ["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" }, |
| "fc1ScalesT": { "kind": "tensor", "semantic": "fc1_scales", "role": "input" }, |
| "fc2T": { "kind": "tensor", "semantic": "fc2_experts_weights", "role": "input" }, |
| "fc2ScalesT": { "kind": "tensor", "semantic": "fc2_scales", "role": "input" }, |
| "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)", |
| "wave32Adapter": "has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize == 32 and device.adapterInfo.subgroupMaxSize == 32", |
| "canPinSubgroupSize32": "device.features.has(\"subgroups\") and device.features.has(\"subgroup-size-control\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize <= 32 and device.adapterInfo.subgroupMaxSize >= 32", |
| "pinSubgroupSize32": "canPinSubgroupSize32 and not wave32Adapter", |
| "wave32Effective": "wave32Adapter or pinSubgroupSize32", |
| "hiddenSize": "dim(shapes.input, ranks.input - 1)", |
| "numTokens": "numel(shapes.input) / max(1, hiddenSize)", |
| "topK": "attrs.k", |
| "weightBits": "attrs.expert_weight_bits", |
| "quantBlockSize": "attrs.block_size if has(attrs, \"block_size\") else 0", |
| "packSize": "2 if weightBits == 4 else 1", |
| "quantMidpoint": "8 if weightBits == 4 else 128", |
| "fusionSize": "2 if attrs.activation_type == \"swiglu\" and attrs.swiglu_fusion == 1 else 1", |
| "interSize": "dim(shapes.fc1_experts_weights, 1) / fusionSize", |
| "fc1PackedCols": "dim(shapes.fc1_experts_weights, 2)", |
| "fc2PackedCols": "dim(shapes.fc2_experts_weights, 2)", |
| "colWiseScales": "quantBlockSize == 0", |
| "fc1ScaleBlocks": "1 if colWiseScales else hiddenSize / max(1, quantBlockSize)", |
| "fc2ScaleBlocks": "1 if colWiseScales else interSize / max(1, quantBlockSize)", |
| "activationSupported": "(attrs.activation_type == \"relu\" and attrs.swiglu_fusion == 0) or (attrs.activation_type == \"swiglu\" and attrs.swiglu_fusion == 1)", |
| "routingModeSupported": "attrs.normalize_routing_weights == 0 or attrs.normalize_routing_weights == 1", |
| "rawWeightLayout": "attrs.weights_prepacked == -1 or attrs.weights_prepacked == 0", |
| "inputOutputShapeOk": "((ranks.input == 2 and ranks.output == 2 and dim(shapes.output, 0) == dim(shapes.input, 0) and dim(shapes.output, 1) == dim(shapes.input, 1)) or (ranks.input == 3 and ranks.output == 3 and dim(shapes.output, 0) == dim(shapes.input, 0) and dim(shapes.output, 1) == dim(shapes.input, 1) and dim(shapes.output, 2) == dim(shapes.input, 2))) and hiddenSize > 0", |
| "quantBlockSizeOk": "colWiseScales or (quantBlockSize >= 16 and pow2ceil(quantBlockSize) == quantBlockSize and hiddenSize % quantBlockSize == 0 and interSize % quantBlockSize == 0)", |
| "quantScalesOk": "tensorDtypes.fc1_scales == \"float32\" and tensorDtypes.fc2_scales == \"float32\" and dim(shapes.fc1_scales, 0) == dim(shapes.router_probs, 1) and dim(shapes.fc2_scales, 0) == dim(shapes.router_probs, 1) and dim(shapes.fc1_scales, 1) == dim(shapes.fc1_experts_weights, 1) and dim(shapes.fc2_scales, 1) == hiddenSize and ((ranks.fc1_scales == 2 and ranks.fc2_scales == 2) if colWiseScales else (ranks.fc1_scales == 3 and ranks.fc2_scales == 3 and dim(shapes.fc1_scales, 2) == fc1ScaleBlocks and dim(shapes.fc2_scales, 2) == fc2ScaleBlocks))", |
| "quantShapeOk": "inputOutputShapeOk and ranks.router_probs == 2 and ranks.fc1_experts_weights == 3 and ranks.fc2_experts_weights == 3 and dim(shapes.router_probs, 0) == numTokens and dim(shapes.fc1_experts_weights, 0) == dim(shapes.router_probs, 1) and dim(shapes.fc2_experts_weights, 0) == dim(shapes.router_probs, 1) and dim(shapes.fc1_experts_weights, 1) % fusionSize == 0 and dim(shapes.fc2_experts_weights, 1) == hiddenSize and dim(shapes.fc1_experts_weights, 2) * packSize == hiddenSize and dim(shapes.fc2_experts_weights, 2) * packSize == interSize and quantBlockSizeOk and quantScalesOk", |
| "quantContract": "activationSupported and routingModeSupported and rawWeightLayout and quantShapeOk and topK >= 1 and topK <= dim(shapes.router_probs, 1)", |
| "workgroupSizeOk": "tunables.workgroupSize >= 1 and tunables.workgroupSize <= deviceWorkgroupCap", |
| "hiddenChunkFits": "topK * interSize * 4 <= device.limits.maxStorageBufferBindingSize and topK * interSize * 4 <= device.limits.maxBufferSize", |
| "hiddenChunkTokens": "numTokens if interSize == 0 else min(numTokens, max(1, floor(min(device.limits.maxStorageBufferBindingSize, device.limits.maxBufferSize) / (topK * interSize * 4))))", |
| "hiddenChunkCount": "max(1, ceilDiv(numTokens, max(1, hiddenChunkTokens)))", |
| "routeScratchBytes": "numTokens * topK * 4", |
| "routedScratchFits": "routeScratchBytes <= device.limits.maxStorageBufferBindingSize and routeScratchBytes <= device.limits.maxBufferSize", |
| "groupTileM": "tunables.groupThreads * tunables.groupRegM", |
| "groupTileN": "tunables.groupThreads * tunables.groupRegN", |
| "groupTileKVec": "ceilDiv(tunables.groupTileK, 4)", |
| "groupThreadCount": "tunables.groupThreads * tunables.groupThreads", |
| "groupSlots": "hiddenChunkTokens * topK", |
| "groupMaxTiles": "ceilDiv(groupSlots, max(1, groupTileM)) + dim(shapes.router_probs, 1)", |
| "groupSharedBytes": "(groupTileM * tunables.groupTileK + 2 * groupTileN * tunables.groupTileK + groupTileM) * 4", |
| "groupSlotOutBytes": "hiddenChunkTokens * topK * hiddenSize * 4", |
| "groupedDeviceOk": "groupThreadCount <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.groupThreads <= device.limits.maxComputeWorkgroupSizeX and tunables.groupThreads <= device.limits.maxComputeWorkgroupSizeY and tunables.groupRouteWorkgroup <= deviceWorkgroupCap and groupSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and dim(shapes.router_probs, 1) * 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(hiddenSize, max(1, groupTileN)) <= device.limits.maxComputeWorkgroupsPerDimension", |
| "groupSlotOutFits": "groupSlotOutBytes <= device.limits.maxStorageBufferBindingSize and groupSlotOutBytes <= device.limits.maxBufferSize", |
| "groupedShapeOk": "numTokens * topK * 4 >= groupTileM * dim(shapes.router_probs, 1)", |
| "groupedContract": "quantContract and workgroupSizeOk and hiddenChunkFits and routedScratchFits and interSize > 0 and groupedDeviceOk and groupedDispatchOk and groupSlotOutFits and groupedShapeOk", |
| "sgmatWorkgroup": "128", |
| "sgmatSubgroups": "4", |
| "sgmatRowSubtiles": "2", |
| "groupedSgmatSharedBytes": "groupTileM * 4 + groupTileM * 32 * 4 + 64 * 32 * 4 + sgmatSubgroups * 4 * 64 * 4", |
| "groupedSgmatOk": "groupedContract and groupTileM == 32 and interSize % 32 == 0 and groupedSgmatSharedBytes <= device.limits.maxComputeWorkgroupStorageSize and ceilDiv(hiddenSize, 64) <= device.limits.maxComputeWorkgroupsPerDimension and hiddenSize % 32 == 0", |
| "decodeLanesOk": "tunables.decodeLanes >= 1 and tunables.decodeLanes <= tunables.workgroupSize and tunables.workgroupSize % tunables.decodeLanes == 0", |
| "decodeRows": "max(1, tunables.workgroupSize / max(1, tunables.decodeLanes))", |
| "decodeDeviceOk": "workgroupSizeOk and decodeLanesOk and tunables.decodeLanes <= device.limits.maxComputeWorkgroupSizeX and decodeRows <= device.limits.maxComputeWorkgroupSizeY and tunables.workgroupSize * 8 <= device.limits.maxComputeWorkgroupStorageSize", |
| "decodeDispatchOk": "ceilDiv(interSize, decodeRows) <= device.limits.maxComputeWorkgroupsPerDimension and ceilDiv(hiddenSize, decodeRows) <= device.limits.maxComputeWorkgroupsPerDimension and hiddenChunkTokens * topK <= device.limits.maxComputeWorkgroupsPerDimension", |
| "decodeLaneDepth": "min(hiddenSize, interSize) / max(1, tunables.decodeLanes)", |
| "decodeDepthOk": "decodeLaneDepth >= tunables.decodeMinLaneTrips", |
| "splitFc1Blocks": "ceilDiv(numTokens * topK * interSize, max(1, tunables.workgroupSize))", |
| "decodeOccupancyOk": "splitFc1Blocks <= tunables.decodeBlockTarget", |
| "decodeContract": "quantContract and decodeDeviceOk and decodeDispatchOk and hiddenChunkFits and routedScratchFits and interSize > 0 and decodeDepthOk and (decodeOccupancyOk or not groupedContract)" |
| }, |
| "constants": { |
| "tokens": "numTokens", |
| "hidden": "hiddenSize", |
| "experts": "dim(shapes.router_probs, 1)", |
| "fc1Rows": "dim(shapes.fc1_experts_weights, 1)", |
| "inter": "interSize", |
| "topK": "topK", |
| "workgroupSize": "tunables.workgroupSize", |
| "activationType": "attrs.activation_type", |
| "hasSwigluLimit": "has(attrs, \"swiglu_limit\")", |
| "swigluLimit": "attrs.swiglu_limit if has(attrs, \"swiglu_limit\") else 0", |
| "quantMidpoint": "quantMidpoint", |
| "quantBlockSize": "quantBlockSize", |
| "packSize": "packSize", |
| "fc1PackedCols": "fc1PackedCols", |
| "fc2PackedCols": "fc2PackedCols", |
| "colWiseScales": "colWiseScales", |
| "fc1ScaleBlocks": "fc1ScaleBlocks", |
| "fc2ScaleBlocks": "fc2ScaleBlocks" |
| }, |
| "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" } |
| ], |
| "fc1Activation": [ |
| { |
| "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": "u32" |
| }, |
| { |
| "name": "fc1_scales", |
| "arg": "fc1ScalesT", |
| "semantic": "fc1_scales", |
| "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, numTokens - repeat.chunk * hiddenChunkTokens)" |
| } |
| ] |
| } |
| } |
| ], |
| "output": [ |
| { "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": "u32" |
| }, |
| { |
| "name": "fc2_scales", |
| "arg": "fc2ScalesT", |
| "semantic": "fc2_scales", |
| "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, numTokens - repeat.chunk * hiddenChunkTokens)" |
| } |
| ] |
| } |
| } |
| ], |
| "zeroInterOutput": [ |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| } |
| ], |
| "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, numTokens - repeat.chunk * hiddenChunkTokens)" |
| } |
| ] |
| } |
| } |
| ], |
| "fc1Grouped": [ |
| { |
| "name": "input", |
| "arg": "inputT", |
| "semantic": "input", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "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": "u32" |
| }, |
| { |
| "name": "fc1_scales", |
| "arg": "fc1ScalesT", |
| "semantic": "fc1_scales", |
| "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" } |
| ] |
| } |
| } |
| ], |
| "outputGrouped": [ |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "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": "u32" |
| }, |
| { |
| "name": "fc2_scales", |
| "arg": "fc2ScalesT", |
| "semantic": "fc2_scales", |
| "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, numTokens - repeat.chunk * hiddenChunkTokens)" |
| } |
| ] |
| } |
| } |
| ], |
| "fc1ActivationIo": [ |
| { |
| "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": "u32" |
| }, |
| { |
| "name": "fc1_scales", |
| "arg": "fc1ScalesT", |
| "semantic": "fc1_scales", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { "name": "hidden_act", "semantic": "hiddenAct", "buffer": { "type": "storage" }, "elementType": "f32" } |
| ], |
| "fc1ActivationScheduled": [ |
| { |
| "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": "u32" |
| }, |
| { |
| "name": "fc1_scales", |
| "arg": "fc1ScalesT", |
| "semantic": "fc1_scales", |
| "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" } |
| ] |
| } |
| } |
| ], |
| "outputIo": [ |
| { "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": "u32" |
| }, |
| { |
| "name": "fc2_scales", |
| "arg": "fc2ScalesT", |
| "semantic": "fc2_scales", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "f32" |
| }, |
| { |
| "name": "output", |
| "arg": "outputT", |
| "semantic": "output", |
| "buffer": { "type": "storage" }, |
| "elementType": "f32" |
| } |
| ], |
| "outputScheduled": [ |
| { "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": "u32" |
| }, |
| { |
| "name": "fc2_scales", |
| "arg": "fc2ScalesT", |
| "semantic": "fc2_scales", |
| "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": "quant_zero_inter", |
| "priority": 20, |
| "when": ["quantContract", "workgroupSizeOk", "interSize == 0"], |
| "passes": [ |
| { |
| "id": "output_stage", |
| "name": "QMoE.OutputStageZeroInter", |
| "source": { |
| "shader": "qmoe-output-zero-inter.wgsl.jinja", |
| "inputs": { "outputElementCount": "numel(shapes.output)" } |
| }, |
| "bindings": "zeroInterOutput", |
| "dispatch": { "threads": "numel(shapes.output)", "workgroupSize": "constants.workgroupSize" } |
| } |
| ] |
| }, |
| { |
| "id": "quant_grouped_sgmat_routed", |
| "priority": 32, |
| "when": ["groupedSgmatOk", "wave32Effective"], |
| "constants": { |
| "groupTileM": "groupTileM", |
| "groupRouteWorkgroup": "tunables.groupRouteWorkgroup", |
| "sgmatWorkgroup": "sgmatWorkgroup", |
| "sgmatSubgroups": "sgmatSubgroups", |
| "sgmatRowSubtiles": "sgmatRowSubtiles" |
| }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[numTokens * topK]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[numTokens * topK]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[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 * hiddenSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "QMoE.RouteStage", |
| "shader": "qmoe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "numTokens", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "group_stage", |
| "name": "QMoE.GroupStage", |
| "shader": "expert-group-slots.wgsl.jinja", |
| "bindings": "group", |
| "dispatch": { |
| "threads": "constants.groupRouteWorkgroup", |
| "workgroupSize": "constants.groupRouteWorkgroup" |
| } |
| }, |
| { |
| "id": "fc1_activation_stage", |
| "name": "QMoE.FC1ActivationStageGroupedSgmat", |
| "shader": "qmoe-fc1-activation-grouped-sgmat.wgsl.jinja", |
| "bindings": "fc1Grouped", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize * fusionSize, 64)" } |
| }, |
| { |
| "id": "output_stage", |
| "name": "QMoE.OutputStageGroupedSgmat", |
| "shader": "qmoe-output-grouped-sgmat.wgsl.jinja", |
| "bindings": "outputGrouped", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenSize, 64)" } |
| }, |
| { |
| "id": "mix_stage", |
| "name": "QMoE.MixStage", |
| "shader": "qmoe-mix-stage.wgsl.jinja", |
| "bindings": "mix", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ], |
| "requires": { |
| "features": ["subgroups", "chromium-experimental-subgroup-matrix"], |
| "limits": { "maxComputeWorkgroupStorageSize": 16512 }, |
| "subgroupMatrixConfigs": [{ "componentType": "f32", "resultComponentType": "f32", "M": 8, "N": 8, "K": 8 }] |
| } |
| }, |
| { |
| "id": "quant_grouped_routed", |
| "priority": 30, |
| "when": ["groupedContract"], |
| "constants": { |
| "groupTileM": "groupTileM", |
| "groupTileN": "groupTileN", |
| "groupTileK": "tunables.groupTileK", |
| "groupTileKVec": "groupTileKVec", |
| "groupThreads": "tunables.groupThreads", |
| "regM": "tunables.groupRegM", |
| "regN": "tunables.groupRegN", |
| "groupRouteWorkgroup": "tunables.groupRouteWorkgroup" |
| }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[numTokens * topK]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[numTokens * topK]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[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 * hiddenSize)]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "QMoE.RouteStage", |
| "shader": "qmoe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "numTokens", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "group_stage", |
| "name": "QMoE.GroupStage", |
| "shader": "expert-group-slots.wgsl.jinja", |
| "bindings": "group", |
| "dispatch": { |
| "threads": "constants.groupRouteWorkgroup", |
| "workgroupSize": "constants.groupRouteWorkgroup" |
| } |
| }, |
| { |
| "id": "fc1_activation_stage", |
| "name": "QMoE.FC1ActivationStageGrouped", |
| "shader": "qmoe-fc1-activation-grouped.wgsl.jinja", |
| "bindings": "fc1Grouped", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(interSize, groupTileN)" } |
| }, |
| { |
| "id": "output_stage", |
| "name": "QMoE.OutputStageGrouped", |
| "shader": "qmoe-output-grouped.wgsl.jinja", |
| "bindings": "outputGrouped", |
| "dispatch": { "x": "groupMaxTiles", "y": "ceilDiv(hiddenSize, groupTileN)" } |
| }, |
| { |
| "id": "mix_stage", |
| "name": "QMoE.MixStage", |
| "shader": "qmoe-mix-stage.wgsl.jinja", |
| "bindings": "mix", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "quant_gemv_routed", |
| "priority": 20, |
| "when": ["decodeContract"], |
| "constants": { "decodeLanes": "tunables.decodeLanes", "decodeRows": "decodeRows" }, |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[numTokens * topK]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[numTokens * topK]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[hiddenChunkTokens * topK * interSize]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "QMoE.RouteStage", |
| "shader": "qmoe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "numTokens", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "fc1_activation_stage", |
| "name": "QMoE.FC1ActivationStageGemv", |
| "shader": "qmoe-fc1-activation-gemv.wgsl.jinja", |
| "bindings": "fc1ActivationScheduled", |
| "dispatch": { |
| "x": "ceilDiv(interSize, decodeRows)", |
| "y": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * topK" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "QMoE.OutputStageGemv", |
| "shader": "qmoe-output-gemv.wgsl.jinja", |
| "bindings": "outputScheduled", |
| "dispatch": { |
| "x": "ceilDiv(hiddenSize, decodeRows)", |
| "y": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens)" |
| } |
| } |
| ] |
| } |
| ] |
| }, |
| { |
| "id": "quant_split_routed", |
| "priority": 10, |
| "when": ["quantContract", "workgroupSizeOk", "hiddenChunkFits", "routedScratchFits", "interSize > 0"], |
| "intermediates": [ |
| { "id": "routeExpert", "dtype": "uint32", "shape": "[numTokens * topK]" }, |
| { "id": "routeMix", "dtype": "float32", "shape": "[numTokens * topK]" }, |
| { "id": "hiddenAct", "dtype": "float32", "shape": "[hiddenChunkTokens * topK * interSize]" } |
| ], |
| "passes": [ |
| { |
| "id": "route_stage", |
| "name": "QMoE.RouteStage", |
| "shader": "qmoe-route-stage.wgsl.jinja", |
| "bindings": "route", |
| "dispatch": { "threads": "numTokens", "workgroupSize": "constants.workgroupSize" } |
| }, |
| { |
| "id": "token_chunks", |
| "repeat": { "count": "hiddenChunkCount", "index": "chunk" }, |
| "passes": [ |
| { |
| "id": "fc1_activation_stage", |
| "name": "QMoE.FC1ActivationStage", |
| "shader": "qmoe-fc1-activation-stage.wgsl.jinja", |
| "bindings": "fc1Activation", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * topK * interSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| }, |
| { |
| "id": "output_stage", |
| "name": "QMoE.OutputStage", |
| "shader": "qmoe-output-stage.wgsl.jinja", |
| "bindings": "output", |
| "dispatch": { |
| "threads": "min(hiddenChunkTokens, numTokens - repeat.chunk * hiddenChunkTokens) * hiddenSize", |
| "workgroupSize": "constants.workgroupSize" |
| } |
| } |
| ] |
| } |
| ] |
| } |
| ] |
| } |
|
|
