File size: 12,641 Bytes
5ce7237 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 | {
"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" }
}
]
}
]
}
|