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