{
  "domain": "ai.onnx",
  "name": "RotaryEmbedding",
  "sinceVersion": 23,
  "description": "Implements ONNX opset-23 RotaryEmbedding for float16 and float32 tensors. Applies rotary positional embeddings (RoPE) by rotating each head's embedding vector using precomputed `cos_cache` and `sin_cache` values. A partial rotation can be applied by setting `rotary_embedding_dim` to rotate only a prefix of the head dimension. `position_ids` keeps its standard logical int64 type; valid positions are non-negative and bounded by the WebGPU-addressable cache, so the backend stores them losslessly as uint32. Other ONNX floating-point input types are not yet implemented.",
  "inputs": [
    {
      "role": "X",
      "dtype": "T",
      "description": "Input token embeddings. Shape is `(batch_size, sequence_length, hidden_size)` for rank 3 or `(batch_size, num_heads, sequence_length, head_size)` for rank 4. `head_size` must be even, and the `num_heads` attribute is required for rank-3 input."
    },
    {
      "role": "cos_cache",
      "dtype": "T",
      "description": "Precomputed cosine values. Without `position_ids`, shape is `(batch_size, sequence_length, rotary_dim/2)`; with `position_ids`, shape is `(max_sequence_length, rotary_dim/2)`."
    },
    {
      "role": "sin_cache",
      "dtype": "T",
      "description": "Precomputed sine values with the same shape and type as `cos_cache`."
    },
    {
      "role": "position_ids",
      "dtype": "M",
      "rank": 2,
      "optional": true,
      "description": "Optional logical int64 per-token position indices of shape `(batch_size, sequence_length)`. Valid positions are non-negative cache-row indices and use uint32 WebGPU storage. When supplied, the 2D cache tables are gathered at these positions."
    }
  ],
  "outputs": [
    {
      "role": "Y",
      "dtype": "T",
      "rank": "ranks.X",
      "shape": "shapes.X",
      "description": "Rotary-position-encoded tensor with the same shape and type as `X`."
    }
  ],
  "attributes": { "interleaved": 0, "rotary_embedding_dim": 0 },
  "attributeDescriptions": {
    "interleaved": "Set to 1 to rotate using an interleaved pattern (even/odd elements), or 0 to split the head dimension into two contiguous halves. Default is 0.",
    "num_heads": "Optional number of attention heads. ONNX requires this attribute when `X` is rank 3; it is unnecessary for rank-4 input because the head count is explicit in the shape.",
    "rotary_embedding_dim": "Number of head-dimension elements to rotate; `0` means rotate the full head dimension. When set, only the leading `rotary_embedding_dim` elements are rotated and the rest are passed through unchanged."
  },
  "attributeConstraints": { "interleaved": { "values": [0, 1] } },
  "typeConstraints": { "T": ["float32", "float16"], "M": ["int64"] },
  "args": {
    "x": { "kind": "tensor", "semantic": "X", "role": "input" },
    "cos": { "kind": "tensor", "semantic": "cos_cache", "role": "input" },
    "sin": { "kind": "tensor", "semantic": "sin_cache", "role": "input" },
    "positionIds": {
      "kind": "tensor",
      "semantic": "position_ids",
      "role": "input",
      "dtype": "uint32",
      "narrowing": "checked",
      "required": false
    },
    "y": { "kind": "tensor", "semantic": "Y", "role": "output" }
  },
  "tunables": { "WORKGROUP_SIZE": 256 },
  "derive": {
    "pairDispatchOk": "ceilDiv(numel(shapes.X) / 2, tunables.WORKGROUP_SIZE) <= device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension",
    "rank3HeadSize": "dim(shapes.X, 2) / attrs.num_heads if attrs.num_heads is defined and attrs.num_heads > 0 else 0",
    "rank3Contract": "f16Ok(dtypes.T) and ranks.X == 3 and ranks.cos_cache == 3 and ranks.sin_cache == 3 and ranks.Y == 3 and sameShape(shapes.X, shapes.Y) and attrs.num_heads is defined and attrs.num_heads >= 1 and dim(shapes.X, 2) % attrs.num_heads == 0 and rank3HeadSize % 2 == 0 and (attrs.rotary_embedding_dim == 0 or attrs.rotary_embedding_dim <= rank3HeadSize) and dim(shapes.cos_cache, 2) * 2 == (attrs.rotary_embedding_dim if attrs.rotary_embedding_dim != 0 else rank3HeadSize) and sameShape(shapes.cos_cache, shapes.sin_cache) and dim(shapes.cos_cache, 0) == dim(shapes.X, 0) and dim(shapes.cos_cache, 1) == dim(shapes.X, 1) and not present.positionIds",
    "rank4Contract": "f16Ok(dtypes.T) and ranks.X == 4 and ranks.cos_cache == 3 and ranks.sin_cache == 3 and ranks.Y == 4 and dim(shapes.X, 3) % 2 == 0 and (attrs.rotary_embedding_dim == 0 or attrs.rotary_embedding_dim <= dim(shapes.X, 3)) and dim(shapes.cos_cache, 2) * 2 == (attrs.rotary_embedding_dim if attrs.rotary_embedding_dim != 0 else dim(shapes.X, 3)) and dim(shapes.sin_cache, 2) == dim(shapes.cos_cache, 2) and sameShape(shapes.X, shapes.Y) and dim(shapes.cos_cache, 0) == dim(shapes.X, 0) and dim(shapes.sin_cache, 0) == dim(shapes.X, 0) and dim(shapes.cos_cache, 1) == dim(shapes.X, 2) and dim(shapes.sin_cache, 1) == dim(shapes.X, 2) and not present.positionIds",
    "rank3PosContract": "f16Ok(dtypes.T) and ranks.X == 3 and ranks.cos_cache == 2 and ranks.sin_cache == 2 and ranks.Y == 3 and sameShape(shapes.X, shapes.Y) and attrs.num_heads is defined and attrs.num_heads >= 1 and dim(shapes.X, 2) % attrs.num_heads == 0 and rank3HeadSize % 2 == 0 and (attrs.rotary_embedding_dim == 0 or attrs.rotary_embedding_dim <= rank3HeadSize) and dim(shapes.cos_cache, 1) * 2 == (attrs.rotary_embedding_dim if attrs.rotary_embedding_dim != 0 else rank3HeadSize) and sameShape(shapes.cos_cache, shapes.sin_cache) and present.positionIds and ranks.position_ids == 2 and dim(shapes.position_ids, 0) == dim(shapes.X, 0) and dim(shapes.position_ids, 1) == dim(shapes.X, 1)",
    "rank4PosContract": "f16Ok(dtypes.T) and ranks.X == 4 and ranks.cos_cache == 2 and ranks.sin_cache == 2 and ranks.Y == 4 and dim(shapes.X, 3) % 2 == 0 and (attrs.rotary_embedding_dim == 0 or attrs.rotary_embedding_dim <= dim(shapes.X, 3)) and dim(shapes.cos_cache, 1) * 2 == (attrs.rotary_embedding_dim if attrs.rotary_embedding_dim != 0 else dim(shapes.X, 3)) and sameShape(shapes.cos_cache, shapes.sin_cache) and sameShape(shapes.X, shapes.Y) and present.positionIds and ranks.position_ids == 2 and dim(shapes.position_ids, 0) == dim(shapes.X, 0) and dim(shapes.position_ids, 1) == dim(shapes.X, 2)"
  },
  "bindingSets": {
    "rank3": [
      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
      {
        "name": "cos_cache",
        "arg": "cos",
        "semantic": "cos_cache",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$T"
      },
      {
        "name": "sin_cache",
        "arg": "sin",
        "semantic": "sin_cache",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$T"
      },
      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$T" },
      {
        "name": "params",
        "semantic": "kernel.params",
        "buffer": { "type": "uniform" },
        "struct": {
          "name": "Params",
          "fields": [
            { "name": "pairCount", "type": "u32", "value": "numel(shapes.X) / 2" },
            { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.X, 1)" },
            { "name": "numHeads", "type": "u32", "value": "attrs.num_heads" },
            { "name": "headSize", "type": "u32", "value": "rank3HeadSize" },
            { "name": "rotaryDim", "type": "u32", "value": "dim(shapes.cos_cache, 2) * 2" },
            { "name": "halfRotaryDim", "type": "u32", "value": "dim(shapes.cos_cache, 2)" },
            {
              "name": "cacheBatchStride",
              "type": "u32",
              "value": "dim(shapes.cos_cache, 1) * dim(shapes.cos_cache, 2)"
            }
          ]
        }
      }
    ],
    "rank3Pos": [
      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
      {
        "name": "cos_cache",
        "arg": "cos",
        "semantic": "cos_cache",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$T"
      },
      {
        "name": "sin_cache",
        "arg": "sin",
        "semantic": "sin_cache",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$T"
      },
      {
        "name": "position_ids",
        "arg": "positionIds",
        "semantic": "position_ids",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$M"
      },
      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$T" },
      {
        "name": "params",
        "semantic": "kernel.params",
        "buffer": { "type": "uniform" },
        "struct": {
          "name": "Params",
          "fields": [
            { "name": "pairCount", "type": "u32", "value": "numel(shapes.X) / 2" },
            { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.X, 1)" },
            { "name": "numHeads", "type": "u32", "value": "attrs.num_heads" },
            { "name": "headSize", "type": "u32", "value": "rank3HeadSize" },
            { "name": "rotaryDim", "type": "u32", "value": "dim(shapes.cos_cache, 1) * 2" },
            { "name": "halfRotaryDim", "type": "u32", "value": "dim(shapes.cos_cache, 1)" }
          ]
        }
      }
    ],
    "rank4Pos": [
      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
      {
        "name": "cos_cache",
        "arg": "cos",
        "semantic": "cos_cache",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$T"
      },
      {
        "name": "sin_cache",
        "arg": "sin",
        "semantic": "sin_cache",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$T"
      },
      {
        "name": "position_ids",
        "arg": "positionIds",
        "semantic": "position_ids",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$M"
      },
      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$T" },
      {
        "name": "params",
        "semantic": "kernel.params",
        "buffer": { "type": "uniform" },
        "struct": {
          "name": "Params",
          "fields": [
            { "name": "pairCount", "type": "u32", "value": "numel(shapes.X) / 2" },
            { "name": "numHeads", "type": "u32", "value": "dim(shapes.X, 1)" },
            { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.X, 2)" },
            { "name": "headSize", "type": "u32", "value": "dim(shapes.X, 3)" },
            { "name": "rotaryDim", "type": "u32", "value": "dim(shapes.cos_cache, 1) * 2" },
            { "name": "halfRotaryDim", "type": "u32", "value": "dim(shapes.cos_cache, 1)" }
          ]
        }
      }
    ],
    "rank4": [
      { "name": "x", "arg": "x", "semantic": "X", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
      {
        "name": "cos_cache",
        "arg": "cos",
        "semantic": "cos_cache",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$T"
      },
      {
        "name": "sin_cache",
        "arg": "sin",
        "semantic": "sin_cache",
        "buffer": { "type": "read-only-storage" },
        "elementType": "$T"
      },
      { "name": "y", "arg": "y", "semantic": "Y", "buffer": { "type": "storage" }, "elementType": "$T" },
      {
        "name": "params",
        "semantic": "kernel.params",
        "buffer": { "type": "uniform" },
        "struct": {
          "name": "Params",
          "fields": [
            { "name": "pairCount", "type": "u32", "value": "numel(shapes.X) / 2" },
            { "name": "numHeads", "type": "u32", "value": "dim(shapes.X, 1)" },
            { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.X, 2)" },
            { "name": "headSize", "type": "u32", "value": "dim(shapes.X, 3)" },
            { "name": "rotaryDim", "type": "u32", "value": "dim(shapes.cos_cache, 2) * 2" },
            { "name": "halfRotaryDim", "type": "u32", "value": "dim(shapes.cos_cache, 2)" },
            {
              "name": "cacheBatchStride",
              "type": "u32",
              "value": "dim(shapes.cos_cache, 1) * dim(shapes.cos_cache, 2)"
            }
          ]
        }
      }
    ]
  },
  "variants": [
    {
      "id": "rank3_cache2_pos",
      "when": ["rank3PosContract", "pairDispatchOk"],
      "constants": { "interleaved": "attrs.interleaved != 0", "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
      "passes": [
        {
          "id": "main",
          "name": "rotary_embedding3d_pos",
          "source": { "shader": "rotary-embedding.wgsl.jinja", "inputs": { "rank": 3, "hasPositionIds": "true" } },
          "bindings": "rank3Pos",
          "dispatch": { "threads": "numel(shapes.X) / 2", "workgroupSize": "tunables.WORKGROUP_SIZE" }
        }
      ]
    },
    {
      "id": "rank4_cache2_pos",
      "when": ["rank4PosContract", "pairDispatchOk"],
      "constants": { "interleaved": "attrs.interleaved != 0", "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
      "passes": [
        {
          "id": "main",
          "name": "rotary_embedding4d_pos",
          "source": { "shader": "rotary-embedding.wgsl.jinja", "inputs": { "rank": 4, "hasPositionIds": "true" } },
          "bindings": "rank4Pos",
          "dispatch": { "threads": "numel(shapes.X) / 2", "workgroupSize": "tunables.WORKGROUP_SIZE" }
        }
      ]
    },
    {
      "id": "rank3_cache3",
      "when": ["rank3Contract", "pairDispatchOk"],
      "constants": { "interleaved": "attrs.interleaved != 0", "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
      "passes": [
        {
          "id": "main",
          "name": "rotary_embedding3d",
          "source": { "shader": "rotary-embedding.wgsl.jinja", "inputs": { "rank": 3 } },
          "bindings": "rank3",
          "dispatch": { "threads": "numel(shapes.X) / 2", "workgroupSize": "tunables.WORKGROUP_SIZE" }
        }
      ]
    },
    {
      "id": "rank4_cache3",
      "when": ["rank4Contract", "pairDispatchOk"],
      "constants": { "interleaved": "attrs.interleaved != 0", "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
      "passes": [
        {
          "id": "main",
          "name": "rotary_embedding4d",
          "source": { "shader": "rotary-embedding.wgsl.jinja", "inputs": { "rank": 4 } },
          "bindings": "rank4",
          "dispatch": { "threads": "numel(shapes.X) / 2", "workgroupSize": "tunables.WORKGROUP_SIZE" }
        }
      ]
    }
  ]
}