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