sync 6fdf6301e2bb

Browse files

Files changed (7) hide show

README.md +88 -0
build/webgpu/bench.json +0 -0
build/webgpu/manifest.json +148 -0
build/webgpu/metadata.json +27 -0
build/webgpu/ngram-hash-mapping-shared-prefix.wgsl.jinja +67 -0
build/webgpu/ngram-hash-mapping.wgsl.jinja +123 -0
build/webgpu/test.json +0 -0

README.md CHANGED Viewed

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

 ---
+library_name: kernels
 license: apache-2.0
+tags:
+- kernel
+- webgpu
+- wgsl
 ---
+# com.microsoft.NGramHashMapping
+`com.microsoft`  ·  ONNX Runtime contrib operator  ·  contrib since_version 1
+## Description
+Engram n-gram hash ids from compressed tokenizer ids. For every order `n` in `[2, max_ngram_size]` it mixes the `n` causal shifts of `input_ids` as `mix = ids[t] * multipliers[0] xor ... xor ids[t-n+1] * multipliers[n-1]`, the products wrapping in two's complement, then emits `mix` modulo each of that order's head vocabulary sizes. `past_ids`/`present_ids` carry the `max_ngram_size - 1` preceding ids, so chunked prefill and decode agree with one full-sequence run. Only `int32` ids are implemented: `int64` is a different answer, not a wider one, since its products wrap at another width.
+See the [ONNX Runtime `NGramHashMapping` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.NGramHashMapping) for the reference semantics.
+## Inputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `inputIdsT` | `input_ids` | `M` | `2` | — | Compressed tokenizer ids with shape `(batch_size, sequence_length)`. A zero-length sequence is accepted: it emits no hash ids and passes the carry state through unchanged. | required |
+| `multipliersT` | `multipliers` | `M` | `1` | — | Per-shift hash multipliers with shape `(max_ngram_size)`. Conventionally odd, but any value is accepted; the product wraps in two's complement rather than saturating. | required |
+| `vocabSizesT` | `vocab_sizes` | `M` | `1` | — | Per-output-head vocabulary sizes with shape `((max_ngram_size - 1) * n_head_per_ngram)`, conventionally prime and strictly positive. A non-positive entry has no meaningful modulo; this kernel emits a hash id of `0` for that head, which is what every GPU implementation of the operator does, while the reference CPU implementation rejects it instead. | required |
+| `pastIdsT` | `past_ids` | `M` | `2` | — | Optional ids for the `max_ngram_size - 1` positions preceding this call, with shape `(batch_size, max_ngram_size - 1)`. Right-aligned, so the last slot is the most recent id. When it is absent the missing history is `pad_id`, which is what a fresh sequence means. | optional |
+## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
+| --- | --- | --- | --- | --- | --- | --- |
+| `hashIdsT` | `hash_ids` | `M` | `3` | derived | Hash ids with shape `(batch_size, sequence_length, (max_ngram_size - 1) * n_head_per_ngram)`. The heads of order `n = 2` come first, then `n = 3`, and so on. | required |
+| `presentIdsT` | `present_ids` | `M` | `2` | derived | The trailing `max_ngram_size - 1` ids of `past_ids` followed by `input_ids`, with shape `(batch_size, max_ngram_size - 1)`. Feed it back as `past_ids` on the next call. It is always written, into a buffer distinct from `past_ids`. | required |
+## Attributes
+Attributes and default values (overridable per request):
+| Attribute | Default | Description |
+| --- | --- | --- |
+| `max_ngram_size` | — | Highest n-gram order, at least 2. It is baked into the rendered kernel, so the shift walk and the head arithmetic are compile-time; this package renders orders up to 16. |
+| `n_head_per_ngram` | — | Number of hash heads emitted for each n-gram order, at least 1. It is baked into the rendered kernel alongside `max_ngram_size`, and their product with `max_ngram_size - 1` is capped at 64 heads. |
+| `pad_id` | — | Compressed tokenizer id used for causal-shift positions before the beginning of the whole sequence. It must be representable as `int32`. |
+## Type constraints
+| Variable | Allowed dtypes |
+| --- | --- |
+| `M` | `int32` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `past_shared_prefix` — Stage each causal prefix once per token in workgroup memory, then distribute contiguous head outputs across lanes. The tile obeys device storage and invocation limits. Small combined coefficient tables retain the compact token-owned path to avoid staging and barrier overhead.
+- `fresh_shared_prefix` — Stage each causal prefix once per token in workgroup memory, then distribute contiguous head outputs across lanes. The tile obeys device storage and invocation limits. Small combined coefficient tables retain the compact token-owned path to avoid staging and barrier overhead.
+## 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 cases
+- [`ngram-hash-mapping-shared-prefix.wgsl.jinja`](build/webgpu/ngram-hash-mapping-shared-prefix.wgsl.jinja)
+- [`ngram-hash-mapping.wgsl.jinja`](build/webgpu/ngram-hash-mapping.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.
+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.
+```js
+import { getKernel } from "@huggingface/kernels";
+const kernel = await getKernel("webgpu-kernels/com.microsoft.NGramHashMapping", { version: 1 });
+const { hashIdsT, presentIdsT } = await kernel({
+  inputIdsT: { data: inputIdsTData, shape: [1, 1] },
+  multipliersT: { data: multipliersTData, shape: [3] },
+  vocabSizesT: { data: vocabSizesTData, shape: [4] },
+}, {
+  attrs: { max_ngram_size: 3, n_head_per_ngram: 2, pad_id: 9 },
+});
+```

build/webgpu/bench.json ADDED Viewed

The diff for this file is too large to render. See raw diff

build/webgpu/manifest.json ADDED Viewed

	@@ -0,0 +1,148 @@

+{
+  "domain": "com.microsoft",
+  "name": "NGramHashMapping",
+  "sinceVersion": 1,
+  "inputs": {
+    "inputIdsT": { "onnx": "input_ids", "dtype": "M", "rank": 2 },
+    "multipliersT": { "onnx": "multipliers", "dtype": "M", "rank": 1 },
+    "vocabSizesT": { "onnx": "vocab_sizes", "dtype": "M", "rank": 1 },
+    "pastIdsT": { "onnx": "past_ids", "dtype": "M", "rank": 2, "optional": true }
+  },
+  "outputs": {
+    "hashIdsT": {
+      "onnx": "hash_ids",
+      "dtype": "M",
+      "rank": 3,
+      "shape": "[dim(shapes.inputIdsT, 0), dim(shapes.inputIdsT, 1), (attrs.max_ngram_size - 1) * attrs.n_head_per_ngram]"
+    },
+    "presentIdsT": {
+      "onnx": "present_ids",
+      "dtype": "M",
+      "rank": 2,
+      "shape": "[dim(shapes.inputIdsT, 0), attrs.max_ngram_size - 1]"
+    }
+  },
+  "attributes": { "max_ngram_size": {}, "n_head_per_ngram": {}, "pad_id": {} },
+  "attributeConstraints": {
+    "max_ngram_size": { "required": true },
+    "n_head_per_ngram": { "required": true },
+    "pad_id": { "required": true }
+  },
+  "typeConstraints": { "M": ["int32"] },
+  "tunables": { "WORKGROUP_SIZE": { "default": 256 }, "ITEMS_PER_LANE": { "default": 4 } },
+  "derive": {
+    "foldedDispatchCapacity": "min(device.limits.maxComputeWorkgroupsPerDimension, 65535) * min(device.limits.maxComputeWorkgroupsPerDimension, 65535)",
+    "maxNgramSize": "attrs.max_ngram_size",
+    "headsPerNgram": "attrs.n_head_per_ngram",
+    "stateLength": "maxNgramSize - 1",
+    "numHeads": "stateLength * headsPerNgram",
+    "stagedTableScalars": "maxNgramSize + numHeads",
+    "padId": "attrs.pad_id",
+    "idsRankOk": "ranks.inputIdsT == 2",
+    "batchSize": "dim(shapes.inputIdsT, 0) if idsRankOk else 0",
+    "seqLength": "dim(shapes.inputIdsT, 1) if idsRankOk else 0",
+    "spanLength": "seqLength + stateLength",
+    "sharedTileSlots": "max(1, min(tunables.WORKGROUP_SIZE, spanLength, floor(device.limits.maxComputeWorkgroupStorageSize / (max(1, stateLength) * 4))))",
+    "sharedSpanTiles": "ceilDiv(spanLength, sharedTileSlots)",
+    "sharedDispatchOk": "sharedSpanTiles <= foldedDispatchCapacity",
+    "sharedMemoryOk": "sharedTileSlots * stateLength * 4 <= device.limits.maxComputeWorkgroupStorageSize",
+    "tileSlots": "tunables.WORKGROUP_SIZE * tunables.ITEMS_PER_LANE",
+    "spanTiles": "ceilDiv(spanLength, max(1, tileSlots))",
+    "idScalar": "dtypes.M",
+    "attrsOk": "maxNgramSize >= 2 and maxNgramSize <= 16 and headsPerNgram >= 1 and numHeads <= 64 and padId >= -2147483648 and padId <= 2147483647",
+    "dtypeOk": "tensorDtypes.inputIdsT == \"int32\" and tensorDtypes.multipliersT == \"int32\" and tensorDtypes.vocabSizesT == \"int32\" and tensorDtypes.hashIdsT == \"int32\" and tensorDtypes.presentIdsT == \"int32\"",
+    "shapeOk": "idsRankOk and ranks.multipliersT == 1 and dim(shapes.multipliersT, 0) == maxNgramSize and ranks.vocabSizesT == 1 and dim(shapes.vocabSizesT, 0) == numHeads and ranks.hashIdsT == 3 and sameShape(shapes.hashIdsT, [batchSize, seqLength, numHeads]) and ranks.presentIdsT == 2 and sameShape(shapes.presentIdsT, [batchSize, stateLength])",
+    "dispatchOk": "batchSize <= min(device.limits.maxComputeWorkgroupsPerDimension, 65535) and spanTiles <= foldedDispatchCapacity",
+    "workgroupOk": "tunables.WORKGROUP_SIZE > 0 and tunables.WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.WORKGROUP_SIZE <= device.limits.maxComputeWorkgroupSizeX and tunables.ITEMS_PER_LANE >= 1 and floor(tunables.ITEMS_PER_LANE) == tunables.ITEMS_PER_LANE",
+    "baseContract": "attrsOk and dtypeOk and shapeOk and dispatchOk and workgroupOk",
+    "noPastContract": "baseContract and not present.pastIdsT",
+    "pastContract": "baseContract and present.pastIdsT and tensorDtypes.pastIdsT == \"int32\" and ranks.pastIdsT == 2 and sameShape(shapes.pastIdsT, [batchSize, stateLength])"
+  },
+  "bindings": {
+    "input_ids": { "arg": "inputIdsT", "elementType": "$idScalar" },
+    "multipliers": { "arg": "multipliersT", "elementType": "$idScalar", "length": "$maxNgramSize" },
+    "vocab_sizes": { "arg": "vocabSizesT", "elementType": "$idScalar", "length": "$numHeads" },
+    "past_ids": { "arg": "pastIdsT", "elementType": "$idScalar" },
+    "hash_ids": { "arg": "hashIdsT", "elementType": "$idScalar" },
+    "present_ids": { "arg": "presentIdsT", "elementType": "$idScalar" },
+    "params": { "struct": [{ "name": "seqLength", "type": "u32", "value": "seqLength" }] }
+  },
+  "variants": [
+    {
+      "id": "past",
+      "when": ["pastContract"],
+      "derive": { "hasPast": true },
+      "passes": [
+        {
+          "id": "main",
+          "name": "NGramHashMapping.Past",
+          "shader": "ngram-hash-mapping.wgsl.jinja",
+          "bindings": ["input_ids", "multipliers", "vocab_sizes", "past_ids", "hash_ids", "present_ids", "params"],
+          "dispatch": {
+            "x": "min(spanTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "batchSize",
+            "z": "ceilDiv(spanTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "past_shared_prefix",
+      "when": ["pastContract", "sharedDispatchOk", "sharedMemoryOk"],
+      "derive": { "hasPast": true },
+      "passes": [
+        {
+          "id": "main",
+          "name": "NGramHashMapping.Past.SharedPrefix",
+          "shader": "ngram-hash-mapping-shared-prefix.wgsl.jinja",
+          "bindings": ["input_ids", "multipliers", "vocab_sizes", "past_ids", "hash_ids", "present_ids", "params"],
+          "dispatch": {
+            "x": "min(sharedSpanTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "batchSize",
+            "z": "ceilDiv(sharedSpanTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "priority": 10,
+      "demoteWhen": ["stagedTableScalars < 59"]
+    },
+    {
+      "id": "fresh",
+      "when": ["noPastContract"],
+      "derive": { "hasPast": false },
+      "passes": [
+        {
+          "id": "main",
+          "name": "NGramHashMapping.Fresh",
+          "shader": "ngram-hash-mapping.wgsl.jinja",
+          "bindings": ["input_ids", "multipliers", "vocab_sizes", "hash_ids", "present_ids", "params"],
+          "dispatch": {
+            "x": "min(spanTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "batchSize",
+            "z": "ceilDiv(spanTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ]
+    },
+    {
+      "id": "fresh_shared_prefix",
+      "when": ["noPastContract", "sharedDispatchOk", "sharedMemoryOk"],
+      "derive": { "hasPast": false },
+      "passes": [
+        {
+          "id": "main",
+          "name": "NGramHashMapping.Fresh.SharedPrefix",
+          "shader": "ngram-hash-mapping-shared-prefix.wgsl.jinja",
+          "bindings": ["input_ids", "multipliers", "vocab_sizes", "hash_ids", "present_ids", "params"],
+          "dispatch": {
+            "x": "min(sharedSpanTiles, DISPATCH_FOLD_WIDTH)",
+            "y": "batchSize",
+            "z": "ceilDiv(sharedSpanTiles, DISPATCH_FOLD_WIDTH)"
+          }
+        }
+      ],
+      "priority": 10,
+      "demoteWhen": ["stagedTableScalars < 59"]
+    }
+  ]
+}

build/webgpu/metadata.json ADDED Viewed

	@@ -0,0 +1,27 @@

+{
+  "name": "com.microsoft.NGramHashMapping",
+  "id": "_com_microsoft_ngramhashmapping_webgpu_2772cbc",
+  "version": 1,
+  "license": "Apache-2.0",
+  "backend": { "type": "webgpu" },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "bench.json": "9TYxNQMMA7MFXU6lsmktq1DuTX+F96McWtO10Wi6Kow=",
+      "manifest.json": "GIvp96ai3vrBQBgRUoQPH44v0GP+JsH/ps8PEWjC0io=",
+      "ngram-hash-mapping-shared-prefix.wgsl.jinja": "AxbE9+l3w7VS9v44/AxMP+ncC4yxdss9T1A44ZQ+9Fw=",
+      "ngram-hash-mapping.wgsl.jinja": "2PNqcut+20bhW9G5cOGwNKBRjzpnSIErvj3ZquDQX14=",
+      "test.json": "zzm2j7sopidQ0yqpDfJEr3BWCj0Z+Vl0okbjsSeEbfw="
+    }
+  },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.1",
+    "variants": {
+      "past": ["ngram-hash-mapping.wgsl.jinja"],
+      "past_shared_prefix": ["ngram-hash-mapping-shared-prefix.wgsl.jinja"],
+      "fresh": ["ngram-hash-mapping.wgsl.jinja"],
+      "fresh_shared_prefix": ["ngram-hash-mapping-shared-prefix.wgsl.jinja"]
+    }
+  }
+}

build/webgpu/ngram-hash-mapping-shared-prefix.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,67 @@

+{{ env.wgsl.resourceDeclarations }}
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+const TILE: u32 = {{ sharedTileSlots }}u;
+const STATE_LENGTH: u32 = {{ stateLength }}u;
+const HEADS: u32 = {{ numHeads }}u;
+const HEADS_PER_ORDER: u32 = {{ headsPerNgram }}u;
+{% if not hasPast %}
+const PAD_ID: i32 = {{ padId }};
+{% endif %}
+// Order-major storage coalesces prefix writes and shares each value across heads.
+var<workgroup> prefixes: array<i32, {{ sharedTileSlots * stateLength }}>;
+fn positive_mod(value: i32, divisor: i32) -> i32 {
+  let r = value % max(divisor, 1);
+  return select(r, r + divisor, r < 0);
+}
+@compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  let b = wg.y;
+  let tile = wg.x + wg.z * {{ DISPATCH_FOLD_WIDTH }}u;
+  let seq = params.seqLength;
+  let tileBase = tile * TILE;
+  let inputBase = b * seq;
+{% if hasPast %}
+  let pastBase = b * STATE_LENGTH;
+{% endif %}
+  if (lid.x < TILE) {
+    let t = tileBase + lid.x;
+    if (t < seq) {
+      // Unsigned products preserve the specified two's-complement wrapping.
+      var mix = bitcast<u32>(input_ids[inputBase + t]) * bitcast<u32>(multipliers[0u]);
+{% for k in range(1, maxNgramSize) %}
+{% if hasPast %}
+      var token{{ k }}: i32;
+      if (t >= {{ k }}u) { token{{ k }} = input_ids[inputBase + t - {{ k }}u]; }
+      else { token{{ k }} = past_ids[pastBase + STATE_LENGTH + t - {{ k }}u]; }
+{% else %}
+      var token{{ k }} = PAD_ID;
+      if (t >= {{ k }}u) { token{{ k }} = input_ids[inputBase + t - {{ k }}u]; }
+{% endif %}
+      mix = mix ^ (bitcast<u32>(token{{ k }}) * bitcast<u32>(multipliers[{{ k }}u]));
+      prefixes[{{ k - 1 }}u * TILE + lid.x] = bitcast<i32>(mix);
+{% endfor %}
+    } else if (t < seq + STATE_LENGTH) {
+{% if hasPast %}
+      var carry: i32;
+      if (t >= STATE_LENGTH) { carry = input_ids[inputBase + t - STATE_LENGTH]; }
+      else { carry = past_ids[pastBase + t]; }
+{% else %}
+      var carry = PAD_ID;
+      if (t >= STATE_LENGTH) { carry = input_ids[inputBase + t - STATE_LENGTH]; }
+{% endif %}
+      present_ids[b * STATE_LENGTH + t - seq] = carry;
+    }
+  }
+  // Every lane reaches the barrier, including carry-only and partial tiles.
+  workgroupBarrier();
+  var tokens = 0u;
+  if (tileBase < seq) { tokens = min(TILE, seq - tileBase); }
+  for (var index = lid.x; index < tokens * HEADS; index = index + WG) {
+    let token = index / HEADS;
+    let head = index % HEADS;
+    let mix = prefixes[(head / HEADS_PER_ORDER) * TILE + token];
+    hash_ids[(inputBase + tileBase) * HEADS + index] = positive_mod(mix, vocab_sizes[head]);
+  }
+}

build/webgpu/ngram-hash-mapping.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,123 @@

+{{ env.wgsl.resourceDeclarations }}
+// Engram n-gram hash ids. For every order n in [2, MAX_NGRAM] the kernel mixes the n causal
+// shifts of the compressed tokenizer ids,
+//   mix_n = ids[t] * m[0] ^ ids[t-1] * m[1] ^ ... ^ ids[t-n+1] * m[n-1],
+// and emits mix_n modulo each of that order's head vocabulary sizes. mix_n extends mix_{n-1}
+// by a single term, so one walk over the window produces every order.
+//
+// Each product is formed in u32 so that overflow wraps in two's complement, which is the
+// arithmetic the operator specifies. A signed multiply that overflows has an indeterminate
+// result in WGSL, so the bitcasts here are load-bearing rather than cosmetic.
+//
+// One dispatch covers both outputs. Workgroup row wg.y owns batch row b, and the tile index
+// folded across wg.x/wg.z walks a span of SEQ + STATE_LENGTH slots: the first SEQ slots hash a
+// token position, and the trailing STATE_LENGTH slots emit the carry state that the next call
+// feeds back as past_ids. A slot past the end of the span exits, so a short sequence costs one
+// workgroup rather than a second dispatch.
+const WG: u32 = {{ tunables.WORKGROUP_SIZE }}u;
+const ITEMS: u32 = {{ tunables.ITEMS_PER_LANE }}u;
+const STATE_LENGTH: u32 = {{ stateLength }}u;
+const PAD_ID: i32 = {{ padId }};
+// Euclidean remainder for a positive divisor: WGSL `%` truncates toward zero exactly as C
+// does, so a negative dividend needs the one correction the operator specifies. A head whose
+// vocabulary size is not strictly positive has no meaningful modulo, and clamping the divisor
+// to 1 makes its remainder 0 -- the value every GPU implementation of this operator emits for
+// that head, and the reason it stays a device-side division rather than a rejected input.
+fn positive_mod(value: i32, divisor: i32) -> i32 {
+  let r = value % max(divisor, 1);
+  return select(r, r + divisor, r < 0);
+}
+{% macro hash_position(interior) %}
+      let idBase = inputBase + t;
+      var mix = bitcast<u32>(input_ids[idBase]) * mult0;
+{% for k in range(1, maxNgramSize) %}
+{% if interior %}
+      mix = mix ^ (bitcast<u32>(input_ids[idBase - {{ k }}u]) * mult{{ k }});
+{% else %}
+      var token{{ k }} = PAD_ID;
+      if (t >= {{ k }}u) {
+        token{{ k }} = input_ids[idBase - {{ k }}u];
+      }
+{% if hasPast %}
+      else {
+        // past_ids is right-aligned, so position -1 is its last slot. t < k bounds the index
+        // inside the window, which is why no second range test is needed here.
+        token{{ k }} = past_ids[pastBase + STATE_LENGTH + t - {{ k }}u];
+      }
+{% endif %}
+      mix = mix ^ (bitcast<u32>(token{{ k }}) * mult{{ k }});
+{% endif %}
+      let mixOrder{{ k }} = bitcast<i32>(mix);
+{% endfor %}
+{% for n in range(1, maxNgramSize) %}
+{% for h in range(headsPerNgram) %}
+      let head{{ (n - 1) * headsPerNgram + h }} = positive_mod(mixOrder{{ n }}, vocab{{ (n - 1) * headsPerNgram + h }});
+{% endfor %}
+{% endfor %}
+      let outBase = idBase * {{ numHeads }}u;
+{% for h in range(numHeads) %}
+      hash_ids[outBase + {{ h }}u] = head{{ h }};
+{% endfor %}
+{% endmacro %}
+@compute @workgroup_size(WG)
+fn main(@builtin(workgroup_id) wg: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) {
+  // wg.y is a real axis -- one workgroup row per batch row -- so only the span tiles are
+  // folded, across wg.x and wg.z at the per-axis workgroup fold width.
+  let b = wg.y;
+  let tile = wg.x + wg.z * {{ DISPATCH_FOLD_WIDTH }}u;
+  let seq = params.seqLength;
+  let span = seq + STATE_LENGTH;
+  let inputBase = b * seq;
+  let tileBase = tile * (WG * ITEMS);
+{% if hasPast %}
+  let pastBase = b * STATE_LENGTH;
+{% endif %}
+  // The multiplier and vocabulary tables are attribute-sized and every position reads all of
+  // them, so an invocation stages them once instead of reloading them inside the mix walk.
+{% for k in range(maxNgramSize) %}
+  let mult{{ k }} = bitcast<u32>(multipliers[{{ k }}u]);
+{% endfor %}
+{% for h in range(numHeads) %}
+  let vocab{{ h }} = vocab_sizes[{{ h }}u];
+{% endfor %}
+  // An interior tile is one whose every slot hashes a token position and whose every causal
+  // shift lands inside input_ids. Both of the per-slot tests below, and the per-shift history
+  // fallback, are then statically false, which is worth one workgroup-uniform branch.
+  if (tileBase >= STATE_LENGTH && tileBase + WG * ITEMS <= seq) {
+    for (var item = 0u; item < ITEMS; item = item + 1u) {
+      let t = tileBase + lid.x + item * WG;
+{{ hash_position(true) }}
+    }
+    return;
+  }
+  for (var item = 0u; item < ITEMS; item = item + 1u) {
+    let slot = tileBase + lid.x + item * WG;
+    if (slot >= span) {
+      break;
+    }
+    if (slot < seq) {
+      let t = slot;
+{{ hash_position(false) }}
+    } else {
+      // Carry slot j holds the id at position seq - STATE_LENGTH + j of this call, and slot j
+      // of past_ids when the call is shorter than the window. Writing it from this dispatch is
+      // safe because present_ids never shares its buffer with past_ids here.
+      var carry = PAD_ID;
+      if (slot >= STATE_LENGTH) {
+        carry = input_ids[inputBase + slot - STATE_LENGTH];
+      }
+{% if hasPast %}
+      else {
+        carry = past_ids[pastBase + slot];
+      }
+{% endif %}
+      present_ids[b * STATE_LENGTH + slot - seq] = carry;
+    }
+  }
+}

build/webgpu/test.json ADDED Viewed

The diff for this file is too large to render. See raw diff