sync 91d990483a17

Browse files

Files changed (6) hide show

README.md +15 -11
build/webgpu/bench.json +0 -1
build/webgpu/linear-attention-gate.wgsl.jinja +2 -3
build/webgpu/manifest.json +62 -154
build/webgpu/metadata.json +15 -7
build/webgpu/test.json +3 -4

README.md CHANGED Viewed

@@ -18,19 +18,19 @@ See the [ONNX Runtime `LinearAttentionGate` contrib-operator spec](https://githu
 ## Inputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `a` | `aT` | `T` | — | — | Decay gate projection with shape `(B, T, H)`. Any rank of at least 1 is accepted; the last axis is the head count and the leading axes are folded. | required |
-| `dt_bias` | `dtBiasT` | `TF` | `1` | — | Per-head float32 bias added to `a`, with shape (H). | required |
-| `decay_scale` | `decayScaleT` | `TF` | `1` | — | Per-head float32 multiplier applied to `softplus(a + dt_bias)`, with shape `(H)`. For gated DeltaNet this is `-exp(A_log)`. | required |
-| `b` | `bT` | `T` | — | — | Update-rate projection with the same shape as `a` when `beta` is requested. It is accepted but unused when `beta` is omitted. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `decay` | `decayT` | `T` | same as `a` | same as `a` | `decay_scale * softplus(a + dt_bias)`, with the same shape as `a`. | required |
-| `beta` | `betaT` | `T` | same as `a` | same as `a` | sigmoid(b), with the same shape as `a`. Requires the `b` input. | optional |
 ## Type constraints
@@ -41,7 +41,7 @@ See the [ONNX Runtime `LinearAttentionGate` contrib-operator spec](https://githu
 ## 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
@@ -49,10 +49,14 @@ See the [ONNX Runtime `LinearAttentionGate` contrib-operator spec](https://githu
 ## 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.

 ## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `aT` | `a` | `T` | — | — | Decay gate projection with shape `(B, T, H)`. Any rank of at least 1 is accepted; the last axis is the head count and the leading axes are folded. | required |
+| `dtBiasT` | `dt_bias` | `TF` | `1` | — | Per-head float32 bias added to `a`, with shape (H). | required |
+| `decayScaleT` | `decay_scale` | `TF` | `1` | — | Per-head float32 multiplier applied to `softplus(a + dt_bias)`, with shape `(H)`. For gated DeltaNet this is `-exp(A_log)`. | required |
+| `bT` | `b` | `T` | — | — | Update-rate projection with the same shape as `a` when `beta` is requested. It is accepted but unused when `beta` is omitted. | optional |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `decayT` | `decay` | `T` | same as `aT` | same as `aT` | `decay_scale * softplus(a + dt_bias)`, with the same shape as `a`. | required |
+| `betaT` | `beta` | `T` | same as `aT` | same as `aT` | sigmoid(b), with the same shape as `a`. Requires the `b` input. | optional |
 ## Type constraints
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, 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
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.LinearAttentionGate",
   "tunableSpace": { "WORKGROUP_SIZE": [32, 64, 128, 256] },
   "cases": [
     {

 {
   "tunableSpace": { "WORKGROUP_SIZE": [32, 64, 128, 256] },
   "cases": [
     {

build/webgpu/linear-attention-gate.wgsl.jinja CHANGED Viewed

@@ -51,11 +51,10 @@ fn softplus(x: f32) -> f32 {
 // so the operands are widened on load and the result is narrowed only on store.
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(
-  @builtin(global_invocation_id) gid: vec3<u32>,
-  @builtin(num_workgroups) nwg: vec3<u32>
 ) {
   // Rebuild the flat invocation index after the 2D dispatch fold.
-  let item = gid.x + gid.y * nwg.x * WORKGROUP_SIZE;
   if (item >= GATE_ITEMS) {
     return;
   }

 // so the operands are widened on load and the result is narrowed only on store.
 @compute @workgroup_size(WORKGROUP_SIZE, 1, 1)
 fn main(
+  @builtin(global_invocation_id) gid: vec3<u32>
 ) {
   // Rebuild the flat invocation index after the 2D dispatch fold.
+  let item = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WORKGROUP_SIZE;
   if (item >= GATE_ITEMS) {
     return;
   }

build/webgpu/manifest.json CHANGED Viewed

@@ -2,62 +2,21 @@
   "domain": "com.microsoft",
   "name": "LinearAttentionGate",
   "sinceVersion": 1,
-  "description": "Fuses the gate projections used by `com.microsoft.LinearAttention`'s gated-delta recurrence: `decay = decay_scale * softplus(a + dt_bias)` and, when requested, `beta = sigmoid(b)`. The last input axis is the head axis; `dt_bias` and `decay_scale` are float32 per-head vectors. Gate arithmetic is performed in float32 and narrowed only on store. Requesting `beta` requires `b`; an unconsumed `b` is permitted when `beta` is omitted.",
-  "inputs": [
-    {
-      "role": "a",
-      "dtype": "T",
-      "description": "Decay gate projection with shape `(B, T, H)`. Any rank of at least 1 is accepted; the last axis is the head count and the leading axes are folded."
-    },
-    {
-      "role": "dt_bias",
-      "dtype": "TF",
-      "rank": 1,
-      "description": "Per-head float32 bias added to `a`, with shape (H)."
-    },
-    {
-      "role": "decay_scale",
-      "dtype": "TF",
-      "rank": 1,
-      "description": "Per-head float32 multiplier applied to `softplus(a + dt_bias)`, with shape `(H)`. For gated DeltaNet this is `-exp(A_log)`."
-    },
-    {
-      "role": "b",
-      "dtype": "T",
-      "optional": true,
-      "description": "Update-rate projection with the same shape as `a` when `beta` is requested. It is accepted but unused when `beta` is omitted."
-    }
-  ],
-  "outputs": [
-    {
-      "role": "decay",
-      "dtype": "T",
-      "rank": "ranks.aT",
-      "shape": "shapes.aT",
-      "description": "`decay_scale * softplus(a + dt_bias)`, with the same shape as `a`."
-    },
-    {
-      "role": "beta",
-      "dtype": "T",
-      "rank": "ranks.aT",
-      "optional": true,
-      "shape": "shapes.aT",
-      "description": "sigmoid(b), with the same shape as `a`. Requires the `b` input."
-    }
-  ],
-  "typeConstraints": { "T": ["float32", "float16"], "TF": ["float32"] },
-  "tunables": { "WORKGROUP_SIZE": 64 },
-  "args": {
-    "aT": { "kind": "tensor", "semantic": "a", "role": "input" },
-    "dtBiasT": { "kind": "tensor", "semantic": "dt_bias", "role": "weights" },
-    "decayScaleT": { "kind": "tensor", "semantic": "decay_scale", "role": "weights" },
-    "bT": { "kind": "tensor", "semantic": "b", "role": "input", "required": false },
-    "decayT": { "kind": "tensor", "semantic": "decay", "role": "output" },
-    "betaT": { "kind": "tensor", "semantic": "beta", "role": "output", "required": false }
   },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
-    "foldedDispatchCapacity": "device.limits.maxComputeWorkgroupsPerDimension * device.limits.maxComputeWorkgroupsPerDimension",
     "numHeads": "dim(shapes.aT, ranks.aT - 1)",
     "gateCount": "numel(shapes.aT)",
     "headsVec4": "numHeads / 4",
@@ -72,98 +31,31 @@
     "scalarDispatchFits": "ceilDiv(gateCount, tunables.WORKGROUP_SIZE) <= foldedDispatchCapacity",
     "vec4DispatchFits": "numHeads % 4 == 0 and ceilDiv(gateVec4Count, tunables.WORKGROUP_SIZE) <= foldedDispatchCapacity"
   },
-  "bindingSets": {
-    "decayOnlyIo": [
-      {
-        "name": "a",
-        "arg": "aT",
-        "semantic": "a",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$gateElement",
-        "length": "$gateItems"
-      },
-      {
-        "name": "dt_bias",
-        "arg": "dtBiasT",
-        "semantic": "dt_bias",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$paramElement",
-        "length": "$headItems"
-      },
-      {
-        "name": "decay_scale",
-        "arg": "decayScaleT",
-        "semantic": "decay_scale",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$paramElement",
-        "length": "$headItems"
-      },
-      {
-        "name": "decay",
-        "arg": "decayT",
-        "semantic": "decay",
-        "buffer": { "type": "storage" },
-        "elementType": "$gateElement",
-        "length": "$gateItems"
-      }
-    ],
-    "withBetaIo": [
-      {
-        "name": "a",
-        "arg": "aT",
-        "semantic": "a",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$gateElement",
-        "length": "$gateItems"
-      },
-      {
-        "name": "dt_bias",
-        "arg": "dtBiasT",
-        "semantic": "dt_bias",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$paramElement",
-        "length": "$headItems"
-      },
-      {
-        "name": "decay_scale",
-        "arg": "decayScaleT",
-        "semantic": "decay_scale",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$paramElement",
-        "length": "$headItems"
-      },
-      {
-        "name": "b",
-        "arg": "bT",
-        "semantic": "b",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$gateElement",
-        "length": "$gateItems"
-      },
-      {
-        "name": "decay",
-        "arg": "decayT",
-        "semantic": "decay",
-        "buffer": { "type": "storage" },
-        "elementType": "$gateElement",
-        "length": "$gateItems"
-      },
-      {
-        "name": "beta",
-        "arg": "betaT",
-        "semantic": "beta",
-        "buffer": { "type": "storage" },
-        "elementType": "$gateElement",
-        "length": "$gateItems"
-      }
-    ]
   },
   "variants": [
     {
       "id": "vec4_gate_beta",
       "priority": 30,
-      "when": ["betaContract", "workgroupFits", "vec4DispatchFits"],
-      "constants": {
         "vectorized": true,
         "hasBeta": true,
         "usesF16": "gateDtype == \"float16\"",
@@ -178,16 +70,20 @@
           "id": "main",
           "name": "LinearAttentionGate.Vec4GateBeta",
           "shader": "linear-attention-gate.wgsl.jinja",
-          "bindings": "withBetaIo",
-          "dispatch": { "threads": "gateVec4Count", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "vec4_gate",
       "priority": 20,
-      "when": ["decayOnlyContract", "workgroupFits", "vec4DispatchFits"],
-      "constants": {
         "vectorized": true,
         "hasBeta": false,
         "usesF16": "gateDtype == \"float16\"",
@@ -202,16 +98,20 @@
           "id": "main",
           "name": "LinearAttentionGate.Vec4Gate",
           "shader": "linear-attention-gate.wgsl.jinja",
-          "bindings": "decayOnlyIo",
-          "dispatch": { "threads": "gateVec4Count", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "scalar_gate_beta",
       "priority": 10,
-      "when": ["betaContract", "workgroupFits", "scalarDispatchFits"],
-      "constants": {
         "vectorized": false,
         "hasBeta": true,
         "usesF16": "gateDtype == \"float16\"",
@@ -226,16 +126,20 @@
           "id": "main",
           "name": "LinearAttentionGate.ScalarGateBeta",
           "shader": "linear-attention-gate.wgsl.jinja",
-          "bindings": "withBetaIo",
-          "dispatch": { "threads": "gateCount", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "scalar_gate",
       "priority": 0,
-      "when": ["decayOnlyContract", "workgroupFits", "scalarDispatchFits"],
-      "constants": {
         "vectorized": false,
         "hasBeta": false,
         "usesF16": "gateDtype == \"float16\"",
@@ -250,8 +154,12 @@
           "id": "main",
           "name": "LinearAttentionGate.ScalarGate",
           "shader": "linear-attention-gate.wgsl.jinja",
-          "bindings": "decayOnlyIo",
-          "dispatch": { "threads": "gateCount", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     }

   "domain": "com.microsoft",
   "name": "LinearAttentionGate",
   "sinceVersion": 1,
+  "inputs": {
+    "aT": { "onnx": "a", "dtype": "T" },
+    "dtBiasT": { "onnx": "dt_bias", "dtype": "TF", "rank": 1 },
+    "decayScaleT": { "onnx": "decay_scale", "dtype": "TF", "rank": 1 },
+    "bT": { "onnx": "b", "dtype": "T", "optional": true }
   },
+  "outputs": {
+    "decayT": { "onnx": "decay", "dtype": "T", "rank": "ranks.aT", "shape": "shapes.aT" },
+    "betaT": { "onnx": "beta", "dtype": "T", "rank": "ranks.aT", "optional": true, "shape": "shapes.aT" }
+  },
+  "typeConstraints": { "T": ["float32", "float16"], "TF": ["float32"] },
+  "tunables": { "WORKGROUP_SIZE": { "default": 64 } },
   "derive": {
     "deviceWorkgroupCap": "min(device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "foldedDispatchCapacity": "min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
     "numHeads": "dim(shapes.aT, ranks.aT - 1)",
     "gateCount": "numel(shapes.aT)",
     "headsVec4": "numHeads / 4",
     "scalarDispatchFits": "ceilDiv(gateCount, tunables.WORKGROUP_SIZE) <= foldedDispatchCapacity",
     "vec4DispatchFits": "numHeads % 4 == 0 and ceilDiv(gateVec4Count, tunables.WORKGROUP_SIZE) <= foldedDispatchCapacity"
   },
+  "when": ["workgroupFits"],
+  "bindings": {
+    "a": { "arg": "aT", "buffer": "read-only-storage", "elementType": "$gateElement", "length": "$gateItems" },
+    "dt_bias": {
+      "arg": "dtBiasT",
+      "buffer": "read-only-storage",
+      "elementType": "$paramElement",
+      "length": "$headItems"
+    },
+    "decay_scale": {
+      "arg": "decayScaleT",
+      "buffer": "read-only-storage",
+      "elementType": "$paramElement",
+      "length": "$headItems"
+    },
+    "b": { "arg": "bT", "buffer": "read-only-storage", "elementType": "$gateElement", "length": "$gateItems" },
+    "decay": { "arg": "decayT", "buffer": "storage", "elementType": "$gateElement", "length": "$gateItems" },
+    "beta": { "arg": "betaT", "buffer": "storage", "elementType": "$gateElement", "length": "$gateItems" }
   },
   "variants": [
     {
       "id": "vec4_gate_beta",
       "priority": 30,
+      "when": ["betaContract", "vec4DispatchFits"],
+      "derive": {
         "vectorized": true,
         "hasBeta": true,
         "usesF16": "gateDtype == \"float16\"",
           "id": "main",
           "name": "LinearAttentionGate.Vec4GateBeta",
           "shader": "linear-attention-gate.wgsl.jinja",
+          "bindings": ["a", "dt_bias", "decay_scale", "b", "decay", "beta"],
+          "dispatch": {
+            "x": "min(ceilDiv((gateVec4Count), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((gateVec4Count), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "vec4_gate",
       "priority": 20,
+      "when": ["decayOnlyContract", "vec4DispatchFits"],
+      "derive": {
         "vectorized": true,
         "hasBeta": false,
         "usesF16": "gateDtype == \"float16\"",
           "id": "main",
           "name": "LinearAttentionGate.Vec4Gate",
           "shader": "linear-attention-gate.wgsl.jinja",
+          "bindings": ["a", "dt_bias", "decay_scale", "decay"],
+          "dispatch": {
+            "x": "min(ceilDiv((gateVec4Count), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((gateVec4Count), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "scalar_gate_beta",
       "priority": 10,
+      "when": ["betaContract", "scalarDispatchFits"],
+      "derive": {
         "vectorized": false,
         "hasBeta": true,
         "usesF16": "gateDtype == \"float16\"",
           "id": "main",
           "name": "LinearAttentionGate.ScalarGateBeta",
           "shader": "linear-attention-gate.wgsl.jinja",
+          "bindings": ["a", "dt_bias", "decay_scale", "b", "decay", "beta"],
+          "dispatch": {
+            "x": "min(ceilDiv((gateCount), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((gateCount), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "scalar_gate",
       "priority": 0,
+      "when": ["decayOnlyContract", "scalarDispatchFits"],
+      "derive": {
         "vectorized": false,
         "hasBeta": false,
         "usesF16": "gateDtype == \"float16\"",
           "id": "main",
           "name": "LinearAttentionGate.ScalarGate",
           "shader": "linear-attention-gate.wgsl.jinja",
+          "bindings": ["a", "dt_bias", "decay_scale", "decay"],
+          "dispatch": {
+            "x": "min(ceilDiv((gateCount), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((gateCount), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,18 +1,26 @@
 {
   "name": "com.microsoft.LinearAttentionGate",
-  "id": "_com_microsoft_linearattentiongate_webgpu_cfa91bb",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "1NTbbJYE62OsV1lkReUeXTrI7rnBiOY3HTg5ldTfPyk=",
-      "linear-attention-gate.wgsl.jinja": "NRzrWvQ2VrRHu5YsGLn5RGMIXHePxHum+umyWdbcWJs=",
-      "manifest.json": "s1kmQ2WU1JJPknyUoE8HJj3gUC7F44ubx+ZkAtbc8To=",
-      "test.json": "mRMxucSUiqKh/upG/J2xGsRxpeEDFhVLOYGYRCMQd2Y="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.LinearAttentionGate" }
 }

 {
   "name": "com.microsoft.LinearAttentionGate",
+  "id": "_com_microsoft_linearattentiongate_webgpu_4bb5397",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "nSMkCE+adLNKSBkCYEFz8YeOJf6YjUsw4vEzFToXs5U=",
+      "linear-attention-gate.wgsl.jinja": "Gi935dqD4NLjZbH6v4gzTYZElbzpeL5mC+ZD91H5pvo=",
+      "manifest.json": "kEskRGoQGGG0erg57qnQc/Sa4hJq5cNnduwqwU5zuDk=",
+      "test.json": "rlpI/FCSMX4yxQoveUCaj13GqK8JA+dadMzKFo129I8="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "vec4_gate_beta": ["linear-attention-gate.wgsl.jinja"],
+      "vec4_gate": ["linear-attention-gate.wgsl.jinja"],
+      "scalar_gate_beta": ["linear-attention-gate.wgsl.jinja"],
+      "scalar_gate": ["linear-attention-gate.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/test.json CHANGED Viewed

@@ -1,10 +1,9 @@
 {
-  "op": "com.microsoft.LinearAttentionGate",
   "cases": [
     {
       "name": "rank3_h8_vec4_gate_beta",
       "provenance": {
-        "notes": "The (B, T, H) shape the schema names, with a head count that is a multiple of four so the vectorized path runs. Covers vec4_gate_beta. Tolerance: the measured kernel-vs-reference max relative deviation over this op's fixtures is 4.4e-7, which is the f32 floor for softplus just above the log1p series crossover; 2e-6 keeps roughly four times that as device margin."
       },
       "inputs": {
         "aT": {
@@ -100,7 +99,7 @@
     {
       "name": "b_without_beta_is_ignored",
       "provenance": {
-        "notes": "The public schema requires b only in the forward direction: requesting beta requires b. ORT CPU and WebGPU accept b when beta is omitted and do not read it; this fixture prevents the manifest from imposing the stronger, non-schema b-implies-beta rule."
       },
       "inputs": {
         "aT": {
@@ -125,7 +124,7 @@
     {
       "name": "rank2_h4_vec4_pinned_head_parameters",
       "provenance": {
-        "notes": "Hand-computed, anchoring the op off its own reference. Row 0 is built so a + dt_bias == 0 in every head, making softplus exactly ln(2) = 0.69314718, so each head's decay is its decay_scale times that one constant -- a wrong head-to-parameter mapping cannot reproduce the column pattern. Row 1 biases to 25 in every head, above the x > 20 softplus threshold where softplus(x) == x exactly, so those four values are exact. beta pins both sigmoid branches (x >= 0 and x < 0) at 0, +/-1, +/-2 and +/-6."
       },
       "inputs": {
         "aT": {

 {
   "cases": [
     {
       "name": "rank3_h8_vec4_gate_beta",
       "provenance": {
+        "notes": "The schema's (B,T,H) layout uses a head count divisible by four, exercising vectorized beta gating. A 2e-6 tolerance covers f32 Softplus rounding near the log1p series crossover."
       },
       "inputs": {
         "aT": {
     {
       "name": "b_without_beta_is_ignored",
       "provenance": {
+        "notes": "The schema requires `b` when beta is requested but permits `b` when beta is omitted. This case supplies unused `b` without beta."
       },
       "inputs": {
         "aT": {
     {
       "name": "rank2_h4_vec4_pinned_head_parameters",
       "provenance": {
+        "notes": "Expected values are derived directly from the gating equation. Row 0 sets `a + dt_bias` to zero in every head, so Softplus equals ln(2) and each decay differs only by `decay_scale`. Row 1 sets the value to 25, where the guarded Softplus branch returns x exactly. Beta inputs cover both sigmoid signs at 0, +/-1, +/-2, and +/-6."
       },
       "inputs": {
         "aT": {