Hugging Face's logo

Kernels:
webgpu-kernels
/
com.microsoft.GemmaRotaryEmbedding

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
1092006
·
verified ·
1 Parent(s): cc4bd23

sync 91d990483a17

Browse files
Files changed (6) hide show
  1. README.md +16 -12
  2. build/webgpu/bench.json +0 -1
  3. build/webgpu/gemma-rotary-embedding.wgsl.jinja +3 -4
  4. build/webgpu/manifest.json +61 -159
  5. build/webgpu/metadata.json +10 -7
  6. build/webgpu/test.json +76 -3
README.md CHANGED
@@ -18,20 +18,20 @@ See the [ONNX Runtime `GemmaRotaryEmbedding` contrib-operator spec](https://gith
18
 
19
  ## Inputs
20
 
21
- | Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
22
  | --- | --- | --- | --- | --- | --- | --- | --- |
23
- | `emb` | `embT` | `U` | `float32` | `3` | — | Rotary angles with shape `(batch_size, seq_len, dim)`, shared by every head. | required |
24
- | `q` | `qT` | `T` | same as logical dtype | `4` | — | Query state with shape `(batch_size, num_heads, seq_len, dim)`. | required |
25
- | `q_rot` | `qRotT` | `T` | same as logical dtype | `4` | — | Half-rotated query state, same shape as `q`. | required |
26
- | `k` | `kT` | `T` | same as logical dtype | `4` | — | Key state, same shape as `q`. | required |
27
- | `k_rot` | `kRotT` | `T` | same as logical dtype | `4` | — | Half-rotated key state, same shape as `q`. | required |
28
 
29
  ## Outputs
30
 
31
- | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
32
  | --- | --- | --- | --- | --- | --- | --- |
33
- | `output1` | `output1T` | `T` | same as `q` | same as `q` | Rotary-embedded query, same shape as `q`. | required |
34
- | `output2` | `output2T` | `T` | same as `q` | same as `q` | Rotary-embedded key, same shape as `q`. | required |
35
 
36
  ## Type constraints
37
 
@@ -46,7 +46,7 @@ Every implementation variant requires `shader-f16`; the package has no variant-l
46
 
47
  ## Files
48
 
49
- - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
50
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
51
  - [`test.json`](build/webgpu/test.json) — correctness cases
52
  - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -54,10 +54,14 @@ Every implementation variant requires `shader-f16`; the package has no variant-l
54
 
55
  ## Use with `@huggingface/kernels`
56
 
57
- The loader derives every required output's shape and logical dtype from the manifest contract and this call.
58
- It then allocates the result tensors automatically.
 
 
 
59
 
60
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 
61
 
62
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
63
 
 
18
 
19
  ## Inputs
20
 
21
+ | Name | Upstream name | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
22
  | --- | --- | --- | --- | --- | --- | --- | --- |
23
+ | `embT` | `emb` | `U` | `float32` | `3` | — | Rotary angles with shape `(batch_size, seq_len, dim)`, shared by every head. | required |
24
+ | `qT` | `q` | `T` | same as logical dtype | `4` | — | Query state with shape `(batch_size, num_heads, seq_len, dim)`. | required |
25
+ | `qRotT` | `q_rot` | `T` | same as logical dtype | `4` | — | Half-rotated query state, same shape as `q`. | required |
26
+ | `kT` | `k` | `T` | same as logical dtype | `4` | — | Key state, same shape as `q`. | required |
27
+ | `kRotT` | `k_rot` | `T` | same as logical dtype | `4` | — | Half-rotated key state, same shape as `q`. | required |
28
 
29
  ## Outputs
30
 
31
+ | Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
32
  | --- | --- | --- | --- | --- | --- | --- |
33
+ | `output1T` | `output1` | `T` | same as `qT` | same as `qT` | Rotary-embedded query, same shape as `q`. | required |
34
+ | `output2T` | `output2` | `T` | same as `qT` | same as `qT` | Rotary-embedded key, same shape as `q`. | required |
35
 
36
  ## Type constraints
37
 
 
46
 
47
  ## Files
48
 
49
+ - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
50
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
51
  - [`test.json`](build/webgpu/test.json) — correctness cases
52
  - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 
54
 
55
  ## Use with `@huggingface/kernels`
56
 
57
+ ```sh
58
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.2
59
+ ```
60
+
61
+ Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
62
 
63
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
64
+ It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
65
 
66
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
67
 
build/webgpu/bench.json CHANGED
@@ -1,5 +1,4 @@
1
  {
2
- "op": "com.microsoft.GemmaRotaryEmbedding",
3
  "tunableSpace": { "workgroupSize": [64, 128, 256] },
4
  "cases": [
5
  {
 
1
  {
 
2
  "tunableSpace": { "workgroupSize": [64, 128, 256] },
3
  "cases": [
4
  {
build/webgpu/gemma-rotary-embedding.wgsl.jinja CHANGED
@@ -12,11 +12,10 @@ const ZERO: {{ scalar }} = {{ scalar }}(0.0);
12
  {% endif %}
13
 
14
  @compute @workgroup_size(WG, 1, 1)
15
- fn main(@builtin(global_invocation_id) gid: vec3<u32>,
16
- @builtin(num_workgroups) nwg: vec3<u32>) {
17
  // 2D-folded flat index: gid.y carries the high bits past the
18
- // maxComputeWorkgroupsPerDimension limit. Reduces to gid.x when nwg.y == 1.
19
- let index = gid.x + gid.y * nwg.x * WG;
20
  if (index >= params.count) {
21
  return;
22
  }
 
12
  {% endif %}
13
 
14
  @compute @workgroup_size(WG, 1, 1)
15
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
16
  // 2D-folded flat index: gid.y carries the high bits past the
17
+ // per-axis dispatch fold width. Reduces to gid.x when the dispatch does not fold.
18
+ let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
19
  if (index >= params.count) {
20
  return;
21
  }
build/webgpu/manifest.json CHANGED
@@ -2,51 +2,19 @@
2
  "domain": "com.microsoft",
3
  "name": "GemmaRotaryEmbedding",
4
  "sinceVersion": 1,
5
- "description": "Fuses the Gemma rotary-embedding tail: computes `sin` and `cos` from float32 `emb`, casts them to float16, then evaluates `q * cos + q_rot * sin` and the corresponding expression for `k`. `emb` has shape `(batch, seq, dim)` and is broadcast over the head axis of the `(batch, heads, seq, dim)` operands. Each product is rounded to float16 before the addition.",
6
- "inputs": [
7
- {
8
- "role": "emb",
9
- "dtype": "U",
10
- "rank": 3,
11
- "description": "Rotary angles with shape `(batch_size, seq_len, dim)`, shared by every head."
12
- },
13
- {
14
- "role": "q",
15
- "dtype": "T",
16
- "rank": 4,
17
- "description": "Query state with shape `(batch_size, num_heads, seq_len, dim)`."
18
- },
19
- { "role": "q_rot", "dtype": "T", "rank": 4, "description": "Half-rotated query state, same shape as `q`." },
20
- { "role": "k", "dtype": "T", "rank": 4, "description": "Key state, same shape as `q`." },
21
- { "role": "k_rot", "dtype": "T", "rank": 4, "description": "Half-rotated key state, same shape as `q`." }
22
- ],
23
- "outputs": [
24
- {
25
- "role": "output1",
26
- "dtype": "T",
27
- "rank": "ranks.qT",
28
- "shape": "shapes.qT",
29
- "description": "Rotary-embedded query, same shape as `q`."
30
- },
31
- {
32
- "role": "output2",
33
- "dtype": "T",
34
- "rank": "ranks.qT",
35
- "shape": "shapes.qT",
36
- "description": "Rotary-embedded key, same shape as `q`."
37
- }
38
- ],
39
- "typeConstraints": { "T": ["float16"], "U": ["float32"] },
40
- "args": {
41
- "embT": { "kind": "tensor", "semantic": "emb", "role": "input", "dtype": "float32" },
42
- "qT": { "kind": "tensor", "semantic": "q", "role": "input" },
43
- "qRotT": { "kind": "tensor", "semantic": "q_rot", "role": "input" },
44
- "kT": { "kind": "tensor", "semantic": "k", "role": "input" },
45
- "kRotT": { "kind": "tensor", "semantic": "k_rot", "role": "input" },
46
- "output1T": { "kind": "tensor", "semantic": "output1", "role": "output" },
47
- "output2T": { "kind": "tensor", "semantic": "output2", "role": "output" }
48
  },
49
- "tunables": { "workgroupSize": 256 },
 
 
 
 
 
50
  "derive": {
51
  "batchSize": "dim(shapes.qT, 0)",
52
  "numHeads": "dim(shapes.qT, 1)",
@@ -55,121 +23,16 @@
55
  "sameAsQ": "ranks.qRotT == 4 and ranks.kT == 4 and ranks.kRotT == 4 and sameShape(shapes.qRotT, shapes.qT) and sameShape(shapes.kT, shapes.qT) and sameShape(shapes.kRotT, shapes.qT)",
56
  "dtypesMatch": "tensorDtypes.qRotT == tensorDtypes.qT and tensorDtypes.kT == tensorDtypes.qT and tensorDtypes.kRotT == tensorDtypes.qT and tensorDtypes.output1T == tensorDtypes.qT and tensorDtypes.output2T == tensorDtypes.qT",
57
  "contract": "ranks.qT == 4 and ranks.embT == 3 and sameAsQ and dtypesMatch and tensorDtypes.qT == \"float16\" and f16Ok(dtypes.T) and tensorDtypes.embT == \"float32\" and dim(shapes.embT, 0) == batchSize and dim(shapes.embT, 1) == seqLen and dim(shapes.embT, 2) == headDim and sameShape(shapes.output1T, shapes.qT) and sameShape(shapes.output2T, shapes.qT)",
58
- "vec4Ok": "headDim % 4 == 0"
59
- },
60
- "bindingSets": {
61
- "scalar": [
62
- {
63
- "name": "emb",
64
- "arg": "embT",
65
- "semantic": "emb",
66
- "buffer": { "type": "read-only-storage" },
67
- "elementType": "f32"
68
- },
69
- { "name": "q", "arg": "qT", "semantic": "q", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
70
- {
71
- "name": "q_rot",
72
- "arg": "qRotT",
73
- "semantic": "q_rot",
74
- "buffer": { "type": "read-only-storage" },
75
- "elementType": "$scalar"
76
- },
77
- { "name": "k", "arg": "kT", "semantic": "k", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
78
- {
79
- "name": "k_rot",
80
- "arg": "kRotT",
81
- "semantic": "k_rot",
82
- "buffer": { "type": "read-only-storage" },
83
- "elementType": "$scalar"
84
- },
85
- {
86
- "name": "output1",
87
- "arg": "output1T",
88
- "semantic": "output1",
89
- "buffer": { "type": "storage" },
90
- "elementType": "$scalar"
91
- },
92
- {
93
- "name": "output2",
94
- "arg": "output2T",
95
- "semantic": "output2",
96
- "buffer": { "type": "storage" },
97
- "elementType": "$scalar"
98
- },
99
- {
100
- "name": "params",
101
- "semantic": "kernel.params",
102
- "buffer": { "type": "uniform" },
103
- "struct": {
104
- "name": "Params",
105
- "fields": [
106
- { "name": "count", "type": "u32", "value": "numel(shapes.qT)" },
107
- { "name": "seqDim", "type": "u32", "value": "(seqLen * headDim)" },
108
- { "name": "headSeqDim", "type": "u32", "value": "(numHeads * seqLen * headDim)" }
109
- ]
110
- }
111
- }
112
- ],
113
- "vec4": [
114
- {
115
- "name": "emb",
116
- "arg": "embT",
117
- "semantic": "emb",
118
- "buffer": { "type": "read-only-storage" },
119
- "elementType": "vec4<f32>"
120
- },
121
- { "name": "q", "arg": "qT", "semantic": "q", "buffer": { "type": "read-only-storage" }, "elementType": "$vector" },
122
- {
123
- "name": "q_rot",
124
- "arg": "qRotT",
125
- "semantic": "q_rot",
126
- "buffer": { "type": "read-only-storage" },
127
- "elementType": "$vector"
128
- },
129
- { "name": "k", "arg": "kT", "semantic": "k", "buffer": { "type": "read-only-storage" }, "elementType": "$vector" },
130
- {
131
- "name": "k_rot",
132
- "arg": "kRotT",
133
- "semantic": "k_rot",
134
- "buffer": { "type": "read-only-storage" },
135
- "elementType": "$vector"
136
- },
137
- {
138
- "name": "output1",
139
- "arg": "output1T",
140
- "semantic": "output1",
141
- "buffer": { "type": "storage" },
142
- "elementType": "$vector"
143
- },
144
- {
145
- "name": "output2",
146
- "arg": "output2T",
147
- "semantic": "output2",
148
- "buffer": { "type": "storage" },
149
- "elementType": "$vector"
150
- },
151
- {
152
- "name": "params",
153
- "semantic": "kernel.params",
154
- "buffer": { "type": "uniform" },
155
- "struct": {
156
- "name": "Params",
157
- "fields": [
158
- { "name": "count", "type": "u32", "value": "numel(shapes.qT) / 4" },
159
- { "name": "seqDim", "type": "u32", "value": "(seqLen * headDim) / 4" },
160
- { "name": "headSeqDim", "type": "u32", "value": "(numHeads * seqLen * headDim) / 4" }
161
- ]
162
- }
163
- }
164
- ]
165
  },
 
166
  "variants": [
167
  {
168
  "id": "vec4",
169
  "priority": 10,
 
170
  "requires": { "features": ["shader-f16"] },
171
- "when": ["contract", "vec4Ok", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
172
- "constants": {
173
  "vec4": true,
174
  "scalar": "dtypes.T",
175
  "vector": "\"vec4<f16>\"",
@@ -180,8 +43,28 @@
180
  "id": "main",
181
  "name": "GemmaRotaryEmbedding.Vec4",
182
  "shader": "gemma-rotary-embedding.wgsl.jinja",
183
- "bindings": "vec4",
184
- "dispatch": { "threads": "numel(shapes.qT) / 4", "workgroupSize": "constants.workgroupSize" }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
185
  }
186
  ]
187
  },
@@ -189,8 +72,7 @@
189
  "id": "scalar",
190
  "priority": 0,
191
  "requires": { "features": ["shader-f16"] },
192
- "when": ["contract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
193
- "constants": {
194
  "vec4": false,
195
  "scalar": "dtypes.T",
196
  "vector": "\"vec4<f16>\"",
@@ -201,8 +83,28 @@
201
  "id": "main",
202
  "name": "GemmaRotaryEmbedding",
203
  "shader": "gemma-rotary-embedding.wgsl.jinja",
204
- "bindings": "scalar",
205
- "dispatch": { "threads": "max(1, numel(shapes.qT))", "workgroupSize": "constants.workgroupSize" }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
206
  }
207
  ]
208
  }
 
2
  "domain": "com.microsoft",
3
  "name": "GemmaRotaryEmbedding",
4
  "sinceVersion": 1,
5
+ "inputs": {
6
+ "embT": { "onnx": "emb", "dtype": "U", "rank": 3, "storage": "float32" },
7
+ "qT": { "onnx": "q", "dtype": "T", "rank": 4 },
8
+ "qRotT": { "onnx": "q_rot", "dtype": "T", "rank": 4 },
9
+ "kT": { "onnx": "k", "dtype": "T", "rank": 4 },
10
+ "kRotT": { "onnx": "k_rot", "dtype": "T", "rank": 4 }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
11
  },
12
+ "outputs": {
13
+ "output1T": { "onnx": "output1", "dtype": "T", "rank": "ranks.qT", "shape": "shapes.qT" },
14
+ "output2T": { "onnx": "output2", "dtype": "T", "rank": "ranks.qT", "shape": "shapes.qT" }
15
+ },
16
+ "typeConstraints": { "T": ["float16"], "U": ["float32"] },
17
+ "tunables": { "workgroupSize": { "default": 256 } },
18
  "derive": {
19
  "batchSize": "dim(shapes.qT, 0)",
20
  "numHeads": "dim(shapes.qT, 1)",
 
23
  "sameAsQ": "ranks.qRotT == 4 and ranks.kT == 4 and ranks.kRotT == 4 and sameShape(shapes.qRotT, shapes.qT) and sameShape(shapes.kT, shapes.qT) and sameShape(shapes.kRotT, shapes.qT)",
24
  "dtypesMatch": "tensorDtypes.qRotT == tensorDtypes.qT and tensorDtypes.kT == tensorDtypes.qT and tensorDtypes.kRotT == tensorDtypes.qT and tensorDtypes.output1T == tensorDtypes.qT and tensorDtypes.output2T == tensorDtypes.qT",
25
  "contract": "ranks.qT == 4 and ranks.embT == 3 and sameAsQ and dtypesMatch and tensorDtypes.qT == \"float16\" and f16Ok(dtypes.T) and tensorDtypes.embT == \"float32\" and dim(shapes.embT, 0) == batchSize and dim(shapes.embT, 1) == seqLen and dim(shapes.embT, 2) == headDim and sameShape(shapes.output1T, shapes.qT) and sameShape(shapes.output2T, shapes.qT)",
26
+ "vec4Ok": "(seqLen * headDim) % 4 == 0"
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
27
  },
28
+ "when": ["contract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
29
  "variants": [
30
  {
31
  "id": "vec4",
32
  "priority": 10,
33
+ "when": ["vec4Ok"],
34
  "requires": { "features": ["shader-f16"] },
35
+ "derive": {
 
36
  "vec4": true,
37
  "scalar": "dtypes.T",
38
  "vector": "\"vec4<f16>\"",
 
43
  "id": "main",
44
  "name": "GemmaRotaryEmbedding.Vec4",
45
  "shader": "gemma-rotary-embedding.wgsl.jinja",
46
+ "bindings": [
47
+ { "arg": "embT", "name": "emb", "elementType": "vec4<f32>" },
48
+ { "arg": "qT", "name": "q", "elementType": "$vector" },
49
+ { "arg": "qRotT", "name": "q_rot", "elementType": "$vector" },
50
+ { "arg": "kT", "name": "k", "elementType": "$vector" },
51
+ { "arg": "kRotT", "name": "k_rot", "elementType": "$vector" },
52
+ { "arg": "output1T", "name": "output1", "elementType": "$vector" },
53
+ { "arg": "output2T", "name": "output2", "elementType": "$vector" },
54
+ {
55
+ "name": "params",
56
+ "struct": [
57
+ { "name": "count", "type": "u32", "value": "numel(shapes.qT) / 4" },
58
+ { "name": "seqDim", "type": "u32", "value": "(seqLen * headDim) / 4" },
59
+ { "name": "headSeqDim", "type": "u32", "value": "(numHeads * seqLen * headDim) / 4" }
60
+ ]
61
+ }
62
+ ],
63
+ "dispatch": {
64
+ "x": "min(ceilDiv((numel(shapes.qT) / 4), (workgroupSize)), 65535)",
65
+ "y": "ceilDiv(ceilDiv((numel(shapes.qT) / 4), (workgroupSize)), 65535)",
66
+ "z": 1
67
+ }
68
  }
69
  ]
70
  },
 
72
  "id": "scalar",
73
  "priority": 0,
74
  "requires": { "features": ["shader-f16"] },
75
+ "derive": {
 
76
  "vec4": false,
77
  "scalar": "dtypes.T",
78
  "vector": "\"vec4<f16>\"",
 
83
  "id": "main",
84
  "name": "GemmaRotaryEmbedding",
85
  "shader": "gemma-rotary-embedding.wgsl.jinja",
86
+ "bindings": [
87
+ { "arg": "embT", "name": "emb" },
88
+ { "arg": "qT", "name": "q", "elementType": "$scalar" },
89
+ { "arg": "qRotT", "name": "q_rot", "elementType": "$scalar" },
90
+ { "arg": "kT", "name": "k", "elementType": "$scalar" },
91
+ { "arg": "kRotT", "name": "k_rot", "elementType": "$scalar" },
92
+ { "arg": "output1T", "name": "output1", "elementType": "$scalar" },
93
+ { "arg": "output2T", "name": "output2", "elementType": "$scalar" },
94
+ {
95
+ "name": "params",
96
+ "struct": [
97
+ { "name": "count", "type": "u32", "value": "numel(shapes.qT)" },
98
+ { "name": "seqDim", "type": "u32", "value": "(seqLen * headDim)" },
99
+ { "name": "headSeqDim", "type": "u32", "value": "(numHeads * seqLen * headDim)" }
100
+ ]
101
+ }
102
+ ],
103
+ "dispatch": {
104
+ "x": "min(ceilDiv((max(1, numel(shapes.qT))), (workgroupSize)), 65535)",
105
+ "y": "ceilDiv(ceilDiv((max(1, numel(shapes.qT))), (workgroupSize)), 65535)",
106
+ "z": 1
107
+ }
108
  }
109
  ]
110
  }
build/webgpu/metadata.json CHANGED
@@ -1,18 +1,21 @@
1
  {
2
  "name": "com.microsoft.GemmaRotaryEmbedding",
3
- "id": "_com_microsoft_gemmarotaryembedding_webgpu_c64044f",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "/3JoilqqNUy3rMDPsbi0npsxk9oZgl6Db/2k+439H0c=",
11
- "gemma-rotary-embedding.wgsl.jinja": "w+t56Q+FUBau7k0eWwrulSaGWFEKYch4aCunpG3sbBs=",
12
- "manifest.json": "FK8aZrRWKnUBl9hUi7MrNAX7nE9/6pz89Wh/AL5kCDQ=",
13
- "test.json": "KnP1fOvNgEvi3e9lrItvga2K6sLswyE5trl1BhX72Js="
14
  }
15
  },
16
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
17
- "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.GemmaRotaryEmbedding" }
 
 
 
18
  }
 
1
  {
2
  "name": "com.microsoft.GemmaRotaryEmbedding",
3
+ "id": "_com_microsoft_gemmarotaryembedding_webgpu_e5a6aca",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "E1026CMMv9Pppl2+JpGoj3dWCSmk8tZUEwBYo3UCiD4=",
11
+ "gemma-rotary-embedding.wgsl.jinja": "t4wKCJ/E3156iIqQTq1d2XGqMBA83kI0EAADYck7NI0=",
12
+ "manifest.json": "41x/WJ1/AyRVHsYgM3nIqECf6Z8zix9MBzVgLIaLILE=",
13
+ "test.json": "WwfG+ffVK8tAHKvZCOkixu97BOEIQjEdR4NSOSLA9dc="
14
  }
15
  },
16
+ "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
17
+ "webgpu": {
18
+ "manifestSpec": "2.0",
19
+ "variants": { "vec4": ["gemma-rotary-embedding.wgsl.jinja"], "scalar": ["gemma-rotary-embedding.wgsl.jinja"] }
20
+ }
21
  }
build/webgpu/test.json CHANGED
@@ -1,10 +1,9 @@
1
  {
2
- "op": "com.microsoft.GemmaRotaryEmbedding",
3
  "cases": [
4
  {
5
  "name": "f16_vec4",
6
  "provenance": {
7
- "notes": "float16 operands with a float32 emb, which is the type pair ONNX Runtime constrains this operator to. Sin and cos are computed at float32 and rounded once."
8
  },
9
  "inputs": {
10
  "embT": {
@@ -134,7 +133,7 @@
134
  {
135
  "name": "pinned_batch2_emb_broadcast",
136
  "provenance": {
137
- "notes": "Hand-computed by an independent model that Unsqueezes emb to (batch, 1, seq, dim) and broadcasts with real array shapes, never forming a flat index -- which is the step ONNX Runtime's own unit test gets wrong: it indexes emb by num_heads where the CUDA kernel uses seq_len, and its only case has batch_size 1, where the wrong term is multiplied by zero. heads (3) and seq_len (2) differ here and batch is 2, so the two disagree."
138
  },
139
  "inputs": {
140
  "embT": {
@@ -198,6 +197,80 @@
198
  "tolerance": 0.005
199
  }
200
  }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
201
  }
202
  ]
203
  }
 
1
  {
 
2
  "cases": [
3
  {
4
  "name": "f16_vec4",
5
  "provenance": {
6
+ "notes": "Float16 operands with float32 emb. Sin and cos are computed in float32 and rounded once."
7
  },
8
  "inputs": {
9
  "embT": {
 
133
  {
134
  "name": "pinned_batch2_emb_broadcast",
135
  "provenance": {
136
+ "notes": "Expected values broadcast `emb` as `(batch, 1, sequence, dimension)`. Batch 2, three heads, and sequence length 2 exercise sequence-based embedding indexing independently of the head index."
137
  },
138
  "inputs": {
139
  "embT": {
 
197
  "tolerance": 0.005
198
  }
199
  }
200
+ },
201
+ {
202
+ "name": "f16_vec4_dim6_seq4_straddles_heads",
203
+ "provenance": {
204
+ "notes": "headDim 6 is not a multiple of four but seqLen * headDim = 24 is: the vec4 route admits it and its vec4s straddle head-dim boundaries within one (batch, seq) block."
205
+ },
206
+ "inputs": {
207
+ "embT": {
208
+ "dtype": "float32",
209
+ "shape": [1, 4, 6],
210
+ "data": { "kind": "fillFloat32", "scale": 2.5, "sinStep": 0.25, "cosStep": 0.38 }
211
+ },
212
+ "qT": {
213
+ "dtype": "float16",
214
+ "shape": [1, 2, 4, 6],
215
+ "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.19, "cosStep": 0.3 }
216
+ },
217
+ "qRotT": {
218
+ "dtype": "float16",
219
+ "shape": [1, 2, 4, 6],
220
+ "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.45, "cosStep": 0.14 }
221
+ },
222
+ "kT": {
223
+ "dtype": "float16",
224
+ "shape": [1, 2, 4, 6],
225
+ "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.37, "cosStep": 0.2 }
226
+ },
227
+ "kRotT": {
228
+ "dtype": "float16",
229
+ "shape": [1, 2, 4, 6],
230
+ "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.49, "cosStep": 0.26 }
231
+ }
232
+ },
233
+ "outputs": {
234
+ "output1T": { "dtype": "float16", "shape": [1, 2, 4, 6], "tolerance": 0.005 },
235
+ "output2T": { "dtype": "float16", "shape": [1, 2, 4, 6], "tolerance": 0.005 }
236
+ }
237
+ },
238
+ {
239
+ "name": "f16_vec4_dim5_seq4_odd_headdim",
240
+ "provenance": {
241
+ "notes": "An odd headDim (5) with seqLen 4: seqLen * headDim = 20 keeps every vec4 of emb angles and of q/k inside one (batch, head) block."
242
+ },
243
+ "inputs": {
244
+ "embT": {
245
+ "dtype": "float32",
246
+ "shape": [2, 4, 5],
247
+ "data": { "kind": "fillFloat32", "scale": 2.5, "sinStep": 0.26, "cosStep": 0.38 }
248
+ },
249
+ "qT": {
250
+ "dtype": "float16",
251
+ "shape": [2, 2, 4, 5],
252
+ "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.19999999999999998, "cosStep": 0.3 }
253
+ },
254
+ "qRotT": {
255
+ "dtype": "float16",
256
+ "shape": [2, 2, 4, 5],
257
+ "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.46, "cosStep": 0.14 }
258
+ },
259
+ "kT": {
260
+ "dtype": "float16",
261
+ "shape": [2, 2, 4, 5],
262
+ "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.38, "cosStep": 0.2 }
263
+ },
264
+ "kRotT": {
265
+ "dtype": "float16",
266
+ "shape": [2, 2, 4, 5],
267
+ "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.5, "cosStep": 0.26 }
268
+ }
269
+ },
270
+ "outputs": {
271
+ "output1T": { "dtype": "float16", "shape": [2, 2, 4, 5], "tolerance": 0.005 },
272
+ "output2T": { "dtype": "float16", "shape": [2, 2, 4, 5], "tolerance": 0.005 }
273
+ }
274
  }
275
  ]
276
  }