sync 91d990483a17

Browse files

Files changed (7) hide show

README.md +17 -12
build/webgpu/bench.json +0 -1
build/webgpu/manifest.json +270 -233
build/webgpu/metadata.json +20 -7
build/webgpu/rotary-embedding-quad.wgsl.jinja +88 -0
build/webgpu/rotary-embedding.wgsl.jinja +17 -17
build/webgpu/test.json +108 -3

README.md CHANGED Viewed

@@ -12,24 +12,24 @@ tags:
 ## 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
@@ -38,8 +38,8 @@ 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
@@ -50,18 +50,23 @@ Attributes and default values (overridable per request):
 ## 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.

 ## 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 unsupported.
 See the [ONNX `RotaryEmbedding` spec](https://onnx.ai/onnx/operators/onnx__RotaryEmbedding.html) for the reference semantics.
 ## Inputs
+| Name | Upstream name | 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` | `cos_cache` | `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` | `sin_cache` | `T` | same as logical dtype | — | — | Precomputed sine values with the same shape and type as `cos_cache`. | required |
+| `positionIds` | `position_ids` | `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 | Upstream name | 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
 | 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. |
 | `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` | `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. |
 ## Type constraints
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, 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-quad.wgsl.jinja`](build/webgpu/rotary-embedding-quad.wgsl.jinja)
 - [`rotary-embedding.wgsl.jinja`](build/webgpu/rotary-embedding.wgsl.jinja)
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.RotaryEmbedding",
   "tunableSpace": { "WORKGROUP_SIZE": [128, 256] },
   "cases": [
     {

 {
   "tunableSpace": { "WORKGROUP_SIZE": [128, 256] },
   "cases": [
     {

build/webgpu/manifest.json CHANGED Viewed

@@ -2,287 +2,324 @@
   "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" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "RotaryEmbedding",
   "sinceVersion": 23,
+  "inputs": {
+    "x": { "onnx": "X", "dtype": "T" },
+    "cos": { "onnx": "cos_cache", "dtype": "T" },
+    "sin": { "onnx": "sin_cache", "dtype": "T" },
+    "positionIds": {
+      "onnx": "position_ids",
       "dtype": "M",
       "rank": 2,
       "optional": true,
+      "storage": "uint32",
+      "narrowing": "checked"
     }
   },
+  "outputs": { "y": { "onnx": "Y", "dtype": "T", "rank": "ranks.x", "shape": "shapes.x" } },
+  "attributes": { "interleaved": { "default": 0 }, "rotary_embedding_dim": { "default": 0 }, "num_heads": {} },
   "attributeConstraints": { "interleaved": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32", "float16"], "M": ["int64"] },
+  "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
   "derive": {
+    "pairDispatchOk": "ceilDiv(numel(shapes.x) / 2, tunables.WORKGROUP_SIZE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "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 == 3 and ranks.sin == 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, 2) * 2 == (attrs.rotary_embedding_dim if attrs.rotary_embedding_dim != 0 else rank3HeadSize) and sameShape(shapes.cos, shapes.sin) and dim(shapes.cos, 0) == dim(shapes.x, 0) and dim(shapes.cos, 1) == dim(shapes.x, 1) and not present.positionIds",
+    "rank4Contract": "f16Ok(dtypes.T) and ranks.x == 4 and ranks.cos == 3 and ranks.sin == 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, 2) * 2 == (attrs.rotary_embedding_dim if attrs.rotary_embedding_dim != 0 else dim(shapes.x, 3)) and dim(shapes.sin, 2) == dim(shapes.cos, 2) and sameShape(shapes.x, shapes.y) and dim(shapes.cos, 0) == dim(shapes.x, 0) and dim(shapes.sin, 0) == dim(shapes.x, 0) and dim(shapes.cos, 1) == dim(shapes.x, 2) and dim(shapes.sin, 1) == dim(shapes.x, 2) and not present.positionIds",
+    "rank3PosContract": "f16Ok(dtypes.T) and ranks.x == 3 and ranks.cos == 2 and ranks.sin == 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, 1) * 2 == (attrs.rotary_embedding_dim if attrs.rotary_embedding_dim != 0 else rank3HeadSize) and sameShape(shapes.cos, shapes.sin) and present.positionIds and ranks.positionIds == 2 and dim(shapes.positionIds, 0) == dim(shapes.x, 0) and dim(shapes.positionIds, 1) == dim(shapes.x, 1)",
+    "rank4PosContract": "f16Ok(dtypes.T) and ranks.x == 4 and ranks.cos == 2 and ranks.sin == 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, 1) * 2 == (attrs.rotary_embedding_dim if attrs.rotary_embedding_dim != 0 else dim(shapes.x, 3)) and sameShape(shapes.cos, shapes.sin) and sameShape(shapes.x, shapes.y) and present.positionIds and ranks.positionIds == 2 and dim(shapes.positionIds, 0) == dim(shapes.x, 0) and dim(shapes.positionIds, 1) == dim(shapes.x, 2)",
+    "quadCacheOk": "ranks.cos >= 1 and dim(shapes.cos, ranks.cos - 1) % 4 == 0",
+    "quad3Ok": "rank3HeadSize % 8 == 0 and quadCacheOk",
+    "quad4Ok": "ranks.x == 4 and dim(shapes.x, 3) % 8 == 0 and quadCacheOk",
+    "quadDispatchOk": "ceilDiv(numel(shapes.x) / 8, tunables.WORKGROUP_SIZE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"
   },
+  "bindings": {
+    "x": { "buffer": "read-only-storage", "elementType": "$vector" },
+    "cos_cache": { "arg": "cos", "buffer": "read-only-storage", "elementType": "$vector" },
+    "sin_cache": { "arg": "sin", "buffer": "read-only-storage", "elementType": "$vector" },
+    "position_ids": { "arg": "positionIds", "buffer": "read-only-storage", "elementType": "$M" },
+    "y": { "buffer": "storage", "elementType": "$vector" },
+    "params": {
+      "buffer": "uniform",
+      "struct": [
+        { "name": "quadCount", "type": "u32", "value": "numel(shapes.x) / 8" },
+        { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.x, 1)" },
+        { "name": "numHeads", "type": "u32", "value": "attrs.num_heads" },
+        { "name": "headVec4", "type": "u32", "value": "rank3HeadSize / 4" },
+        { "name": "headQuads", "type": "u32", "value": "rank3HeadSize / 8" },
+        { "name": "rotaryVec4", "type": "u32", "value": "dim(shapes.cos, 1) / 2" },
+        { "name": "halfRotaryVec4", "type": "u32", "value": "dim(shapes.cos, 1) / 4" }
+      ]
+    },
+    "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$T" },
+    "cos_cache_2": { "arg": "cos", "name": "cos_cache", "buffer": "read-only-storage", "elementType": "$T" },
+    "sin_cache_2": { "arg": "sin", "name": "sin_cache", "buffer": "read-only-storage", "elementType": "$T" },
+    "y_2": { "name": "y", "buffer": "storage", "elementType": "$T" },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "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, 1) * 2" },
+        { "name": "halfRotaryDim", "type": "u32", "value": "dim(shapes.cos, 1)" }
+      ]
+    },
+    "params_3": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "quadCount", "type": "u32", "value": "numel(shapes.x) / 8" },
+        { "name": "numHeads", "type": "u32", "value": "dim(shapes.x, 1)" },
+        { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.x, 2)" },
+        { "name": "headVec4", "type": "u32", "value": "dim(shapes.x, 3) / 4" },
+        { "name": "headQuads", "type": "u32", "value": "dim(shapes.x, 3) / 8" },
+        { "name": "rotaryVec4", "type": "u32", "value": "dim(shapes.cos, 1) / 2" },
+        { "name": "halfRotaryVec4", "type": "u32", "value": "dim(shapes.cos, 1) / 4" }
+      ]
+    },
+    "params_4": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "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, 1) * 2" },
+        { "name": "halfRotaryDim", "type": "u32", "value": "dim(shapes.cos, 1)" }
+      ]
+    },
+    "params_5": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "quadCount", "type": "u32", "value": "numel(shapes.x) / 8" },
+        { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.x, 1)" },
+        { "name": "numHeads", "type": "u32", "value": "attrs.num_heads" },
+        { "name": "headVec4", "type": "u32", "value": "rank3HeadSize / 4" },
+        { "name": "headQuads", "type": "u32", "value": "rank3HeadSize / 8" },
+        { "name": "rotaryVec4", "type": "u32", "value": "dim(shapes.cos, 2) / 2" },
+        { "name": "halfRotaryVec4", "type": "u32", "value": "dim(shapes.cos, 2) / 4" },
+        { "name": "cacheBatchStrideVec4", "type": "u32", "value": "dim(shapes.cos, 1) * dim(shapes.cos, 2) / 4" }
+      ]
+    },
+    "params_6": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "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, 2) * 2" },
+        { "name": "halfRotaryDim", "type": "u32", "value": "dim(shapes.cos, 2)" },
+        { "name": "cacheBatchStride", "type": "u32", "value": "dim(shapes.cos, 1) * dim(shapes.cos, 2)" }
+      ]
+    },
+    "params_7": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "quadCount", "type": "u32", "value": "numel(shapes.x) / 8" },
+        { "name": "numHeads", "type": "u32", "value": "dim(shapes.x, 1)" },
+        { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.x, 2)" },
+        { "name": "headVec4", "type": "u32", "value": "dim(shapes.x, 3) / 4" },
+        { "name": "headQuads", "type": "u32", "value": "dim(shapes.x, 3) / 8" },
+        { "name": "rotaryVec4", "type": "u32", "value": "dim(shapes.cos, 2) / 2" },
+        { "name": "halfRotaryVec4", "type": "u32", "value": "dim(shapes.cos, 2) / 4" },
+        { "name": "cacheBatchStrideVec4", "type": "u32", "value": "dim(shapes.cos, 1) * dim(shapes.cos, 2) / 4" }
+      ]
+    },
+    "params_8": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "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, 2) * 2" },
+        { "name": "halfRotaryDim", "type": "u32", "value": "dim(shapes.cos, 2)" },
+        { "name": "cacheBatchStride", "type": "u32", "value": "dim(shapes.cos, 1) * dim(shapes.cos, 2)" }
+      ]
+    }
   },
   "variants": [
+    {
+      "id": "rank3_cache2_pos_quad",
+      "priority": 10,
+      "when": ["rank3PosContract", "quad3Ok", "quadDispatchOk"],
+      "derive": {
+        "interleaved": "attrs.interleaved != 0",
+        "usesF16": "dtypes.T == \"f16\"",
+        "vector": "\"vec4<\" ~ dtypes.T ~ \">\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "rotary_embedding3d_pos_quad",
+          "shader": "rotary-embedding-quad.wgsl.jinja",
+          "derive": { "rank": 3, "hasPositionIds": "true" },
+          "bindings": ["x", "cos_cache", "sin_cache", "position_ids", "y", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
+        }
+      ]
+    },
     {
       "id": "rank3_cache2_pos",
+      "priority": 0,
       "when": ["rank3PosContract", "pairDispatchOk"],
+      "derive": { "interleaved": "attrs.interleaved != 0", "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "rotary_embedding3d_pos",
+          "shader": "rotary-embedding.wgsl.jinja",
+          "derive": { "rank": 3, "hasPositionIds": "true" },
+          "bindings": ["x_2", "cos_cache_2", "sin_cache_2", "position_ids", "y_2", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
+        }
+      ]
+    },
+    {
+      "id": "rank4_cache2_pos_quad",
+      "priority": 10,
+      "when": ["rank4PosContract", "quad4Ok", "quadDispatchOk"],
+      "derive": {
+        "interleaved": "attrs.interleaved != 0",
+        "usesF16": "dtypes.T == \"f16\"",
+        "vector": "\"vec4<\" ~ dtypes.T ~ \">\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "rotary_embedding4d_pos_quad",
+          "shader": "rotary-embedding-quad.wgsl.jinja",
+          "derive": { "rank": 4, "hasPositionIds": "true" },
+          "bindings": ["x", "cos_cache", "sin_cache", "position_ids", "y", "params_3"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "rank4_cache2_pos",
+      "priority": 0,
       "when": ["rank4PosContract", "pairDispatchOk"],
+      "derive": { "interleaved": "attrs.interleaved != 0", "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "rotary_embedding4d_pos",
+          "shader": "rotary-embedding.wgsl.jinja",
+          "derive": { "rank": 4, "hasPositionIds": "true" },
+          "bindings": ["x_2", "cos_cache_2", "sin_cache_2", "position_ids", "y_2", "params_4"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
+        }
+      ]
+    },
+    {
+      "id": "rank3_cache3_quad",
+      "priority": 10,
+      "when": ["rank3Contract", "quad3Ok", "quadDispatchOk"],
+      "derive": {
+        "interleaved": "attrs.interleaved != 0",
+        "usesF16": "dtypes.T == \"f16\"",
+        "vector": "\"vec4<\" ~ dtypes.T ~ \">\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "rotary_embedding3d_quad",
+          "shader": "rotary-embedding-quad.wgsl.jinja",
+          "derive": { "rank": 3 },
+          "bindings": ["x", "cos_cache", "sin_cache", "y", "params_5"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "rank3_cache3",
+      "priority": 0,
       "when": ["rank3Contract", "pairDispatchOk"],
+      "derive": { "interleaved": "attrs.interleaved != 0", "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "rotary_embedding3d",
+          "shader": "rotary-embedding.wgsl.jinja",
+          "derive": { "rank": 3 },
+          "bindings": ["x_2", "cos_cache_2", "sin_cache_2", "y_2", "params_6"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
+        }
+      ]
+    },
+    {
+      "id": "rank4_cache3_quad",
+      "priority": 10,
+      "when": ["rank4Contract", "quad4Ok", "quadDispatchOk"],
+      "derive": {
+        "interleaved": "attrs.interleaved != 0",
+        "usesF16": "dtypes.T == \"f16\"",
+        "vector": "\"vec4<\" ~ dtypes.T ~ \">\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "rotary_embedding4d_quad",
+          "shader": "rotary-embedding-quad.wgsl.jinja",
+          "derive": { "rank": 4 },
+          "bindings": ["x", "cos_cache", "sin_cache", "y", "params_7"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "rank4_cache3",
+      "priority": 0,
       "when": ["rank4Contract", "pairDispatchOk"],
+      "derive": { "interleaved": "attrs.interleaved != 0", "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
       "passes": [
         {
           "id": "main",
           "name": "rotary_embedding4d",
+          "shader": "rotary-embedding.wgsl.jinja",
+          "derive": { "rank": 4 },
+          "bindings": ["x_2", "cos_cache_2", "sin_cache_2", "y_2", "params_8"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,18 +1,31 @@
 {
   "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" }
 }

 {
   "name": "ai.onnx.RotaryEmbedding",
+  "id": "_ai_onnx_rotaryembedding_webgpu_5abb823",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "Ii8Ve/7TIfVkcj95zDTEbsfND6RUQkTEHTKUFQhHaEk=",
+      "manifest.json": "OeUM+N0V8ex5eCGPP+KmT0LDzPNRCc50xnAuaYDS8o0=",
+      "rotary-embedding-quad.wgsl.jinja": "vGNZeSSXSpPhhBz/0Wh9Z1zY5lSqp0VlCsfWmIM98CQ=",
+      "rotary-embedding.wgsl.jinja": "DNauIDqC9Gs5N07X/rqra/Q72p6LnkcYueWx0UAaXlo=",
+      "test.json": "394nZIOecEWid/Qm/+4eUAtAw0l+WnhmDdAuuFb96IA="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "rank3_cache2_pos_quad": ["rotary-embedding-quad.wgsl.jinja"],
+      "rank3_cache2_pos": ["rotary-embedding.wgsl.jinja"],
+      "rank4_cache2_pos_quad": ["rotary-embedding-quad.wgsl.jinja"],
+      "rank4_cache2_pos": ["rotary-embedding.wgsl.jinja"],
+      "rank3_cache3_quad": ["rotary-embedding-quad.wgsl.jinja"],
+      "rank3_cache3": ["rotary-embedding.wgsl.jinja"],
+      "rank4_cache3_quad": ["rotary-embedding-quad.wgsl.jinja"],
+      "rank4_cache3": ["rotary-embedding.wgsl.jinja"]
+    }
+  }
 }

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

	@@ -0,0 +1,88 @@

+{% 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 dispatch's
+  // per-axis workgroup fold width (outputs > 16.7M elements).
+{% elif note == "limit" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+{% elif note == "device-axis" %}
+  // The flat dispatch is folded across x/y at a fixed per-axis workgroup
+  // width; 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 dispatch's
+  // per-axis workgroup fold width (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
+  // dispatch's per-axis workgroup fold width.
+{% elif note == "dispatch" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+{% endif %}
+{% if bound == "" %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
+{%- elif guardInline %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
+  if ({{ name }} >= {{ bound }}) { return; }
+{%- else %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }
+{%- endif %}
+{% endmacro %}
+{% if usesF16 %}enable f16;
+{% endif %}{{ env.wgsl.resourceDeclarations }}
+// One invocation owns four consecutive rotated pairs, so x, cos and sin are read as vec4s and
+// both rotated outputs are written as vec4s; the head size and rotary dim are multiples of
+// eight, so no quad straddles a head or the rotated/tail boundary. Partial rotations use the
+// remaining quad lanes to copy two unchanged tail vec4s. Every index below counts vec4s.
+@compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("q", "params.quadCount", note="") }}
+  let task = q % params.headQuads;
+  let headFlat = q / params.headQuads;
+  let base = headFlat * params.headVec4;
+  if (task >= params.halfRotaryVec4) {
+    let tail = params.rotaryVec4 + (task - params.halfRotaryVec4) * 2u;
+    y[base + tail] = x[base + tail];
+    y[base + tail + 1u] = x[base + tail + 1u];
+    return;
+  }
+  let pair = task;
+{% if 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 hasPositionIds is defined and hasPositionIds %}
+  // Gather the 2D cache tables through this token's position index.
+  let cache = u32(position_ids[batch * params.sequenceLength + token]) * params.halfRotaryVec4 + pair;
+{% else %}
+  // The standard no-index cache owns one row per (batch, token).
+  let cache = batch * params.cacheBatchStrideVec4 + token * params.halfRotaryVec4 + pair;
+{% endif %}
+  let cf = vec4<f32>(cos_cache[cache]);
+  let sf = vec4<f32>(sin_cache[cache]);
+{% if interleaved %}
+  // Four interleaved pairs span two vec4s: (x0.x, x0.y), (x0.z, x0.w), (x1.x, x1.y), (x1.z, x1.w).
+  let x0 = vec4<f32>(x[base + pair * 2u]);
+  let x1 = vec4<f32>(x[base + pair * 2u + 1u]);
+  let a = vec4<f32>(x0.x, x0.z, x1.x, x1.z);
+  let b = vec4<f32>(x0.y, x0.w, x1.y, x1.w);
+  let ra = a * cf - b * sf;
+  let rb = a * sf + b * cf;
+  y[base + pair * 2u] = {{ vector }}(vec4<f32>(ra.x, rb.x, ra.y, rb.y));
+  y[base + pair * 2u + 1u] = {{ vector }}(vec4<f32>(ra.z, rb.z, ra.w, rb.w));
+{% else %}
+  let a = vec4<f32>(x[base + pair]);
+  let b = vec4<f32>(x[base + params.halfRotaryVec4 + pair]);
+  y[base + pair] = {{ vector }}(a * cf - b * sf);
+  y[base + params.halfRotaryVec4 + pair] = {{ vector }}(a * sf + b * cf);
+{% endif %}
+}

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

@@ -1,30 +1,30 @@
 {% 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;
   }
@@ -36,7 +36,7 @@
 // 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;
@@ -52,14 +52,14 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
   }
   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 %}

 {% 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 dispatch's
+  // per-axis workgroup fold width (outputs > 16.7M elements).
 {% elif note == "limit" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
 {% elif note == "device-axis" %}
+  // The flat dispatch is folded across x/y at a fixed per-axis workgroup
+  // width; 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 dispatch's
+  // per-axis workgroup fold width (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
+  // dispatch's per-axis workgroup fold width.
 {% elif note == "dispatch" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
 {% endif %}
 {% if bound == "" %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
 {%- elif guardInline %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) { return; }
 {%- else %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
   }
 // 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>) {
 {{ flat_index_2d("p", "params.pairCount", note="") }}
   let tasksPerHead = params.headSize / 2u;
   }
   let pair = task;
+{% if 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 hasPositionIds is defined and 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 %}

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "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],
@@ -34,7 +33,7 @@
       "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] } },
@@ -628,7 +627,7 @@
     {
       "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": {
@@ -825,6 +824,112 @@
         "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."
       }
     }
   ]
 }

 {
   "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],
       "provenance": {
         "source": "onnxruntime/test/providers/cpu/llm/rotary_embedding_op_test.cc",
         "test": "RotaryEmbedding",
+        "notes": "A rank-4 identity-angle rotary embedding must preserve subnormal activations."
       },
       "inputs": {
         "x": { "dtype": "float32", "shape": [1, 1, 1, 2], "data": { "kind": "values", "values": [1e-40, -1e-40] } },
     {
       "name": "dispatch_fold_rank4_pair_last_row_seq65540",
       "provenance": {
+        "notes": "Shape [1,1,65540,512] requires 65,540 workgroups for paired rotation, so dispatch folds into two dimensions. The highest pair indices must still be written with the correct periodic values."
       },
       "inputs": {
         "x": {
         "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."
       }
+    },
+    {
+      "name": "quad_interleaved_rank3_head16_multihead_f16",
+      "provenance": {
+        "notes": "Interleaved rotation on the quad route: head size 16 over 2 heads with f16 storage, so each invocation rotates four interleaved pairs read as two vec4s and writes them back in the same interleaving."
+      },
+      "attrs": { "interleaved": 1, "num_heads": 2 },
+      "inputs": {
+        "x": {
+          "dtype": "float16",
+          "shape": [1, 3, 32],
+          "data": { "kind": "fillFloat32", "scale": 0.3, "sinStep": 0.13, "cosStep": 0.21 }
+        },
+        "cos": {
+          "dtype": "float16",
+          "shape": [1, 3, 8],
+          "data": { "kind": "rotaryCos", "thetaStart": 0.05, "thetaStep": 0.07 }
+        },
+        "sin": {
+          "dtype": "float16",
+          "shape": [1, 3, 8],
+          "data": { "kind": "rotarySin", "thetaStart": 0.05, "thetaStep": 0.07 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float16", "shape": [1, 3, 32], "tolerance": 0.002 } }
+    },
+    {
+      "name": "quad_partial_rotary_tail_rank4_head24_rd8",
+      "provenance": {
+        "notes": "Partial rotation on the quad route: rotary dim 8 of head size 24 leaves a 16-element tail that the remaining quad lanes copy as two vec4s each, with a batch-strided cache."
+      },
+      "attrs": { "rotary_embedding_dim": 8 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 2, 3, 24],
+          "data": { "kind": "fillFloat32", "scale": 0.45, "sinStep": 0.19, "cosStep": 0.11 }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "fillFloat32", "scale": 1.0, "sinStep": 0.31, "cosStep": 0.17 }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [2, 3, 4],
+          "data": { "kind": "fillFloat32", "scale": 1.0, "sinStep": 0.23, "cosStep": 0.29 }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 2, 3, 24], "tolerance": 0.000001 } }
+    },
+    {
+      "name": "quad_pos_ids_rank3_head16_gather",
+      "provenance": {
+        "notes": "Position-id gather on the quad route: 2D cache tables of half width 8 are gathered per (batch, token) as vec4 rows, head size 16 over 2 heads."
+      },
+      "attrs": { "num_heads": 2 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [2, 3, 32],
+          "data": { "kind": "fillFloat32", "scale": 0.4, "sinStep": 0.17, "cosStep": 0.23 }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [12, 8],
+          "data": { "kind": "fillFloat32", "scale": 1.0, "sinStep": 0.27, "cosStep": 0.13 }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [12, 8],
+          "data": { "kind": "fillFloat32", "scale": 1.0, "sinStep": 0.19, "cosStep": 0.31 }
+        },
+        "positionIds": {
+          "dtype": "uint32",
+          "shape": [2, 3],
+          "data": { "kind": "values", "values": [4, 5, 11, 0, 2, 9] }
+        }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 32], "tolerance": 0.0001 } }
+    },
+    {
+      "name": "quad_pos_ids_rank4_interleaved_head8",
+      "provenance": {
+        "notes": "Position-id gather with interleaved rotation on the rank-4 quad route: head size 8 is exactly one quad per head, so the rotate lanes cover every element and the tail branch never runs."
+      },
+      "attrs": { "interleaved": 1 },
+      "inputs": {
+        "x": {
+          "dtype": "float32",
+          "shape": [1, 2, 3, 8],
+          "data": { "kind": "fillFloat32", "scale": 0.4, "sinStep": 0.11, "cosStep": 0.29 }
+        },
+        "cos": {
+          "dtype": "float32",
+          "shape": [8, 4],
+          "data": { "kind": "fillFloat32", "scale": 1.0, "sinStep": 0.33, "cosStep": 0.21 }
+        },
+        "sin": {
+          "dtype": "float32",
+          "shape": [8, 4],
+          "data": { "kind": "fillFloat32", "scale": 1.0, "sinStep": 0.15, "cosStep": 0.37 }
+        },
+        "positionIds": { "dtype": "uint32", "shape": [1, 3], "data": { "kind": "values", "values": [7, 0, 3] } }
+      },
+      "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 3, 8], "tolerance": 0.0001 } }
     }
   ]
 }