{
  "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"
          }
        ]
      }
    }
  ]
}