| { |
| "domain": "ai.onnx", |
| "name": "RMSNormalization", |
| "sinceVersion": 23, |
| "description": "Computes RMS normalization over the suffix dimensions of `X` starting at `axis`: `Y = X / sqrt(mean(X^2) + epsilon) * scale`. The normalization stage supports TensorProto `stash_type` values `1` (float32) and `10` (float16), and is cast back to the dtype of `X` before `scale` is applied. The input type `T` and scale/output type `V` may independently be float16 or float32; ONNX's bfloat16 and double cases are not yet implemented.", |
| "inputs": [ |
| { |
| "role": "X", |
| "dtype": "T", |
| "description": "Input tensor to be normalized; the RMS is taken over the last dimensions starting at `axis`." |
| }, |
| { |
| "role": "scale", |
| "dtype": "V", |
| "description": "Scale tensor, unidirectionally broadcastable to `X`; its dtype `V` may differ from the input dtype `T`." |
| } |
| ], |
| "outputs": [ |
| { |
| "role": "Y", |
| "dtype": "V", |
| "rank": "ranks.X", |
| "shape": "shapes.X", |
| "description": "Normalized and scaled output tensor; same shape as `X` and same dtype `V` as `scale`." |
| } |
| ], |
| "attributes": { "axis": -1, "epsilon": 0.00001, "stash_type": 1 }, |
| "attributeDescriptions": { |
| "axis": "The first dimension of the normalization suffix; negative values count from the end, so the default `-1` normalizes over only the last dimension.", |
| "epsilon": "Small constant added to the mean square before taking the square root to avoid division by zero.", |
| "stash_type": "TensorProto element type used for normalization: `1` computes in float32, while `10` computes in float16." |
| }, |
| "attributeConstraints": { "stash_type": { "values": [1, 10] } }, |
| "typeConstraints": { "T": ["float32", "float16"], "V": ["float32", "float16"] }, |
| "args": { |
| "x": { "kind": "tensor", "semantic": "X", "role": "input" }, |
| "scale": { "kind": "tensor", "semantic": "scale", "role": "input" }, |
| "y": { "kind": "tensor", "semantic": "Y", "role": "output" } |
| }, |
| "tunables": { |
| "WORKGROUP_SIZE": 256, |
| "SPLIT_MAX_ROWS": 256, |
| "SPLIT_MIN_HIDDEN": 16384, |
| "SPLIT_TARGET_ELEMENTS": 4096, |
| "MAX_SPLITS": 64 |
| }, |
| "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", |
| "reportedNonWave32Adapter": "not wave32Adapter and (has(device.adapterInfo, \"subgroupMinSize\") or has(device.adapterInfo, \"subgroupMaxSize\"))", |
| "normMaxWorkgroup": "min(tunables.WORKGROUP_SIZE, deviceWorkgroupCap)", |
| "hasSubgroupId": "device.features.has(\"subgroups\") and device.wgslLanguageFeatures.has(\"subgroup_id\")", |
| "axisNorm": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.X", |
| "normalizedRows": "outer(shapes.X, axisNorm)", |
| "normalizedHidden": "dim(shapes.X, axisNorm) * inner(shapes.X, axisNorm)", |
| "normalizedDispatchRows": "0 if normalizedHidden == 0 else normalizedRows", |
| "normalizedWorkgroupHidden": "max(1, normalizedHidden)", |
| "normalizationShapeOk": "ranks.X >= 1 and ranks.scale >= 0 and ranks.scale <= ranks.X and sameShape(shapes.Y, shapes.X) and attrs.axis + ranks.X >= 0 and attrs.axis < ranks.X and broadcastable(shapes.scale, shapes.X) and f16Ok(dtypes.T) and f16Ok(dtypes.V)", |
| "baseOk": "normalizationShapeOk and attrs.stash_type == onnxDtypeCode(\"float32\")", |
| "stashF16Ok": "normalizationShapeOk and attrs.stash_type == onnxDtypeCode(\"float16\")", |
| "lastAxisOk": "baseOk and (attrs.axis == -1 or attrs.axis == ranks.X - 1)", |
| "suffixAxisOk": "baseOk and ranks.X >= 2 and not (attrs.axis == -1 or attrs.axis == ranks.X - 1)" |
| }, |
| "bindingSets": { |
| "rows": [ |
| { |
| "name": "x", |
| "arg": "x", |
| "semantic": "X", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$xElement" |
| }, |
| { |
| "name": "scale", |
| "arg": "scale", |
| "semantic": "scale", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$ioElement" |
| }, |
| { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$ioElement" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "rows", "type": "u32", "value": "normalizedRows" }, |
| { |
| "name": "rowStride", |
| "type": "u32", |
| "value": "max(1, min(normalizedRows, device.limits.maxComputeWorkgroupsPerDimension))" |
| } |
| ] |
| } |
| } |
| ], |
| "splitPartials": [ |
| { |
| "name": "x", |
| "arg": "x", |
| "semantic": "X", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$xElement" |
| }, |
| { "name": "partials", "semantic": "partials", "buffer": { "type": "storage" }, "elementType": "f32" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "rows", "type": "u32", "value": "splitRows" }, |
| { |
| "name": "rowStride", |
| "type": "u32", |
| "value": "max(1, min(splitRows, device.limits.maxComputeWorkgroupsPerDimension))" |
| } |
| ] |
| } |
| } |
| ], |
| "splitNormalize": [ |
| { |
| "name": "x", |
| "arg": "x", |
| "semantic": "X", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$xElement" |
| }, |
| { |
| "name": "scale", |
| "arg": "scale", |
| "semantic": "scale", |
| "buffer": { "type": "read-only-storage" }, |
| "elementType": "$ioElement" |
| }, |
| { "name": "partials", "semantic": "partials", "buffer": { "type": "read-only-storage" }, "elementType": "f32" }, |
| { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$ioElement" }, |
| { |
| "name": "params", |
| "semantic": "kernel.params", |
| "buffer": { "type": "uniform" }, |
| "struct": { |
| "name": "Params", |
| "fields": [ |
| { "name": "rows", "type": "u32", "value": "splitRows" }, |
| { |
| "name": "rowStride", |
| "type": "u32", |
| "value": "max(1, min(splitRows, device.limits.maxComputeWorkgroupsPerDimension))" |
| } |
| ] |
| } |
| } |
| ] |
| }, |
| "variants": [ |
| { |
| "id": "stash_f16_serial", |
| "priority": 1000, |
| "when": "stashF16Ok", |
| "constants": { |
| "scalar": "dtypes.V", |
| "xElement": "dtypes.T", |
| "ioElement": "dtypes.V", |
| "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"", |
| "hiddenSize": "normalizedHidden", |
| "epsilon": "attrs.epsilon" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RMSNormalization.StashF16Serial", |
| "source": { |
| "shader": "rms-normalization-stash-f16-serial.wgsl.jinja", |
| "inputs": { |
| "xShape": "shapes.X", |
| "scaleShape": "shapes.scale", |
| "xRank": "ranks.X", |
| "scaleRank": "ranks.scale" |
| } |
| }, |
| "bindings": "rows", |
| "dispatch": { "workgroups": "normalizedDispatchRows" } |
| } |
| ] |
| }, |
| { |
| "id": "suffix_axis_splitk", |
| "priority": 15, |
| "derive": { |
| "splitRows": "normalizedRows", |
| "splitHidden": "normalizedHidden", |
| "split": "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(splitHidden, tunables.SPLIT_TARGET_ELEMENTS)))" |
| }, |
| "when": ["baseOk", "ranks.X >= 2", "normalizedRows <= tunables.SPLIT_MAX_ROWS", "normalizedHidden >= tunables.SPLIT_MIN_HIDDEN", "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) <= device.limits.maxComputeWorkgroupsPerDimension", "normalizedRows * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) * 4 <= device.limits.maxStorageBufferBindingSize", "normalizedRows * min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normalizedHidden, tunables.SPLIT_TARGET_ELEMENTS))) * 4 <= device.limits.maxBufferSize"], |
| "demoteWhen": ["reportedNonWave32Adapter"], |
| "constants": { |
| "scalar": "dtypes.V", |
| "xElement": "dtypes.T", |
| "ioElement": "dtypes.V", |
| "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"", |
| "hiddenSize": "splitHidden", |
| "workgroupSize": "normMaxWorkgroup", |
| "split": "split", |
| "epsilon": "attrs.epsilon" |
| }, |
| "intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[splitRows * split]" }], |
| "passes": [ |
| { |
| "id": "partials", |
| "name": "RMSNormalization.SplitKPartials", |
| "source": { "shader": "rms-normalization-splitk-partials.wgsl.jinja", "inputs": {} }, |
| "bindings": "splitPartials", |
| "dispatch": { "workgroups": "splitRows", "z": "split" } |
| }, |
| { |
| "id": "normalize", |
| "name": "RMSNormalization.SplitKNormalize", |
| "source": { |
| "shader": "rms-normalization-splitk-normalize.wgsl.jinja", |
| "inputs": { |
| "xShape": "shapes.X", |
| "scaleShape": "shapes.scale", |
| "xRank": "ranks.X", |
| "scaleRank": "ranks.scale", |
| "writeStats": false, |
| "rmsScaleAfterCast": true |
| } |
| }, |
| "bindings": "splitNormalize", |
| "dispatch": { "workgroups": "splitRows", "z": "split" } |
| } |
| ] |
| }, |
| { |
| "id": "last_axis", |
| "priority": 0, |
| "when": "lastAxisOk", |
| "constants": { |
| "scalar": "dtypes.V", |
| "xElement": "dtypes.T", |
| "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"", |
| "ioElement": "dtypes.V", |
| "hiddenSize": "normalizedHidden", |
| "workgroupSize": "min(normMaxWorkgroup, pow2ceil(normalizedWorkgroupHidden))", |
| "epsilon": "attrs.epsilon" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RMSNormalization", |
| "source": { |
| "shader": "rms-normalization.wgsl.jinja", |
| "inputs": { |
| "xShape": "shapes.X", |
| "scaleShape": "shapes.scale", |
| "xRank": "ranks.X", |
| "scaleRank": "ranks.scale", |
| "writeStats": false, |
| "rmsScaleAfterCast": true |
| } |
| }, |
| "bindings": "rows", |
| "dispatch": { "workgroups": "normalizedDispatchRows" } |
| } |
| ] |
| }, |
| { |
| "id": "suffix_axis", |
| "priority": 10, |
| "when": "suffixAxisOk", |
| "constants": { |
| "scalar": "dtypes.V", |
| "xElement": "dtypes.T", |
| "usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"", |
| "ioElement": "dtypes.V", |
| "hiddenSize": "normalizedHidden", |
| "workgroupSize": "min(normMaxWorkgroup, pow2ceil(normalizedWorkgroupHidden))", |
| "epsilon": "attrs.epsilon" |
| }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RMSNormalization.SuffixAxis", |
| "source": { |
| "shader": "rms-normalization.wgsl.jinja", |
| "inputs": { |
| "xShape": "shapes.X", |
| "scaleShape": "shapes.scale", |
| "xRank": "ranks.X", |
| "scaleRank": "ranks.scale", |
| "writeStats": false, |
| "rmsScaleAfterCast": true |
| } |
| }, |
| "bindings": "rows", |
| "dispatch": { "workgroups": "normalizedDispatchRows" } |
| } |
| ] |
| }, |
| { |
| "id": "last_axis_row_vec4", |
| "priority": 110, |
| "when": ["lastAxisOk", "dtypes.T == dtypes.V", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.X, -1)", "dim(shapes.scale, -1) == dim(shapes.X, -1)", "dim(shapes.X, -1) % 4 == 0"], |
| "constants": { "xElement": "\"vec4<\" ~ dtypes.T ~ \">\"", "ioElement": "\"vec4<\" ~ dtypes.T ~ \">\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RMSNormalization.LastAxisRow", |
| "source": { |
| "shader": "norm-row-stats.wgsl.jinja", |
| "inputs": { |
| "mode": "\"rms\"", |
| "vec4": true, |
| "writeStats": false, |
| "rmsScaleAfterCast": true, |
| "scalar": "dtypes.T", |
| "usesF16": "dtypes.T == \"f16\"", |
| "hidden": "dim(shapes.X, -1)", |
| "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.X, -1) / 4)))", |
| "epsilon": "attrs.epsilon", |
| "hiddenVec": "dim(shapes.X, -1) / 4", |
| "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "combineSubgroups": "hasSubgroupId" |
| } |
| }, |
| "subgroupCollectivesWidth": "portable", |
| "bindings": "rows", |
| "dispatch": { "workgroups": "normalizedDispatchRows" } |
| } |
| ] |
| }, |
| { |
| "id": "last_axis_row", |
| "priority": 100, |
| "when": ["lastAxisOk", "dtypes.T == dtypes.V", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.X, -1)", "dim(shapes.scale, -1) == dim(shapes.X, -1)", "true"], |
| "constants": { "xElement": "dtypes.T", "ioElement": "dtypes.T" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RMSNormalization.LastAxisRow", |
| "source": { |
| "shader": "norm-row-stats.wgsl.jinja", |
| "inputs": { |
| "mode": "\"rms\"", |
| "vec4": false, |
| "writeStats": false, |
| "rmsScaleAfterCast": true, |
| "scalar": "dtypes.T", |
| "usesF16": "dtypes.T == \"f16\"", |
| "hidden": "dim(shapes.X, -1)", |
| "wg": "min(normMaxWorkgroup, pow2ceil(max(1, dim(shapes.X, -1))))", |
| "epsilon": "attrs.epsilon", |
| "hiddenVec": 1, |
| "vecType": "\"vec4<\" ~ dtypes.T ~ \">\"", |
| "combineSubgroups": "hasSubgroupId" |
| } |
| }, |
| "subgroupCollectivesWidth": "portable", |
| "bindings": "rows", |
| "dispatch": { "workgroups": "normalizedDispatchRows" } |
| } |
| ] |
| } |
| ] |
| } |
|
|
