Xenova's picture
Xenova HF Staff
sync 2e7068faf55e
48220c7 verified
Raw
History Blame
22.3 kB
{
"domain": "ai.onnx",
"name": "SimplifiedLayerNormalization",
"conformance": "legacy-default-domain",
"sinceVersion": 1,
"description": "Implements ONNX Runtime's legacy RMS normalization for models that serialize `SimplifiedLayerNormalization` in the default ONNX domain: `Y = (X / sqrt(mean(X^2) + epsilon)) * scale`. It shares kernels with [ONNX `RMSNormalization`](https://onnx.ai/onnx/operators/onnx__RMSNormalization.html), but retains scale-before-output-cast ordering and optional `inv_std_var`. See the [compatibility discussion](https://github.com/onnx/onnx/issues/6582#issuecomment-3591862327). Supports float16/float32, float32 statistics, and a nonempty normalization suffix.",
"inputs": [
{
"role": "X",
"dtype": "T",
"description": "Input tensor; the mean square is taken over the suffix dimensions starting at `axis`."
},
{
"role": "scale",
"dtype": "V",
"description": "Scale tensor, unidirectionally broadcastable to `X`; output `Y` has this tensor's dtype."
}
],
"outputs": [
{
"role": "Y",
"dtype": "V",
"rank": "ranks.X",
"shape": "shapes.X",
"description": "Normalized and scaled output; same shape as `X` and same dtype as `scale`."
},
{
"role": "inv_std_var",
"dtype": "U",
"rank": "ranks.X",
"shape": "prefix(shapes.X, axisNorm) + fill(1, ranks.X - axisNorm)",
"optional": true,
"description": "Optional inverse RMS statistic `1 / sqrt(mean(X^2) + epsilon)`; same rank as `X`, with the dimensions from `axis` onward set to one."
}
],
"attributes": { "axis": -1, "epsilon": 0.00001, "stash_type": 1, "keep_dims": 1 },
"attributeDescriptions": {
"axis": "The first dimension of the normalization suffix; negative values count from the end, so the default `-1` normalizes 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 of `inv_std_var`; this package implements the legacy float32 value `1`.",
"keep_dims": "Legacy compatibility attribute accepted and ignored by ONNX Runtime. The supported value `1` records the runtime's fixed behavior: optional statistics retain the input rank with normalized dimensions set to one."
},
"attributeConstraints": { "stash_type": { "values": [1] }, "keep_dims": { "values": [1] } },
"typeConstraints": { "T": ["float32", "float16"], "V": ["float32", "float16"], "U": ["float32"] },
"args": {
"x": { "kind": "tensor", "semantic": "X", "role": "input" },
"scale": { "kind": "tensor", "semantic": "scale", "role": "input" },
"y": { "kind": "tensor", "semantic": "Y", "role": "output" },
"invStdVar": { "kind": "tensor", "semantic": "inv_std_var", "role": "output", "required": false }
},
"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",
"normRows": "outer(shapes.X, axisNorm)",
"normHidden": "dim(shapes.X, axisNorm) * inner(shapes.X, axisNorm)",
"normRowStride": "max(1, min(normRows, device.limits.maxComputeWorkgroupsPerDimension))",
"rowWg": "min(normMaxWorkgroup, pow2ceil(max(1, normHidden)))",
"baseOk": "ranks.X >= 1 and sameShape(shapes.Y, shapes.X) and ranks.scale >= 0 and ranks.scale <= ranks.X and broadcastable(shapes.scale, shapes.X) and attrs.axis + ranks.X >= 0 and attrs.axis < ranks.X and normHidden > 0 and attrs.stash_type == onnxDtypeCode(\"float32\") and f16Ok(dtypes.T) and f16Ok(dtypes.V)",
"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)",
"noStats": "not present.invStdVar",
"statsOk": "present.invStdVar and ranks.inv_std_var == ranks.X and sameShape(prefix(shapes.inv_std_var, axisNorm), prefix(shapes.X, axisNorm)) and numel(suffix(shapes.inv_std_var, axisNorm)) == 1",
"sameDtype": "dtypes.T == dtypes.V",
"splitCount": "min(tunables.MAX_SPLITS, pow2ceil(ceilDiv(normHidden, tunables.SPLIT_TARGET_ELEMENTS)))",
"splitScratchBytes": "normRows * splitCount * 4",
"splitFits": "normRows <= tunables.SPLIT_MAX_ROWS and splitCount <= device.limits.maxComputeWorkgroupsPerDimension and splitScratchBytes <= device.limits.maxStorageBufferBindingSize and splitScratchBytes <= device.limits.maxBufferSize"
},
"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": "normRows" },
{ "name": "rowStride", "type": "u32", "value": "normRowStride" }
]
}
}
],
"rowsStats": [
{
"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": "inv_std_out",
"arg": "invStdVar",
"semantic": "inv_std_var",
"buffer": { "type": "storage" },
"elementType": "f32"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": {
"name": "Params",
"fields": [
{ "name": "rows", "type": "u32", "value": "normRows" },
{ "name": "rowStride", "type": "u32", "value": "normRowStride" }
]
}
}
],
"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": "normRows" },
{ "name": "rowStride", "type": "u32", "value": "normRowStride" }
]
}
}
],
"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": "normRows" },
{ "name": "rowStride", "type": "u32", "value": "normRowStride" }
]
}
}
],
"splitNormalizeStats": [
{
"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": "inv_std_out",
"arg": "invStdVar",
"semantic": "inv_std_var",
"buffer": { "type": "storage" },
"elementType": "f32"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": {
"name": "Params",
"fields": [
{ "name": "rows", "type": "u32", "value": "normRows" },
{ "name": "rowStride", "type": "u32", "value": "normRowStride" }
]
}
}
]
},
"variants": [
{
"id": "last_axis",
"priority": 1,
"when": ["lastAxisOk", "noStats"],
"constants": {
"scalar": "dtypes.V",
"xElement": "dtypes.T",
"ioElement": "dtypes.V",
"usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
"hiddenSize": "normHidden",
"workgroupSize": "rowWg",
"epsilon": "attrs.epsilon"
},
"passes": [
{
"id": "main",
"name": "SimplifiedLayerNormalization.Row",
"source": {
"shader": "rms-normalization.wgsl.jinja",
"inputs": {
"xShape": "shapes.X",
"scaleShape": "shapes.scale",
"xRank": "ranks.X",
"scaleRank": "ranks.scale",
"writeStats": false,
"rmsScaleAfterCast": false
}
},
"bindings": "rows",
"dispatch": { "workgroups": "normRows" }
}
]
},
{
"id": "last_axis_stats",
"priority": 2,
"when": ["lastAxisOk", "statsOk"],
"constants": {
"scalar": "dtypes.V",
"xElement": "dtypes.T",
"ioElement": "dtypes.V",
"usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
"hiddenSize": "normHidden",
"workgroupSize": "rowWg",
"epsilon": "attrs.epsilon"
},
"passes": [
{
"id": "main",
"name": "SimplifiedLayerNormalization.Row",
"source": {
"shader": "rms-normalization.wgsl.jinja",
"inputs": {
"xShape": "shapes.X",
"scaleShape": "shapes.scale",
"xRank": "ranks.X",
"scaleRank": "ranks.scale",
"writeStats": true,
"rmsScaleAfterCast": false
}
},
"bindings": "rowsStats",
"dispatch": { "workgroups": "normRows" }
}
]
},
{
"id": "suffix_axis",
"priority": 10,
"when": ["suffixAxisOk", "noStats"],
"constants": {
"scalar": "dtypes.V",
"xElement": "dtypes.T",
"ioElement": "dtypes.V",
"usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
"hiddenSize": "normHidden",
"workgroupSize": "rowWg",
"epsilon": "attrs.epsilon"
},
"passes": [
{
"id": "main",
"name": "SimplifiedLayerNormalization.Row",
"source": {
"shader": "rms-normalization.wgsl.jinja",
"inputs": {
"xShape": "shapes.X",
"scaleShape": "shapes.scale",
"xRank": "ranks.X",
"scaleRank": "ranks.scale",
"writeStats": false,
"rmsScaleAfterCast": false
}
},
"bindings": "rows",
"dispatch": { "workgroups": "normRows" }
}
]
},
{
"id": "suffix_axis_stats",
"priority": 11,
"when": ["suffixAxisOk", "statsOk"],
"constants": {
"scalar": "dtypes.V",
"xElement": "dtypes.T",
"ioElement": "dtypes.V",
"usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
"hiddenSize": "normHidden",
"workgroupSize": "rowWg",
"epsilon": "attrs.epsilon"
},
"passes": [
{
"id": "main",
"name": "SimplifiedLayerNormalization.Row",
"source": {
"shader": "rms-normalization.wgsl.jinja",
"inputs": {
"xShape": "shapes.X",
"scaleShape": "shapes.scale",
"xRank": "ranks.X",
"scaleRank": "ranks.scale",
"writeStats": true,
"rmsScaleAfterCast": false
}
},
"bindings": "rowsStats",
"dispatch": { "workgroups": "normRows" }
}
]
},
{
"id": "suffix_axis_splitk",
"priority": 15,
"when": ["baseOk", "ranks.X >= 2", "noStats", "splitFits"],
"demoteWhen": ["reportedNonWave32Adapter", "normHidden < tunables.SPLIT_MIN_HIDDEN"],
"constants": {
"scalar": "dtypes.V",
"xElement": "dtypes.T",
"ioElement": "dtypes.V",
"usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
"hiddenSize": "normHidden",
"workgroupSize": "normMaxWorkgroup",
"split": "splitCount",
"epsilon": "attrs.epsilon"
},
"intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[normRows * splitCount]" }],
"passes": [
{
"id": "partials",
"name": "SimplifiedLayerNormalization.SplitKPartials",
"source": { "shader": "rms-normalization-splitk-partials.wgsl.jinja", "inputs": {} },
"bindings": "splitPartials",
"dispatch": { "workgroups": "normRows", "z": "splitCount" }
},
{
"id": "normalize",
"name": "SimplifiedLayerNormalization.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": false
}
},
"bindings": "splitNormalize",
"dispatch": { "workgroups": "normRows", "z": "splitCount" }
}
]
},
{
"id": "suffix_axis_splitk_stats",
"priority": 16,
"when": ["baseOk", "ranks.X >= 2", "statsOk", "splitFits"],
"demoteWhen": ["reportedNonWave32Adapter", "normHidden < tunables.SPLIT_MIN_HIDDEN"],
"constants": {
"scalar": "dtypes.V",
"xElement": "dtypes.T",
"ioElement": "dtypes.V",
"usesF16": "dtypes.T == \"f16\" or dtypes.V == \"f16\"",
"hiddenSize": "normHidden",
"workgroupSize": "normMaxWorkgroup",
"split": "splitCount",
"epsilon": "attrs.epsilon"
},
"intermediates": [{ "id": "partials", "dtype": "float32", "shape": "[normRows * splitCount]" }],
"passes": [
{
"id": "partials",
"name": "SimplifiedLayerNormalization.SplitKPartials",
"source": { "shader": "rms-normalization-splitk-partials.wgsl.jinja", "inputs": {} },
"bindings": "splitPartials",
"dispatch": { "workgroups": "normRows", "z": "splitCount" }
},
{
"id": "normalize",
"name": "SimplifiedLayerNormalization.SplitKNormalize",
"source": {
"shader": "rms-normalization-splitk-normalize.wgsl.jinja",
"inputs": {
"xShape": "shapes.X",
"scaleShape": "shapes.scale",
"xRank": "ranks.X",
"scaleRank": "ranks.scale",
"writeStats": true,
"rmsScaleAfterCast": false
}
},
"bindings": "splitNormalizeStats",
"dispatch": { "workgroups": "normRows", "z": "splitCount" }
}
]
},
{
"id": "last_axis_row_vec4",
"priority": 110,
"when": ["lastAxisOk", "sameDtype", "noStats", "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": {
"scalar": "dtypes.T",
"xElement": "\"vec4<\" ~ dtypes.T ~ \">\"",
"ioElement": "\"vec4<\" ~ dtypes.T ~ \">\""
},
"passes": [
{
"id": "main",
"name": "SimplifiedLayerNormalization.LastAxisRow",
"source": {
"shader": "norm-row-stats.wgsl.jinja",
"inputs": {
"mode": "\"rms\"",
"vec4": true,
"writeStats": false,
"rmsScaleAfterCast": false,
"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": "normRows" }
}
]
},
{
"id": "last_axis_row",
"priority": 100,
"when": ["lastAxisOk", "sameDtype", "noStats", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.X, -1)", "dim(shapes.scale, -1) == dim(shapes.X, -1)"],
"constants": { "scalar": "dtypes.T", "xElement": "dtypes.T", "ioElement": "dtypes.T" },
"passes": [
{
"id": "main",
"name": "SimplifiedLayerNormalization.LastAxisRow",
"source": {
"shader": "norm-row-stats.wgsl.jinja",
"inputs": {
"mode": "\"rms\"",
"vec4": false,
"writeStats": false,
"rmsScaleAfterCast": false,
"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": "normRows" }
}
]
},
{
"id": "last_axis_row_vec4_stats",
"priority": 112,
"when": ["lastAxisOk", "sameDtype", "statsOk", "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": {
"scalar": "dtypes.T",
"xElement": "\"vec4<\" ~ dtypes.T ~ \">\"",
"ioElement": "\"vec4<\" ~ dtypes.T ~ \">\""
},
"passes": [
{
"id": "main",
"name": "SimplifiedLayerNormalization.LastAxisRow",
"source": {
"shader": "norm-row-stats.wgsl.jinja",
"inputs": {
"mode": "\"rms\"",
"vec4": true,
"writeStats": true,
"rmsScaleAfterCast": false,
"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": "rowsStats",
"dispatch": { "workgroups": "normRows" }
}
]
},
{
"id": "last_axis_row_stats",
"priority": 102,
"when": ["lastAxisOk", "sameDtype", "statsOk", "ranks.scale >= 1", "numel(shapes.scale) == dim(shapes.X, -1)", "dim(shapes.scale, -1) == dim(shapes.X, -1)"],
"constants": { "scalar": "dtypes.T", "xElement": "dtypes.T", "ioElement": "dtypes.T" },
"passes": [
{
"id": "main",
"name": "SimplifiedLayerNormalization.LastAxisRow",
"source": {
"shader": "norm-row-stats.wgsl.jinja",
"inputs": {
"mode": "\"rms\"",
"vec4": false,
"writeStats": true,
"rmsScaleAfterCast": false,
"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": "rowsStats",
"dispatch": { "workgroups": "normRows" }
}
]
}
]
}