sync 6fdf6301e2bb

Browse files

Files changed (8) hide show

README.md +23 -8
build/webgpu/bench.json +329 -10
build/webgpu/manifest.json +345 -65
build/webgpu/metadata.json +21 -12
build/webgpu/test.json +0 -0
build/webgpu/varlen-causal-conv-stream.wgsl.jinja +12 -9
build/webgpu/varlen-causal-conv.wgsl.jinja +17 -14
build/webgpu/varlen-unit-sequence.wgsl.jinja +59 -0

README.md CHANGED Viewed

@@ -24,7 +24,7 @@ See the [ONNX Runtime `VarlenCausalConvWithState` contrib-operator spec](https:/
 | `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
@@ -32,7 +32,7 @@ See the [ONNX Runtime `VarlenCausalConvWithState` contrib-operator spec](https:/
 | 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
@@ -42,6 +42,7 @@ Default values (overridable per request):
 | Attribute | Default | Description |
 | --- | --- | --- |
 | `activation` | `"none"` | Fused activation applied after convolution and bias. One of `none`, `silu`, or `swish`; the standard default is `none`. |
 | `state_update_capacity` | `0` | Static number of compact per-request prefix transition values to expose, in `[0, 8]`. The standard default is 0. |
 ## Type constraints
@@ -51,20 +52,34 @@ Default values (overridable per request):
 | `T` | `float32`, `float16` |
 | `M` | `int32` |
 ## 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
 - [`varlen-causal-conv-stream.wgsl.jinja`](build/webgpu/varlen-causal-conv-stream.wgsl.jinja)
 - [`varlen-causal-conv.wgsl.jinja`](build/webgpu/varlen-causal-conv.wgsl.jinja)
 - [`varlen-state-update.wgsl.jinja`](build/webgpu/varlen-state-update.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
-npm install --save-exact @huggingface/kernels@0.0.1-preview.2
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
@@ -79,9 +94,9 @@ import { getKernel } from "@huggingface/kernels";
 const kernel = await getKernel("webgpu-kernels/com.microsoft.VarlenCausalConvWithState", { version: 1 });
 const { outputT, finalStateT } = await kernel({
-  inputT: { data: inputTData, shape: [5, 6] },
-  weightT: { data: weightTData, shape: [6, 1, 4] },
-  cumulativeSequenceLengthT: { data: cumulativeSequenceLengthTData, shape: [3] },
-  initialStateT: { data: initialStateTData, shape: [2, 6, 3] },
 });
 ```

 | `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) * dilation)`, 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) * dilation)`. 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
 | Attribute | Default | Description |
 | --- | --- | --- |
 | `activation` | `"none"` | Fused activation applied after convolution and bias. One of `none`, `silu`, or `swish`; the standard default is `none`. |
+| `dilation` | `1` | Positive integer spacing between taps; the initial and final states hold (`kernel_size` - 1) * dilation raw samples. |
 | `state_update_capacity` | `0` | Static number of compact per-request prefix transition values to expose, in `[0, 8]`. The standard default is 0. |
 ## Type constraints
 | `T` | `float32`, `float16` |
 | `M` | `int32` |
+## Implementation variants
+One implementation is selected per call from the device capabilities, the request shapes and the dtypes; these notes say what each one covers.
+- `unit_plain` — Direct single-token sequence convolution with contiguous vector state traffic and fused capture writes. Strictly increasing schedules with `total_tokens == batch_size` imply one token per sequence. Channel divisibility selects the vector width; the workgroup respects device limits. Dilated and f16 requests retain their existing paths.
+- `unit_bias` — Direct single-token sequence convolution with contiguous vector state traffic and fused capture writes. Strictly increasing schedules with `total_tokens == batch_size` imply one token per sequence. Channel divisibility selects the vector width; the workgroup respects device limits. Dilated and f16 requests retain their existing paths.
+- `unit_state_update` — Direct single-token sequence convolution with contiguous vector state traffic and fused capture writes. Strictly increasing schedules with `total_tokens == batch_size` imply one token per sequence. Channel divisibility selects the vector width; the workgroup respects device limits. Dilated and f16 requests retain their existing paths.
+- `unit_bias_state_update` — Direct single-token sequence convolution with contiguous vector state traffic and fused capture writes. Strictly increasing schedules with `total_tokens == batch_size` imply one token per sequence. Channel divisibility selects the vector width; the workgroup respects device limits. Dilated and f16 requests retain their existing paths.
+- `capture_copy_state_update_stream` — Copy the capture prefix in contiguous channel vectors while preserving the convolution path and exact zero-fill semantics. The existing copy template operates on vector groups; no extra storage bindings or optional GPU features are needed.
+- `capture_copy_state_update` — Copy the capture prefix in contiguous channel vectors while preserving the convolution path and exact zero-fill semantics. The existing copy template operates on vector groups; no extra storage bindings or optional GPU features are needed.
+- `capture_copy_bias_state_update_stream` — Copy the capture prefix in contiguous channel vectors while preserving the convolution path and exact zero-fill semantics. The existing copy template operates on vector groups; no extra storage bindings or optional GPU features are needed.
+- `capture_copy_bias_state_update` — Copy the capture prefix in contiguous channel vectors while preserving the convolution path and exact zero-fill semantics. The existing copy template operates on vector groups; no extra storage bindings or optional GPU features are needed.
 ## Files
 - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
 - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
 - [`test.json`](build/webgpu/test.json) — correctness cases
+- [`bench.json`](build/webgpu/bench.json) — benchmark cases
 - [`varlen-causal-conv-stream.wgsl.jinja`](build/webgpu/varlen-causal-conv-stream.wgsl.jinja)
 - [`varlen-causal-conv.wgsl.jinja`](build/webgpu/varlen-causal-conv.wgsl.jinja)
 - [`varlen-state-update.wgsl.jinja`](build/webgpu/varlen-state-update.wgsl.jinja)
+- [`varlen-unit-sequence.wgsl.jinja`](build/webgpu/varlen-unit-sequence.wgsl.jinja)
 ## Use with `@huggingface/kernels`
 ```sh
+npm install --save-exact @huggingface/kernels@0.0.1-preview.3
 ```
 Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
 const kernel = await getKernel("webgpu-kernels/com.microsoft.VarlenCausalConvWithState", { version: 1 });
 const { outputT, finalStateT } = await kernel({
+  inputT: { data: inputTData, shape: [1, 4] },
+  weightT: { data: weightTData, shape: [4, 1, 4] },
+  cumulativeSequenceLengthT: { data: cumulativeSequenceLengthTData, shape: [2] },
+  initialStateT: { data: initialStateTData, shape: [1, 4, 3] },
 });
 ```

build/webgpu/bench.json CHANGED Viewed

@@ -4,6 +4,9 @@
     "streamWorkgroupSize": [64, 128, 256],
     "streamChunk": [16, 32, 64, 128, 256]
   },
   "cases": [
     {
       "name": "varlenconv-decode-b32-c4096-k4",
@@ -27,7 +30,16 @@
         "outputT": { "shape": [32, 4096], "dtype": "float32" },
         "finalStateT": { "shape": [32, 4096, 3], "dtype": "float32" }
       },
-      "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 32 * 4096 * dtypeBytes(args.dtype)" }] },
       "attrs": { "activation": "silu" }
     },
     {
@@ -49,7 +61,16 @@
         "outputT": { "shape": [2048, 4096], "dtype": "float32" },
         "finalStateT": { "shape": [1, 4096, 3], "dtype": "float32" }
       },
-      "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 2048 * 4096 * dtypeBytes(args.dtype)" }] },
       "attrs": { "activation": "silu" }
     },
     {
@@ -71,7 +92,16 @@
         "outputT": { "shape": [512, 4096], "dtype": "float32" },
         "finalStateT": { "shape": [8, 4096, 3], "dtype": "float32" }
       },
-      "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 512 * 4096 * dtypeBytes(args.dtype)" }] },
       "attrs": { "activation": "silu" }
     },
     {
@@ -95,10 +125,19 @@
         "finalStateT": { "shape": [6, 2048, 3], "dtype": "float32" },
         "stateUpdateT": { "dtype": "float32", "shape": [6, 4, 2048] }
       },
-      "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 378 * 2048 * dtypeBytes(args.dtype)" }] },
       "attrs": { "activation": "silu", "state_update_capacity": 4 },
       "provenance": {
-        "notes": "Benchmarks the schema-defined compact state_update alongside convolution; every request captures up to the configured contiguous-prefix capacity."
       }
     },
     {
@@ -122,10 +161,19 @@
         "finalStateT": { "shape": [6, 2048, 3], "dtype": "float32" },
         "stateUpdateT": { "dtype": "float32", "shape": [6, 4, 2048] }
       },
-      "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 378 * 2048 * dtypeBytes(args.dtype)" }] },
       "attrs": { "activation": "silu", "state_update_capacity": 4 },
       "provenance": {
-        "notes": "Benchmarks the schema-defined compact state_update alongside convolution; every request captures up to the configured contiguous-prefix capacity."
       }
     },
     {
@@ -147,7 +195,16 @@
         "outputT": { "shape": [512, 4096], "dtype": "float32" },
         "finalStateT": { "shape": [4, 4096, 3], "dtype": "float32" }
       },
-      "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 512 * 4096 * dtypeBytes(args.dtype)" }] },
       "attrs": { "activation": "silu" }
     },
     {
@@ -169,7 +226,16 @@
         "outputT": { "shape": [512, 4096], "dtype": "float16" },
         "finalStateT": { "shape": [8, 4096, 3], "dtype": "float16" }
       },
-      "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 512 * 4096 * dtypeBytes(args.dtype)" }] },
       "attrs": { "activation": "silu" }
     },
     {
@@ -191,8 +257,261 @@
         "outputT": { "shape": [512, 4096], "dtype": "float32" },
         "finalStateT": { "shape": [8, 4096, 2], "dtype": "float32" }
       },
-      "bench": { "metrics": [{ "type": "bandwidth", "value": "2 * 512 * 4096 * dtypeBytes(args.dtype)" }] },
       "attrs": { "activation": "silu" }
     }
   ]
 }

     "streamWorkgroupSize": [64, 128, 256],
     "streamChunk": [16, 32, 64, 128, 256]
   },
+  "fixtureArrays": {
+    "varlenconv-unit-capture-b64-c4096-k4-cap4_input_captureCountT": [2, 0, -3, 99, 2, 0, -3, 99, 2, 0, -3, 99, 2, 0, -3, 99, 2, 0, -3, 99, 2, 0, -3, 99, 2, 0, -3, 99, 2, 0, -3, 99, 2, 0, -3, 99, 2, 0, -3, 99, 2, 0, -3, 99, 2, 0, -3, 99, 2, 0, -3, 99, 2, 0, -3, 99, 2, 0, -3, 99, 2, 0, -3, 99]
+  },
   "cases": [
     {
       "name": "varlenconv-decode-b32-c4096-k4",
         "outputT": { "shape": [32, 4096], "dtype": "float32" },
         "finalStateT": { "shape": [32, 4096, 3], "dtype": "float32" }
       },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "Logical tensor traffic",
+            "provenance": "logical",
+            "value": "(numel(shapes.inputT) + numel(shapes.weightT) + numel(shapes.initialStateT) + numel(shapes.outputT) + numel(shapes.finalStateT) + (numel(shapes.biasT) if has(shapes, \"biasT\") else 0) + (numel(shapes.stateUpdateT) if has(shapes, \"stateUpdateT\") else 0)) * dtypeBytes(args.dtype) + 4 * (numel(shapes.cumulativeSequenceLengthT) + (numel(shapes.captureCountT) if has(shapes, \"captureCountT\") else 0))"
+          }
+        ]
+      },
       "attrs": { "activation": "silu" }
     },
     {
         "outputT": { "shape": [2048, 4096], "dtype": "float32" },
         "finalStateT": { "shape": [1, 4096, 3], "dtype": "float32" }
       },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "Logical tensor traffic",
+            "provenance": "logical",
+            "value": "(numel(shapes.inputT) + numel(shapes.weightT) + numel(shapes.initialStateT) + numel(shapes.outputT) + numel(shapes.finalStateT) + (numel(shapes.biasT) if has(shapes, \"biasT\") else 0) + (numel(shapes.stateUpdateT) if has(shapes, \"stateUpdateT\") else 0)) * dtypeBytes(args.dtype) + 4 * (numel(shapes.cumulativeSequenceLengthT) + (numel(shapes.captureCountT) if has(shapes, \"captureCountT\") else 0))"
+          }
+        ]
+      },
       "attrs": { "activation": "silu" }
     },
     {
         "outputT": { "shape": [512, 4096], "dtype": "float32" },
         "finalStateT": { "shape": [8, 4096, 3], "dtype": "float32" }
       },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "Logical tensor traffic",
+            "provenance": "logical",
+            "value": "(numel(shapes.inputT) + numel(shapes.weightT) + numel(shapes.initialStateT) + numel(shapes.outputT) + numel(shapes.finalStateT) + (numel(shapes.biasT) if has(shapes, \"biasT\") else 0) + (numel(shapes.stateUpdateT) if has(shapes, \"stateUpdateT\") else 0)) * dtypeBytes(args.dtype) + 4 * (numel(shapes.cumulativeSequenceLengthT) + (numel(shapes.captureCountT) if has(shapes, \"captureCountT\") else 0))"
+          }
+        ]
+      },
       "attrs": { "activation": "silu" }
     },
     {
         "finalStateT": { "shape": [6, 2048, 3], "dtype": "float32" },
         "stateUpdateT": { "dtype": "float32", "shape": [6, 4, 2048] }
       },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "Logical tensor traffic",
+            "provenance": "logical",
+            "value": "(numel(shapes.inputT) + numel(shapes.weightT) + numel(shapes.initialStateT) + numel(shapes.outputT) + numel(shapes.finalStateT) + (numel(shapes.biasT) if has(shapes, \"biasT\") else 0) + (numel(shapes.stateUpdateT) if has(shapes, \"stateUpdateT\") else 0)) * dtypeBytes(args.dtype) + 4 * (numel(shapes.cumulativeSequenceLengthT) + (numel(shapes.captureCountT) if has(shapes, \"captureCountT\") else 0))"
+          }
+        ]
+      },
       "attrs": { "activation": "silu", "state_update_capacity": 4 },
       "provenance": {
+        "notes": "The compact state_update accompanies convolution and captures up to the configured contiguous-prefix capacity."
       }
     },
     {
         "finalStateT": { "shape": [6, 2048, 3], "dtype": "float32" },
         "stateUpdateT": { "dtype": "float32", "shape": [6, 4, 2048] }
       },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "Logical tensor traffic",
+            "provenance": "logical",
+            "value": "(numel(shapes.inputT) + numel(shapes.weightT) + numel(shapes.initialStateT) + numel(shapes.outputT) + numel(shapes.finalStateT) + (numel(shapes.biasT) if has(shapes, \"biasT\") else 0) + (numel(shapes.stateUpdateT) if has(shapes, \"stateUpdateT\") else 0)) * dtypeBytes(args.dtype) + 4 * (numel(shapes.cumulativeSequenceLengthT) + (numel(shapes.captureCountT) if has(shapes, \"captureCountT\") else 0))"
+          }
+        ]
+      },
       "attrs": { "activation": "silu", "state_update_capacity": 4 },
       "provenance": {
+        "notes": "The compact state_update accompanies convolution and captures up to the configured contiguous-prefix capacity."
       }
     },
     {
         "outputT": { "shape": [512, 4096], "dtype": "float32" },
         "finalStateT": { "shape": [4, 4096, 3], "dtype": "float32" }
       },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "Logical tensor traffic",
+            "provenance": "logical",
+            "value": "(numel(shapes.inputT) + numel(shapes.weightT) + numel(shapes.initialStateT) + numel(shapes.outputT) + numel(shapes.finalStateT) + (numel(shapes.biasT) if has(shapes, \"biasT\") else 0) + (numel(shapes.stateUpdateT) if has(shapes, \"stateUpdateT\") else 0)) * dtypeBytes(args.dtype) + 4 * (numel(shapes.cumulativeSequenceLengthT) + (numel(shapes.captureCountT) if has(shapes, \"captureCountT\") else 0))"
+          }
+        ]
+      },
       "attrs": { "activation": "silu" }
     },
     {
         "outputT": { "shape": [512, 4096], "dtype": "float16" },
         "finalStateT": { "shape": [8, 4096, 3], "dtype": "float16" }
       },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "Logical tensor traffic",
+            "provenance": "logical",
+            "value": "(numel(shapes.inputT) + numel(shapes.weightT) + numel(shapes.initialStateT) + numel(shapes.outputT) + numel(shapes.finalStateT) + (numel(shapes.biasT) if has(shapes, \"biasT\") else 0) + (numel(shapes.stateUpdateT) if has(shapes, \"stateUpdateT\") else 0)) * dtypeBytes(args.dtype) + 4 * (numel(shapes.cumulativeSequenceLengthT) + (numel(shapes.captureCountT) if has(shapes, \"captureCountT\") else 0))"
+          }
+        ]
+      },
       "attrs": { "activation": "silu" }
     },
     {
         "outputT": { "shape": [512, 4096], "dtype": "float32" },
         "finalStateT": { "shape": [8, 4096, 2], "dtype": "float32" }
       },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "Logical tensor traffic",
+            "provenance": "logical",
+            "value": "(numel(shapes.inputT) + numel(shapes.weightT) + numel(shapes.initialStateT) + numel(shapes.outputT) + numel(shapes.finalStateT) + (numel(shapes.biasT) if has(shapes, \"biasT\") else 0) + (numel(shapes.stateUpdateT) if has(shapes, \"stateUpdateT\") else 0)) * dtypeBytes(args.dtype) + 4 * (numel(shapes.cumulativeSequenceLengthT) + (numel(shapes.captureCountT) if has(shapes, \"captureCountT\") else 0))"
+          }
+        ]
+      },
       "attrs": { "activation": "silu" }
+    },
+    {
+      "name": "varlenconv-unit-capture-b1-c4096-k4-cap4",
+      "preset": "smoke",
+      "vars": { "dtype": "float32" },
+      "inputs": {
+        "inputT": { "shape": [1, 4096], "dtype": "float32", "dist": "normal", "seed": 4100, "scale": 1 },
+        "weightT": { "shape": [4096, 1, 4], "dtype": "float32", "dist": "normal", "seed": 4101, "scale": 1 },
+        "cumulativeSequenceLengthT": { "dtype": "int32", "shape": [2], "data": { "kind": "values", "values": [0, 1] } },
+        "biasT": { "shape": [4096], "dtype": "float32", "dist": "normal", "seed": 4102, "scale": 1 },
+        "initialStateT": { "shape": [1, 4096, 3], "dtype": "float32", "dist": "normal", "seed": 4103, "scale": 1 },
+        "captureCountT": { "dtype": "int32", "shape": [1], "data": { "kind": "values", "values": [2] } }
+      },
+      "outputs": {
+        "outputT": { "shape": [1, 4096], "dtype": "float32" },
+        "finalStateT": { "shape": [1, 4096, 3], "dtype": "float32" },
+        "stateUpdateT": { "dtype": "float32", "shape": [1, 4, 4096] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "Logical tensor traffic",
+            "provenance": "logical",
+            "value": "(numel(shapes.inputT) + numel(shapes.weightT) + numel(shapes.initialStateT) + numel(shapes.outputT) + numel(shapes.finalStateT) + (numel(shapes.biasT) if has(shapes, \"biasT\") else 0) + (numel(shapes.stateUpdateT) if has(shapes, \"stateUpdateT\") else 0)) * dtypeBytes(args.dtype) + 4 * (numel(shapes.cumulativeSequenceLengthT) + (numel(shapes.captureCountT) if has(shapes, \"captureCountT\") else 0))"
+          }
+        ]
+      },
+      "attrs": { "activation": "silu", "state_update_capacity": 4 }
+    },
+    {
+      "name": "varlenconv-unit-capture-b64-c4096-k4-cap4",
+      "preset": "smoke",
+      "vars": { "dtype": "float32" },
+      "inputs": {
+        "inputT": { "shape": [64, 4096], "dtype": "float32", "dist": "normal", "seed": 4100, "scale": 1 },
+        "weightT": { "shape": [4096, 1, 4], "dtype": "float32", "dist": "normal", "seed": 4101, "scale": 1 },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [65],
+          "data": {
+            "kind": "values",
+            "values": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64]
+          }
+        },
+        "biasT": { "shape": [4096], "dtype": "float32", "dist": "normal", "seed": 4102, "scale": 1 },
+        "initialStateT": { "shape": [64, 4096, 3], "dtype": "float32", "dist": "normal", "seed": 4103, "scale": 1 },
+        "captureCountT": {
+          "dtype": "int32",
+          "shape": [64],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/varlenconv-unit-capture-b64-c4096-k4-cap4_input_captureCountT" }
+          }
+        }
+      },
+      "outputs": {
+        "outputT": { "shape": [64, 4096], "dtype": "float32" },
+        "finalStateT": { "shape": [64, 4096, 3], "dtype": "float32" },
+        "stateUpdateT": { "dtype": "float32", "shape": [64, 4, 4096] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "Logical tensor traffic",
+            "provenance": "logical",
+            "value": "(numel(shapes.inputT) + numel(shapes.weightT) + numel(shapes.initialStateT) + numel(shapes.outputT) + numel(shapes.finalStateT) + (numel(shapes.biasT) if has(shapes, \"biasT\") else 0) + (numel(shapes.stateUpdateT) if has(shapes, \"stateUpdateT\") else 0)) * dtypeBytes(args.dtype) + 4 * (numel(shapes.cumulativeSequenceLengthT) + (numel(shapes.captureCountT) if has(shapes, \"captureCountT\") else 0))"
+          }
+        ]
+      },
+      "attrs": { "activation": "silu", "state_update_capacity": 4 }
+    },
+    {
+      "name": "varlenconv-unit-alignment-b64-c4098-k3-cap4",
+      "preset": "smoke",
+      "vars": { "dtype": "float32" },
+      "inputs": {
+        "inputT": { "shape": [64, 4098], "dtype": "float32", "dist": "normal", "seed": 4100, "scale": 1 },
+        "weightT": { "shape": [4098, 1, 3], "dtype": "float32", "dist": "normal", "seed": 4101, "scale": 1 },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [65],
+          "dist": "linearMod",
+          "mod": 65,
+          "step": 1,
+          "offset": 0
+        },
+        "biasT": { "shape": [4098], "dtype": "float32", "dist": "normal", "seed": 4102, "scale": 1 },
+        "initialStateT": { "shape": [64, 4098, 2], "dtype": "float32", "dist": "normal", "seed": 4103, "scale": 1 },
+        "captureCountT": {
+          "dtype": "int32",
+          "shape": [64],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/varlenconv-unit-capture-b64-c4096-k4-cap4_input_captureCountT" }
+          }
+        }
+      },
+      "outputs": {
+        "outputT": { "shape": [64, 4098], "dtype": "float32" },
+        "finalStateT": { "shape": [64, 4098, 2], "dtype": "float32" },
+        "stateUpdateT": { "dtype": "float32", "shape": [64, 4, 4098] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "Logical tensor traffic",
+            "provenance": "logical",
+            "value": "(numel(shapes.inputT) + numel(shapes.weightT) + numel(shapes.initialStateT) + numel(shapes.outputT) + numel(shapes.finalStateT) + (numel(shapes.biasT) if has(shapes, \"biasT\") else 0) + (numel(shapes.stateUpdateT) if has(shapes, \"stateUpdateT\") else 0)) * dtypeBytes(args.dtype) + 4 * (numel(shapes.cumulativeSequenceLengthT) + (numel(shapes.captureCountT) if has(shapes, \"captureCountT\") else 0))"
+          }
+        ]
+      },
+      "attrs": { "activation": "silu", "state_update_capacity": 4 }
+    },
+    {
+      "name": "varlenconv-unit-alignment-b64-c4097-k8-cap8",
+      "preset": "smoke",
+      "vars": { "dtype": "float32" },
+      "inputs": {
+        "inputT": { "shape": [64, 4097], "dtype": "float32", "dist": "normal", "seed": 4100, "scale": 1 },
+        "weightT": { "shape": [4097, 1, 8], "dtype": "float32", "dist": "normal", "seed": 4101, "scale": 1 },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [65],
+          "dist": "linearMod",
+          "mod": 65,
+          "step": 1,
+          "offset": 0
+        },
+        "biasT": { "shape": [4097], "dtype": "float32", "dist": "normal", "seed": 4102, "scale": 1 },
+        "initialStateT": { "shape": [64, 4097, 7], "dtype": "float32", "dist": "normal", "seed": 4103, "scale": 1 },
+        "captureCountT": {
+          "dtype": "int32",
+          "shape": [64],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/varlenconv-unit-capture-b64-c4096-k4-cap4_input_captureCountT" }
+          }
+        }
+      },
+      "outputs": {
+        "outputT": { "shape": [64, 4097], "dtype": "float32" },
+        "finalStateT": { "shape": [64, 4097, 7], "dtype": "float32" },
+        "stateUpdateT": { "dtype": "float32", "shape": [64, 8, 4097] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "Logical tensor traffic",
+            "provenance": "logical",
+            "value": "(numel(shapes.inputT) + numel(shapes.weightT) + numel(shapes.initialStateT) + numel(shapes.outputT) + numel(shapes.finalStateT) + (numel(shapes.biasT) if has(shapes, \"biasT\") else 0) + (numel(shapes.stateUpdateT) if has(shapes, \"stateUpdateT\") else 0)) * dtypeBytes(args.dtype) + 4 * (numel(shapes.cumulativeSequenceLengthT) + (numel(shapes.captureCountT) if has(shapes, \"captureCountT\") else 0))"
+          }
+        ]
+      },
+      "attrs": { "activation": "silu", "state_update_capacity": 8 }
+    },
+    {
+      "name": "varlenconv-capture-b64-t128-c4096-k4-d1",
+      "preset": "smoke",
+      "vars": { "dtype": "float32" },
+      "inputs": {
+        "inputT": { "shape": [128, 4096], "dtype": "float32", "dist": "normal", "seed": 4100, "scale": 1 },
+        "weightT": { "shape": [4096, 1, 4], "dtype": "float32", "dist": "normal", "seed": 4101, "scale": 1 },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [65],
+          "dist": "linearMod",
+          "mod": 130,
+          "step": 2,
+          "offset": 0
+        },
+        "biasT": { "shape": [4096], "dtype": "float32", "dist": "normal", "seed": 4102, "scale": 1 },
+        "initialStateT": { "shape": [64, 4096, 3], "dtype": "float32", "dist": "normal", "seed": 4103, "scale": 1 },
+        "captureCountT": {
+          "dtype": "int32",
+          "shape": [64],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/varlenconv-unit-capture-b64-c4096-k4-cap4_input_captureCountT" }
+          }
+        }
+      },
+      "outputs": {
+        "outputT": { "shape": [128, 4096], "dtype": "float32" },
+        "finalStateT": { "shape": [64, 4096, 3], "dtype": "float32" },
+        "stateUpdateT": { "dtype": "float32", "shape": [64, 4, 4096] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "Logical tensor traffic",
+            "provenance": "logical",
+            "value": "(numel(shapes.inputT) + numel(shapes.weightT) + numel(shapes.initialStateT) + numel(shapes.outputT) + numel(shapes.finalStateT) + (numel(shapes.biasT) if has(shapes, \"biasT\") else 0) + (numel(shapes.stateUpdateT) if has(shapes, \"stateUpdateT\") else 0)) * dtypeBytes(args.dtype) + 4 * (numel(shapes.cumulativeSequenceLengthT) + (numel(shapes.captureCountT) if has(shapes, \"captureCountT\") else 0))"
+          }
+        ]
+      },
+      "attrs": { "activation": "silu", "state_update_capacity": 4, "dilation": 1 }
+    },
+    {
+      "name": "varlenconv-capture-b64-t128-c4096-k4-d2",
+      "preset": "smoke",
+      "vars": { "dtype": "float32" },
+      "inputs": {
+        "inputT": { "shape": [128, 4096], "dtype": "float32", "dist": "normal", "seed": 4100, "scale": 1 },
+        "weightT": { "shape": [4096, 1, 4], "dtype": "float32", "dist": "normal", "seed": 4101, "scale": 1 },
+        "cumulativeSequenceLengthT": {
+          "dtype": "int32",
+          "shape": [65],
+          "dist": "linearMod",
+          "mod": 130,
+          "step": 2,
+          "offset": 0
+        },
+        "biasT": { "shape": [4096], "dtype": "float32", "dist": "normal", "seed": 4102, "scale": 1 },
+        "initialStateT": { "shape": [64, 4096, 6], "dtype": "float32", "dist": "normal", "seed": 4103, "scale": 1 },
+        "captureCountT": {
+          "dtype": "int32",
+          "shape": [64],
+          "data": {
+            "kind": "values",
+            "values": { "$ref": "#/fixtureArrays/varlenconv-unit-capture-b64-c4096-k4-cap4_input_captureCountT" }
+          }
+        }
+      },
+      "outputs": {
+        "outputT": { "shape": [128, 4096], "dtype": "float32" },
+        "finalStateT": { "shape": [64, 4096, 6], "dtype": "float32" },
+        "stateUpdateT": { "dtype": "float32", "shape": [64, 4, 4096] }
+      },
+      "bench": {
+        "metrics": [
+          {
+            "type": "bandwidth",
+            "name": "Logical tensor traffic",
+            "provenance": "logical",
+            "value": "(numel(shapes.inputT) + numel(shapes.weightT) + numel(shapes.initialStateT) + numel(shapes.outputT) + numel(shapes.finalStateT) + (numel(shapes.biasT) if has(shapes, \"biasT\") else 0) + (numel(shapes.stateUpdateT) if has(shapes, \"stateUpdateT\") else 0)) * dtypeBytes(args.dtype) + 4 * (numel(shapes.cumulativeSequenceLengthT) + (numel(shapes.captureCountT) if has(shapes, \"captureCountT\") else 0))"
+          }
+        ]
+      },
+      "attrs": { "activation": "silu", "state_update_capacity": 4, "dilation": 2 }
     }
   ]
 }

build/webgpu/manifest.json CHANGED Viewed

@@ -21,7 +21,11 @@
       "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": {
@@ -35,33 +39,31 @@
     "batchSize": "dim(shapes.cumulativeSequenceLengthT, 0) - 1",
     "weightRankOk": "ranks.weightT == 3 and dim(shapes.weightT, 1) == 1",
     "kernelSize": "dim(shapes.weightT, 2)",
-    "stateLength": "kernelSize - 1",
     "stateUpdateCapacity": "attrs.state_update_capacity",
     "capturePairOk": "present.captureCountT == (stateUpdateCapacity > 0)",
     "stateShape": "ranks.initialStateT == 3 and dim(shapes.initialStateT, 0) == batchSize and dim(shapes.initialStateT, 1) == channels and dim(shapes.initialStateT, 2) == stateLength",
     "finalStateShape": "ranks.finalStateT == 3 and dim(shapes.finalStateT, 0) == batchSize and dim(shapes.finalStateT, 1) == channels and dim(shapes.finalStateT, 2) == stateLength",
     "captureShape": "not present.captureCountT or (ranks.captureCountT == 1 and dim(shapes.captureCountT, 0) == batchSize and tensorDtypes.captureCountT == \"int32\")",
     "stateUpdateShape": "not present.stateUpdateT or (ranks.stateUpdateT == 3 and dim(shapes.stateUpdateT, 0) == batchSize and dim(shapes.stateUpdateT, 1) == stateUpdateCapacity and dim(shapes.stateUpdateT, 2) == channels and tensorDtypes.stateUpdateT == tensorDtypes.inputT)",
-    "commonContract": "ranks.inputT == 2 and weightRankOk and ranks.outputT == 2 and ranks.cumulativeSequenceLengthT == 1 and channels >= 1 and kernelSize >= 1 and batchSize >= 1 and totalTokens >= batchSize and stateUpdateCapacity >= 0 and stateUpdateCapacity <= 8 and floor(stateUpdateCapacity) == stateUpdateCapacity and capturePairOk and captureShape and stateUpdateShape and (tensorDtypes.inputT == \"float32\" or tensorDtypes.inputT == \"float16\") and tensorDtypes.weightT == tensorDtypes.inputT and tensorDtypes.outputT == tensorDtypes.inputT and tensorDtypes.initialStateT == tensorDtypes.inputT and tensorDtypes.finalStateT == tensorDtypes.inputT and tensorDtypes.cumulativeSequenceLengthT == \"int32\" and f16Ok(dtypes.T) and dim(shapes.weightT, 0) == channels and dim(shapes.outputT, 0) == totalTokens and dim(shapes.outputT, 1) == channels and stateShape and finalStateShape",
     "biasOk": "ranks.biasT == 1 and tensorDtypes.biasT == tensorDtypes.inputT and dim(shapes.biasT, 0) == channels",
     "plainContract": "commonContract and not present.biasT and (not present.stateUpdateT or stateUpdateCapacity == 0)",
     "biasContract": "commonContract and present.biasT and biasOk and (not present.stateUpdateT or stateUpdateCapacity == 0)",
     "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" },
@@ -69,34 +71,173 @@
         { "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",
@@ -111,38 +252,51 @@
       ]
     },
     {
-      "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",
@@ -153,18 +307,23 @@
           "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)",
@@ -174,34 +333,62 @@
       ]
     },
     {
-      "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)",
@@ -215,8 +402,7 @@
       "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"
@@ -235,13 +421,37 @@
         }
       ]
     },
     {
       "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"
@@ -260,13 +470,48 @@
         }
       ]
     },
     {
       "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"
@@ -276,7 +521,7 @@
           "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)",
@@ -296,13 +541,48 @@
         }
       ]
     },
     {
       "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"
@@ -312,7 +592,7 @@
           "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)",

       "shape": "[batchSize, stateUpdateCapacity, channels]"
     }
   },
+  "attributes": {
+    "activation": { "default": "none" },
+    "dilation": { "default": 1 },
+    "state_update_capacity": { "default": 0 }
+  },
   "attributeConstraints": { "activation": { "values": ["none", "silu", "swish"] } },
   "typeConstraints": { "T": ["float32", "float16"], "M": ["int32"] },
   "tunables": {
     "batchSize": "dim(shapes.cumulativeSequenceLengthT, 0) - 1",
     "weightRankOk": "ranks.weightT == 3 and dim(shapes.weightT, 1) == 1",
     "kernelSize": "dim(shapes.weightT, 2)",
+    "stateLength": "(kernelSize - 1) * attrs.dilation",
     "stateUpdateCapacity": "attrs.state_update_capacity",
     "capturePairOk": "present.captureCountT == (stateUpdateCapacity > 0)",
     "stateShape": "ranks.initialStateT == 3 and dim(shapes.initialStateT, 0) == batchSize and dim(shapes.initialStateT, 1) == channels and dim(shapes.initialStateT, 2) == stateLength",
     "finalStateShape": "ranks.finalStateT == 3 and dim(shapes.finalStateT, 0) == batchSize and dim(shapes.finalStateT, 1) == channels and dim(shapes.finalStateT, 2) == stateLength",
     "captureShape": "not present.captureCountT or (ranks.captureCountT == 1 and dim(shapes.captureCountT, 0) == batchSize and tensorDtypes.captureCountT == \"int32\")",
     "stateUpdateShape": "not present.stateUpdateT or (ranks.stateUpdateT == 3 and dim(shapes.stateUpdateT, 0) == batchSize and dim(shapes.stateUpdateT, 1) == stateUpdateCapacity and dim(shapes.stateUpdateT, 2) == channels and tensorDtypes.stateUpdateT == tensorDtypes.inputT)",
+    "commonContract": "ranks.inputT == 2 and weightRankOk and ranks.outputT == 2 and ranks.cumulativeSequenceLengthT == 1 and channels >= 1 and kernelSize >= 1 and batchSize >= 1 and totalTokens >= batchSize and stateUpdateCapacity >= 0 and stateUpdateCapacity <= 8 and floor(stateUpdateCapacity) == stateUpdateCapacity and capturePairOk and captureShape and stateUpdateShape and (tensorDtypes.inputT == \"float32\" or tensorDtypes.inputT == \"float16\") and tensorDtypes.weightT == tensorDtypes.inputT and tensorDtypes.outputT == tensorDtypes.inputT and tensorDtypes.initialStateT == tensorDtypes.inputT and tensorDtypes.finalStateT == tensorDtypes.inputT and tensorDtypes.cumulativeSequenceLengthT == \"int32\" and f16Ok(dtypes.T) and dim(shapes.weightT, 0) == channels and dim(shapes.outputT, 0) == totalTokens and dim(shapes.outputT, 1) == channels and stateShape and finalStateShape and attrs.dilation >= 1 and floor(attrs.dilation) == attrs.dilation",
     "biasOk": "ranks.biasT == 1 and tensorDtypes.biasT == tensorDtypes.inputT and dim(shapes.biasT, 0) == channels",
     "plainContract": "commonContract and not present.biasT and (not present.stateUpdateT or stateUpdateCapacity == 0)",
     "biasContract": "commonContract and present.biasT and biasOk and (not present.stateUpdateT or stateUpdateCapacity == 0)",
     "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",
+    "causalDilation": "attrs.dilation",
+    "unitWg": "min(tunables.streamWorkgroupSize, device.limits.maxComputeInvocationsPerWorkgroup, device.limits.maxComputeWorkgroupSizeX)",
+    "useSilu": "attrs.activation != \"none\""
   },
   "bindings": {
+    "input": { "arg": "inputT", "elementType": "$inputScalar" },
+    "weight": { "arg": "weightT", "elementType": "$inputScalar" },
+    "cumulative_sequence_length": { "arg": "cumulativeSequenceLengthT", "elementType": "i32" },
+    "initial_state": { "arg": "initialStateT", "elementType": "$inputScalar" },
+    "output": { "arg": "outputT", "elementType": "$outputScalar" },
+    "final_state": { "arg": "finalStateT", "elementType": "$outputScalar" },
     "params": {
       "struct": [
         { "name": "batchSize", "type": "u32", "value": "batchSize" },
         { "name": "channels", "type": "u32", "value": "channels" },
         { "name": "kernelSize", "type": "u32", "value": "kernelSize" }
       ]
     },
+    "bias": { "arg": "biasT", "elementType": "$inputScalar" },
+    "capture_count": { "arg": "captureCountT", "elementType": "i32" },
+    "state_update": { "arg": "stateUpdateT", "elementType": "$outputScalar" },
+    "params__uniform": {
       "name": "params",
       "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" }
       ]
+    },
+    "input_unit": { "arg": "inputT", "name": "input", "elementType": "$unitStorageType" },
+    "bias_unit": { "arg": "biasT", "name": "bias", "elementType": "$unitStorageType" },
+    "output_unit": { "arg": "outputT", "name": "output", "elementType": "$unitStorageType" },
+    "state_update_unit": { "arg": "stateUpdateT", "name": "state_update", "elementType": "$unitStorageType" },
+    "params_unit": {
+      "name": "params",
+      "struct": [
+        { "name": "channels", "type": "u32", "value": "channels" },
+        { "name": "totalTokens", "type": "u32", "value": "totalTokens" }
+      ]
+    },
+    "weight_unit": { "arg": "weightT", "elementType": "$unitStorageType", "name": "weight" },
+    "initial_state_unit": { "arg": "initialStateT", "elementType": "$unitStorageType", "name": "initial_state" },
+    "final_state_unit": { "arg": "finalStateT", "elementType": "$unitStorageType", "name": "final_state" },
+    "input_capture": { "arg": "inputT", "name": "input", "elementType": "$captureStorageType" },
+    "state_update_capture": { "arg": "stateUpdateT", "name": "state_update", "elementType": "$captureStorageType" },
+    "params_capture": {
+      "name": "params",
+      "struct": [
+        { "name": "batchSize", "type": "u32", "value": "batchSize" },
+        { "name": "channels", "type": "u32", "value": "channels / captureVectorWidth" },
+        { "name": "totalTokens", "type": "u32", "value": "totalTokens" },
+        { "name": "stateUpdateCapacity", "type": "u32", "value": "stateUpdateCapacity" }
+      ]
     }
   },
   "variants": [
     {
+      "id": "unit_plain",
+      "priority": 50,
+      "when": ["plainContract", "tensorDtypes.inputT == \"float32\"", "totalTokens == batchSize", "attrs.dilation == 1", "kernelSize >= 2", "kernelSize <= 8", "unitWg >= 1", "floor(unitWg) == unitWg"],
+      "derive": {
+        "hasBias": false,
+        "hasCapture": false,
+        "unitVectorWidth": "4 if channels % 4 == 0 else (2 if channels % 2 == 0 else 1)",
+        "unitStorageType": "dtypes.T if unitVectorWidth == 1 else \"vec\" ~ unitVectorWidth ~ \"<\" ~ dtypes.T ~ \">\"",
+        "unitValueType": "\"f32\" if unitVectorWidth == 1 else \"vec\" ~ unitVectorWidth ~ \"<f32>\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "VarlenCausalConvWithState.UnitSequence",
+          "shader": "varlen-unit-sequence.wgsl.jinja",
+          "bindings": ["input_unit", "weight_unit", "initial_state_unit", "output_unit", "final_state_unit", "params_unit"],
+          "dispatch": {
+            "x": "min(ceilDiv((totalTokens * channels / unitVectorWidth), (unitWg)), 65535)",
+            "y": "ceilDiv(ceilDiv((totalTokens * channels / unitVectorWidth), (unitWg)), 65535)",
+            "z": 1
+          }
+        }
+      ]
+    },
+    {
+      "id": "unit_bias",
+      "priority": 50,
+      "when": ["biasContract", "tensorDtypes.inputT == \"float32\"", "totalTokens == batchSize", "attrs.dilation == 1", "kernelSize >= 2", "kernelSize <= 8", "unitWg >= 1", "floor(unitWg) == unitWg"],
+      "derive": {
+        "hasBias": true,
+        "hasCapture": false,
+        "unitVectorWidth": "4 if channels % 4 == 0 else (2 if channels % 2 == 0 else 1)",
+        "unitStorageType": "dtypes.T if unitVectorWidth == 1 else \"vec\" ~ unitVectorWidth ~ \"<\" ~ dtypes.T ~ \">\"",
+        "unitValueType": "\"f32\" if unitVectorWidth == 1 else \"vec\" ~ unitVectorWidth ~ \"<f32>\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "VarlenCausalConvWithState.UnitSequence",
+          "shader": "varlen-unit-sequence.wgsl.jinja",
+          "bindings": ["input_unit", "weight_unit", "initial_state_unit", "bias_unit", "output_unit", "final_state_unit", "params_unit"],
+          "dispatch": {
+            "x": "min(ceilDiv((totalTokens * channels / unitVectorWidth), (unitWg)), 65535)",
+            "y": "ceilDiv(ceilDiv((totalTokens * channels / unitVectorWidth), (unitWg)), 65535)",
+            "z": 1
+          }
+        }
+      ]
+    },
+    {
+      "id": "unit_state_update",
+      "priority": 50,
+      "when": ["stateUpdateContract", "tensorDtypes.inputT == \"float32\"", "totalTokens == batchSize", "attrs.dilation == 1", "kernelSize >= 2", "kernelSize <= 8", "unitWg >= 1", "floor(unitWg) == unitWg"],
+      "derive": {
+        "hasBias": false,
+        "hasCapture": true,
+        "unitVectorWidth": "4 if channels % 4 == 0 else (2 if channels % 2 == 0 else 1)",
+        "unitStorageType": "dtypes.T if unitVectorWidth == 1 else \"vec\" ~ unitVectorWidth ~ \"<\" ~ dtypes.T ~ \">\"",
+        "unitValueType": "\"f32\" if unitVectorWidth == 1 else \"vec\" ~ unitVectorWidth ~ \"<f32>\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "VarlenCausalConvWithState.UnitSequence",
+          "shader": "varlen-unit-sequence.wgsl.jinja",
+          "bindings": ["input_unit", "weight_unit", "initial_state_unit", "output_unit", "final_state_unit", "params_unit", "capture_count", "state_update_unit"],
+          "dispatch": {
+            "x": "min(ceilDiv((totalTokens * channels / unitVectorWidth), (unitWg)), 65535)",
+            "y": "ceilDiv(ceilDiv((totalTokens * channels / unitVectorWidth), (unitWg)), 65535)",
+            "z": 1
+          }
+        }
+      ]
+    },
+    {
+      "id": "unit_bias_state_update",
+      "priority": 50,
+      "when": ["biasStateUpdateContract", "tensorDtypes.inputT == \"float32\"", "totalTokens == batchSize", "attrs.dilation == 1", "kernelSize >= 2", "kernelSize <= 8", "unitWg >= 1", "floor(unitWg) == unitWg"],
+      "derive": {
+        "hasBias": true,
+        "hasCapture": true,
+        "unitVectorWidth": "4 if channels % 4 == 0 else (2 if channels % 2 == 0 else 1)",
+        "unitStorageType": "dtypes.T if unitVectorWidth == 1 else \"vec\" ~ unitVectorWidth ~ \"<\" ~ dtypes.T ~ \">\"",
+        "unitValueType": "\"f32\" if unitVectorWidth == 1 else \"vec\" ~ unitVectorWidth ~ \"<f32>\""
+      },
+      "passes": [
+        {
+          "id": "main",
+          "name": "VarlenCausalConvWithState.UnitSequence",
+          "shader": "varlen-unit-sequence.wgsl.jinja",
+          "bindings": ["input_unit", "weight_unit", "initial_state_unit", "bias_unit", "output_unit", "final_state_unit", "params_unit", "capture_count", "state_update_unit"],
+          "dispatch": {
+            "x": "min(ceilDiv((totalTokens * channels / unitVectorWidth), (unitWg)), 65535)",
+            "y": "ceilDiv(ceilDiv((totalTokens * channels / unitVectorWidth), (unitWg)), 65535)",
+            "z": 1
+          }
+        }
+      ]
+    },
+    {
+      "id": "capture_copy_state_update_stream",
+      "priority": 25,
+      "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", "attrs.dilation == 1", "tensorDtypes.inputT == \"float32\"", "channels % 2 == 0"],
       "derive": {
         "hasBias": false,
         "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_capture", "cumulative_sequence_length", "capture_count", "state_update_capture", "params_capture"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * stateUpdateCapacity * channels / captureVectorWidth), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * stateUpdateCapacity * channels / captureVectorWidth), (workgroupSize)), 65535)",
+            "z": 1
+          },
+          "derive": {
+            "captureVectorWidth": "4 if channels % 4 == 0 else 2",
+            "captureStorageType": "\"vec\" ~ captureVectorWidth ~ \"<\" ~ dtypes.T ~ \">\"",
+            "outputScalar": "captureStorageType"
+          }
+        },
         {
           "id": "main",
           "name": "VarlenCausalConvWithState.Stream",
       ]
     },
     {
+      "id": "capture_copy_state_update",
+      "priority": 5,
+      "when": ["stateUpdateContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX", "tensorDtypes.inputT == \"float32\"", "channels % 2 == 0"],
       "derive": {
+        "hasBias": false,
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
+        "workgroupSize": "tunables.workgroupSize"
       },
       "passes": [
+        {
+          "id": "state_update",
+          "name": "VarlenCausalConvWithState.StateUpdate",
+          "shader": "varlen-state-update.wgsl.jinja",
+          "bindings": ["input_capture", "cumulative_sequence_length", "capture_count", "state_update_capture", "params_capture"],
+          "dispatch": {
+            "x": "min(ceilDiv((batchSize * stateUpdateCapacity * channels / captureVectorWidth), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * stateUpdateCapacity * channels / captureVectorWidth), (workgroupSize)), 65535)",
+            "z": 1
+          },
+          "derive": {
+            "captureVectorWidth": "4 if channels % 4 == 0 else 2",
+            "captureStorageType": "\"vec\" ~ captureVectorWidth ~ \"<\" ~ dtypes.T ~ \">\"",
+            "outputScalar": "captureStorageType"
+          }
+        },
         {
           "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": "capture_copy_bias_state_update_stream",
+      "priority": 25,
+      "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", "attrs.dilation == 1", "tensorDtypes.inputT == \"float32\"", "channels % 2 == 0"],
       "derive": {
+        "hasBias": true,
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
         "workgroupSize": "tunables.streamWorkgroupSize",
           "id": "state_update",
           "name": "VarlenCausalConvWithState.StateUpdate",
           "shader": "varlen-state-update.wgsl.jinja",
+          "bindings": ["input_capture", "cumulative_sequence_length", "capture_count", "state_update_capture", "params_capture"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * stateUpdateCapacity * channels / captureVectorWidth), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * stateUpdateCapacity * channels / captureVectorWidth), (workgroupSize)), 65535)",
             "z": 1
+          },
+          "derive": {
+            "captureVectorWidth": "4 if channels % 4 == 0 else 2",
+            "captureStorageType": "\"vec\" ~ captureVectorWidth ~ \"<\" ~ dtypes.T ~ \">\"",
+            "outputScalar": "captureStorageType"
           }
         },
         {
           "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)",
       ]
     },
     {
+      "id": "capture_copy_bias_state_update",
+      "priority": 5,
+      "when": ["biasStateUpdateContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX", "tensorDtypes.inputT == \"float32\"", "channels % 2 == 0"],
       "derive": {
         "hasBias": true,
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
+        "workgroupSize": "tunables.workgroupSize"
       },
       "passes": [
         {
           "id": "state_update",
           "name": "VarlenCausalConvWithState.StateUpdate",
           "shader": "varlen-state-update.wgsl.jinja",
+          "bindings": ["input_capture", "cumulative_sequence_length", "capture_count", "state_update_capture", "params_capture"],
           "dispatch": {
+            "x": "min(ceilDiv((batchSize * stateUpdateCapacity * channels / captureVectorWidth), (workgroupSize)), 65535)",
+            "y": "ceilDiv(ceilDiv((batchSize * stateUpdateCapacity * channels / captureVectorWidth), (workgroupSize)), 65535)",
             "z": 1
+          },
+          "derive": {
+            "captureVectorWidth": "4 if channels % 4 == 0 else 2",
+            "captureStorageType": "\"vec\" ~ captureVectorWidth ~ \"<\" ~ dtypes.T ~ \">\"",
+            "outputScalar": "captureStorageType"
           }
         },
+        {
+          "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": "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", "attrs.dilation == 1"],
+      "derive": {
+        "hasBias": "present.biasT",
+        "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)",
       "priority": 0,
       "when": ["plainContract", "tunables.workgroupSize >= 1", "floor(tunables.workgroupSize) == tunables.workgroupSize", "tunables.workgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.workgroupSize <= device.limits.maxComputeWorkgroupSizeX"],
       "derive": {
+        "hasBias": "present.biasT",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
         "workgroupSize": "tunables.workgroupSize"
         }
       ]
     },
+    {
+      "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", "attrs.dilation == 1"],
+      "derive": {
+        "hasBias": "present.biasT",
+        "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": "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": "present.biasT",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
         "workgroupSize": "tunables.workgroupSize"
         }
       ]
     },
+    {
+      "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", "attrs.dilation == 1"],
+      "derive": {
+        "hasBias": "present.biasT",
+        "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__uniform"],
+          "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": "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": "present.biasT",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
         "workgroupSize": "tunables.workgroupSize"
           "id": "state_update",
           "name": "VarlenCausalConvWithState.StateUpdate",
           "shader": "varlen-state-update.wgsl.jinja",
+          "bindings": ["input", "cumulative_sequence_length", "capture_count", "state_update", "params__uniform"],
           "dispatch": {
             "x": "min(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
         }
       ]
     },
+    {
+      "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", "attrs.dilation == 1"],
+      "derive": {
+        "hasBias": "present.biasT",
+        "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__uniform"],
+          "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": "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": "present.biasT",
         "inputScalar": "dtypes.T",
         "outputScalar": "dtypes.T",
         "workgroupSize": "tunables.workgroupSize"
           "id": "state_update",
           "name": "VarlenCausalConvWithState.StateUpdate",
           "shader": "varlen-state-update.wgsl.jinja",
+          "bindings": ["input", "cumulative_sequence_length", "capture_count", "state_update", "params__uniform"],
           "dispatch": {
             "x": "min(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",
             "y": "ceilDiv(ceilDiv((batchSize * stateUpdateCapacity * channels), (workgroupSize)), 65535)",

build/webgpu/metadata.json CHANGED Viewed

@@ -1,31 +1,40 @@
 {
   "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"]
     }
   }

 {
   "name": "com.microsoft.VarlenCausalConvWithState",
+  "id": "_com_microsoft_varlencausalconvwithstate_webgpu_a11f8c1",
   "version": 1,
   "license": "Apache-2.0",
   "backend": { "type": "webgpu" },
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "bench.json": "0VIH7EDZLjXt1ZgOmAA4oDeWfhadYrY8VJOkvCIbDlg=",
+      "manifest.json": "uEm8UogYuIyHJRhVKAebGdlv7yxdEIRJs7rLubPctf0=",
+      "test.json": "o32EnBOjkjEdcFA/wkWFt1ywXTCCN20uEndKoBIoYB8=",
+      "varlen-causal-conv-stream.wgsl.jinja": "th2IihEJq0gHEpdRm2PNN8MleOvgSZlAstRE0W6wXmo=",
+      "varlen-causal-conv.wgsl.jinja": "t9K8tccFYN5I9uDeOAUs+gtNaHrQS6iy2S0IbQtp4ik=",
+      "varlen-state-update.wgsl.jinja": "aTGir5h/LQkg5ami41pKaH4EG8ybaltnnM0oJSJ6Neg=",
+      "varlen-unit-sequence.wgsl.jinja": "g2tYe2vdYUlDWySCvrljSsqUAuodUmjYnOW/6tr51jw="
     }
   },
+  "provenance": { "kernel": { "sha": "6fdf6301e2bbcc2f03bf1eaf493b7ad55ef33afc", "dirty": false } },
   "webgpu": {
+    "manifestSpec": "2.1",
     "variants": {
+      "unit_plain": ["varlen-unit-sequence.wgsl.jinja"],
+      "unit_bias": ["varlen-unit-sequence.wgsl.jinja"],
+      "unit_state_update": ["varlen-unit-sequence.wgsl.jinja"],
+      "unit_bias_state_update": ["varlen-unit-sequence.wgsl.jinja"],
+      "capture_copy_state_update_stream": ["varlen-causal-conv-stream.wgsl.jinja", "varlen-state-update.wgsl.jinja"],
+      "capture_copy_state_update": ["varlen-causal-conv.wgsl.jinja", "varlen-state-update.wgsl.jinja"],
+      "capture_copy_bias_state_update_stream": ["varlen-causal-conv-stream.wgsl.jinja", "varlen-state-update.wgsl.jinja"],
+      "capture_copy_bias_state_update": ["varlen-causal-conv.wgsl.jinja", "varlen-state-update.wgsl.jinja"],
       "plain_stream": ["varlen-causal-conv-stream.wgsl.jinja"],
       "plain": ["varlen-causal-conv.wgsl.jinja"],
+      "bias_stream": ["varlen-causal-conv-stream.wgsl.jinja"],
       "bias": ["varlen-causal-conv.wgsl.jinja"],
+      "state_update_stream": ["varlen-causal-conv-stream.wgsl.jinja", "varlen-state-update.wgsl.jinja"],
       "state_update": ["varlen-causal-conv.wgsl.jinja", "varlen-state-update.wgsl.jinja"],
+      "bias_state_update_stream": ["varlen-causal-conv-stream.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

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

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

@@ -6,9 +6,7 @@ fn sequence_of_token(token: u32, batch_size: u32) -> u32 {
     }
   }
   return batch_size - 1u;
-}
-{%- endmacro %}
 {% macro varlen_conv_helpers(useSilu) %}
 // virtual[pos] is one sequence's raw sample stream: the kernelSize-1 carried samples
 // from initial_state, then the sequence's own tokens. Reading the window through this
@@ -21,15 +19,22 @@ fn virtual_sample(start: u32, state_base: u32, channel: u32, state_length: u32,
   return f32(initial_state[state_base + pos]);
 }
-fn activate(value: f32) -> f32 {
 {% if useSilu %}
   return value / (1.0 + exp(-value));
 {% else %}
   return value;
 {% endif %}
-}
-{%- endmacro %}
 {% set stateLength = kernelSize - 1 %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.VarlenCausalConvWithState, streaming pass.
@@ -49,9 +54,7 @@ const STATE_LENGTH: u32 = {{ stateLength }}u;
 @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;

     }
   }
   return batch_size - 1u;
+}{% endmacro %}
 {% macro varlen_conv_helpers(useSilu) %}
 // virtual[pos] is one sequence's raw sample stream: the kernelSize-1 carried samples
 // from initial_state, then the sequence's own tokens. Reading the window through this
   return f32(initial_state[state_base + pos]);
 }
+{{ varlen_conv_activation(useSilu) }}{% endmacro %}
+{% macro varlen_conv_activation(useSilu, width=1) %}
+{% set valueType = "f32" if width == 1 else "vec" ~ width ~ "<f32>" %}
+fn activate(value: {{ valueType }}) -> {{ valueType }} {
 {% if useSilu %}
   return value / (1.0 + exp(-value));
 {% else %}
   return value;
 {% endif %}
+}{% endmacro %}
 {% set stateLength = kernelSize - 1 %}
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.VarlenCausalConvWithState, streaming pass.
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WG", "index", "") }}
   let blocks = (params.totalTokens + CHUNK - 1u) / CHUNK;
   if (index >= blocks * params.channels) {
     return;

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

@@ -6,9 +6,7 @@ fn sequence_of_token(token: u32, batch_size: u32) -> u32 {
     }
   }
   return batch_size - 1u;
-}
-{%- endmacro %}
 {% macro varlen_conv_helpers(useSilu) %}
 // virtual[pos] is one sequence's raw sample stream: the kernelSize-1 carried samples
 // from initial_state, then the sequence's own tokens. Reading the window through this
@@ -21,14 +19,24 @@ fn virtual_sample(start: u32, state_base: u32, channel: u32, state_length: u32,
   return f32(initial_state[state_base + pos]);
 }
-fn activate(value: f32) -> f32 {
 {% if useSilu %}
   return value / (1.0 + exp(-value));
 {% else %}
   return value;
 {% endif %}
-}
-{%- endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.VarlenCausalConvWithState, portable pass.
@@ -45,15 +53,10 @@ const WG: u32 = {{ workgroupSize }}u;
 @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;
-  }
   let token = index / params.channels;
   let channel = index % params.channels;
-  let state_length = params.kernelSize - 1u;
   // A malformed schedule returns before any buffer access. This is memory-safety
   // containment, not validation; outputs are unspecified for such a request.
@@ -75,7 +78,7 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
   acc = f32(bias[channel]);
 {% endif %}
   for (var k: u32 = 0u; k < params.kernelSize; k = k + 1u) {
-    let value = virtual_sample(start, state_base, channel, state_length, local + k);
     acc = acc + value * f32(weight[channel * params.kernelSize + k]);
   }
   output[index] = {{ outputScalar }}(activate(acc));

     }
   }
   return batch_size - 1u;
+}{% endmacro %}
 {% macro varlen_conv_helpers(useSilu) %}
 // virtual[pos] is one sequence's raw sample stream: the kernelSize-1 carried samples
 // from initial_state, then the sequence's own tokens. Reading the window through this
   return f32(initial_state[state_base + pos]);
 }
+{{ varlen_conv_activation(useSilu) }}{% endmacro %}
+{% macro varlen_conv_activation(useSilu, width=1) %}
+{% set valueType = "f32" if width == 1 else "vec" ~ width ~ "<f32>" %}
+fn activate(value: {{ valueType }}) -> {{ valueType }} {
 {% if useSilu %}
   return value / (1.0 + exp(-value));
 {% else %}
   return value;
 {% endif %}
+}{% endmacro %}
+{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
+{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
+  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
+  // per-axis workgroup fold width.
+  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};
+  if ({{ name }} >= {{ bound }}) {
+    return;
+  }{% endmacro %}
 {{ env.wgsl.resourceDeclarations }}
 // com.microsoft.VarlenCausalConvWithState, portable pass.
 @compute @workgroup_size(WG, 1, 1)
 fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+{{ flat_index_2d("WG", "index", "params.totalTokens * params.channels") }}
   let token = index / params.channels;
   let channel = index % params.channels;
+  let state_length = {% if causalDilation != 1 %}(params.kernelSize - 1u) * {{ causalDilation }}u{% else %}params.kernelSize - 1u{% endif %};
   // A malformed schedule returns before any buffer access. This is memory-safety
   // containment, not validation; outputs are unspecified for such a request.
   acc = f32(bias[channel]);
 {% endif %}
   for (var k: u32 = 0u; k < params.kernelSize; k = k + 1u) {
+    let value = virtual_sample(start, state_base, channel, state_length, local + k{% if causalDilation != 1 %} * {{ causalDilation }}u{% endif %});
     acc = acc + value * f32(weight[channel * params.kernelSize + k]);
   }
   output[index] = {{ outputScalar }}(activate(acc));

build/webgpu/varlen-unit-sequence.wgsl.jinja ADDED Viewed

	@@ -0,0 +1,59 @@

+{% macro varlen_conv_activation(useSilu, width=1) %}
+{% set valueType = "f32" if width == 1 else "vec" ~ width ~ "<f32>" %}
+fn activate(value: {{ valueType }}) -> {{ valueType }} {
+{% if useSilu %}
+  return value / ({% if width == 1 %}1.0{% else %}{{ valueType }}(1.0){% endif %} + exp(-value));
+{% else %}
+  return value;
+{% endif %}
+}{% endmacro %}
+{% set stateLength = kernelSize - 1 %}
+{{ env.wgsl.resourceDeclarations }}
+{{ varlen_conv_activation(useSilu, unitVectorWidth) }}
+const WG: u32 = {{ unitWg }}u;
+@compute @workgroup_size(WG)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
+  let channelGroups = params.channels / {{ unitVectorWidth }}u;
+  if (index >= params.totalTokens * channelGroups) { return; }
+  // Strictly increasing integer prefix sums over T == batch imply one token per sequence.
+  let sequence = index / channelGroups;
+  let channel = index % channelGroups;
+  // Loads follow the contiguous channel-major storage layout. Their components
+  // are rearranged into channel vectors at compile time for the original
+  // ordered tap sum.
+  let weightBase = channel * {{ kernelSize }}u;
+  let stateBase = (sequence * channelGroups + channel) * {{ stateLength }}u;
+{% for i in range(kernelSize) %}
+  let packedW{{ i }} = weight[weightBase + {{ i }}u];
+{% endfor %}
+{% for i in range(stateLength) %}
+  let packedS{{ i }} = initial_state[stateBase + {{ i }}u];
+{% endfor %}
+{% for i in range(kernelSize) %}
+  let w{{ i }} = {{ unitValueType }}({% for c in range(unitVectorWidth) %}packedW{{ ((c * kernelSize + i) / unitVectorWidth)|int }}{% if unitVectorWidth > 1 %}[{{ (c * kernelSize + i) % unitVectorWidth }}]{% endif %}{% if not loop.last %}, {% endif %}{% endfor %});
+{% endfor %}
+{% for i in range(stateLength) %}
+  let s{{ i }} = {{ unitValueType }}({% for c in range(unitVectorWidth) %}packedS{{ ((c * stateLength + i) / unitVectorWidth)|int }}{% if unitVectorWidth > 1 %}[{{ (c * stateLength + i) % unitVectorWidth }}]{% endif %}{% if not loop.last %}, {% endif %}{% endfor %});
+{% endfor %}
+{% if hasBias %}
+  let biasValue = {{ unitValueType }}(bias[channel]);
+{% else %}
+  let biasValue = {{ unitValueType }}(0.0);
+{% endif %}
+  let x = {{ unitValueType }}(input[index]);
+  output[index] = {{ unitStorageType }}(activate(biasValue{% for i in range(stateLength) %} + s{{ i }} * w{{ i }}{% endfor %} + x * w{{ kernelSize - 1 }}));
+// Repack the shifted state without scalar scatter stores.
+{% for i in range(stateLength) %}
+  final_state[stateBase + {{ i }}u] = {{ unitStorageType }}({% for c in range(unitVectorWidth) %}{% set scalarOffset = i * unitVectorWidth + c %}{% set tap = scalarOffset % stateLength %}{% if tap == stateLength - 1 %}x{% else %}s{{ tap + 1 }}{% endif %}{% if unitVectorWidth > 1 %}[{{ (scalarOffset / stateLength)|int }}]{% endif %}{% if not loop.last %}, {% endif %}{% endfor %});
+{% endfor %}
+{% if hasCapture %}
+  var captured = {{ unitStorageType }}(0.0);
+  if (capture_count[sequence] > 0) { captured = {{ unitStorageType }}(x); }
+  let updateBase = (sequence * {{ stateUpdateCapacity }}u) * channelGroups + channel;
+  state_update[updateBase] = captured;
+{% for slot in range(1, stateUpdateCapacity) %}
+  state_update[updateBase + {{ slot }}u * channelGroups] = {{ unitStorageType }}(0.0);
+{% endfor %}
+{% endif %}
+}