sync 2e7068faf55e

Browse files

Files changed (6) hide show

README.md +75 -0
build/webgpu/bench.json +247 -0
build/webgpu/manifest.json +290 -0
build/webgpu/metadata.json +18 -0
build/webgpu/rotary-embedding.wgsl.jinja +82 -0
build/webgpu/test.json +830 -0

README.md CHANGED Viewed

@@ -1,3 +1,78 @@
 ---
 license: apache-2.0
 ---

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# ai.onnx.RotaryEmbedding
+`ai.onnx`  ·  standard ONNX operator  ·  ONNX opset ≥ 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.
+See the [ONNX `RotaryEmbedding` spec](https://onnx.ai/onnx/operators/onnx__RotaryEmbedding.html) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- | --- |
+| `X` | `x` | `T` | same as logical dtype | — | — | 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. | required |
+| `cos_cache` | `cos` | `T` | same as logical dtype | — | — | 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)`. | required |
+| `sin_cache` | `sin` | `T` | same as logical dtype | — | — | Precomputed sine values with the same shape and type as `cos_cache`. | required |
+| `position_ids` | `positionIds` | `M` | `uint32` | `2` | — | 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. | optional |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `Y` | `y` | `T` | same as `X` | same as `X` | Rotary-position-encoded tensor with the same shape and type as `X`. | required |
+## Attributes
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `interleaved` | `0` | 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. |
+| `rotary_embedding_dim` | `0` | 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. |
+| `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. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+| `M` | `int64` |
+## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
+- [`rotary-embedding.wgsl.jinja`](build/webgpu/rotary-embedding.wgsl.jinja)
+## Use with `@huggingface/kernels`
+The loader derives every required output's shape and logical dtype from the manifest contract and this call.
+It then allocates the result tensors automatically.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/ai.onnx.RotaryEmbedding", { version: 1 });
+const { y } = await kernel({
+  x: { data: xData, shape: [1, 2, 1, 4] },
+  cos: { data: cosData, shape: [16, 2] },
+  sin: { data: sinData, shape: [16, 2] },
+  positionIds: { data: positionIdsData, shape: [1, 1] },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,247 @@

+{
+  "op": "ai.onnx.RotaryEmbedding",
+  "tunableSpace": { "WORKGROUP_SIZE": [128, 256] },
+  "cases": [
+    {
+      "name": "rotary-f32-b4s128h64",
+      "attrs": { "num_heads": 1 },
+      "preset": "smoke",
+      "vars": { "batch": 4, "seq": 128, "headSize": 64 },
+      "inputs": {
+        "x": { "shape": [4, 128, 64], "dtype": "float32", "dist": "normal", "seed": 115, "scale": 0.2 },
+        "cos": {
+          "shape": [4, 128, 32],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 116,
+          "scale": 0.1,
+          "offset": 0.9
+        },
+        "sin": { "shape": [4, 128, 32], "dtype": "float32", "dist": "uniform", "seed": 117, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [4, 128, 64], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.batch * args.seq * args.headSize * 4 * 2 + args.batch * args.seq * (args.headSize / 2) * 4 * 2"
+          }
+        ]
+      }
+    },
+    {
+      "name": "rotary-f32-b8s8192h128-large",
+      "attrs": { "num_heads": 1 },
+      "vars": { "batch": 8, "seq": 8192, "headSize": 128 },
+      "inputs": {
+        "x": { "shape": [8, 8192, 128], "dtype": "float32", "dist": "normal", "seed": 115, "scale": 0.2 },
+        "cos": {
+          "shape": [8, 8192, 64],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 116,
+          "scale": 0.1,
+          "offset": 0.9
+        },
+        "sin": { "shape": [8, 8192, 64], "dtype": "float32", "dist": "uniform", "seed": 117, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [8, 8192, 128], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.batch * args.seq * args.headSize * 4 * 2 + args.batch * args.seq * (args.headSize / 2) * 4 * 2"
+          }
+        ]
+      }
+    },
+    {
+      "name": "rotary-f32-rank4-b8h32s256h128-large",
+      "vars": { "batch": 8, "heads": 32, "seq": 256, "headSize": 128 },
+      "inputs": {
+        "x": { "shape": [8, 32, 256, 128], "dtype": "float32", "dist": "normal", "seed": 115, "scale": 0.2 },
+        "cos": {
+          "shape": [8, 256, 64],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 116,
+          "scale": 0.1,
+          "offset": 0.9
+        },
+        "sin": { "shape": [8, 256, 64], "dtype": "float32", "dist": "uniform", "seed": 117, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [8, 32, 256, 128], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.batch * args.heads * args.seq * args.headSize * 4 * 2 + args.batch * args.seq * (args.headSize / 2) * 4 * 2"
+          }
+        ]
+      }
+    },
+    {
+      "name": "rotary-f32-decode-q1-rank4-b1h32s1h128-single-row",
+      "preset": "smoke",
+      "vars": { "batch": 1, "heads": 32, "seq": 1, "headSize": 128 },
+      "inputs": {
+        "x": { "shape": [1, 32, 1, 128], "dtype": "float32", "dist": "normal", "seed": 115, "scale": 0.2 },
+        "cos": { "shape": [1, 1, 64], "dtype": "float32", "dist": "uniform", "seed": 116, "scale": 0.1, "offset": 0.9 },
+        "sin": { "shape": [1, 1, 64], "dtype": "float32", "dist": "uniform", "seed": 117, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [1, 32, 1, 128], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.batch * args.heads * args.seq * args.headSize * 4 * 2 + args.batch * args.seq * (args.headSize / 2) * 4 * 2"
+          }
+        ]
+      }
+    },
+    {
+      "name": "rotary-f32-partial-rotary-tail-b8s8192h128-read-amplification",
+      "preset": "smoke",
+      "vars": { "batch": 8, "seq": 8192, "headSize": 128, "rotaryDim": 64 },
+      "attrs": { "rotary_embedding_dim": 64, "num_heads": 1 },
+      "inputs": {
+        "x": { "shape": [8, 8192, 128], "dtype": "float32", "dist": "normal", "seed": 115, "scale": 0.2 },
+        "cos": {
+          "shape": [8, 8192, 32],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 116,
+          "scale": 0.1,
+          "offset": 0.9
+        },
+        "sin": { "shape": [8, 8192, 32], "dtype": "float32", "dist": "uniform", "seed": 117, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [8, 8192, 128], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.batch * args.seq * args.headSize * 4 * 2 + args.batch * args.seq * (args.rotaryDim / 2) * 4 * 2"
+          }
+        ]
+      }
+    },
+    {
+      "name": "rotary-f16-b8s8192h128-large",
+      "attrs": { "num_heads": 1 },
+      "preset": "smoke",
+      "vars": { "batch": 8, "seq": 8192, "headSize": 128 },
+      "inputs": {
+        "x": { "shape": [8, 8192, 128], "dtype": "float16", "dist": "normal", "seed": 115, "scale": 0.2 },
+        "cos": {
+          "shape": [8, 8192, 64],
+          "dtype": "float16",
+          "dist": "uniform",
+          "seed": 116,
+          "scale": 0.1,
+          "offset": 0.9
+        },
+        "sin": { "shape": [8, 8192, 64], "dtype": "float16", "dist": "uniform", "seed": 117, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [8, 8192, 128], "dtype": "float16" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.batch * args.seq * args.headSize * 2 * 2 + args.batch * args.seq * (args.headSize / 2) * 2 * 2"
+          }
+        ]
+      }
+    },
+    {
+      "name": "rotary-f32-multihead-b4s1024h2048-nh32-packed-layout",
+      "attrs": { "num_heads": 32 },
+      "preset": "smoke",
+      "vars": { "batch": 4, "seq": 1024, "hidden": 2048, "headSize": 64 },
+      "inputs": {
+        "x": { "shape": [4, 1024, 2048], "dtype": "float32", "dist": "normal", "seed": 115, "scale": 0.2 },
+        "cos": {
+          "shape": [4, 1024, 32],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 116,
+          "scale": 0.1,
+          "offset": 0.9
+        },
+        "sin": { "shape": [4, 1024, 32], "dtype": "float32", "dist": "uniform", "seed": 117, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [4, 1024, 2048], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.batch * args.seq * args.hidden * 4 * 2 + args.batch * args.seq * (args.headSize / 2) * 4 * 2"
+          }
+        ]
+      }
+    },
+    {
+      "name": "rotary-f32-rank3-partial-rd126-tail-copy-full-read-amplification",
+      "preset": "stress",
+      "vars": { "batch": 8, "seq": 8192, "headSize": 128, "rotaryDim": 126 },
+      "attrs": { "rotary_embedding_dim": 126, "num_heads": 1 },
+      "inputs": {
+        "x": { "shape": [8, 8192, 128], "dtype": "float32", "dist": "normal", "seed": 215, "scale": 0.2 },
+        "cos": {
+          "shape": [8, 8192, 63],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 216,
+          "scale": 0.1,
+          "offset": 0.9
+        },
+        "sin": { "shape": [8, 8192, 63], "dtype": "float32", "dist": "uniform", "seed": 217, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [8, 8192, 128], "dtype": "float32", "dist": "empty" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.batch * args.seq * args.headSize * 4 * 2 + args.batch * args.seq * (args.rotaryDim / 2) * 4 * 2"
+          }
+        ]
+      }
+    },
+    {
+      "name": "rotary-f32-rank4-partial-rd126-tail-copy-b4h16s512h128",
+      "preset": "stress",
+      "vars": { "batch": 4, "heads": 16, "seq": 512, "headSize": 128, "rotaryDim": 126 },
+      "attrs": { "rotary_embedding_dim": 126 },
+      "inputs": {
+        "x": { "shape": [4, 16, 512, 128], "dtype": "float32", "dist": "normal", "seed": 315, "scale": 0.2 },
+        "cos": {
+          "shape": [4, 512, 63],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 316,
+          "scale": 0.1,
+          "offset": 0.9
+        },
+        "sin": { "shape": [4, 512, 63], "dtype": "float32", "dist": "uniform", "seed": 317, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [4, 16, 512, 128], "dtype": "float32", "dist": "empty" } },
+      "bench": {
+        "primary": true,
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "value": "args.batch * args.heads * args.seq * args.headSize * 4 * 2 + args.batch * args.seq * (args.rotaryDim / 2) * 4 * 2"
+          }
+        ]
+      }
+    }
+  ]
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,290 @@

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

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,18 @@

+{
+  "name": "ai.onnx.RotaryEmbedding",
+  "id": "_ai_onnx_rotaryembedding_webgpu_23c1d92",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "Uu/zPgVC3ubArN3FjvwWkoSP9bbEqLHRRppel4wsIeA=",
+      "manifest.json": "yfsgfGknbiGBH5nzsmJmqJWPq1CskrxPgTz0vc+Qe/A=",
+      "rotary-embedding.wgsl.jinja": "dkJWgoDB2IWxw7espG+q4D1MX04J2ofR9xHvntuxg9w=",
+      "test.json": "p4hC6UXkzEdV+pQlTtZ0t2LcfXq3psczZWXmN0fRGes="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.RotaryEmbedding" }
+}

build/webgpu/rotary-embedding.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,82 @@

+{% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
+{% if note == "dispatch-limit" %}
+  // 2D-folded flat index: gid.y carries the high bits past the
+  // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
+{% elif note == "limit" %}
+  // 2D-folded flat index: gid.y carries the high bits past the
+  // maxComputeWorkgroupsPerDimension limit.
+{% elif note == "device-axis" %}
+  // The flat dispatch is folded across x/y at the device's per-axis workgroup
+  // limit; gid.y carries the high portion of the output index.
+{% elif note == "vec4-limit" %}
+  // 2D-folded flat vec4 index: gid.y carries the high bits past the
+  // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills into y).
+{% elif note == "element-limit" %}
+  // 2D-folded flat element index: gid.y carries the high bits past the
+  // maxComputeWorkgroupsPerDimension limit.
+{% elif note == "dispatch" %}
+  // 2D-folded flat index: gid.y carries the high bits past the
+  // maxComputeWorkgroupsPerDimension dispatch limit.
+{% endif %}
+{% if bound == "" %}
+  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+{%- elif guardInline %}
+  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+  if ({{ name }} >= {{ bound }}) { return; }
+{%- else %}
+  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }
+{%- endif %}
+{% endmacro %}
+{% if usesF16 %}enable f16;
+{% endif %}{{ env.wgsl.resourceDeclarations }}
+// One invocation owns both outputs of a rotated pair, so x/cos/sin are read once.
+// Partial rotations use the remaining pair lanes to copy two unchanged tail values.
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
+{{ flat_index_2d("p", "params.pairCount", note="") }}
+  let tasksPerHead = params.headSize / 2u;
+  let task = p % tasksPerHead;
+  let headFlat = p / tasksPerHead;
+  let base = headFlat * params.headSize;
+  if (task >= params.halfRotaryDim) {
+    let tail = params.rotaryDim + (task - params.halfRotaryDim) * 2u;
+    y[base + tail] = x[base + tail];
+    y[base + tail + 1u] = x[base + tail + 1u];
+    return;
+  }
+  let pair = task;
+{% if source.rank == 3 %}
+  let token = (headFlat / params.numHeads) % params.sequenceLength;
+  let batch = headFlat / (params.numHeads * params.sequenceLength);
+{% else %}
+  let token = headFlat % params.sequenceLength;
+  let batch = (headFlat / params.sequenceLength) / params.numHeads;
+{% endif %}
+{% if source.hasPositionIds is defined and source.hasPositionIds %}
+  // Gather the 2D cache tables through this token's position index.
+  let cache = u32(position_ids[batch * params.sequenceLength + token]) * params.halfRotaryDim + pair;
+{% else %}
+  // The standard no-index cache owns one row per (batch, token).
+  let cache = batch * params.cacheBatchStride + token * params.halfRotaryDim + pair;
+{% endif %}
+  let cf = f32(cos_cache[cache]);
+  let sf = f32(sin_cache[cache]);
+{% if interleaved %}
+  let aOffset = pair * 2u;
+  let bOffset = aOffset + 1u;
+{% else %}
+  let aOffset = pair;
+  let bOffset = pair + params.halfRotaryDim;
+{% endif %}
+  let a = f32(x[base + aOffset]);
+  let b = f32(x[base + bOffset]);
+  y[base + aOffset] = {{ scalar }}(a * cf - b * sf);
+  y[base + bOffset] = {{ scalar }}(a * sf + b * cf);
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,830 @@

+{
+  "op": "ai.onnx.RotaryEmbedding",
+  "fixtureArrays": {
+    "position_id_cos_table": [1, 0.994, 0.9553, 0.9171, 0.8253, 0.7584, 0.6216, 0.5319, 0.3624, 0.2579, 0.0707, -0.0392, -0.2272, -0.3327, -0.5048, -0.5966, -0.7374, -0.8071, -0.9041, -0.9455, -0.99, -0.9995, -0.9875, -0.9642, -0.8968, -0.8428, -0.7259, -0.646, -0.4903, -0.3916, -0.2108, -0.1022],
+    "position_id_sin_table": [0, 0.1692, 0.2955, 0.4529, 0.5646, 0.6961, 0.7833, 0.8772, 0.932, 0.9799, 0.9975, 0.9951, 0.9738, 0.9214, 0.8632, 0.7654, 0.6755, 0.541, 0.4274, 0.2683, 0.1411, -0.0284, -0.1577, -0.3225, -0.4425, -0.5879, -0.6878, -0.8007, -0.8716, -0.942, -0.9775, -0.9991],
+    "pos_ids_rank3_pair_gather_input_x": [0, 0.6442, 0.9854, 0.8632, 0.335, -0.3508, -0.8716, -0.9825, -0.6313, 0.0168, 0.657, 0.9882, 0.8546, 0.3191, -0.3665, -0.8797, -0.9792, -0.6181, 0.0336, 0.6696, 0.9906, 0.8457, 0.3031, -0.3821, -0.8876, -0.9756, -0.6048, 0.0504, 0.682, 0.9928, 0.8367, 0.2871, -0.3976, -0.8952, -0.9718, -0.5914, 0.0672, 0.6942, 0.9946, 0.8273, 0.2709, -0.4129, -0.9026, -0.9677, -0.5777, 0.084, 0.7062, 0.9962]
+  },
+  "cases": [
+    {
+      "name": "f32_identity_angle_preserves_subnormal_pair_gpu_gap",
+      "attrs": { "num_heads": 1 },
+      "skipGpu": {
+        "category": "permanent",
+        "reason": "Portable WGSL floating-point semantics do not guarantee preservation of the subnormal values required by this fixture. Backend evidence: Metal flushes denormals to zero in the ALU; this op's f32 multiply/FMA flushes the subnormal operand on GPU even at identity angle. Permanent FTZ limitation."
+      },
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc",
+        "test": "RotaryEmbedding",
+        "notes": "With cos=1 and sin=0, rotary embedding is an identity transform; positive and negative subnormal activations should survive."
+      },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 1, 2], "data": { "kind": "values", "values": [1e-40, -1e-40] } },
+        "cos": { "dtype": "float32", "shape": [1, 1, 1], "data": { "kind": "values", "values": [1.0] } },
+        "sin": { "dtype": "float32", "shape": [1, 1, 1], "data": { "kind": "values", "values": [0.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "f32_identity_angle_preserves_subnormal_rank4_pair_gpu_gap",
+      "skipGpu": {
+        "category": "permanent",
+        "reason": "Portable WGSL floating-point semantics do not guarantee preservation of the subnormal values required by this fixture. Backend evidence: Metal flushes denormals to zero in the ALU; this op's f32 multiply/FMA flushes the subnormal operand on GPU even at identity angle. Permanent FTZ limitation."
+      },
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc",
+        "test": "RotaryEmbedding",
+        "notes": "Rank-4 companion for identity-angle rotary embedding preserving subnormal activations."
+      },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 1, 1, 2], "data": { "kind": "values", "values": [1e-40, -1e-40] } },
+        "cos": { "dtype": "float32", "shape": [1, 1, 1], "data": { "kind": "values", "values": [1.0] } },
+        "sin": { "dtype": "float32", "shape": [1, 1, 1], "data": { "kind": "values", "values": [0.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 1, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "dispatch_cliff_rank3_partial_rotary",
+      "attrs": { "rotary_embedding_dim": 2, "num_heads": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [4194241, 1, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29 }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [4194241, 1, 1],
+          "data": { "kind": "rotaryCos", "thetaStart": 0.37, "thetaStep": 0.2 }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [4194241, 1, 1],
+          "data": { "kind": "rotarySin", "thetaStart": 0.37, "thetaStep": 0.2 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [4194241, 1, 4], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "batch2_seq3_head4",
+      "attrs": { "num_heads": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [2, 3, 2],
+          "data": { "kind": "rotaryCos", "thetaStart": 0.1, "thetaStep": 0.2 }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [2, 3, 2],
+          "data": { "kind": "rotarySin", "thetaStart": 0.1, "thetaStep": 0.2 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 4], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "seq1_head2_identity_angle_zero",
+      "attrs": { "num_heads": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 1, 2], "data": { "kind": "values", "values": [3.0, -4.0] } },
+        "cos": { "dtype": "float32", "shape": [1, 1, 1], "data": { "kind": "values", "values": [1.0] } },
+        "sin": { "dtype": "float32", "shape": [1, 1, 1], "data": { "kind": "values", "values": [0.0] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 2], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "batch1_seq2_head8_distinct_pairs",
+      "attrs": { "num_heads": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 8],
+          "data": {
+            "kind": "values",
+            "values": [1.0, 2.0, 3.0, 4.0, -1.0, -2.0, -3.0, -4.0, 5.0, 6.0, 7.0, 8.0, -5.0, -6.0, -7.0, -8.0]
+          }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [1, 2, 4],
+          "data": { "kind": "values", "values": [1.0, 0.0, 0.5, -0.5, 0.0, 1.0, -0.5, 0.5] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [1, 2, 4],
+          "data": { "kind": "values", "values": [0.0, 1.0, 0.5, 0.5, 1.0, 0.0, 0.5, -0.5] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 8], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "backend_style_batch2_seq3_head8",
+      "attrs": { "num_heads": 1 },
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/case/node/rotaryembedding.py",
+        "notes": "Standard rank-3, no-position-ids case with per-batch 3D cosine and sine caches."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 8],
+          "data": { "kind": "fillFloat32", "scale": 0.6, "sinStep": 0.17, "cosStep": 0.31 }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "rotaryCos", "thetaStart": 0.07, "thetaStep": 0.11 }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "rotarySin", "thetaStart": 0.07, "thetaStep": 0.11 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 8], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "ort_not_interleaved_no_position_ids_multi_head",
+      "attrs": { "num_heads": 3 },
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc",
+        "test": "RotaryEmbeddingTest.RotaryEmbedding_NotInterleaved_NoPosIds_SmallData_LlamaMSFT",
+        "notes": "Preserves ORT's standard batch-1 rank-3 cos/sin cache shape."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 18],
+          "data": {
+            "kind": "values",
+            "values": [-1.0408, 0.9166, -1.3042, -1.1097, -1.2188, 1.1676, 1.0076, -0.7529, -0.225, -0.4327, -1.5071, -0.4586, -0.8663, -0.2656, 0.1665, 0.7911, -0.932, -0.8579, -1.0574, -0.1188, -0.9078, 0.3452, -0.5713, -0.2351, -0.848, 0.5266, -1.2944, -0.0243, -0.2354, -0.7087, -0.9647, -0.0991, -0.2994, -0.065, -1.572, -1.3211]
+          }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [1, 2, 3],
+          "data": { "kind": "values", "values": [1.0, 1.0, 1.0, 0.5403, 0.9989, 1.0] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [1, 2, 3],
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.0, 0.8415, 0.0464, 0.0022] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 18], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "non_unit_cache_values_are_used_directly",
+      "attrs": { "num_heads": 1 },
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc",
+        "notes": "The operator consumes supplied cos/sin cache values directly; they need not form unit rotations."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 3, 4],
+          "data": {
+            "kind": "values",
+            "values": [2.0, -3.0, 5.0, -7.0, -11.0, 13.0, -17.0, 19.0, 23.0, -29.0, 31.0, -37.0]
+          }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [1, 3, 2],
+          "data": { "kind": "values", "values": [1.25, -0.5, 0.75, 0.125, -1.5, 2.0] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [1, 3, 2],
+          "data": { "kind": "values", "values": [-0.25, 0.5, 1.5, -0.75, 0.25, -1.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 3, 4], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "head12_non_power_of_two",
+      "attrs": { "num_heads": 1 },
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc",
+        "notes": "Covers an even head size that is not a power of two."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 12],
+          "data": { "kind": "fillFloat32", "scale": 0.4, "sinStep": 0.23, "cosStep": 0.19 }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [1, 2, 6],
+          "data": { "kind": "rotaryCos", "thetaStart": 0.03, "thetaStep": 0.17 }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [1, 2, 6],
+          "data": { "kind": "rotarySin", "thetaStart": 0.03, "thetaStep": 0.17 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 12], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "ort_interleaved_no_position_ids_multi_head",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc",
+        "test": "RotaryEmbeddingTest.RotaryEmbedding_Interleaved_NoPosIds_SmallData_LlamaMSFT"
+      },
+      "attrs": { "interleaved": 1, "num_heads": 2 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 3, 8],
+          "data": {
+            "kind": "values",
+            "values": [-1.0408, 0.9166, -1.3042, -1.1097, -0.132, -0.2751, -0.235, 0.0937, -1.2188, 1.1676, -1.0574, -0.1188, -0.7396, -1.2425, -0.1752, 0.699, -0.811, 0.6737, -1.1233, -0.0919, -0.6861, 0.7202, 0.1963, 0.6142]
+          }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [1, 3, 2],
+          "data": { "kind": "values", "values": [1.0, 1.0, 0.5403, 0.9999, -0.4161, 0.9998] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [1, 3, 2],
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.8415, 0.01, 0.9093, 0.02] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 3, 8], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "ort_interleaved_rank4_small_data",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc",
+        "test": "RotaryEmbeddingTest.RotaryEmbedding_Interleaved_SmallData_LlamaMSFT_4D_Input",
+        "notes": "ORT uses sequential position_ids [0, 1, 2], represented directly by the rank-4 cache rows."
+      },
+      "attrs": { "interleaved": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 3, 4],
+          "data": {
+            "kind": "values",
+            "values": [-1.0408, 0.9166, -1.3042, -1.1097, -1.2188, 1.1676, -1.0574, -0.1188, -0.811, 0.6737, -1.1233, -0.0919, -0.132, -0.2751, -0.235, 0.0937, -0.7396, -1.2425, -0.1752, 0.699, -0.6861, 0.7202, 0.1963, 0.6142]
+          }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [1, 3, 2],
+          "data": { "kind": "values", "values": [1.0, 1.0, 0.5403, 0.9999, -0.4161, 0.9998] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [1, 3, 2],
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.8415, 0.01, 0.9093, 0.02] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [1, 2, 3, 4],
+          "tolerance": 0.0001,
+          "data": {
+            "kind": "values",
+            "values": [-1.0408, 0.9166, -1.3042, -1.1097, -1.6411, -0.3948, -1.0561, -0.1294, -0.2751, -1.0178, -1.1212, -0.1143, -0.132, -0.2751, -0.235, 0.0937, 0.646, -1.2937, -0.1822, 0.6972, -0.3694, -0.9235, 0.184, 0.618]
+          }
+        }
+      }
+    },
+    {
+      "name": "ort_partial_rotary_dim_preserves_tail",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc",
+        "test": "RotaryEmbeddingTest.RotaryEmbedding_CustomRotaryDim_SmallData_Phi"
+      },
+      "attrs": { "rotary_embedding_dim": 4, "num_heads": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 6],
+          "data": {
+            "kind": "values",
+            "values": [-1.0408, 0.9166, -1.3042, -1.1097, -1.2188, 1.1676, 1.0076, -0.7529, -0.225, -0.4327, -1.5071, -0.4586]
+          }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [1, 2, 2],
+          "data": { "kind": "values", "values": [1.0, 1.0, 1.0, 0.5403] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [1, 2, 2],
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.0, 0.8415] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 6], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "ort_partial_rotary_dim_preserves_tail_f16",
+      "provenance": {
+        "source": "onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc",
+        "test": "RotaryEmbeddingTest.RotaryEmbedding_CustomRotaryDim_SmallData_Phi",
+        "notes": "Same Phi-style partial rotary-dimension case as ORT, exercising the fp16 path and confirming tail values stay unrotated."
+      },
+      "attrs": { "rotary_embedding_dim": 4, "num_heads": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 2, 6],
+          "data": {
+            "kind": "values",
+            "values": [-1.0408, 0.9166, -1.3042, -1.1097, -1.2188, 1.1676, 1.0076, -0.7529, -0.225, -0.4327, -1.5071, -0.4586]
+          }
+        },
+        "cos": {
+          "dtype": "float16",
+          "shape": [1, 2, 2],
+          "data": { "kind": "values", "values": [1.0, 1.0, 1.0, 0.5403] }
+        },
+        "sin": {
+          "dtype": "float16",
+          "shape": [1, 2, 2],
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.0, 0.8415] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float16",
+          "shape": [1, 2, 6],
+          "tolerance": 0.002,
+          "data": {
+            "kind": "values",
+            "values": [-1.0408, 0.9166, -1.3042, -1.1097, -1.2188, 1.1676, 1.0076, -0.0427, -0.225, -0.8673, -1.5071, -0.4586]
+          }
+        }
+      }
+    },
+    {
+      "name": "ort_contrib_packed_batching_repeated_cache_rows",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/rotary_embedding_op_test.cc",
+        "test": "ContribOpRotaryEmbeddingTest.RotaryEmbedding_CustomRotaryDim_SmallData_Phi_Packed_Batching",
+        "notes": "Materializes the indexed ORT case's gathered rows [0, 0, 1] into the standard per-batch rank-3 cache used when position_ids is omitted."
+      },
+      "attrs": { "rotary_embedding_dim": 4, "num_heads": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 3, 6],
+          "data": {
+            "kind": "values",
+            "values": [-1.0408, 0.9166, -1.3042, -1.1097, -1.2188, 1.1676, -1.0408, 0.9166, -1.3042, -1.1097, -1.2188, 1.1676, 1.0076, -0.7529, -0.225, -0.4327, -1.5071, -0.4586]
+          }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [1, 3, 2],
+          "data": { "kind": "values", "values": [1.0, 1.0, 1.0, 1.0, 1.0, 0.5403] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [1, 3, 2],
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.0, 0.0, 0.0, 0.8415] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [1, 3, 6],
+          "tolerance": 0.0001,
+          "data": {
+            "kind": "values",
+            "values": [-1.0408, 0.9166, -1.3042, -1.1097, -1.2188, 1.1676, -1.0408, 0.9166, -1.3042, -1.1097, -1.2188, 1.1676, 1.0076, -0.0427, -0.225, -0.8673, -1.5071, -0.4586]
+          }
+        }
+      }
+    },
+    {
+      "name": "onnx_backend_rank4_no_position_ids_batch_cache",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_rotary_embedding_no_position_ids",
+        "notes": "Compact deterministic data using the official no-position-id rank-4 cache layout."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 2, 3, 4],
+          "data": { "kind": "fillFloat32", "scale": 0.55, "sinStep": 0.13, "cosStep": 0.29 }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [2, 3, 2],
+          "data": { "kind": "values", "values": [1.0, 0.5, 0.25, -0.75, -1.5, 2.0, 0.0, 1.0, 0.75, 0.125, -0.5, -1.0] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [2, 3, 2],
+          "data": { "kind": "values", "values": [0.0, 1.0, -0.25, 0.5, 1.5, -0.75, -1.0, 0.0, 0.25, -0.5, 0.5, 0.25] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 2, 3, 4], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "onnx_backend_rank4_no_position_ids_interleaved",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_rotary_embedding_no_position_ids_interleaved",
+        "notes": "Compact deterministic data using the official no-position-id rank-4 interleaved layout."
+      },
+      "attrs": { "interleaved": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 2, 3, 4],
+          "data": { "kind": "fillFloat32", "scale": 0.55, "sinStep": 0.13, "cosStep": 0.29 }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [2, 3, 2],
+          "data": { "kind": "values", "values": [1.0, 0.5, 0.25, -0.75, -1.5, 2.0, 0.0, 1.0, 0.75, 0.125, -0.5, -1.0] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [2, 3, 2],
+          "data": { "kind": "values", "values": [0.0, 1.0, -0.25, 0.5, 1.5, -0.75, -1.0, 0.0, 0.25, -0.5, 0.5, 0.25] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 2, 3, 4], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "onnx_backend_rank4_no_position_ids_rotary_dim_tail",
+      "provenance": {
+        "source": "cmake/external/onnx/onnx/backend/test/data/node/test_rotary_embedding_no_position_ids_rotary_dim",
+        "notes": "Compact deterministic data covering partial rotary dimension with unrotated tail values."
+      },
+      "attrs": { "rotary_embedding_dim": 4 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 2, 2, 6],
+          "data": { "kind": "fillFloat32", "scale": 0.45, "sinStep": 0.19, "cosStep": 0.11 }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [2, 2, 2],
+          "data": { "kind": "values", "values": [1.0, 0.5, 0.25, -0.75, -1.5, 2.0, 0.0, 1.0] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [2, 2, 2],
+          "data": { "kind": "values", "values": [0.0, 1.0, -0.25, 0.5, 1.5, -0.75, -1.0, 0.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 2, 2, 6], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "empty_input_zero_dim",
+      "attrs": { "num_heads": 1 },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [2, 3, 0], "data": { "kind": "values", "values": [] } },
+        "cos": { "dtype": "float32", "shape": [2, 3, 0], "data": { "kind": "values", "values": [] } },
+        "sin": { "dtype": "float32", "shape": [2, 3, 0], "data": { "kind": "values", "values": [] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 0], "tolerance": 0 } }
+    },
+    {
+      "name": "interleaved_partial_rotary_dim_rank3",
+      "provenance": {
+        "notes": "Exercises the rank3_cache3 kernel's interleaved + partial-rotary branch: leading pairs rotate while the d>=rotary_embedding_dim tail is copied through unchanged."
+      },
+      "attrs": { "interleaved": 1, "rotary_embedding_dim": 4, "num_heads": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 8],
+          "data": {
+            "kind": "values",
+            "values": [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, -1.0, -2.0, -3.0, -4.0, -5.0, -6.0, -7.0, -8.0]
+          }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [1, 2, 2],
+          "data": { "kind": "values", "values": [1.0, 1.0, 0.5403, 0.9999] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [1, 2, 2],
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.8415, 0.01] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 8], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "decode_q1_rank4_head80_multihead_f16",
+      "provenance": {
+        "notes": "f16 autoregressive decode step (seq=1) with non-power-of-two head_dim 80 over 8 heads, the dominant inference path; exercises the rank4_cache3 kernel under the production fp16 dtype with the single-row decode shape."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 8, 1, 80],
+          "data": { "kind": "fillFloat32", "scale": 0.3, "sinStep": 0.13, "cosStep": 0.21 }
+        },
+        "cos": {
+          "dtype": "float16",
+          "shape": [1, 1, 40],
+          "data": { "kind": "rotaryCos", "thetaStart": 0.05, "thetaStep": 0.07 }
+        },
+        "sin": {
+          "dtype": "float16",
+          "shape": [1, 1, 40],
+          "data": { "kind": "rotarySin", "thetaStart": 0.05, "thetaStep": 0.07 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [1, 8, 1, 80], "tolerance": 0.02 } }
+    },
+    {
+      "name": "f16_multihead_packed_head96_not_interleaved",
+      "attrs": { "num_heads": 4 },
+      "provenance": {
+        "notes": "f16 rank-3 packed multi-head layout (hidden = num_heads * head_size) with non-power-of-two head_dim 96, exercising the rank3_cache3 kernel's per-head indexing under fp16."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 3, 384],
+          "data": { "kind": "fillFloat32", "scale": 0.25, "sinStep": 0.11, "cosStep": 0.19 }
+        },
+        "cos": {
+          "dtype": "float16",
+          "shape": [1, 3, 48],
+          "data": { "kind": "rotaryCos", "thetaStart": 0.04, "thetaStep": 0.09 }
+        },
+        "sin": {
+          "dtype": "float16",
+          "shape": [1, 3, 48],
+          "data": { "kind": "rotarySin", "thetaStart": 0.04, "thetaStep": 0.09 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [1, 3, 384], "tolerance": 0.02 } }
+    },
+    {
+      "name": "interleaved_partial_rotary_rank4_multibatch_cache_stride",
+      "attrs": { "interleaved": 1, "rotary_embedding_dim": 4 },
+      "provenance": {
+        "notes": "Exercises the rank-4 interleaved + partial-rotary path with a per-batch rank-3 cache and batch>1, verifying the leading rotated pairs, unchanged tail, and correct per-batch cache-row selection."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 2, 2, 6],
+          "data": { "kind": "fillFloat32", "scale": 0.45, "sinStep": 0.17, "cosStep": 0.23 }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [2, 2, 2],
+          "data": { "kind": "values", "values": [1.0, 0.5, 0.25, -0.75, -1.5, 2.0, 0.0, 1.0] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [2, 2, 2],
+          "data": { "kind": "values", "values": [0.0, 1.0, -0.25, 0.5, 1.5, -0.75, -1.0, 0.0] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 2, 2, 6], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "interleaved_partial_rotary_multihead_packed_rank3",
+      "attrs": { "interleaved": 1, "rotary_embedding_dim": 4, "num_heads": 4 },
+      "provenance": {
+        "notes": "Interleaved pairing + partial rotary_embedding_dim=4 + packed multi-head layout (hidden=64, num_heads=4 -> head_size=16). Verifies per-head rotation of the leading 4 lanes, unchanged tails, and correct per-head base indexing."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 3, 64],
+          "data": { "kind": "fillFloat32", "scale": 0.3, "sinStep": 0.11, "cosStep": 0.19 }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [1, 3, 2],
+          "data": { "kind": "values", "values": [1.0, 1.0, 0.5403, 0.9999, -0.4161, 0.9998] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [1, 3, 2],
+          "data": { "kind": "values", "values": [0.0, 0.0, 0.8415, 0.01, 0.9093, 0.02] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 3, 64], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "dispatch_fold_rank4_pair_last_row_seq65540",
+      "provenance": {
+        "notes": "Regression guard for the 2D dispatch fold: pairCount = numel(X)/2 = 16,781,440 -> ceil(/256)=65540 > 65535 forces the gid.y high-bit fold. Confirms the highest pair indices are written at the maxComputeWorkgroupsPerDimension cliff. Periodic inputs preserve coverage without evaluating 33M host transcendentals."
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 1, 65540, 512],
+          "data": { "kind": "cycle", "values": [0.5, -0.75, 1.25, -1.5] }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [1, 65540, 256],
+          "data": { "kind": "cycle", "values": [0.5, 0.75, -0.25, 1.0] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [1, 65540, 256],
+          "data": { "kind": "cycle", "values": [0.25, -0.5, 0.75, 0.125] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 65540, 512], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "pos_ids_rank3_pair_gather",
+      "attrs": { "num_heads": 2 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 8],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/pos_ids_rank3_pair_gather_input_x" } }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [16, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/position_id_cos_table" } }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [16, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/position_id_sin_table" } }
+        },
+        "positionIds": {
+          "dtype": "uint32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [4, 5, 6, 0, 2, 9] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 8], "tolerance": 0.0001 } },
+      "provenance": {
+        "source": "onnxruntime/core/providers/webgpu/llm/rotary_embedding.cc",
+        "test": "RotaryEmbedding position_ids gather (opset-23 input 3)",
+        "notes": "When position_ids is supplied the cos/sin caches are 2D (max_position, half) tables gathered per (batch, token); without it the caches are addressed by the token index directly."
+      }
+    },
+    {
+      "name": "pos_ids_rank3_interleaved_partial",
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 4, 12],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/pos_ids_rank3_pair_gather_input_x" } }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [16, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/position_id_cos_table" } }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [16, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/position_id_sin_table" } }
+        },
+        "positionIds": { "dtype": "uint32", "shape": [1, 4], "data": { "kind": "values", "values": [9, 0, 9, 2] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 4, 12], "tolerance": 0.0001 } },
+      "provenance": {
+        "source": "onnxruntime/core/providers/webgpu/llm/rotary_embedding.cc",
+        "test": "RotaryEmbedding position_ids gather (opset-23 input 3)",
+        "notes": "When position_ids is supplied the cos/sin caches are 2D (max_position, half) tables gathered per (batch, token); without it the caches are addressed by the token index directly."
+      },
+      "attrs": { "interleaved": 1, "rotary_embedding_dim": 4, "num_heads": 2 }
+    },
+    {
+      "name": "pos_ids_rank4_pair_gather",
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 3, 4],
+          "data": {
+            "kind": "values",
+            "values": [0.0, 0.6442, 0.9854, 0.8632, 0.335, -0.3508, -0.8716, -0.9825, -0.6313, 0.0168, 0.657, 0.9882, 0.8546, 0.3191, -0.3665, -0.8797, -0.9792, -0.6181, 0.0336, 0.6696, 0.9906, 0.8457, 0.3031, -0.3821]
+          }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [16, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/position_id_cos_table" } }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [16, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/position_id_sin_table" } }
+        },
+        "positionIds": { "dtype": "uint32", "shape": [1, 3], "data": { "kind": "values", "values": [7, 3, 11] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 3, 4], "tolerance": 0.0001 } },
+      "provenance": {
+        "source": "onnxruntime/core/providers/webgpu/llm/rotary_embedding.cc",
+        "test": "RotaryEmbedding position_ids gather (opset-23 input 3)",
+        "notes": "When position_ids is supplied the cos/sin caches are 2D (max_position, half) tables gathered per (batch, token); without it the caches are addressed by the token index directly."
+      }
+    },
+    {
+      "name": "pos_ids_rank3_decode_step",
+      "attrs": { "num_heads": 2 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 1, 8],
+          "data": { "kind": "values", "values": [0.0, 0.6442, 0.9854, 0.8632, 0.335, -0.3508, -0.8716, -0.9825] }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [16, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/position_id_cos_table" } }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [16, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/position_id_sin_table" } }
+        },
+        "positionIds": { "dtype": "uint32", "shape": [1, 1], "data": { "kind": "values", "values": [5] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 8], "tolerance": 0.0001 } },
+      "provenance": {
+        "source": "onnxruntime/core/providers/webgpu/llm/rotary_embedding.cc",
+        "test": "RotaryEmbedding position_ids gather (opset-23 input 3)",
+        "notes": "When position_ids is supplied the cos/sin caches are 2D (max_position, half) tables gathered per (batch, token); without it the caches are addressed by the token index directly."
+      }
+    },
+    {
+      "name": "position_ids_rank3_prefill_from_past_len",
+      "attrs": { "num_heads": 2 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.37, "scale": 0.7 }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [16, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/position_id_cos_table" } }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [16, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/position_id_sin_table" } }
+        },
+        "positionIds": {
+          "dtype": "uint32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [5, 6, 7, 5, 6, 7] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 8], "tolerance": 0.0001 } },
+      "provenance": {
+        "source": "onnxruntime/core/providers/webgpu/llm/rotary_embedding.cc",
+        "test": "RotaryEmbedding position_ids gather (opset-23 input 3)",
+        "notes": "The standard [batch, sequence] position table selects cache rows 5, 6, and 7 for both batches, covering a prefill that begins after an existing cache prefix."
+      }
+    },
+    {
+      "name": "position_ids_rank4_decode_step_past_len",
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 1, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.41, "scale": 0.7 }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [16, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/position_id_cos_table" } }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [16, 2],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/position_id_sin_table" } }
+        },
+        "positionIds": { "dtype": "uint32", "shape": [1, 1], "data": { "kind": "values", "values": [9] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 1, 4], "tolerance": 0.0001 } },
+      "provenance": {
+        "source": "onnxruntime/core/providers/webgpu/llm/rotary_embedding.cc",
+        "test": "RotaryEmbedding position_ids gather (opset-23 input 3)",
+        "notes": "A standard [1,1] position table selects cache row 9 for the single decode token; both heads use that row."
+      }
+    }
+  ]
+}