sync 2e7068faf55e

Browse files

Files changed (6) hide show

README.md +72 -0
build/webgpu/bench.json +74 -0
build/webgpu/gemma-rotary-embedding.wgsl.jinja +42 -0
build/webgpu/manifest.json +210 -0
build/webgpu/metadata.json +18 -0
build/webgpu/test.json +203 -0

README.md CHANGED Viewed

@@ -1,3 +1,75 @@
 ---
 license: apache-2.0
 ---

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.GemmaRotaryEmbedding
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Fuses the Gemma rotary-embedding tail: computes `sin` and `cos` from float32 `emb`, casts them to float16, then evaluates `q * cos + q_rot * sin` and the corresponding expression for `k`. `emb` has shape `(batch, seq, dim)` and is broadcast over the head axis of the `(batch, heads, seq, dim)` operands. Each product is rounded to float16 before the addition.
+See the [ONNX Runtime `GemmaRotaryEmbedding` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.GemmaRotaryEmbedding) for the reference semantics.
+## Inputs
+| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- | --- |
+| `emb` | `embT` | `U` | `float32` | `3` | — | Rotary angles with shape `(batch_size, seq_len, dim)`, shared by every head. | required |
+| `q` | `qT` | `T` | same as logical dtype | `4` | — | Query state with shape `(batch_size, num_heads, seq_len, dim)`. | required |
+| `q_rot` | `qRotT` | `T` | same as logical dtype | `4` | — | Half-rotated query state, same shape as `q`. | required |
+| `k` | `kT` | `T` | same as logical dtype | `4` | — | Key state, same shape as `q`. | required |
+| `k_rot` | `kRotT` | `T` | same as logical dtype | `4` | — | Half-rotated key state, same shape as `q`. | required |
+## Outputs
+| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `output1` | `output1T` | `T` | same as `q` | same as `q` | Rotary-embedded query, same shape as `q`. | required |
+| `output2` | `output2T` | `T` | same as `q` | same as `q` | Rotary-embedded key, same shape as `q`. | required |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float16` |
+| `U` | `float32` |
+## Device requirements
+Every implementation variant requires `shader-f16`; the package has no variant-level fallback without that capability.
+## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
+- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
+- [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
+- [`gemma-rotary-embedding.wgsl.jinja`](build/webgpu/gemma-rotary-embedding.wgsl.jinja)
+## Use with `@huggingface/kernels`
+The loader derives every required output's shape and logical dtype from the manifest contract and this call.
+It then allocates the result tensors automatically.
+The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+Replace each `*Data` placeholder with a typed array containing the corresponding input data.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/com.microsoft.GemmaRotaryEmbedding", { version: 1 });
+const { output1T, output2T } = await kernel({
+  embT: { data: embTData, shape: [1, 2, 4] },
+  qT: { data: qTData, shape: [1, 1, 2, 4] },
+  qRotT: { data: qRotTData, shape: [1, 1, 2, 4] },
+  kT: { data: kTData, shape: [1, 1, 2, 4] },
+  kRotT: { data: kRotTData, shape: [1, 1, 2, 4] },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,74 @@

+{
+  "op": "com.microsoft.GemmaRotaryEmbedding",
+  "tunableSpace": { "workgroupSize": [64, 128, 256] },
+  "cases": [
+    {
+      "name": "gemma-rotary-f16-b1-h16-s2048-d256",
+      "preset": "smoke",
+      "vars": { "dtype": "float16" },
+      "inputs": {
+        "embT": { "shape": [1, 2048, 256], "dtype": "float32", "dist": "normal", "seed": 8100, "scale": 2 },
+        "qT": { "shape": [1, 16, 2048, 256], "dtype": "float16", "dist": "normal", "seed": 8101, "scale": 1 },
+        "qRotT": { "shape": [1, 16, 2048, 256], "dtype": "float16", "dist": "normal", "seed": 8102, "scale": 1 },
+        "kT": { "shape": [1, 16, 2048, 256], "dtype": "float16", "dist": "normal", "seed": 8103, "scale": 1 },
+        "kRotT": { "shape": [1, 16, 2048, 256], "dtype": "float16", "dist": "normal", "seed": 8104, "scale": 1 }
+      },
+      "outputs": {
+        "output1T": { "shape": [1, 16, 2048, 256], "dtype": "float16" },
+        "output2T": { "shape": [1, 16, 2048, 256], "dtype": "float16" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "6 * 8388608 * dtypeBytes(args.dtype) + 524288 * 4" }] }
+    },
+    {
+      "name": "gemma-rotary-f16-b1-h16-s1-d256",
+      "preset": "smoke",
+      "vars": { "dtype": "float16" },
+      "inputs": {
+        "embT": { "shape": [1, 1, 256], "dtype": "float32", "dist": "normal", "seed": 8110, "scale": 2 },
+        "qT": { "shape": [1, 16, 1, 256], "dtype": "float16", "dist": "normal", "seed": 8111, "scale": 1 },
+        "qRotT": { "shape": [1, 16, 1, 256], "dtype": "float16", "dist": "normal", "seed": 8112, "scale": 1 },
+        "kT": { "shape": [1, 16, 1, 256], "dtype": "float16", "dist": "normal", "seed": 8113, "scale": 1 },
+        "kRotT": { "shape": [1, 16, 1, 256], "dtype": "float16", "dist": "normal", "seed": 8114, "scale": 1 }
+      },
+      "outputs": {
+        "output1T": { "shape": [1, 16, 1, 256], "dtype": "float16" },
+        "output2T": { "shape": [1, 16, 1, 256], "dtype": "float16" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "6 * 4096 * dtypeBytes(args.dtype) + 256 * 4" }] }
+    },
+    {
+      "name": "gemma-rotary-f16-b4-h8-s512-d256",
+      "preset": "model",
+      "vars": { "dtype": "float16" },
+      "inputs": {
+        "embT": { "shape": [4, 512, 256], "dtype": "float32", "dist": "normal", "seed": 8130, "scale": 2 },
+        "qT": { "shape": [4, 8, 512, 256], "dtype": "float16", "dist": "normal", "seed": 8131, "scale": 1 },
+        "qRotT": { "shape": [4, 8, 512, 256], "dtype": "float16", "dist": "normal", "seed": 8132, "scale": 1 },
+        "kT": { "shape": [4, 8, 512, 256], "dtype": "float16", "dist": "normal", "seed": 8133, "scale": 1 },
+        "kRotT": { "shape": [4, 8, 512, 256], "dtype": "float16", "dist": "normal", "seed": 8134, "scale": 1 }
+      },
+      "outputs": {
+        "output1T": { "shape": [4, 8, 512, 256], "dtype": "float16" },
+        "output2T": { "shape": [4, 8, 512, 256], "dtype": "float16" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "6 * 4194304 * dtypeBytes(args.dtype) + 524288 * 4" }] }
+    },
+    {
+      "name": "gemma-rotary-f16-b2-h16-s512-d250-scalar",
+      "preset": "model",
+      "vars": { "dtype": "float16" },
+      "inputs": {
+        "embT": { "shape": [2, 512, 250], "dtype": "float32", "dist": "normal", "seed": 8140, "scale": 2 },
+        "qT": { "shape": [2, 16, 512, 250], "dtype": "float16", "dist": "normal", "seed": 8141, "scale": 1 },
+        "qRotT": { "shape": [2, 16, 512, 250], "dtype": "float16", "dist": "normal", "seed": 8142, "scale": 1 },
+        "kT": { "shape": [2, 16, 512, 250], "dtype": "float16", "dist": "normal", "seed": 8143, "scale": 1 },
+        "kRotT": { "shape": [2, 16, 512, 250], "dtype": "float16", "dist": "normal", "seed": 8144, "scale": 1 }
+      },
+      "outputs": {
+        "output1T": { "shape": [2, 16, 512, 250], "dtype": "float16" },
+        "output2T": { "shape": [2, 16, 512, 250], "dtype": "float16" }
+      },
+      "bench": { "metrics": [{ "type": "bandwidth", "value": "6 * 4096000 * dtypeBytes(args.dtype) + 256000 * 4" }] }
+    }
+  ]
+}

build/webgpu/gemma-rotary-embedding.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,42 @@

+enable f16;
+{{ env.wgsl.resourceDeclarations }}
+// output1 = q * cos(emb) + q_rot * sin(emb); output2 applies the same
+// rotation to K. The operator is elementwise over (batch, heads, seq, dim),
+// while emb omits the head axis and is shared by every head.
+const WG: u32 = {{ workgroupSize }}u;
+{% if vec4 %}
+const ZERO: {{ vector }} = {{ vector }}(0.0);
+{% else %}
+const ZERO: {{ scalar }} = {{ scalar }}(0.0);
+{% endif %}
+@compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>,
+        @builtin(num_workgroups) nwg: vec3<u32>) {
+  // 2D-folded flat index: gid.y carries the high bits past the
+  // maxComputeWorkgroupsPerDimension limit. Reduces to gid.x when nwg.y == 1.
+  let index = gid.x + gid.y * nwg.x * WG;
+  if (index >= params.count) {
+    return;
+  }
+  // (batch, head, seq, dim) -> (batch, seq, dim): divide out the head-major stride and
+  // keep the remainder, which is exactly the (seq, dim) offset shared by every head.
+  let emb_index = (index / params.headSeqDim) * params.seqDim + index % params.seqDim;
+  // emb is float32 for every T. Compute sine and cosine at full width, then
+  // round them to T before multiplication as required by the operator graph.
+  let angle = emb[emb_index];
+{% if vec4 %}
+  let sin_val = {{ vector }}(sin(angle));
+  let cos_val = {{ vector }}(cos(angle));
+{% else %}
+  let sin_val = {{ scalar }}(sin(angle));
+  let cos_val = {{ scalar }}(cos(angle));
+{% endif %}
+  // Round both products to T before addition. Explicit fma(..., 0) prevents a
+  // backend from contracting the final add and retaining wider intermediates.
+  output1[index] = fma(q[index], cos_val, ZERO) + fma(q_rot[index], sin_val, ZERO);
+  output2[index] = fma(k[index], cos_val, ZERO) + fma(k_rot[index], sin_val, ZERO);
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,210 @@

+{
+  "domain": "com.microsoft",
+  "name": "GemmaRotaryEmbedding",
+  "sinceVersion": 1,
+  "description": "Fuses the Gemma rotary-embedding tail: computes `sin` and `cos` from float32 `emb`, casts them to float16, then evaluates `q * cos + q_rot * sin` and the corresponding expression for `k`. `emb` has shape `(batch, seq, dim)` and is broadcast over the head axis of the `(batch, heads, seq, dim)` operands. Each product is rounded to float16 before the addition.",
+  "inputs": [
+    {
+      "role": "emb",
+      "dtype": "U",
+      "rank": 3,
+      "description": "Rotary angles with shape `(batch_size, seq_len, dim)`, shared by every head."
+    },
+    {
+      "role": "q",
+      "dtype": "T",
+      "rank": 4,
+      "description": "Query state with shape `(batch_size, num_heads, seq_len, dim)`."
+    },
+    { "role": "q_rot", "dtype": "T", "rank": 4, "description": "Half-rotated query state, same shape as `q`." },
+    { "role": "k", "dtype": "T", "rank": 4, "description": "Key state, same shape as `q`." },
+    { "role": "k_rot", "dtype": "T", "rank": 4, "description": "Half-rotated key state, same shape as `q`." }
+  ],
+  "outputs": [
+    {
+      "role": "output1",
+      "dtype": "T",
+      "rank": "ranks.qT",
+      "shape": "shapes.qT",
+      "description": "Rotary-embedded query, same shape as `q`."
+    },
+    {
+      "role": "output2",
+      "dtype": "T",
+      "rank": "ranks.qT",
+      "shape": "shapes.qT",
+      "description": "Rotary-embedded key, same shape as `q`."
+    }
+  ],
+  "typeConstraints": { "T": ["float16"], "U": ["float32"] },
+  "args": {
+    "embT": { "kind": "tensor", "semantic": "emb", "role": "input", "dtype": "float32" },
+    "qT": { "kind": "tensor", "semantic": "q", "role": "input" },
+    "qRotT": { "kind": "tensor", "semantic": "q_rot", "role": "input" },
+    "kT": { "kind": "tensor", "semantic": "k", "role": "input" },
+    "kRotT": { "kind": "tensor", "semantic": "k_rot", "role": "input" },
+    "output1T": { "kind": "tensor", "semantic": "output1", "role": "output" },
+    "output2T": { "kind": "tensor", "semantic": "output2", "role": "output" }
+  },
+  "tunables": { "workgroupSize": 256 },
+  "derive": {
+    "batchSize": "dim(shapes.qT, 0)",
+    "numHeads": "dim(shapes.qT, 1)",
+    "seqLen": "dim(shapes.qT, 2)",
+    "headDim": "dim(shapes.qT, 3)",
+    "sameAsQ": "ranks.qRotT == 4 and ranks.kT == 4 and ranks.kRotT == 4 and sameShape(shapes.qRotT, shapes.qT) and sameShape(shapes.kT, shapes.qT) and sameShape(shapes.kRotT, shapes.qT)",
+    "dtypesMatch": "tensorDtypes.qRotT == tensorDtypes.qT and tensorDtypes.kT == tensorDtypes.qT and tensorDtypes.kRotT == tensorDtypes.qT and tensorDtypes.output1T == tensorDtypes.qT and tensorDtypes.output2T == tensorDtypes.qT",
+    "contract": "ranks.qT == 4 and ranks.embT == 3 and sameAsQ and dtypesMatch and tensorDtypes.qT == \"float16\" and f16Ok(dtypes.T) and tensorDtypes.embT == \"float32\" and dim(shapes.embT, 0) == batchSize and dim(shapes.embT, 1) == seqLen and dim(shapes.embT, 2) == headDim and sameShape(shapes.output1T, shapes.qT) and sameShape(shapes.output2T, shapes.qT)",
+    "vec4Ok": "headDim % 4 == 0"
+  },
+  "bindingSets": {
+    "scalar": [
+      {
+        "name": "emb",
+        "arg": "embT",
+        "semantic": "emb",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "f32"
+      },
+      { "name": "q", "arg": "qT", "semantic": "q", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      {
+        "name": "q_rot",
+        "arg": "qRotT",
+        "semantic": "q_rot",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$scalar"
+      },
+      { "name": "k", "arg": "kT", "semantic": "k", "buffer": { "type": "read-only-storage" }, "elementType": "$scalar" },
+      {
+        "name": "k_rot",
+        "arg": "kRotT",
+        "semantic": "k_rot",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$scalar"
+      },
+      {
+        "name": "output1",
+        "arg": "output1T",
+        "semantic": "output1",
+        "buffer": { "type": "storage" },
+        "elementType": "$scalar"
+      },
+      {
+        "name": "output2",
+        "arg": "output2T",
+        "semantic": "output2",
+        "buffer": { "type": "storage" },
+        "elementType": "$scalar"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "count", "type": "u32", "value": "numel(shapes.qT)" },
+            { "name": "seqDim", "type": "u32", "value": "(seqLen * headDim)" },
+            { "name": "headSeqDim", "type": "u32", "value": "(numHeads * seqLen * headDim)" }
+          ]
+        }
+      }
+    ],
+    "vec4": [
+      {
+        "name": "emb",
+        "arg": "embT",
+        "semantic": "emb",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "vec4<f32>"
+      },
+      { "name": "q", "arg": "qT", "semantic": "q", "buffer": { "type": "read-only-storage" }, "elementType": "$vector" },
+      {
+        "name": "q_rot",
+        "arg": "qRotT",
+        "semantic": "q_rot",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$vector"
+      },
+      { "name": "k", "arg": "kT", "semantic": "k", "buffer": { "type": "read-only-storage" }, "elementType": "$vector" },
+      {
+        "name": "k_rot",
+        "arg": "kRotT",
+        "semantic": "k_rot",
+        "buffer": { "type": "read-only-storage" },
+        "elementType": "$vector"
+      },
+      {
+        "name": "output1",
+        "arg": "output1T",
+        "semantic": "output1",
+        "buffer": { "type": "storage" },
+        "elementType": "$vector"
+      },
+      {
+        "name": "output2",
+        "arg": "output2T",
+        "semantic": "output2",
+        "buffer": { "type": "storage" },
+        "elementType": "$vector"
+      },
+      {
+        "name": "params",
+        "semantic": "kernel.params",
+        "buffer": { "type": "uniform" },
+        "struct": {
+          "name": "Params",
+          "fields": [
+            { "name": "count", "type": "u32", "value": "numel(shapes.qT) / 4" },
+            { "name": "seqDim", "type": "u32", "value": "(seqLen * headDim) / 4" },
+            { "name": "headSeqDim", "type": "u32", "value": "(numHeads * seqLen * headDim) / 4" }
+          ]
+        }
+      }
+    ]
+  },
+  "variants": [
+    {
+      "id": "vec4",
+      "priority": 10,
+      "requires": { "features": ["shader-f16"] },
+      "when": ["contract", "vec4Ok", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
+      "constants": {
+        "vec4": true,
+        "scalar": "dtypes.T",
+        "vector": "\"vec4<f16>\"",
+        "workgroupSize": "tunables.workgroupSize"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmaRotaryEmbedding.Vec4",
+          "shader": "gemma-rotary-embedding.wgsl.jinja",
+          "bindings": "vec4",
+          "dispatch": { "threads": "numel(shapes.qT) / 4", "workgroupSize": "constants.workgroupSize" }
+        }
+      ]
+    },
+    {
+      "id": "scalar",
+      "priority": 0,
+      "requires": { "features": ["shader-f16"] },
+      "when": ["contract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
+      "constants": {
+        "vec4": false,
+        "scalar": "dtypes.T",
+        "vector": "\"vec4<f16>\"",
+        "workgroupSize": "tunables.workgroupSize"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "GemmaRotaryEmbedding",
+          "shader": "gemma-rotary-embedding.wgsl.jinja",
+          "bindings": "scalar",
+          "dispatch": { "threads": "max(1, numel(shapes.qT))", "workgroupSize": "constants.workgroupSize" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,18 @@

+{
+  "name": "com.microsoft.GemmaRotaryEmbedding",
+  "id": "_com_microsoft_gemmarotaryembedding_webgpu_c64044f",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "/3JoilqqNUy3rMDPsbi0npsxk9oZgl6Db/2k+439H0c=",
+      "gemma-rotary-embedding.wgsl.jinja": "w+t56Q+FUBau7k0eWwrulSaGWFEKYch4aCunpG3sbBs=",
+      "manifest.json": "FK8aZrRWKnUBl9hUi7MrNAX7nE9/6pz89Wh/AL5kCDQ=",
+      "test.json": "KnP1fOvNgEvi3e9lrItvga2K6sLswyE5trl1BhX72Js="
+    }
+  },
+  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
+  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.GemmaRotaryEmbedding" }
+}

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,203 @@

+{
+  "op": "com.microsoft.GemmaRotaryEmbedding",
+  "cases": [
+    {
+      "name": "f16_vec4",
+      "provenance": {
+        "notes": "float16 operands with a float32 emb, which is the type pair ONNX Runtime constrains this operator to. Sin and cos are computed at float32 and rounded once."
+      },
+      "inputs": {
+        "embT": {
+          "dtype": "float32",
+          "shape": [2, 4, 8],
+          "data": { "kind": "fillFloat32", "scale": 2.5, "sinStep": 0.23, "cosStep": 0.37 }
+        },
+        "qT": {
+          "dtype": "float16",
+          "shape": [2, 2, 4, 8],
+          "data": {
+            "kind": "fillFloat32",
+            "scale": 0.8,
+            "sinStep": 0.16999999999999998,
+            "cosStep": 0.29000000000000004
+          }
+        },
+        "qRotT": {
+          "dtype": "float16",
+          "shape": [2, 2, 4, 8],
+          "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.43, "cosStep": 0.13 }
+        },
+        "kT": {
+          "dtype": "float16",
+          "shape": [2, 2, 4, 8],
+          "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.35, "cosStep": 0.19 }
+        },
+        "kRotT": {
+          "dtype": "float16",
+          "shape": [2, 2, 4, 8],
+          "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.47, "cosStep": 0.25 }
+        }
+      },
+      "outputs": {
+        "output1T": { "dtype": "float16", "shape": [2, 2, 4, 8], "tolerance": 0.005 },
+        "output2T": { "dtype": "float16", "shape": [2, 2, 4, 8], "tolerance": 0.005 }
+      }
+    },
+    {
+      "name": "f16_scalar",
+      "provenance": { "notes": "float16 on the scalar variant, dim 5." },
+      "inputs": {
+        "embT": {
+          "dtype": "float32",
+          "shape": [2, 3, 5],
+          "data": { "kind": "fillFloat32", "scale": 2.5, "sinStep": 0.24000000000000002, "cosStep": 0.38 }
+        },
+        "qT": {
+          "dtype": "float16",
+          "shape": [2, 2, 3, 5],
+          "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.18, "cosStep": 0.30000000000000004 }
+        },
+        "qRotT": {
+          "dtype": "float16",
+          "shape": [2, 2, 3, 5],
+          "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.44, "cosStep": 0.14 }
+        },
+        "kT": {
+          "dtype": "float16",
+          "shape": [2, 2, 3, 5],
+          "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.36, "cosStep": 0.2 }
+        },
+        "kRotT": {
+          "dtype": "float16",
+          "shape": [2, 2, 3, 5],
+          "data": { "kind": "fillFloat32", "scale": 0.8, "sinStep": 0.48, "cosStep": 0.26 }
+        }
+      },
+      "outputs": {
+        "output1T": { "dtype": "float16", "shape": [2, 2, 3, 5], "tolerance": 0.005 },
+        "output2T": { "dtype": "float16", "shape": [2, 2, 3, 5], "tolerance": 0.005 }
+      }
+    },
+    {
+      "name": "pinned_f16_round_before_add",
+      "provenance": {
+        "notes": "Pins the round-to-T-before-add rule at float16. Every element is a near-cancellation: the two products are around 100 in magnitude and their sum is under 1, so rounding each product to float16 before the add moves the result by 150 to 900 ulps, while a backend that contracted the multiply into an fma and kept a float32 product would land somewhere else entirely. The angles were chosen so their float16 sin and cos survive several ulps of perturbation, so this does not test the GPU's transcendental."
+      },
+      "inputs": {
+        "embT": {
+          "dtype": "float32",
+          "shape": [1, 2, 4],
+          "data": {
+            "kind": "values",
+            "values": [0.896741, -0.7510284, 2.9392502, 2.4113605, 2.7308347, -2.6883833, -0.7908616, -0.4171633]
+          }
+        },
+        "qT": {
+          "dtype": "float16",
+          "shape": [1, 1, 2, 4],
+          "data": { "kind": "values", "values": [-105.625, 248.375, 232.875, 109.75, 214.75, 205.75, -98.5625, -142.0] }
+        },
+        "qRotT": {
+          "dtype": "float16",
+          "shape": [1, 1, 2, 4],
+          "data": { "kind": "values", "values": [84.4375, 266.0, 1135.0, 122.25, 493.25, -421.0, -98.0625, -322.25] }
+        },
+        "kT": {
+          "dtype": "float16",
+          "shape": [1, 1, 2, 4],
+          "data": {
+            "kind": "values",
+            "values": [-239.625, 217.875, 249.125, -202.125, -127.8125, -135.75, 118.4375, 208.75]
+          }
+        },
+        "kRotT": {
+          "dtype": "float16",
+          "shape": [1, 1, 2, 4],
+          "data": { "kind": "values", "values": [-165.25, 194.25, -411.75, 281.25, 252.375, 151.25, -194.0, -455.5] }
+        }
+      },
+      "outputs": {
+        "output1T": {
+          "dtype": "float16",
+          "shape": [1, 1, 2, 4],
+          "data": { "kind": "values", "values": [0.0, -0.125, -0.125, -0.1875, 0.125, -0.75, 0.375, 0.875] },
+          "tolerance": 0.005
+        },
+        "output2T": {
+          "dtype": "float16",
+          "shape": [1, 1, 2, 4],
+          "data": { "kind": "values", "values": [-278.5, 26.625, -326.75, 338.25, 218.0, 55.8125, 221.125, 375.5] },
+          "tolerance": 0.005
+        }
+      }
+    },
+    {
+      "name": "pinned_batch2_emb_broadcast",
+      "provenance": {
+        "notes": "Hand-computed by an independent model that Unsqueezes emb to (batch, 1, seq, dim) and broadcasts with real array shapes, never forming a flat index -- which is the step ONNX Runtime's own unit test gets wrong: it indexes emb by num_heads where the CUDA kernel uses seq_len, and its only case has batch_size 1, where the wrong term is multiplied by zero. heads (3) and seq_len (2) differ here and batch is 2, so the two disagree."
+      },
+      "inputs": {
+        "embT": {
+          "dtype": "float32",
+          "shape": [2, 2, 4],
+          "data": {
+            "kind": "values",
+            "values": [-2.75, -2.383333, -2.016667, -1.65, -1.283333, -0.916667, -0.55, -0.183333, 0.183333, 0.55, 0.916667, 1.283333, 1.65, 2.016667, 2.383333, 2.75]
+          }
+        },
+        "qT": {
+          "dtype": "float16",
+          "shape": [2, 3, 2, 4],
+          "data": {
+            "kind": "values",
+            "values": [0.686035, -0.222046, -0.056, -0.728027, 0.777832, -0.014, 0.219971, -0.970215, 0.512207, -0.260986, -0.564941, 0.10498, -0.898926, -0.805176, -0.737793, 0.758789, 0.099976, -0.890137, 0.366943, 0.452881, -0.427979, 0.220947, -0.284912, -0.829102, 0.36499, 0.771973, 0.406006, 0.899902, 0.604004, 0.124023, 0.341064, 0.088013, 0.03299, -0.908203, 0.726074, 0.077026, -0.650879, -0.657227, 0.290039, 0.905762, 0.600098, -0.36499, -0.568848, 0.399902, -0.366943, 0.542969, 0.333984, -0.650879]
+          }
+        },
+        "qRotT": {
+          "dtype": "float16",
+          "shape": [2, 3, 2, 4],
+          "data": {
+            "kind": "values",
+            "values": [0.368896, 0.120972, 0.058014, 0.370117, -0.188965, 0.052002, 0.64209, -0.974121, -0.845215, -0.566895, 0.748047, 0.915039, -0.520996, -0.717773, 0.74707, -0.085022, -0.570801, 0.869141, 0.112976, -0.671875, -0.346924, -0.794922, -0.644043, -0.182007, 0.862793, -0.569824, -0.189941, -0.253906, -0.791016, -0.698242, 0.540039, 0.003, -0.426025, 0.733887, -0.112, -0.364014, -0.724121, 0.167969, -0.540039, 0.088989, 0.011002, 0.203003, -0.213989, -0.292969, -0.151001, 0.895996, 0.682129, 0.910156]
+          }
+        },
+        "kT": {
+          "dtype": "float16",
+          "shape": [2, 3, 2, 4],
+          "data": {
+            "kind": "values",
+            "values": [-0.395996, -0.052002, 0.120972, -0.955078, -0.782227, 0.232056, -0.702148, 0.06897, 0.581055, 0.923828, -0.251953, 0.258057, -0.813965, -0.799805, 0.491943, -0.418945, 0.583008, -0.597168, 0.333008, 0.731934, 0.852051, -0.009003, -0.476074, 0.206055, -0.527832, 0.714844, 0.948242, -0.370117, 0.729004, 0.481934, -0.527832, -0.467041, -0.711914, 0.515137, 0.551758, -0.099976, -0.858887, -0.544922, 0.580078, 0.621094, -0.52002, -0.331055, 0.326904, -0.25293, 0.430908, -0.585938, -0.588867, -0.177002]
+          }
+        },
+        "kRotT": {
+          "dtype": "float16",
+          "shape": [2, 3, 2, 4],
+          "data": {
+            "kind": "values",
+            "values": [0.106018, 0.551758, 0.952148, -0.405029, 0.814941, -0.312012, -0.819824, 0.955078, -0.234009, 0.085022, 0.008003, 0.471924, 0.798828, -0.641113, 0.679199, -0.859863, -0.974121, -0.881836, -0.122986, -0.63623, 0.18396, 0.156006, 0.094971, 0.619141, -0.366943, 0.041992, 0.945801, -0.298096, 0.036011, 0.471924, 0.223999, 0.214966, 0.890137, 0.748047, 0.523926, 0.665039, 0.189941, 0.49292, -0.488037, -0.437012, 0.37207, -0.284912, 0.765137, 0.020996, -0.462891, -0.111023, -0.63916, 0.460938]
+          }
+        }
+      },
+      "outputs": {
+        "output1T": {
+          "dtype": "float16",
+          "shape": [2, 3, 2, 4],
+          "data": {
+            "kind": "values",
+            "values": [-0.774902, 0.078064, -0.028198, -0.311523, 0.401855, -0.049774, -0.147949, -0.776367, -0.150879, 0.579102, -0.431152, -0.920898, 0.244629, 0.07959, -1.019531, 0.761719, 0.125366, 0.048828, -0.260254, 0.634277, 0.211426, 0.765137, 0.093506, -0.782227, 0.516113, 0.360596, 0.096313, 0.011597, -0.836426, -0.683594, 0.123657, -0.0802, -0.045197, -0.390869, 0.352783, -0.327393, -0.670898, 0.435059, -0.582031, -0.803223, 0.592285, -0.2052, -0.516113, -0.167725, -0.12146, 0.574219, 0.22644, 0.949219]
+          },
+          "tolerance": 0.005
+        },
+        "output2T": {
+          "dtype": "float16",
+          "shape": [2, 3, 2, 4],
+          "data": {
+            "kind": "values",
+            "values": [0.325439, -0.341553, -0.911621, 0.479492, -1.003906, 0.388672, -0.17041, -0.106262, -0.447754, -0.729492, 0.10144, -0.490967, -0.99707, 0.022217, 0.064697, -0.255371, -0.167236, 1.040039, -0.032593, 0.576172, 0.065063, -0.129272, -0.455322, 0.089783, -0.585938, 0.631348, 1.328125, -0.390869, -0.021759, 0.218018, 0.537109, 0.513672, -0.538086, 0.830078, 0.751465, 0.609375, 0.257324, 0.679688, -0.756836, -0.741211, -0.443359, -0.431152, 0.805664, -0.051575, -0.495605, 0.152588, -0.011963, 0.339355]
+          },
+          "tolerance": 0.005
+        }
+      }
+    }
+  ]
+}