| { |
| "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", |
| "quadOk": "quad3Ok if ranks.x == 3 else quad4Ok", |
| "quadDispatchOk": "ceilDiv(numel(shapes.x) / 8, tunables.WORKGROUP_SIZE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)", |
| "headElements": "rank3HeadSize if ranks.x == 3 else dim(shapes.x, 3)", |
| "headCount": "attrs.num_heads if ranks.x == 3 else dim(shapes.x, 1)", |
| "tokenCount": "dim(shapes.x, 1) if ranks.x == 3 else dim(shapes.x, 2)", |
| "rotationPairs": "dim(shapes.cos, -1)", |
| "rotaryOutputWg": "max(1, min(tunables.WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX))" |
| }, |
| "bindings": { |
| "x": { "elementType": "$vector" }, |
| "cos_cache": { "arg": "cos", "elementType": "$vector" }, |
| "sin_cache": { "arg": "sin", "elementType": "$vector" }, |
| "position_ids": { "arg": "positionIds", "elementType": "$M" }, |
| "y": { "elementType": "$vector" }, |
| "params_quad": { |
| "name": "params", |
| "struct": [ |
| { "name": "quadCount", "type": "u32", "value": "numel(shapes.x) / 8" }, |
| { "name": "numHeads", "type": "u32", "value": "headCount" }, |
| { "name": "sequenceLength", "type": "u32", "value": "tokenCount" }, |
| { "name": "headVec4", "type": "u32", "value": "headElements / 4" }, |
| { "name": "headQuads", "type": "u32", "value": "headElements / 8" }, |
| { "name": "rotaryVec4", "type": "u32", "value": "rotationPairs / 2" }, |
| { "name": "halfRotaryVec4", "type": "u32", "value": "rotationPairs / 4" } |
| ] |
| }, |
| "params_pair": { |
| "name": "params", |
| "struct": [ |
| { "name": "pairCount", "type": "u32", "value": "numel(shapes.x) / 2" }, |
| { "name": "numHeads", "type": "u32", "value": "headCount" }, |
| { "name": "sequenceLength", "type": "u32", "value": "tokenCount" }, |
| { "name": "headSize", "type": "u32", "value": "headElements" }, |
| { "name": "rotaryDim", "type": "u32", "value": "rotationPairs * 2" }, |
| { "name": "halfRotaryDim", "type": "u32", "value": "rotationPairs" } |
| ] |
| }, |
| "params_quad_cache3": { |
| "name": "params", |
| "struct": [ |
| { "name": "quadCount", "type": "u32", "value": "numel(shapes.x) / 8" }, |
| { "name": "numHeads", "type": "u32", "value": "headCount" }, |
| { "name": "sequenceLength", "type": "u32", "value": "tokenCount" }, |
| { "name": "headVec4", "type": "u32", "value": "headElements / 4" }, |
| { "name": "headQuads", "type": "u32", "value": "headElements / 8" }, |
| { "name": "rotaryVec4", "type": "u32", "value": "rotationPairs / 2" }, |
| { "name": "halfRotaryVec4", "type": "u32", "value": "rotationPairs / 4" }, |
| { "name": "cacheBatchStrideVec4", "type": "u32", "value": "dim(shapes.cos, 1) * dim(shapes.cos, 2) / 4" } |
| ] |
| }, |
| "params_pair_cache3": { |
| "name": "params", |
| "struct": [ |
| { "name": "pairCount", "type": "u32", "value": "numel(shapes.x) / 2" }, |
| { "name": "numHeads", "type": "u32", "value": "headCount" }, |
| { "name": "sequenceLength", "type": "u32", "value": "tokenCount" }, |
| { "name": "headSize", "type": "u32", "value": "headElements" }, |
| { "name": "rotaryDim", "type": "u32", "value": "rotationPairs * 2" }, |
| { "name": "halfRotaryDim", "type": "u32", "value": "rotationPairs" }, |
| { "name": "cacheBatchStride", "type": "u32", "value": "dim(shapes.cos, 1) * dim(shapes.cos, 2)" } |
| ] |
| }, |
| "x_t": { "name": "x", "elementType": "$T" }, |
| "cos_cache_t": { "arg": "cos", "name": "cos_cache", "elementType": "$T" }, |
| "sin_cache_t": { "arg": "sin", "name": "sin_cache", "elementType": "$T" }, |
| "y_t": { "name": "y", "elementType": "$T" }, |
| "output_vec_params": { |
| "name": "params", |
| "struct": [{ "name": "quadCount", "type": "u32", "value": "numel(shapes.x) / 4" }] |
| } |
| }, |
| "variants": [ |
| { |
| "id": "rank3_cache2_pos_quad", |
| "priority": 10, |
| "when": ["rank3PosContract", "quadOk", "quadDispatchOk"], |
| "derive": { "interleaved": "attrs.interleaved != 0", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RotaryEmbedding.Rank3PosQuad", |
| "shader": "rotary-embedding-quad.wgsl.jinja", |
| "derive": { "rank": 3, "hasPositionIds": "true" }, |
| "bindings": ["x", "cos_cache", "sin_cache", "position_ids", "y", "params_quad"], |
| "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", "scalar": "dtypes.T" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RotaryEmbedding.Rank3Pos", |
| "shader": "rotary-embedding.wgsl.jinja", |
| "derive": { "rank": 3, "hasPositionIds": "true" }, |
| "bindings": ["x_t", "cos_cache_t", "sin_cache_t", "position_ids", "y_t", "params_pair"], |
| "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_cache2_pos_output_vec4", |
| "priority": 20, |
| "when": ["rank3PosContract", "headElements > 0", "headElements % 4 == 0", "rotationPairs % 4 != 0", "attrs.interleaved == 0", "ceilDiv(numel(shapes.x) / 4, rotaryOutputWg) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"], |
| "derive": { "scalar": "dtypes.T", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RotaryEmbedding.Rank3PosOutputVec4", |
| "shader": "rotary-embedding-output-vec4.wgsl.jinja", |
| "derive": { "rank": 3, "hasPositionIds": "true" }, |
| "bindings": ["x", "cos_cache_t", "sin_cache_t", "position_ids", "y", "output_vec_params"], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank4_cache2_pos_quad", |
| "priority": 10, |
| "when": ["rank4PosContract", "quadOk", "quadDispatchOk"], |
| "derive": { "interleaved": "attrs.interleaved != 0", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RotaryEmbedding.Rank4PosQuad", |
| "shader": "rotary-embedding-quad.wgsl.jinja", |
| "derive": { "rank": 4, "hasPositionIds": "true" }, |
| "bindings": ["x", "cos_cache", "sin_cache", "position_ids", "y", "params_quad"], |
| "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", "scalar": "dtypes.T" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RotaryEmbedding.Rank4Pos", |
| "shader": "rotary-embedding.wgsl.jinja", |
| "derive": { "rank": 4, "hasPositionIds": "true" }, |
| "bindings": ["x_t", "cos_cache_t", "sin_cache_t", "position_ids", "y_t", "params_pair"], |
| "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_output_vec4", |
| "priority": 20, |
| "when": ["rank4PosContract", "headElements > 0", "headElements % 4 == 0", "rotationPairs % 4 != 0", "attrs.interleaved == 0", "ceilDiv(numel(shapes.x) / 4, rotaryOutputWg) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"], |
| "derive": { "scalar": "dtypes.T", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RotaryEmbedding.Rank4PosOutputVec4", |
| "shader": "rotary-embedding-output-vec4.wgsl.jinja", |
| "derive": { "rank": 4, "hasPositionIds": "true" }, |
| "bindings": ["x", "cos_cache_t", "sin_cache_t", "position_ids", "y", "output_vec_params"], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank3_cache3_quad", |
| "priority": 10, |
| "when": ["rank3Contract", "quadOk", "quadDispatchOk"], |
| "derive": { "interleaved": "attrs.interleaved != 0", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RotaryEmbedding.Rank3Quad", |
| "shader": "rotary-embedding-quad.wgsl.jinja", |
| "derive": { "rank": 3 }, |
| "bindings": ["x", "cos_cache", "sin_cache", "y", "params_quad_cache3"], |
| "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", "scalar": "dtypes.T" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RotaryEmbedding.Rank3", |
| "shader": "rotary-embedding.wgsl.jinja", |
| "derive": { "rank": 3 }, |
| "bindings": ["x_t", "cos_cache_t", "sin_cache_t", "y_t", "params_pair_cache3"], |
| "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_output_vec4", |
| "priority": 20, |
| "when": ["rank3Contract", "headElements > 0", "headElements % 4 == 0", "rotationPairs % 4 != 0", "attrs.interleaved == 0", "ceilDiv(numel(shapes.x) / 4, rotaryOutputWg) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"], |
| "derive": { "scalar": "dtypes.T", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RotaryEmbedding.Rank3OutputVec4", |
| "shader": "rotary-embedding-output-vec4.wgsl.jinja", |
| "derive": { "rank": 3 }, |
| "bindings": ["x", "cos_cache_t", "sin_cache_t", "y", "output_vec_params"], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| }, |
| { |
| "id": "rank4_cache3_quad", |
| "priority": 10, |
| "when": ["rank4Contract", "quadOk", "quadDispatchOk"], |
| "derive": { "interleaved": "attrs.interleaved != 0", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RotaryEmbedding.Rank4Quad", |
| "shader": "rotary-embedding-quad.wgsl.jinja", |
| "derive": { "rank": 4 }, |
| "bindings": ["x", "cos_cache", "sin_cache", "y", "params_quad_cache3"], |
| "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", "scalar": "dtypes.T" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RotaryEmbedding.Rank4", |
| "shader": "rotary-embedding.wgsl.jinja", |
| "derive": { "rank": 4 }, |
| "bindings": ["x_t", "cos_cache_t", "sin_cache_t", "y_t", "params_pair_cache3"], |
| "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_output_vec4", |
| "priority": 20, |
| "when": ["rank4Contract", "headElements > 0", "headElements % 4 == 0", "rotationPairs % 4 != 0", "attrs.interleaved == 0", "ceilDiv(numel(shapes.x) / 4, rotaryOutputWg) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"], |
| "derive": { "scalar": "dtypes.T", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" }, |
| "passes": [ |
| { |
| "id": "main", |
| "name": "RotaryEmbedding.Rank4OutputVec4", |
| "shader": "rotary-embedding-output-vec4.wgsl.jinja", |
| "derive": { "rank": 4 }, |
| "bindings": ["x", "cos_cache_t", "sin_cache_t", "y", "output_vec_params"], |
| "dispatch": { |
| "x": "min(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)", |
| "y": "ceilDiv(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)", |
| "z": 1 |
| } |
| } |
| ] |
| } |
| ] |
| } |
|
|
