{
  "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" }
        }
      ]
    }
  ]
}