sync 91d990483a17

Browse files

Files changed (7) hide show

README.md +13 -9
build/webgpu/bench.json +0 -1
build/webgpu/manifest.json +74 -162
build/webgpu/metadata.json +14 -8
build/webgpu/scatter-flat-copy.wgsl.jinja +30 -30
build/webgpu/tensor-scatter.wgsl.jinja +40 -31
build/webgpu/test.json +50 -3

README.md CHANGED Viewed

@@ -18,17 +18,17 @@ See the [ONNX `TensorScatter` spec](https://onnx.ai/onnx/operators/onnx__TensorS
 ## Inputs
-| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
-| `past_cache` | `past` | `T` | runtime-selected; narrow integers and bool use 32-bit slots | — | — | Existing cache tensor with shape `(batch_size, ..., max_sequence_length, ...)`. | required |
-| `update` | `update` | `T` | runtime-selected; narrow integers and bool use 32-bit slots | — | — | New values to scatter in, with the same shape as `past_cache` except the sequence dimension equals `sequence_length`. | required |
-| `write_indices` | `writeIndices` | `I` | `uint32` | `1` | — | Logical int64 per-sample write offset into the cache sequence dimension; shape `(batch_size,)`, stored as uint32 by WebGPU, and assumed all zeros if absent. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `present_cache` | `present` | `T` | same as `past_cache` | same as `past_cache` | Updated cache; same shape as `past_cache`. | required |
 ## Attributes
@@ -48,7 +48,7 @@ Default values (overridable per request):
 ## 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
@@ -57,10 +57,14 @@ Default values (overridable per request):
 ## 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 | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
+| `past` | `past_cache` | `T` | runtime-selected; narrow integers and bool use 32-bit slots | — | — | Existing cache tensor with shape `(batch_size, ..., max_sequence_length, ...)`. | required |
+| `update` | — | `T` | runtime-selected; narrow integers and bool use 32-bit slots | — | — | New values to scatter in, with the same shape as `past_cache` except the sequence dimension equals `sequence_length`. | required |
+| `writeIndices` | `write_indices` | `I` | `uint32` | `1` | — | Logical int64 per-sample write offset into the cache sequence dimension; shape `(batch_size,)`, stored as uint32 by WebGPU, and assumed all zeros if absent. | optional |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `present` | `present_cache` | `T` | same as `past` | same as `past` | Updated cache; same shape as `past_cache`. | required |
 ## Attributes
 ## 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": "ai.onnx.TensorScatter",
   "cases": [
     {
       "name": "decode-kv-update-f32",

 {
   "cases": [
     {
       "name": "decode-kv-update-f32",

build/webgpu/manifest.json CHANGED Viewed

@@ -2,135 +2,39 @@
   "domain": "ai.onnx",
   "name": "TensorScatter",
   "sinceVersion": 24,
-  "description": "Functionally updates a KV cache tensor by scattering an `update` tensor into the `past_cache` along a sequence axis, producing `present_cache` with the same shape. Each batch sample's update is written at the offset given by `write_indices` (zero if omitted), either linearly or in wrap-around `circular` fashion.",
-  "inputs": [
-    {
-      "role": "past_cache",
-      "dtype": "T",
-      "description": "Existing cache tensor with shape `(batch_size, ..., max_sequence_length, ...)`."
-    },
-    {
-      "role": "update",
-      "dtype": "T",
-      "description": "New values to scatter in, with the same shape as `past_cache` except the sequence dimension equals `sequence_length`."
-    },
-    {
-      "role": "write_indices",
       "dtype": "I",
       "rank": 1,
       "optional": true,
-      "description": "Logical int64 per-sample write offset into the cache sequence dimension; shape `(batch_size,)`, stored as uint32 by WebGPU, and assumed all zeros if absent."
-    }
-  ],
-  "outputs": [
-    {
-      "role": "present_cache",
-      "dtype": "T",
-      "rank": "ranks.past_cache",
-      "description": "Updated cache; same shape as `past_cache`.",
-      "shape": "shapes.past_cache"
     }
-  ],
-  "attributes": { "axis": -2, "mode": "linear" },
-  "attributeDescriptions": {
-    "axis": "Sequence dimension of `past_cache` and `update`; cannot be 0 (the batch dimension). Default is `-2`.",
-    "mode": "Write mode: `linear` requires `write_indices + sequence_length <= max_sequence_length`; `circular` wraps the write index modulo `max_sequence_length`."
   },
   "attributeConstraints": { "mode": { "values": ["linear", "circular"] } },
   "typeConstraints": {
     "T": ["float32", "float16", "int32", "int16", "int8", "uint32", "uint8", "bool"],
     "I": ["int64"]
   },
-  "args": {
-    "past": { "kind": "tensor", "semantic": "past_cache", "role": "input" },
-    "update": { "kind": "tensor", "semantic": "update", "role": "input" },
-    "writeIndices": {
-      "kind": "tensor",
-      "semantic": "write_indices",
-      "role": "indices",
-      "dtype": "uint32",
-      "narrowing": "checked",
-      "required": false
-    },
-    "present": { "kind": "tensor", "semantic": "present_cache", "role": "output" }
-  },
-  "tunables": { "WORKGROUP_SIZE": 256 },
-  "derive": {},
-  "bindingSets": {
-    "copyCache": [
-      {
-        "name": "src",
-        "arg": "past",
-        "semantic": "past_cache",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$copyElementType"
-      },
-      {
-        "name": "dst",
-        "arg": "present",
-        "semantic": "present_cache",
-        "buffer": { "type": "storage" },
-        "elementType": "$copyElementType"
-      }
-    ],
-    "contiguousUpdate": [
-      {
-        "name": "update",
-        "arg": "update",
-        "semantic": "update",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "present",
-        "arg": "present",
-        "semantic": "present_cache",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.update)" }] }
-      }
-    ],
-    "indexedUpdate": [
-      {
-        "name": "update",
-        "arg": "update",
-        "semantic": "update",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "writeIndices",
-        "arg": "writeIndices",
-        "semantic": "write_indices",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "u32"
-      },
-      {
-        "name": "present",
-        "arg": "present",
-        "semantic": "present_cache",
-        "buffer": { "type": "storage" },
-        "elementType": "$scalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "count", "type": "u32", "value": "numel(shapes.update)" }] }
-      }
-    ]
   },
   "variants": [
     {
       "id": "no_write_indices",
       "priority": 0,
-      "when": ["not present.writeIndices", "ranks.past == ranks.update", "ranks.past == ranks.present", "ranks.past >= 2", "(attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.past) >= 1", "(attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.past) < ranks.past", "numel(shapes.past) == numel(shapes.present)", "f16Ok(dtypes.T)"],
-      "constants": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "copyElementType": "(\"vec4<\" ~ dtypes.T ~ \">\") if numel(shapes.present) % 4 == 0 else dtypes.T"
@@ -139,46 +43,50 @@
         {
           "id": "copy",
           "name": "TensorScatter.Copy",
-          "source": {
-            "shader": "scatter-flat-copy.wgsl.jinja",
-            "inputs": {
-              "count": "ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))",
-              "copyWidth": "16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4)",
-              "elementCount": "numel(shapes.present)",
-              "scalarTailSafe": "numel(shapes.present) % 4 != 0"
-            }
           },
-          "bindings": "copyCache",
           "dispatch": {
-            "gridStride": "ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           },
           "viewAlias": [{ "input": "src", "output": "dst" }]
         },
         {
           "id": "scatter",
           "name": "TensorScatter.Scatter",
-          "source": {
-            "shader": "tensor-scatter.wgsl.jinja",
-            "inputs": {
-              "pastShape": "shapes.past",
-              "updateShape": "shapes.update",
-              "rank": "ranks.past",
-              "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.past",
-              "circular": "attrs.mode == \"circular\"",
-              "hasIndices": "false"
-            }
           },
-          "bindings": "contiguousUpdate",
-          "dispatch": { "threads": "numel(shapes.update)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "with_write_indices",
       "priority": 20,
-      "when": ["present.writeIndices", "dtypes.I == \"u32\"", "ranks.past == ranks.update", "ranks.past == ranks.present", "ranks.past >= 2", "(attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.past) >= 1", "(attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.past) < ranks.past", "numel(shapes.past) == numel(shapes.present)", "f16Ok(dtypes.T)"],
-      "constants": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "copyElementType": "(\"vec4<\" ~ dtypes.T ~ \">\") if numel(shapes.present) % 4 == 0 else dtypes.T"
@@ -187,38 +95,42 @@
         {
           "id": "copy",
           "name": "TensorScatter.Copy",
-          "source": {
-            "shader": "scatter-flat-copy.wgsl.jinja",
-            "inputs": {
-              "count": "ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))",
-              "copyWidth": "16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4)",
-              "elementCount": "numel(shapes.present)",
-              "scalarTailSafe": "numel(shapes.present) % 4 != 0"
-            }
           },
-          "bindings": "copyCache",
           "dispatch": {
-            "gridStride": "ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))",
-            "workgroupSize": "tunables.WORKGROUP_SIZE"
           },
           "viewAlias": [{ "input": "src", "output": "dst" }]
         },
         {
           "id": "scatter",
           "name": "TensorScatter.Scatter",
-          "source": {
-            "shader": "tensor-scatter.wgsl.jinja",
-            "inputs": {
-              "pastShape": "shapes.past",
-              "updateShape": "shapes.update",
-              "rank": "ranks.past",
-              "axis": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.past",
-              "circular": "attrs.mode == \"circular\"",
-              "hasIndices": "true"
-            }
           },
-          "bindings": "indexedUpdate",
-          "dispatch": { "threads": "numel(shapes.update)", "workgroupSize": "tunables.WORKGROUP_SIZE" }
         }
       ]
     }

   "domain": "ai.onnx",
   "name": "TensorScatter",
   "sinceVersion": 24,
+  "inputs": {
+    "past": { "onnx": "past_cache", "dtype": "T" },
+    "update": { "dtype": "T" },
+    "writeIndices": {
+      "onnx": "write_indices",
       "dtype": "I",
       "rank": 1,
       "optional": true,
+      "storage": "uint32",
+      "narrowing": "checked"
     }
   },
+  "outputs": { "present": { "onnx": "present_cache", "dtype": "T", "rank": "ranks.past", "shape": "shapes.past" } },
+  "attributes": { "axis": { "default": -2 }, "mode": { "default": "linear" } },
   "attributeConstraints": { "mode": { "values": ["linear", "circular"] } },
   "typeConstraints": {
     "T": ["float32", "float16", "int32", "int16", "int8", "uint32", "uint8", "bool"],
     "I": ["int64"]
   },
+  "tunables": { "WORKGROUP_SIZE": { "default": 256 } },
+  "when": ["(attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.past) < ranks.past", "(attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.past) >= 1", "f16Ok(dtypes.T)", "numel(shapes.past) == numel(shapes.present)", "ranks.past == ranks.present", "ranks.past == ranks.update", "ranks.past >= 2"],
+  "bindings": {
+    "update": { "buffer": "read-only-storage", "elementType": "$scalar" },
+    "present": { "buffer": "storage", "elementType": "$scalar" },
+    "params": { "buffer": "uniform", "struct": [{ "name": "count", "type": "u32", "value": "numel(shapes.update)" }] },
+    "writeIndices": { "buffer": "read-only-storage", "elementType": "u32" }
   },
   "variants": [
     {
       "id": "no_write_indices",
       "priority": 0,
+      "when": ["not present.writeIndices"],
+      "derive": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "copyElementType": "(\"vec4<\" ~ dtypes.T ~ \">\") if numel(shapes.present) % 4 == 0 else dtypes.T"
         {
           "id": "copy",
           "name": "TensorScatter.Copy",
+          "shader": "scatter-flat-copy.wgsl.jinja",
+          "derive": {
+            "count": "ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))",
+            "copyWidth": "16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4)",
+            "elementCount": "numel(shapes.present)",
+            "scalarTailSafe": "numel(shapes.present) % 4 != 0"
           },
+          "bindings": [
+            { "arg": "past", "name": "src", "elementType": "$copyElementType" },
+            { "arg": "present", "name": "dst", "elementType": "$copyElementType" }
+          ],
           "dispatch": {
+            "x": "min(ceilDiv((ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           },
           "viewAlias": [{ "input": "src", "output": "dst" }]
         },
         {
           "id": "scatter",
           "name": "TensorScatter.Scatter",
+          "shader": "tensor-scatter.wgsl.jinja",
+          "derive": {
+            "pastShape": "shapes.past",
+            "updateShape": "shapes.update",
+            "rank": "ranks.past",
+            "axisSpec": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.past",
+            "circular": "attrs.mode == \"circular\"",
+            "hasIndices": "false"
           },
+          "bindings": ["update", "present", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.update)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.update)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "with_write_indices",
       "priority": 20,
+      "when": ["present.writeIndices", "dtypes.I == \"u32\""],
+      "derive": {
         "scalar": "dtypes.T",
         "usesF16": "dtypes.T == \"f16\"",
         "copyElementType": "(\"vec4<\" ~ dtypes.T ~ \">\") if numel(shapes.present) % 4 == 0 else dtypes.T"
         {
           "id": "copy",
           "name": "TensorScatter.Copy",
+          "shader": "scatter-flat-copy.wgsl.jinja",
+          "derive": {
+            "count": "ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))",
+            "copyWidth": "16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4)",
+            "elementCount": "numel(shapes.present)",
+            "scalarTailSafe": "numel(shapes.present) % 4 != 0"
           },
+          "bindings": [
+            { "arg": "past", "name": "src", "elementType": "$copyElementType" },
+            { "arg": "present", "name": "dst", "elementType": "$copyElementType" }
+          ],
           "dispatch": {
+            "x": "min(ceilDiv((ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((ceilDiv(numel(shapes.present), 16 if dtypes.T == \"f16\" and numel(shapes.present) % 4 != 0 else (8 if numel(shapes.present) % 4 != 0 else 4))), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
           },
           "viewAlias": [{ "input": "src", "output": "dst" }]
         },
         {
           "id": "scatter",
           "name": "TensorScatter.Scatter",
+          "shader": "tensor-scatter.wgsl.jinja",
+          "derive": {
+            "pastShape": "shapes.past",
+            "updateShape": "shapes.update",
+            "rank": "ranks.past",
+            "axisSpec": "attrs.axis if attrs.axis >= 0 else attrs.axis + ranks.past",
+            "circular": "attrs.mode == \"circular\"",
+            "hasIndices": "true"
           },
+          "bindings": ["update", "writeIndices", "present", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((numel(shapes.update)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((numel(shapes.update)), (tunables.WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,19 +1,25 @@
 {
   "name": "ai.onnx.TensorScatter",
-  "id": "_ai_onnx_tensorscatter_webgpu_d5b6d1c",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "27NgakY5cT0I8768syqpQg8+XMBWBQIPNUfS8FN4GHM=",
-      "manifest.json": "hIT+BgOyanYiB3FJJjnjjheedXZaLszhHIe7buQVVAM=",
-      "scatter-flat-copy.wgsl.jinja": "TcsxU7ecJvkki1Bo9K7kjXnovGgU7IKumZI6bDsxwUM=",
-      "tensor-scatter.wgsl.jinja": "Ys/UVlAuBJYHUJ0DwrUbTj/T2+2JBokOnwPIqMO7P6U=",
-      "test.json": "AiLV0oDQiO121/e1IpqXRfz1PNTcC9khxqUrBdw3/yY="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/ai.onnx.TensorScatter" }
 }

 {
   "name": "ai.onnx.TensorScatter",
+  "id": "_ai_onnx_tensorscatter_webgpu_fae7a8d",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "8pIi4/PPMqyk12pePykM6uwawjM0N6tsexFwiHFqySw=",
+      "manifest.json": "RY7RXbSuGDdCHSr8e4qqS5wG4oCULwzNV5C9/ToI6Ak=",
+      "scatter-flat-copy.wgsl.jinja": "mXcfAXV6Nxe+Y+cvem8Wwwt0BeqRQXaiCNhQoiynX0M=",
+      "tensor-scatter.wgsl.jinja": "cDtR3E7EOm41/YKXMPKfPP3YCakYj/pbU85wt2obnY0=",
+      "test.json": "H/VHqdsOkO3A3dM1LSKrAmk5XO2vlrDyry0EXKm/2X4="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "no_write_indices": ["scatter-flat-copy.wgsl.jinja", "tensor-scatter.wgsl.jinja"],
+      "with_write_indices": ["scatter-flat-copy.wgsl.jinja", "tensor-scatter.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/scatter-flat-copy.wgsl.jinja CHANGED Viewed

@@ -2,46 +2,46 @@
 // pass. The element width is owned by the binding element type: when
 // numel(output) % 4 == 0 the bindings rebind as vec4<T> (128-bit loads).
 // Tail-safe modes keep scalar bindings but copy a compile-time group of adjacent
-// elements per invocation, guarding only the final partial group. A grid-stride
-// loop lets the dispatch stay within maxComputeWorkgroupsPerDimension.
 {% if usesF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% set wg = workgroupSize if workgroupSize is defined else tunables.WORKGROUP_SIZE %}
-const COUNT: u32 = {{ source.count }}u;
 const WG: u32 = {{ wg }}u;
 @compute @workgroup_size({{ wg }})
-fn main(
-  @builtin(global_invocation_id) gid: vec3<u32>,
-  @builtin(num_workgroups) nwg: vec3<u32>
-) {
-  let stride = nwg.x * WG;
-  for (var i = gid.x; i < COUNT; i += stride) {
-{% if source.scalarTailSafe %}
-    let base = i * {{ source.copyWidth }}u;
-    dst[base] = src[base];
-    if (base + 1u < {{ source.elementCount }}u) { dst[base + 1u] = src[base + 1u]; }
-    if (base + 2u < {{ source.elementCount }}u) { dst[base + 2u] = src[base + 2u]; }
-    if (base + 3u < {{ source.elementCount }}u) { dst[base + 3u] = src[base + 3u]; }
-    if (base + 4u < {{ source.elementCount }}u) { dst[base + 4u] = src[base + 4u]; }
-    if (base + 5u < {{ source.elementCount }}u) { dst[base + 5u] = src[base + 5u]; }
-    if (base + 6u < {{ source.elementCount }}u) { dst[base + 6u] = src[base + 6u]; }
-    if (base + 7u < {{ source.elementCount }}u) { dst[base + 7u] = src[base + 7u]; }
-{% if source.copyWidth == 16 %}
-    if (base + 8u < {{ source.elementCount }}u) { dst[base + 8u] = src[base + 8u]; }
-    if (base + 9u < {{ source.elementCount }}u) { dst[base + 9u] = src[base + 9u]; }
-    if (base + 10u < {{ source.elementCount }}u) { dst[base + 10u] = src[base + 10u]; }
-    if (base + 11u < {{ source.elementCount }}u) { dst[base + 11u] = src[base + 11u]; }
-    if (base + 12u < {{ source.elementCount }}u) { dst[base + 12u] = src[base + 12u]; }
-    if (base + 13u < {{ source.elementCount }}u) { dst[base + 13u] = src[base + 13u]; }
-    if (base + 14u < {{ source.elementCount }}u) { dst[base + 14u] = src[base + 14u]; }
-    if (base + 15u < {{ source.elementCount }}u) { dst[base + 15u] = src[base + 15u]; }
 {% endif %}
 {% else %}
-    dst[i] = src[i];
 {% endif %}
-  }
 }

 // pass. The element width is owned by the binding element type: when
 // numel(output) % 4 == 0 the bindings rebind as vec4<T> (128-bit loads).
 // Tail-safe modes keep scalar bindings but copy a compile-time group of adjacent
+// elements per invocation, guarding only the final partial group.
+{% set scalarTailSafe = scalarTailSafe is defined and scalarTailSafe %}
+{% set copyWidth = copyWidth | default(0) %}
+{% set elementCount = elementCount | default(0) %}
 {% if usesF16 %}
 enable f16;
 {% endif %}
 {{ env.wgsl.resourceDeclarations }}
 {% set wg = workgroupSize if workgroupSize is defined else tunables.WORKGROUP_SIZE %}
+const COUNT: u32 = {{ count }}u;
 const WG: u32 = {{ wg }}u;
 @compute @workgroup_size({{ wg }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let i = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
+  if (i >= COUNT) {
+    return;
+  }
+{% if scalarTailSafe %}
+  let base = i * {{ copyWidth }}u;
+  dst[base] = src[base];
+  if (base + 1u < {{ elementCount }}u) { dst[base + 1u] = src[base + 1u]; }
+  if (base + 2u < {{ elementCount }}u) { dst[base + 2u] = src[base + 2u]; }
+  if (base + 3u < {{ elementCount }}u) { dst[base + 3u] = src[base + 3u]; }
+  if (base + 4u < {{ elementCount }}u) { dst[base + 4u] = src[base + 4u]; }
+  if (base + 5u < {{ elementCount }}u) { dst[base + 5u] = src[base + 5u]; }
+  if (base + 6u < {{ elementCount }}u) { dst[base + 6u] = src[base + 6u]; }
+  if (base + 7u < {{ elementCount }}u) { dst[base + 7u] = src[base + 7u]; }
+{% if copyWidth == 16 %}
+  if (base + 8u < {{ elementCount }}u) { dst[base + 8u] = src[base + 8u]; }
+  if (base + 9u < {{ elementCount }}u) { dst[base + 9u] = src[base + 9u]; }
+  if (base + 10u < {{ elementCount }}u) { dst[base + 10u] = src[base + 10u]; }
+  if (base + 11u < {{ elementCount }}u) { dst[base + 11u] = src[base + 11u]; }
+  if (base + 12u < {{ elementCount }}u) { dst[base + 12u] = src[base + 12u]; }
+  if (base + 13u < {{ elementCount }}u) { dst[base + 13u] = src[base + 13u]; }
+  if (base + 14u < {{ elementCount }}u) { dst[base + 14u] = src[base + 14u]; }
+  if (base + 15u < {{ elementCount }}u) { dst[base + 15u] = src[base + 15u]; }
 {% endif %}
 {% else %}
+  dst[i] = src[i];
 {% endif %}
 }

build/webgpu/tensor-scatter.wgsl.jinja CHANGED Viewed

@@ -1,63 +1,60 @@
 {% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
 {% if note == "dispatch-limit" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension dispatch limit (outputs > 16.7M elements).
 {% elif note == "limit" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
 {% elif note == "device-axis" %}
-  // The flat dispatch is folded across x/y at the device's per-axis workgroup
-  // limit; gid.y carries the high portion of the output index.
 {% elif note == "vec4-limit" %}
-  // 2D-folded flat vec4 index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit (the dispatch caps x and spills into y).
 {% elif note == "element-limit" %}
   // 2D-folded flat element index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit.
 {% elif note == "dispatch" %}
-  // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension dispatch limit.
 {% endif %}
 {% if bound == "" %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
 {%- elif guardInline %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) { return; }
 {%- else %}
-  let {{ name }} = gid.x + gid.y * nwg.x * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
   }
 {%- endif %}
 {% endmacro %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
-{% set max_seq = source.pastShape[source.axis] %}
-{% set seq_len = source.updateShape[source.axis] %}
 {% set dim0_stride = namespace(value=1) %}
-{% for j in range(1, source.rank) %}{% set dim0_stride.value = dim0_stride.value * source.updateShape[j] %}{% endfor %}
 {% set axis_stride = namespace(value=1) %}
-{% for j in range(source.axis + 1, source.rank) %}{% set axis_stride.value = axis_stride.value * source.updateShape[j] %}{% endfor %}
 // Map a flat index into the `update` tensor onto a flat index into `present`,
 // using `cache_axis` as the sequence-axis coordinate.
 fn present_offset(update_index: u32, cache_axis: u32) -> u32 {
   var rem = update_index;
   var offset = 0u;
-{% for d in range(source.rank) %}
 {% set u_stride = namespace(value=1) %}
-{% for j in range(d + 1, source.rank) %}{% set u_stride.value = u_stride.value * source.updateShape[j] %}{% endfor %}
 {% set p_stride = namespace(value=1) %}
-{% for j in range(d + 1, source.rank) %}{% set p_stride.value = p_stride.value * source.pastShape[j] %}{% endfor %}
-{% if d != source.axis %}
   let coord{{ d }} = rem / {{ max(u_stride.value, 1) }}u;
 {% endif %}
   rem = rem % {{ max(u_stride.value, 1) }}u;
-{% if d == source.axis %}
   offset = offset + cache_axis * {{ p_stride.value }}u;
 {% else %}
   offset = offset + coord{{ d }} * {{ p_stride.value }}u;
@@ -67,21 +64,33 @@ fn present_offset(update_index: u32, cache_axis: u32) -> u32 {
 }
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
-fn main(@builtin(global_invocation_id) gid: vec3<u32>, @builtin(num_workgroups) nwg: vec3<u32>) {
 {{ flat_index_2d("j") }}
   // batch_idx is the leading (dim-0) coordinate; s is the coordinate along the
   // sequence axis. write_indices is indexed per batch sample.
-{% if source.hasIndices %}
   let batch_idx = j / {{ max(dim0_stride.value, 1) }}u;
 {% endif %}
   let s = (j / {{ max(axis_stride.value, 1) }}u) % {{ max(seq_len, 1) }}u;
-{% if source.hasIndices %}
   let w = writeIndices[batch_idx];
 {% else %}
   let w = 0u;
 {% endif %}
-{% if source.circular %}
-  let cache_axis = (w + s) % {{ max_seq }}u;
 {% else %}
   let cache_axis = w + s;
 {% endif %}

 {% macro flat_index_2d(name="i", bound="params.count", guardInline=false, note="dispatch-limit") %}
 {% if note == "dispatch-limit" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width (outputs > 16.7M elements).
 {% elif note == "limit" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
 {% elif note == "device-axis" %}
+  // The flat dispatch is folded across x/y at a fixed per-axis workgroup
+  // width; gid.y carries the high portion of the output index.
 {% elif note == "vec4-limit" %}
+  // 2D-folded flat vec4 index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width (the dispatch caps x and spills into y).
 {% elif note == "element-limit" %}
   // 2D-folded flat element index: gid.y carries the high bits past the
+  // dispatch's per-axis workgroup fold width.
 {% elif note == "dispatch" %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
 {% endif %}
 {% if bound == "" %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
 {%- elif guardInline %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) { return; }
 {%- else %}
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ tunables.WORKGROUP_SIZE }}u;
   if ({{ name }} >= {{ bound }}) {
     return;
   }
 {%- endif %}
 {% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
+{% set max_seq = pastShape[axisSpec] %}
+{% set seq_len = updateShape[axisSpec] %}
 {% set dim0_stride = namespace(value=1) %}
+{% for j in range(1, rank) %}{% set dim0_stride.value = dim0_stride.value * updateShape[j] %}{% endfor %}
 {% set axis_stride = namespace(value=1) %}
+{% for j in range(axisSpec + 1, rank) %}{% set axis_stride.value = axis_stride.value * updateShape[j] %}{% endfor %}
 // Map a flat index into the `update` tensor onto a flat index into `present`,
 // using `cache_axis` as the sequence-axis coordinate.
 fn present_offset(update_index: u32, cache_axis: u32) -> u32 {
   var rem = update_index;
   var offset = 0u;
+{% for d in range(rank) %}
 {% set u_stride = namespace(value=1) %}
+{% for j in range(d + 1, rank) %}{% set u_stride.value = u_stride.value * updateShape[j] %}{% endfor %}
 {% set p_stride = namespace(value=1) %}
+{% for j in range(d + 1, rank) %}{% set p_stride.value = p_stride.value * pastShape[j] %}{% endfor %}
+{% if d != axisSpec %}
   let coord{{ d }} = rem / {{ max(u_stride.value, 1) }}u;
 {% endif %}
   rem = rem % {{ max(u_stride.value, 1) }}u;
+{% if d == axisSpec %}
   offset = offset + cache_axis * {{ p_stride.value }}u;
 {% else %}
   offset = offset + coord{{ d }} * {{ p_stride.value }}u;
 }
 @compute @workgroup_size({{ tunables.WORKGROUP_SIZE }})
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 {{ flat_index_2d("j") }}
   // batch_idx is the leading (dim-0) coordinate; s is the coordinate along the
   // sequence axis. write_indices is indexed per batch sample.
+{% if hasIndices %}
   let batch_idx = j / {{ max(dim0_stride.value, 1) }}u;
 {% endif %}
   let s = (j / {{ max(axis_stride.value, 1) }}u) % {{ max(seq_len, 1) }}u;
+{% if hasIndices %}
   let w = writeIndices[batch_idx];
 {% else %}
   let w = 0u;
 {% endif %}
+{% if circular %}
+{% if seq_len > max_seq %}
+  // More updates than cache slots: positions wrap onto each other and the
+  // operator's sequential order makes the LAST update to a slot win. Only the
+  // final max_seq updates can be that last writer; earlier ones stay silent so
+  // the result does not depend on invocation scheduling.
+  if (s + {{ max_seq }}u < {{ seq_len }}u) {
+    return;
+  }
+{% endif %}
+  // Reduce both terms before adding them: a write index near the u32 ceiling
+  // makes the sum wrap, and a wrapped sum lands on a different slot than the
+  // one the operator names unless the capacity happens to divide 2^32.
+  let cache_axis = (w % {{ max_seq }}u + s % {{ max_seq }}u) % {{ max_seq }}u;
 {% else %}
   let cache_axis = w + s;
 {% endif %}

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "ai.onnx.TensorScatter",
   "fixtureArrays": {
     "ort_linear_4d_spec_example_input_past": [1, 2, 3, 4, 5, 5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 4, 3, 2, 1, 0, 1, 2, 3, 4, 5, 5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 4, 3, 2, 1, 0],
     "ort_linear_3d_default_axis_input_past": [1, 2, 3, 4, 5, 5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 5, 4, 3, 2, 1, 1, 2, 3, 4, 5, 5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 5, 4, 3, 2, 1, 1, 2, 3, 4, 5, 5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 5, 4, 3, 2, 1],
@@ -45,7 +44,7 @@
     {
       "name": "decode_kv_update_f32_axis_minus2_compact",
       "provenance": {
-        "notes": "Compact sibling for the decode-style TensorScatter benchmark; preserves axis=-2, per-batch write indices, full-cache copy, and one-token update rows."
       },
       "attrs": { "mode": "linear", "axis": -2 },
       "inputs": {
@@ -436,7 +435,7 @@
       "name": "ort_standard_int16_linear_extremes",
       "provenance": {
         "source": "onnxruntime/core/providers/cpu/llm/tensorscatter.cc",
-        "notes": "ONNX TensorScatter-24 permits int16; both signed extrema guard the runtime's widened i32 copy-and-overwrite route."
       },
       "attrs": { "axis": 2, "mode": "linear" },
       "inputs": {
@@ -520,6 +519,54 @@
         "writeIndices": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [1] } }
       },
       "outputs": { "present": { "dtype": "bool", "shape": [1, 1, 4], "tolerance": 0 } }
     }
   ]
 }

 {
   "fixtureArrays": {
     "ort_linear_4d_spec_example_input_past": [1, 2, 3, 4, 5, 5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 4, 3, 2, 1, 0, 1, 2, 3, 4, 5, 5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 4, 3, 2, 1, 0],
     "ort_linear_3d_default_axis_input_past": [1, 2, 3, 4, 5, 5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 5, 4, 3, 2, 1, 1, 2, 3, 4, 5, 5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 5, 4, 3, 2, 1, 1, 2, 3, 4, 5, 5, 6, 7, 8, 9, 8, 7, 6, 5, 4, 5, 4, 3, 2, 1],
     {
       "name": "decode_kv_update_f32_axis_minus2_compact",
       "provenance": {
+        "notes": "A compact decode-style TensorScatter uses axis=-2, per-batch write indices, a full-cache copy, and one-token update rows."
       },
       "attrs": { "mode": "linear", "axis": -2 },
       "inputs": {
       "name": "ort_standard_int16_linear_extremes",
       "provenance": {
         "source": "onnxruntime/core/providers/cpu/llm/tensorscatter.cc",
+        "notes": "ONNX TensorScatter-24 permits int16; both signed extrema verify logical int16 values stored in i32 slots during copy and overwrite."
       },
       "attrs": { "axis": 2, "mode": "linear" },
       "inputs": {
         "writeIndices": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [1] } }
       },
       "outputs": { "present": { "dtype": "bool", "shape": [1, 1, 4], "tolerance": 0 } }
+    },
+    {
+      "name": "circular_update_longer_than_capacity_last_write_wins",
+      "provenance": {
+        "notes": "Six update rows enter a four-slot circular cache from write index 3: rows 0 and 4 both land on slot 3, while rows 1 and 5 land on slot 0. Sequence order makes the later row determine each collision, producing [[60,61],[30,31],[40,41],[50,51]] regardless of GPU scheduling."
+      },
+      "attrs": { "mode": "circular" },
+      "inputs": {
+        "past": { "dtype": "float32", "shape": [1, 4, 2], "data": { "kind": "constant", "value": 0.0 } },
+        "update": {
+          "dtype": "float32",
+          "shape": [1, 6, 2],
+          "data": {
+            "kind": "values",
+            "values": [10.0, 11.0, 20.0, 21.0, 30.0, 31.0, 40.0, 41.0, 50.0, 51.0, 60.0, 61.0]
+          }
+        },
+        "writeIndices": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [3] } }
+      },
+      "outputs": { "present": { "dtype": "float32", "shape": [1, 4, 2], "tolerance": 0 } }
+    },
+    {
+      "name": "circular_write_index_at_u32_ceiling",
+      "provenance": {
+        "notes": "A circular cache reduces the write index modulo the capacity, and the sum of index and offset was formed before the reduction. At a write index near the u32 ceiling that sum wraps, so the second update lands on the slot the first one already took instead of the next one, and the two race. Slots 0 and 1 here are what int64 arithmetic names."
+      },
+      "attrs": { "mode": "circular" },
+      "inputs": {
+        "past": {
+          "dtype": "float32",
+          "shape": [1, 3, 2],
+          "data": { "kind": "values", "values": [-1.0, -2.0, -3.0, -4.0, -5.0, -6.0] }
+        },
+        "update": {
+          "dtype": "float32",
+          "shape": [1, 2, 2],
+          "data": { "kind": "values", "values": [10.0, 11.0, 20.0, 21.0] }
+        },
+        "writeIndices": { "dtype": "uint32", "shape": [1], "data": { "kind": "values", "values": [4294967295] } }
+      },
+      "outputs": {
+        "present": {
+          "dtype": "float32",
+          "shape": [1, 3, 2],
+          "data": { "kind": "values", "values": [10.0, 11.0, 20.0, 21.0, -5.0, -6.0] },
+          "tolerance": 0
+        }
+      }
     }
   ]
 }