| { |
| "op": "ai.onnx.RotaryEmbedding", |
| "tunableSpace": { "WORKGROUP_SIZE": [128, 256] }, |
| "cases": [ |
| { |
| "name": "rotary-f32-b4s128h64", |
| "attrs": { "num_heads": 1 }, |
| "preset": "smoke", |
| "vars": { "batch": 4, "seq": 128, "headSize": 64 }, |
| "inputs": { |
| "x": { "shape": [4, 128, 64], "dtype": "float32", "dist": "normal", "seed": 115, "scale": 0.2 }, |
| "cos": { |
| "shape": [4, 128, 32], |
| "dtype": "float32", |
| "dist": "uniform", |
| "seed": 116, |
| "scale": 0.1, |
| "offset": 0.9 |
| }, |
| "sin": { "shape": [4, 128, 32], "dtype": "float32", "dist": "uniform", "seed": 117, "scale": 0.1 } |
| }, |
| "outputs": { "y": { "shape": [4, 128, 64], "dtype": "float32" } }, |
| "bench": { |
| "primary": true, |
| "metrics": [ |
| { |
| "type": "bandwidth", |
| "value": "args.batch * args.seq * args.headSize * 4 * 2 + args.batch * args.seq * (args.headSize / 2) * 4 * 2" |
| } |
| ] |
| } |
| }, |
| { |
| "name": "rotary-f32-b8s8192h128-large", |
| "attrs": { "num_heads": 1 }, |
| "vars": { "batch": 8, "seq": 8192, "headSize": 128 }, |
| "inputs": { |
| "x": { "shape": [8, 8192, 128], "dtype": "float32", "dist": "normal", "seed": 115, "scale": 0.2 }, |
| "cos": { |
| "shape": [8, 8192, 64], |
| "dtype": "float32", |
| "dist": "uniform", |
| "seed": 116, |
| "scale": 0.1, |
| "offset": 0.9 |
| }, |
| "sin": { "shape": [8, 8192, 64], "dtype": "float32", "dist": "uniform", "seed": 117, "scale": 0.1 } |
| }, |
| "outputs": { "y": { "shape": [8, 8192, 128], "dtype": "float32" } }, |
| "bench": { |
| "primary": true, |
| "metrics": [ |
| { |
| "type": "bandwidth", |
| "value": "args.batch * args.seq * args.headSize * 4 * 2 + args.batch * args.seq * (args.headSize / 2) * 4 * 2" |
| } |
| ] |
| } |
| }, |
| { |
| "name": "rotary-f32-rank4-b8h32s256h128-large", |
| "vars": { "batch": 8, "heads": 32, "seq": 256, "headSize": 128 }, |
| "inputs": { |
| "x": { "shape": [8, 32, 256, 128], "dtype": "float32", "dist": "normal", "seed": 115, "scale": 0.2 }, |
| "cos": { |
| "shape": [8, 256, 64], |
| "dtype": "float32", |
| "dist": "uniform", |
| "seed": 116, |
| "scale": 0.1, |
| "offset": 0.9 |
| }, |
| "sin": { "shape": [8, 256, 64], "dtype": "float32", "dist": "uniform", "seed": 117, "scale": 0.1 } |
| }, |
| "outputs": { "y": { "shape": [8, 32, 256, 128], "dtype": "float32" } }, |
| "bench": { |
| "primary": true, |
| "metrics": [ |
| { |
| "type": "bandwidth", |
| "value": "args.batch * args.heads * args.seq * args.headSize * 4 * 2 + args.batch * args.seq * (args.headSize / 2) * 4 * 2" |
| } |
| ] |
| } |
| }, |
| { |
| "name": "rotary-f32-decode-q1-rank4-b1h32s1h128-single-row", |
| "preset": "smoke", |
| "vars": { "batch": 1, "heads": 32, "seq": 1, "headSize": 128 }, |
| "inputs": { |
| "x": { "shape": [1, 32, 1, 128], "dtype": "float32", "dist": "normal", "seed": 115, "scale": 0.2 }, |
| "cos": { "shape": [1, 1, 64], "dtype": "float32", "dist": "uniform", "seed": 116, "scale": 0.1, "offset": 0.9 }, |
| "sin": { "shape": [1, 1, 64], "dtype": "float32", "dist": "uniform", "seed": 117, "scale": 0.1 } |
| }, |
| "outputs": { "y": { "shape": [1, 32, 1, 128], "dtype": "float32" } }, |
| "bench": { |
| "primary": true, |
| "metrics": [ |
| { |
| "type": "bandwidth", |
| "value": "args.batch * args.heads * args.seq * args.headSize * 4 * 2 + args.batch * args.seq * (args.headSize / 2) * 4 * 2" |
| } |
| ] |
| } |
| }, |
| { |
| "name": "rotary-f32-partial-rotary-tail-b8s8192h128-read-amplification", |
| "preset": "smoke", |
| "vars": { "batch": 8, "seq": 8192, "headSize": 128, "rotaryDim": 64 }, |
| "attrs": { "rotary_embedding_dim": 64, "num_heads": 1 }, |
| "inputs": { |
| "x": { "shape": [8, 8192, 128], "dtype": "float32", "dist": "normal", "seed": 115, "scale": 0.2 }, |
| "cos": { |
| "shape": [8, 8192, 32], |
| "dtype": "float32", |
| "dist": "uniform", |
| "seed": 116, |
| "scale": 0.1, |
| "offset": 0.9 |
| }, |
| "sin": { "shape": [8, 8192, 32], "dtype": "float32", "dist": "uniform", "seed": 117, "scale": 0.1 } |
| }, |
| "outputs": { "y": { "shape": [8, 8192, 128], "dtype": "float32" } }, |
| "bench": { |
| "primary": true, |
| "metrics": [ |
| { |
| "type": "bandwidth", |
| "value": "args.batch * args.seq * args.headSize * 4 * 2 + args.batch * args.seq * (args.rotaryDim / 2) * 4 * 2" |
| } |
| ] |
| } |
| }, |
| { |
| "name": "rotary-f16-b8s8192h128-large", |
| "attrs": { "num_heads": 1 }, |
| "preset": "smoke", |
| "vars": { "batch": 8, "seq": 8192, "headSize": 128 }, |
| "inputs": { |
| "x": { "shape": [8, 8192, 128], "dtype": "float16", "dist": "normal", "seed": 115, "scale": 0.2 }, |
| "cos": { |
| "shape": [8, 8192, 64], |
| "dtype": "float16", |
| "dist": "uniform", |
| "seed": 116, |
| "scale": 0.1, |
| "offset": 0.9 |
| }, |
| "sin": { "shape": [8, 8192, 64], "dtype": "float16", "dist": "uniform", "seed": 117, "scale": 0.1 } |
| }, |
| "outputs": { "y": { "shape": [8, 8192, 128], "dtype": "float16" } }, |
| "bench": { |
| "primary": true, |
| "metrics": [ |
| { |
| "type": "bandwidth", |
| "value": "args.batch * args.seq * args.headSize * 2 * 2 + args.batch * args.seq * (args.headSize / 2) * 2 * 2" |
| } |
| ] |
| } |
| }, |
| { |
| "name": "rotary-f32-multihead-b4s1024h2048-nh32-packed-layout", |
| "attrs": { "num_heads": 32 }, |
| "preset": "smoke", |
| "vars": { "batch": 4, "seq": 1024, "hidden": 2048, "headSize": 64 }, |
| "inputs": { |
| "x": { "shape": [4, 1024, 2048], "dtype": "float32", "dist": "normal", "seed": 115, "scale": 0.2 }, |
| "cos": { |
| "shape": [4, 1024, 32], |
| "dtype": "float32", |
| "dist": "uniform", |
| "seed": 116, |
| "scale": 0.1, |
| "offset": 0.9 |
| }, |
| "sin": { "shape": [4, 1024, 32], "dtype": "float32", "dist": "uniform", "seed": 117, "scale": 0.1 } |
| }, |
| "outputs": { "y": { "shape": [4, 1024, 2048], "dtype": "float32" } }, |
| "bench": { |
| "primary": true, |
| "metrics": [ |
| { |
| "type": "bandwidth", |
| "value": "args.batch * args.seq * args.hidden * 4 * 2 + args.batch * args.seq * (args.headSize / 2) * 4 * 2" |
| } |
| ] |
| } |
| }, |
| { |
| "name": "rotary-f32-rank3-partial-rd126-tail-copy-full-read-amplification", |
| "preset": "stress", |
| "vars": { "batch": 8, "seq": 8192, "headSize": 128, "rotaryDim": 126 }, |
| "attrs": { "rotary_embedding_dim": 126, "num_heads": 1 }, |
| "inputs": { |
| "x": { "shape": [8, 8192, 128], "dtype": "float32", "dist": "normal", "seed": 215, "scale": 0.2 }, |
| "cos": { |
| "shape": [8, 8192, 63], |
| "dtype": "float32", |
| "dist": "uniform", |
| "seed": 216, |
| "scale": 0.1, |
| "offset": 0.9 |
| }, |
| "sin": { "shape": [8, 8192, 63], "dtype": "float32", "dist": "uniform", "seed": 217, "scale": 0.1 } |
| }, |
| "outputs": { "y": { "shape": [8, 8192, 128], "dtype": "float32", "dist": "empty" } }, |
| "bench": { |
| "primary": true, |
| "metrics": [ |
| { |
| "type": "bandwidth", |
| "value": "args.batch * args.seq * args.headSize * 4 * 2 + args.batch * args.seq * (args.rotaryDim / 2) * 4 * 2" |
| } |
| ] |
| } |
| }, |
| { |
| "name": "rotary-f32-rank4-partial-rd126-tail-copy-b4h16s512h128", |
| "preset": "stress", |
| "vars": { "batch": 4, "heads": 16, "seq": 512, "headSize": 128, "rotaryDim": 126 }, |
| "attrs": { "rotary_embedding_dim": 126 }, |
| "inputs": { |
| "x": { "shape": [4, 16, 512, 128], "dtype": "float32", "dist": "normal", "seed": 315, "scale": 0.2 }, |
| "cos": { |
| "shape": [4, 512, 63], |
| "dtype": "float32", |
| "dist": "uniform", |
| "seed": 316, |
| "scale": 0.1, |
| "offset": 0.9 |
| }, |
| "sin": { "shape": [4, 512, 63], "dtype": "float32", "dist": "uniform", "seed": 317, "scale": 0.1 } |
| }, |
| "outputs": { "y": { "shape": [4, 16, 512, 128], "dtype": "float32", "dist": "empty" } }, |
| "bench": { |
| "primary": true, |
| "metrics": [ |
| { |
| "type": "bandwidth", |
| "value": "args.batch * args.heads * args.seq * args.headSize * 4 * 2 + args.batch * args.seq * (args.rotaryDim / 2) * 4 * 2" |
| } |
| ] |
| } |
| } |
| ] |
| } |
|
|
