sync 91d990483a17

Browse files

Files changed (8) hide show

README.md +18 -14
build/webgpu/bench.json +0 -1
build/webgpu/manifest.json +117 -308
build/webgpu/metadata.json +21 -9
build/webgpu/test.json +0 -1
build/webgpu/varlen-causal-conv-stream.wgsl.jinja +3 -7
build/webgpu/varlen-causal-conv.wgsl.jinja +6 -11
build/webgpu/varlen-state-update.wgsl.jinja +2 -6

README.md CHANGED Viewed

@@ -18,22 +18,22 @@ See the [ONNX Runtime `VarlenCausalConvWithState` contrib-operator spec](https:/
 ## Inputs
-| Name | Bind key | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
-| `input` | `inputT` | `T` | same as logical dtype | `2` | — | Token-major packed input with shape `(total_tokens, channels)`. | required |
-| `weight` | `weightT` | `T` | same as logical dtype | `3` | — | Depthwise kernel with shape `(channels, 1, kernel_size)`. | required |
-| `cumulative_sequence_length` | `cumulativeSequenceLengthT` | `M` | `int32` | `1` | — | Exclusive prefix sums with shape `(batch_size + 1)`; sequence `i` owns tokens `[cum[i], cum[i + 1])`. | required |
-| `bias` | `biasT` | `T` | same as logical dtype | `1` | — | Optional per-channel bias with shape `(channels,)`. In an ONNX graph an omitted bias must still occupy input index 3 as an empty name so `initial_state` stays at index 4. | optional |
-| `initial_state` | `initialStateT` | `T` | same as logical dtype | `3` | — | Required committed carry state with shape `(batch_size, channels, kernel_size - 1)`, holding the raw samples immediately preceding this call. | required |
-| `capture_count` | `captureCountT` | `M` | `int32` | `1` | — | Optional int32 vector with shape `(batch_size)`. Required exactly when `state_update_capacity` is positive; each value is clamped to `[0, min(state_update_capacity, sequence_length)]`. | optional |
 ## Outputs
-| Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
-| `output` | `outputT` | `T` | same as `input` | same as `input` | Convolution output with the same shape as `input`. | required |
-| `final_state` | `finalStateT` | `T` | `3` | derived; see description | State after each sequence's final token, shape `(batch_size, channels, kernel_size - 1)`. Always fully written. | required |
-| `state_update` | `stateUpdateT` | `T` | `3` | derived; see description | Optional compact transition values with shape `(batch_size, state_update_capacity, channels)`. Active slots contain the original local input tokens and all other slots are zero. | optional |
 ## Attributes
@@ -53,7 +53,7 @@ Default values (overridable per request):
 ## Files
-- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -63,10 +63,14 @@ Default values (overridable per request):
 ## Use with `@huggingface/kernels`
-The loader derives every required output's shape and logical dtype from the manifest contract and this call.
-It then allocates the result tensors automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

 ## Inputs
+| Name | Upstream name | Logical dtype | WebGPU storage | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- | --- |
+| `inputT` | `input` | `T` | same as logical dtype | `2` | — | Token-major packed input with shape `(total_tokens, channels)`. | required |
+| `weightT` | `weight` | `T` | same as logical dtype | `3` | — | Depthwise kernel with shape `(channels, 1, kernel_size)`. | required |
+| `cumulativeSequenceLengthT` | `cumulative_sequence_length` | `M` | `int32` | `1` | — | Exclusive prefix sums with shape `(batch_size + 1)`, starting at 0, ending at `total_tokens`, and strictly increasing so every sequence is non-empty. Sequence `i` owns tokens `[cum[i], cum[i + 1])`. Outputs are unspecified for a malformed schedule. | required |
+| `biasT` | `bias` | `T` | same as logical dtype | `1` | — | Optional per-channel bias with shape `(channels,)`. In an ONNX graph an omitted bias must still occupy input index 3 as an empty name so `initial_state` stays at index 4. | optional |
+| `initialStateT` | `initial_state` | `T` | same as logical dtype | `3` | — | Required committed carry state with shape `(batch_size, channels, kernel_size - 1)`, holding the raw samples immediately preceding this call. | required |
+| `captureCountT` | `capture_count` | `M` | `int32` | `1` | — | Optional int32 vector with shape `(batch_size)`. Required exactly when `state_update_capacity` is positive; each value is clamped to `[0, min(state_update_capacity, sequence_length)]`. | optional |
 ## Outputs
+| Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
 | --- | --- | --- | --- | --- | --- | --- |
+| `outputT` | `output` | `T` | same as `inputT` | same as `inputT` | Convolution output with the same shape as `input`. | required |
+| `finalStateT` | `final_state` | `T` | `3` | derived | State after each sequence's final token, shape `(batch_size, channels, kernel_size - 1)`. Always fully written. | required |
+| `stateUpdateT` | `state_update` | `T` | `3` | derived | Optional compact transition values with shape `(batch_size, state_update_capacity, channels)`. Active slots contain the original local input tokens and all other slots are zero. | optional |
 ## Attributes
 ## Files
+- [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
 - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 ## Use with `@huggingface/kernels`
+```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.2
+```
+Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
+It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
 Replace each `*Data` placeholder with a typed array containing the corresponding input data.

build/webgpu/bench.json CHANGED Viewed

@@ -1,5 +1,4 @@
 {
-  "op": "com.microsoft.VarlenCausalConvWithState",
   "tunableSpace": {
     "workgroupSize": [64, 128, 256],
     "streamWorkgroupSize": [64, 128, 256],

 {
   "tunableSpace": {
     "workgroupSize": [64, 128, 256],
     "streamWorkgroupSize": [64, 128, 256],

build/webgpu/manifest.json CHANGED Viewed

@@ -2,101 +2,33 @@
   "domain": "com.microsoft",
   "name": "VarlenCausalConvWithState",
   "sinceVersion": 1,
-  "description": "Stateful causal depthwise convolution over packed token-major variable-length sequences, without reads across sequence boundaries. `initial_state` carries preceding raw samples and `final_state` is fully written. At positive `state_update_capacity`, `capture_count` selects a clamped prefix of raw input tokens for compact `state_update`; inactive slots are zero. SiLU and Swish are aliases. This implementation supports float16 and float32 with float32 accumulation; bfloat16 is not implemented.",
-  "inputs": [
-    {
-      "role": "input",
-      "dtype": "T",
-      "rank": 2,
-      "description": "Token-major packed input with shape `(total_tokens, channels)`."
-    },
-    {
-      "role": "weight",
-      "dtype": "T",
-      "rank": 3,
-      "description": "Depthwise kernel with shape `(channels, 1, kernel_size)`."
-    },
-    {
-      "role": "cumulative_sequence_length",
-      "dtype": "M",
-      "rank": 1,
-      "description": "Exclusive prefix sums with shape `(batch_size + 1)`; sequence `i` owns tokens `[cum[i], cum[i + 1])`."
-    },
-    {
-      "role": "bias",
-      "dtype": "T",
-      "rank": 1,
-      "optional": true,
-      "description": "Optional per-channel bias with shape `(channels,)`. In an ONNX graph an omitted bias must still occupy input index 3 as an empty name so `initial_state` stays at index 4."
-    },
-    {
-      "role": "initial_state",
-      "dtype": "T",
-      "rank": 3,
-      "description": "Required committed carry state with shape `(batch_size, channels, kernel_size - 1)`, holding the raw samples immediately preceding this call."
-    },
-    {
-      "role": "capture_count",
-      "dtype": "M",
-      "rank": 1,
-      "optional": true,
-      "description": "Optional int32 vector with shape `(batch_size)`. Required exactly when `state_update_capacity` is positive; each value is clamped to `[0, min(state_update_capacity, sequence_length)]`."
-    }
-  ],
-  "outputs": [
-    {
-      "role": "output",
-      "dtype": "T",
-      "rank": "ranks.inputT",
-      "shape": "shapes.inputT",
-      "description": "Convolution output with the same shape as `input`."
-    },
-    {
-      "role": "final_state",
-      "dtype": "T",
-      "rank": 3,
-      "shape": "[batchSize, channels, stateLength]",
-      "description": "State after each sequence's final token, shape `(batch_size, channels, kernel_size - 1)`. Always fully written."
-    },
-    {
-      "role": "state_update",
       "dtype": "T",
       "rank": 3,
       "optional": true,
-      "shape": "[batchSize, stateUpdateCapacity, channels]",
-      "description": "Optional compact transition values with shape `(batch_size, state_update_capacity, channels)`. Active slots contain the original local input tokens and all other slots are zero."
     }
-  ],
-  "typeConstraints": { "T": ["float32", "float16"], "M": ["int32"] },
-  "attributes": { "activation": "none", "state_update_capacity": 0 },
-  "attributeDescriptions": {
-    "activation": "Fused activation applied after convolution and bias. One of `none`, `silu`, or `swish`; the standard default is `none`.",
-    "state_update_capacity": "Static number of compact per-request prefix transition values to expose, in `[0, 8]`. The standard default is 0."
   },
   "attributeConstraints": { "activation": { "values": ["none", "silu", "swish"] } },
-  "args": {
-    "inputT": { "kind": "tensor", "semantic": "input", "role": "input" },
-    "weightT": { "kind": "tensor", "semantic": "weight", "role": "input" },
-    "cumulativeSequenceLengthT": {
-      "kind": "tensor",
-      "semantic": "cumulative_sequence_length",
-      "role": "input",
-      "dtype": "int32"
-    },
-    "biasT": { "kind": "tensor", "semantic": "bias", "role": "input", "required": false },
-    "initialStateT": { "kind": "tensor", "semantic": "initial_state", "role": "input" },
-    "captureCountT": {
-      "kind": "tensor",
-      "semantic": "capture_count",
-      "role": "input",
-      "dtype": "int32",
-      "required": false
-    },
-    "outputT": { "kind": "tensor", "semantic": "output", "role": "output" },
-    "finalStateT": { "kind": "tensor", "semantic": "final_state", "role": "output" },
-    "stateUpdateT": { "kind": "tensor", "semantic": "state_update", "role": "output", "required": false }
   },
-  "tunables": { "workgroupSize": 256, "streamWorkgroupSize": 64, "streamChunk": 16 },
   "derive": {
     "totalTokens": "dim(shapes.inputT, 0)",
     "channels": "dim(shapes.inputT, 1)",
@@ -117,188 +49,51 @@
     "stateUpdateContract": "commonContract and not present.biasT and present.stateUpdateT and stateUpdateCapacity > 0",
     "biasStateUpdateContract": "commonContract and present.biasT and biasOk and present.stateUpdateT and stateUpdateCapacity > 0"
   },
-  "bindingSets": {
-    "plain": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$inputScalar"
-      },
-      {
-        "name": "weight",
-        "arg": "weightT",
-        "semantic": "weight",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$inputScalar"
-      },
-      {
-        "name": "cumulative_sequence_length",
-        "arg": "cumulativeSequenceLengthT",
-        "semantic": "cumulative_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "initial_state",
-        "arg": "initialStateT",
-        "semantic": "initial_state",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$inputScalar"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$outputScalar"
-      },
-      {
-        "name": "final_state",
-        "arg": "finalStateT",
-        "semantic": "final_state",
-        "buffer": { "type": "storage" },
-        "elementType": "$outputScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "batchSize", "type": "u32", "value": "batchSize" },
-            { "name": "channels", "type": "u32", "value": "channels" },
-            { "name": "totalTokens", "type": "u32", "value": "totalTokens" },
-            { "name": "kernelSize", "type": "u32", "value": "kernelSize" }
-          ]
-        }
-      }
-    ],
-    "bias": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$inputScalar"
-      },
-      {
-        "name": "weight",
-        "arg": "weightT",
-        "semantic": "weight",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$inputScalar"
-      },
-      {
-        "name": "cumulative_sequence_length",
-        "arg": "cumulativeSequenceLengthT",
-        "semantic": "cumulative_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "initial_state",
-        "arg": "initialStateT",
-        "semantic": "initial_state",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$inputScalar"
-      },
-      {
-        "name": "bias",
-        "arg": "biasT",
-        "semantic": "bias",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$inputScalar"
-      },
-      {
-        "name": "output",
-        "arg": "outputT",
-        "semantic": "output",
-        "buffer": { "type": "storage" },
-        "elementType": "$outputScalar"
-      },
-      {
-        "name": "final_state",
-        "arg": "finalStateT",
-        "semantic": "final_state",
-        "buffer": { "type": "storage" },
-        "elementType": "$outputScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "batchSize", "type": "u32", "value": "batchSize" },
-            { "name": "channels", "type": "u32", "value": "channels" },
-            { "name": "totalTokens", "type": "u32", "value": "totalTokens" },
-            { "name": "kernelSize", "type": "u32", "value": "kernelSize" }
-          ]
-        }
-      }
-    ],
-    "state_update": [
-      {
-        "name": "input",
-        "arg": "inputT",
-        "semantic": "input",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "$inputScalar"
-      },
-      {
-        "name": "cumulative_sequence_length",
-        "arg": "cumulativeSequenceLengthT",
-        "semantic": "cumulative_sequence_length",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "capture_count",
-        "arg": "captureCountT",
-        "semantic": "capture_count",
-        "buffer": { "type": "read-only-storage" },
-        "elementType": "i32"
-      },
-      {
-        "name": "state_update",
-        "arg": "stateUpdateT",
-        "semantic": "state_update",
-        "buffer": { "type": "storage" },
-        "elementType": "$outputScalar"
-      },
-      {
-        "name": "params",
-        "semantic": "kernel.state_update_params",
-        "buffer": { "type": "uniform" },
-        "struct": {
-          "name": "Params",
-          "fields": [
-            { "name": "batchSize", "type": "u32", "value": "batchSize" },
-            { "name": "channels", "type": "u32", "value": "channels" },
-            { "name": "totalTokens", "type": "u32", "value": "totalTokens" },
-            { "name": "stateUpdateCapacity", "type": "u32", "value": "stateUpdateCapacity" }
-          ]
-        }
-      }
-    ]
   },
   "variants": [
     {
       "id": "plain_stream",
       "priority": 20,
       "when": ["plainContract", "kernelSize >= 2", "kernelSize <= 8", "tunables.streamChunk >= 1", "floor(tunables.streamChunk) == tunables.streamChunk", "tunables.streamWorkgroupSize >= 1", "floor(tunables.streamWorkgroupSize) == tunables.streamWorkgroupSize", "tunables.streamWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.streamWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
-      "constants": {
         "hasBias": false,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
-        "usesF16": "tensorDtypes.inputT == \"float16\"",
         "workgroupSize": "tunables.streamWorkgroupSize",
-        "kernelSize": "kernelSize",
         "chunkSize": "min(tunables.streamChunk, ceilDiv(totalTokens, batchSize))"
       },
       "passes": [
@@ -306,10 +101,11 @@
           "id": "main",
           "name": "VarlenCausalConvWithState.Stream",
           "shader": "varlen-causal-conv-stream.wgsl.jinja",
-          "bindings": "plain",
           "dispatch": {
-            "threads": "ceilDiv(totalTokens, constants.chunkSize) * channels",
-            "workgroupSize": "constants.workgroupSize"
           }
         }
       ]
@@ -318,14 +114,12 @@
       "id": "bias_stream",
       "priority": 20,
       "when": ["biasContract", "kernelSize >= 2", "kernelSize <= 8", "tunables.streamChunk >= 1", "floor(tunables.streamChunk) == tunables.streamChunk", "tunables.streamWorkgroupSize >= 1", "floor(tunables.streamWorkgroupSize) == tunables.streamWorkgroupSize", "tunables.streamWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.streamWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
-      "constants": {
         "hasBias": true,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
-        "usesF16": "tensorDtypes.inputT == \"float16\"",
         "workgroupSize": "tunables.streamWorkgroupSize",
-        "kernelSize": "kernelSize",
         "chunkSize": "min(tunables.streamChunk, ceilDiv(totalTokens, batchSize))"
       },
       "passes": [
@@ -333,10 +127,11 @@
           "id": "main",
           "name": "VarlenCausalConvWithState.Stream",
           "shader": "varlen-causal-conv-stream.wgsl.jinja",
-          "bindings": "bias",
           "dispatch": {
-            "threads": "ceilDiv(totalTokens, constants.chunkSize) * channels",
-            "workgroupSize": "constants.workgroupSize"
           }
         }
       ]
@@ -345,14 +140,12 @@
       "id": "state_update_stream",
       "priority": 20,
       "when": ["stateUpdateContract", "kernelSize >= 2", "kernelSize <= 8", "tunables.streamChunk >= 1", "floor(tunables.streamChunk) == tunables.streamChunk", "tunables.streamWorkgroupSize >= 1", "floor(tunables.streamWorkgroupSize) == tunables.streamWorkgroupSize", "tunables.streamWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.streamWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
-      "constants": {
         "hasBias": false,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
-        "usesF16": "tensorDtypes.inputT == \"float16\"",
         "workgroupSize": "tunables.streamWorkgroupSize",
-        "kernelSize": "kernelSize",
         "chunkSize": "min(tunables.streamChunk, ceilDiv(totalTokens, batchSize))"
       },
       "passes": [
@@ -360,20 +153,22 @@
           "id": "state_update",
           "name": "VarlenCausalConvWithState.StateUpdate",
           "shader": "varlen-state-update.wgsl.jinja",
-          "bindings": "state_update",
           "dispatch": {
-            "threads": "batchSize * stateUpdateCapacity * channels",
-            "workgroupSize": "constants.workgroupSize"
           }
         },
         {
           "id": "main",
           "name": "VarlenCausalConvWithState.Stream",
           "shader": "varlen-causal-conv-stream.wgsl.jinja",
-          "bindings": "plain",
           "dispatch": {
-            "threads": "ceilDiv(totalTokens, constants.chunkSize) * channels",
-            "workgroupSize": "constants.workgroupSize"
           }
         }
       ]
@@ -382,14 +177,12 @@
       "id": "bias_state_update_stream",
       "priority": 20,
       "when": ["biasStateUpdateContract", "kernelSize >= 2", "kernelSize <= 8", "tunables.streamChunk >= 1", "floor(tunables.streamChunk) == tunables.streamChunk", "tunables.streamWorkgroupSize >= 1", "floor(tunables.streamWorkgroupSize) == tunables.streamWorkgroupSize", "tunables.streamWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.streamWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
-      "constants": {
         "hasBias": true,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
-        "usesF16": "tensorDtypes.inputT == \"float16\"",
         "workgroupSize": "tunables.streamWorkgroupSize",
-        "kernelSize": "kernelSize",
         "chunkSize": "min(tunables.streamChunk, ceilDiv(totalTokens, batchSize))"
       },
       "passes": [
@@ -397,20 +190,22 @@
           "id": "state_update",
           "name": "VarlenCausalConvWithState.StateUpdate",
           "shader": "varlen-state-update.wgsl.jinja",
-          "bindings": "state_update",
           "dispatch": {
-            "threads": "batchSize * stateUpdateCapacity * channels",
-            "workgroupSize": "constants.workgroupSize"
           }
         },
         {
           "id": "main",
           "name": "VarlenCausalConvWithState.Stream",
           "shader": "varlen-causal-conv-stream.wgsl.jinja",
-          "bindings": "bias",
           "dispatch": {
-            "threads": "ceilDiv(totalTokens, constants.chunkSize) * channels",
-            "workgroupSize": "constants.workgroupSize"
           }
         }
       ]
@@ -419,12 +214,11 @@
       "id": "plain",
       "priority": 0,
       "when": ["plainContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
-      "constants": {
         "hasBias": false,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
-        "usesF16": "tensorDtypes.inputT == \"float16\"",
         "workgroupSize": "tunables.workgroupSize"
       },
       "passes": [
@@ -432,8 +226,12 @@
           "id": "main",
           "name": "VarlenCausalConvWithState",
           "shader": "varlen-causal-conv.wgsl.jinja",
-          "bindings": "plain",
-          "dispatch": { "threads": "totalTokens * channels", "workgroupSize": "constants.workgroupSize" }
         }
       ]
     },
@@ -441,12 +239,11 @@
       "id": "bias",
       "priority": 0,
       "when": ["biasContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
-      "constants": {
         "hasBias": true,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
-        "usesF16": "tensorDtypes.inputT == \"float16\"",
         "workgroupSize": "tunables.workgroupSize"
       },
       "passes": [
@@ -454,8 +251,12 @@
           "id": "main",
           "name": "VarlenCausalConvWithState",
           "shader": "varlen-causal-conv.wgsl.jinja",
-          "bindings": "bias",
-          "dispatch": { "threads": "totalTokens * channels", "workgroupSize": "constants.workgroupSize" }
         }
       ]
     },
@@ -463,12 +264,11 @@
       "id": "state_update",
       "priority": 0,
       "when": ["stateUpdateContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
-      "constants": {
         "hasBias": false,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
-        "usesF16": "tensorDtypes.inputT == \"float16\"",
         "workgroupSize": "tunables.workgroupSize"
       },
       "passes": [
@@ -476,18 +276,23 @@
           "id": "state_update",
           "name": "VarlenCausalConvWithState.StateUpdate",
           "shader": "varlen-state-update.wgsl.jinja",
-          "bindings": "state_update",
           "dispatch": {
-            "threads": "batchSize * stateUpdateCapacity * channels",
-            "workgroupSize": "constants.workgroupSize"
           }
         },
         {
           "id": "main",
           "name": "VarlenCausalConvWithState",
           "shader": "varlen-causal-conv.wgsl.jinja",
-          "bindings": "plain",
-          "dispatch": { "threads": "totalTokens * channels", "workgroupSize": "constants.workgroupSize" }
         }
       ]
     },
@@ -495,12 +300,11 @@
       "id": "bias_state_update",
       "priority": 0,
       "when": ["biasStateUpdateContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
-      "constants": {
         "hasBias": true,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
-        "usesF16": "tensorDtypes.inputT == \"float16\"",
         "workgroupSize": "tunables.workgroupSize"
       },
       "passes": [
@@ -508,18 +312,23 @@
           "id": "state_update",
           "name": "VarlenCausalConvWithState.StateUpdate",
           "shader": "varlen-state-update.wgsl.jinja",
-          "bindings": "state_update",
           "dispatch": {
-            "threads": "batchSize * stateUpdateCapacity * channels",
-            "workgroupSize": "constants.workgroupSize"
           }
         },
         {
           "id": "main",
           "name": "VarlenCausalConvWithState",
           "shader": "varlen-causal-conv.wgsl.jinja",
-          "bindings": "bias",
-          "dispatch": { "threads": "totalTokens * channels", "workgroupSize": "constants.workgroupSize" }
         }
       ]
     }

   "domain": "com.microsoft",
   "name": "VarlenCausalConvWithState",
   "sinceVersion": 1,
+  "inputs": {
+    "inputT": { "onnx": "input", "dtype": "T", "rank": 2 },
+    "weightT": { "onnx": "weight", "dtype": "T", "rank": 3 },
+    "cumulativeSequenceLengthT": { "onnx": "cumulative_sequence_length", "dtype": "M", "rank": 1, "storage": "int32" },
+    "biasT": { "onnx": "bias", "dtype": "T", "rank": 1, "optional": true },
+    "initialStateT": { "onnx": "initial_state", "dtype": "T", "rank": 3 },
+    "captureCountT": { "onnx": "capture_count", "dtype": "M", "rank": 1, "optional": true, "storage": "int32" }
+  },
+  "outputs": {
+    "outputT": { "onnx": "output", "dtype": "T", "rank": "ranks.inputT", "shape": "shapes.inputT" },
+    "finalStateT": { "onnx": "final_state", "dtype": "T", "rank": 3, "shape": "[batchSize, channels, stateLength]" },
+    "stateUpdateT": {
+      "onnx": "state_update",
       "dtype": "T",
       "rank": 3,
       "optional": true,
+      "shape": "[batchSize, stateUpdateCapacity, channels]"
     }
   },
+  "attributes": { "activation": { "default": "none" }, "state_update_capacity": { "default": 0 } },
   "attributeConstraints": { "activation": { "values": ["none", "silu", "swish"] } },
+  "typeConstraints": { "T": ["float32", "float16"], "M": ["int32"] },
+  "tunables": {
+    "workgroupSize": { "default": 256 },
+    "streamWorkgroupSize": { "default": 64 },
+    "streamChunk": { "default": 16 }
   },
   "derive": {
     "totalTokens": "dim(shapes.inputT, 0)",
     "channels": "dim(shapes.inputT, 1)",
     "stateUpdateContract": "commonContract and not present.biasT and present.stateUpdateT and stateUpdateCapacity > 0",
     "biasStateUpdateContract": "commonContract and present.biasT and biasOk and present.stateUpdateT and stateUpdateCapacity > 0"
   },
+  "bindings": {
+    "input": { "arg": "inputT", "buffer": "read-only-storage", "elementType": "$inputScalar" },
+    "weight": { "arg": "weightT", "buffer": "read-only-storage", "elementType": "$inputScalar" },
+    "cumulative_sequence_length": {
+      "arg": "cumulativeSequenceLengthT",
+      "buffer": "read-only-storage",
+      "elementType": "i32"
+    },
+    "initial_state": { "arg": "initialStateT", "buffer": "read-only-storage", "elementType": "$inputScalar" },
+    "output": { "arg": "outputT", "buffer": "storage", "elementType": "$outputScalar" },
+    "final_state": { "arg": "finalStateT", "buffer": "storage", "elementType": "$outputScalar" },
+    "params": {
+      "buffer": "uniform",
+      "struct": [
+        { "name": "batchSize", "type": "u32", "value": "batchSize" },
+        { "name": "channels", "type": "u32", "value": "channels" },
+        { "name": "totalTokens", "type": "u32", "value": "totalTokens" },
+        { "name": "kernelSize", "type": "u32", "value": "kernelSize" }
+      ]
+    },
+    "bias": { "arg": "biasT", "buffer": "read-only-storage", "elementType": "$inputScalar" },
+    "capture_count": { "arg": "captureCountT", "buffer": "read-only-storage", "elementType": "i32" },
+    "state_update": { "arg": "stateUpdateT", "buffer": "storage", "elementType": "$outputScalar" },
+    "params_2": {
+      "name": "params",
+      "buffer": "uniform",
+      "struct": [
+        { "name": "batchSize", "type": "u32", "value": "batchSize" },
+        { "name": "channels", "type": "u32", "value": "channels" },
+        { "name": "totalTokens", "type": "u32", "value": "totalTokens" },
+        { "name": "stateUpdateCapacity", "type": "u32", "value": "stateUpdateCapacity" }
+      ]
+    }
   },
   "variants": [
     {
       "id": "plain_stream",
       "priority": 20,
       "when": ["plainContract", "kernelSize >= 2", "kernelSize <= 8", "tunables.streamChunk >= 1", "floor(tunables.streamChunk) == tunables.streamChunk", "tunables.streamWorkgroupSize >= 1", "floor(tunables.streamWorkgroupSize) == tunables.streamWorkgroupSize", "tunables.streamWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.streamWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
+      "derive": {
         "hasBias": false,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
         "workgroupSize": "tunables.streamWorkgroupSize",
         "chunkSize": "min(tunables.streamChunk, ceilDiv(totalTokens, batchSize))"
       },
       "passes": [
           "id": "main",
           "name": "VarlenCausalConvWithState.Stream",
           "shader": "varlen-causal-conv-stream.wgsl.jinja",
+          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "output", "final_state", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
+            "z": 1
           }
         }
       ]
       "id": "bias_stream",
       "priority": 20,
       "when": ["biasContract", "kernelSize >= 2", "kernelSize <= 8", "tunables.streamChunk >= 1", "floor(tunables.streamChunk) == tunables.streamChunk", "tunables.streamWorkgroupSize >= 1", "floor(tunables.streamWorkgroupSize) == tunables.streamWorkgroupSize", "tunables.streamWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.streamWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
+      "derive": {
         "hasBias": true,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
         "workgroupSize": "tunables.streamWorkgroupSize",
         "chunkSize": "min(tunables.streamChunk, ceilDiv(totalTokens, batchSize))"
       },
       "passes": [
           "id": "main",
           "name": "VarlenCausalConvWithState.Stream",
           "shader": "varlen-causal-conv-stream.wgsl.jinja",
+          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "bias", "output", "final_state", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
+            "z": 1
           }
         }
       ]
       "id": "state_update_stream",
       "priority": 20,
       "when": ["stateUpdateContract", "kernelSize >= 2", "kernelSize <= 8", "tunables.streamChunk >= 1", "floor(tunables.streamChunk) == tunables.streamChunk", "tunables.streamWorkgroupSize >= 1", "floor(tunables.streamWorkgroupSize) == tunables.streamWorkgroupSize", "tunables.streamWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.streamWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
+      "derive": {
         "hasBias": false,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
         "workgroupSize": "tunables.streamWorkgroupSize",
         "chunkSize": "min(tunables.streamChunk, ceilDiv(totalTokens, batchSize))"
       },
       "passes": [
           "id": "state_update",
           "name": "VarlenCausalConvWithState.StateUpdate",
           "shader": "varlen-state-update.wgsl.jinja",
+          "bindings": ["input", "cumulative_sequence_length", "capture_count", "state_update", "params_2"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "main",
           "name": "VarlenCausalConvWithState.Stream",
           "shader": "varlen-causal-conv-stream.wgsl.jinja",
+          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "output", "final_state", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
+            "z": 1
           }
         }
       ]
       "id": "bias_state_update_stream",
       "priority": 20,
       "when": ["biasStateUpdateContract", "kernelSize >= 2", "kernelSize <= 8", "tunables.streamChunk >= 1", "floor(tunables.streamChunk) == tunables.streamChunk", "tunables.streamWorkgroupSize >= 1", "floor(tunables.streamWorkgroupSize) == tunables.streamWorkgroupSize", "tunables.streamWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.streamWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
+      "derive": {
         "hasBias": true,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
         "workgroupSize": "tunables.streamWorkgroupSize",
         "chunkSize": "min(tunables.streamChunk, ceilDiv(totalTokens, batchSize))"
       },
       "passes": [
           "id": "state_update",
           "name": "VarlenCausalConvWithState.StateUpdate",
           "shader": "varlen-state-update.wgsl.jinja",
+          "bindings": ["input", "cumulative_sequence_length", "capture_count", "state_update", "params_2"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "main",
           "name": "VarlenCausalConvWithState.Stream",
           "shader": "varlen-causal-conv-stream.wgsl.jinja",
+          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "bias", "output", "final_state", "params"],
           "dispatch": {
+            "x": "min(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((ceilDiv(totalTokens, chunkSize) * channels), (workgroupSize)), 65535)",
+            "z": 1
           }
         }
       ]
       "id": "plain",
       "priority": 0,
       "when": ["plainContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
+      "derive": {
         "hasBias": false,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
         "workgroupSize": "tunables.workgroupSize"
       },
       "passes": [
           "id": "main",
           "name": "VarlenCausalConvWithState",
           "shader": "varlen-causal-conv.wgsl.jinja",
+          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "output", "final_state", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
       "id": "bias",
       "priority": 0,
       "when": ["biasContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
+      "derive": {
         "hasBias": true,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
         "workgroupSize": "tunables.workgroupSize"
       },
       "passes": [
           "id": "main",
           "name": "VarlenCausalConvWithState",
           "shader": "varlen-causal-conv.wgsl.jinja",
+          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "bias", "output", "final_state", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
       "id": "state_update",
       "priority": 0,
       "when": ["stateUpdateContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
+      "derive": {
         "hasBias": false,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
         "workgroupSize": "tunables.workgroupSize"
       },
       "passes": [
           "id": "state_update",
           "name": "VarlenCausalConvWithState.StateUpdate",
           "shader": "varlen-state-update.wgsl.jinja",
+          "bindings": ["input", "cumulative_sequence_length", "capture_count", "state_update", "params_2"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "main",
           "name": "VarlenCausalConvWithState",
           "shader": "varlen-causal-conv.wgsl.jinja",
+          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "output", "final_state", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     },
       "id": "bias_state_update",
       "priority": 0,
       "when": ["biasStateUpdateContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
+      "derive": {
         "hasBias": true,
         "useSilu": "attrs.activation != \"none\"",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
         "workgroupSize": "tunables.workgroupSize"
       },
       "passes": [
           "id": "state_update",
           "name": "VarlenCausalConvWithState.StateUpdate",
           "shader": "varlen-state-update.wgsl.jinja",
+          "bindings": ["input", "cumulative_sequence_length", "capture_count", "state_update", "params_2"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
+            "z": 1
           }
         },
         {
           "id": "main",
           "name": "VarlenCausalConvWithState",
           "shader": "varlen-causal-conv.wgsl.jinja",
+          "bindings": ["input", "weight", "cumulative_sequence_length", "initial_state", "bias", "output", "final_state", "params"],
+          "dispatch": {
+            "x": "min(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((totalTokens * channels), (workgroupSize)), 65535)",
+            "z": 1
+          }
         }
       ]
     }

build/webgpu/metadata.json CHANGED Viewed

@@ -1,20 +1,32 @@
 {
   "name": "com.microsoft.VarlenCausalConvWithState",
-  "id": "_com_microsoft_varlencausalconvwithstate_webgpu_9c4b81e",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "bench.json": "z/PVF6DFL2GAEJzuVOeObMuyIEskbcAJRlEve4V0LNo=",
-      "manifest.json": "nToIgIEAIC9CFMInjTS/RJQFVy0mUbtVqoNZeTxYLcU=",
-      "test.json": "c7BMrgBYnFuFqXjL+U+ZbNZoDMfq9O/jhQRWxP8wdk4=",
-      "varlen-causal-conv-stream.wgsl.jinja": "HFjAOPhuaBvGbATRz/a85lgoLzZJobT3aHhw6Ao0BBc=",
-      "varlen-causal-conv.wgsl.jinja": "KURdzm1/gZYSl8q//HtKsQ2F8ymw3T4Uge/4kDDrQZI=",
-      "varlen-state-update.wgsl.jinja": "SDAR9CKpClXi5LZR1G3z2QxBw6Fbbey6N3XuSuTeWvs="
     }
   },
-  "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
-  "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.VarlenCausalConvWithState" }
 }

 {
   "name": "com.microsoft.VarlenCausalConvWithState",
+  "id": "_com_microsoft_varlencausalconvwithstate_webgpu_7c6635e",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "y4qlCvMEobqusx4BHjuqJFrTxWGJfwFiE3ck+AejJQw=",
+      "manifest.json": "lcISl3PiMOgRa2gUETuiKNwm2RcbERhUNsPKM4yAfbc=",
+      "test.json": "66AJhdzvoC4y/XYc3VfVW0rgbaNmqskI834QnP591ng=",
+      "varlen-causal-conv-stream.wgsl.jinja": "0FWRyGR35U5ZB7Fb4Q1MFQHx90PALdOSNk1JoyeQCdQ=",
+      "varlen-causal-conv.wgsl.jinja": "OhLAVTlZId7AFw0Hfp8DWRcNZsgoMZgxpWdKxGdp4Ks=",
+      "varlen-state-update.wgsl.jinja": "aTGir5h/LQkg5ami41pKaH4EG8ybaltnnM0oJSJ6Neg="
     }
   },
+  "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
+  "webgpu": {
+    "manifestSpec": "2.0",
+    "variants": {
+      "plain_stream": ["varlen-causal-conv-stream.wgsl.jinja"],
+      "bias_stream": ["varlen-causal-conv-stream.wgsl.jinja"],
+      "state_update_stream": ["varlen-causal-conv-stream.wgsl.jinja", "varlen-state-update.wgsl.jinja"],
+      "bias_state_update_stream": ["varlen-causal-conv-stream.wgsl.jinja", "varlen-state-update.wgsl.jinja"],
+      "plain": ["varlen-causal-conv.wgsl.jinja"],
+      "bias": ["varlen-causal-conv.wgsl.jinja"],
+      "state_update": ["varlen-causal-conv.wgsl.jinja", "varlen-state-update.wgsl.jinja"],
+      "bias_state_update": ["varlen-causal-conv.wgsl.jinja", "varlen-state-update.wgsl.jinja"]
+    }
+  }
 }

build/webgpu/test.json CHANGED Viewed

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

 {
   "cases": [
     {
       "name": "k4_varlen_plain",

build/webgpu/varlen-causal-conv-stream.wgsl.jinja CHANGED Viewed

@@ -30,9 +30,6 @@ fn activate(value: f32) -> f32 {
 }
 {%- endmacro %}
 {% set stateLength = kernelSize - 1 %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.VarlenCausalConvWithState, streaming pass.
@@ -51,11 +48,10 @@ const STATE_LENGTH: u32 = {{ stateLength }}u;
 {{ varlen_conv_helpers(useSilu) }}
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit. Reduces to gid.x when nwg.y == 1.
-  let index = gid.x + gid.y * nwg.x * WG;
   let blocks = (params.totalTokens + CHUNK - 1u) / CHUNK;
   if (index >= blocks * params.channels) {
     return;

 }
 {%- endmacro %}
 {% set stateLength = kernelSize - 1 %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.VarlenCausalConvWithState, streaming pass.
 {{ varlen_conv_helpers(useSilu) }}
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
+  // per-axis dispatch fold width. Reduces to gid.x when the dispatch does not fold.
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   let blocks = (params.totalTokens + CHUNK - 1u) / CHUNK;
   if (index >= blocks * params.channels) {
     return;

build/webgpu/varlen-causal-conv.wgsl.jinja CHANGED Viewed

@@ -29,18 +29,14 @@ fn activate(value: f32) -> f32 {
 {% endif %}
 }
 {%- endmacro %}
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.VarlenCausalConvWithState, portable pass.
 // One invocation per (token, channel) output element, for any kernel extent.
 // `channels` is the contiguous axis of the token-major layout, so consecutive
-// lanes read consecutive addresses and every load coalesces. Each element
-// re-reads its kernel window, but the reuse distance is kernelSize token rows,
-// which stays cache-resident; the streaming form removes the reread where
-// the kernel extent is known.
 const WG: u32 = {{ workgroupSize }}u;
 {{ varlen_sequence_of_token() }}
@@ -48,11 +44,10 @@ const WG: u32 = {{ workgroupSize }}u;
 {{ varlen_conv_helpers(useSilu) }}
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
-  // maxComputeWorkgroupsPerDimension limit. Reduces to gid.x when nwg.y == 1.
-  let index = gid.x + gid.y * nwg.x * WG;
   if (index >= params.totalTokens * params.channels) {
     return;
   }

 {% endif %}
 }
 {%- endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.VarlenCausalConvWithState, portable pass.
 // One invocation per (token, channel) output element, for any kernel extent.
 // `channels` is the contiguous axis of the token-major layout, so consecutive
+// lanes address consecutive channel values. Each element reads its kernel
+// window; specializations with a fixed kernel extent use the streaming form to
+// avoid rereading window values.
 const WG: u32 = {{ workgroupSize }}u;
 {{ varlen_sequence_of_token() }}
 {{ varlen_conv_helpers(useSilu) }}
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   // 2D-folded flat index: gid.y carries the high bits past the
+  // per-axis dispatch fold width. Reduces to gid.x when the dispatch does not fold.
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   if (index >= params.totalTokens * params.channels) {
     return;
   }

build/webgpu/varlen-state-update.wgsl.jinja CHANGED Viewed

@@ -1,6 +1,3 @@
-{% if usesF16 %}
-enable f16;
-{% endif %}
 {{ env.wgsl.resourceDeclarations }}
 // Emit the compact transition payload defined by the contrib schema. One
@@ -9,9 +6,8 @@ enable f16;
 const WG: u32 = {{ workgroupSize }}u;
 @compute @workgroup_size(WG, 1, 1)
-fn main(@builtin(global_invocation_id) gid: vec3<u32>,
-        @builtin(num_workgroups) nwg: vec3<u32>) {
-  let index = gid.x + gid.y * nwg.x * WG;
   let count = params.batchSize * params.stateUpdateCapacity * params.channels;
   if (index >= count) {
     return;

 {{ env.wgsl.resourceDeclarations }}
 // Emit the compact transition payload defined by the contrib schema. One
 const WG: u32 = {{ workgroupSize }}u;
 @compute @workgroup_size(WG, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
   let count = params.batchSize * params.stateUpdateCapacity * params.channels;
   if (index >= count) {
     return;