| { |
| "domain": "com.microsoft", |
| "name": "GatedRMSNorm", |
| "sinceVersion": 1, |
| "inputs": { |
| "xT": { "onnx": "X", "dtype": "T" }, |
| "scaleT": { "onnx": "scale", "dtype": "T" }, |
| "gateT": { "onnx": "gate", "dtype": "T" } |
| }, |
| "outputs": { "yT": { "onnx": "Y", "dtype": "T", "rank": "ranks.xT", "shape": "shapes.xT" } }, |
| "attributes": { "epsilon": { "default": 0.00001 } }, |
| "typeConstraints": { "T": ["float32", "float16"] }, |
| "tunables": { "WORKGROUP_SIZE": { "default": 256 }, "LANE_WORDS": { "default": 2 } }, |
| "derive": { |
| "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)", |
| "foldedDispatchCapacity": "min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "normSize": "numel(shapes.scaleT)", |
| "lastDim": "dim(shapes.xT, ranks.xT - 1)", |
| "elementCount": "numel(shapes.xT)", |
| "rowCount": "elementCount / max(1, normSize)", |
| "dtypeOk": "tensorDtypes.xT == tensorDtypes.scaleT and tensorDtypes.xT == tensorDtypes.gateT and tensorDtypes.xT == tensorDtypes.yT and f16Ok(dtypes.T)", |
| "shapeOk": "ranks.xT >= 1 and ranks.scaleT >= 1 and normSize > 0 and lastDim % normSize == 0 and sameShape(shapes.gateT, shapes.xT) and sameShape(shapes.yT, shapes.xT)", |
| "contractOk": "dtypeOk and shapeOk", |
| "workgroupOk": "tunables.WORKGROUP_SIZE > 0 and tunables.WORKGROUP_SIZE <= deviceWorkgroupCap and pow2ceil(tunables.WORKGROUP_SIZE) == tunables.WORKGROUP_SIZE and tunables.WORKGROUP_SIZE * 4 <= device.limits.maxComputeWorkgroupStorageSize and tunables.LANE_WORDS > 0", |
| "vec4Ok": "normSize % 4 == 0", |
| "vecWords": "normSize / 4", |
| "vecPreferredLanes": "max(1, min(tunables.WORKGROUP_SIZE / 2, pow2ceil(max(1, ceilDiv(vecWords, max(1, tunables.LANE_WORDS))))))", |
| "vecWideLanes": "max(1, min(tunables.WORKGROUP_SIZE, pow2ceil(max(1, ceilDiv(vecWords, max(1, tunables.LANE_WORDS))))))", |
| "vecLanesPerRow": "pick([[ceilDiv(vecWords, vecPreferredLanes) <= 8, vecPreferredLanes]], vecWideLanes)", |
| "vecRowsPerGroup": "tunables.WORKGROUP_SIZE / max(1, vecLanesPerRow)", |
| "vecWordsPerLane": "ceilDiv(vecWords, max(1, vecLanesPerRow))", |
| "vecRowGroups": "ceilDiv(rowCount, max(1, vecRowsPerGroup))", |
| "scalarPreferredLanes": "1 if normSize <= 8 else (max(1, min(tunables.WORKGROUP_SIZE / 2, pow2ceil(max(1, ceilDiv(normSize, max(1, tunables.LANE_WORDS)))))))", |
| "scalarWideLanes": "max(1, min(tunables.WORKGROUP_SIZE, pow2ceil(max(1, ceilDiv(normSize, max(1, tunables.LANE_WORDS))))))", |
| "scalarLanesPerRow": "pick([[ceilDiv(normSize, scalarPreferredLanes) <= 8, scalarPreferredLanes]], scalarWideLanes)", |
| "scalarRowsPerGroup": "tunables.WORKGROUP_SIZE / max(1, scalarLanesPerRow)", |
| "scalarWordsPerLane": "ceilDiv(normSize, max(1, scalarLanesPerRow))", |
| "scalarRowGroups": "ceilDiv(rowCount, max(1, scalarRowsPerGroup))", |
| "subgroupShapeKnown": "device.features.has(\"subgroups\") and has(device.adapterInfo, \"subgroupMinSize\") and has(device.adapterInfo, \"subgroupMaxSize\") and device.adapterInfo.subgroupMinSize > 0 and device.adapterInfo.subgroupMinSize <= device.adapterInfo.subgroupMaxSize and device.adapterInfo.subgroupMaxSize <= tunables.WORKGROUP_SIZE and tunables.WORKGROUP_SIZE % device.adapterInfo.subgroupMaxSize == 0", |
| "sgMin": "device.adapterInfo.subgroupMinSize if subgroupShapeKnown else 0", |
| "sgFullWidth": "subgroupShapeKnown and device.adapterInfo.subgroupMinSize == device.adapterInfo.subgroupMaxSize", |
| "sgVecLanesPerRow": "min(sgMin, vecLanesPerRow) if subgroupShapeKnown else 1", |
| "sgVecRowsPerGroup": "tunables.WORKGROUP_SIZE / max(1, sgVecLanesPerRow)", |
| "sgVecWordsPerLane": "ceilDiv(vecWords, max(1, sgVecLanesPerRow))", |
| "sgVecRowGroups": "ceilDiv(rowCount, max(1, sgVecRowsPerGroup))", |
| "sgScalarLanesPerRow": "min(sgMin, scalarLanesPerRow) if subgroupShapeKnown else 1", |
| "sgScalarRowsPerGroup": "tunables.WORKGROUP_SIZE / max(1, sgScalarLanesPerRow)", |
| "sgScalarWordsPerLane": "ceilDiv(normSize, max(1, sgScalarLanesPerRow))", |
| "sgScalarRowGroups": "ceilDiv(rowCount, max(1, sgScalarRowsPerGroup))", |
| "vectorElement": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "workgroupSize": "tunables.WORKGROUP_SIZE", |
| "epsilon": "attrs.epsilon", |
| "scalar": "dtypes.T", |
| "vectorScalar": "vectorElement" |
| }, |
| "when": ["contractOk", "workgroupOk"], |
| "bindings": { |
| "x": { "arg": "xT", "elementType": "$vectorElement" }, |
| "scale": { "arg": "scaleT", "elementType": "$vectorElement" }, |
| "gate": { "arg": "gateT", "elementType": "$vectorElement" }, |
| "y": { "arg": "yT", "elementType": "$vectorElement" }, |
| "x_x_t": { "arg": "xT", "name": "x", "elementType": "$scalar" }, |
| "scale_main": { "arg": "scaleT", "name": "scale", "elementType": "$scalar" }, |
| "gate_main": { "arg": "gateT", "name": "gate", "elementType": "$scalar" }, |
| "y_y_t": { "arg": "yT", "name": "y", "elementType": "$scalar" }, |
| "params": { "struct": [{ "name": "rows", "type": "u32", "value": "rowCount" }] } |
| }, |
| "variants": [ |
| { |
| "id": "subgroup_rows_vec4", |
| "priority": 30, |
| "when": ["vec4Ok", "subgroupShapeKnown", "sgVecWordsPerLane >= 1", "sgVecWordsPerLane <= 8", "sgVecRowGroups <= foldedDispatchCapacity"], |
| "requires": { "features": ["subgroups"] }, |
| "derive": { |
| "foldKind": "\"none\" if sgVecLanesPerRow <= 1 else (\"subgroup_full\" if (sgFullWidth and sgVecLanesPerRow == sgMin) else \"subgroup_butterfly\")", |
| "foldSteps": "log2ceil(sgVecLanesPerRow)", |
| "useVec4": "true", |
| "stageRow": "true", |
| "rowWords": "vecWords", |
| "wordsPerLane": "sgVecWordsPerLane", |
| "lanesPerRow": "sgVecLanesPerRow", |
| "rowsPerGroup": "sgVecRowsPerGroup" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GatedRMSNorm.SubgroupRowsVec4", |
| "shader": "gated-rms-norm-rows.wgsl.jinja", |
| "bindings": ["x", "scale", "gate", "y", "params"], |
| "dispatch": { |
| "x": "min(sgVecRowGroups, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(sgVecRowGroups, DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "packed_rows_vec4", |
| "priority": 20, |
| "when": ["vec4Ok", "vecRowGroups <= foldedDispatchCapacity"], |
| "derive": { |
| "foldKind": "\"none\" if vecLanesPerRow <= 1 else \"workgroup_tree\"", |
| "foldSteps": "log2ceil(vecLanesPerRow)", |
| "useVec4": "true", |
| "stageRow": "vecWordsPerLane <= 8", |
| "rowWords": "vecWords", |
| "wordsPerLane": "vecWordsPerLane", |
| "lanesPerRow": "vecLanesPerRow", |
| "rowsPerGroup": "vecRowsPerGroup" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GatedRMSNorm.PackedRowsVec4", |
| "shader": "gated-rms-norm-rows.wgsl.jinja", |
| "bindings": ["x", "scale", "gate", "y", "params"], |
| "dispatch": { |
| "x": "min(vecRowGroups, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(vecRowGroups, DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "subgroup_rows_scalar", |
| "priority": 10, |
| "when": ["subgroupShapeKnown", "sgScalarWordsPerLane >= 1", "sgScalarWordsPerLane <= 8", "sgScalarRowGroups <= foldedDispatchCapacity"], |
| "requires": { "features": ["subgroups"] }, |
| "derive": { |
| "foldKind": "\"none\" if sgScalarLanesPerRow <= 1 else (\"subgroup_full\" if (sgFullWidth and sgScalarLanesPerRow == sgMin) else \"subgroup_butterfly\")", |
| "foldSteps": "log2ceil(sgScalarLanesPerRow)", |
| "useVec4": "false", |
| "stageRow": "true", |
| "rowWords": "normSize", |
| "wordsPerLane": "sgScalarWordsPerLane", |
| "lanesPerRow": "sgScalarLanesPerRow", |
| "rowsPerGroup": "sgScalarRowsPerGroup" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GatedRMSNorm.SubgroupRowsScalar", |
| "shader": "gated-rms-norm-rows.wgsl.jinja", |
| "bindings": ["x_x_t", "scale_main", "gate_main", "y_y_t", "params"], |
| "dispatch": { |
| "x": "min(sgScalarRowGroups, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(sgScalarRowGroups, DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "packed_rows_scalar", |
| "priority": 0, |
| "when": ["scalarRowGroups <= foldedDispatchCapacity"], |
| "derive": { |
| "foldKind": "\"none\" if scalarLanesPerRow <= 1 else \"workgroup_tree\"", |
| "foldSteps": "log2ceil(scalarLanesPerRow)", |
| "useVec4": "false", |
| "stageRow": "scalarWordsPerLane <= 8", |
| "rowWords": "normSize", |
| "wordsPerLane": "scalarWordsPerLane", |
| "lanesPerRow": "scalarLanesPerRow", |
| "rowsPerGroup": "scalarRowsPerGroup" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GatedRMSNorm.PackedRowsScalar", |
| "shader": "gated-rms-norm-rows.wgsl.jinja", |
| "bindings": ["x_x_t", "scale_main", "gate_main", "y_y_t", "params"], |
| "dispatch": { |
| "x": "min(scalarRowGroups, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(scalarRowGroups, DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "subgroup_wide_rows_vec4", |
| "priority": 25, |
| "when": ["vec4Ok", "vecRowGroups <= foldedDispatchCapacity", "dtypes.T == \"f32\"", "sgFullWidth", "device.wgslLanguageFeatures.has(\"subgroup_id\")", "vecLanesPerRow > sgMin", "vecLanesPerRow % max(1, sgMin) == 0"], |
| "requires": { "features": ["subgroups"] }, |
| "derive": { |
| "foldKind": "\"subgroup_tree\"", |
| "useVec4": "true", |
| "stageRow": "vecWordsPerLane <= 8", |
| "rowWords": "vecWords", |
| "wordsPerLane": "vecWordsPerLane", |
| "lanesPerRow": "vecLanesPerRow", |
| "rowsPerGroup": "vecRowsPerGroup", |
| "subgroupWidth": "sgMin", |
| "subgroupsPerRow": "vecLanesPerRow / max(1, sgMin)" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GatedRMSNorm.SubgroupWideRowsVec4", |
| "shader": "gated-rms-norm-rows.wgsl.jinja", |
| "bindings": ["x", "scale", "gate", "y", "params"], |
| "dispatch": { |
| "x": "min(vecRowGroups, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(vecRowGroups, DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| }, |
| { |
| "id": "subgroup_wide_rows_scalar", |
| "priority": 5, |
| "when": ["scalarRowGroups <= foldedDispatchCapacity", "dtypes.T == \"f32\"", "sgFullWidth", "device.wgslLanguageFeatures.has(\"subgroup_id\")", "scalarLanesPerRow > sgMin", "scalarLanesPerRow % max(1, sgMin) == 0"], |
| "requires": { "features": ["subgroups"] }, |
| "derive": { |
| "foldKind": "\"subgroup_tree\"", |
| "useVec4": "false", |
| "stageRow": "scalarWordsPerLane <= 8", |
| "rowWords": "normSize", |
| "wordsPerLane": "scalarWordsPerLane", |
| "lanesPerRow": "scalarLanesPerRow", |
| "rowsPerGroup": "scalarRowsPerGroup", |
| "subgroupWidth": "sgMin", |
| "subgroupsPerRow": "scalarLanesPerRow / max(1, sgMin)" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "GatedRMSNorm.SubgroupWideRowsScalar", |
| "shader": "gated-rms-norm-rows.wgsl.jinja", |
| "bindings": ["x_x_t", "scale_main", "gate_main", "y_y_t", "params"], |
| "dispatch": { |
| "x": "min(scalarRowGroups, DISPATCH_FOLD_WIDTH)", |
| "y": "ceilDiv(scalarRowGroups, DISPATCH_FOLD_WIDTH)" |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|
