sync 6fdf6301e2bb

Browse files

Files changed (5) hide show

README.md +2 -2
build/webgpu/linear-attention-gate.wgsl.jinja +9 -8
build/webgpu/manifest.json +24 -41
build/webgpu/metadata.json +6 -6
build/webgpu/test.json +120 -2

README.md CHANGED Viewed

@@ -44,13 +44,13 @@ See the [ONNX Runtime `LinearAttentionGate` contrib-operator spec](https://githu
 - [`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
 - [`linear-attention-gate.wgsl.jinja`](build/webgpu/linear-attention-gate.wgsl.jinja)
 ## 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.

 - [`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 cases
 - [`linear-attention-gate.wgsl.jinja`](build/webgpu/linear-attention-gate.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

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

@@ -1,6 +1,11 @@
-{% if usesF16 %}
-enable f16;
-{% endif -%}
 {{ env.wgsl.resourceDeclarations }}
 {% if vectorized %}
@@ -53,11 +58,7 @@ fn softplus(x: f32) -> f32 {
 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;
-  }
   // 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

+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 {% if vectorized %}
 fn main(
   @builtin(global_invocation_id) gid: vec3<u32>
 ) {
+{{ flat_index_2d("WORKGROUP_SIZE", "item", "GATE_ITEMS") }}
   // 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

build/webgpu/manifest.json CHANGED Viewed

@@ -29,26 +29,17 @@
     "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"
   },
   "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": [
     {
@@ -57,13 +48,11 @@
       "when": ["betaContract", "vec4DispatchFits"],
       "derive": {
         "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": [
         {
@@ -72,8 +61,8 @@
           "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
           }
         }
@@ -85,13 +74,11 @@
       "when": ["decayOnlyContract", "vec4DispatchFits"],
       "derive": {
         "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": [
         {
@@ -100,8 +87,8 @@
           "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
           }
         }
@@ -113,13 +100,11 @@
       "when": ["betaContract", "scalarDispatchFits"],
       "derive": {
         "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": [
         {
@@ -128,8 +113,8 @@
           "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
           }
         }
@@ -141,13 +126,11 @@
       "when": ["decayOnlyContract", "scalarDispatchFits"],
       "derive": {
         "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": [
         {
@@ -156,8 +139,8 @@
           "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
           }
         }

     "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",
+    "workgroupSize": "tunables.WORKGROUP_SIZE"
   },
   "when": ["workgroupFits"],
   "bindings": {
+    "a": { "arg": "aT", "elementType": "$gateElement", "length": "$gateItems" },
+    "dt_bias": { "arg": "dtBiasT", "elementType": "$paramElement", "length": "$headItems" },
+    "decay_scale": { "arg": "decayScaleT", "elementType": "$paramElement", "length": "$headItems" },
+    "b": { "arg": "bT", "elementType": "$gateElement", "length": "$gateItems" },
+    "decay": { "arg": "decayT", "elementType": "$gateElement", "length": "$gateItems" },
+    "beta": { "arg": "betaT", "elementType": "$gateElement", "length": "$gateItems" }
   },
   "variants": [
     {
       "when": ["betaContract", "vec4DispatchFits"],
       "derive": {
         "vectorized": true,
+        "hasBeta": "present.betaT",
         "gateElement": "\"vec4<f16>\" if gateDtype == \"float16\" else \"vec4<f32>\"",
         "paramElement": "\"vec4<f32>\"",
         "headItems": "headsVec4",
+        "gateItems": "gateVec4Count"
       },
       "passes": [
         {
           "shader": "linear-attention-gate.wgsl.jinja",
           "bindings": ["a", "dt_bias", "decay_scale", "b", "decay", "beta"],
           "dispatch": {
+            "x": "min(ceilDiv((gateItems), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((gateItems), (tunables.WORKGROUP_SIZE)), 65535)",
             "z": 1
           }
         }
       "when": ["decayOnlyContract", "vec4DispatchFits"],
       "derive": {
         "vectorized": true,
+        "hasBeta": "present.betaT",
         "gateElement": "\"vec4<f16>\" if gateDtype == \"float16\" else \"vec4<f32>\"",
         "paramElement": "\"vec4<f32>\"",
         "headItems": "headsVec4",
+        "gateItems": "gateVec4Count"
       },
       "passes": [
         {
           "shader": "linear-attention-gate.wgsl.jinja",
           "bindings": ["a", "dt_bias", "decay_scale", "decay"],
           "dispatch": {
+            "x": "min(ceilDiv((gateItems), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((gateItems), (tunables.WORKGROUP_SIZE)), 65535)",
             "z": 1
           }
         }
       "when": ["betaContract", "scalarDispatchFits"],
       "derive": {
         "vectorized": false,
+        "hasBeta": "present.betaT",
         "gateElement": "\"f16\" if gateDtype == \"float16\" else \"f32\"",
         "paramElement": "\"f32\"",
         "headItems": "numHeads",
+        "gateItems": "gateCount"
       },
       "passes": [
         {
           "shader": "linear-attention-gate.wgsl.jinja",
           "bindings": ["a", "dt_bias", "decay_scale", "b", "decay", "beta"],
           "dispatch": {
+            "x": "min(ceilDiv((gateItems), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((gateItems), (tunables.WORKGROUP_SIZE)), 65535)",
             "z": 1
           }
         }
       "when": ["decayOnlyContract", "scalarDispatchFits"],
       "derive": {
         "vectorized": false,
+        "hasBeta": "present.betaT",
         "gateElement": "\"f16\" if gateDtype == \"float16\" else \"f32\"",
         "paramElement": "\"f32\"",
         "headItems": "numHeads",
+        "gateItems": "gateCount"
       },
       "passes": [
         {
           "shader": "linear-attention-gate.wgsl.jinja",
           "bindings": ["a", "dt_bias", "decay_scale", "decay"],
           "dispatch": {
+            "x": "min(ceilDiv((gateItems), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((gateItems), (tunables.WORKGROUP_SIZE)), 65535)",
             "z": 1
           }
         }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "com.microsoft.LinearAttentionGate",
-  "id": "_com_microsoft_linearattentiongate_webgpu_4bb5397",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
@@ -8,14 +8,14 @@
     "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"],

 {
   "name": "com.microsoft.LinearAttentionGate",
+  "id": "_com_microsoft_linearattentiongate_webgpu_502eb17",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
     "algorithm": "sha256",
     "files": {
       "bench.json": "nSMkCE+adLNKSBkCYEFz8YeOJf6YjUsw4vEzFToXs5U=",
+      "linear-attention-gate.wgsl.jinja": "/e9pU3UP1X3ScwIjKeEPLVvqcznnSXVYQAvll/xqigE=",
+      "manifest.json": "SjgmqEe/oy28CeKd29htvOfS0sW/t7WqolTgm5l1U6k=",
+      "test.json": "NLi6ptg6CNprO6zHHud1Vxkf7JZduFkizDYvb6CttXA="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
       "vec4_gate_beta": ["linear-attention-gate.wgsl.jinja"],
       "vec4_gate": ["linear-attention-gate.wgsl.jinja"],

build/webgpu/test.json CHANGED Viewed

@@ -3,7 +3,7 @@
     {
       "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": {
@@ -27,7 +27,7 @@
     {
       "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": {
@@ -262,6 +262,124 @@
         "decayT": { "dtype": "float16", "shape": [5, 6], "tolerance": 0, "relTolerance": 0.002 },
         "betaT": { "dtype": "float16", "shape": [5, 6], "tolerance": 0, "relTolerance": 0.002 }
       }
     }
   ]
 }

     {
       "name": "rank3_h8_vec4_gate_beta",
       "provenance": {
+        "notes": "Head count 8 (divisible by four) exercises four-wide beta computation in the (B,T,H) layout; a 2e-6 tolerance covers float32 Softplus rounding near where its log1p-based series approximation switches."
       },
       "inputs": {
         "aT": {
     {
       "name": "rank2_h6_scalar_gate_beta",
       "provenance": {
+        "notes": "Head count 6 is not a multiple of four, so four-wide parameter mapping does not apply; a rank-2 (no batch dimension) input checks per-head computation one head at a time."
       },
       "inputs": {
         "aT": {
         "decayT": { "dtype": "float16", "shape": [5, 6], "tolerance": 0, "relTolerance": 0.002 },
         "betaT": { "dtype": "float16", "shape": [5, 6], "tolerance": 0, "relTolerance": 0.002 }
       }
+    },
+    {
+      "name": "ort_gate_float_decode_h32",
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.17, "scale": 6.0 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [32], "data": { "kind": "linspace", "start": -2.0, "end": 2.0 } },
+        "decayScaleT": { "dtype": "float32", "shape": [32], "data": { "kind": "linspace", "start": -4.0, "end": -0.1 } },
+        "bT": {
+          "dtype": "float32",
+          "shape": [1, 1, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.13, "scale": 6.0 }
+        }
+      },
+      "outputs": {
+        "decayT": { "dtype": "float32", "shape": [1, 1, 32], "tolerance": 0, "relTolerance": 0.000002 },
+        "betaT": { "dtype": "float32", "shape": [1, 1, 32], "tolerance": 0, "relTolerance": 0.000002 }
+      }
+    },
+    {
+      "name": "ort_gate_float_speculative_decode_tile_h32",
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [1, 4, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.17, "scale": 6.0 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [32], "data": { "kind": "linspace", "start": -2.0, "end": 2.0 } },
+        "decayScaleT": { "dtype": "float32", "shape": [32], "data": { "kind": "linspace", "start": -4.0, "end": -0.1 } },
+        "bT": {
+          "dtype": "float32",
+          "shape": [1, 4, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.13, "scale": 6.0 }
+        }
+      },
+      "outputs": {
+        "decayT": { "dtype": "float32", "shape": [1, 4, 32], "tolerance": 0, "relTolerance": 0.000002 },
+        "betaT": { "dtype": "float32", "shape": [1, 4, 32], "tolerance": 0, "relTolerance": 0.000002 }
+      }
+    },
+    {
+      "name": "ort_gate_float_decay_only_h16",
+      "inputs": {
+        "aT": {
+          "dtype": "float32",
+          "shape": [2, 3, 16],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.17, "scale": 6.0 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [16], "data": { "kind": "linspace", "start": -2.0, "end": 2.0 } },
+        "decayScaleT": { "dtype": "float32", "shape": [16], "data": { "kind": "linspace", "start": -4.0, "end": -0.1 } }
+      },
+      "outputs": { "decayT": { "dtype": "float32", "shape": [2, 3, 16], "tolerance": 0, "relTolerance": 0.000002 } }
+    },
+    {
+      "name": "ort_gate_float16_speculative_decode_tile_h32",
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [1, 4, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.17, "scale": 6.0 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [32], "data": { "kind": "linspace", "start": -2.0, "end": 2.0 } },
+        "decayScaleT": { "dtype": "float32", "shape": [32], "data": { "kind": "linspace", "start": -4.0, "end": -0.1 } },
+        "bT": {
+          "dtype": "float16",
+          "shape": [1, 4, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.13, "scale": 6.0 }
+        }
+      },
+      "outputs": {
+        "decayT": { "dtype": "float16", "shape": [1, 4, 32], "tolerance": 0, "relTolerance": 0.002 },
+        "betaT": { "dtype": "float16", "shape": [1, 4, 32], "tolerance": 0, "relTolerance": 0.002 }
+      }
+    },
+    {
+      "name": "ort_gate_float16_prefill_h32",
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [2, 37, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.17, "scale": 6.0 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [32], "data": { "kind": "linspace", "start": -2.0, "end": 2.0 } },
+        "decayScaleT": { "dtype": "float32", "shape": [32], "data": { "kind": "linspace", "start": -4.0, "end": -0.1 } },
+        "bT": {
+          "dtype": "float16",
+          "shape": [2, 37, 32],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.13, "scale": 6.0 }
+        }
+      },
+      "outputs": {
+        "decayT": { "dtype": "float16", "shape": [2, 37, 32], "tolerance": 0, "relTolerance": 0.002 },
+        "betaT": { "dtype": "float16", "shape": [2, 37, 32], "tolerance": 0, "relTolerance": 0.002 }
+      }
+    },
+    {
+      "name": "ort_gate_float16_ragged_tail_h7",
+      "inputs": {
+        "aT": {
+          "dtype": "float16",
+          "shape": [1, 5, 7],
+          "data": { "kind": "fillFloat32", "sinStep": 0.23, "cosStep": 0.17, "scale": 6.0 }
+        },
+        "dtBiasT": { "dtype": "float32", "shape": [7], "data": { "kind": "linspace", "start": -2.0, "end": 2.0 } },
+        "decayScaleT": { "dtype": "float32", "shape": [7], "data": { "kind": "linspace", "start": -4.0, "end": -0.1 } },
+        "bT": {
+          "dtype": "float16",
+          "shape": [1, 5, 7],
+          "data": { "kind": "fillFloat32", "sinStep": 0.31, "cosStep": 0.13, "scale": 6.0 }
+        }
+      },
+      "outputs": {
+        "decayT": { "dtype": "float16", "shape": [1, 5, 7], "tolerance": 0, "relTolerance": 0.002 },
+        "betaT": { "dtype": "float16", "shape": [1, 5, 7], "tolerance": 0, "relTolerance": 0.002 }
+      }
     }
   ]
 }