Hugging Face's logo

Kernels:
webgpu-kernels
/
ai.onnx.RotaryEmbedding

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
6332f1b
·
verified ·
1 Parent(s): 1844ba1

sync 6fdf6301e2bb

Browse files
Files changed (8) hide show
  1. README.md +20 -2
  2. build/webgpu/bench.json +346 -0
  3. build/webgpu/manifest.json +152 -133
  4. build/webgpu/metadata.json +14 -9
  5. build/webgpu/rotary-embedding-output-vec4.wgsl.jinja +63 -0
  6. build/webgpu/rotary-embedding-quad.wgsl.jinja +6 -32
  7. build/webgpu/rotary-embedding.wgsl.jinja +6 -32
  8. build/webgpu/test.json +253 -1
README.md CHANGED
@@ -48,19 +48,37 @@ Attributes and default values (overridable per request):
48
  | `T` | `float32`, `float16` |
49
  | `M` | `int64` |
50
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
51
  ## Files
52
 
53
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
54
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
55
  - [`test.json`](build/webgpu/test.json) — correctness cases
56
- - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 
57
  - [`rotary-embedding-quad.wgsl.jinja`](build/webgpu/rotary-embedding-quad.wgsl.jinja)
58
  - [`rotary-embedding.wgsl.jinja`](build/webgpu/rotary-embedding.wgsl.jinja)
59
 
60
  ## Use with `@huggingface/kernels`
61
 
62
  ```sh
63
- npm install --save-exact @huggingface/kernels@0.0.1-preview.2
64
  ```
65
 
66
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 
48
  | `T` | `float32`, `float16` |
49
  | `M` | `int64` |
50
 
51
+ ## Implementation variants
52
+
53
+ One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
54
+
55
+ - `rank3_cache2_pos_quad` — Processes four pairs with vector loads and stores when head and cache alignment keep each group within a head and rotation region.
56
+ - `rank3_cache2_pos` — Processes one rotated pair per invocation with scalar storage, or copies an unchanged pair from the tail.
57
+ - `rank3_cache2_pos_output_vec4` — Writes four contiguous outputs per invocation with gathered pair values. Non-rotated components retain their original storage bits.
58
+ - `rank4_cache2_pos_quad` — Processes four pairs with vector loads and stores when head and cache alignment keep each group within a head and rotation region.
59
+ - `rank4_cache2_pos` — Processes one rotated pair per invocation with scalar storage, or copies an unchanged pair from the tail.
60
+ - `rank4_cache2_pos_output_vec4` — Writes four contiguous outputs per invocation with gathered pair values. Non-rotated components retain their original storage bits.
61
+ - `rank3_cache3_quad` — Processes four pairs with vector loads and stores when head and cache alignment keep each group within a head and rotation region.
62
+ - `rank3_cache3` — Processes one rotated pair per invocation with scalar storage, or copies an unchanged pair from the tail.
63
+ - `rank3_cache3_output_vec4` — Writes four contiguous outputs per invocation with gathered pair values. Non-rotated components retain their original storage bits.
64
+ - `rank4_cache3_quad` — Processes four pairs with vector loads and stores when head and cache alignment keep each group within a head and rotation region.
65
+ - `rank4_cache3` — Processes one rotated pair per invocation with scalar storage, or copies an unchanged pair from the tail.
66
+ - `rank4_cache3_output_vec4` — Writes four contiguous outputs per invocation with gathered pair values. Non-rotated components retain their original storage bits.
67
+
68
  ## Files
69
 
70
  - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
71
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
72
  - [`test.json`](build/webgpu/test.json) — correctness cases
73
+ - [`bench.json`](build/webgpu/bench.json) — benchmark cases
74
+ - [`rotary-embedding-output-vec4.wgsl.jinja`](build/webgpu/rotary-embedding-output-vec4.wgsl.jinja)
75
  - [`rotary-embedding-quad.wgsl.jinja`](build/webgpu/rotary-embedding-quad.wgsl.jinja)
76
  - [`rotary-embedding.wgsl.jinja`](build/webgpu/rotary-embedding.wgsl.jinja)
77
 
78
  ## Use with `@huggingface/kernels`
79
 
80
  ```sh
81
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.3
82
  ```
83
 
84
  Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
build/webgpu/bench.json CHANGED
@@ -241,6 +241,352 @@
241
  }
242
  ]
243
  }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
244
  }
245
  ]
246
  }
 
241
  }
242
  ]
243
  }
244
+ },
245
+ {
246
+ "name": "geometry_float32_rank3_head100_rot94_interleaved0_indexedfalse",
247
+ "preset": "edge",
248
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 94, "num_heads": 3 },
249
+ "inputs": {
250
+ "x": { "dtype": "float32", "shape": [2, 257, 300], "dist": "normal", "seed": 1601, "scale": 0.2 },
251
+ "cos": { "dtype": "float32", "shape": [2, 257, 47], "dist": "uniform", "seed": 1602, "scale": 1 },
252
+ "sin": { "dtype": "float32", "shape": [2, 257, 47], "dist": "uniform", "seed": 1603, "scale": 1 }
253
+ },
254
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 257, 300] } }
255
+ },
256
+ {
257
+ "name": "geometry_float32_rank3_head100_rot94_interleaved0_indexedtrue",
258
+ "preset": "edge",
259
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 94, "num_heads": 3 },
260
+ "inputs": {
261
+ "x": { "dtype": "float32", "shape": [2, 257, 300], "dist": "normal", "seed": 1601, "scale": 0.2 },
262
+ "cos": { "dtype": "float32", "shape": [262, 47], "dist": "uniform", "seed": 1602, "scale": 1 },
263
+ "sin": { "dtype": "float32", "shape": [262, 47], "dist": "uniform", "seed": 1603, "scale": 1 },
264
+ "positionIds": { "dtype": "uint32", "shape": [2, 257], "dist": "uniform", "seed": 1604, "min": 0, "max": 262 }
265
+ },
266
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 257, 300] } }
267
+ },
268
+ {
269
+ "name": "geometry_float32_rank3_head100_rot100_interleaved0_indexedfalse",
270
+ "preset": "edge",
271
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 100, "num_heads": 3 },
272
+ "inputs": {
273
+ "x": { "dtype": "float32", "shape": [2, 257, 300], "dist": "normal", "seed": 1601, "scale": 0.2 },
274
+ "cos": { "dtype": "float32", "shape": [2, 257, 50], "dist": "uniform", "seed": 1602, "scale": 1 },
275
+ "sin": { "dtype": "float32", "shape": [2, 257, 50], "dist": "uniform", "seed": 1603, "scale": 1 }
276
+ },
277
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 257, 300] } }
278
+ },
279
+ {
280
+ "name": "geometry_float32_rank3_head100_rot100_interleaved0_indexedtrue",
281
+ "preset": "edge",
282
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 100, "num_heads": 3 },
283
+ "inputs": {
284
+ "x": { "dtype": "float32", "shape": [2, 257, 300], "dist": "normal", "seed": 1601, "scale": 0.2 },
285
+ "cos": { "dtype": "float32", "shape": [262, 50], "dist": "uniform", "seed": 1602, "scale": 1 },
286
+ "sin": { "dtype": "float32", "shape": [262, 50], "dist": "uniform", "seed": 1603, "scale": 1 },
287
+ "positionIds": { "dtype": "uint32", "shape": [2, 257], "dist": "uniform", "seed": 1604, "min": 0, "max": 262 }
288
+ },
289
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 257, 300] } }
290
+ },
291
+ {
292
+ "name": "geometry_float32_rank3_head100_rot92_interleaved0_indexedfalse",
293
+ "preset": "edge",
294
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 92, "num_heads": 3 },
295
+ "inputs": {
296
+ "x": { "dtype": "float32", "shape": [2, 257, 300], "dist": "normal", "seed": 1601, "scale": 0.2 },
297
+ "cos": { "dtype": "float32", "shape": [2, 257, 46], "dist": "uniform", "seed": 1602, "scale": 1 },
298
+ "sin": { "dtype": "float32", "shape": [2, 257, 46], "dist": "uniform", "seed": 1603, "scale": 1 }
299
+ },
300
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 257, 300] } }
301
+ },
302
+ {
303
+ "name": "geometry_float32_rank3_head100_rot92_interleaved0_indexedtrue",
304
+ "preset": "edge",
305
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 92, "num_heads": 3 },
306
+ "inputs": {
307
+ "x": { "dtype": "float32", "shape": [2, 257, 300], "dist": "normal", "seed": 1601, "scale": 0.2 },
308
+ "cos": { "dtype": "float32", "shape": [262, 46], "dist": "uniform", "seed": 1602, "scale": 1 },
309
+ "sin": { "dtype": "float32", "shape": [262, 46], "dist": "uniform", "seed": 1603, "scale": 1 },
310
+ "positionIds": { "dtype": "uint32", "shape": [2, 257], "dist": "uniform", "seed": 1604, "min": 0, "max": 262 }
311
+ },
312
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 257, 300] } }
313
+ },
314
+ {
315
+ "name": "geometry_float32_rank4_head100_rot94_interleaved0_indexedfalse",
316
+ "preset": "edge",
317
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 94 },
318
+ "inputs": {
319
+ "x": { "dtype": "float32", "shape": [2, 3, 257, 100], "dist": "normal", "seed": 1601, "scale": 0.2 },
320
+ "cos": { "dtype": "float32", "shape": [2, 257, 47], "dist": "uniform", "seed": 1602, "scale": 1 },
321
+ "sin": { "dtype": "float32", "shape": [2, 257, 47], "dist": "uniform", "seed": 1603, "scale": 1 }
322
+ },
323
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 257, 100] } }
324
+ },
325
+ {
326
+ "name": "geometry_float32_rank4_head100_rot94_interleaved0_indexedtrue",
327
+ "preset": "edge",
328
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 94 },
329
+ "inputs": {
330
+ "x": { "dtype": "float32", "shape": [2, 3, 257, 100], "dist": "normal", "seed": 1601, "scale": 0.2 },
331
+ "cos": { "dtype": "float32", "shape": [262, 47], "dist": "uniform", "seed": 1602, "scale": 1 },
332
+ "sin": { "dtype": "float32", "shape": [262, 47], "dist": "uniform", "seed": 1603, "scale": 1 },
333
+ "positionIds": { "dtype": "uint32", "shape": [2, 257], "dist": "uniform", "seed": 1604, "min": 0, "max": 262 }
334
+ },
335
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 257, 100] } }
336
+ },
337
+ {
338
+ "name": "geometry_float32_rank4_head100_rot100_interleaved0_indexedfalse",
339
+ "preset": "edge",
340
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 100 },
341
+ "inputs": {
342
+ "x": { "dtype": "float32", "shape": [2, 3, 257, 100], "dist": "normal", "seed": 1601, "scale": 0.2 },
343
+ "cos": { "dtype": "float32", "shape": [2, 257, 50], "dist": "uniform", "seed": 1602, "scale": 1 },
344
+ "sin": { "dtype": "float32", "shape": [2, 257, 50], "dist": "uniform", "seed": 1603, "scale": 1 }
345
+ },
346
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 257, 100] } }
347
+ },
348
+ {
349
+ "name": "geometry_float32_rank4_head100_rot100_interleaved0_indexedtrue",
350
+ "preset": "edge",
351
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 100 },
352
+ "inputs": {
353
+ "x": { "dtype": "float32", "shape": [2, 3, 257, 100], "dist": "normal", "seed": 1601, "scale": 0.2 },
354
+ "cos": { "dtype": "float32", "shape": [262, 50], "dist": "uniform", "seed": 1602, "scale": 1 },
355
+ "sin": { "dtype": "float32", "shape": [262, 50], "dist": "uniform", "seed": 1603, "scale": 1 },
356
+ "positionIds": { "dtype": "uint32", "shape": [2, 257], "dist": "uniform", "seed": 1604, "min": 0, "max": 262 }
357
+ },
358
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 257, 100] } }
359
+ },
360
+ {
361
+ "name": "geometry_float32_rank4_head100_rot92_interleaved0_indexedfalse",
362
+ "preset": "edge",
363
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 92 },
364
+ "inputs": {
365
+ "x": { "dtype": "float32", "shape": [2, 3, 257, 100], "dist": "normal", "seed": 1601, "scale": 0.2 },
366
+ "cos": { "dtype": "float32", "shape": [2, 257, 46], "dist": "uniform", "seed": 1602, "scale": 1 },
367
+ "sin": { "dtype": "float32", "shape": [2, 257, 46], "dist": "uniform", "seed": 1603, "scale": 1 }
368
+ },
369
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 257, 100] } }
370
+ },
371
+ {
372
+ "name": "geometry_float32_rank4_head100_rot92_interleaved0_indexedtrue",
373
+ "preset": "edge",
374
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 92 },
375
+ "inputs": {
376
+ "x": { "dtype": "float32", "shape": [2, 3, 257, 100], "dist": "normal", "seed": 1601, "scale": 0.2 },
377
+ "cos": { "dtype": "float32", "shape": [262, 46], "dist": "uniform", "seed": 1602, "scale": 1 },
378
+ "sin": { "dtype": "float32", "shape": [262, 46], "dist": "uniform", "seed": 1603, "scale": 1 },
379
+ "positionIds": { "dtype": "uint32", "shape": [2, 257], "dist": "uniform", "seed": 1604, "min": 0, "max": 262 }
380
+ },
381
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 257, 100] } }
382
+ },
383
+ {
384
+ "name": "geometry_float16_rank3_head128_rot126_interleaved0_indexedfalse",
385
+ "preset": "edge",
386
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 126, "num_heads": 3 },
387
+ "inputs": {
388
+ "x": { "dtype": "float16", "shape": [2, 257, 384], "dist": "normal", "seed": 1601, "scale": 0.2 },
389
+ "cos": { "dtype": "float16", "shape": [2, 257, 63], "dist": "uniform", "seed": 1602, "scale": 1 },
390
+ "sin": { "dtype": "float16", "shape": [2, 257, 63], "dist": "uniform", "seed": 1603, "scale": 1 }
391
+ },
392
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 257, 384] } }
393
+ },
394
+ {
395
+ "name": "geometry_float16_rank3_head128_rot126_interleaved0_indexedtrue",
396
+ "preset": "edge",
397
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 126, "num_heads": 3 },
398
+ "inputs": {
399
+ "x": { "dtype": "float16", "shape": [2, 257, 384], "dist": "normal", "seed": 1601, "scale": 0.2 },
400
+ "cos": { "dtype": "float16", "shape": [262, 63], "dist": "uniform", "seed": 1602, "scale": 1 },
401
+ "sin": { "dtype": "float16", "shape": [262, 63], "dist": "uniform", "seed": 1603, "scale": 1 },
402
+ "positionIds": { "dtype": "uint32", "shape": [2, 257], "dist": "uniform", "seed": 1604, "min": 0, "max": 262 }
403
+ },
404
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 257, 384] } }
405
+ },
406
+ {
407
+ "name": "geometry_float16_rank3_head100_rot94_interleaved0_indexedfalse",
408
+ "preset": "edge",
409
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 94, "num_heads": 3 },
410
+ "inputs": {
411
+ "x": { "dtype": "float16", "shape": [2, 257, 300], "dist": "normal", "seed": 1601, "scale": 0.2 },
412
+ "cos": { "dtype": "float16", "shape": [2, 257, 47], "dist": "uniform", "seed": 1602, "scale": 1 },
413
+ "sin": { "dtype": "float16", "shape": [2, 257, 47], "dist": "uniform", "seed": 1603, "scale": 1 }
414
+ },
415
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 257, 300] } }
416
+ },
417
+ {
418
+ "name": "geometry_float16_rank3_head100_rot94_interleaved0_indexedtrue",
419
+ "preset": "edge",
420
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 94, "num_heads": 3 },
421
+ "inputs": {
422
+ "x": { "dtype": "float16", "shape": [2, 257, 300], "dist": "normal", "seed": 1601, "scale": 0.2 },
423
+ "cos": { "dtype": "float16", "shape": [262, 47], "dist": "uniform", "seed": 1602, "scale": 1 },
424
+ "sin": { "dtype": "float16", "shape": [262, 47], "dist": "uniform", "seed": 1603, "scale": 1 },
425
+ "positionIds": { "dtype": "uint32", "shape": [2, 257], "dist": "uniform", "seed": 1604, "min": 0, "max": 262 }
426
+ },
427
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 257, 300] } }
428
+ },
429
+ {
430
+ "name": "geometry_float16_rank3_head100_rot100_interleaved0_indexedfalse",
431
+ "preset": "edge",
432
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 100, "num_heads": 3 },
433
+ "inputs": {
434
+ "x": { "dtype": "float16", "shape": [2, 257, 300], "dist": "normal", "seed": 1601, "scale": 0.2 },
435
+ "cos": { "dtype": "float16", "shape": [2, 257, 50], "dist": "uniform", "seed": 1602, "scale": 1 },
436
+ "sin": { "dtype": "float16", "shape": [2, 257, 50], "dist": "uniform", "seed": 1603, "scale": 1 }
437
+ },
438
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 257, 300] } }
439
+ },
440
+ {
441
+ "name": "geometry_float16_rank3_head100_rot100_interleaved0_indexedtrue",
442
+ "preset": "edge",
443
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 100, "num_heads": 3 },
444
+ "inputs": {
445
+ "x": { "dtype": "float16", "shape": [2, 257, 300], "dist": "normal", "seed": 1601, "scale": 0.2 },
446
+ "cos": { "dtype": "float16", "shape": [262, 50], "dist": "uniform", "seed": 1602, "scale": 1 },
447
+ "sin": { "dtype": "float16", "shape": [262, 50], "dist": "uniform", "seed": 1603, "scale": 1 },
448
+ "positionIds": { "dtype": "uint32", "shape": [2, 257], "dist": "uniform", "seed": 1604, "min": 0, "max": 262 }
449
+ },
450
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 257, 300] } }
451
+ },
452
+ {
453
+ "name": "geometry_float16_rank3_head100_rot92_interleaved0_indexedfalse",
454
+ "preset": "edge",
455
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 92, "num_heads": 3 },
456
+ "inputs": {
457
+ "x": { "dtype": "float16", "shape": [2, 257, 300], "dist": "normal", "seed": 1601, "scale": 0.2 },
458
+ "cos": { "dtype": "float16", "shape": [2, 257, 46], "dist": "uniform", "seed": 1602, "scale": 1 },
459
+ "sin": { "dtype": "float16", "shape": [2, 257, 46], "dist": "uniform", "seed": 1603, "scale": 1 }
460
+ },
461
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 257, 300] } }
462
+ },
463
+ {
464
+ "name": "geometry_float16_rank3_head100_rot92_interleaved0_indexedtrue",
465
+ "preset": "edge",
466
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 92, "num_heads": 3 },
467
+ "inputs": {
468
+ "x": { "dtype": "float16", "shape": [2, 257, 300], "dist": "normal", "seed": 1601, "scale": 0.2 },
469
+ "cos": { "dtype": "float16", "shape": [262, 46], "dist": "uniform", "seed": 1602, "scale": 1 },
470
+ "sin": { "dtype": "float16", "shape": [262, 46], "dist": "uniform", "seed": 1603, "scale": 1 },
471
+ "positionIds": { "dtype": "uint32", "shape": [2, 257], "dist": "uniform", "seed": 1604, "min": 0, "max": 262 }
472
+ },
473
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 257, 300] } }
474
+ },
475
+ {
476
+ "name": "geometry_float16_rank4_head128_rot126_interleaved0_indexedfalse",
477
+ "preset": "edge",
478
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 126 },
479
+ "inputs": {
480
+ "x": { "dtype": "float16", "shape": [2, 3, 257, 128], "dist": "normal", "seed": 1601, "scale": 0.2 },
481
+ "cos": { "dtype": "float16", "shape": [2, 257, 63], "dist": "uniform", "seed": 1602, "scale": 1 },
482
+ "sin": { "dtype": "float16", "shape": [2, 257, 63], "dist": "uniform", "seed": 1603, "scale": 1 }
483
+ },
484
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 3, 257, 128] } }
485
+ },
486
+ {
487
+ "name": "geometry_float16_rank4_head128_rot126_interleaved0_indexedtrue",
488
+ "preset": "edge",
489
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 126 },
490
+ "inputs": {
491
+ "x": { "dtype": "float16", "shape": [2, 3, 257, 128], "dist": "normal", "seed": 1601, "scale": 0.2 },
492
+ "cos": { "dtype": "float16", "shape": [262, 63], "dist": "uniform", "seed": 1602, "scale": 1 },
493
+ "sin": { "dtype": "float16", "shape": [262, 63], "dist": "uniform", "seed": 1603, "scale": 1 },
494
+ "positionIds": { "dtype": "uint32", "shape": [2, 257], "dist": "uniform", "seed": 1604, "min": 0, "max": 262 }
495
+ },
496
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 3, 257, 128] } }
497
+ },
498
+ {
499
+ "name": "geometry_float16_rank4_head100_rot94_interleaved0_indexedfalse",
500
+ "preset": "edge",
501
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 94 },
502
+ "inputs": {
503
+ "x": { "dtype": "float16", "shape": [2, 3, 257, 100], "dist": "normal", "seed": 1601, "scale": 0.2 },
504
+ "cos": { "dtype": "float16", "shape": [2, 257, 47], "dist": "uniform", "seed": 1602, "scale": 1 },
505
+ "sin": { "dtype": "float16", "shape": [2, 257, 47], "dist": "uniform", "seed": 1603, "scale": 1 }
506
+ },
507
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 3, 257, 100] } }
508
+ },
509
+ {
510
+ "name": "geometry_float16_rank4_head100_rot94_interleaved0_indexedtrue",
511
+ "preset": "edge",
512
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 94 },
513
+ "inputs": {
514
+ "x": { "dtype": "float16", "shape": [2, 3, 257, 100], "dist": "normal", "seed": 1601, "scale": 0.2 },
515
+ "cos": { "dtype": "float16", "shape": [262, 47], "dist": "uniform", "seed": 1602, "scale": 1 },
516
+ "sin": { "dtype": "float16", "shape": [262, 47], "dist": "uniform", "seed": 1603, "scale": 1 },
517
+ "positionIds": { "dtype": "uint32", "shape": [2, 257], "dist": "uniform", "seed": 1604, "min": 0, "max": 262 }
518
+ },
519
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 3, 257, 100] } }
520
+ },
521
+ {
522
+ "name": "geometry_float16_rank4_head100_rot100_interleaved0_indexedfalse",
523
+ "preset": "edge",
524
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 100 },
525
+ "inputs": {
526
+ "x": { "dtype": "float16", "shape": [2, 3, 257, 100], "dist": "normal", "seed": 1601, "scale": 0.2 },
527
+ "cos": { "dtype": "float16", "shape": [2, 257, 50], "dist": "uniform", "seed": 1602, "scale": 1 },
528
+ "sin": { "dtype": "float16", "shape": [2, 257, 50], "dist": "uniform", "seed": 1603, "scale": 1 }
529
+ },
530
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 3, 257, 100] } }
531
+ },
532
+ {
533
+ "name": "geometry_float16_rank4_head100_rot100_interleaved0_indexedtrue",
534
+ "preset": "edge",
535
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 100 },
536
+ "inputs": {
537
+ "x": { "dtype": "float16", "shape": [2, 3, 257, 100], "dist": "normal", "seed": 1601, "scale": 0.2 },
538
+ "cos": { "dtype": "float16", "shape": [262, 50], "dist": "uniform", "seed": 1602, "scale": 1 },
539
+ "sin": { "dtype": "float16", "shape": [262, 50], "dist": "uniform", "seed": 1603, "scale": 1 },
540
+ "positionIds": { "dtype": "uint32", "shape": [2, 257], "dist": "uniform", "seed": 1604, "min": 0, "max": 262 }
541
+ },
542
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 3, 257, 100] } }
543
+ },
544
+ {
545
+ "name": "geometry_float16_rank4_head100_rot92_interleaved0_indexedfalse",
546
+ "preset": "edge",
547
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 92 },
548
+ "inputs": {
549
+ "x": { "dtype": "float16", "shape": [2, 3, 257, 100], "dist": "normal", "seed": 1601, "scale": 0.2 },
550
+ "cos": { "dtype": "float16", "shape": [2, 257, 46], "dist": "uniform", "seed": 1602, "scale": 1 },
551
+ "sin": { "dtype": "float16", "shape": [2, 257, 46], "dist": "uniform", "seed": 1603, "scale": 1 }
552
+ },
553
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 3, 257, 100] } }
554
+ },
555
+ {
556
+ "name": "geometry_float16_rank4_head100_rot92_interleaved0_indexedtrue",
557
+ "preset": "edge",
558
+ "attrs": { "interleaved": 0, "rotary_embedding_dim": 92 },
559
+ "inputs": {
560
+ "x": { "dtype": "float16", "shape": [2, 3, 257, 100], "dist": "normal", "seed": 1601, "scale": 0.2 },
561
+ "cos": { "dtype": "float16", "shape": [262, 46], "dist": "uniform", "seed": 1602, "scale": 1 },
562
+ "sin": { "dtype": "float16", "shape": [262, 46], "dist": "uniform", "seed": 1603, "scale": 1 },
563
+ "positionIds": { "dtype": "uint32", "shape": [2, 257], "dist": "uniform", "seed": 1604, "min": 0, "max": 262 }
564
+ },
565
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 3, 257, 100] } }
566
+ },
567
+ {
568
+ "name": "geometry_output_fold_float32_rot126_wg8",
569
+ "preset": "stress",
570
+ "attrs": { "num_heads": 1, "rotary_embedding_dim": 126 },
571
+ "tunables": { "WORKGROUP_SIZE": 8 },
572
+ "inputs": {
573
+ "x": { "dtype": "float32", "shape": [2, 8193, 128], "dist": "normal", "seed": 1611, "scale": 0.2 },
574
+ "cos": { "dtype": "float32", "shape": [2, 8193, 63], "dist": "uniform", "seed": 1612, "scale": 1 },
575
+ "sin": { "dtype": "float32", "shape": [2, 8193, 63], "dist": "uniform", "seed": 1613, "scale": 1 }
576
+ },
577
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 8193, 128] } }
578
+ },
579
+ {
580
+ "name": "geometry_output_fold_float16_rot126_wg8",
581
+ "preset": "stress",
582
+ "attrs": { "num_heads": 1, "rotary_embedding_dim": 126 },
583
+ "tunables": { "WORKGROUP_SIZE": 8 },
584
+ "inputs": {
585
+ "x": { "dtype": "float16", "shape": [2, 8193, 128], "dist": "normal", "seed": 1611, "scale": 0.2 },
586
+ "cos": { "dtype": "float16", "shape": [2, 8193, 63], "dist": "uniform", "seed": 1612, "scale": 1 },
587
+ "sin": { "dtype": "float16", "shape": [2, 8193, 63], "dist": "uniform", "seed": 1613, "scale": 1 }
588
+ },
589
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 8193, 128] } }
590
  }
591
  ]
592
  }
build/webgpu/manifest.json CHANGED
@@ -30,139 +30,90 @@
30
  "quadCacheOk": "ranks.cos >= 1 and dim(shapes.cos, ranks.cos - 1) % 4 == 0",
31
  "quad3Ok": "rank3HeadSize % 8 == 0 and quadCacheOk",
32
  "quad4Ok": "ranks.x == 4 and dim(shapes.x, 3) % 8 == 0 and quadCacheOk",
33
- "quadDispatchOk": "ceilDiv(numel(shapes.x) / 8, tunables.WORKGROUP_SIZE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)"
 
 
 
 
 
 
34
  },
35
  "bindings": {
36
- "x": { "buffer": "read-only-storage", "elementType": "$vector" },
37
- "cos_cache": { "arg": "cos", "buffer": "read-only-storage", "elementType": "$vector" },
38
- "sin_cache": { "arg": "sin", "buffer": "read-only-storage", "elementType": "$vector" },
39
- "position_ids": { "arg": "positionIds", "buffer": "read-only-storage", "elementType": "$M" },
40
- "y": { "buffer": "storage", "elementType": "$vector" },
41
- "params": {
42
- "buffer": "uniform",
43
- "struct": [
44
- { "name": "quadCount", "type": "u32", "value": "numel(shapes.x) / 8" },
45
- { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.x, 1)" },
46
- { "name": "numHeads", "type": "u32", "value": "attrs.num_heads" },
47
- { "name": "headVec4", "type": "u32", "value": "rank3HeadSize / 4" },
48
- { "name": "headQuads", "type": "u32", "value": "rank3HeadSize / 8" },
49
- { "name": "rotaryVec4", "type": "u32", "value": "dim(shapes.cos, 1) / 2" },
50
- { "name": "halfRotaryVec4", "type": "u32", "value": "dim(shapes.cos, 1) / 4" }
51
- ]
52
- },
53
- "x_2": { "name": "x", "buffer": "read-only-storage", "elementType": "$T" },
54
- "cos_cache_2": { "arg": "cos", "name": "cos_cache", "buffer": "read-only-storage", "elementType": "$T" },
55
- "sin_cache_2": { "arg": "sin", "name": "sin_cache", "buffer": "read-only-storage", "elementType": "$T" },
56
- "y_2": { "name": "y", "buffer": "storage", "elementType": "$T" },
57
- "params_2": {
58
- "name": "params",
59
- "buffer": "uniform",
60
- "struct": [
61
- { "name": "pairCount", "type": "u32", "value": "numel(shapes.x) / 2" },
62
- { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.x, 1)" },
63
- { "name": "numHeads", "type": "u32", "value": "attrs.num_heads" },
64
- { "name": "headSize", "type": "u32", "value": "rank3HeadSize" },
65
- { "name": "rotaryDim", "type": "u32", "value": "dim(shapes.cos, 1) * 2" },
66
- { "name": "halfRotaryDim", "type": "u32", "value": "dim(shapes.cos, 1)" }
67
- ]
68
- },
69
- "params_3": {
70
  "name": "params",
71
- "buffer": "uniform",
72
  "struct": [
73
  { "name": "quadCount", "type": "u32", "value": "numel(shapes.x) / 8" },
74
- { "name": "numHeads", "type": "u32", "value": "dim(shapes.x, 1)" },
75
- { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.x, 2)" },
76
- { "name": "headVec4", "type": "u32", "value": "dim(shapes.x, 3) / 4" },
77
- { "name": "headQuads", "type": "u32", "value": "dim(shapes.x, 3) / 8" },
78
- { "name": "rotaryVec4", "type": "u32", "value": "dim(shapes.cos, 1) / 2" },
79
- { "name": "halfRotaryVec4", "type": "u32", "value": "dim(shapes.cos, 1) / 4" }
80
  ]
81
  },
82
- "params_4": {
83
  "name": "params",
84
- "buffer": "uniform",
85
  "struct": [
86
  { "name": "pairCount", "type": "u32", "value": "numel(shapes.x) / 2" },
87
- { "name": "numHeads", "type": "u32", "value": "dim(shapes.x, 1)" },
88
- { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.x, 2)" },
89
- { "name": "headSize", "type": "u32", "value": "dim(shapes.x, 3)" },
90
- { "name": "rotaryDim", "type": "u32", "value": "dim(shapes.cos, 1) * 2" },
91
- { "name": "halfRotaryDim", "type": "u32", "value": "dim(shapes.cos, 1)" }
92
  ]
93
  },
94
- "params_5": {
95
  "name": "params",
96
- "buffer": "uniform",
97
  "struct": [
98
  { "name": "quadCount", "type": "u32", "value": "numel(shapes.x) / 8" },
99
- { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.x, 1)" },
100
- { "name": "numHeads", "type": "u32", "value": "attrs.num_heads" },
101
- { "name": "headVec4", "type": "u32", "value": "rank3HeadSize / 4" },
102
- { "name": "headQuads", "type": "u32", "value": "rank3HeadSize / 8" },
103
- { "name": "rotaryVec4", "type": "u32", "value": "dim(shapes.cos, 2) / 2" },
104
- { "name": "halfRotaryVec4", "type": "u32", "value": "dim(shapes.cos, 2) / 4" },
105
  { "name": "cacheBatchStrideVec4", "type": "u32", "value": "dim(shapes.cos, 1) * dim(shapes.cos, 2) / 4" }
106
  ]
107
  },
108
- "params_6": {
109
  "name": "params",
110
- "buffer": "uniform",
111
  "struct": [
112
  { "name": "pairCount", "type": "u32", "value": "numel(shapes.x) / 2" },
113
- { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.x, 1)" },
114
- { "name": "numHeads", "type": "u32", "value": "attrs.num_heads" },
115
- { "name": "headSize", "type": "u32", "value": "rank3HeadSize" },
116
- { "name": "rotaryDim", "type": "u32", "value": "dim(shapes.cos, 2) * 2" },
117
- { "name": "halfRotaryDim", "type": "u32", "value": "dim(shapes.cos, 2)" },
118
  { "name": "cacheBatchStride", "type": "u32", "value": "dim(shapes.cos, 1) * dim(shapes.cos, 2)" }
119
  ]
120
  },
121
- "params_7": {
122
- "name": "params",
123
- "buffer": "uniform",
124
- "struct": [
125
- { "name": "quadCount", "type": "u32", "value": "numel(shapes.x) / 8" },
126
- { "name": "numHeads", "type": "u32", "value": "dim(shapes.x, 1)" },
127
- { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.x, 2)" },
128
- { "name": "headVec4", "type": "u32", "value": "dim(shapes.x, 3) / 4" },
129
- { "name": "headQuads", "type": "u32", "value": "dim(shapes.x, 3) / 8" },
130
- { "name": "rotaryVec4", "type": "u32", "value": "dim(shapes.cos, 2) / 2" },
131
- { "name": "halfRotaryVec4", "type": "u32", "value": "dim(shapes.cos, 2) / 4" },
132
- { "name": "cacheBatchStrideVec4", "type": "u32", "value": "dim(shapes.cos, 1) * dim(shapes.cos, 2) / 4" }
133
- ]
134
- },
135
- "params_8": {
136
  "name": "params",
137
- "buffer": "uniform",
138
- "struct": [
139
- { "name": "pairCount", "type": "u32", "value": "numel(shapes.x) / 2" },
140
- { "name": "numHeads", "type": "u32", "value": "dim(shapes.x, 1)" },
141
- { "name": "sequenceLength", "type": "u32", "value": "dim(shapes.x, 2)" },
142
- { "name": "headSize", "type": "u32", "value": "dim(shapes.x, 3)" },
143
- { "name": "rotaryDim", "type": "u32", "value": "dim(shapes.cos, 2) * 2" },
144
- { "name": "halfRotaryDim", "type": "u32", "value": "dim(shapes.cos, 2)" },
145
- { "name": "cacheBatchStride", "type": "u32", "value": "dim(shapes.cos, 1) * dim(shapes.cos, 2)" }
146
- ]
147
  }
148
  },
149
  "variants": [
150
  {
151
  "id": "rank3_cache2_pos_quad",
152
  "priority": 10,
153
- "when": ["rank3PosContract", "quad3Ok", "quadDispatchOk"],
154
- "derive": {
155
- "interleaved": "attrs.interleaved != 0",
156
- "usesF16": "dtypes.T == \"f16\"",
157
- "vector": "\"vec4<\" ~ dtypes.T ~ \">\""
158
- },
159
  "passes": [
160
  {
161
  "id": "main",
162
- "name": "rotary_embedding3d_pos_quad",
163
  "shader": "rotary-embedding-quad.wgsl.jinja",
164
  "derive": { "rank": 3, "hasPositionIds": "true" },
165
- "bindings": ["x", "cos_cache", "sin_cache", "position_ids", "y", "params"],
166
  "dispatch": {
167
  "x": "min(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
168
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -175,14 +126,14 @@
175
  "id": "rank3_cache2_pos",
176
  "priority": 0,
177
  "when": ["rank3PosContract", "pairDispatchOk"],
178
- "derive": { "interleaved": "attrs.interleaved != 0", "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
179
  "passes": [
180
  {
181
  "id": "main",
182
- "name": "rotary_embedding3d_pos",
183
  "shader": "rotary-embedding.wgsl.jinja",
184
  "derive": { "rank": 3, "hasPositionIds": "true" },
185
- "bindings": ["x_2", "cos_cache_2", "sin_cache_2", "position_ids", "y_2", "params_2"],
186
  "dispatch": {
187
  "x": "min(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
188
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -191,22 +142,38 @@
191
  }
192
  ]
193
  },
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
194
  {
195
  "id": "rank4_cache2_pos_quad",
196
  "priority": 10,
197
- "when": ["rank4PosContract", "quad4Ok", "quadDispatchOk"],
198
- "derive": {
199
- "interleaved": "attrs.interleaved != 0",
200
- "usesF16": "dtypes.T == \"f16\"",
201
- "vector": "\"vec4<\" ~ dtypes.T ~ \">\""
202
- },
203
  "passes": [
204
  {
205
  "id": "main",
206
- "name": "rotary_embedding4d_pos_quad",
207
  "shader": "rotary-embedding-quad.wgsl.jinja",
208
  "derive": { "rank": 4, "hasPositionIds": "true" },
209
- "bindings": ["x", "cos_cache", "sin_cache", "position_ids", "y", "params_3"],
210
  "dispatch": {
211
  "x": "min(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
212
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -219,14 +186,14 @@
219
  "id": "rank4_cache2_pos",
220
  "priority": 0,
221
  "when": ["rank4PosContract", "pairDispatchOk"],
222
- "derive": { "interleaved": "attrs.interleaved != 0", "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
223
  "passes": [
224
  {
225
  "id": "main",
226
- "name": "rotary_embedding4d_pos",
227
  "shader": "rotary-embedding.wgsl.jinja",
228
  "derive": { "rank": 4, "hasPositionIds": "true" },
229
- "bindings": ["x_2", "cos_cache_2", "sin_cache_2", "position_ids", "y_2", "params_4"],
230
  "dispatch": {
231
  "x": "min(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
232
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -235,22 +202,38 @@
235
  }
236
  ]
237
  },
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
238
  {
239
  "id": "rank3_cache3_quad",
240
  "priority": 10,
241
- "when": ["rank3Contract", "quad3Ok", "quadDispatchOk"],
242
- "derive": {
243
- "interleaved": "attrs.interleaved != 0",
244
- "usesF16": "dtypes.T == \"f16\"",
245
- "vector": "\"vec4<\" ~ dtypes.T ~ \">\""
246
- },
247
  "passes": [
248
  {
249
  "id": "main",
250
- "name": "rotary_embedding3d_quad",
251
  "shader": "rotary-embedding-quad.wgsl.jinja",
252
  "derive": { "rank": 3 },
253
- "bindings": ["x", "cos_cache", "sin_cache", "y", "params_5"],
254
  "dispatch": {
255
  "x": "min(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
256
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -263,14 +246,14 @@
263
  "id": "rank3_cache3",
264
  "priority": 0,
265
  "when": ["rank3Contract", "pairDispatchOk"],
266
- "derive": { "interleaved": "attrs.interleaved != 0", "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
267
  "passes": [
268
  {
269
  "id": "main",
270
- "name": "rotary_embedding3d",
271
  "shader": "rotary-embedding.wgsl.jinja",
272
  "derive": { "rank": 3 },
273
- "bindings": ["x_2", "cos_cache_2", "sin_cache_2", "y_2", "params_6"],
274
  "dispatch": {
275
  "x": "min(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
276
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -279,22 +262,38 @@
279
  }
280
  ]
281
  },
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
282
  {
283
  "id": "rank4_cache3_quad",
284
  "priority": 10,
285
- "when": ["rank4Contract", "quad4Ok", "quadDispatchOk"],
286
- "derive": {
287
- "interleaved": "attrs.interleaved != 0",
288
- "usesF16": "dtypes.T == \"f16\"",
289
- "vector": "\"vec4<\" ~ dtypes.T ~ \">\""
290
- },
291
  "passes": [
292
  {
293
  "id": "main",
294
- "name": "rotary_embedding4d_quad",
295
  "shader": "rotary-embedding-quad.wgsl.jinja",
296
  "derive": { "rank": 4 },
297
- "bindings": ["x", "cos_cache", "sin_cache", "y", "params_7"],
298
  "dispatch": {
299
  "x": "min(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
300
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -307,14 +306,14 @@
307
  "id": "rank4_cache3",
308
  "priority": 0,
309
  "when": ["rank4Contract", "pairDispatchOk"],
310
- "derive": { "interleaved": "attrs.interleaved != 0", "usesF16": "dtypes.T == \"f16\"", "scalar": "dtypes.T" },
311
  "passes": [
312
  {
313
  "id": "main",
314
- "name": "rotary_embedding4d",
315
  "shader": "rotary-embedding.wgsl.jinja",
316
  "derive": { "rank": 4 },
317
- "bindings": ["x_2", "cos_cache_2", "sin_cache_2", "y_2", "params_8"],
318
  "dispatch": {
319
  "x": "min(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
320
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
@@ -322,6 +321,26 @@
322
  }
323
  }
324
  ]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
325
  }
326
  ]
327
  }
 
30
  "quadCacheOk": "ranks.cos >= 1 and dim(shapes.cos, ranks.cos - 1) % 4 == 0",
31
  "quad3Ok": "rank3HeadSize % 8 == 0 and quadCacheOk",
32
  "quad4Ok": "ranks.x == 4 and dim(shapes.x, 3) % 8 == 0 and quadCacheOk",
33
+ "quadOk": "quad3Ok if ranks.x == 3 else quad4Ok",
34
+ "quadDispatchOk": "ceilDiv(numel(shapes.x) / 8, tunables.WORKGROUP_SIZE) <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
35
+ "headElements": "rank3HeadSize if ranks.x == 3 else dim(shapes.x, 3)",
36
+ "headCount": "attrs.num_heads if ranks.x == 3 else dim(shapes.x, 1)",
37
+ "tokenCount": "dim(shapes.x, 1) if ranks.x == 3 else dim(shapes.x, 2)",
38
+ "rotationPairs": "dim(shapes.cos, -1)",
39
+ "rotaryOutputWg": "max(1, min(tunables.WORKGROUP_SIZE, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX))"
40
  },
41
  "bindings": {
42
+ "x": { "elementType": "$vector" },
43
+ "cos_cache": { "arg": "cos", "elementType": "$vector" },
44
+ "sin_cache": { "arg": "sin", "elementType": "$vector" },
45
+ "position_ids": { "arg": "positionIds", "elementType": "$M" },
46
+ "y": { "elementType": "$vector" },
47
+ "params_quad": {
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
48
  "name": "params",
 
49
  "struct": [
50
  { "name": "quadCount", "type": "u32", "value": "numel(shapes.x) / 8" },
51
+ { "name": "numHeads", "type": "u32", "value": "headCount" },
52
+ { "name": "sequenceLength", "type": "u32", "value": "tokenCount" },
53
+ { "name": "headVec4", "type": "u32", "value": "headElements / 4" },
54
+ { "name": "headQuads", "type": "u32", "value": "headElements / 8" },
55
+ { "name": "rotaryVec4", "type": "u32", "value": "rotationPairs / 2" },
56
+ { "name": "halfRotaryVec4", "type": "u32", "value": "rotationPairs / 4" }
57
  ]
58
  },
59
+ "params_pair": {
60
  "name": "params",
 
61
  "struct": [
62
  { "name": "pairCount", "type": "u32", "value": "numel(shapes.x) / 2" },
63
+ { "name": "numHeads", "type": "u32", "value": "headCount" },
64
+ { "name": "sequenceLength", "type": "u32", "value": "tokenCount" },
65
+ { "name": "headSize", "type": "u32", "value": "headElements" },
66
+ { "name": "rotaryDim", "type": "u32", "value": "rotationPairs * 2" },
67
+ { "name": "halfRotaryDim", "type": "u32", "value": "rotationPairs" }
68
  ]
69
  },
70
+ "params_quad_cache3": {
71
  "name": "params",
 
72
  "struct": [
73
  { "name": "quadCount", "type": "u32", "value": "numel(shapes.x) / 8" },
74
+ { "name": "numHeads", "type": "u32", "value": "headCount" },
75
+ { "name": "sequenceLength", "type": "u32", "value": "tokenCount" },
76
+ { "name": "headVec4", "type": "u32", "value": "headElements / 4" },
77
+ { "name": "headQuads", "type": "u32", "value": "headElements / 8" },
78
+ { "name": "rotaryVec4", "type": "u32", "value": "rotationPairs / 2" },
79
+ { "name": "halfRotaryVec4", "type": "u32", "value": "rotationPairs / 4" },
80
  { "name": "cacheBatchStrideVec4", "type": "u32", "value": "dim(shapes.cos, 1) * dim(shapes.cos, 2) / 4" }
81
  ]
82
  },
83
+ "params_pair_cache3": {
84
  "name": "params",
 
85
  "struct": [
86
  { "name": "pairCount", "type": "u32", "value": "numel(shapes.x) / 2" },
87
+ { "name": "numHeads", "type": "u32", "value": "headCount" },
88
+ { "name": "sequenceLength", "type": "u32", "value": "tokenCount" },
89
+ { "name": "headSize", "type": "u32", "value": "headElements" },
90
+ { "name": "rotaryDim", "type": "u32", "value": "rotationPairs * 2" },
91
+ { "name": "halfRotaryDim", "type": "u32", "value": "rotationPairs" },
92
  { "name": "cacheBatchStride", "type": "u32", "value": "dim(shapes.cos, 1) * dim(shapes.cos, 2)" }
93
  ]
94
  },
95
+ "x_t": { "name": "x", "elementType": "$T" },
96
+ "cos_cache_t": { "arg": "cos", "name": "cos_cache", "elementType": "$T" },
97
+ "sin_cache_t": { "arg": "sin", "name": "sin_cache", "elementType": "$T" },
98
+ "y_t": { "name": "y", "elementType": "$T" },
99
+ "output_vec_params": {
 
 
 
 
 
 
 
 
 
 
100
  "name": "params",
101
+ "struct": [{ "name": "quadCount", "type": "u32", "value": "numel(shapes.x) / 4" }]
 
 
 
 
 
 
 
 
 
102
  }
103
  },
104
  "variants": [
105
  {
106
  "id": "rank3_cache2_pos_quad",
107
  "priority": 10,
108
+ "when": ["rank3PosContract", "quadOk", "quadDispatchOk"],
109
+ "derive": { "interleaved": "attrs.interleaved != 0", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" },
 
 
 
 
110
  "passes": [
111
  {
112
  "id": "main",
113
+ "name": "RotaryEmbedding.Rank3PosQuad",
114
  "shader": "rotary-embedding-quad.wgsl.jinja",
115
  "derive": { "rank": 3, "hasPositionIds": "true" },
116
+ "bindings": ["x", "cos_cache", "sin_cache", "position_ids", "y", "params_quad"],
117
  "dispatch": {
118
  "x": "min(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
119
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
 
126
  "id": "rank3_cache2_pos",
127
  "priority": 0,
128
  "when": ["rank3PosContract", "pairDispatchOk"],
129
+ "derive": { "interleaved": "attrs.interleaved != 0", "scalar": "dtypes.T" },
130
  "passes": [
131
  {
132
  "id": "main",
133
+ "name": "RotaryEmbedding.Rank3Pos",
134
  "shader": "rotary-embedding.wgsl.jinja",
135
  "derive": { "rank": 3, "hasPositionIds": "true" },
136
+ "bindings": ["x_t", "cos_cache_t", "sin_cache_t", "position_ids", "y_t", "params_pair"],
137
  "dispatch": {
138
  "x": "min(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
139
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
 
142
  }
143
  ]
144
  },
145
+ {
146
+ "id": "rank3_cache2_pos_output_vec4",
147
+ "priority": 20,
148
+ "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)"],
149
+ "derive": { "scalar": "dtypes.T", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" },
150
+ "passes": [
151
+ {
152
+ "id": "main",
153
+ "name": "RotaryEmbedding.Rank3PosOutputVec4",
154
+ "shader": "rotary-embedding-output-vec4.wgsl.jinja",
155
+ "derive": { "rank": 3, "hasPositionIds": "true" },
156
+ "bindings": ["x", "cos_cache_t", "sin_cache_t", "position_ids", "y", "output_vec_params"],
157
+ "dispatch": {
158
+ "x": "min(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)",
159
+ "y": "ceilDiv(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)",
160
+ "z": 1
161
+ }
162
+ }
163
+ ]
164
+ },
165
  {
166
  "id": "rank4_cache2_pos_quad",
167
  "priority": 10,
168
+ "when": ["rank4PosContract", "quadOk", "quadDispatchOk"],
169
+ "derive": { "interleaved": "attrs.interleaved != 0", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" },
 
 
 
 
170
  "passes": [
171
  {
172
  "id": "main",
173
+ "name": "RotaryEmbedding.Rank4PosQuad",
174
  "shader": "rotary-embedding-quad.wgsl.jinja",
175
  "derive": { "rank": 4, "hasPositionIds": "true" },
176
+ "bindings": ["x", "cos_cache", "sin_cache", "position_ids", "y", "params_quad"],
177
  "dispatch": {
178
  "x": "min(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
179
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
 
186
  "id": "rank4_cache2_pos",
187
  "priority": 0,
188
  "when": ["rank4PosContract", "pairDispatchOk"],
189
+ "derive": { "interleaved": "attrs.interleaved != 0", "scalar": "dtypes.T" },
190
  "passes": [
191
  {
192
  "id": "main",
193
+ "name": "RotaryEmbedding.Rank4Pos",
194
  "shader": "rotary-embedding.wgsl.jinja",
195
  "derive": { "rank": 4, "hasPositionIds": "true" },
196
+ "bindings": ["x_t", "cos_cache_t", "sin_cache_t", "position_ids", "y_t", "params_pair"],
197
  "dispatch": {
198
  "x": "min(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
199
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
 
202
  }
203
  ]
204
  },
205
+ {
206
+ "id": "rank4_cache2_pos_output_vec4",
207
+ "priority": 20,
208
+ "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)"],
209
+ "derive": { "scalar": "dtypes.T", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" },
210
+ "passes": [
211
+ {
212
+ "id": "main",
213
+ "name": "RotaryEmbedding.Rank4PosOutputVec4",
214
+ "shader": "rotary-embedding-output-vec4.wgsl.jinja",
215
+ "derive": { "rank": 4, "hasPositionIds": "true" },
216
+ "bindings": ["x", "cos_cache_t", "sin_cache_t", "position_ids", "y", "output_vec_params"],
217
+ "dispatch": {
218
+ "x": "min(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)",
219
+ "y": "ceilDiv(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)",
220
+ "z": 1
221
+ }
222
+ }
223
+ ]
224
+ },
225
  {
226
  "id": "rank3_cache3_quad",
227
  "priority": 10,
228
+ "when": ["rank3Contract", "quadOk", "quadDispatchOk"],
229
+ "derive": { "interleaved": "attrs.interleaved != 0", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" },
 
 
 
 
230
  "passes": [
231
  {
232
  "id": "main",
233
+ "name": "RotaryEmbedding.Rank3Quad",
234
  "shader": "rotary-embedding-quad.wgsl.jinja",
235
  "derive": { "rank": 3 },
236
+ "bindings": ["x", "cos_cache", "sin_cache", "y", "params_quad_cache3"],
237
  "dispatch": {
238
  "x": "min(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
239
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
 
246
  "id": "rank3_cache3",
247
  "priority": 0,
248
  "when": ["rank3Contract", "pairDispatchOk"],
249
+ "derive": { "interleaved": "attrs.interleaved != 0", "scalar": "dtypes.T" },
250
  "passes": [
251
  {
252
  "id": "main",
253
+ "name": "RotaryEmbedding.Rank3",
254
  "shader": "rotary-embedding.wgsl.jinja",
255
  "derive": { "rank": 3 },
256
+ "bindings": ["x_t", "cos_cache_t", "sin_cache_t", "y_t", "params_pair_cache3"],
257
  "dispatch": {
258
  "x": "min(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
259
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
 
262
  }
263
  ]
264
  },
265
+ {
266
+ "id": "rank3_cache3_output_vec4",
267
+ "priority": 20,
268
+ "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)"],
269
+ "derive": { "scalar": "dtypes.T", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" },
270
+ "passes": [
271
+ {
272
+ "id": "main",
273
+ "name": "RotaryEmbedding.Rank3OutputVec4",
274
+ "shader": "rotary-embedding-output-vec4.wgsl.jinja",
275
+ "derive": { "rank": 3 },
276
+ "bindings": ["x", "cos_cache_t", "sin_cache_t", "y", "output_vec_params"],
277
+ "dispatch": {
278
+ "x": "min(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)",
279
+ "y": "ceilDiv(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)",
280
+ "z": 1
281
+ }
282
+ }
283
+ ]
284
+ },
285
  {
286
  "id": "rank4_cache3_quad",
287
  "priority": 10,
288
+ "when": ["rank4Contract", "quadOk", "quadDispatchOk"],
289
+ "derive": { "interleaved": "attrs.interleaved != 0", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" },
 
 
 
 
290
  "passes": [
291
  {
292
  "id": "main",
293
+ "name": "RotaryEmbedding.Rank4Quad",
294
  "shader": "rotary-embedding-quad.wgsl.jinja",
295
  "derive": { "rank": 4 },
296
+ "bindings": ["x", "cos_cache", "sin_cache", "y", "params_quad_cache3"],
297
  "dispatch": {
298
  "x": "min(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
299
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 8), (tunables.WORKGROUP_SIZE)), 65535)",
 
306
  "id": "rank4_cache3",
307
  "priority": 0,
308
  "when": ["rank4Contract", "pairDispatchOk"],
309
+ "derive": { "interleaved": "attrs.interleaved != 0", "scalar": "dtypes.T" },
310
  "passes": [
311
  {
312
  "id": "main",
313
+ "name": "RotaryEmbedding.Rank4",
314
  "shader": "rotary-embedding.wgsl.jinja",
315
  "derive": { "rank": 4 },
316
+ "bindings": ["x_t", "cos_cache_t", "sin_cache_t", "y_t", "params_pair_cache3"],
317
  "dispatch": {
318
  "x": "min(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
319
  "y": "ceilDiv(ceilDiv((numel(shapes.x) / 2), (tunables.WORKGROUP_SIZE)), 65535)",
 
321
  }
322
  }
323
  ]
324
+ },
325
+ {
326
+ "id": "rank4_cache3_output_vec4",
327
+ "priority": 20,
328
+ "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)"],
329
+ "derive": { "scalar": "dtypes.T", "vector": "\"vec4<\" ~ dtypes.T ~ \">\"" },
330
+ "passes": [
331
+ {
332
+ "id": "main",
333
+ "name": "RotaryEmbedding.Rank4OutputVec4",
334
+ "shader": "rotary-embedding-output-vec4.wgsl.jinja",
335
+ "derive": { "rank": 4 },
336
+ "bindings": ["x", "cos_cache_t", "sin_cache_t", "y", "output_vec_params"],
337
+ "dispatch": {
338
+ "x": "min(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)",
339
+ "y": "ceilDiv(ceilDiv((numel(shapes.x) / 4), (rotaryOutputWg)), 65535)",
340
+ "z": 1
341
+ }
342
+ }
343
+ ]
344
  }
345
  ]
346
  }
build/webgpu/metadata.json CHANGED
@@ -1,31 +1,36 @@
1
  {
2
  "name": "ai.onnx.RotaryEmbedding",
3
- "id": "_ai_onnx_rotaryembedding_webgpu_5abb823",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "Ii8Ve/7TIfVkcj95zDTEbsfND6RUQkTEHTKUFQhHaEk=",
11
- "manifest.json": "OeUM+N0V8ex5eCGPP+KmT0LDzPNRCc50xnAuaYDS8o0=",
12
- "rotary-embedding-quad.wgsl.jinja": "vGNZeSSXSpPhhBz/0Wh9Z1zY5lSqp0VlCsfWmIM98CQ=",
13
- "rotary-embedding.wgsl.jinja": "DNauIDqC9Gs5N07X/rqra/Q72p6LnkcYueWx0UAaXlo=",
14
- "test.json": "394nZIOecEWid/Qm/+4eUAtAw0l+WnhmDdAuuFb96IA="
 
15
  }
16
  },
17
- "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
18
  "webgpu": {
19
- "manifestSpec": "2.0",
20
  "variants": {
21
  "rank3_cache2_pos_quad": ["rotary-embedding-quad.wgsl.jinja"],
22
  "rank3_cache2_pos": ["rotary-embedding.wgsl.jinja"],
 
23
  "rank4_cache2_pos_quad": ["rotary-embedding-quad.wgsl.jinja"],
24
  "rank4_cache2_pos": ["rotary-embedding.wgsl.jinja"],
 
25
  "rank3_cache3_quad": ["rotary-embedding-quad.wgsl.jinja"],
26
  "rank3_cache3": ["rotary-embedding.wgsl.jinja"],
 
27
  "rank4_cache3_quad": ["rotary-embedding-quad.wgsl.jinja"],
28
- "rank4_cache3": ["rotary-embedding.wgsl.jinja"]
 
29
  }
30
  }
31
  }
 
1
  {
2
  "name": "ai.onnx.RotaryEmbedding",
3
+ "id": "_ai_onnx_rotaryembedding_webgpu_2e1ee52",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "mMX9JqS6iVQEHMbCdXVsxZV48DXuX+8WkKxxiGVl0mg=",
11
+ "manifest.json": "qCFRSaYUyD+0a5+txe0/dY+zhynVlrMAnM/E74PcdHE=",
12
+ "rotary-embedding-output-vec4.wgsl.jinja": "BuRzwqt+bM1PS7pcjLCRyFjssKqVhr4W0OieFaQoU4Y=",
13
+ "rotary-embedding-quad.wgsl.jinja": "JGnYSFIWX7ouZBKzSpxgs+tLyzHghsQvrJbRKNf/7AY=",
14
+ "rotary-embedding.wgsl.jinja": "1VFzkd+u+AZPNHZ3KsKYDm5/dexHYnIy5SoFa7LdU4k=",
15
+ "test.json": "A1nlp8NNygye9SwTjT9CiPHFBUfzHAMNigbFlPYFWq8="
16
  }
17
  },
18
+ "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
19
  "webgpu": {
20
+ "manifestSpec": "2.1",
21
  "variants": {
22
  "rank3_cache2_pos_quad": ["rotary-embedding-quad.wgsl.jinja"],
23
  "rank3_cache2_pos": ["rotary-embedding.wgsl.jinja"],
24
+ "rank3_cache2_pos_output_vec4": ["rotary-embedding-output-vec4.wgsl.jinja"],
25
  "rank4_cache2_pos_quad": ["rotary-embedding-quad.wgsl.jinja"],
26
  "rank4_cache2_pos": ["rotary-embedding.wgsl.jinja"],
27
+ "rank4_cache2_pos_output_vec4": ["rotary-embedding-output-vec4.wgsl.jinja"],
28
  "rank3_cache3_quad": ["rotary-embedding-quad.wgsl.jinja"],
29
  "rank3_cache3": ["rotary-embedding.wgsl.jinja"],
30
+ "rank3_cache3_output_vec4": ["rotary-embedding-output-vec4.wgsl.jinja"],
31
  "rank4_cache3_quad": ["rotary-embedding-quad.wgsl.jinja"],
32
+ "rank4_cache3": ["rotary-embedding.wgsl.jinja"],
33
+ "rank4_cache3_output_vec4": ["rotary-embedding-output-vec4.wgsl.jinja"]
34
  }
35
  }
36
  }
build/webgpu/rotary-embedding-output-vec4.wgsl.jinja ADDED
@@ -0,0 +1,63 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
2
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
3
+ // 2D-folded flat index: gid.y carries the high bits past the dispatch's
4
+ // per-axis workgroup fold width.
5
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
6
+ if ({{ name }} >= {{ bound }}) {
7
+ return;
8
+ }{% endmacro %}
9
+ {{ env.wgsl.resourceDeclarations }}
10
+
11
+ fn read_x(index: u32) -> f32 {
12
+ return f32(x[index / 4u][index % 4u]);
13
+ }
14
+
15
+ @compute @workgroup_size({{ rotaryOutputWg }})
16
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
17
+ {{ flat_index_2d(rotaryOutputWg, "q", "params.quadCount") }}
18
+ let head = q / {{ (headElements / 4) | int }}u;
19
+ let d0 = (q % {{ (headElements / 4) | int }}u) * 4u;
20
+ let base = head * {{ headElements }}u;
21
+ {% if (rotationPairs * 2) % 4 != 0 %}
22
+ var result = x[q];
23
+ {% else %}
24
+ var result: {{ vector }};
25
+ {% endif %}
26
+ {% if rotationPairs * 2 < headElements %}
27
+ if (d0 >= {{ rotationPairs * 2 }}u) {
28
+ y[q] = x[q];
29
+ return;
30
+ }
31
+ {% endif %}
32
+ {% if rank == 3 %}
33
+ let token = (head / {{ max(1, headCount) }}u) % {{ max(1, tokenCount) }}u;
34
+ let batch = head / {{ max(1, headCount * tokenCount) }}u;
35
+ {% else %}
36
+ let token = head % {{ max(1, tokenCount) }}u;
37
+ let batch = head / {{ max(1, tokenCount * headCount) }}u;
38
+ {% endif %}
39
+ {% if hasPositionIds %}
40
+ let cache = u32(position_ids[batch * {{ tokenCount }}u + token]) * {{ rotationPairs }}u;
41
+ {% else %}
42
+ let cache = (batch * {{ tokenCount }}u + token) * {{ rotationPairs }}u;
43
+ {% endif %}
44
+ {% for lane in range(4) %}
45
+ {
46
+ let d = d0 + {{ lane }}u;
47
+ {% if (rotationPairs * 2) % 4 != 0 %}
48
+ if (d < {{ rotationPairs * 2 }}u) {
49
+ {% endif %}
50
+ let first = d < {{ rotationPairs }}u;
51
+ let pair = select(d - {{ rotationPairs }}u, d, first);
52
+ let a = read_x(base + pair);
53
+ let b = read_x(base + pair + {{ rotationPairs }}u);
54
+ let c = f32(cos_cache[cache + pair]);
55
+ let s = f32(sin_cache[cache + pair]);
56
+ result[{{ lane }}] = {{ scalar }}(select(a * s + b * c, a * c - b * s, first));
57
+ {% if (rotationPairs * 2) % 4 != 0 %}
58
+ }
59
+ {% endif %}
60
+ }
61
+ {% endfor %}
62
+ y[q] = result;
63
+ }
build/webgpu/rotary-embedding-quad.wgsl.jinja CHANGED
@@ -1,45 +1,19 @@
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
- }
31
- {%- endif %}
32
- {% endmacro %}
33
-
34
- {% if usesF16 %}enable f16;
35
- {% endif %}{{ env.wgsl.resourceDeclarations }}
36
  // One invocation owns four consecutive rotated pairs, so x, cos and sin are read as vec4s and
37
  // both rotated outputs are written as vec4s; the head size and rotary dim are multiples of
38
  // eight, so no quad straddles a head or the rotated/tail boundary. Partial rotations use the
39
  // remaining quad lanes to copy two unchanged tail vec4s. Every index below counts vec4s.
40
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
41
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
42
- {{ flat_index_2d("q", "params.quadCount", note="") }}
43
 
44
  let task = q % params.headQuads;
45
  let headFlat = q / params.headQuads;
 
1
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
2
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3
  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
4
  // per-axis workgroup fold width.
5
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
 
 
 
 
 
 
 
6
  if ({{ name }} >= {{ bound }}) {
7
  return;
8
+ }{% endmacro %}
9
+ {{ env.wgsl.resourceDeclarations }}
 
 
 
 
10
  // One invocation owns four consecutive rotated pairs, so x, cos and sin are read as vec4s and
11
  // both rotated outputs are written as vec4s; the head size and rotary dim are multiples of
12
  // eight, so no quad straddles a head or the rotated/tail boundary. Partial rotations use the
13
  // remaining quad lanes to copy two unchanged tail vec4s. Every index below counts vec4s.
14
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
15
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
16
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "q", "params.quadCount") }}
17
 
18
  let task = q % params.headQuads;
19
  let headFlat = q / params.headQuads;
build/webgpu/rotary-embedding.wgsl.jinja CHANGED
@@ -1,43 +1,17 @@
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
- }
31
- {%- endif %}
32
- {% endmacro %}
33
-
34
- {% if usesF16 %}enable f16;
35
- {% endif %}{{ env.wgsl.resourceDeclarations }}
36
  // One invocation owns both outputs of a rotated pair, so x/cos/sin are read once.
37
  // Partial rotations use the remaining pair lanes to copy two unchanged tail values.
38
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
39
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
40
- {{ flat_index_2d("p", "params.pairCount", note="") }}
41
 
42
  let tasksPerHead = params.headSize / 2u;
43
  let task = p % tasksPerHead;
 
1
+ {% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
2
+ {% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3
  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
4
  // per-axis workgroup fold width.
5
+ let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
 
 
 
 
 
 
 
6
  if ({{ name }} >= {{ bound }}) {
7
  return;
8
+ }{% endmacro %}
9
+ {{ env.wgsl.resourceDeclarations }}
 
 
 
 
10
  // One invocation owns both outputs of a rotated pair, so x/cos/sin are read once.
11
  // Partial rotations use the remaining pair lanes to copy two unchanged tail values.
12
  @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
13
  fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
14
+ {{ flat_index_2d(tunables.WORKGROUP_SIZE, "p", "params.pairCount") }}
15
 
16
  let tasksPerHead = params.headSize / 2u;
17
  let task = p % tasksPerHead;
build/webgpu/test.json CHANGED
@@ -528,7 +528,7 @@
528
  {
529
  "name": "decode_q1_rank4_head80_multihead_f16",
530
  "provenance": {
531
- "notes": "f16 autoregressive decode step (seq=1) with non-power-of-two head_dim 80 over 8 heads, the dominant inference path; exercises the rank4_cache3 kernel under the production fp16 dtype with the single-row decode shape."
532
  },
533
  "inputs": {
534
  "x": {
@@ -930,6 +930,258 @@
930
  "positionIds": { "dtype": "uint32", "shape": [1, 3], "data": { "kind": "values", "values": [7, 0, 3] } }
931
  },
932
  "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 3, 8], "tolerance": 0.0001 } }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
933
  }
934
  ]
935
  }
 
528
  {
529
  "name": "decode_q1_rank4_head80_multihead_f16",
530
  "provenance": {
531
+ "notes": "A float16 decode step (sequence length 1) with 8 heads and a non-power-of-two head_dim of 80 checks rotary embedding at the single-token generation shape."
532
  },
533
  "inputs": {
534
  "x": {
 
930
  "positionIds": { "dtype": "uint32", "shape": [1, 3], "data": { "kind": "values", "values": [7, 0, 3] } }
931
  },
932
  "outputs": { "y": { "dtype": "float32", "shape": [1, 2, 3, 8], "tolerance": 0.0001 } }
933
+ },
934
+ {
935
+ "name": "contiguous_output_float32_rank3_head100_rot94_indexedfalse",
936
+ "provenance": {
937
+ "notes": "Exercises contiguous output groups across an unaligned half-rotation, with varying dyadic inputs and cache values. Full rotations use the default rotation attribute."
938
+ },
939
+ "attrs": { "num_heads": 3, "rotary_embedding_dim": 94 },
940
+ "inputs": {
941
+ "x": {
942
+ "dtype": "float32",
943
+ "shape": [2, 3, 300],
944
+ "data": { "kind": "cycle", "values": [0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, -0.125, -0.5, -1.0, -2.0] }
945
+ },
946
+ "cos": {
947
+ "dtype": "float32",
948
+ "shape": [2, 3, 47],
949
+ "data": { "kind": "cycle", "values": [0.25, 0.5, 0.75, 1.0, -0.25, -0.5, -0.75, 0.125, 0.375, 0.625, 0.875] }
950
+ },
951
+ "sin": { "dtype": "float32", "shape": [2, 3, 47], "data": { "kind": "cycle", "values": [-0.5, 0.25, 0.75] } }
952
+ },
953
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 300], "tolerance": 0 } }
954
+ },
955
+ {
956
+ "name": "contiguous_output_float32_rank4_head100_rot100_indexedfalse",
957
+ "provenance": {
958
+ "notes": "Exercises contiguous output groups across an unaligned half-rotation, with varying dyadic inputs and cache values. Full rotations use the default rotation attribute."
959
+ },
960
+ "attrs": {},
961
+ "inputs": {
962
+ "x": {
963
+ "dtype": "float32",
964
+ "shape": [2, 3, 3, 100],
965
+ "data": { "kind": "cycle", "values": [0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, -0.125, -0.5, -1.0, -2.0] }
966
+ },
967
+ "cos": {
968
+ "dtype": "float32",
969
+ "shape": [2, 3, 50],
970
+ "data": { "kind": "cycle", "values": [0.25, 0.5, 0.75, 1.0, -0.25, -0.5, -0.75, 0.125, 0.375, 0.625, 0.875] }
971
+ },
972
+ "sin": { "dtype": "float32", "shape": [2, 3, 50], "data": { "kind": "cycle", "values": [-0.5, 0.25, 0.75] } }
973
+ },
974
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 3, 100], "tolerance": 0 } }
975
+ },
976
+ {
977
+ "name": "contiguous_output_float32_rank3_head100_rot92_indexedtrue",
978
+ "provenance": {
979
+ "notes": "Exercises contiguous output groups across an unaligned half-rotation, with varying dyadic inputs and cache values. Full rotations use the default rotation attribute."
980
+ },
981
+ "attrs": { "num_heads": 3, "rotary_embedding_dim": 92 },
982
+ "inputs": {
983
+ "x": {
984
+ "dtype": "float32",
985
+ "shape": [2, 3, 300],
986
+ "data": { "kind": "cycle", "values": [0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, -0.125, -0.5, -1.0, -2.0] }
987
+ },
988
+ "cos": {
989
+ "dtype": "float32",
990
+ "shape": [5, 46],
991
+ "data": { "kind": "cycle", "values": [0.25, 0.5, 0.75, 1.0, -0.25, -0.5, -0.75, 0.125, 0.375, 0.625, 0.875] }
992
+ },
993
+ "sin": { "dtype": "float32", "shape": [5, 46], "data": { "kind": "cycle", "values": [-0.5, 0.25, 0.75] } },
994
+ "positionIds": {
995
+ "dtype": "uint32",
996
+ "shape": [2, 3],
997
+ "data": { "kind": "values", "values": [0, 3, 1, 4, 2, 0] }
998
+ }
999
+ },
1000
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 300], "tolerance": 0 } }
1001
+ },
1002
+ {
1003
+ "name": "contiguous_output_float32_rank4_head128_rot126_indexedtrue",
1004
+ "provenance": {
1005
+ "notes": "Exercises contiguous output groups across an unaligned half-rotation, with varying dyadic inputs and cache values. Full rotations use the default rotation attribute."
1006
+ },
1007
+ "attrs": { "rotary_embedding_dim": 126 },
1008
+ "inputs": {
1009
+ "x": {
1010
+ "dtype": "float32",
1011
+ "shape": [2, 3, 3, 128],
1012
+ "data": { "kind": "cycle", "values": [0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, -0.125, -0.5, -1.0, -2.0] }
1013
+ },
1014
+ "cos": {
1015
+ "dtype": "float32",
1016
+ "shape": [5, 63],
1017
+ "data": { "kind": "cycle", "values": [0.25, 0.5, 0.75, 1.0, -0.25, -0.5, -0.75, 0.125, 0.375, 0.625, 0.875] }
1018
+ },
1019
+ "sin": { "dtype": "float32", "shape": [5, 63], "data": { "kind": "cycle", "values": [-0.5, 0.25, 0.75] } },
1020
+ "positionIds": {
1021
+ "dtype": "uint32",
1022
+ "shape": [2, 3],
1023
+ "data": { "kind": "values", "values": [0, 3, 1, 4, 2, 0] }
1024
+ }
1025
+ },
1026
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 3, 3, 128], "tolerance": 0 } }
1027
+ },
1028
+ {
1029
+ "name": "contiguous_output_mixed_special_tail_float32",
1030
+ "attrs": { "num_heads": 2, "rotary_embedding_dim": 2 },
1031
+ "inputs": {
1032
+ "x": { "dtype": "float32", "shape": [1, 3, 8], "data": { "kind": "cycle", "values": [1.0, 2.0, "NaN", -0.0] } },
1033
+ "cos": { "dtype": "float32", "shape": [1, 3, 1], "data": { "kind": "values", "values": [0.25, 0.5, 0.75] } },
1034
+ "sin": { "dtype": "float32", "shape": [1, 3, 1], "data": { "kind": "values", "values": [-0.5, 0.25, 0.75] } }
1035
+ },
1036
+ "outputs": { "y": { "dtype": "float32", "shape": [1, 3, 8], "tolerance": 0, "allowNaN": true } }
1037
+ },
1038
+ {
1039
+ "name": "contiguous_output_fold_float32_wg8",
1040
+ "provenance": {
1041
+ "notes": "The contiguous output grid crosses the dispatch fold with eight-invocation workgroups. Different input and cache periods expose incorrect row and cache offsets."
1042
+ },
1043
+ "attrs": { "num_heads": 1, "rotary_embedding_dim": 126 },
1044
+ "tunables": { "WORKGROUP_SIZE": 8 },
1045
+ "inputs": {
1046
+ "x": {
1047
+ "dtype": "float32",
1048
+ "shape": [2, 8193, 128],
1049
+ "data": { "kind": "cycle", "values": [0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, -0.125, -0.5, -1.0, -2.0] }
1050
+ },
1051
+ "cos": {
1052
+ "dtype": "float32",
1053
+ "shape": [2, 8193, 63],
1054
+ "data": { "kind": "cycle", "values": [0.25, 0.5, 0.75, 1.0, -0.25, -0.5, -0.75, 0.125, 0.375, 0.625, 0.875] }
1055
+ },
1056
+ "sin": { "dtype": "float32", "shape": [2, 8193, 63], "data": { "kind": "cycle", "values": [-0.5, 0.25, 0.75] } }
1057
+ },
1058
+ "outputs": { "y": { "dtype": "float32", "shape": [2, 8193, 128], "tolerance": 0 } }
1059
+ },
1060
+ {
1061
+ "name": "contiguous_output_float16_rank3_head100_rot94_indexedfalse",
1062
+ "provenance": {
1063
+ "notes": "Exercises contiguous output groups across an unaligned half-rotation, with varying dyadic inputs and cache values. Full rotations use the default rotation attribute."
1064
+ },
1065
+ "attrs": { "num_heads": 3, "rotary_embedding_dim": 94 },
1066
+ "inputs": {
1067
+ "x": {
1068
+ "dtype": "float16",
1069
+ "shape": [2, 3, 300],
1070
+ "data": { "kind": "cycle", "values": [0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, -0.125, -0.5, -1.0, -2.0] }
1071
+ },
1072
+ "cos": {
1073
+ "dtype": "float16",
1074
+ "shape": [2, 3, 47],
1075
+ "data": { "kind": "cycle", "values": [0.25, 0.5, 0.75, 1.0, -0.25, -0.5, -0.75, 0.125, 0.375, 0.625, 0.875] }
1076
+ },
1077
+ "sin": { "dtype": "float16", "shape": [2, 3, 47], "data": { "kind": "cycle", "values": [-0.5, 0.25, 0.75] } }
1078
+ },
1079
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 3, 300], "tolerance": 0 } }
1080
+ },
1081
+ {
1082
+ "name": "contiguous_output_float16_rank4_head100_rot100_indexedfalse",
1083
+ "provenance": {
1084
+ "notes": "Exercises contiguous output groups across an unaligned half-rotation, with varying dyadic inputs and cache values. Full rotations use the default rotation attribute."
1085
+ },
1086
+ "attrs": {},
1087
+ "inputs": {
1088
+ "x": {
1089
+ "dtype": "float16",
1090
+ "shape": [2, 3, 3, 100],
1091
+ "data": { "kind": "cycle", "values": [0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, -0.125, -0.5, -1.0, -2.0] }
1092
+ },
1093
+ "cos": {
1094
+ "dtype": "float16",
1095
+ "shape": [2, 3, 50],
1096
+ "data": { "kind": "cycle", "values": [0.25, 0.5, 0.75, 1.0, -0.25, -0.5, -0.75, 0.125, 0.375, 0.625, 0.875] }
1097
+ },
1098
+ "sin": { "dtype": "float16", "shape": [2, 3, 50], "data": { "kind": "cycle", "values": [-0.5, 0.25, 0.75] } }
1099
+ },
1100
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 3, 3, 100], "tolerance": 0 } }
1101
+ },
1102
+ {
1103
+ "name": "contiguous_output_float16_rank3_head100_rot92_indexedtrue",
1104
+ "provenance": {
1105
+ "notes": "Exercises contiguous output groups across an unaligned half-rotation, with varying dyadic inputs and cache values. Full rotations use the default rotation attribute."
1106
+ },
1107
+ "attrs": { "num_heads": 3, "rotary_embedding_dim": 92 },
1108
+ "inputs": {
1109
+ "x": {
1110
+ "dtype": "float16",
1111
+ "shape": [2, 3, 300],
1112
+ "data": { "kind": "cycle", "values": [0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, -0.125, -0.5, -1.0, -2.0] }
1113
+ },
1114
+ "cos": {
1115
+ "dtype": "float16",
1116
+ "shape": [5, 46],
1117
+ "data": { "kind": "cycle", "values": [0.25, 0.5, 0.75, 1.0, -0.25, -0.5, -0.75, 0.125, 0.375, 0.625, 0.875] }
1118
+ },
1119
+ "sin": { "dtype": "float16", "shape": [5, 46], "data": { "kind": "cycle", "values": [-0.5, 0.25, 0.75] } },
1120
+ "positionIds": {
1121
+ "dtype": "uint32",
1122
+ "shape": [2, 3],
1123
+ "data": { "kind": "values", "values": [0, 3, 1, 4, 2, 0] }
1124
+ }
1125
+ },
1126
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 3, 300], "tolerance": 0 } }
1127
+ },
1128
+ {
1129
+ "name": "contiguous_output_float16_rank4_head128_rot126_indexedtrue",
1130
+ "provenance": {
1131
+ "notes": "Exercises contiguous output groups across an unaligned half-rotation, with varying dyadic inputs and cache values. Full rotations use the default rotation attribute."
1132
+ },
1133
+ "attrs": { "rotary_embedding_dim": 126 },
1134
+ "inputs": {
1135
+ "x": {
1136
+ "dtype": "float16",
1137
+ "shape": [2, 3, 3, 128],
1138
+ "data": { "kind": "cycle", "values": [0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, -0.125, -0.5, -1.0, -2.0] }
1139
+ },
1140
+ "cos": {
1141
+ "dtype": "float16",
1142
+ "shape": [5, 63],
1143
+ "data": { "kind": "cycle", "values": [0.25, 0.5, 0.75, 1.0, -0.25, -0.5, -0.75, 0.125, 0.375, 0.625, 0.875] }
1144
+ },
1145
+ "sin": { "dtype": "float16", "shape": [5, 63], "data": { "kind": "cycle", "values": [-0.5, 0.25, 0.75] } },
1146
+ "positionIds": {
1147
+ "dtype": "uint32",
1148
+ "shape": [2, 3],
1149
+ "data": { "kind": "values", "values": [0, 3, 1, 4, 2, 0] }
1150
+ }
1151
+ },
1152
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 3, 3, 128], "tolerance": 0 } }
1153
+ },
1154
+ {
1155
+ "name": "contiguous_output_mixed_special_tail_float16",
1156
+ "attrs": { "num_heads": 2, "rotary_embedding_dim": 2 },
1157
+ "inputs": {
1158
+ "x": { "dtype": "float16", "shape": [1, 3, 8], "data": { "kind": "cycle", "values": [1.0, 2.0, "NaN", -0.0] } },
1159
+ "cos": { "dtype": "float16", "shape": [1, 3, 1], "data": { "kind": "values", "values": [0.25, 0.5, 0.75] } },
1160
+ "sin": { "dtype": "float16", "shape": [1, 3, 1], "data": { "kind": "values", "values": [-0.5, 0.25, 0.75] } }
1161
+ },
1162
+ "outputs": { "y": { "dtype": "float16", "shape": [1, 3, 8], "tolerance": 0, "allowNaN": true } }
1163
+ },
1164
+ {
1165
+ "name": "contiguous_output_fold_float16_wg8",
1166
+ "provenance": {
1167
+ "notes": "The contiguous output grid crosses the dispatch fold with eight-invocation workgroups. Different input and cache periods expose incorrect row and cache offsets."
1168
+ },
1169
+ "attrs": { "num_heads": 1, "rotary_embedding_dim": 126 },
1170
+ "tunables": { "WORKGROUP_SIZE": 8 },
1171
+ "inputs": {
1172
+ "x": {
1173
+ "dtype": "float16",
1174
+ "shape": [2, 8193, 128],
1175
+ "data": { "kind": "cycle", "values": [0.125, 0.25, 0.5, 0.75, 1.0, 1.25, 1.5, -0.125, -0.5, -1.0, -2.0] }
1176
+ },
1177
+ "cos": {
1178
+ "dtype": "float16",
1179
+ "shape": [2, 8193, 63],
1180
+ "data": { "kind": "cycle", "values": [0.25, 0.5, 0.75, 1.0, -0.25, -0.5, -0.75, 0.125, 0.375, 0.625, 0.875] }
1181
+ },
1182
+ "sin": { "dtype": "float16", "shape": [2, 8193, 63], "data": { "kind": "cycle", "values": [-0.5, 0.25, 0.75] } }
1183
+ },
1184
+ "outputs": { "y": { "dtype": "float16", "shape": [2, 8193, 128], "tolerance": 0 } }
1185
  }
1186
  ]
1187
  }