sync 2e7068faf55e

Browse files

Files changed (6) hide show

README.md +65 -0
build/webgpu/bench.json +167 -0
build/webgpu/linear-attention-gate.wgsl.jinja +92 -0
build/webgpu/manifest.json +259 -0
build/webgpu/metadata.json +18 -0
build/webgpu/test.json +268 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.LinearAttentionGate
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 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.
+See the [ONNX Runtime `LinearAttentionGate` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.LinearAttentionGate) for the reference semantics.
+## 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
+| Variable | Allowed dtypes |
+| --- | --- |
+| `T` | `float32`, `float16` |
+| `TF` | `float32` |
+## 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
+- [`linear-attention-gate.wgsl.jinja`](build/webgpu/linear-attention-gate.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.LinearAttentionGate", { version: 1 });
+const { decayT } = await kernel({
+  aT: { data: aTData, shape: [5] },
+  dtBiasT: { data: dtBiasTData, shape: [5] },
+  decayScaleT: { data: decayScaleTData, shape: [5] },
+});
+```

build/webgpu/bench.json ADDED Viewed

	@@ -0,0 +1,167 @@

+{
+  "op": "com.microsoft.LinearAttentionGate",
+  "tunableSpace": { "WORKGROUP_SIZE": [32, 64, 128, 256] },
+  "cases": [
+    {
+      "name": "gate-smoke-decode-h64",
+      "preset": "smoke",
+      "provenance": {
+        "notes": "One token: the decode shape, where the whole op is a single tiny dispatch and launch overhead dominates."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 2.0 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [64], "data": { "kind": "linspace", "start": -1.5, "end": 1.0 } },
+        "decayScaleT": { "dtype": "float32", "shape": [64], "data": { "kind": "linspace", "start": -3.0, "end": -0.2 } },
+        "bT": {
+          "dtype": "float32",
+          "shape": [1, 1, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 2.5 }
+        }
+      },
+      "outputs": {
+        "decayT": { "dtype": "float32", "shape": [1, 1, 64] },
+        "betaT": { "dtype": "float32", "shape": [1, 1, 64] }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "4 * (4 * numel(shapes.aT) + 2 * numel(shapes.dtBiasT))" }]
+      }
+    },
+    {
+      "name": "gate-smoke-prefill-t512-h32",
+      "preset": "smoke",
+      "provenance": { "notes": "Prefill over 512 tokens on the vectorized path." },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 512, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 2.0 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [32], "data": { "kind": "linspace", "start": -1.5, "end": 1.0 } },
+        "decayScaleT": { "dtype": "float32", "shape": [32], "data": { "kind": "linspace", "start": -3.0, "end": -0.2 } },
+        "bT": {
+          "dtype": "float32",
+          "shape": [1, 512, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 2.5 }
+        }
+      },
+      "outputs": {
+        "decayT": { "dtype": "float32", "shape": [1, 512, 32] },
+        "betaT": { "dtype": "float32", "shape": [1, 512, 32] }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "4 * (4 * numel(shapes.aT) + 2 * numel(shapes.dtBiasT))" }]
+      }
+    },
+    {
+      "name": "gate-smoke-prefill-t256-h30",
+      "preset": "smoke",
+      "provenance": { "notes": "Head count 30 is not a multiple of four, so this is the scalar path at prefill width." },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 256, 30],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 2.0 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [30], "data": { "kind": "linspace", "start": -1.5, "end": 1.0 } },
+        "decayScaleT": { "dtype": "float32", "shape": [30], "data": { "kind": "linspace", "start": -3.0, "end": -0.2 } },
+        "bT": {
+          "dtype": "float32",
+          "shape": [1, 256, 30],
+          "data": { "kind": "fillFloat32", "sinStep": 0.19, "cosStep": 0.23, "scale": 2.5 }
+        }
+      },
+      "outputs": {
+        "decayT": { "dtype": "float32", "shape": [1, 256, 30] },
+        "betaT": { "dtype": "float32", "shape": [1, 256, 30] }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "4 * (4 * numel(shapes.aT) + 2 * numel(shapes.dtBiasT))" }]
+      }
+    },
+    {
+      "name": "gate-smoke-decay-only-b2-t128-h64",
+      "preset": "smoke",
+      "provenance": {
+        "notes": "No b input: two of the four tensors disappear, which is the bandwidth the optional-pair binding actually saves."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [2, 128, 64],
+          "data": { "kind": "fillFloat32", "sinStep": 0.13, "cosStep": 0.29, "scale": 2.0 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [64], "data": { "kind": "linspace", "start": -1.5, "end": 1.0 } },
+        "decayScaleT": { "dtype": "float32", "shape": [64], "data": { "kind": "linspace", "start": -3.0, "end": -0.2 } }
+      },
+      "outputs": { "decayT": { "dtype": "float32", "shape": [2, 128, 64] } },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "4 * (2 * numel(shapes.aT) + 2 * numel(shapes.dtBiasT))" }]
+      }
+    },
+    {
+      "name": "gate-qwen3next-decode-s1-h32",
+      "preset": "model",
+      "provenance": {
+        "notes": "Qwen3-Next class defaults (linear_num_value_heads 32): the decay/beta gate for one decode step."
+      },
+      "vars": { "batch": 1, "seq": 1, "heads": 32 },
+      "inputs": {
+        "aT": { "shape": [1, 1, 32], "dtype": "float32", "dist": "normal", "seed": 8400, "scale": 0.5 },
+        "dtBiasT": { "shape": [32], "dtype": "float32", "dist": "normal", "seed": 8401, "scale": 0.2 },
+        "decayScaleT": { "shape": [32], "dtype": "float32", "dist": "uniform", "seed": 8402, "min": 0.5, "max": 1.5 },
+        "bT": { "shape": [1, 1, 32], "dtype": "float32", "dist": "normal", "seed": 8403, "scale": 0.5 }
+      },
+      "outputs": {
+        "decayT": { "shape": [1, 1, 32], "dtype": "float32" },
+        "betaT": { "shape": [1, 1, 32], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "4 * (4 * numel(shapes.aT) + 2 * numel(shapes.dtBiasT))" }]
+      }
+    },
+    {
+      "name": "gate-qwen3next-prefill-s2048-h32",
+      "preset": "model",
+      "provenance": { "notes": "Qwen3-Next class defaults over a 2048-token prefill chunk." },
+      "vars": { "batch": 1, "seq": 2048, "heads": 32 },
+      "inputs": {
+        "aT": { "shape": [1, 2048, 32], "dtype": "float32", "dist": "normal", "seed": 8500, "scale": 0.5 },
+        "dtBiasT": { "shape": [32], "dtype": "float32", "dist": "normal", "seed": 8501, "scale": 0.2 },
+        "decayScaleT": { "shape": [32], "dtype": "float32", "dist": "uniform", "seed": 8502, "min": 0.5, "max": 1.5 },
+        "bT": { "shape": [1, 2048, 32], "dtype": "float32", "dist": "normal", "seed": 8503, "scale": 0.5 }
+      },
+      "outputs": {
+        "decayT": { "shape": [1, 2048, 32], "dtype": "float32" },
+        "betaT": { "shape": [1, 2048, 32], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "4 * (4 * numel(shapes.aT) + 2 * numel(shapes.dtBiasT))" }]
+      }
+    },
+    {
+      "name": "gate-qwen3next-prefill-s8192-h32",
+      "preset": "model",
+      "provenance": {
+        "notes": "Qwen3-Next class defaults over an 8192-token prefill chunk, where the gate is purely bandwidth-bound."
+      },
+      "vars": { "batch": 1, "seq": 8192, "heads": 32 },
+      "inputs": {
+        "aT": { "shape": [1, 8192, 32], "dtype": "float32", "dist": "normal", "seed": 8600, "scale": 0.5 },
+        "dtBiasT": { "shape": [32], "dtype": "float32", "dist": "normal", "seed": 8601, "scale": 0.2 },
+        "decayScaleT": { "shape": [32], "dtype": "float32", "dist": "uniform", "seed": 8602, "min": 0.5, "max": 1.5 },
+        "bT": { "shape": [1, 8192, 32], "dtype": "float32", "dist": "normal", "seed": 8603, "scale": 0.5 }
+      },
+      "outputs": {
+        "decayT": { "shape": [1, 8192, 32], "dtype": "float32" },
+        "betaT": { "shape": [1, 8192, 32], "dtype": "float32" }
+      },
+      "bench": {
+        "metrics": [{ "type": "bandwidth", "value": "4 * (4 * numel(shapes.aT) + 2 * numel(shapes.dtBiasT))" }]
+      }
+    }
+  ]
+}

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

	@@ -0,0 +1,92 @@

+{% if usesF16 %}
+enable f16;
+{% endif -%}
+{{ env.wgsl.resourceDeclarations }}
+{% if vectorized %}
+// Number of vec4 parameter slots across the head axis.
+{% else %}
+// Number of scalar parameter slots across the head axis.
+{% endif %}
+const HEAD_ITEMS: u32 = {{ headItems }}u;
+const GATE_ITEMS: u32 = {{ gateItems }}u;
+const WORKGROUP_SIZE: u32 = {{ workgroupSize }}u;
+{% if hasBeta %}
+fn sigmoid_safe(x: f32) -> f32 {
+  if (x >= 0.0) {
+    let z = exp(-x);
+    return 1.0 / (1.0 + z);
+  }
+  let z = exp(x);
+  return z / (1.0 + z);
+}
+{% endif %}
+// WGSL has no log1p builtin, and log(1.0 + y) loses relative accuracy for small y: rounding
+// 1.0 + y costs up to half an ulp of 1.0, which is enormous next to y itself. Below the
+// crossover the series is the accurate form -- its truncation is y^7 / 8, under 6e-9 relative
+// at y = 0.125 -- and above it the same half-ulp is under 6e-7 relative because log1p(y) has
+// grown past 0.118. The algebraic (1.0 + y) - 1.0 correction is not usable here: a backend is
+// free to fold that subtraction back to y, which silently restores the error it removes.
+fn log1p_pos(y: f32) -> f32 {
+  if (y < 0.125) {
+    return y * (1.0 + y * (-0.5 + y * (0.33333334 + y * (-0.25 + y * (0.2 + y * (-0.16666667 + y * 0.14285714))))));
+  }
+  return log(1.0 + y);
+}
+fn softplus(x: f32) -> f32 {
+  // Both arms feed log1p an exp() of a non-positive argument, so the operand stays in (0, 1]
+  // and never overflows; for large x the increment rounds away and the result is exactly x,
+  // which is the correct limit.
+  if (x > 0.0) {
+    return x + log1p_pos(exp(-x));
+  }
+  return log1p_pos(exp(x));
+}
+// The gate arithmetic is FP32 for every tensor type. decay feeds exp() inside the
+// LinearAttention recurrence, where any precision lost here is amplified exponentially,
+// 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;
+  }
+  // The last axis is the head axis, so the per-head parameter index is the flat index
+  // modulo the head count. The vectorized path holds because the head count is a multiple
+  // of four: a vec4 of adjacent gates never straddles two heads' parameter slots.
+  let head = item % HEAD_ITEMS;
+{% if vectorized %}
+  let biased = vec4<f32>(a[item]) + dt_bias[head];
+  let activated = vec4<f32>(
+    softplus(biased.x),
+    softplus(biased.y),
+    softplus(biased.z),
+    softplus(biased.w)
+  );
+  decay[item] = {{ gateElement }}(decay_scale[head] * activated);
+{% if hasBeta %}
+  let raw_beta = vec4<f32>(b[item]);
+  beta[item] = {{ gateElement }}(vec4<f32>(
+    sigmoid_safe(raw_beta.x),
+    sigmoid_safe(raw_beta.y),
+    sigmoid_safe(raw_beta.z),
+    sigmoid_safe(raw_beta.w)
+  ));
+{% endif %}
+{% else %}
+  let biased = f32(a[item]) + dt_bias[head];
+  decay[item] = {{ gateElement }}(decay_scale[head] * softplus(biased));
+{% if hasBeta %}
+  beta[item] = {{ gateElement }}(sigmoid_safe(f32(b[item])));
+{% endif %}
+{% endif %}
+}

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,259 @@

+{
+  "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",
+    "gateVec4Count": "gateCount / 4",
+    "gateDtype": "tensorDtypes.aT",
+    "gateDtypeOk": "(gateDtype == \"float32\" or gateDtype == \"float16\") and f16Ok(dtypes.T)",
+    "paramsOk": "ranks.aT >= 1 and numHeads > 0 and ranks.dtBiasT == 1 and ranks.decayScaleT == 1 and tensorDtypes.dtBiasT == \"float32\" and tensorDtypes.decayScaleT == \"float32\" and dim(shapes.dtBiasT, 0) == numHeads and dim(shapes.decayScaleT, 0) == numHeads",
+    "tensorContract": "gateDtypeOk and paramsOk and sameShape(shapes.decayT, shapes.aT) and tensorDtypes.decayT == gateDtype",
+    "betaContract": "tensorContract and present.bT and present.betaT and sameShape(shapes.bT, shapes.aT) and sameShape(shapes.betaT, shapes.aT) and tensorDtypes.bT == gateDtype and tensorDtypes.betaT == gateDtype",
+    "decayOnlyContract": "tensorContract and not present.betaT",
+    "workgroupFits": "tunables.WORKGROUP_SIZE > 0 and tunables.WORKGROUP_SIZE <= deviceWorkgroupCap",
+    "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\"",
+        "gateElement": "\"vec4<f16>\" if gateDtype == \"float16\" else \"vec4<f32>\"",
+        "paramElement": "\"vec4<f32>\"",
+        "headItems": "headsVec4",
+        "gateItems": "gateVec4Count",
+        "workgroupSize": "tunables.WORKGROUP_SIZE"
+      },
+      "passes": [
+        {
+          "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\"",
+        "gateElement": "\"vec4<f16>\" if gateDtype == \"float16\" else \"vec4<f32>\"",
+        "paramElement": "\"vec4<f32>\"",
+        "headItems": "headsVec4",
+        "gateItems": "gateVec4Count",
+        "workgroupSize": "tunables.WORKGROUP_SIZE"
+      },
+      "passes": [
+        {
+          "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\"",
+        "gateElement": "\"f16\" if gateDtype == \"float16\" else \"f32\"",
+        "paramElement": "\"f32\"",
+        "headItems": "numHeads",
+        "gateItems": "gateCount",
+        "workgroupSize": "tunables.WORKGROUP_SIZE"
+      },
+      "passes": [
+        {
+          "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\"",
+        "gateElement": "\"f16\" if gateDtype == \"float16\" else \"f32\"",
+        "paramElement": "\"f32\"",
+        "headItems": "numHeads",
+        "gateItems": "gateCount",
+        "workgroupSize": "tunables.WORKGROUP_SIZE"
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "LinearAttentionGate.ScalarGate",
+          "shader": "linear-attention-gate.wgsl.jinja",
+          "bindings": "decayOnlyIo",
+          "dispatch": { "threads": "gateCount", "workgroupSize": "tunables.WORKGROUP_SIZE" }
+        }
+      ]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

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

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

build/webgpu/test.json ADDED Viewed

	@@ -0,0 +1,268 @@

+{
+  "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": {
+          "dtype": "float32",
+          "shape": [2, 3, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.17, "scale": 2.5 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [8], "data": { "kind": "linspace", "start": -1.5, "end": 1.1 } },
+        "decayScaleT": { "dtype": "float32", "shape": [8], "data": { "kind": "linspace", "start": -3.2, "end": -0.15 } },
+        "bT": {
+          "dtype": "float32",
+          "shape": [2, 3, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.13, "scale": 3.0 }
+        }
+      },
+      "outputs": {
+        "decayT": { "dtype": "float32", "shape": [2, 3, 8], "tolerance": 0, "relTolerance": 0.000002 },
+        "betaT": { "dtype": "float32", "shape": [2, 3, 8], "tolerance": 0, "relTolerance": 0.000002 }
+      }
+    },
+    {
+      "name": "rank2_h6_scalar_gate_beta",
+      "provenance": {
+        "notes": "Head count 6 is not a multiple of four, so the vectorized head-to-parameter mapping does not hold and the scalar path is the only eligible one. Covers scalar_gate_beta."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [5, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 2.2 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [6], "data": { "kind": "linspace", "start": -2.0, "end": 0.9 } },
+        "decayScaleT": { "dtype": "float32", "shape": [6], "data": { "kind": "linspace", "start": -2.6, "end": -0.2 } },
+        "bT": {
+          "dtype": "float32",
+          "shape": [5, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.11, "scale": 2.8 }
+        }
+      },
+      "outputs": {
+        "decayT": { "dtype": "float32", "shape": [5, 6], "tolerance": 0, "relTolerance": 0.000002 },
+        "betaT": { "dtype": "float32", "shape": [5, 6], "tolerance": 0, "relTolerance": 0.000002 }
+      }
+    },
+    {
+      "name": "rank4_h4_vec4_gate_only",
+      "provenance": {
+        "notes": "Rank 4 proves the leading axes are folded rather than parsed: only the last axis is the head axis. No b input, so no beta output. Covers vec4_gate."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [2, 2, 3, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.41, "cosStep": 0.23, "scale": 3.1 }
+        },
+        "dtBiasT": {
+          "dtype": "float32",
+          "shape": [4],
+          "data": { "kind": "values", "values": [-1.25, 0.4, 1.75, -0.6] }
+        },
+        "decayScaleT": {
+          "dtype": "float32",
+          "shape": [4],
+          "data": { "kind": "values", "values": [-1.5, -0.35, -2.8, -0.9] }
+        }
+      },
+      "outputs": { "decayT": { "dtype": "float32", "shape": [2, 2, 3, 4], "tolerance": 0, "relTolerance": 0.000002 } }
+    },
+    {
+      "name": "rank1_h5_scalar_gate_only",
+      "provenance": {
+        "notes": "Rank 1 is the single-token decode shape: the whole tensor is one head row, so every element reads a different parameter slot. Covers scalar_gate."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [5],
+          "data": { "kind": "values", "values": [-9.5, -1.25, 0.0, 2.5, 14.0] }
+        },
+        "dtBiasT": {
+          "dtype": "float32",
+          "shape": [5],
+          "data": { "kind": "values", "values": [0.75, -0.5, 1.25, -1.75, 0.3] }
+        },
+        "decayScaleT": {
+          "dtype": "float32",
+          "shape": [5],
+          "data": { "kind": "values", "values": [-0.4, -1.6, -3.05, -0.85, -2.2] }
+        }
+      },
+      "outputs": { "decayT": { "dtype": "float32", "shape": [5], "tolerance": 0, "relTolerance": 0.000002 } }
+    },
+    {
+      "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": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.27, "cosStep": 0.19, "scale": 1.7 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [4], "data": { "kind": "linspace", "start": -0.7, "end": 0.8 } },
+        "decayScaleT": {
+          "dtype": "float32",
+          "shape": [4],
+          "data": { "kind": "values", "values": [-0.5, -1.0, -1.5, -2.0] }
+        },
+        "bT": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "fillFloat32", "sinStep": 0.43, "cosStep": 0.11, "scale": 4.0 }
+        }
+      },
+      "outputs": { "decayT": { "dtype": "float32", "shape": [2, 4], "tolerance": 0, "relTolerance": 0.000002 } }
+    },
+    {
+      "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": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": [0.0, -1.0, 1.0, -0.5, 25.0, 24.0, 26.0, 24.5] }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [4], "data": { "kind": "values", "values": [0.0, 1.0, -1.0, 0.5] } },
+        "decayScaleT": {
+          "dtype": "float32",
+          "shape": [4],
+          "data": { "kind": "values", "values": [-1.0, -2.0, -0.5, -4.0] }
+        },
+        "bT": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": { "kind": "values", "values": [0.0, 6.0, -6.0, 1.0, -1.0, 2.0, -2.0, 0.0] }
+        }
+      },
+      "outputs": {
+        "decayT": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": {
+            "kind": "values",
+            "values": [-0.69314718, -1.3862944, -0.34657359, -2.7725887, -25.0, -50.0, -12.5, -100.0]
+          },
+          "tolerance": 0,
+          "relTolerance": 0.000001
+        },
+        "betaT": {
+          "dtype": "float32",
+          "shape": [2, 4],
+          "data": {
+            "kind": "values",
+            "values": [0.5, 0.99752738, 0.0024726232, 0.73105858, 0.26894142, 0.88079708, 0.11920292, 0.5]
+          },
+          "tolerance": 0,
+          "relTolerance": 0.000001
+        }
+      }
+    },
+    {
+      "name": "rank2_h5_scalar_pinned_saturation",
+      "provenance": {
+        "notes": "Hand-computed on the scalar path (head count 5). Both rows bias to one shared value -- 25 in row 0 (the x > 20 softplus identity) and 0 in row 1 (ln 2) -- so the five distinct decay values in each row are exactly the five decay_scale entries times one constant. That is the per-head parameter index, index % head_count, stated as arithmetic."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [24.5, 25.5, 25.0, 23.0, 27.0, -0.5, 0.5, 0.0, -2.0, 2.0] }
+        },
+        "dtBiasT": {
+          "dtype": "float32",
+          "shape": [5],
+          "data": { "kind": "values", "values": [0.5, -0.5, 0.0, 2.0, -2.0] }
+        },
+        "decayScaleT": {
+          "dtype": "float32",
+          "shape": [5],
+          "data": { "kind": "values", "values": [-1.0, -2.0, -0.25, -4.0, -0.5] }
+        },
+        "bT": {
+          "dtype": "float32",
+          "shape": [2, 5],
+          "data": { "kind": "values", "values": [-6.0, -2.0, 0.0, 2.0, 6.0, -1.0, -0.25, 0.25, 1.0, 3.0] }
+        }
+      },
+      "outputs": {
+        "decayT": {
+          "dtype": "float32",
+          "shape": [2, 5],
+          "data": {
+            "kind": "values",
+            "values": [-25.0, -50.0, -6.25, -100.0, -12.5, -0.69314718, -1.3862944, -0.1732868, -2.7725887, -0.34657359]
+          },
+          "tolerance": 0,
+          "relTolerance": 0.000001
+        },
+        "betaT": {
+          "dtype": "float32",
+          "shape": [2, 5],
+          "data": {
+            "kind": "values",
+            "values": [0.0024726232, 0.11920292, 0.5, 0.88079708, 0.99752738, 0.26894142, 0.4378235, 0.5621765, 0.73105858, 0.95257413]
+          },
+          "tolerance": 0,
+          "relTolerance": 0.000001
+        }
+      }
+    },
+    {
+      "name": "f16_rank3_h8_vec4_gate_beta",
+      "provenance": {
+        "notes": "float16 a/b/decay/beta with float32 dt_bias and decay_scale -- the schema's T/TF split. The gate arithmetic still runs in f32 and only the store narrows."
+      },
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [2, 3, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.17, "scale": 2.5 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [8], "data": { "kind": "linspace", "start": -1.5, "end": 1.1 } },
+        "decayScaleT": { "dtype": "float32", "shape": [8], "data": { "kind": "linspace", "start": -3.2, "end": -0.15 } },
+        "bT": {
+          "dtype": "float16",
+          "shape": [2, 3, 8],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.13, "scale": 3.0 }
+        }
+      },
+      "outputs": {
+        "decayT": { "dtype": "float16", "shape": [2, 3, 8], "tolerance": 0, "relTolerance": 0.002 },
+        "betaT": { "dtype": "float16", "shape": [2, 3, 8], "tolerance": 0, "relTolerance": 0.002 }
+      }
+    },
+    {
+      "name": "f16_rank2_h6_scalar_gate_beta",
+      "provenance": { "notes": "float16 gate tensors on the scalar path (head count 6)." },
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [5, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.29, "cosStep": 0.19, "scale": 2.2 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [6], "data": { "kind": "linspace", "start": -2.0, "end": 0.9 } },
+        "decayScaleT": { "dtype": "float32", "shape": [6], "data": { "kind": "linspace", "start": -2.6, "end": -0.2 } },
+        "bT": {
+          "dtype": "float16",
+          "shape": [5, 6],
+          "data": { "kind": "fillFloat32", "sinStep": 0.37, "cosStep": 0.11, "scale": 2.8 }
+        }
+      },
+      "outputs": {
+        "decayT": { "dtype": "float16", "shape": [5, 6], "tolerance": 0, "relTolerance": 0.002 },
+        "betaT": { "dtype": "float16", "shape": [5, 6], "tolerance": 0, "relTolerance": 0.002 }
+      }
+    }
+  ]
+}