sync 91d990483a17

Browse files

Files changed (7) hide show

README.md +22 -18
build/webgpu/bench.json +0 -1
build/webgpu/embed-normalize.wgsl.jinja +2 -6
build/webgpu/embed-sum.wgsl.jinja +4 -8
build/webgpu/manifest.json +368 -993
build/webgpu/metadata.json +48 -8
build/webgpu/test.json +0 -1

README.md CHANGED Viewed

@@ -18,25 +18,25 @@ See the [ONNX Runtime `EmbedLayerNormalization` contrib-operator spec](https://g
 ## Inputs
-| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
-| `input_ids` | `inputIdsT` | `T1` | `int32` | `2` | — | Word ids of shape `(batch_size, sequence_length)`. | required |
-| `segment_ids` | `segmentIdsT` | `T1` | `int32` | `2` | — | Segment ids `(batch_size, sequence_length)`. Requires `segment_embedding`; when omitted with that table present, every token uses row 0. Values must be valid non-negative table-row indices. | optional |
-| `word_embedding` | `wordEmbeddingT` | `T` | same as logical dtype | `2` | — | Non-empty word embedding table `(vocab, hidden_size)`. Every `input_ids` value must be a valid non-negative row index. | required |
-| `position_embedding` | `positionEmbeddingT` | `T` | same as logical dtype | `2` | — | Non-empty position embedding table `(max_positions, hidden_size)`. Without `position_ids`, it must contain at least `sequence_length` rows. | required |
-| `segment_embedding` | `segmentEmbeddingT` | `T` | same as logical dtype | `2` | — | Non-empty segment embedding table `(segments, hidden_size)`. If `segment_ids` is absent, row 0 is used for every token. | optional |
-| `gamma` | `gammaT` | `T` | same as logical dtype | `1` | — | Layer-normalization scale of shape `(hidden_size)`. | required |
-| `beta` | `betaT` | `T` | same as logical dtype | `1` | — | Layer-normalization bias of shape `(hidden_size)`. | required |
-| `mask` | `maskT` | `T1` | `int32` | `2` | — | Attention mask of shape `(batch_size, sequence_length)`. Only used to produce `mask_index`. | optional |
-| `position_ids` | `positionIdsT` | `T1` | `int32` | `2` | — | Position ids `(batch_size, sequence_length)`, or `(1, sequence_length)` to share one row across the batch. Values must be valid non-negative table-row indices; absent uses the position within the sequence. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
-| `output` | `outputT` | `T` | same as logical dtype | `3` | derived; see description | Normalized embeddings of shape `(batch_size, sequence_length, hidden_size)`. | required |
-| `mask_index` | `maskIndexT` | `T1` | `int32` | `1` | `[input_ids[0]]` | Position of the first zero in each mask row, or `sequence_length` when no zero exists; shape `(batch_size)`. It is zero when the optional mask input is absent. | optional |
-| `embedding_sum` | `embeddingSumT` | `T` | same as logical dtype | `3` | derived; see description | The summed embeddings before normalization, including the segment term when present. Float16 uses staged `(word + segment) + position`; float32 uses `(word + position) + segment`. | optional |
 ## Attributes
@@ -45,7 +45,7 @@ Attributes and default values (overridable per request):
 | Attribute | Default | Description |
 | --- | --- | --- |
 | `epsilon` | `9.999999960041972e-13` | Non-negative epsilon added to the layer-normalization variance before taking the square root. |
-| `mask_index_type` | — | Optional shape-inference hint for the `mask_index` output type. The schema's `T1` constraint fixes the runtime tensor type to int32. |
 ## Type constraints
@@ -56,7 +56,7 @@ Attributes and 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
@@ -66,10 +66,14 @@ Attributes and 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 |
 | --- | --- | --- | --- | --- | --- | --- | --- |
+| `inputIdsT` | `input_ids` | `T1` | `int32` | `2` | — | Word ids of shape `(batch_size, sequence_length)`. | required |
+| `segmentIdsT` | `segment_ids` | `T1` | `int32` | `2` | — | Segment ids `(batch_size, sequence_length)`. Requires `segment_embedding`; when omitted with that table present, every token uses row 0. Values must be valid non-negative table-row indices. | optional |
+| `wordEmbeddingT` | `word_embedding` | `T` | same as logical dtype | `2` | — | Non-empty word embedding table `(vocab, hidden_size)`. Every `input_ids` value must be a valid non-negative row index. | required |
+| `positionEmbeddingT` | `position_embedding` | `T` | same as logical dtype | `2` | — | Non-empty position embedding table `(max_positions, hidden_size)`. Without `position_ids`, it must contain at least `sequence_length` rows. | required |
+| `segmentEmbeddingT` | `segment_embedding` | `T` | same as logical dtype | `2` | — | Non-empty segment embedding table `(segments, hidden_size)`. If `segment_ids` is absent, row 0 is used for every token. | optional |
+| `gammaT` | `gamma` | `T` | same as logical dtype | `1` | — | Layer-normalization scale of shape `(hidden_size)`. | required |
+| `betaT` | `beta` | `T` | same as logical dtype | `1` | — | Layer-normalization bias of shape `(hidden_size)`. | required |
+| `maskT` | `mask` | `T1` | `int32` | `2` | — | Attention mask of shape `(batch_size, sequence_length)`. Only used to produce `mask_index`. | optional |
+| `positionIdsT` | `position_ids` | `T1` | `int32` | `2` | — | Position ids `(batch_size, sequence_length)`, or `(1, sequence_length)` to share one row across the batch. Values must be valid non-negative table-row indices; absent uses the position within the sequence. | optional |
 ## Outputs
+| Name | Upstream name | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
+| `outputT` | `output` | `T` | same as logical dtype | `3` | derived | Normalized embeddings of shape `(batch_size, sequence_length, hidden_size)`. | required |
+| `maskIndexT` | `mask_index` | `T1` | `int32` | `1` | `[inputIdsT[0]]` | Position of the first zero in each mask row, or `sequence_length` when no zero exists; shape `(batch_size)`. It is zero when the optional mask input is absent. | optional |
+| `embeddingSumT` | `embedding_sum` | `T` | same as logical dtype | `3` | derived | The summed embeddings before normalization, including the segment term when present. Float16 uses staged `(word + segment) + position`; float32 uses `(word + position) + segment`. | optional |
 ## Attributes
 | Attribute | Default | Description |
 | --- | --- | --- |
 | `epsilon` | `9.999999960041972e-13` | Non-negative epsilon added to the layer-normalization variance before taking the square root. |
+| `mask_index_type` | — | Optional shape-inference hint for the `mask_index` output type. The schema's `T1` constraint fixes the output tensor data type to int32. |
 ## Type constraints
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

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

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

build/webgpu/embed-normalize.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.EmbedLayerNormalization, normalization pass.
@@ -79,11 +76,10 @@ fn reduce_sum(value: f32, tid: u32) -> f32 {
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
   // 2D-folded row index: wg.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension dispatch limit.
-  let token = wg.x + wg.y * nwg.x;
   if (token >= params.tokens) {
     return;
   }

 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.EmbedLayerNormalization, normalization pass.
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
   // 2D-folded row index: wg.y carries the high bits past the
+  // per-axis dispatch fold width.
+  let token = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (token >= params.tokens) {
     return;
   }

build/webgpu/embed-sum.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.EmbedLayerNormalization, embedding-sum pass.
@@ -20,8 +17,8 @@ const HIDDEN: u32 = {{ hidden }}u;
 const SEQUENCE: u32 = {{ sequenceLength }}u;
 {% endif %}
 const WG: u32 = {{ workgroupSize }}u;
-// Defensive row clamps keep an invalid id from reading outside an embedding
-// table. Request validation rejects such ids before a valid dispatch.
 const WORD_ROWS: u32 = {{ wordRows }}u;
 const POSITION_ROWS: u32 = {{ positionRows }}u;
 {% if hasSegment and hasSegmentIds %}
@@ -30,11 +27,10 @@ const SEGMENT_ROWS: u32 = {{ segmentRows }}u;
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
   // 2D-folded row index: wg.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension dispatch limit.
-  let token = wg.x + wg.y * nwg.x;
   if (token >= params.tokens) {
     return;
   }

 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.EmbedLayerNormalization, embedding-sum pass.
 const SEQUENCE: u32 = {{ sequenceLength }}u;
 {% endif %}
 const WG: u32 = {{ workgroupSize }}u;
+// IDs must index valid embedding rows. Defensive clamps prevent out-of-bounds
+// reads if that precondition is violated.
 const WORD_ROWS: u32 = {{ wordRows }}u;
 const POSITION_ROWS: u32 = {{ positionRows }}u;
 {% if hasSegment and hasSegmentIds %}
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(workgroup_id) wg: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
   // 2D-folded row index: wg.y carries the high bits past the
+  // per-axis dispatch fold width.
+  let token = wg.x + wg.y * {{ DISPATCH_FOLD_WIDTH }}u;
   if (token >= params.tokens) {
     return;
   }

build/webgpu/manifest.json CHANGED Viewed

@@ -2,110 +2,44 @@
   "domain": "com.microsoft",
   "name": "EmbedLayerNormalization",
   "sinceVersion": 1,
-  "description": "BERT embedding fusion: looks up word and position tables, optionally adds a segment table, then applies layer normalization. A segment table without IDs uses row 0. `embedding_sum` is the pre-normalization sum. `mask_index` is the first zero or the sequence length; without `mask`, it is zero. Batch and sequence dimensions must be non-empty.",
-  "inputs": [
-    {
-      "role": "input_ids",
-      "dtype": "T1",
-      "rank": 2,
-      "description": "Word ids of shape `(batch_size, sequence_length)`."
-    },
-    {
-      "role": "segment_ids",
-      "dtype": "T1",
-      "rank": 2,
-      "optional": true,
-      "description": "Segment ids `(batch_size, sequence_length)`. Requires `segment_embedding`; when omitted with that table present, every token uses row 0. Values must be valid non-negative table-row indices."
-    },
-    {
-      "role": "word_embedding",
-      "dtype": "T",
-      "rank": 2,
-      "description": "Non-empty word embedding table `(vocab, hidden_size)`. Every `input_ids` value must be a valid non-negative row index."
-    },
-    {
-      "role": "position_embedding",
-      "dtype": "T",
-      "rank": 2,
-      "description": "Non-empty position embedding table `(max_positions, hidden_size)`. Without `position_ids`, it must contain at least `sequence_length` rows."
-    },
-    {
-      "role": "segment_embedding",
-      "dtype": "T",
-      "rank": 2,
-      "optional": true,
-      "description": "Non-empty segment embedding table `(segments, hidden_size)`. If `segment_ids` is absent, row 0 is used for every token."
-    },
-    { "role": "gamma", "dtype": "T", "rank": 1, "description": "Layer-normalization scale of shape `(hidden_size)`." },
-    { "role": "beta", "dtype": "T", "rank": 1, "description": "Layer-normalization bias of shape `(hidden_size)`." },
-    {
-      "role": "mask",
-      "dtype": "T1",
-      "rank": 2,
-      "optional": true,
-      "description": "Attention mask of shape `(batch_size, sequence_length)`. Only used to produce `mask_index`."
-    },
-    {
-      "role": "position_ids",
-      "dtype": "T1",
-      "rank": 2,
-      "optional": true,
-      "description": "Position ids `(batch_size, sequence_length)`, or `(1, sequence_length)` to share one row across the batch. Values must be valid non-negative table-row indices; absent uses the position within the sequence."
-    }
-  ],
-  "outputs": [
-    {
-      "role": "output",
       "dtype": "T",
       "rank": 3,
-      "shape": "[dim(shapes.inputIdsT, 0), dim(shapes.inputIdsT, 1), hidden]",
-      "description": "Normalized embeddings of shape `(batch_size, sequence_length, hidden_size)`."
     },
-    {
-      "role": "mask_index",
       "dtype": "T1",
       "rank": 1,
       "optional": true,
       "shape": "[dim(shapes.inputIdsT, 0)]",
-      "description": "Position of the first zero in each mask row, or `sequence_length` when no zero exists; shape `(batch_size)`. It is zero when the optional mask input is absent."
     },
-    {
-      "role": "embedding_sum",
       "dtype": "T",
       "rank": 3,
       "optional": true,
-      "shape": "[dim(shapes.inputIdsT, 0), dim(shapes.inputIdsT, 1), hidden]",
-      "description": "The summed embeddings before normalization, including the segment term when present. Float16 uses staged `(word + segment) + position`; float32 uses `(word + position) + segment`."
     }
-  ],
-  "attributes": { "epsilon": 9.999999960041972e-13 },
-  "attributeDescriptions": {
-    "epsilon": "Non-negative epsilon added to the layer-normalization variance before taking the square root.",
-    "mask_index_type": "Optional shape-inference hint for the `mask_index` output type. The schema's `T1` constraint fixes the runtime tensor type to int32."
   },
   "attributeConstraints": { "mask_index_type": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32", "float16"], "T1": ["int32"] },
-  "args": {
-    "inputIdsT": { "kind": "tensor", "semantic": "input_ids", "role": "input", "dtype": "int32" },
-    "segmentIdsT": { "kind": "tensor", "semantic": "segment_ids", "role": "input", "dtype": "int32", "required": false },
-    "wordEmbeddingT": { "kind": "tensor", "semantic": "word_embedding", "role": "weights" },
-    "positionEmbeddingT": { "kind": "tensor", "semantic": "position_embedding", "role": "weights" },
-    "segmentEmbeddingT": { "kind": "tensor", "semantic": "segment_embedding", "role": "weights", "required": false },
-    "gammaT": { "kind": "tensor", "semantic": "gamma", "role": "weights" },
-    "betaT": { "kind": "tensor", "semantic": "beta", "role": "weights" },
-    "maskT": { "kind": "tensor", "semantic": "mask", "role": "input", "dtype": "int32", "required": false },
-    "positionIdsT": {
-      "kind": "tensor",
-      "semantic": "position_ids",
-      "role": "input",
-      "dtype": "int32",
-      "required": false
-    },
-    "outputT": { "kind": "tensor", "semantic": "output", "role": "output" },
-    "maskIndexT": { "kind": "tensor", "semantic": "mask_index", "role": "output", "dtype": "int32", "required": false },
-    "embeddingSumT": { "kind": "tensor", "semantic": "embedding_sum", "role": "output", "required": false }
-  },
-  "tunables": { "WORKGROUP_SIZE": 128, "MASK_WORKGROUP_SIZE": 64 },
   "derive": {
     "batchSize": "dim(shapes.inputIdsT, 0)",
     "sequenceLength": "dim(shapes.inputIdsT, 1)",
@@ -123,14 +57,9 @@
     "embeddingSumContract": "ranks.embeddingSumT == 3 and sameShape(shapes.embeddingSumT, shapes.outputT) and tensorDtypes.embeddingSumT == tensorDtypes.wordEmbeddingT if present.embeddingSumT else true",
     "maskIndexShapeOk": "ranks.maskIndexT == 1 and dim(shapes.maskIndexT, 0) == batchSize if present.maskIndexT else true",
     "embedContractOk": "epsilonOk and tableShapeOk and segmentContract and positionIdsContract and maskContract and maskIndexTypeOk and ioShapeOk and embeddingSumContract and maskIndexShapeOk and batchSize > 0 and sequenceLength > 0",
-    "dispatchFits": "tunables.WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.MASK_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup"
-  },
-  "constants": {
     "aScalar": "dtypes.T",
     "scalar": "dtypes.T",
-    "usesF16": "dtypes.T == \"f16\"",
-    "hidden": "hidden",
-    "sequenceLength": "sequenceLength",
     "epsilon": "epsilonValue",
     "workgroupSize": "tunables.WORKGROUP_SIZE",
     "maskWorkgroupSize": "tunables.MASK_WORKGROUP_SIZE",
@@ -140,1567 +69,1013 @@
     "hasSegment": "present.segmentEmbeddingT",
     "hasSegmentIds": "present.segmentIdsT",
     "hasPositionIds": "present.positionIdsT",
-    "broadcastPositionIds": "broadcastPositionIds",
     "writeEmbeddingSum": "present.embeddingSumT",
     "hasMask": "present.maskT",
     "HIDDEN_LEN": "hidden"
   },
-  "bindingSets": {
-    "embed_noseg_nopos_nosum": [
-      {
-        "name": "input_ids",
-        "arg": "inputIdsT",
-        "semantic": "input_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "word_embedding",
-        "arg": "wordEmbeddingT",
-        "semantic": "word_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "position_embedding",
-        "arg": "positionEmbeddingT",
-        "semantic": "position_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "tokens", "type": "u32", "value": "tokens" }] }
-      }
-    ],
-    "embed_noseg_nopos_sum": [
-      {
-        "name": "input_ids",
-        "arg": "inputIdsT",
-        "semantic": "input_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "word_embedding",
-        "arg": "wordEmbeddingT",
-        "semantic": "word_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "position_embedding",
-        "arg": "positionEmbeddingT",
-        "semantic": "position_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "embedding_sum",
-        "arg": "embeddingSumT",
-        "semantic": "embedding_sum",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "tokens", "type": "u32", "value": "tokens" }] }
-      }
-    ],
-    "embed_noseg_posids_nosum": [
-      {
-        "name": "input_ids",
-        "arg": "inputIdsT",
-        "semantic": "input_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "position_ids",
-        "arg": "positionIdsT",
-        "semantic": "position_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "word_embedding",
-        "arg": "wordEmbeddingT",
-        "semantic": "word_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "position_embedding",
-        "arg": "positionEmbeddingT",
-        "semantic": "position_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "tokens", "type": "u32", "value": "tokens" }] }
-      }
-    ],
-    "embed_noseg_posids_sum": [
-      {
-        "name": "input_ids",
-        "arg": "inputIdsT",
-        "semantic": "input_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "position_ids",
-        "arg": "positionIdsT",
-        "semantic": "position_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "word_embedding",
-        "arg": "wordEmbeddingT",
-        "semantic": "word_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "position_embedding",
-        "arg": "positionEmbeddingT",
-        "semantic": "position_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "embedding_sum",
-        "arg": "embeddingSumT",
-        "semantic": "embedding_sum",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "tokens", "type": "u32", "value": "tokens" }] }
-      }
-    ],
-    "embed_segdefault_nopos_nosum": [
-      {
-        "name": "input_ids",
-        "arg": "inputIdsT",
-        "semantic": "input_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "word_embedding",
-        "arg": "wordEmbeddingT",
-        "semantic": "word_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "position_embedding",
-        "arg": "positionEmbeddingT",
-        "semantic": "position_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "segment_embedding",
-        "arg": "segmentEmbeddingT",
-        "semantic": "segment_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "tokens", "type": "u32", "value": "tokens" }] }
-      }
-    ],
-    "embed_segdefault_nopos_sum": [
-      {
-        "name": "input_ids",
-        "arg": "inputIdsT",
-        "semantic": "input_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "word_embedding",
-        "arg": "wordEmbeddingT",
-        "semantic": "word_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "position_embedding",
-        "arg": "positionEmbeddingT",
-        "semantic": "position_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "segment_embedding",
-        "arg": "segmentEmbeddingT",
-        "semantic": "segment_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "embedding_sum",
-        "arg": "embeddingSumT",
-        "semantic": "embedding_sum",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "tokens", "type": "u32", "value": "tokens" }] }
-      }
-    ],
-    "embed_segdefault_posids_nosum": [
-      {
-        "name": "input_ids",
-        "arg": "inputIdsT",
-        "semantic": "input_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "position_ids",
-        "arg": "positionIdsT",
-        "semantic": "position_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "word_embedding",
-        "arg": "wordEmbeddingT",
-        "semantic": "word_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "position_embedding",
-        "arg": "positionEmbeddingT",
-        "semantic": "position_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "segment_embedding",
-        "arg": "segmentEmbeddingT",
-        "semantic": "segment_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "tokens", "type": "u32", "value": "tokens" }] }
-      }
-    ],
-    "embed_segdefault_posids_sum": [
-      {
-        "name": "input_ids",
-        "arg": "inputIdsT",
-        "semantic": "input_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "position_ids",
-        "arg": "positionIdsT",
-        "semantic": "position_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "word_embedding",
-        "arg": "wordEmbeddingT",
-        "semantic": "word_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "position_embedding",
-        "arg": "positionEmbeddingT",
-        "semantic": "position_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "segment_embedding",
-        "arg": "segmentEmbeddingT",
-        "semantic": "segment_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "embedding_sum",
-        "arg": "embeddingSumT",
-        "semantic": "embedding_sum",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "tokens", "type": "u32", "value": "tokens" }] }
-      }
-    ],
-    "embed_seg_nopos_nosum": [
-      {
-        "name": "input_ids",
-        "arg": "inputIdsT",
-        "semantic": "input_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "segment_ids",
-        "arg": "segmentIdsT",
-        "semantic": "segment_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "word_embedding",
-        "arg": "wordEmbeddingT",
-        "semantic": "word_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "position_embedding",
-        "arg": "positionEmbeddingT",
-        "semantic": "position_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "segment_embedding",
-        "arg": "segmentEmbeddingT",
-        "semantic": "segment_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "tokens", "type": "u32", "value": "tokens" }] }
-      }
-    ],
-    "embed_seg_nopos_sum": [
-      {
-        "name": "input_ids",
-        "arg": "inputIdsT",
-        "semantic": "input_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "segment_ids",
-        "arg": "segmentIdsT",
-        "semantic": "segment_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "word_embedding",
-        "arg": "wordEmbeddingT",
-        "semantic": "word_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "position_embedding",
-        "arg": "positionEmbeddingT",
-        "semantic": "position_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "segment_embedding",
-        "arg": "segmentEmbeddingT",
-        "semantic": "segment_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "embedding_sum",
-        "arg": "embeddingSumT",
-        "semantic": "embedding_sum",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "tokens", "type": "u32", "value": "tokens" }] }
-      }
-    ],
-    "embed_seg_posids_nosum": [
-      {
-        "name": "input_ids",
-        "arg": "inputIdsT",
-        "semantic": "input_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "segment_ids",
-        "arg": "segmentIdsT",
-        "semantic": "segment_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "position_ids",
-        "arg": "positionIdsT",
-        "semantic": "position_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "word_embedding",
-        "arg": "wordEmbeddingT",
-        "semantic": "word_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "position_embedding",
-        "arg": "positionEmbeddingT",
-        "semantic": "position_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "segment_embedding",
-        "arg": "segmentEmbeddingT",
-        "semantic": "segment_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "tokens", "type": "u32", "value": "tokens" }] }
-      }
-    ],
-    "embed_seg_posids_sum": [
-      {
-        "name": "input_ids",
-        "arg": "inputIdsT",
-        "semantic": "input_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "segment_ids",
-        "arg": "segmentIdsT",
-        "semantic": "segment_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "position_ids",
-        "arg": "positionIdsT",
-        "semantic": "position_ids",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "word_embedding",
-        "arg": "wordEmbeddingT",
-        "semantic": "word_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "position_embedding",
-        "arg": "positionEmbeddingT",
-        "semantic": "position_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "segment_embedding",
-        "arg": "segmentEmbeddingT",
-        "semantic": "segment_embedding",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "embedding_sum",
-        "arg": "embeddingSumT",
-        "semantic": "embedding_sum",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "tokens", "type": "u32", "value": "tokens" }] }
-      }
-    ],
-    "normalizeSet": [
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$aScalar"
-      },
-      {
-        "name": "gamma",
-        "arg": "gammaT",
-        "semantic": "gamma",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "beta",
-        "arg": "betaT",
-        "semantic": "beta",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$aScalar",
-        "length": "$HIDDEN_LEN"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "tokens", "type": "u32", "value": "tokens" }] }
-      }
-    ],
-    "maskIndexSet": [
-      {
-        "name": "mask",
-        "arg": "maskT",
-        "semantic": "mask",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "mask_index",
-        "arg": "maskIndexT",
-        "semantic": "mask_index",
-        "buffer": { "type": "storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "batch", "type": "u32", "value": "batchSize" }] }
-      }
-    ],
-    "zeroMaskIndexSet": [
-      {
-        "name": "mask_index",
-        "arg": "maskIndexT",
-        "semantic": "mask_index",
-        "buffer": { "type": "storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": { "name": "Params", "fields": [{ "name": "batch", "type": "u32", "value": "batchSize" }] }
-      }
-    ]
   },
   "variants": [
     {
       "id": "noseg_nopos_nosum_nomask",
-      "when": ["embedContractOk", "dispatchFits", "not present.segmentEmbeddingT", "not present.positionIdsT", "not present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_noseg_nopos_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         }
       ]
     },
     {
       "id": "noseg_nopos_nosum_mask",
-      "when": ["embedContractOk", "dispatchFits", "not present.segmentEmbeddingT", "not present.positionIdsT", "not present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_noseg_nopos_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "maskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "noseg_nopos_sum_nomask",
-      "when": ["embedContractOk", "dispatchFits", "not present.segmentEmbeddingT", "not present.positionIdsT", "present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_noseg_nopos_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         }
       ]
     },
     {
       "id": "noseg_nopos_sum_mask",
-      "when": ["embedContractOk", "dispatchFits", "not present.segmentEmbeddingT", "not present.positionIdsT", "present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_noseg_nopos_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "maskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "noseg_posids_nosum_nomask",
-      "when": ["embedContractOk", "dispatchFits", "not present.segmentEmbeddingT", "present.positionIdsT", "not present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_noseg_posids_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         }
       ]
     },
     {
       "id": "noseg_posids_nosum_mask",
-      "when": ["embedContractOk", "dispatchFits", "not present.segmentEmbeddingT", "present.positionIdsT", "not present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_noseg_posids_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "maskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "noseg_posids_sum_nomask",
-      "when": ["embedContractOk", "dispatchFits", "not present.segmentEmbeddingT", "present.positionIdsT", "present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_noseg_posids_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         }
       ]
     },
     {
       "id": "noseg_posids_sum_mask",
-      "when": ["embedContractOk", "dispatchFits", "not present.segmentEmbeddingT", "present.positionIdsT", "present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_noseg_posids_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "maskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "seg_nopos_nosum_nomask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "present.segmentIdsT", "not present.positionIdsT", "not present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_seg_nopos_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         }
       ]
     },
     {
       "id": "seg_nopos_nosum_mask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "present.segmentIdsT", "not present.positionIdsT", "not present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_seg_nopos_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "maskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "seg_nopos_sum_nomask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "present.segmentIdsT", "not present.positionIdsT", "present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_seg_nopos_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         }
       ]
     },
     {
       "id": "seg_nopos_sum_mask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "present.segmentIdsT", "not present.positionIdsT", "present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_seg_nopos_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "maskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "seg_posids_nosum_nomask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "present.segmentIdsT", "present.positionIdsT", "not present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_seg_posids_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         }
       ]
     },
     {
       "id": "seg_posids_nosum_mask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "present.segmentIdsT", "present.positionIdsT", "not present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_seg_posids_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "maskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "seg_posids_sum_nomask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "present.segmentIdsT", "present.positionIdsT", "present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_seg_posids_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         }
       ]
     },
     {
       "id": "seg_posids_sum_mask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "present.segmentIdsT", "present.positionIdsT", "present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_seg_posids_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "maskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "segdefault_nopos_nosum_nomask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_segdefault_nopos_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         }
       ]
     },
     {
       "id": "segdefault_nopos_nosum_mask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_segdefault_nopos_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "maskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "segdefault_nopos_sum_nomask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_segdefault_nopos_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         }
       ]
     },
     {
       "id": "segdefault_nopos_sum_mask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_segdefault_nopos_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "maskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "segdefault_posids_nosum_nomask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_segdefault_posids_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         }
       ]
     },
     {
       "id": "segdefault_posids_nosum_mask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_segdefault_posids_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "maskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "segdefault_posids_sum_nomask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_segdefault_posids_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         }
       ]
     },
     {
       "id": "segdefault_posids_sum_mask",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_segdefault_posids_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "maskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "noseg_nopos_nosum_mask_without_input",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT == (\"noseg\" != \"noseg\")", "present.segmentIdsT == (\"noseg\" == \"seg\")", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_noseg_nopos_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "zeroMaskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "noseg_nopos_sum_mask_without_input",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT == (\"noseg\" != \"noseg\")", "present.segmentIdsT == (\"noseg\" == \"seg\")", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_noseg_nopos_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "zeroMaskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "noseg_posids_nosum_mask_without_input",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT == (\"noseg\" != \"noseg\")", "present.segmentIdsT == (\"noseg\" == \"seg\")", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_noseg_posids_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "zeroMaskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "noseg_posids_sum_mask_without_input",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT == (\"noseg\" != \"noseg\")", "present.segmentIdsT == (\"noseg\" == \"seg\")", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_noseg_posids_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "zeroMaskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "segdefault_nopos_nosum_mask_without_input",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT == (\"segdefault\" != \"noseg\")", "present.segmentIdsT == (\"segdefault\" == \"seg\")", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_segdefault_nopos_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "zeroMaskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "segdefault_nopos_sum_mask_without_input",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT == (\"segdefault\" != \"noseg\")", "present.segmentIdsT == (\"segdefault\" == \"seg\")", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_segdefault_nopos_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "zeroMaskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "segdefault_posids_nosum_mask_without_input",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT == (\"segdefault\" != \"noseg\")", "present.segmentIdsT == (\"segdefault\" == \"seg\")", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_segdefault_posids_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "zeroMaskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "segdefault_posids_sum_mask_without_input",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT == (\"segdefault\" != \"noseg\")", "present.segmentIdsT == (\"segdefault\" == \"seg\")", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_segdefault_posids_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "zeroMaskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "seg_nopos_nosum_mask_without_input",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT == (\"seg\" != \"noseg\")", "present.segmentIdsT == (\"seg\" == \"seg\")", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_seg_nopos_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "zeroMaskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "seg_nopos_sum_mask_without_input",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT == (\"seg\" != \"noseg\")", "present.segmentIdsT == (\"seg\" == \"seg\")", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_seg_nopos_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "zeroMaskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "seg_posids_nosum_mask_without_input",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT == (\"seg\" != \"noseg\")", "present.segmentIdsT == (\"seg\" == \"seg\")", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_seg_posids_nosum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "zeroMaskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     },
     {
       "id": "seg_posids_sum_mask_without_input",
-      "when": ["embedContractOk", "dispatchFits", "present.segmentEmbeddingT == (\"seg\" != \"noseg\")", "present.segmentIdsT == (\"seg\" == \"seg\")", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
-          "bindings": "embed_seg_posids_sum",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
-          "bindings": "normalizeSet",
-          "dispatch": { "workgroups": "tokens" }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
-          "bindings": "zeroMaskIndexSet",
-          "dispatch": { "threads": "batchSize", "workgroupSize": "tunables.MASK_WORKGROUP_SIZE" }
         }
       ]
     }

   "domain": "com.microsoft",
   "name": "EmbedLayerNormalization",
   "sinceVersion": 1,
+  "inputs": {
+    "inputIdsT": { "onnx": "input_ids", "dtype": "T1", "rank": 2, "storage": "int32" },
+    "segmentIdsT": { "onnx": "segment_ids", "dtype": "T1", "rank": 2, "optional": true, "storage": "int32" },
+    "wordEmbeddingT": { "onnx": "word_embedding", "dtype": "T", "rank": 2 },
+    "positionEmbeddingT": { "onnx": "position_embedding", "dtype": "T", "rank": 2 },
+    "segmentEmbeddingT": { "onnx": "segment_embedding", "dtype": "T", "rank": 2, "optional": true },
+    "gammaT": { "onnx": "gamma", "dtype": "T", "rank": 1 },
+    "betaT": { "onnx": "beta", "dtype": "T", "rank": 1 },
+    "maskT": { "onnx": "mask", "dtype": "T1", "rank": 2, "optional": true, "storage": "int32" },
+    "positionIdsT": { "onnx": "position_ids", "dtype": "T1", "rank": 2, "optional": true, "storage": "int32" }
+  },
+  "outputs": {
+    "outputT": {
+      "onnx": "output",
       "dtype": "T",
       "rank": 3,
+      "shape": "[dim(shapes.inputIdsT, 0), dim(shapes.inputIdsT, 1), hidden]"
     },
+    "maskIndexT": {
+      "onnx": "mask_index",
       "dtype": "T1",
       "rank": 1,
       "optional": true,
       "shape": "[dim(shapes.inputIdsT, 0)]",
+      "storage": "int32"
     },
+    "embeddingSumT": {
+      "onnx": "embedding_sum",
       "dtype": "T",
       "rank": 3,
       "optional": true,
+      "shape": "[dim(shapes.inputIdsT, 0), dim(shapes.inputIdsT, 1), hidden]"
     }
   },
+  "attributes": { "epsilon": { "default": 9.999999960041972e-13 }, "mask_index_type": {} },
   "attributeConstraints": { "mask_index_type": { "values": [0, 1] } },
   "typeConstraints": { "T": ["float32", "float16"], "T1": ["int32"] },
+  "tunables": { "WORKGROUP_SIZE": { "default": 128 }, "MASK_WORKGROUP_SIZE": { "default": 64 } },
   "derive": {
     "batchSize": "dim(shapes.inputIdsT, 0)",
     "sequenceLength": "dim(shapes.inputIdsT, 1)",
     "embeddingSumContract": "ranks.embeddingSumT == 3 and sameShape(shapes.embeddingSumT, shapes.outputT) and tensorDtypes.embeddingSumT == tensorDtypes.wordEmbeddingT if present.embeddingSumT else true",
     "maskIndexShapeOk": "ranks.maskIndexT == 1 and dim(shapes.maskIndexT, 0) == batchSize if present.maskIndexT else true",
     "embedContractOk": "epsilonOk and tableShapeOk and segmentContract and positionIdsContract and maskContract and maskIndexTypeOk and ioShapeOk and embeddingSumContract and maskIndexShapeOk and batchSize > 0 and sequenceLength > 0",
+    "dispatchFits": "tunables.WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup and tunables.MASK_WORKGROUP_SIZE <= device.limits.maxComputeInvocationsPerWorkgroup",
     "aScalar": "dtypes.T",
     "scalar": "dtypes.T",
     "epsilon": "epsilonValue",
     "workgroupSize": "tunables.WORKGROUP_SIZE",
     "maskWorkgroupSize": "tunables.MASK_WORKGROUP_SIZE",
     "hasSegment": "present.segmentEmbeddingT",
     "hasSegmentIds": "present.segmentIdsT",
     "hasPositionIds": "present.positionIdsT",
     "writeEmbeddingSum": "present.embeddingSumT",
     "hasMask": "present.maskT",
     "HIDDEN_LEN": "hidden"
   },
+  "when": ["embedContractOk", "dispatchFits"],
+  "bindings": {
+    "input_ids": { "arg": "inputIdsT", "buffer": "read-only-storage", "elementType": "i32" },
+    "word_embedding": { "arg": "wordEmbeddingT", "buffer": "read-only-storage", "elementType": "$aScalar" },
+    "position_embedding": { "arg": "positionEmbeddingT", "buffer": "read-only-storage", "elementType": "$aScalar" },
+    "output": { "arg": "outputT", "buffer": "storage", "elementType": "$aScalar" },
+    "params": { "buffer": "uniform", "struct": [{ "name": "tokens", "type": "u32", "value": "tokens" }] },
+    "embedding_sum": { "arg": "embeddingSumT", "buffer": "storage", "elementType": "$aScalar" },
+    "position_ids": { "arg": "positionIdsT", "buffer": "read-only-storage", "elementType": "i32" },
+    "segment_ids": { "arg": "segmentIdsT", "buffer": "read-only-storage", "elementType": "i32" },
+    "segment_embedding": { "arg": "segmentEmbeddingT", "buffer": "read-only-storage", "elementType": "$aScalar" },
+    "gamma": { "arg": "gammaT", "buffer": "read-only-storage", "elementType": "$aScalar", "length": "$HIDDEN_LEN" },
+    "beta": { "arg": "betaT", "buffer": "read-only-storage", "elementType": "$aScalar", "length": "$HIDDEN_LEN" },
+    "mask": { "arg": "maskT", "buffer": "read-only-storage", "elementType": "i32" },
+    "mask_index": { "arg": "maskIndexT", "buffer": "storage", "elementType": "i32" },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [{ "name": "batch", "type": "u32", "value": "batchSize" }]
+    }
   },
   "variants": [
     {
       "id": "noseg_nopos_nosum_nomask",
+      "when": ["not present.segmentEmbeddingT", "not present.positionIdsT", "not present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "word_embedding", "position_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "noseg_nopos_nosum_mask",
+      "when": ["not present.segmentEmbeddingT", "not present.positionIdsT", "not present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "word_embedding", "position_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask", "mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "noseg_nopos_sum_nomask",
+      "when": ["not present.segmentEmbeddingT", "not present.positionIdsT", "present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "word_embedding", "position_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "noseg_nopos_sum_mask",
+      "when": ["not present.segmentEmbeddingT", "not present.positionIdsT", "present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "word_embedding", "position_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask", "mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "noseg_posids_nosum_nomask",
+      "when": ["not present.segmentEmbeddingT", "present.positionIdsT", "not present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "position_ids", "word_embedding", "position_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "noseg_posids_nosum_mask",
+      "when": ["not present.segmentEmbeddingT", "present.positionIdsT", "not present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "position_ids", "word_embedding", "position_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask", "mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "noseg_posids_sum_nomask",
+      "when": ["not present.segmentEmbeddingT", "present.positionIdsT", "present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "position_ids", "word_embedding", "position_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "noseg_posids_sum_mask",
+      "when": ["not present.segmentEmbeddingT", "present.positionIdsT", "present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "position_ids", "word_embedding", "position_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask", "mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "seg_nopos_nosum_nomask",
+      "when": ["present.segmentEmbeddingT", "present.segmentIdsT", "not present.positionIdsT", "not present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "segment_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "seg_nopos_nosum_mask",
+      "when": ["present.segmentEmbeddingT", "present.segmentIdsT", "not present.positionIdsT", "not present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "segment_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask", "mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "seg_nopos_sum_nomask",
+      "when": ["present.segmentEmbeddingT", "present.segmentIdsT", "not present.positionIdsT", "present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "segment_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "seg_nopos_sum_mask",
+      "when": ["present.segmentEmbeddingT", "present.segmentIdsT", "not present.positionIdsT", "present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "segment_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask", "mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "seg_posids_nosum_nomask",
+      "when": ["present.segmentEmbeddingT", "present.segmentIdsT", "present.positionIdsT", "not present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "segment_ids", "position_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "seg_posids_nosum_mask",
+      "when": ["present.segmentEmbeddingT", "present.segmentIdsT", "present.positionIdsT", "not present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "segment_ids", "position_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask", "mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "seg_posids_sum_nomask",
+      "when": ["present.segmentEmbeddingT", "present.segmentIdsT", "present.positionIdsT", "present.embeddingSumT", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "segment_ids", "position_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "seg_posids_sum_mask",
+      "when": ["present.segmentEmbeddingT", "present.segmentIdsT", "present.positionIdsT", "present.embeddingSumT", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "segment_ids", "position_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask", "mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "segdefault_nopos_nosum_nomask",
+      "when": ["present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "segdefault_nopos_nosum_mask",
+      "when": ["present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask", "mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "segdefault_nopos_sum_nomask",
+      "when": ["present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "segdefault_nopos_sum_mask",
+      "when": ["present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask", "mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "segdefault_posids_nosum_nomask",
+      "when": ["present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "position_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "segdefault_posids_nosum_mask",
+      "when": ["present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "position_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask", "mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "segdefault_posids_sum_nomask",
+      "when": ["present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "not present.maskIndexT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "position_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         }
       ]
     },
     {
       "id": "segdefault_posids_sum_mask",
+      "when": ["present.segmentEmbeddingT", "not present.segmentIdsT", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "position_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.MaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask", "mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "noseg_nopos_nosum_mask_without_input",
+      "when": ["present.segmentEmbeddingT == (\"noseg\" != \"noseg\")", "present.segmentIdsT == (\"noseg\" == \"seg\")", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "word_embedding", "position_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "noseg_nopos_sum_mask_without_input",
+      "when": ["present.segmentEmbeddingT == (\"noseg\" != \"noseg\")", "present.segmentIdsT == (\"noseg\" == \"seg\")", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "word_embedding", "position_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "noseg_posids_nosum_mask_without_input",
+      "when": ["present.segmentEmbeddingT == (\"noseg\" != \"noseg\")", "present.segmentIdsT == (\"noseg\" == \"seg\")", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "position_ids", "word_embedding", "position_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "noseg_posids_sum_mask_without_input",
+      "when": ["present.segmentEmbeddingT == (\"noseg\" != \"noseg\")", "present.segmentIdsT == (\"noseg\" == \"seg\")", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "position_ids", "word_embedding", "position_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "segdefault_nopos_nosum_mask_without_input",
+      "when": ["present.segmentEmbeddingT == (\"segdefault\" != \"noseg\")", "present.segmentIdsT == (\"segdefault\" == \"seg\")", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "segdefault_nopos_sum_mask_without_input",
+      "when": ["present.segmentEmbeddingT == (\"segdefault\" != \"noseg\")", "present.segmentIdsT == (\"segdefault\" == \"seg\")", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "segdefault_posids_nosum_mask_without_input",
+      "when": ["present.segmentEmbeddingT == (\"segdefault\" != \"noseg\")", "present.segmentIdsT == (\"segdefault\" == \"seg\")", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "position_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "segdefault_posids_sum_mask_without_input",
+      "when": ["present.segmentEmbeddingT == (\"segdefault\" != \"noseg\")", "present.segmentIdsT == (\"segdefault\" == \"seg\")", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "position_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "seg_nopos_nosum_mask_without_input",
+      "when": ["present.segmentEmbeddingT == (\"seg\" != \"noseg\")", "present.segmentIdsT == (\"seg\" == \"seg\")", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "segment_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "seg_nopos_sum_mask_without_input",
+      "when": ["present.segmentEmbeddingT == (\"seg\" != \"noseg\")", "present.segmentIdsT == (\"seg\" == \"seg\")", "present.positionIdsT == (\"nopos\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "segment_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "seg_posids_nosum_mask_without_input",
+      "when": ["present.segmentEmbeddingT == (\"seg\" != \"noseg\")", "present.segmentIdsT == (\"seg\" == \"seg\")", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"nosum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "segment_ids", "position_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
     {
       "id": "seg_posids_sum_mask_without_input",
+      "when": ["present.segmentEmbeddingT == (\"seg\" != \"noseg\")", "present.segmentIdsT == (\"seg\" == \"seg\")", "present.positionIdsT == (\"posids\" == \"posids\")", "present.embeddingSumT == (\"sum\" == \"sum\")", "present.maskIndexT", "not present.maskT"],
       "passes": [
         {
           "id": "sum",
           "name": "EmbedLayerNormalization.EmbeddingSum",
           "shader": "embed-sum.wgsl.jinja",
+          "bindings": ["input_ids", "segment_ids", "position_ids", "word_embedding", "position_embedding", "segment_embedding", "output", "embedding_sum", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "normalize",
           "name": "EmbedLayerNormalization.Normalize",
           "shader": "embed-normalize.wgsl.jinja",
+          "bindings": ["output", "gamma", "beta", "params"],
+          "dispatch": { "x": "min(tokens, 65535)", "y": "ceilDiv(tokens, 65535)", "z": 1 }
         },
         {
           "id": "maskIndex",
           "name": "EmbedLayerNormalization.ZeroMaskIndex",
           "shader": "embed-mask-index.wgsl.jinja",
+          "bindings": ["mask_index", "params_2"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize), (tunables.MASK_WORKGROUP_SIZE)), 65535)",
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,20 +1,60 @@
 {
   "name": "com.microsoft.EmbedLayerNormalization",
-  "id": "_com_microsoft_embedlayernormalization_webgpu_b297617",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "LyhXY2RrY86niUBQGOouS5wCSVNuVu0RylHfXPXwZRc=",
       "embed-mask-index.wgsl.jinja": "7wT7/LzdrOnanlTU5kmneEYpGb6D8v0H9/9mra8eF5Q=",
-      "embed-normalize.wgsl.jinja": "YObvmdAvDuyx8f4pwTmVxphDSivMh4Y2oovmT7eGm1s=",
-      "embed-sum.wgsl.jinja": "cCsjv9CmNWgQLtMOAvEQOveXEO6jpeX7NMD4PiaaIW0=",
-      "manifest.json": "C/rakYKSOpdBjI3xKiqWRE18Dg002pKoEVcA/yRtprM=",
-      "test.json": "UE2iD8tAA2HxgRcTihjDe6LYfamqEWUopaGGER7vNQc="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.EmbedLayerNormalization" }
 }

 {
   "name": "com.microsoft.EmbedLayerNormalization",
+  "id": "_com_microsoft_embedlayernormalization_webgpu_1487059",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "Hs0qFEnXlwW07Hr/0rT+Bq7eqznmLpk7Zf1zpHE3qvw=",
       "embed-mask-index.wgsl.jinja": "7wT7/LzdrOnanlTU5kmneEYpGb6D8v0H9/9mra8eF5Q=",
+      "embed-normalize.wgsl.jinja": "U2vsL1A7Y88g0DgeHpCON1vrTL7Dzc9Z2+CUmAFtt+0=",
+      "embed-sum.wgsl.jinja": "RJVXdkHRkytYbkJ8wtcuZpB7TG/ubEOZPmIOYpOXMCw=",
+      "manifest.json": "ocqCMw8JDV5GTOF6rXRhscDQ+z2G9NjzZDfxlwAznP8=",
+      "test.json": "D+ZuT3RBcLxlldTtFpYDKm1J2fRt64z2CvSv0vP8wbk="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "noseg_nopos_nosum_nomask": ["embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "noseg_nopos_nosum_mask": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "noseg_nopos_sum_nomask": ["embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "noseg_nopos_sum_mask": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "noseg_posids_nosum_nomask": ["embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "noseg_posids_nosum_mask": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "noseg_posids_sum_nomask": ["embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "noseg_posids_sum_mask": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "seg_nopos_nosum_nomask": ["embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "seg_nopos_nosum_mask": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "seg_nopos_sum_nomask": ["embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "seg_nopos_sum_mask": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "seg_posids_nosum_nomask": ["embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "seg_posids_nosum_mask": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "seg_posids_sum_nomask": ["embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "seg_posids_sum_mask": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "segdefault_nopos_nosum_nomask": ["embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "segdefault_nopos_nosum_mask": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "segdefault_nopos_sum_nomask": ["embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "segdefault_nopos_sum_mask": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "segdefault_posids_nosum_nomask": ["embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "segdefault_posids_nosum_mask": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "segdefault_posids_sum_nomask": ["embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "segdefault_posids_sum_mask": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "noseg_nopos_nosum_mask_without_input": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "noseg_nopos_sum_mask_without_input": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "noseg_posids_nosum_mask_without_input": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "noseg_posids_sum_mask_without_input": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "segdefault_nopos_nosum_mask_without_input": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "segdefault_nopos_sum_mask_without_input": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "segdefault_posids_nosum_mask_without_input": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "segdefault_posids_sum_mask_without_input": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "seg_nopos_nosum_mask_without_input": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "seg_nopos_sum_mask_without_input": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "seg_posids_nosum_mask_without_input": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"],
+      "seg_posids_sum_mask_without_input": ["embed-mask-index.wgsl.jinja", "embed-normalize.wgsl.jinja", "embed-sum.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/test.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.EmbedLayerNormalization",
   "cases": [
     {
       "name": "segment_embedding_without_segment_ids_uses_row_zero",

 {
   "cases": [
     {
       "name": "segment_embedding_without_segment_ids_uses_row_zero",