sync 2e7068faf55e

Browse files

Files changed (6) hide show

README.md +81 -0
build/webgpu/bench.json +229 -0
build/webgpu/manifest.json +251 -0
build/webgpu/metadata.json +18 -0
build/webgpu/mrotary-embedding.wgsl.jinja +108 -0
build/webgpu/test.json +246 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.MRotaryEmbedding
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Multimodal rotary position embedding (M-RoPE) for Qwen models. Each token has temporal, height, and width position streams; `mrope_section` partitions the half-rotary axis and `mrope_layout` assigns them. Text-only tokens set all streams equal, reducing the op to `RotaryEmbedding`. The effective rotary dimension must be positive and even; an odd head size is supported with a smaller even `rotary_embedding_dim`. This package supports float16/float32 and non-packed mode; bfloat16 and packed batching are not implemented. Position ids must be valid non-negative cache-row indices.
+See the [ONNX Runtime `MRotaryEmbedding` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.MRotaryEmbedding) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- | --- |
+| `input` | `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. The effective rotary dimension must be even, and `num_heads` is required for rank-3 input. | required |
+| `position_ids` | `positionIds` | `M` | `uint32` | `3` | — | Logical int64 position indices of shape `(3, batch_size, sequence_length)`, holding the temporal, height and width streams in that order along the first axis. Valid positions are non-negative cache-row indices and use uint32 WebGPU storage. | required |
+| `cos_cache` | `cos` | `T` | same as logical dtype | `2` | — | Precomputed cosine values of shape `(max_sequence_length, rotary_dim/2)`, shared by all three position streams. | required |
+| `sin_cache` | `sin` | `T` | same as logical dtype | `2` | — | Precomputed sine values with the same shape and type as `cos_cache`. | required |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `output` | `y` | `T` | same as `input` | same as `input` | Rotary-position-encoded tensor with the same shape and type as `input`. | 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. This is the rotation pairing and is independent of `mrope_layout`. |
+| `is_packed_batching` | `0` | Whether `position_ids` uses packed-batch metadata. The default and only supported value is 0; packed batching (1) is not implemented. |
+| `mrope_layout` | `0` | How the three sections are combined into one per-token cos/sin vector: `0` for the sectioned/chunked layout (Qwen2-VL, Qwen2.5-VL) or `1` for the interleaved layout (Qwen3-VL, Qwen3.5). Default is 0. |
+| `num_heads` | `0` | Number of attention heads. The schema default is 0. A positive value is required for rank-3 `input` and whenever `rotary_embedding_dim` is nonzero; rank-4 execution otherwise infers the head count from `input`. |
+| `rotary_embedding_dim` | `0` | Positive even number of head-dimension elements to rotate; `0` means the full head dimension, which must then be even. A smaller even value permits an odd head size and copies the remaining tail unchanged. |
+| `scale` | `1` | Scale applied to the gathered cosine and sine values before the rotation. Default is 1.0. |
+| `mrope_section` | — | Three non-negative integers `[section_t, section_h, section_w]` dividing the half-rotary axis among the temporal, height and width streams. They must sum to `rotary_embedding_dim / 2`, or to `head_size / 2` when `rotary_embedding_dim` is 0. Required. |
+## 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
+- [`mrotary-embedding.wgsl.jinja`](build/webgpu/mrotary-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/com.microsoft.MRotaryEmbedding", { version: 1 });
+const { y } = await kernel({
+  x: { data: xData, shape: [1, 4, 16] },
+  positionIds: { data: positionIdsData, shape: [3, 1, 4] },
+  cos: { data: cosData, shape: [8, 4] },
+  sin: { data: sinData, shape: [8, 4] },
+}, {
+  attrs: { num_heads: 2, mrope_section: [2, 1, 1] },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,229 @@

+{
+  "op": "com.microsoft.MRotaryEmbedding",
+  "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
+  "cases": [
+    {
+      "name": "mrope-f32-smoke-b1h4s64d64",
+      "preset": "smoke",
+      "attrs": {
+        "mrope_layout": 0,
+        "interleaved": 0,
+        "mrope_section": [16, 8, 8],
+        "rotary_embedding_dim": 0,
+        "scale": 1
+      },
+      "vars": { "elems": 16384, "rotPairs": 8192 },
+      "inputs": {
+        "x": { "shape": [1, 4, 64, 64], "dtype": "float32", "dist": "normal", "seed": 1000, "scale": 0.2 },
+        "positionIds": { "shape": [3, 1, 64], "dtype": "uint32", "dist": "uniform", "seed": 1001, "scale": 127 },
+        "cos": { "shape": [128, 32], "dtype": "float32", "dist": "uniform", "seed": 1002, "scale": 0.1, "offset": 0.9 },
+        "sin": { "shape": [128, 32], "dtype": "float32", "dist": "uniform", "seed": 1003, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [1, 4, 64, 64], "dtype": "float32" } },
+      "bench": {
+        "primary": false,
+        "metrics": [{ "type": "bandwidth", "value": "args.elems * 4 * 2 + args.rotPairs * 4 * 2" }]
+      }
+    },
+    {
+      "name": "qwen3-vl-2b-decode-q-h16-d128",
+      "preset": "model",
+      "attrs": {
+        "mrope_layout": 1,
+        "interleaved": 0,
+        "mrope_section": [24, 20, 20],
+        "rotary_embedding_dim": 0,
+        "scale": 1
+      },
+      "vars": { "elems": 2048, "rotPairs": 1024 },
+      "provenance": {
+        "model": "Qwen/Qwen3-VL-2B-Instruct",
+        "notes": "Qwen3-VL interleaved M-RoPE over the query projection: 16 heads x head_dim 128, mrope_section [24, 20, 20] summing to head_dim/2. Cache rows cover the model's position budget."
+      },
+      "inputs": {
+        "x": { "shape": [1, 16, 1, 128], "dtype": "float32", "dist": "normal", "seed": 1000, "scale": 0.2 },
+        "positionIds": { "shape": [3, 1, 1], "dtype": "uint32", "dist": "uniform", "seed": 1001, "scale": 32767 },
+        "cos": {
+          "shape": [32768, 64],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 1002,
+          "scale": 0.1,
+          "offset": 0.9
+        },
+        "sin": { "shape": [32768, 64], "dtype": "float32", "dist": "uniform", "seed": 1003, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [1, 16, 1, 128], "dtype": "float32" } },
+      "bench": {
+        "primary": true,
+        "metrics": [{ "type": "bandwidth", "value": "args.elems * 4 * 2 + args.rotPairs * 4 * 2" }]
+      }
+    },
+    {
+      "name": "qwen3-vl-2b-prefill-t512-q-h16-d128",
+      "preset": "model",
+      "attrs": {
+        "mrope_layout": 1,
+        "interleaved": 0,
+        "mrope_section": [24, 20, 20],
+        "rotary_embedding_dim": 0,
+        "scale": 1
+      },
+      "vars": { "elems": 1048576, "rotPairs": 524288 },
+      "provenance": {
+        "model": "Qwen/Qwen3-VL-2B-Instruct",
+        "notes": "Qwen3-VL interleaved M-RoPE over the query projection: 16 heads x head_dim 128, mrope_section [24, 20, 20] summing to head_dim/2. Cache rows cover the model's position budget."
+      },
+      "inputs": {
+        "x": { "shape": [1, 16, 512, 128], "dtype": "float32", "dist": "normal", "seed": 1000, "scale": 0.2 },
+        "positionIds": { "shape": [3, 1, 512], "dtype": "uint32", "dist": "uniform", "seed": 1001, "scale": 32767 },
+        "cos": {
+          "shape": [32768, 64],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 1002,
+          "scale": 0.1,
+          "offset": 0.9
+        },
+        "sin": { "shape": [32768, 64], "dtype": "float32", "dist": "uniform", "seed": 1003, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [1, 16, 512, 128], "dtype": "float32" } },
+      "bench": {
+        "primary": false,
+        "metrics": [{ "type": "bandwidth", "value": "args.elems * 4 * 2 + args.rotPairs * 4 * 2" }]
+      }
+    },
+    {
+      "name": "qwen3-vl-2b-decode-k-h8-d128",
+      "preset": "model",
+      "attrs": {
+        "mrope_layout": 1,
+        "interleaved": 0,
+        "mrope_section": [24, 20, 20],
+        "rotary_embedding_dim": 0,
+        "scale": 1
+      },
+      "vars": { "elems": 1024, "rotPairs": 512 },
+      "provenance": {
+        "model": "Qwen/Qwen3-VL-2B-Instruct",
+        "notes": "Same layer's key projection: 8 kv heads, so half the query work and the smallest dispatch M-RoPE sees in decode."
+      },
+      "inputs": {
+        "x": { "shape": [1, 8, 1, 128], "dtype": "float32", "dist": "normal", "seed": 1000, "scale": 0.2 },
+        "positionIds": { "shape": [3, 1, 1], "dtype": "uint32", "dist": "uniform", "seed": 1001, "scale": 32767 },
+        "cos": {
+          "shape": [32768, 64],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 1002,
+          "scale": 0.1,
+          "offset": 0.9
+        },
+        "sin": { "shape": [32768, 64], "dtype": "float32", "dist": "uniform", "seed": 1003, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [1, 8, 1, 128], "dtype": "float32" } },
+      "bench": {
+        "primary": false,
+        "metrics": [{ "type": "bandwidth", "value": "args.elems * 4 * 2 + args.rotPairs * 4 * 2" }]
+      }
+    },
+    {
+      "name": "qwen2-5-vl-3b-prefill-t512-q-h16-d128",
+      "preset": "model",
+      "attrs": {
+        "mrope_layout": 0,
+        "interleaved": 0,
+        "mrope_section": [16, 24, 24],
+        "rotary_embedding_dim": 0,
+        "scale": 1
+      },
+      "vars": { "elems": 1048576, "rotPairs": 524288 },
+      "provenance": {
+        "model": "Qwen/Qwen2.5-VL-3B-Instruct",
+        "notes": "Qwen2.5-VL sectioned M-RoPE, mrope_section [16, 24, 24]. The sectioned branch reads three contiguous chunks where the interleaved branch strides by 3, so the two layouts have different cache access patterns at the same shape."
+      },
+      "inputs": {
+        "x": { "shape": [1, 16, 512, 128], "dtype": "float32", "dist": "normal", "seed": 1000, "scale": 0.2 },
+        "positionIds": { "shape": [3, 1, 512], "dtype": "uint32", "dist": "uniform", "seed": 1001, "scale": 32767 },
+        "cos": {
+          "shape": [32768, 64],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 1002,
+          "scale": 0.1,
+          "offset": 0.9
+        },
+        "sin": { "shape": [32768, 64], "dtype": "float32", "dist": "uniform", "seed": 1003, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [1, 16, 512, 128], "dtype": "float32" } },
+      "bench": {
+        "primary": false,
+        "metrics": [{ "type": "bandwidth", "value": "args.elems * 4 * 2 + args.rotPairs * 4 * 2" }]
+      }
+    },
+    {
+      "name": "qwen3-vl-2b-prefill-t512-q-h16-d128-f16",
+      "preset": "model",
+      "attrs": {
+        "mrope_layout": 1,
+        "interleaved": 0,
+        "mrope_section": [24, 20, 20],
+        "rotary_embedding_dim": 0,
+        "scale": 1
+      },
+      "vars": { "elems": 1048576, "rotPairs": 524288 },
+      "provenance": {
+        "model": "Qwen/Qwen3-VL-2B-Instruct",
+        "notes": "f16 activations at the same Qwen3-VL prefill shape."
+      },
+      "inputs": {
+        "x": { "shape": [1, 16, 512, 128], "dtype": "float16", "dist": "normal", "seed": 1000, "scale": 0.2 },
+        "positionIds": { "shape": [3, 1, 512], "dtype": "uint32", "dist": "uniform", "seed": 1001, "scale": 32767 },
+        "cos": {
+          "shape": [32768, 64],
+          "dtype": "float16",
+          "dist": "uniform",
+          "seed": 1002,
+          "scale": 0.1,
+          "offset": 0.9
+        },
+        "sin": { "shape": [32768, 64], "dtype": "float16", "dist": "uniform", "seed": 1003, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [1, 16, 512, 128], "dtype": "float16" } },
+      "bench": {
+        "primary": false,
+        "metrics": [{ "type": "bandwidth", "value": "args.elems * 2 * 2 + args.rotPairs * 2 * 2" }]
+      }
+    },
+    {
+      "name": "mrope-f32-stress-b4h32s2048d128",
+      "preset": "stress",
+      "attrs": {
+        "mrope_layout": 1,
+        "interleaved": 0,
+        "mrope_section": [24, 20, 20],
+        "rotary_embedding_dim": 0,
+        "scale": 1
+      },
+      "vars": { "elems": 33554432, "rotPairs": 16777216 },
+      "inputs": {
+        "x": { "shape": [4, 32, 2048, 128], "dtype": "float32", "dist": "normal", "seed": 1000, "scale": 0.2 },
+        "positionIds": { "shape": [3, 4, 2048], "dtype": "uint32", "dist": "uniform", "seed": 1001, "scale": 32767 },
+        "cos": {
+          "shape": [32768, 64],
+          "dtype": "float32",
+          "dist": "uniform",
+          "seed": 1002,
+          "scale": 0.1,
+          "offset": 0.9
+        },
+        "sin": { "shape": [32768, 64], "dtype": "float32", "dist": "uniform", "seed": 1003, "scale": 0.1 }
+      },
+      "outputs": { "y": { "shape": [4, 32, 2048, 128], "dtype": "float32" } },
+      "bench": {
+        "primary": false,
+        "metrics": [{ "type": "bandwidth", "value": "args.elems * 4 * 2 + args.rotPairs * 4 * 2" }]
+      }
+    }
+  ]
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,251 @@

+{
+  "domain": "com.microsoft",
+  "name": "MRotaryEmbedding",
+  "sinceVersion": 1,
+  "description": "Multimodal rotary position embedding (M-RoPE) for Qwen models. Each token has temporal, height, and width position streams; `mrope_section` partitions the half-rotary axis and `mrope_layout` assigns them. Text-only tokens set all streams equal, reducing the op to `RotaryEmbedding`. The effective rotary dimension must be positive and even; an odd head size is supported with a smaller even `rotary_embedding_dim`. This package supports float16/float32 and non-packed mode; bfloat16 and packed batching are not implemented. Position ids must be valid non-negative cache-row indices.",
+  "inputs": [
+    {
+      "role": "input",
+      "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. The effective rotary dimension must be even, and `num_heads` is required for rank-3 input."
+    },
+    {
+      "role": "position_ids",
+      "dtype": "M",
+      "rank": 3,
+      "description": "Logical int64 position indices of shape `(3, batch_size, sequence_length)`, holding the temporal, height and width streams in that order along the first axis. Valid positions are non-negative cache-row indices and use uint32 WebGPU storage."
+    },
+    {
+      "role": "cos_cache",
+      "dtype": "T",
+      "rank": 2,
+      "description": "Precomputed cosine values of shape `(max_sequence_length, rotary_dim/2)`, shared by all three position streams."
+    },
+    {
+      "role": "sin_cache",
+      "dtype": "T",
+      "rank": 2,
+      "description": "Precomputed sine values with the same shape and type as `cos_cache`."
+    }
+  ],
+  "outputs": [
+    {
+      "role": "output",
+      "dtype": "T",
+      "rank": "ranks.input",
+      "shape": "shapes.input",
+      "description": "Rotary-position-encoded tensor with the same shape and type as `input`."
+    }
+  ],
+  "attributes": {
+    "interleaved": 0,
+    "is_packed_batching": 0,
+    "mrope_layout": 0,
+    "num_heads": 0,
+    "rotary_embedding_dim": 0,
+    "scale": 1
+  },
+  "attributeConstraints": {
+    "interleaved": { "values": [0, 1] },
+    "is_packed_batching": { "values": [0] },
+    "mrope_layout": { "values": [0, 1] },
+    "mrope_section": { "required": true }
+  },
+  "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. This is the rotation pairing and is independent of `mrope_layout`.",
+    "is_packed_batching": "Whether `position_ids` uses packed-batch metadata. The default and only supported value is 0; packed batching (1) is not implemented.",
+    "mrope_layout": "How the three sections are combined into one per-token cos/sin vector: `0` for the sectioned/chunked layout (Qwen2-VL, Qwen2.5-VL) or `1` for the interleaved layout (Qwen3-VL, Qwen3.5). Default is 0.",
+    "mrope_section": "Three non-negative integers `[section_t, section_h, section_w]` dividing the half-rotary axis among the temporal, height and width streams. They must sum to `rotary_embedding_dim / 2`, or to `head_size / 2` when `rotary_embedding_dim` is 0. Required.",
+    "num_heads": "Number of attention heads. The schema default is 0. A positive value is required for rank-3 `input` and whenever `rotary_embedding_dim` is nonzero; rank-4 execution otherwise infers the head count from `input`.",
+    "rotary_embedding_dim": "Positive even number of head-dimension elements to rotate; `0` means the full head dimension, which must then be even. A smaller even value permits an odd head size and copies the remaining tail unchanged.",
+    "scale": "Scale applied to the gathered cosine and sine values before the rotation. Default is 1.0."
+  },
+  "typeConstraints": { "T": ["float32", "float16"], "M": ["int64"] },
+  "args": {
+    "x": { "kind": "tensor", "semantic": "input", "role": "input" },
+    "positionIds": {
+      "kind": "tensor",
+      "semantic": "position_ids",
+      "role": "input",
+      "dtype": "uint32",
+      "narrowing": "checked"
+    },
+    "cos": { "kind": "tensor", "semantic": "cos_cache", "role": "input" },
+    "sin": { "kind": "tensor", "semantic": "sin_cache", "role": "input" },
+    "y": { "kind": "tensor", "semantic": "output", "role": "output" }
+  },
+  "tunables": { "WORKGROUP_SIZE": 256 },
+  "derive": {
+    "rank3HeadSize": "dim(shapes.input, 2) / attrs.num_heads if attrs.num_heads is defined and attrs.num_heads > 0 else 0",
+    "headSize": "rank3HeadSize if ranks.input == 3 else dim(shapes.input, 3)",
+    "effectiveRotaryDim": "attrs.rotary_embedding_dim if attrs.rotary_embedding_dim != 0 else headSize",
+    "tasksPerHead": "ceilDiv(headSize, 2)",
+    "pairCount": "(numel(shapes.input) / max(1, headSize)) * tasksPerHead",
+    "pairDispatchOk": "ceilDiv(pairCount, tunables.WORKGROUP_SIZE) <= device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension",
+    "halfRotaryDim": "dim(shapes.cos_cache, 1)",
+    "sectionsDefined": "attrs.mrope_section is defined and (attrs.mrope_section | length) == 3",
+    "sectionsValid": "sectionsDefined and attrs.mrope_section[0] >= 0 and attrs.mrope_section[1] >= 0 and attrs.mrope_section[2] >= 0 and attrs.mrope_section[0] + attrs.mrope_section[1] + attrs.mrope_section[2] == halfRotaryDim",
+    "commonContract": "f16Ok(dtypes.T) and sameShape(shapes.input, shapes.output) and (attrs.interleaved == 0 or attrs.interleaved == 1) and (attrs.mrope_layout == 0 or attrs.mrope_layout == 1) and attrs.num_heads >= 0 and attrs.rotary_embedding_dim >= 0 and (attrs.rotary_embedding_dim == 0 or attrs.num_heads > 0) and sameShape(shapes.cos_cache, shapes.sin_cache) and ranks.cos_cache == 2 and ranks.sin_cache == 2 and sectionsValid and ranks.position_ids == 3 and dim(shapes.position_ids, 0) == 3 and dim(shapes.position_ids, 1) == dim(shapes.input, 0) and dim(shapes.position_ids, 2) == dim(shapes.input, 1 if ranks.input == 3 else 2)",
+    "rank3Contract": "commonContract and ranks.input == 3 and ranks.output == 3 and attrs.num_heads is defined and attrs.num_heads >= 1 and dim(shapes.input, 2) % attrs.num_heads == 0 and rank3HeadSize > 0 and effectiveRotaryDim % 2 == 0 and (attrs.rotary_embedding_dim == 0 or attrs.rotary_embedding_dim <= rank3HeadSize) and halfRotaryDim * 2 == effectiveRotaryDim",
+    "rank4Contract": "commonContract and ranks.input == 4 and ranks.output == 4 and dim(shapes.input, 3) > 0 and effectiveRotaryDim % 2 == 0 and (attrs.rotary_embedding_dim == 0 or attrs.rotary_embedding_dim <= dim(shapes.input, 3)) and halfRotaryDim * 2 == effectiveRotaryDim"
+  },
+  "bindingSets": {
+    "$common": [
+      { "name": "x", "arg": "x", "semantic": "input", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      {
+        "name": "position_ids",
+        "arg": "positionIds",
+        "semantic": "position_ids",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "u32"
+      },
+      {
+        "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": "output", "buffer": { "type": "storage" }, "elementType": "$T" }
+    ],
+    "rank3": [
+      { "name": "x", "arg": "x", "semantic": "input", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      {
+        "name": "position_ids",
+        "arg": "positionIds",
+        "semantic": "position_ids",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "u32"
+      },
+      {
+        "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": "output", "buffer": { "type": "storage" }, "elementType": "$T" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "pairCount", "type": "u32", "value": "pairCount" },
+            { "name": "batchSize", "type": "u32", "value": "dim(shapes.input, 0)" },
+            { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.input, 1)" },
+            { "name": "numHeads", "type": "u32", "value": "attrs.num_heads" },
+            { "name": "headSize", "type": "u32", "value": "rank3HeadSize" },
+            { "name": "rotaryDim", "type": "u32", "value": "halfRotaryDim * 2" },
+            { "name": "halfRotaryDim", "type": "u32", "value": "halfRotaryDim" },
+            { "name": "section0", "type": "u32", "value": "attrs.mrope_section[0]" },
+            { "name": "section1", "type": "u32", "value": "attrs.mrope_section[1]" },
+            { "name": "section2", "type": "u32", "value": "attrs.mrope_section[2]" },
+            { "name": "scale", "type": "f32", "value": "attrs.scale" }
+          ]
+        }
+      }
+    ],
+    "rank4": [
+      { "name": "x", "arg": "x", "semantic": "input", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
+      {
+        "name": "position_ids",
+        "arg": "positionIds",
+        "semantic": "position_ids",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "u32"
+      },
+      {
+        "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": "output", "buffer": { "type": "storage" }, "elementType": "$T" },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "pairCount", "type": "u32", "value": "pairCount" },
+            { "name": "batchSize", "type": "u32", "value": "dim(shapes.input, 0)" },
+            { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.input, 2)" },
+            { "name": "numHeads", "type": "u32", "value": "dim(shapes.input, 1)" },
+            { "name": "headSize", "type": "u32", "value": "dim(shapes.input, 3)" },
+            { "name": "rotaryDim", "type": "u32", "value": "halfRotaryDim * 2" },
+            { "name": "halfRotaryDim", "type": "u32", "value": "halfRotaryDim" },
+            { "name": "section0", "type": "u32", "value": "attrs.mrope_section[0]" },
+            { "name": "section1", "type": "u32", "value": "attrs.mrope_section[1]" },
+            { "name": "section2", "type": "u32", "value": "attrs.mrope_section[2]" },
+            { "name": "scale", "type": "f32", "value": "attrs.scale" }
+          ]
+        }
+      }
+    ]
+  },
+  "variants": [
+    {
+      "id": "rank3",
+      "when": ["rank3Contract", "pairDispatchOk"],
+      "constants": {
+        "interleaved": "attrs.interleaved != 0",
+        "mropeSectioned": "attrs.mrope_layout == 0",
+        "usesF16": "dtypes.T == \"f16\"",
+        "scalar": "dtypes.T"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "mrotary_embedding3d",
+          "source": { "shader": "mrotary-embedding.wgsl.jinja", "inputs": { "rank": 3 } },
+          "bindings": "rank3",
+          "dispatch": { "threads": "pairCount", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    },
+    {
+      "id": "rank4",
+      "when": ["rank4Contract", "pairDispatchOk"],
+      "constants": {
+        "interleaved": "attrs.interleaved != 0",
+        "mropeSectioned": "attrs.mrope_layout == 0",
+        "usesF16": "dtypes.T == \"f16\"",
+        "scalar": "dtypes.T"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "mrotary_embedding4d",
+          "source": { "shader": "mrotary-embedding.wgsl.jinja", "inputs": { "rank": 4 } },
+          "bindings": "rank4",
+          "dispatch": { "threads": "pairCount", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

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

+{
+  "name": "com.microsoft.MRotaryEmbedding",
+  "id": "_com_microsoft_mrotaryembedding_webgpu_2097c43",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "RRGYpzSKyDx1HHSFAzgf4u0UdW+WMXLIo+cRg29orLg=",
+      "manifest.json": "rNhF7mp6GzPe1XKu5djfgpCv5MKyJdPU7G06AvbgLXE=",
+      "mrotary-embedding.wgsl.jinja": "ywEaYPPZeG6tSo/Ulh3lIOBwXLhdu70TAZJAgus88cA=",
+      "test.json": "0mcf/8uDq9fTEq082qLPOGoRnjowxCk+vqzBBFVA0Ls="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.MRotaryEmbedding" }
+}

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

	@@ -0,0 +1,108 @@

+{% 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 }}
+// M-RoPE: three position streams (temporal, height, width) index one cos/sin cache.
+// The half-rotary axis is partitioned among the streams, so each rotated pair picks
+// its own stream and gathers that stream's position for this token. 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 + 1u) / 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];
+    if (tail + 1u < params.headSize) {
+      y[base + tail + 1u] = x[base + tail + 1u];
+    }
+    return;
+  }
+  // `pair` indexes the half-rotary axis, which is exactly the cos/sin cache column
+  // under both rotation pairings: interleaved reads column channel/2 and the split
+  // form reads channel % halfRotaryDim, and this lane owns channel == pair in each.
+  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 %}
+  // Which of the three position streams owns this cache column.
+  var stream = 0u;
+{% if mropeSectioned %}
+  // Sectioned (Qwen2-VL, Qwen2.5-VL): three contiguous chunks [T | H | W].
+  if (pair >= params.section0 && pair < params.section0 + params.section1) {
+    stream = 1u;
+  } else if (pair >= params.section0 + params.section1) {
+    stream = 2u;
+  }
+{% else %}
+  // Interleaved (Qwen3-VL, Qwen3.5): T everywhere, then H over every third column
+  // from offset 1 and W over every third column from offset 2, each bounded by its
+  // own section length.
+  if (pair % 3u == 1u && pair / 3u < params.section1) {
+    stream = 1u;
+  } else if (pair % 3u == 2u && pair / 3u < params.section2) {
+    stream = 2u;
+  }
+{% endif %}
+  let pos = position_ids[(stream * params.batchSize + batch) * params.sequenceLength + token];
+  let cache = pos * params.halfRotaryDim + pair;
+  // Round scaled cache values to T before rotation. This intermediate cast is
+  // observable for f16 and must not be deferred to the output store.
+  let cf = f32({{ scalar }}(f32(cos_cache[cache]) * params.scale));
+  let sf = f32({{ scalar }}(f32(sin_cache[cache]) * params.scale));
+{% 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,246 @@

+{
+  "op": "com.microsoft.MRotaryEmbedding",
+  "fixtureArrays": {
+    "ort_sectioned_rank3_input_x": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]
+  },
+  "cases": [
+    {
+      "name": "rank4_odd_head_even_partial_rotation",
+      "provenance": {
+        "source": "onnxruntime/contrib_ops/cuda/bert/mrotary_embedding_impl.cu",
+        "notes": "The provider requires the effective rotary dimension, not the full head size, to be even. The singleton tail element is copied unchanged."
+      },
+      "attrs": {
+        "num_heads": 1,
+        "rotary_embedding_dim": 4,
+        "mrope_layout": 0,
+        "interleaved": 0,
+        "mrope_section": [1, 1, 0]
+      },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 1, 2, 5], "data": { "kind": "linspace", "start": -1.25, "end": 1.5 } },
+        "positionIds": {
+          "dtype": "uint32",
+          "shape": [3, 1, 2],
+          "data": { "kind": "values", "values": [0, 1, 1, 2, 2, 0] }
+        },
+        "cos": { "dtype": "float32", "shape": [3, 2], "data": { "kind": "linspace", "start": 0.25, "end": 1.0 } },
+        "sin": { "dtype": "float32", "shape": [3, 2], "data": { "kind": "linspace", "start": -0.5, "end": 0.5 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 1, 2, 5], "tolerance": 0.00001 } }
+    },
+    {
+      "name": "ort_sectioned_rank3",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/mrotary_embedding_op_test.cc",
+        "test": "ContribOpMRotaryEmbeddingTest.SectionedRank3",
+        "notes": "Pinned from ORT. Sectioned layout over three one-column sections, so each cache column is fed by a different position stream (T=[0,1], H=[1,2], W=[2,3]) and a stream mix-up changes every output."
+      },
+      "attrs": {
+        "num_heads": 2,
+        "rotary_embedding_dim": 6,
+        "mrope_layout": 0,
+        "interleaved": 0,
+        "scale": 1,
+        "mrope_section": [1, 1, 1]
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 12],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_sectioned_rank3_input_x" } }
+        },
+        "positionIds": {
+          "dtype": "uint32",
+          "shape": [3, 1, 2],
+          "data": { "kind": "values", "values": [0, 1, 1, 2, 2, 3] }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [4, 3],
+          "data": { "kind": "values", "values": [1.0, 1.01, 1.02, 1.1, 1.11, 1.12, 1.2, 1.21, 1.22, 1.3, 1.31, 1.32] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [4, 3],
+          "data": { "kind": "values", "values": [0.1, 0.11, 0.12, 0.15, 0.16, 0.17, 0.2, 0.21, 0.22, 0.25, 0.26, 0.27] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [1, 2, 12],
+          "tolerance": 0.00001,
+          "data": {
+            "kind": "values",
+            "values": [0.6, 1.42, 2.34, 4.1, 5.87, 7.98, 6.0, 7.12, 8.34, 10.7, 13.49, 16.62, 11.9, 13.37, 14.94, 19.55, 23.51, 27.81, 17.6, 19.37, 21.24, 27.05, 32.03, 37.35]
+          }
+        }
+      }
+    },
+    {
+      "name": "ort_interleaved_rank4",
+      "provenance": {
+        "source": "onnxruntime/test/contrib_ops/mrotary_embedding_op_test.cc",
+        "test": "ContribOpMRotaryEmbeddingTest.InterleavedRank4",
+        "notes": "Pinned from ORT. Interleaved mrope_layout together with interleaved rotation pairing and a non-unit scale, so the two independent 'interleaved' knobs and the scale are exercised at once."
+      },
+      "attrs": {
+        "num_heads": 2,
+        "rotary_embedding_dim": 6,
+        "mrope_layout": 1,
+        "interleaved": 1,
+        "scale": 0.5,
+        "mrope_section": [1, 1, 1]
+      },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 2, 6],
+          "data": { "kind": "values", "values": { "$ref": "#/fixtureArrays/ort_sectioned_rank3_input_x" } }
+        },
+        "positionIds": {
+          "dtype": "uint32",
+          "shape": [3, 1, 2],
+          "data": { "kind": "values", "values": [0, 1, 1, 2, 2, 3] }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [4, 3],
+          "data": { "kind": "values", "values": [1.0, 1.01, 1.02, 1.1, 1.11, 1.12, 1.2, 1.21, 1.22, 1.3, 1.31, 1.32] }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [4, 3],
+          "data": { "kind": "values", "values": [0.1, 0.11, 0.12, 0.15, 0.16, 0.17, 0.2, 0.21, 0.22, 0.25, 0.26, 0.27] }
+        }
+      },
+      "outputs": {
+        "y": {
+          "dtype": "float32",
+          "shape": [1, 2, 2, 6],
+          "tolerance": 0.00001,
+          "data": {
+            "kind": "values",
+            "values": [0.4, 1.05, 1.345, 2.46, 2.39, 4.21, 3.25, 4.925, 4.395, 6.995, 5.64, 9.405, 5.8, 7.65, 7.045, 10.08, 8.39, 12.85, 8.95, 12.425, 10.395, 15.515, 11.94, 18.945]
+          }
+        }
+      }
+    },
+    {
+      "name": "rank3_layout1_interleaved0",
+      "provenance": {
+        "notes": "Interleaved mrope_layout with the split (non-interleaved) rotation pairing — the layout/pairing knobs are independent, and this is the combination neither pinned case covers."
+      },
+      "attrs": {
+        "num_heads": 4,
+        "rotary_embedding_dim": 0,
+        "mrope_layout": 1,
+        "interleaved": 0,
+        "scale": 1,
+        "mrope_section": [2, 1, 1]
+      },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [2, 3, 32], "data": { "kind": "linspace", "start": -2.0, "end": 2.0 } },
+        "positionIds": {
+          "dtype": "uint32",
+          "shape": [3, 2, 3],
+          "data": { "kind": "values", "values": [0, 1, 2, 1, 2, 3, 2, 3, 4, 3, 4, 5, 4, 5, 6, 5, 6, 7] }
+        },
+        "cos": { "dtype": "float32", "shape": [8, 4], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } },
+        "sin": { "dtype": "float32", "shape": [8, 4], "data": { "kind": "linspace", "start": 1.0, "end": -1.0 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 32], "tolerance": 0.00001, "relTolerance": 0.00001 } }
+    },
+    {
+      "name": "rank4_layout0_partial_rotary_tail",
+      "provenance": {
+        "notes": "Sectioned layout with rotary_embedding_dim < head_size, so the lanes past the rotary window copy their tail values unchanged."
+      },
+      "attrs": {
+        "num_heads": 2,
+        "rotary_embedding_dim": 8,
+        "mrope_layout": 0,
+        "interleaved": 0,
+        "scale": 1,
+        "mrope_section": [1, 2, 1]
+      },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 2, 4, 12], "data": { "kind": "linspace", "start": -2.0, "end": 2.0 } },
+        "positionIds": {
+          "dtype": "uint32",
+          "shape": [3, 1, 4],
+          "data": { "kind": "values", "values": [0, 1, 2, 3, 1, 2, 3, 4, 2, 3, 4, 5] }
+        },
+        "cos": { "dtype": "float32", "shape": [6, 4], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } },
+        "sin": { "dtype": "float32", "shape": [6, 4], "data": { "kind": "linspace", "start": 1.0, "end": -1.0 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 4, 12], "tolerance": 0.00001, "relTolerance": 0.00001 } }
+    },
+    {
+      "name": "rank3_text_only_streams_collapse",
+      "provenance": {
+        "notes": "Text-only tokens set T == H == W, which ORT documents as reducing M-RoPE to standard RoPE; the section split must then become unobservable."
+      },
+      "attrs": { "num_heads": 2, "mrope_section": [2, 1, 1] },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [1, 4, 16], "data": { "kind": "linspace", "start": -2.0, "end": 2.0 } },
+        "positionIds": {
+          "dtype": "uint32",
+          "shape": [3, 1, 4],
+          "data": { "kind": "values", "values": [0, 1, 2, 3, 0, 1, 2, 3, 0, 1, 2, 3] }
+        },
+        "cos": { "dtype": "float32", "shape": [8, 4], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } },
+        "sin": { "dtype": "float32", "shape": [8, 4], "data": { "kind": "linspace", "start": 1.0, "end": -1.0 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 4, 16], "tolerance": 0.00001, "relTolerance": 0.00001 } }
+    },
+    {
+      "name": "f16_rank4_layout1_interleaved",
+      "provenance": {
+        "notes": "f16 storage, including the provider's T-rounded scaled-cache intermediate, on the interleaved layout and pairing."
+      },
+      "attrs": {
+        "rotary_embedding_dim": 0,
+        "mrope_layout": 1,
+        "interleaved": 1,
+        "scale": 0.3,
+        "mrope_section": [2, 2, 2]
+      },
+      "inputs": {
+        "x": { "dtype": "float16", "shape": [1, 2, 3, 12], "data": { "kind": "linspace", "start": -2.0, "end": 2.0 } },
+        "positionIds": {
+          "dtype": "uint32",
+          "shape": [3, 1, 3],
+          "data": { "kind": "values", "values": [0, 1, 2, 1, 2, 3, 2, 3, 4] }
+        },
+        "cos": { "dtype": "float16", "shape": [6, 6], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } },
+        "sin": { "dtype": "float16", "shape": [6, 6], "data": { "kind": "linspace", "start": 1.0, "end": -1.0 } }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [1, 2, 3, 12], "tolerance": 0.01, "relTolerance": 0.01 } }
+    },
+    {
+      "name": "rank3_scale_half_sectioned",
+      "provenance": { "notes": "Non-unit scale on the sectioned layout; scale multiplies both gathered cos and sin." },
+      "attrs": {
+        "num_heads": 2,
+        "rotary_embedding_dim": 0,
+        "mrope_layout": 0,
+        "interleaved": 0,
+        "scale": 0.5,
+        "mrope_section": [1, 1, 2]
+      },
+      "inputs": {
+        "x": { "dtype": "float32", "shape": [2, 2, 16], "data": { "kind": "linspace", "start": -2.0, "end": 2.0 } },
+        "positionIds": {
+          "dtype": "uint32",
+          "shape": [3, 2, 2],
+          "data": { "kind": "values", "values": [0, 1, 1, 2, 2, 3, 3, 4, 4, 0, 0, 1] }
+        },
+        "cos": { "dtype": "float32", "shape": [5, 4], "data": { "kind": "linspace", "start": -1.0, "end": 1.0 } },
+        "sin": { "dtype": "float32", "shape": [5, 4], "data": { "kind": "linspace", "start": 1.0, "end": -1.0 } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 2, 16], "tolerance": 0.00001, "relTolerance": 0.00001 } }
+    }
+  ]
+}