Hugging Face's logo

Kernels:
webgpu-kernels
/
com.microsoft.MRotaryEmbedding

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
296e50f
·
verified ·
1 Parent(s): 92e5044

sync 91d990483a17

Browse files
Files changed (6) hide show
  1. README.md +15 -11
  2. build/webgpu/bench.json +0 -1
  3. build/webgpu/manifest.json +83 -200
  4. build/webgpu/metadata.json +10 -7
  5. build/webgpu/mrotary-embedding.wgsl.jinja +16 -16
  6. build/webgpu/test.json +2 -3
README.md CHANGED
@@ -18,18 +18,18 @@ See the [ONNX Runtime `MRotaryEmbedding` contrib-operator spec](https://github.c
18
 
19
  ## Inputs
20
 
21
- | Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
22
  | --- | --- | --- | --- | --- | --- | --- | --- |
23
- | `input` | `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. The effective rotary dimension must be even, and `num_heads` is required for rank-3 input. | required |
24
- | `position_ids` | `positionIds` | `M` | `uint32` | `3` | — | 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. | required |
25
- | `cos_cache` | `cos` | `T` | same as logical dtype | `2` | — | Precomputed cosine values of shape `(max_sequence_length, rotary_dim/2)`, shared by all three position streams. | required |
26
- | `sin_cache` | `sin` | `T` | same as logical dtype | `2` | — | Precomputed sine values with the same shape and type as `cos_cache`. | required |
27
 
28
  ## Outputs
29
 
30
- | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
31
  | --- | --- | --- | --- | --- | --- | --- |
32
- | `output` | `y` | `T` | same as `input` | same as `input` | Rotary-position-encoded tensor with the same shape and type as `input`. | required |
33
 
34
  ## Attributes
35
 
@@ -40,10 +40,10 @@ Attributes and default values (overridable per request):
40
  | `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. This is the rotation pairing and is independent of `mrope_layout`. |
41
  | `is_packed_batching` | `0` | Whether `position_ids` uses packed-batch metadata. The default and only supported value is 0; packed batching (1) is not implemented. |
42
  | `mrope_layout` | `0` | 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. |
 
43
  | `num_heads` | `0` | 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`. |
44
  | `rotary_embedding_dim` | `0` | 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. |
45
  | `scale` | `1` | Scale applied to the gathered cosine and sine values before the rotation. Default is 1.0. |
46
- | `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. |
47
 
48
  ## Type constraints
49
 
@@ -54,7 +54,7 @@ Attributes and default values (overridable per request):
54
 
55
  ## Files
56
 
57
- - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
58
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
59
  - [`test.json`](build/webgpu/test.json) — correctness cases
60
  - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -62,10 +62,14 @@ Attributes and default values (overridable per request):
62
 
63
  ## Use with `@huggingface/kernels`
64
 
65
- The loader derives every required output's shape and logical dtype from the manifest contract and this call.
66
- It then allocates the result tensors automatically.
 
 
 
67
 
68
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 
69
 
70
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
71
 
 
18
 
19
  ## Inputs
20
 
21
+ | Name | Upstream name | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
22
  | --- | --- | --- | --- | --- | --- | --- | --- |
23
+ | `x` | `input` | `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. The effective rotary dimension must be even, and `num_heads` is required for rank-3 input. | required |
24
+ | `positionIds` | `position_ids` | `M` | `uint32` | `3` | — | 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. | required |
25
+ | `cos` | `cos_cache` | `T` | same as logical dtype | `2` | — | Precomputed cosine values of shape `(max_sequence_length, rotary_dim/2)`, shared by all three position streams. | required |
26
+ | `sin` | `sin_cache` | `T` | same as logical dtype | `2` | — | Precomputed sine values with the same shape and type as `cos_cache`. | required |
27
 
28
  ## Outputs
29
 
30
+ | Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
31
  | --- | --- | --- | --- | --- | --- | --- |
32
+ | `y` | `output` | `T` | same as `x` | same as `x` | Rotary-position-encoded tensor with the same shape and type as `input`. | required |
33
 
34
  ## Attributes
35
 
 
40
  | `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. This is the rotation pairing and is independent of `mrope_layout`. |
41
  | `is_packed_batching` | `0` | Whether `position_ids` uses packed-batch metadata. The default and only supported value is 0; packed batching (1) is not implemented. |
42
  | `mrope_layout` | `0` | 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. |
43
+ | `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. |
44
  | `num_heads` | `0` | 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`. |
45
  | `rotary_embedding_dim` | `0` | 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. |
46
  | `scale` | `1` | Scale applied to the gathered cosine and sine values before the rotation. Default is 1.0. |
 
47
 
48
  ## Type constraints
49
 
 
54
 
55
  ## Files
56
 
57
+ - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
58
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
59
  - [`test.json`](build/webgpu/test.json) — correctness cases
60
  - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 
62
 
63
  ## Use with `@huggingface/kernels`
64
 
65
+ ```sh
66
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.2
67
+ ```
68
+
69
+ Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
70
 
71
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
72
+ It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
73
 
74
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
75
 
build/webgpu/bench.json CHANGED
@@ -1,5 +1,4 @@
1
  {
2
- "op": "com.microsoft.MRotaryEmbedding",
3
  "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
4
  "cases": [
5
  {
 
1
  {
 
2
  "tunableSpace": { "WORKGROUP_SIZE": [64, 128, 256] },
3
  "cases": [
4
  {
build/webgpu/manifest.json CHANGED
@@ -2,48 +2,21 @@
2
  "domain": "com.microsoft",
3
  "name": "MRotaryEmbedding",
4
  "sinceVersion": 1,
5
- "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.",
6
- "inputs": [
7
- {
8
- "role": "input",
9
- "dtype": "T",
10
- "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."
11
- },
12
- {
13
- "role": "position_ids",
14
- "dtype": "M",
15
- "rank": 3,
16
- "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."
17
- },
18
- {
19
- "role": "cos_cache",
20
- "dtype": "T",
21
- "rank": 2,
22
- "description": "Precomputed cosine values of shape `(max_sequence_length, rotary_dim/2)`, shared by all three position streams."
23
- },
24
- {
25
- "role": "sin_cache",
26
- "dtype": "T",
27
- "rank": 2,
28
- "description": "Precomputed sine values with the same shape and type as `cos_cache`."
29
- }
30
- ],
31
- "outputs": [
32
- {
33
- "role": "output",
34
- "dtype": "T",
35
- "rank": "ranks.input",
36
- "shape": "shapes.input",
37
- "description": "Rotary-position-encoded tensor with the same shape and type as `input`."
38
- }
39
- ],
40
  "attributes": {
41
- "interleaved": 0,
42
- "is_packed_batching": 0,
43
- "mrope_layout": 0,
44
- "num_heads": 0,
45
- "rotary_embedding_dim": 0,
46
- "scale": 1
 
47
  },
48
  "attributeConstraints": {
49
  "interleaved": { "values": [0, 1] },
@@ -51,168 +24,68 @@
51
  "mrope_layout": { "values": [0, 1] },
52
  "mrope_section": { "required": true }
53
  },
54
- "attributeDescriptions": {
55
- "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`.",
56
- "is_packed_batching": "Whether `position_ids` uses packed-batch metadata. The default and only supported value is 0; packed batching (1) is not implemented.",
57
- "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.",
58
- "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.",
59
- "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`.",
60
- "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.",
61
- "scale": "Scale applied to the gathered cosine and sine values before the rotation. Default is 1.0."
62
- },
63
  "typeConstraints": { "T": ["float32", "float16"], "M": ["int64"] },
64
- "args": {
65
- "x": { "kind": "tensor", "semantic": "input", "role": "input" },
66
- "positionIds": {
67
- "kind": "tensor",
68
- "semantic": "position_ids",
69
- "role": "input",
70
- "dtype": "uint32",
71
- "narrowing": "checked"
72
- },
73
- "cos": { "kind": "tensor", "semantic": "cos_cache", "role": "input" },
74
- "sin": { "kind": "tensor", "semantic": "sin_cache", "role": "input" },
75
- "y": { "kind": "tensor", "semantic": "output", "role": "output" }
76
- },
77
- "tunables": { "WORKGROUP_SIZE": 256 },
78
  "derive": {
79
- "rank3HeadSize": "dim(shapes.input, 2) / attrs.num_heads if attrs.num_heads is defined and attrs.num_heads > 0 else 0",
80
- "headSize": "rank3HeadSize if ranks.input == 3 else dim(shapes.input, 3)",
81
  "effectiveRotaryDim": "attrs.rotary_embedding_dim if attrs.rotary_embedding_dim != 0 else headSize",
82
  "tasksPerHead": "ceilDiv(headSize, 2)",
83
- "pairCount": "(numel(shapes.input) / max(1, headSize)) * tasksPerHead",
84
- "pairDispatchOk": "ceilDiv(pairCount, tunables.WORKGROUP_SIZE) <= device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension",
85
- "halfRotaryDim": "dim(shapes.cos_cache, 1)",
86
  "sectionsDefined": "attrs.mrope_section is defined and (attrs.mrope_section | length) == 3",
87
  "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",
88
- "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)",
89
- "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",
90
- "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"
91
  },
92
- "bindingSets": {
93
- "$common": [
94
- { "name": "x", "arg": "x", "semantic": "input", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
95
- {
96
- "name": "position_ids",
97
- "arg": "positionIds",
98
- "semantic": "position_ids",
99
- "buffer": { "type": "read-only-storage" },
100
- "elementType": "u32"
101
- },
102
- {
103
- "name": "cos_cache",
104
- "arg": "cos",
105
- "semantic": "cos_cache",
106
- "buffer": { "type": "read-only-storage" },
107
- "elementType": "$T"
108
- },
109
- {
110
- "name": "sin_cache",
111
- "arg": "sin",
112
- "semantic": "sin_cache",
113
- "buffer": { "type": "read-only-storage" },
114
- "elementType": "$T"
115
- },
116
- { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$T" }
117
- ],
118
- "rank3": [
119
- { "name": "x", "arg": "x", "semantic": "input", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
120
- {
121
- "name": "position_ids",
122
- "arg": "positionIds",
123
- "semantic": "position_ids",
124
- "buffer": { "type": "read-only-storage" },
125
- "elementType": "u32"
126
- },
127
- {
128
- "name": "cos_cache",
129
- "arg": "cos",
130
- "semantic": "cos_cache",
131
- "buffer": { "type": "read-only-storage" },
132
- "elementType": "$T"
133
- },
134
- {
135
- "name": "sin_cache",
136
- "arg": "sin",
137
- "semantic": "sin_cache",
138
- "buffer": { "type": "read-only-storage" },
139
- "elementType": "$T"
140
- },
141
- { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$T" },
142
- {
143
- "name": "params",
144
- "semantic": "kernel.params",
145
- "buffer": { "type": "uniform" },
146
- "struct": {
147
- "name": "Params",
148
- "fields": [
149
- { "name": "pairCount", "type": "u32", "value": "pairCount" },
150
- { "name": "batchSize", "type": "u32", "value": "dim(shapes.input, 0)" },
151
- { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.input, 1)" },
152
- { "name": "numHeads", "type": "u32", "value": "attrs.num_heads" },
153
- { "name": "headSize", "type": "u32", "value": "rank3HeadSize" },
154
- { "name": "rotaryDim", "type": "u32", "value": "halfRotaryDim * 2" },
155
- { "name": "halfRotaryDim", "type": "u32", "value": "halfRotaryDim" },
156
- { "name": "section0", "type": "u32", "value": "attrs.mrope_section[0]" },
157
- { "name": "section1", "type": "u32", "value": "attrs.mrope_section[1]" },
158
- { "name": "section2", "type": "u32", "value": "attrs.mrope_section[2]" },
159
- { "name": "scale", "type": "f32", "value": "attrs.scale" }
160
- ]
161
- }
162
- }
163
- ],
164
- "rank4": [
165
- { "name": "x", "arg": "x", "semantic": "input", "buffer": { "type": "read-only-storage" }, "elementType": "$T" },
166
- {
167
- "name": "position_ids",
168
- "arg": "positionIds",
169
- "semantic": "position_ids",
170
- "buffer": { "type": "read-only-storage" },
171
- "elementType": "u32"
172
- },
173
- {
174
- "name": "cos_cache",
175
- "arg": "cos",
176
- "semantic": "cos_cache",
177
- "buffer": { "type": "read-only-storage" },
178
- "elementType": "$T"
179
- },
180
- {
181
- "name": "sin_cache",
182
- "arg": "sin",
183
- "semantic": "sin_cache",
184
- "buffer": { "type": "read-only-storage" },
185
- "elementType": "$T"
186
- },
187
- { "name": "y", "arg": "y", "semantic": "output", "buffer": { "type": "storage" }, "elementType": "$T" },
188
- {
189
- "name": "params",
190
- "semantic": "kernel.params",
191
- "buffer": { "type": "uniform" },
192
- "struct": {
193
- "name": "Params",
194
- "fields": [
195
- { "name": "pairCount", "type": "u32", "value": "pairCount" },
196
- { "name": "batchSize", "type": "u32", "value": "dim(shapes.input, 0)" },
197
- { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.input, 2)" },
198
- { "name": "numHeads", "type": "u32", "value": "dim(shapes.input, 1)" },
199
- { "name": "headSize", "type": "u32", "value": "dim(shapes.input, 3)" },
200
- { "name": "rotaryDim", "type": "u32", "value": "halfRotaryDim * 2" },
201
- { "name": "halfRotaryDim", "type": "u32", "value": "halfRotaryDim" },
202
- { "name": "section0", "type": "u32", "value": "attrs.mrope_section[0]" },
203
- { "name": "section1", "type": "u32", "value": "attrs.mrope_section[1]" },
204
- { "name": "section2", "type": "u32", "value": "attrs.mrope_section[2]" },
205
- { "name": "scale", "type": "f32", "value": "attrs.scale" }
206
- ]
207
- }
208
- }
209
- ]
210
  },
211
  "variants": [
212
  {
213
  "id": "rank3",
214
- "when": ["rank3Contract", "pairDispatchOk"],
215
- "constants": {
216
  "interleaved": "attrs.interleaved != 0",
217
  "mropeSectioned": "attrs.mrope_layout == 0",
218
  "usesF16": "dtypes.T == \"f16\"",
@@ -222,16 +95,21 @@
222
  {
223
  "id": "main",
224
  "name": "mrotary_embedding3d",
225
- "source": { "shader": "mrotary-embedding.wgsl.jinja", "inputs": { "rank": 3 } },
226
- "bindings": "rank3",
227
- "dispatch": { "threads": "pairCount", "workgroupSize": "tunables.WORKGROUP_SIZE" }
 
 
 
 
 
228
  }
229
  ]
230
  },
231
  {
232
  "id": "rank4",
233
- "when": ["rank4Contract", "pairDispatchOk"],
234
- "constants": {
235
  "interleaved": "attrs.interleaved != 0",
236
  "mropeSectioned": "attrs.mrope_layout == 0",
237
  "usesF16": "dtypes.T == \"f16\"",
@@ -241,9 +119,14 @@
241
  {
242
  "id": "main",
243
  "name": "mrotary_embedding4d",
244
- "source": { "shader": "mrotary-embedding.wgsl.jinja", "inputs": { "rank": 4 } },
245
- "bindings": "rank4",
246
- "dispatch": { "threads": "pairCount", "workgroupSize": "tunables.WORKGROUP_SIZE" }
 
 
 
 
 
247
  }
248
  ]
249
  }
 
2
  "domain": "com.microsoft",
3
  "name": "MRotaryEmbedding",
4
  "sinceVersion": 1,
5
+ "inputs": {
6
+ "x": { "onnx": "input", "dtype": "T" },
7
+ "positionIds": { "onnx": "position_ids", "dtype": "M", "rank": 3, "storage": "uint32", "narrowing": "checked" },
8
+ "cos": { "onnx": "cos_cache", "dtype": "T", "rank": 2 },
9
+ "sin": { "onnx": "sin_cache", "dtype": "T", "rank": 2 }
10
+ },
11
+ "outputs": { "y": { "onnx": "output", "dtype": "T", "rank": "ranks.x", "shape": "shapes.x" } },
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
12
  "attributes": {
13
+ "interleaved": { "default": 0 },
14
+ "is_packed_batching": { "default": 0 },
15
+ "mrope_layout": { "default": 0 },
16
+ "num_heads": { "default": 0 },
17
+ "rotary_embedding_dim": { "default": 0 },
18
+ "scale": { "default": 1 },
19
+ "mrope_section": {}
20
  },
21
  "attributeConstraints": {
22
  "interleaved": { "values": [0, 1] },
 
24
  "mrope_layout": { "values": [0, 1] },
25
  "mrope_section": { "required": true }
26
  },
 
 
 
 
 
 
 
 
 
27
  "typeConstraints": { "T": ["float32", "float16"], "M": ["int64"] },
28
+ "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
 
 
 
 
 
 
 
 
 
 
 
 
 
29
  "derive": {
30
+ "rank3HeadSize": "dim(shapes.x, 2) / attrs.num_heads if attrs.num_heads is defined and attrs.num_heads > 0 else 0",
31
+ "headSize": "rank3HeadSize if ranks.x == 3 else dim(shapes.x, 3)",
32
  "effectiveRotaryDim": "attrs.rotary_embedding_dim if attrs.rotary_embedding_dim != 0 else headSize",
33
  "tasksPerHead": "ceilDiv(headSize, 2)",
34
+ "pairCount": "(numel(shapes.x) / max(1, headSize)) * tasksPerHead",
35
+ "pairDispatchOk": "ceilDiv(pairCount, tunables.WORKGROUP_SIZE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
36
+ "halfRotaryDim": "dim(shapes.cos, 1)",
37
  "sectionsDefined": "attrs.mrope_section is defined and (attrs.mrope_section | length) == 3",
38
  "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",
39
+ "commonContract": "f16Ok(dtypes.T) and sameShape(shapes.x, shapes.y) 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, shapes.sin) and ranks.cos == 2 and ranks.sin == 2 and sectionsValid and ranks.positionIds == 3 and dim(shapes.positionIds, 0) == 3 and dim(shapes.positionIds, 1) == dim(shapes.x, 0) and dim(shapes.positionIds, 2) == dim(shapes.x, 1 if ranks.x == 3 else 2)",
40
+ "rank3Contract": "commonContract and ranks.x == 3 and ranks.y == 3 and attrs.num_heads is defined and attrs.num_heads >= 1 and dim(shapes.x, 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",
41
+ "rank4Contract": "commonContract and ranks.x == 4 and ranks.y == 4 and dim(shapes.x, 3) > 0 and effectiveRotaryDim % 2 == 0 and (attrs.rotary_embedding_dim == 0 or attrs.rotary_embedding_dim <= dim(shapes.x, 3)) and halfRotaryDim * 2 == effectiveRotaryDim"
42
  },
43
+ "when": ["pairDispatchOk"],
44
+ "bindings": {
45
+ "x": { "buffer": "read-only-storage", "elementType": "$T" },
46
+ "position_ids": { "arg": "positionIds", "buffer": "read-only-storage", "elementType": "u32" },
47
+ "cos_cache": { "arg": "cos", "buffer": "read-only-storage", "elementType": "$T" },
48
+ "sin_cache": { "arg": "sin", "buffer": "read-only-storage", "elementType": "$T" },
49
+ "y": { "buffer": "storage", "elementType": "$T" },
50
+ "params": {
51
+ "buffer": "uniform",
52
+ "struct": [
53
+ { "name": "pairCount", "type": "u32", "value": "pairCount" },
54
+ { "name": "batchSize", "type": "u32", "value": "dim(shapes.x, 0)" },
55
+ { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.x, 1)" },
56
+ { "name": "numHeads", "type": "u32", "value": "attrs.num_heads" },
57
+ { "name": "headSize", "type": "u32", "value": "rank3HeadSize" },
58
+ { "name": "rotaryDim", "type": "u32", "value": "halfRotaryDim * 2" },
59
+ { "name": "halfRotaryDim", "type": "u32", "value": "halfRotaryDim" },
60
+ { "name": "section0", "type": "u32", "value": "attrs.mrope_section[0]" },
61
+ { "name": "section1", "type": "u32", "value": "attrs.mrope_section[1]" },
62
+ { "name": "section2", "type": "u32", "value": "attrs.mrope_section[2]" },
63
+ { "name": "scale", "type": "f32", "value": "attrs.scale" }
64
+ ]
65
+ },
66
+ "params_2": {
67
+ "name": "params",
68
+ "buffer": "uniform",
69
+ "struct": [
70
+ { "name": "pairCount", "type": "u32", "value": "pairCount" },
71
+ { "name": "batchSize", "type": "u32", "value": "dim(shapes.x, 0)" },
72
+ { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.x, 2)" },
73
+ { "name": "numHeads", "type": "u32", "value": "dim(shapes.x, 1)" },
74
+ { "name": "headSize", "type": "u32", "value": "dim(shapes.x, 3)" },
75
+ { "name": "rotaryDim", "type": "u32", "value": "halfRotaryDim * 2" },
76
+ { "name": "halfRotaryDim", "type": "u32", "value": "halfRotaryDim" },
77
+ { "name": "section0", "type": "u32", "value": "attrs.mrope_section[0]" },
78
+ { "name": "section1", "type": "u32", "value": "attrs.mrope_section[1]" },
79
+ { "name": "section2", "type": "u32", "value": "attrs.mrope_section[2]" },
80
+ { "name": "scale", "type": "f32", "value": "attrs.scale" }
81
+ ]
82
+ }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
83
  },
84
  "variants": [
85
  {
86
  "id": "rank3",
87
+ "when": ["rank3Contract"],
88
+ "derive": {
89
  "interleaved": "attrs.interleaved != 0",
90
  "mropeSectioned": "attrs.mrope_layout == 0",
91
  "usesF16": "dtypes.T == \"f16\"",
 
95
  {
96
  "id": "main",
97
  "name": "mrotary_embedding3d",
98
+ "shader": "mrotary-embedding.wgsl.jinja",
99
+ "derive": { "rank": 3 },
100
+ "bindings": ["x", "position_ids", "cos_cache", "sin_cache", "y", "params"],
101
+ "dispatch": {
102
+ "x": "min(ceilDiv((pairCount), (tunables.WORKGROUP_SIZE)), 65535)",
103
+ "y": "ceilDiv(ceilDiv((pairCount), (tunables.WORKGROUP_SIZE)), 65535)",
104
+ "z": 1
105
+ }
106
  }
107
  ]
108
  },
109
  {
110
  "id": "rank4",
111
+ "when": ["rank4Contract"],
112
+ "derive": {
113
  "interleaved": "attrs.interleaved != 0",
114
  "mropeSectioned": "attrs.mrope_layout == 0",
115
  "usesF16": "dtypes.T == \"f16\"",
 
119
  {
120
  "id": "main",
121
  "name": "mrotary_embedding4d",
122
+ "shader": "mrotary-embedding.wgsl.jinja",
123
+ "derive": { "rank": 4 },
124
+ "bindings": ["x", "position_ids", "cos_cache", "sin_cache", "y", "params_2"],
125
+ "dispatch": {
126
+ "x": "min(ceilDiv((pairCount), (tunables.WORKGROUP_SIZE)), 65535)",
127
+ "y": "ceilDiv(ceilDiv((pairCount), (tunables.WORKGROUP_SIZE)), 65535)",
128
+ "z": 1
129
+ }
130
  }
131
  ]
132
  }
build/webgpu/metadata.json CHANGED
@@ -1,18 +1,21 @@
1
  {
2
  "name": "com.microsoft.MRotaryEmbedding",
3
- "id": "_com_microsoft_mrotaryembedding_webgpu_396a9ed",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "FKPQK85F0NG70BrG2nvS4WDZWmkLLUtw0hjkaNqyfWI=",
11
- "manifest.json": "rNhF7mp6GzPe1XKu5djfgpCv5MKyJdPU7G06AvbgLXE=",
12
- "mrotary-embedding.wgsl.jinja": "ywEaYPPZeG6tSo/Ulh3lIOBwXLhdu70TAZJAgus88cA=",
13
- "test.json": "0mcf/8uDq9fTEq082qLPOGoRnjowxCk+vqzBBFVA0Ls="
14
  }
15
  },
16
- "provenance": { "kernel": { "sha": "c928d21e6cc1310861cba3bafb75f5f679ecf5f3", "dirty": false } },
17
- "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.MRotaryEmbedding" }
 
 
 
18
  }
 
1
  {
2
  "name": "com.microsoft.MRotaryEmbedding",
3
+ "id": "_com_microsoft_mrotaryembedding_webgpu_83feafb",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "46hHXwN4zErrNbI5Kkvb00wSEt7LAjuixrnpiOUnuC8=",
11
+ "manifest.json": "n2gNx3yadn8wpDXYWYTOGWFQVAiDogbtn03dXyQEAfE=",
12
+ "mrotary-embedding.wgsl.jinja": "meXniZyv3KTe/4K0hd7HSYDjxS0ADty/sNevs/l3es8=",
13
+ "test.json": "mQy2yFi3T6iEpqSliA+KGAipEQeV8moe0jnzSwH9ICI="
14
  }
15
  },
16
+ "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
17
+ "webgpu": {
18
+ "manifestSpec": "2.0",
19
+ "variants": { "rank3": ["mrotary-embedding.wgsl.jinja"], "rank4": ["mrotary-embedding.wgsl.jinja"] }
20
+ }
21
  }
build/webgpu/mrotary-embedding.wgsl.jinja CHANGED
@@ -1,30 +1,30 @@
1
  {% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
2
  {% if note == "dispatch-limit" %}
3
- // 2D-folded flat index: gid.y carries the high bits past the
4
- // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
5
  {% elif note == "limit" %}
6
- // 2D-folded flat index: gid.y carries the high bits past the
7
- // maxComputeWorkgroupsPerDimension limit.
8
  {% elif note == "device-axis" %}
9
- // The flat dispatch is folded across x/y at the device's per-axis workgroup
10
- // limit; gid.y carries the high portion of the output index.
11
  {% elif note == "vec4-limit" %}
12
- // 2D-folded flat vec4 index: gid.y carries the high bits past the
13
- // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills into y).
14
  {% elif note == "element-limit" %}
15
  // 2D-folded flat element index: gid.y carries the high bits past the
16
- // maxComputeWorkgroupsPerDimension limit.
17
  {% elif note == "dispatch" %}
18
- // 2D-folded flat index: gid.y carries the high bits past the
19
- // maxComputeWorkgroupsPerDimension dispatch limit.
20
  {% endif %}
21
  {% if bound == "" %}
22
- let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
23
  {%- elif guardInline %}
24
- let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
25
  if ({{ name }} >= {{ bound }}) { return; }
26
  {%- else %}
27
- let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
28
  if ({{ name }} >= {{ bound }}) {
29
  return;
30
  }
@@ -39,7 +39,7 @@
39
  // owns both outputs of a rotated pair, so x/cos/sin are read once; partial rotations
40
  // use the remaining pair lanes to copy two unchanged tail values.
41
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
42
- fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
43
  {{ flat_index_2d("p", "params.pairCount", note="") }}
44
 
45
  let tasksPerHead = (params.headSize + 1u) / 2u;
@@ -60,7 +60,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups)
60
  // under both rotation pairings: interleaved reads column channel/2 and the split
61
  // form reads channel % halfRotaryDim, and this lane owns channel == pair in each.
62
  let pair = task;
63
- {% if source.rank == 3 %}
64
  let token = (headFlat / params.numHeads) % params.sequenceLength;
65
  let batch = headFlat / (params.numHeads * params.sequenceLength);
66
  {% else %}
 
1
  {% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
2
  {% if note == "dispatch-limit" %}
3
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
4
+ // per-axis workgroup fold width (outputs > 16.7M elements).
5
  {% elif note == "limit" %}
6
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
7
+ // per-axis workgroup fold width.
8
  {% elif note == "device-axis" %}
9
+ // The flat dispatch is folded across x/y at a fixed per-axis workgroup
10
+ // width; gid.y carries the high portion of the output index.
11
  {% elif note == "vec4-limit" %}
12
+ // 2D-folded flat vec4 index: gid.y carries the high bits past the dispatch's
13
+ // per-axis workgroup fold width (the dispatch caps x and spills into y).
14
  {% elif note == "element-limit" %}
15
  // 2D-folded flat element index: gid.y carries the high bits past the
16
+ // dispatch's per-axis workgroup fold width.
17
  {% elif note == "dispatch" %}
18
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
19
+ // per-axis workgroup fold width.
20
  {% endif %}
21
  {% if bound == "" %}
22
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
23
  {%- elif guardInline %}
24
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
25
  if ({{ name }} >= {{ bound }}) { return; }
26
  {%- else %}
27
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
28
  if ({{ name }} >= {{ bound }}) {
29
  return;
30
  }
 
39
  // owns both outputs of a rotated pair, so x/cos/sin are read once; partial rotations
40
  // use the remaining pair lanes to copy two unchanged tail values.
41
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
42
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
43
  {{ flat_index_2d("p", "params.pairCount", note="") }}
44
 
45
  let tasksPerHead = (params.headSize + 1u) / 2u;
 
60
  // under both rotation pairings: interleaved reads column channel/2 and the split
61
  // form reads channel % halfRotaryDim, and this lane owns channel == pair in each.
62
  let pair = task;
63
+ {% if rank == 3 %}
64
  let token = (headFlat / params.numHeads) % params.sequenceLength;
65
  let batch = headFlat / (params.numHeads * params.sequenceLength);
66
  {% else %}
build/webgpu/test.json CHANGED
@@ -1,5 +1,4 @@
1
  {
2
- "op": "com.microsoft.MRotaryEmbedding",
3
  "fixtureArrays": {
4
  "ort_sectioned_rank3_input_x": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]
5
  },
@@ -8,7 +7,7 @@
8
  "name": "rank4_odd_head_even_partial_rotation",
9
  "provenance": {
10
  "source": "onnxruntime/contrib_ops/cuda/bert/mrotary_embedding_impl.cu",
11
- "notes": "The provider requires the effective rotary dimension, not the full head size, to be even. The singleton tail element is copied unchanged."
12
  },
13
  "attrs": {
14
  "num_heads": 1,
@@ -198,7 +197,7 @@
198
  {
199
  "name": "f16_rank4_layout1_interleaved",
200
  "provenance": {
201
- "notes": "f16 storage, including the provider's T-rounded scaled-cache intermediate, on the interleaved layout and pairing."
202
  },
203
  "attrs": {
204
  "rotary_embedding_dim": 0,
 
1
  {
 
2
  "fixtureArrays": {
3
  "ort_sectioned_rank3_input_x": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24]
4
  },
 
7
  "name": "rank4_odd_head_even_partial_rotation",
8
  "provenance": {
9
  "source": "onnxruntime/contrib_ops/cuda/bert/mrotary_embedding_impl.cu",
10
+ "notes": "The effective rotary dimension is even while the full head size is odd, leaving one unrotated tail element that must be copied unchanged."
11
  },
12
  "attrs": {
13
  "num_heads": 1,
 
197
  {
198
  "name": "f16_rank4_layout1_interleaved",
199
  "provenance": {
200
+ "notes": "The interleaved layout uses float16 storage, including float16 rounding of the scaled cache intermediate before pairing."
201
  },
202
  "attrs": {
203
  "rotary_embedding_dim": 0,