Hugging Face's logo

Kernels:
webgpu-kernels
/
com.microsoft.CausalConvWithState

Kernel card Files
xet
 Community
Xenova HF Staff commited on
Commit
476435a
·
verified ·
1 Parent(s): 53a7bdc

sync 91d990483a17

Browse files
Files changed (8) hide show
  1. README.md +16 -12
  2. build/webgpu/bench.json +4 -5
  3. build/webgpu/causal-conv-with-state-tiled.wgsl.jinja +2 -3
  4. build/webgpu/causal-conv-with-state-vec4.wgsl.jinja +2 -3
  5. build/webgpu/causal-conv-with-state.wgsl.jinja +2 -3
  6. build/webgpu/manifest.json +132 -939
  7. build/webgpu/metadata.json +25 -9
  8. build/webgpu/test.json +73 -13
README.md CHANGED
@@ -12,25 +12,25 @@ tags:
12
 
13
  ## Description
14
 
15
- Legacy Microsoft contrib form of stateful 1-D causal depthwise convolution. Each channel uses its own `(channels, 1, kernel)` weight over current and past positions, with optional activation and `past_state`/`present_state` tensors for incremental decoding. The contrib-only `state_window` attribute may retain several rollback states. This inference implementation preserves the existing contrib ABI with `ndim = 1`, float16 or float32 tensors, and float32 accumulation; spatial ranks 2 and 3 and bfloat16 are not implemented.
16
 
17
  See the [ONNX Runtime `CausalConvWithState` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.CausalConvWithState) for the reference semantics.
18
 
19
  ## Inputs
20
 
21
- | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
22
  | --- | --- | --- | --- | --- | --- | --- |
23
- | `input` | `inputT` | `T` | `3` | — | Channels-first input tensor with shape `(batch_size, channels, sequence_length)` for the supported 1-D mode. | required |
24
- | `weight` | `weightT` | `T` | `3` | — | Depthwise convolution kernel with shape `(channels, 1, kernel_size)` for the supported 1-D mode. | required |
25
- | `bias` | `biasT` | `T` | `1` | — | Optional per-channel bias with shape `(channels,)`. | optional |
26
- | `past_state` | `pastStateT` | `T` | derived | — | Carry state from the previous step; shape `(batch_size, channels, k_1 - 1)`, or `(W, batch_size, channels, k_1 - 1)` when `state_window = W > 0`, in which case only slot `W - 1` is read. If absent, the left-side padding is zero. | optional |
27
 
28
  ## Outputs
29
 
30
- | Name | Bind key | Logical dtype | Rank | Shape | Description | Presence |
31
  | --- | --- | --- | --- | --- | --- | --- |
32
- | `output` | `outputT` | `T` | `3` | same as `input` | Convolution output with the same shape as `input`. | required |
33
- | `present_state` | `presentStateT` | `T` | derived | derived; see description | Updated carry state; shape `(batch_size, channels, k_1 - 1)`, or `(W, batch_size, channels, k_1 - 1)` when `state_window = W > 0`. Slot `W - 1` holds the last `k - 1` values along the causal axis; slot `j` holds the same for the prefix ending at position `seq_len - W + j`. | required |
34
 
35
  ## Attributes
36
 
@@ -50,7 +50,7 @@ Default values (overridable per request):
50
 
51
  ## Files
52
 
53
- - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, provenance)
54
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
55
  - [`test.json`](build/webgpu/test.json) — correctness cases
56
  - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
@@ -60,10 +60,14 @@ Default values (overridable per request):
60
 
61
  ## Use with `@huggingface/kernels`
62
 
63
- The loader derives every required output's shape and logical dtype from the manifest contract and this call.
64
- It then allocates the result tensors automatically.
 
 
 
65
 
66
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
 
67
 
68
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
69
 
 
12
 
13
  ## Description
14
 
15
+ Microsoft contrib stateful 1-D causal depthwise convolution. Each channel uses its own `(channels, 1, kernel_size)` weight over current and past positions, with optional activation and `past_state`/`present_state` tensors for incremental decoding. The `state_window` attribute may retain several rollback states. This package supports `ndim = 1`, float16 or float32 tensors, and float32 accumulation; spatial ranks 2 and 3 and bfloat16 are unsupported.
16
 
17
  See the [ONNX Runtime `CausalConvWithState` contrib-operator spec](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.CausalConvWithState) for the reference semantics.
18
 
19
  ## Inputs
20
 
21
+ | Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
22
  | --- | --- | --- | --- | --- | --- | --- |
23
+ | `inputT` | `input` | `T` | `3` | — | Channels-first input tensor with shape `(batch_size, channels, sequence_length)` for the supported 1-D mode. | required |
24
+ | `weightT` | `weight` | `T` | `3` | — | Depthwise convolution kernel with shape `(channels, 1, kernel_size)` for the supported 1-D mode. | required |
25
+ | `biasT` | `bias` | `T` | `1` | — | Optional per-channel bias with shape `(channels,)`. | optional |
26
+ | `pastStateT` | `past_state` | `T` | derived | — | Carry state from the previous step; shape `(batch_size, channels, kernel_size - 1)`, or `(W, batch_size, channels, kernel_size - 1)` when `state_window = W > 0`, in which case only slot `W - 1` is read. If absent, the left-side padding is zero. | optional |
27
 
28
  ## Outputs
29
 
30
+ | Name | Upstream name | Logical dtype | Rank | Shape | Description | Presence |
31
  | --- | --- | --- | --- | --- | --- | --- |
32
+ | `outputT` | `output` | `T` | `3` | same as `inputT` | Convolution output with the same shape as `input`. | required |
33
+ | `presentStateT` | `present_state` | `T` | derived | derived | Updated carry state; shape `(batch_size, channels, kernel_size - 1)`, or `(W, batch_size, channels, kernel_size - 1)` when `state_window = W > 0`. Slot `W - 1` holds the last `kernel_size - 1` values along the causal axis; slot `j` holds the same values for the prefix ending at position `sequence_length - W + j`. | required |
34
 
35
  ## Attributes
36
 
 
50
 
51
  ## Files
52
 
53
+ - [`metadata.json`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
54
  - [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
55
  - [`test.json`](build/webgpu/test.json) — correctness cases
56
  - [`bench.json`](build/webgpu/bench.json) — benchmark + tuning cases
 
60
 
61
  ## Use with `@huggingface/kernels`
62
 
63
+ ```sh
64
+ npm install --save-exact @huggingface/kernels@0.0.1-preview.2
65
+ ```
66
+
67
+ Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.
68
 
69
  The `version: 1` option selects the published kernel contract; it is independent of any operator opset, contrib `since_version`, or model version.
70
+ It follows the `v1` branch as fixes land. To pin exact artifact bytes, pass a 40-character commit `revision` instead of `version`.
71
 
72
  Replace each `*Data` placeholder with a typed array containing the corresponding input data.
73
 
build/webgpu/bench.json CHANGED
@@ -1,5 +1,4 @@
1
  {
2
- "op": "com.microsoft.CausalConvWithState",
3
  "tunableSpace": { "workgroupSize": [64, 128, 256] },
4
  "cases": [
5
  {
@@ -81,7 +80,7 @@
81
  "name": "causal-conv-f32-large-kernel127-prefill-b2c1024t512-alignment-pathology",
82
  "preset": "stress",
83
  "provenance": {
84
- "notes": "A realistic odd-width causal kernel adjacent to the K=128 tiled control. K=127 is valid but fails the tiled kernel's K % 4 contract and falls back to the scalar zero-state implementation."
85
  },
86
  "vars": { "batch": 2, "channels": 1024, "length": 512, "kernel": 127 },
87
  "attrs": { "activation": "none" },
@@ -108,7 +107,7 @@
108
  "preset": "stress",
109
  "tunableSpace": { "tiledWorkgroupSize": [64, 128, 256] },
110
  "provenance": {
111
- "notes": "Regression lock for the historical large-kernel bias cliff: this common per-channel bias path previously excluded the tiled specialization and selected the scalar kernel."
112
  },
113
  "vars": { "batch": 2, "channels": 1024, "length": 512, "kernel": 128 },
114
  "attrs": { "activation": "silu" },
@@ -136,7 +135,7 @@
136
  "preset": "stress",
137
  "tunableSpace": { "tiledWorkgroupSize": [64, 128, 256] },
138
  "provenance": {
139
- "notes": "Regression lock for the historical continuation cliff: valid carry state previously excluded the tiled specialization even without bias or activation."
140
  },
141
  "vars": { "batch": 2, "channels": 1024, "length": 512, "kernel": 128 },
142
  "attrs": { "activation": "none" },
@@ -164,7 +163,7 @@
164
  "preset": "stress",
165
  "tunableSpace": { "tiledWorkgroupSize": [64, 128, 256] },
166
  "provenance": {
167
- "notes": "A realistic chunked continuation and regression lock: valid carry state plus bias previously excluded the tiled specialization and selected the scalar state+bias implementation."
168
  },
169
  "vars": { "batch": 2, "channels": 1024, "length": 512, "kernel": 128 },
170
  "attrs": { "activation": "silu" },
 
1
  {
 
2
  "tunableSpace": { "workgroupSize": [64, 128, 256] },
3
  "cases": [
4
  {
 
80
  "name": "causal-conv-f32-large-kernel127-prefill-b2c1024t512-alignment-pathology",
81
  "preset": "stress",
82
  "provenance": {
83
+ "notes": "A valid odd-width causal kernel with `K = 127` fails the tiled path's four-tap alignment requirement and selects the scalar zero-state implementation."
84
  },
85
  "vars": { "batch": 2, "channels": 1024, "length": 512, "kernel": 127 },
86
  "attrs": { "activation": "none" },
 
107
  "preset": "stress",
108
  "tunableSpace": { "tiledWorkgroupSize": [64, 128, 256] },
109
  "provenance": {
110
+ "notes": "A large-kernel prefill with per-channel bias exercises the tiled specialization at batch 2, 1,024 channels, 512 time steps, and kernel size 128."
111
  },
112
  "vars": { "batch": 2, "channels": 1024, "length": 512, "kernel": 128 },
113
  "attrs": { "activation": "silu" },
 
135
  "preset": "stress",
136
  "tunableSpace": { "tiledWorkgroupSize": [64, 128, 256] },
137
  "provenance": {
138
+ "notes": "A large-kernel continuation with valid carry state and no bias or activation exercises the tiled specialization."
139
  },
140
  "vars": { "batch": 2, "channels": 1024, "length": 512, "kernel": 128 },
141
  "attrs": { "activation": "none" },
 
163
  "preset": "stress",
164
  "tunableSpace": { "tiledWorkgroupSize": [64, 128, 256] },
165
  "provenance": {
166
+ "notes": "A chunked large-kernel continuation with valid carry state and per-channel bias exercises the tiled state-and-bias specialization."
167
  },
168
  "vars": { "batch": 2, "channels": 1024, "length": 512, "kernel": 128 },
169
  "attrs": { "activation": "silu" },
build/webgpu/causal-conv-with-state-tiled.wgsl.jinja CHANGED
@@ -29,12 +29,11 @@ fn activate(value: f32) -> f32 {
29
 
30
  @compute @workgroup_size(WG, 1, 1)
31
  fn main(@builtin(local_invocation_id) lid3: vec3<u32>,
32
- @builtin(workgroup_id) wid: vec3<u32>,
33
- @builtin(num_workgroups) nwg: vec3<u32>) {
34
  let lane = lid3.x;
35
  let tiles_per_row = (params.length + TILE - 1u) / TILE;
36
  // Recover the logical workgroup index after an oversized grid folds into y.
37
- let flat_wg = wid.x + wid.y * nwg.x;
38
  let total_wg = params.batchSize * params.channels * tiles_per_row;
39
  if (flat_wg >= total_wg) {
40
  return;
 
29
 
30
  @compute @workgroup_size(WG, 1, 1)
31
  fn main(@builtin(local_invocation_id) lid3: vec3<u32>,
32
+ @builtin(workgroup_id) wid: vec3<u32>) {
 
33
  let lane = lid3.x;
34
  let tiles_per_row = (params.length + TILE - 1u) / TILE;
35
  // Recover the logical workgroup index after an oversized grid folds into y.
36
+ let flat_wg = wid.x + wid.y * {{ DISPATCH_FOLD_WIDTH }}u;
37
  let total_wg = params.batchSize * params.channels * tiles_per_row;
38
  if (flat_wg >= total_wg) {
39
  return;
build/webgpu/causal-conv-with-state-vec4.wgsl.jinja CHANGED
@@ -21,12 +21,11 @@ fn activate4(value: vec4<f32>) -> vec4<f32> {
21
  }
22
 
23
  @compute @workgroup_size(WG, 1, 1)
24
- fn main(@builtin(global_invocation_id) gid: vec3<u32>,
25
- @builtin(num_workgroups) nwg: vec3<u32>) {
26
  let row_vecs = params.length / 4u;
27
  let work_size = params.batchSize * params.channels * row_vecs;
28
  // Recover the logical 1D index after an oversized dispatch is folded into y.
29
- let index = gid.x + gid.y * nwg.x * WG;
30
  if (index >= work_size) {
31
  return;
32
  }
 
21
  }
22
 
23
  @compute @workgroup_size(WG, 1, 1)
24
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
25
  let row_vecs = params.length / 4u;
26
  let work_size = params.batchSize * params.channels * row_vecs;
27
  // Recover the logical 1D index after an oversized dispatch is folded into y.
28
+ let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
29
  if (index >= work_size) {
30
  return;
31
  }
build/webgpu/causal-conv-with-state.wgsl.jinja CHANGED
@@ -15,14 +15,13 @@ fn activate(value: f32) -> f32 {
15
  }
16
 
17
  @compute @workgroup_size(WG, 1, 1)
18
- fn main(@builtin(global_invocation_id) gid: vec3<u32>,
19
- @builtin(num_workgroups) nwg: vec3<u32>) {
20
  // length == 0 still writes the present_state carryover. The nonzero grid keeps
21
  // one state-writing thread per (batch, channel) even when there is no output.
22
  let len_nz = max(1u, params.length);
23
  let work_size = params.batchSize * params.channels * len_nz;
24
  // gid.y carries the high bits past the per-dimension dispatch limit.
25
- let index = gid.x + gid.y * nwg.x * WG;
26
  if (index >= work_size) {
27
  return;
28
  }
 
15
  }
16
 
17
  @compute @workgroup_size(WG, 1, 1)
18
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 
19
  // length == 0 still writes the present_state carryover. The nonzero grid keeps
20
  // one state-writing thread per (batch, channel) even when there is no output.
21
  let len_nz = max(1u, params.length);
22
  let work_size = params.batchSize * params.channels * len_nz;
23
  // gid.y carries the high bits past the per-dimension dispatch limit.
24
+ let index = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * WG;
25
  if (index >= work_size) {
26
  return;
27
  }
build/webgpu/manifest.json CHANGED
@@ -2,68 +2,30 @@
2
  "domain": "com.microsoft",
3
  "name": "CausalConvWithState",
4
  "sinceVersion": 1,
5
- "description": "Legacy Microsoft contrib form of stateful 1-D causal depthwise convolution. Each channel uses its own `(channels, 1, kernel)` weight over current and past positions, with optional activation and `past_state`/`present_state` tensors for incremental decoding. The contrib-only `state_window` attribute may retain several rollback states. This inference implementation preserves the existing contrib ABI with `ndim = 1`, float16 or float32 tensors, and float32 accumulation; spatial ranks 2 and 3 and bfloat16 are not implemented.",
6
- "inputs": [
7
- {
8
- "role": "input",
9
- "dtype": "T",
10
- "rank": 3,
11
- "description": "Channels-first input tensor with shape `(batch_size, channels, sequence_length)` for the supported 1-D mode."
12
- },
13
- {
14
- "role": "weight",
15
- "dtype": "T",
16
- "rank": 3,
17
- "description": "Depthwise convolution kernel with shape `(channels, 1, kernel_size)` for the supported 1-D mode."
18
- },
19
- {
20
- "role": "bias",
21
- "dtype": "T",
22
- "rank": 1,
23
- "optional": true,
24
- "description": "Optional per-channel bias with shape `(channels,)`."
25
- },
26
- {
27
- "role": "past_state",
28
  "dtype": "T",
29
  "rank": "3 if attrs.state_window == 0 else 4",
30
- "optional": true,
31
- "description": "Carry state from the previous step; shape `(batch_size, channels, k_1 - 1)`, or `(W, batch_size, channels, k_1 - 1)` when `state_window = W > 0`, in which case only slot `W - 1` is read. If absent, the left-side padding is zero."
32
  }
33
- ],
34
- "outputs": [
35
- {
36
- "role": "output",
37
- "dtype": "T",
38
- "rank": 3,
39
- "shape": "shapes.input",
40
- "description": "Convolution output with the same shape as `input`."
41
- },
42
- {
43
- "role": "present_state",
44
  "dtype": "T",
45
  "rank": "3 if attrs.state_window == 0 else 4",
46
- "shape": "[dim(shapes.input, 0), dim(shapes.input, 1), dim(shapes.weight, 2) - 1] if attrs.state_window == 0 else [attrs.state_window, dim(shapes.input, 0), dim(shapes.input, 1), dim(shapes.weight, 2) - 1]",
47
- "description": "Updated carry state; shape `(batch_size, channels, k_1 - 1)`, or `(W, batch_size, channels, k_1 - 1)` when `state_window = W > 0`. Slot `W - 1` holds the last `k - 1` values along the causal axis; slot `j` holds the same for the prefix ending at position `seq_len - W + j`."
48
  }
49
- ],
50
- "attributes": { "activation": "none", "ndim": 1, "state_window": 0 },
51
- "attributeConstraints": { "activation": { "values": ["none", "silu", "swish"] }, "ndim": { "values": [1] } },
52
- "attributeDescriptions": {
53
- "activation": "Activation applied after convolution and bias. Defaults to `none`; `swish` is an alias of SiLU.",
54
- "ndim": "Number of spatial dimensions. This implementation supports the contrib 1D mode (`ndim = 1`).",
55
- "state_window": "Contrib extension selecting the number of rollback state slots to retain, in the range 0 through 8. Defaults to 0."
56
  },
 
 
57
  "typeConstraints": { "T": ["float32", "float16"] },
58
- "args": {
59
- "inputT": { "kind": "tensor", "semantic": "input", "role": "input" },
60
- "weightT": { "kind": "tensor", "semantic": "weight", "role": "input" },
61
- "biasT": { "kind": "tensor", "semantic": "bias", "role": "input", "required": false },
62
- "pastStateT": { "kind": "tensor", "semantic": "past_state", "role": "input", "required": false },
63
- "outputT": { "kind": "tensor", "semantic": "output", "role": "output" },
64
- "presentStateT": { "kind": "tensor", "semantic": "present_state", "role": "output" }
65
- },
66
- "tunables": { "workgroupSize": 256, "tiledWorkgroupSize": 128 },
67
  "derive": {
68
  "stateWindow": "attrs.state_window",
69
  "windowed": "stateWindow > 0",
@@ -82,790 +44,45 @@
82
  "stateNoBiasContract": "commonContract and present.pastStateT and not present.biasT and tensorDtypes.pastStateT == tensorDtypes.inputT and pastStateShapeOk",
83
  "stateBiasContract": "commonContract and present.pastStateT and present.biasT and ranks.biasT == 1 and tensorDtypes.pastStateT == tensorDtypes.inputT and tensorDtypes.biasT == tensorDtypes.inputT and pastStateShapeOk and dim(shapes.biasT, 0) == dim(shapes.inputT, 1)"
84
  },
85
- "bindingSets": {
86
- "zeroScalar": [
87
- {
88
- "name": "input",
89
- "arg": "inputT",
90
- "semantic": "input",
91
- "buffer": { "type": "read-only-storage" },
92
- "elementType": "$inputScalar"
93
- },
94
- {
95
- "name": "weight",
96
- "arg": "weightT",
97
- "semantic": "weight",
98
- "buffer": { "type": "read-only-storage" },
99
- "elementType": "$inputScalar"
100
- },
101
- {
102
- "name": "output",
103
- "arg": "outputT",
104
- "semantic": "output",
105
- "buffer": { "type": "storage" },
106
- "elementType": "$outputScalar"
107
- },
108
- {
109
- "name": "present_state",
110
- "arg": "presentStateT",
111
- "semantic": "present_state",
112
- "buffer": { "type": "storage" },
113
- "elementType": "$outputScalar"
114
- },
115
- {
116
- "name": "params",
117
- "semantic": "kernel.params",
118
- "buffer": { "type": "uniform" },
119
- "struct": {
120
- "name": "Params",
121
- "fields": [
122
- { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" },
123
- { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" },
124
- { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" },
125
- { "name": "kernelSize", "type": "u32", "value": "kernelSize" },
126
- { "name": "stateWindow", "type": "u32", "value": "stateWindow" },
127
- { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" }
128
- ]
129
- }
130
- }
131
- ],
132
- "zeroVec4": [
133
- {
134
- "name": "input",
135
- "arg": "inputT",
136
- "semantic": "input",
137
- "buffer": { "type": "read-only-storage" },
138
- "elementType": "$inputVec4"
139
- },
140
- {
141
- "name": "weight",
142
- "arg": "weightT",
143
- "semantic": "weight",
144
- "buffer": { "type": "read-only-storage" },
145
- "elementType": "$weightElem"
146
- },
147
- {
148
- "name": "output",
149
- "arg": "outputT",
150
- "semantic": "output",
151
- "buffer": { "type": "storage" },
152
- "elementType": "$outputVec4"
153
- },
154
- {
155
- "name": "present_state",
156
- "arg": "presentStateT",
157
- "semantic": "present_state",
158
- "buffer": { "type": "storage" },
159
- "elementType": "$outputScalar"
160
- },
161
- {
162
- "name": "params",
163
- "semantic": "kernel.params",
164
- "buffer": { "type": "uniform" },
165
- "struct": {
166
- "name": "Params",
167
- "fields": [
168
- { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" },
169
- { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" },
170
- { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" },
171
- { "name": "stateWindow", "type": "u32", "value": "stateWindow" },
172
- { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" }
173
- ]
174
- }
175
- }
176
- ],
177
- "biasNoState": [
178
- {
179
- "name": "input",
180
- "arg": "inputT",
181
- "semantic": "input",
182
- "buffer": { "type": "read-only-storage" },
183
- "elementType": "$inputScalar"
184
- },
185
- {
186
- "name": "weight",
187
- "arg": "weightT",
188
- "semantic": "weight",
189
- "buffer": { "type": "read-only-storage" },
190
- "elementType": "$inputScalar"
191
- },
192
- {
193
- "name": "bias",
194
- "arg": "biasT",
195
- "semantic": "bias",
196
- "buffer": { "type": "read-only-storage" },
197
- "elementType": "$inputScalar"
198
- },
199
- {
200
- "name": "output",
201
- "arg": "outputT",
202
- "semantic": "output",
203
- "buffer": { "type": "storage" },
204
- "elementType": "$outputScalar"
205
- },
206
- {
207
- "name": "present_state",
208
- "arg": "presentStateT",
209
- "semantic": "present_state",
210
- "buffer": { "type": "storage" },
211
- "elementType": "$outputScalar"
212
- },
213
- {
214
- "name": "params",
215
- "semantic": "kernel.params",
216
- "buffer": { "type": "uniform" },
217
- "struct": {
218
- "name": "Params",
219
- "fields": [
220
- { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" },
221
- { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" },
222
- { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" },
223
- { "name": "kernelSize", "type": "u32", "value": "kernelSize" },
224
- { "name": "stateWindow", "type": "u32", "value": "stateWindow" },
225
- { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" }
226
- ]
227
- }
228
- }
229
- ],
230
- "stateNoBias": [
231
- {
232
- "name": "input",
233
- "arg": "inputT",
234
- "semantic": "input",
235
- "buffer": { "type": "read-only-storage" },
236
- "elementType": "$inputScalar"
237
- },
238
- {
239
- "name": "weight",
240
- "arg": "weightT",
241
- "semantic": "weight",
242
- "buffer": { "type": "read-only-storage" },
243
- "elementType": "$inputScalar"
244
- },
245
- {
246
- "name": "past_state",
247
- "arg": "pastStateT",
248
- "semantic": "past_state",
249
- "buffer": { "type": "read-only-storage" },
250
- "elementType": "$inputScalar"
251
- },
252
- {
253
- "name": "output",
254
- "arg": "outputT",
255
- "semantic": "output",
256
- "buffer": { "type": "storage" },
257
- "elementType": "$outputScalar"
258
- },
259
- {
260
- "name": "present_state",
261
- "arg": "presentStateT",
262
- "semantic": "present_state",
263
- "buffer": { "type": "storage" },
264
- "elementType": "$outputScalar"
265
- },
266
- {
267
- "name": "params",
268
- "semantic": "kernel.params",
269
- "buffer": { "type": "uniform" },
270
- "struct": {
271
- "name": "Params",
272
- "fields": [
273
- { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" },
274
- { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" },
275
- { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" },
276
- { "name": "kernelSize", "type": "u32", "value": "kernelSize" },
277
- { "name": "stateWindow", "type": "u32", "value": "stateWindow" },
278
- { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" }
279
- ]
280
- }
281
- }
282
- ],
283
- "stateBias": [
284
- {
285
- "name": "input",
286
- "arg": "inputT",
287
- "semantic": "input",
288
- "buffer": { "type": "read-only-storage" },
289
- "elementType": "$inputScalar"
290
- },
291
- {
292
- "name": "weight",
293
- "arg": "weightT",
294
- "semantic": "weight",
295
- "buffer": { "type": "read-only-storage" },
296
- "elementType": "$inputScalar"
297
- },
298
- {
299
- "name": "bias",
300
- "arg": "biasT",
301
- "semantic": "bias",
302
- "buffer": { "type": "read-only-storage" },
303
- "elementType": "$inputScalar"
304
- },
305
- {
306
- "name": "past_state",
307
- "arg": "pastStateT",
308
- "semantic": "past_state",
309
- "buffer": { "type": "read-only-storage" },
310
- "elementType": "$inputScalar"
311
- },
312
- {
313
- "name": "output",
314
- "arg": "outputT",
315
- "semantic": "output",
316
- "buffer": { "type": "storage" },
317
- "elementType": "$outputScalar"
318
- },
319
- {
320
- "name": "present_state",
321
- "arg": "presentStateT",
322
- "semantic": "present_state",
323
- "buffer": { "type": "storage" },
324
- "elementType": "$outputScalar"
325
- },
326
- {
327
- "name": "params",
328
- "semantic": "kernel.params",
329
- "buffer": { "type": "uniform" },
330
- "struct": {
331
- "name": "Params",
332
- "fields": [
333
- { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" },
334
- { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" },
335
- { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" },
336
- { "name": "kernelSize", "type": "u32", "value": "kernelSize" },
337
- { "name": "stateWindow", "type": "u32", "value": "stateWindow" },
338
- { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" }
339
- ]
340
- }
341
- }
342
- ],
343
- "zeroScalarIo": [
344
- {
345
- "name": "input",
346
- "arg": "inputT",
347
- "semantic": "input",
348
- "buffer": { "type": "read-only-storage" },
349
- "elementType": "$inputScalar"
350
- },
351
- {
352
- "name": "weight",
353
- "arg": "weightT",
354
- "semantic": "weight",
355
- "buffer": { "type": "read-only-storage" },
356
- "elementType": "$inputScalar"
357
- },
358
- {
359
- "name": "output",
360
- "arg": "outputT",
361
- "semantic": "output",
362
- "buffer": { "type": "storage" },
363
- "elementType": "$outputScalar"
364
- },
365
- {
366
- "name": "present_state",
367
- "arg": "presentStateT",
368
- "semantic": "present_state",
369
- "buffer": { "type": "storage" },
370
- "elementType": "$outputScalar"
371
- }
372
- ],
373
- "biasNoStateIo": [
374
- {
375
- "name": "input",
376
- "arg": "inputT",
377
- "semantic": "input",
378
- "buffer": { "type": "read-only-storage" },
379
- "elementType": "$inputScalar"
380
- },
381
- {
382
- "name": "weight",
383
- "arg": "weightT",
384
- "semantic": "weight",
385
- "buffer": { "type": "read-only-storage" },
386
- "elementType": "$inputScalar"
387
- },
388
- {
389
- "name": "bias",
390
- "arg": "biasT",
391
- "semantic": "bias",
392
- "buffer": { "type": "read-only-storage" },
393
- "elementType": "$inputScalar"
394
- },
395
- {
396
- "name": "output",
397
- "arg": "outputT",
398
- "semantic": "output",
399
- "buffer": { "type": "storage" },
400
- "elementType": "$outputScalar"
401
- },
402
- {
403
- "name": "present_state",
404
- "arg": "presentStateT",
405
- "semantic": "present_state",
406
- "buffer": { "type": "storage" },
407
- "elementType": "$outputScalar"
408
- }
409
- ],
410
- "stateNoBiasIo": [
411
- {
412
- "name": "input",
413
- "arg": "inputT",
414
- "semantic": "input",
415
- "buffer": { "type": "read-only-storage" },
416
- "elementType": "$inputScalar"
417
- },
418
- {
419
- "name": "weight",
420
- "arg": "weightT",
421
- "semantic": "weight",
422
- "buffer": { "type": "read-only-storage" },
423
- "elementType": "$inputScalar"
424
- },
425
- {
426
- "name": "past_state",
427
- "arg": "pastStateT",
428
- "semantic": "past_state",
429
- "buffer": { "type": "read-only-storage" },
430
- "elementType": "$inputScalar"
431
- },
432
- {
433
- "name": "output",
434
- "arg": "outputT",
435
- "semantic": "output",
436
- "buffer": { "type": "storage" },
437
- "elementType": "$outputScalar"
438
- },
439
- {
440
- "name": "present_state",
441
- "arg": "presentStateT",
442
- "semantic": "present_state",
443
- "buffer": { "type": "storage" },
444
- "elementType": "$outputScalar"
445
- }
446
- ],
447
- "stateBiasIo": [
448
- {
449
- "name": "input",
450
- "arg": "inputT",
451
- "semantic": "input",
452
- "buffer": { "type": "read-only-storage" },
453
- "elementType": "$inputScalar"
454
- },
455
- {
456
- "name": "weight",
457
- "arg": "weightT",
458
- "semantic": "weight",
459
- "buffer": { "type": "read-only-storage" },
460
- "elementType": "$inputScalar"
461
- },
462
- {
463
- "name": "bias",
464
- "arg": "biasT",
465
- "semantic": "bias",
466
- "buffer": { "type": "read-only-storage" },
467
- "elementType": "$inputScalar"
468
- },
469
- {
470
- "name": "past_state",
471
- "arg": "pastStateT",
472
- "semantic": "past_state",
473
- "buffer": { "type": "read-only-storage" },
474
- "elementType": "$inputScalar"
475
- },
476
- {
477
- "name": "output",
478
- "arg": "outputT",
479
- "semantic": "output",
480
- "buffer": { "type": "storage" },
481
- "elementType": "$outputScalar"
482
- },
483
- {
484
- "name": "present_state",
485
- "arg": "presentStateT",
486
- "semantic": "present_state",
487
- "buffer": { "type": "storage" },
488
- "elementType": "$outputScalar"
489
- }
490
- ],
491
- "zeroTiled": [
492
- {
493
- "name": "input",
494
- "arg": "inputT",
495
- "semantic": "input",
496
- "buffer": { "type": "read-only-storage" },
497
- "elementType": "$inputScalar"
498
- },
499
- {
500
- "name": "weight",
501
- "arg": "weightT",
502
- "semantic": "weight",
503
- "buffer": { "type": "read-only-storage" },
504
- "elementType": "$inputScalar"
505
- },
506
- {
507
- "name": "output",
508
- "arg": "outputT",
509
- "semantic": "output",
510
- "buffer": { "type": "storage" },
511
- "elementType": "$outputScalar"
512
- },
513
- {
514
- "name": "present_state",
515
- "arg": "presentStateT",
516
- "semantic": "present_state",
517
- "buffer": { "type": "storage" },
518
- "elementType": "$outputScalar"
519
- },
520
- {
521
- "name": "params",
522
- "semantic": "kernel.params",
523
- "buffer": { "type": "uniform" },
524
- "struct": {
525
- "name": "Params",
526
- "fields": [
527
- { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" },
528
- { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" },
529
- { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" },
530
- { "name": "stateWindow", "type": "u32", "value": "stateWindow" },
531
- { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" }
532
- ]
533
- }
534
- }
535
- ],
536
- "biasNoStateTiled": [
537
- {
538
- "name": "input",
539
- "arg": "inputT",
540
- "semantic": "input",
541
- "buffer": { "type": "read-only-storage" },
542
- "elementType": "$inputScalar"
543
- },
544
- {
545
- "name": "weight",
546
- "arg": "weightT",
547
- "semantic": "weight",
548
- "buffer": { "type": "read-only-storage" },
549
- "elementType": "$inputScalar"
550
- },
551
- {
552
- "name": "bias",
553
- "arg": "biasT",
554
- "semantic": "bias",
555
- "buffer": { "type": "read-only-storage" },
556
- "elementType": "$inputScalar"
557
- },
558
- {
559
- "name": "output",
560
- "arg": "outputT",
561
- "semantic": "output",
562
- "buffer": { "type": "storage" },
563
- "elementType": "$outputScalar"
564
- },
565
- {
566
- "name": "present_state",
567
- "arg": "presentStateT",
568
- "semantic": "present_state",
569
- "buffer": { "type": "storage" },
570
- "elementType": "$outputScalar"
571
- },
572
- {
573
- "name": "params",
574
- "semantic": "kernel.params",
575
- "buffer": { "type": "uniform" },
576
- "struct": {
577
- "name": "Params",
578
- "fields": [
579
- { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" },
580
- { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" },
581
- { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" },
582
- { "name": "stateWindow", "type": "u32", "value": "stateWindow" },
583
- { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" }
584
- ]
585
- }
586
- }
587
- ],
588
- "stateNoBiasTiled": [
589
- {
590
- "name": "input",
591
- "arg": "inputT",
592
- "semantic": "input",
593
- "buffer": { "type": "read-only-storage" },
594
- "elementType": "$inputScalar"
595
- },
596
- {
597
- "name": "weight",
598
- "arg": "weightT",
599
- "semantic": "weight",
600
- "buffer": { "type": "read-only-storage" },
601
- "elementType": "$inputScalar"
602
- },
603
- {
604
- "name": "past_state",
605
- "arg": "pastStateT",
606
- "semantic": "past_state",
607
- "buffer": { "type": "read-only-storage" },
608
- "elementType": "$inputScalar"
609
- },
610
- {
611
- "name": "output",
612
- "arg": "outputT",
613
- "semantic": "output",
614
- "buffer": { "type": "storage" },
615
- "elementType": "$outputScalar"
616
- },
617
- {
618
- "name": "present_state",
619
- "arg": "presentStateT",
620
- "semantic": "present_state",
621
- "buffer": { "type": "storage" },
622
- "elementType": "$outputScalar"
623
- },
624
- {
625
- "name": "params",
626
- "semantic": "kernel.params",
627
- "buffer": { "type": "uniform" },
628
- "struct": {
629
- "name": "Params",
630
- "fields": [
631
- { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" },
632
- { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" },
633
- { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" },
634
- { "name": "stateWindow", "type": "u32", "value": "stateWindow" },
635
- { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" }
636
- ]
637
- }
638
- }
639
- ],
640
- "stateBiasTiled": [
641
- {
642
- "name": "input",
643
- "arg": "inputT",
644
- "semantic": "input",
645
- "buffer": { "type": "read-only-storage" },
646
- "elementType": "$inputScalar"
647
- },
648
- {
649
- "name": "weight",
650
- "arg": "weightT",
651
- "semantic": "weight",
652
- "buffer": { "type": "read-only-storage" },
653
- "elementType": "$inputScalar"
654
- },
655
- {
656
- "name": "bias",
657
- "arg": "biasT",
658
- "semantic": "bias",
659
- "buffer": { "type": "read-only-storage" },
660
- "elementType": "$inputScalar"
661
- },
662
- {
663
- "name": "past_state",
664
- "arg": "pastStateT",
665
- "semantic": "past_state",
666
- "buffer": { "type": "read-only-storage" },
667
- "elementType": "$inputScalar"
668
- },
669
- {
670
- "name": "output",
671
- "arg": "outputT",
672
- "semantic": "output",
673
- "buffer": { "type": "storage" },
674
- "elementType": "$outputScalar"
675
- },
676
- {
677
- "name": "present_state",
678
- "arg": "presentStateT",
679
- "semantic": "present_state",
680
- "buffer": { "type": "storage" },
681
- "elementType": "$outputScalar"
682
- },
683
- {
684
- "name": "params",
685
- "semantic": "kernel.params",
686
- "buffer": { "type": "uniform" },
687
- "struct": {
688
- "name": "Params",
689
- "fields": [
690
- { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" },
691
- { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" },
692
- { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" },
693
- { "name": "stateWindow", "type": "u32", "value": "stateWindow" },
694
- { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" }
695
- ]
696
- }
697
- }
698
- ],
699
- "biasNoStateVec4": [
700
- {
701
- "name": "input",
702
- "arg": "inputT",
703
- "semantic": "input",
704
- "buffer": { "type": "read-only-storage" },
705
- "elementType": "$inputVec4"
706
- },
707
- {
708
- "name": "weight",
709
- "arg": "weightT",
710
- "semantic": "weight",
711
- "buffer": { "type": "read-only-storage" },
712
- "elementType": "$weightElem"
713
- },
714
- {
715
- "name": "bias",
716
- "arg": "biasT",
717
- "semantic": "bias",
718
- "buffer": { "type": "read-only-storage" },
719
- "elementType": "$inputScalar"
720
- },
721
- {
722
- "name": "output",
723
- "arg": "outputT",
724
- "semantic": "output",
725
- "buffer": { "type": "storage" },
726
- "elementType": "$outputVec4"
727
- },
728
- {
729
- "name": "present_state",
730
- "arg": "presentStateT",
731
- "semantic": "present_state",
732
- "buffer": { "type": "storage" },
733
- "elementType": "$outputScalar"
734
- },
735
- {
736
- "name": "params",
737
- "semantic": "kernel.params",
738
- "buffer": { "type": "uniform" },
739
- "struct": {
740
- "name": "Params",
741
- "fields": [
742
- { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" },
743
- { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" },
744
- { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" },
745
- { "name": "stateWindow", "type": "u32", "value": "stateWindow" },
746
- { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" }
747
- ]
748
- }
749
- }
750
- ],
751
- "stateNoBiasVec4": [
752
- {
753
- "name": "input",
754
- "arg": "inputT",
755
- "semantic": "input",
756
- "buffer": { "type": "read-only-storage" },
757
- "elementType": "$inputVec4"
758
- },
759
- {
760
- "name": "weight",
761
- "arg": "weightT",
762
- "semantic": "weight",
763
- "buffer": { "type": "read-only-storage" },
764
- "elementType": "$weightElem"
765
- },
766
- {
767
- "name": "past_state",
768
- "arg": "pastStateT",
769
- "semantic": "past_state",
770
- "buffer": { "type": "read-only-storage" },
771
- "elementType": "$inputScalar"
772
- },
773
- {
774
- "name": "output",
775
- "arg": "outputT",
776
- "semantic": "output",
777
- "buffer": { "type": "storage" },
778
- "elementType": "$outputVec4"
779
- },
780
- {
781
- "name": "present_state",
782
- "arg": "presentStateT",
783
- "semantic": "present_state",
784
- "buffer": { "type": "storage" },
785
- "elementType": "$outputScalar"
786
- },
787
- {
788
- "name": "params",
789
- "semantic": "kernel.params",
790
- "buffer": { "type": "uniform" },
791
- "struct": {
792
- "name": "Params",
793
- "fields": [
794
- { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" },
795
- { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" },
796
- { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" },
797
- { "name": "stateWindow", "type": "u32", "value": "stateWindow" },
798
- { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" }
799
- ]
800
- }
801
- }
802
- ],
803
- "stateBiasVec4": [
804
- {
805
- "name": "input",
806
- "arg": "inputT",
807
- "semantic": "input",
808
- "buffer": { "type": "read-only-storage" },
809
- "elementType": "$inputVec4"
810
- },
811
- {
812
- "name": "weight",
813
- "arg": "weightT",
814
- "semantic": "weight",
815
- "buffer": { "type": "read-only-storage" },
816
- "elementType": "$weightElem"
817
- },
818
- {
819
- "name": "bias",
820
- "arg": "biasT",
821
- "semantic": "bias",
822
- "buffer": { "type": "read-only-storage" },
823
- "elementType": "$inputScalar"
824
- },
825
- {
826
- "name": "past_state",
827
- "arg": "pastStateT",
828
- "semantic": "past_state",
829
- "buffer": { "type": "read-only-storage" },
830
- "elementType": "$inputScalar"
831
- },
832
- {
833
- "name": "output",
834
- "arg": "outputT",
835
- "semantic": "output",
836
- "buffer": { "type": "storage" },
837
- "elementType": "$outputVec4"
838
- },
839
- {
840
- "name": "present_state",
841
- "arg": "presentStateT",
842
- "semantic": "present_state",
843
- "buffer": { "type": "storage" },
844
- "elementType": "$outputScalar"
845
- },
846
- {
847
- "name": "params",
848
- "semantic": "kernel.params",
849
- "buffer": { "type": "uniform" },
850
- "struct": {
851
- "name": "Params",
852
- "fields": [
853
- { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" },
854
- { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" },
855
- { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" },
856
- { "name": "stateWindow", "type": "u32", "value": "stateWindow" },
857
- { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" }
858
- ]
859
- }
860
- }
861
- ]
862
  },
863
  "variants": [
864
  {
865
  "id": "zero_state_vec4",
866
  "priority": 20,
867
  "when": ["zeroStateContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"],
868
- "constants": {
869
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
870
  "workgroupSize": 256,
871
  "hasStateWindow": "windowed",
@@ -876,22 +93,19 @@
876
  "outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
877
  "hasBias": false,
878
  "hasState": false,
879
- "kernelSize": "kernelSize",
880
- "kernelSizePadded": "kernelSizePadded",
881
  "weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T"
882
  },
883
  "passes": [
884
  {
885
  "id": "main",
886
  "name": "CausalConvWithState.Vec4",
887
- "source": {
888
- "shader": "causal-conv-with-state-vec4.wgsl.jinja",
889
- "inputs": { "materializeConvBeforeActivation": false }
890
- },
891
- "bindings": "zeroVec4",
892
  "dispatch": {
893
- "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)",
894
- "workgroupSize": "constants.workgroupSize"
 
895
  }
896
  }
897
  ]
@@ -900,7 +114,7 @@
900
  "id": "state_bias_vec4",
901
  "priority": 20,
902
  "when": ["stateBiasContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"],
903
- "constants": {
904
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
905
  "workgroupSize": 256,
906
  "hasStateWindow": "windowed",
@@ -911,22 +125,19 @@
911
  "outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
912
  "hasBias": true,
913
  "hasState": true,
914
- "kernelSize": "kernelSize",
915
- "kernelSizePadded": "kernelSizePadded",
916
  "weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T"
917
  },
918
  "passes": [
919
  {
920
  "id": "main",
921
  "name": "CausalConvWithState.Vec4",
922
- "source": {
923
- "shader": "causal-conv-with-state-vec4.wgsl.jinja",
924
- "inputs": { "materializeConvBeforeActivation": false }
925
- },
926
- "bindings": "stateBiasVec4",
927
  "dispatch": {
928
- "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)",
929
- "workgroupSize": "constants.workgroupSize"
 
930
  }
931
  }
932
  ]
@@ -935,7 +146,7 @@
935
  "id": "bias_no_state_vec4",
936
  "priority": 20,
937
  "when": ["biasNoStateContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"],
938
- "constants": {
939
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
940
  "workgroupSize": 256,
941
  "hasStateWindow": "windowed",
@@ -946,22 +157,19 @@
946
  "outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
947
  "hasBias": true,
948
  "hasState": false,
949
- "kernelSize": "kernelSize",
950
- "kernelSizePadded": "kernelSizePadded",
951
  "weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T"
952
  },
953
  "passes": [
954
  {
955
  "id": "main",
956
  "name": "CausalConvWithState.Vec4",
957
- "source": {
958
- "shader": "causal-conv-with-state-vec4.wgsl.jinja",
959
- "inputs": { "materializeConvBeforeActivation": false }
960
- },
961
- "bindings": "biasNoStateVec4",
962
  "dispatch": {
963
- "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)",
964
- "workgroupSize": "constants.workgroupSize"
 
965
  }
966
  }
967
  ]
@@ -970,7 +178,7 @@
970
  "id": "state_no_bias_vec4",
971
  "priority": 20,
972
  "when": ["stateNoBiasContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"],
973
- "constants": {
974
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
975
  "workgroupSize": 256,
976
  "hasStateWindow": "windowed",
@@ -981,22 +189,19 @@
981
  "outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
982
  "hasBias": false,
983
  "hasState": true,
984
- "kernelSize": "kernelSize",
985
- "kernelSizePadded": "kernelSizePadded",
986
  "weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T"
987
  },
988
  "passes": [
989
  {
990
  "id": "main",
991
  "name": "CausalConvWithState.Vec4",
992
- "source": {
993
- "shader": "causal-conv-with-state-vec4.wgsl.jinja",
994
- "inputs": { "materializeConvBeforeActivation": false }
995
- },
996
- "bindings": "stateNoBiasVec4",
997
  "dispatch": {
998
- "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)",
999
- "workgroupSize": "constants.workgroupSize"
 
1000
  }
1001
  }
1002
  ]
@@ -1005,7 +210,7 @@
1005
  "id": "zero_state_tiled_large_kernel",
1006
  "priority": 10,
1007
  "when": ["zeroStateContract", "kernelSize >= 32", "dim(shapes.inputT, 2) >= 256", "dim(shapes.inputT, 2) % 8 == 0", "tunables.tiledWorkgroupSize >= 1", "floor(tunables.tiledWorkgroupSize) == tunables.tiledWorkgroupSize", "tunables.tiledWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.tiledWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "(tunables.tiledWorkgroupSize * 8 + 2 * kernelSizePadded - 1) * 4 <= device.limits.maxComputeWorkgroupStorageSize"],
1008
- "constants": {
1009
  "hasBias": false,
1010
  "hasState": false,
1011
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
@@ -1013,8 +218,6 @@
1013
  "outputScalar": "dtypes.T",
1014
  "workgroupSize": "tunables.tiledWorkgroupSize",
1015
  "tileSize": "tunables.tiledWorkgroupSize * 8",
1016
- "kernelSize": "kernelSize",
1017
- "kernelSizePadded": "kernelSizePadded",
1018
  "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1",
1019
  "hasStateWindow": "windowed",
1020
  "usesF16": "tensorDtypes.inputT == \"float16\""
@@ -1023,13 +226,13 @@
1023
  {
1024
  "id": "main",
1025
  "name": "CausalConvWithState.TiledLargeKernel",
1026
- "source": {
1027
- "shader": "causal-conv-with-state-tiled.wgsl.jinja",
1028
- "inputs": { "materializeConvBeforeActivation": false }
1029
- },
1030
- "bindings": "zeroTiled",
1031
  "dispatch": {
1032
- "workgroups": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), constants.tileSize)"
 
 
1033
  }
1034
  }
1035
  ]
@@ -1038,7 +241,7 @@
1038
  "id": "state_bias_tiled_large_kernel",
1039
  "priority": 10,
1040
  "when": ["stateBiasContract", "kernelSize >= 32", "dim(shapes.inputT, 2) >= 256", "dim(shapes.inputT, 2) % 8 == 0", "tunables.tiledWorkgroupSize >= 1", "floor(tunables.tiledWorkgroupSize) == tunables.tiledWorkgroupSize", "tunables.tiledWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.tiledWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "(tunables.tiledWorkgroupSize * 8 + 2 * kernelSizePadded - 1) * 4 <= device.limits.maxComputeWorkgroupStorageSize"],
1041
- "constants": {
1042
  "hasBias": true,
1043
  "hasState": true,
1044
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
@@ -1046,8 +249,6 @@
1046
  "outputScalar": "dtypes.T",
1047
  "workgroupSize": "tunables.tiledWorkgroupSize",
1048
  "tileSize": "tunables.tiledWorkgroupSize * 8",
1049
- "kernelSize": "kernelSize",
1050
- "kernelSizePadded": "kernelSizePadded",
1051
  "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1",
1052
  "hasStateWindow": "windowed",
1053
  "usesF16": "tensorDtypes.inputT == \"float16\""
@@ -1056,13 +257,13 @@
1056
  {
1057
  "id": "main",
1058
  "name": "CausalConvWithState.TiledLargeKernel",
1059
- "source": {
1060
- "shader": "causal-conv-with-state-tiled.wgsl.jinja",
1061
- "inputs": { "materializeConvBeforeActivation": false }
1062
- },
1063
- "bindings": "stateBiasTiled",
1064
  "dispatch": {
1065
- "workgroups": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), constants.tileSize)"
 
 
1066
  }
1067
  }
1068
  ]
@@ -1071,7 +272,7 @@
1071
  "id": "bias_no_state_tiled_large_kernel",
1072
  "priority": 10,
1073
  "when": ["biasNoStateContract", "kernelSize >= 32", "dim(shapes.inputT, 2) >= 256", "dim(shapes.inputT, 2) % 8 == 0", "tunables.tiledWorkgroupSize >= 1", "floor(tunables.tiledWorkgroupSize) == tunables.tiledWorkgroupSize", "tunables.tiledWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.tiledWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "(tunables.tiledWorkgroupSize * 8 + 2 * kernelSizePadded - 1) * 4 <= device.limits.maxComputeWorkgroupStorageSize"],
1074
- "constants": {
1075
  "hasBias": true,
1076
  "hasState": false,
1077
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
@@ -1079,8 +280,6 @@
1079
  "outputScalar": "dtypes.T",
1080
  "workgroupSize": "tunables.tiledWorkgroupSize",
1081
  "tileSize": "tunables.tiledWorkgroupSize * 8",
1082
- "kernelSize": "kernelSize",
1083
- "kernelSizePadded": "kernelSizePadded",
1084
  "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1",
1085
  "hasStateWindow": "windowed",
1086
  "usesF16": "tensorDtypes.inputT == \"float16\""
@@ -1089,13 +288,13 @@
1089
  {
1090
  "id": "main",
1091
  "name": "CausalConvWithState.TiledLargeKernel",
1092
- "source": {
1093
- "shader": "causal-conv-with-state-tiled.wgsl.jinja",
1094
- "inputs": { "materializeConvBeforeActivation": false }
1095
- },
1096
- "bindings": "biasNoStateTiled",
1097
  "dispatch": {
1098
- "workgroups": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), constants.tileSize)"
 
 
1099
  }
1100
  }
1101
  ]
@@ -1104,7 +303,7 @@
1104
  "id": "state_no_bias_tiled_large_kernel",
1105
  "priority": 10,
1106
  "when": ["stateNoBiasContract", "kernelSize >= 32", "dim(shapes.inputT, 2) >= 256", "dim(shapes.inputT, 2) % 8 == 0", "tunables.tiledWorkgroupSize >= 1", "floor(tunables.tiledWorkgroupSize) == tunables.tiledWorkgroupSize", "tunables.tiledWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.tiledWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "(tunables.tiledWorkgroupSize * 8 + 2 * kernelSizePadded - 1) * 4 <= device.limits.maxComputeWorkgroupStorageSize"],
1107
- "constants": {
1108
  "hasBias": false,
1109
  "hasState": true,
1110
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
@@ -1112,8 +311,6 @@
1112
  "outputScalar": "dtypes.T",
1113
  "workgroupSize": "tunables.tiledWorkgroupSize",
1114
  "tileSize": "tunables.tiledWorkgroupSize * 8",
1115
- "kernelSize": "kernelSize",
1116
- "kernelSizePadded": "kernelSizePadded",
1117
  "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1",
1118
  "hasStateWindow": "windowed",
1119
  "usesF16": "tensorDtypes.inputT == \"float16\""
@@ -1122,13 +319,13 @@
1122
  {
1123
  "id": "main",
1124
  "name": "CausalConvWithState.TiledLargeKernel",
1125
- "source": {
1126
- "shader": "causal-conv-with-state-tiled.wgsl.jinja",
1127
- "inputs": { "materializeConvBeforeActivation": false }
1128
- },
1129
- "bindings": "stateNoBiasTiled",
1130
  "dispatch": {
1131
- "workgroups": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), constants.tileSize)"
 
 
1132
  }
1133
  }
1134
  ]
@@ -1137,7 +334,7 @@
1137
  "id": "zero_state",
1138
  "priority": 0,
1139
  "when": ["zeroStateContract"],
1140
- "constants": {
1141
  "hasBias": false,
1142
  "hasState": false,
1143
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
@@ -1151,14 +348,13 @@
1151
  {
1152
  "id": "main",
1153
  "name": "CausalConvWithState",
1154
- "source": {
1155
- "shader": "causal-conv-with-state.wgsl.jinja",
1156
- "inputs": { "materializeConvBeforeActivation": false }
1157
- },
1158
- "bindings": "zeroScalar",
1159
  "dispatch": {
1160
- "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))",
1161
- "workgroupSize": "constants.workgroupSize"
 
1162
  }
1163
  }
1164
  ]
@@ -1167,7 +363,7 @@
1167
  "id": "state_bias",
1168
  "priority": 0,
1169
  "when": ["stateBiasContract"],
1170
- "constants": {
1171
  "hasBias": true,
1172
  "hasState": true,
1173
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
@@ -1181,14 +377,13 @@
1181
  {
1182
  "id": "main",
1183
  "name": "CausalConvWithState",
1184
- "source": {
1185
- "shader": "causal-conv-with-state.wgsl.jinja",
1186
- "inputs": { "materializeConvBeforeActivation": false }
1187
- },
1188
- "bindings": "stateBias",
1189
  "dispatch": {
1190
- "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))",
1191
- "workgroupSize": "constants.workgroupSize"
 
1192
  }
1193
  }
1194
  ]
@@ -1197,7 +392,7 @@
1197
  "id": "bias_no_state",
1198
  "priority": 0,
1199
  "when": ["biasNoStateContract"],
1200
- "constants": {
1201
  "hasBias": true,
1202
  "hasState": false,
1203
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
@@ -1211,14 +406,13 @@
1211
  {
1212
  "id": "main",
1213
  "name": "CausalConvWithState",
1214
- "source": {
1215
- "shader": "causal-conv-with-state.wgsl.jinja",
1216
- "inputs": { "materializeConvBeforeActivation": false }
1217
- },
1218
- "bindings": "biasNoState",
1219
  "dispatch": {
1220
- "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))",
1221
- "workgroupSize": "constants.workgroupSize"
 
1222
  }
1223
  }
1224
  ]
@@ -1227,7 +421,7 @@
1227
  "id": "state_no_bias",
1228
  "priority": 0,
1229
  "when": ["stateNoBiasContract"],
1230
- "constants": {
1231
  "hasBias": false,
1232
  "hasState": true,
1233
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
@@ -1241,14 +435,13 @@
1241
  {
1242
  "id": "main",
1243
  "name": "CausalConvWithState",
1244
- "source": {
1245
- "shader": "causal-conv-with-state.wgsl.jinja",
1246
- "inputs": { "materializeConvBeforeActivation": false }
1247
- },
1248
- "bindings": "stateNoBias",
1249
  "dispatch": {
1250
- "threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))",
1251
- "workgroupSize": "constants.workgroupSize"
 
1252
  }
1253
  }
1254
  ]
 
2
  "domain": "com.microsoft",
3
  "name": "CausalConvWithState",
4
  "sinceVersion": 1,
5
+ "inputs": {
6
+ "inputT": { "onnx": "input", "dtype": "T", "rank": 3 },
7
+ "weightT": { "onnx": "weight", "dtype": "T", "rank": 3 },
8
+ "biasT": { "onnx": "bias", "dtype": "T", "rank": 1, "optional": true },
9
+ "pastStateT": {
10
+ "onnx": "past_state",
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
11
  "dtype": "T",
12
  "rank": "3 if attrs.state_window == 0 else 4",
13
+ "optional": true
 
14
  }
15
+ },
16
+ "outputs": {
17
+ "outputT": { "onnx": "output", "dtype": "T", "rank": 3, "shape": "shapes.inputT" },
18
+ "presentStateT": {
19
+ "onnx": "present_state",
 
 
 
 
 
 
20
  "dtype": "T",
21
  "rank": "3 if attrs.state_window == 0 else 4",
22
+ "shape": "[dim(shapes.inputT, 0), dim(shapes.inputT, 1), dim(shapes.weightT, 2) - 1] if attrs.state_window == 0 else [attrs.state_window, dim(shapes.inputT, 0), dim(shapes.inputT, 1), dim(shapes.weightT, 2) - 1]"
 
23
  }
 
 
 
 
 
 
 
24
  },
25
+ "attributes": { "activation": { "default": "none" }, "ndim": { "default": 1 }, "state_window": { "default": 0 } },
26
+ "attributeConstraints": { "activation": { "values": ["none", "silu", "swish"] }, "ndim": { "values": [1] } },
27
  "typeConstraints": { "T": ["float32", "float16"] },
28
+ "tunables": { "workgroupSize": { "default": 256 }, "tiledWorkgroupSize": { "default": 128 } },
 
 
 
 
 
 
 
 
29
  "derive": {
30
  "stateWindow": "attrs.state_window",
31
  "windowed": "stateWindow > 0",
 
44
  "stateNoBiasContract": "commonContract and present.pastStateT and not present.biasT and tensorDtypes.pastStateT == tensorDtypes.inputT and pastStateShapeOk",
45
  "stateBiasContract": "commonContract and present.pastStateT and present.biasT and ranks.biasT == 1 and tensorDtypes.pastStateT == tensorDtypes.inputT and tensorDtypes.biasT == tensorDtypes.inputT and pastStateShapeOk and dim(shapes.biasT, 0) == dim(shapes.inputT, 1)"
46
  },
47
+ "bindings": {
48
+ "input": { "arg": "inputT", "buffer": "read-only-storage", "elementType": "$inputVec4" },
49
+ "weight": { "arg": "weightT", "buffer": "read-only-storage", "elementType": "$weightElem" },
50
+ "output": { "arg": "outputT", "buffer": "storage", "elementType": "$outputVec4" },
51
+ "present_state": { "arg": "presentStateT", "buffer": "storage", "elementType": "$outputScalar" },
52
+ "params": {
53
+ "buffer": "uniform",
54
+ "struct": [
55
+ { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" },
56
+ { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" },
57
+ { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" },
58
+ { "name": "stateWindow", "type": "u32", "value": "stateWindow" },
59
+ { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" }
60
+ ]
61
+ },
62
+ "bias": { "arg": "biasT", "buffer": "read-only-storage", "elementType": "$inputScalar" },
63
+ "past_state": { "arg": "pastStateT", "buffer": "read-only-storage", "elementType": "$inputScalar" },
64
+ "input_2": { "arg": "inputT", "name": "input", "buffer": "read-only-storage", "elementType": "$inputScalar" },
65
+ "weight_2": { "arg": "weightT", "name": "weight", "buffer": "read-only-storage", "elementType": "$inputScalar" },
66
+ "output_2": { "arg": "outputT", "name": "output", "buffer": "storage", "elementType": "$outputScalar" },
67
+ "params_2": {
68
+ "name": "params",
69
+ "buffer": "uniform",
70
+ "struct": [
71
+ { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" },
72
+ { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" },
73
+ { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" },
74
+ { "name": "kernelSize", "type": "u32", "value": "kernelSize" },
75
+ { "name": "stateWindow", "type": "u32", "value": "stateWindow" },
76
+ { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" }
77
+ ]
78
+ }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
79
  },
80
  "variants": [
81
  {
82
  "id": "zero_state_vec4",
83
  "priority": 20,
84
  "when": ["zeroStateContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"],
85
+ "derive": {
86
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
87
  "workgroupSize": 256,
88
  "hasStateWindow": "windowed",
 
93
  "outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
94
  "hasBias": false,
95
  "hasState": false,
 
 
96
  "weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T"
97
  },
98
  "passes": [
99
  {
100
  "id": "main",
101
  "name": "CausalConvWithState.Vec4",
102
+ "shader": "causal-conv-with-state-vec4.wgsl.jinja",
103
+ "derive": { "materializeConvBeforeActivation": false },
104
+ "bindings": ["input", "weight", "output", "present_state", "params"],
 
 
105
  "dispatch": {
106
+ "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)",
107
+ "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)",
108
+ "z": 1
109
  }
110
  }
111
  ]
 
114
  "id": "state_bias_vec4",
115
  "priority": 20,
116
  "when": ["stateBiasContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"],
117
+ "derive": {
118
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
119
  "workgroupSize": 256,
120
  "hasStateWindow": "windowed",
 
125
  "outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
126
  "hasBias": true,
127
  "hasState": true,
 
 
128
  "weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T"
129
  },
130
  "passes": [
131
  {
132
  "id": "main",
133
  "name": "CausalConvWithState.Vec4",
134
+ "shader": "causal-conv-with-state-vec4.wgsl.jinja",
135
+ "derive": { "materializeConvBeforeActivation": false },
136
+ "bindings": ["input", "weight", "bias", "past_state", "output", "present_state", "params"],
 
 
137
  "dispatch": {
138
+ "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)",
139
+ "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)",
140
+ "z": 1
141
  }
142
  }
143
  ]
 
146
  "id": "bias_no_state_vec4",
147
  "priority": 20,
148
  "when": ["biasNoStateContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"],
149
+ "derive": {
150
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
151
  "workgroupSize": 256,
152
  "hasStateWindow": "windowed",
 
157
  "outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
158
  "hasBias": true,
159
  "hasState": false,
 
 
160
  "weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T"
161
  },
162
  "passes": [
163
  {
164
  "id": "main",
165
  "name": "CausalConvWithState.Vec4",
166
+ "shader": "causal-conv-with-state-vec4.wgsl.jinja",
167
+ "derive": { "materializeConvBeforeActivation": false },
168
+ "bindings": ["input", "weight", "bias", "output", "present_state", "params"],
 
 
169
  "dispatch": {
170
+ "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)",
171
+ "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)",
172
+ "z": 1
173
  }
174
  }
175
  ]
 
178
  "id": "state_no_bias_vec4",
179
  "priority": 20,
180
  "when": ["stateNoBiasContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"],
181
+ "derive": {
182
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
183
  "workgroupSize": 256,
184
  "hasStateWindow": "windowed",
 
189
  "outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
190
  "hasBias": false,
191
  "hasState": true,
 
 
192
  "weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T"
193
  },
194
  "passes": [
195
  {
196
  "id": "main",
197
  "name": "CausalConvWithState.Vec4",
198
+ "shader": "causal-conv-with-state-vec4.wgsl.jinja",
199
+ "derive": { "materializeConvBeforeActivation": false },
200
+ "bindings": ["input", "weight", "past_state", "output", "present_state", "params"],
 
 
201
  "dispatch": {
202
+ "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)",
203
+ "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)",
204
+ "z": 1
205
  }
206
  }
207
  ]
 
210
  "id": "zero_state_tiled_large_kernel",
211
  "priority": 10,
212
  "when": ["zeroStateContract", "kernelSize >= 32", "dim(shapes.inputT, 2) >= 256", "dim(shapes.inputT, 2) % 8 == 0", "tunables.tiledWorkgroupSize >= 1", "floor(tunables.tiledWorkgroupSize) == tunables.tiledWorkgroupSize", "tunables.tiledWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.tiledWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "(tunables.tiledWorkgroupSize * 8 + 2 * kernelSizePadded - 1) * 4 <= device.limits.maxComputeWorkgroupStorageSize"],
213
+ "derive": {
214
  "hasBias": false,
215
  "hasState": false,
216
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
 
218
  "outputScalar": "dtypes.T",
219
  "workgroupSize": "tunables.tiledWorkgroupSize",
220
  "tileSize": "tunables.tiledWorkgroupSize * 8",
 
 
221
  "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1",
222
  "hasStateWindow": "windowed",
223
  "usesF16": "tensorDtypes.inputT == \"float16\""
 
226
  {
227
  "id": "main",
228
  "name": "CausalConvWithState.TiledLargeKernel",
229
+ "shader": "causal-conv-with-state-tiled.wgsl.jinja",
230
+ "derive": { "materializeConvBeforeActivation": false },
231
+ "bindings": ["input_2", "weight_2", "output_2", "present_state", "params"],
 
 
232
  "dispatch": {
233
+ "x": "min(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)",
234
+ "y": "ceilDiv(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)",
235
+ "z": 1
236
  }
237
  }
238
  ]
 
241
  "id": "state_bias_tiled_large_kernel",
242
  "priority": 10,
243
  "when": ["stateBiasContract", "kernelSize >= 32", "dim(shapes.inputT, 2) >= 256", "dim(shapes.inputT, 2) % 8 == 0", "tunables.tiledWorkgroupSize >= 1", "floor(tunables.tiledWorkgroupSize) == tunables.tiledWorkgroupSize", "tunables.tiledWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.tiledWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "(tunables.tiledWorkgroupSize * 8 + 2 * kernelSizePadded - 1) * 4 <= device.limits.maxComputeWorkgroupStorageSize"],
244
+ "derive": {
245
  "hasBias": true,
246
  "hasState": true,
247
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
 
249
  "outputScalar": "dtypes.T",
250
  "workgroupSize": "tunables.tiledWorkgroupSize",
251
  "tileSize": "tunables.tiledWorkgroupSize * 8",
 
 
252
  "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1",
253
  "hasStateWindow": "windowed",
254
  "usesF16": "tensorDtypes.inputT == \"float16\""
 
257
  {
258
  "id": "main",
259
  "name": "CausalConvWithState.TiledLargeKernel",
260
+ "shader": "causal-conv-with-state-tiled.wgsl.jinja",
261
+ "derive": { "materializeConvBeforeActivation": false },
262
+ "bindings": ["input_2", "weight_2", "bias", "past_state", "output_2", "present_state", "params"],
 
 
263
  "dispatch": {
264
+ "x": "min(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)",
265
+ "y": "ceilDiv(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)",
266
+ "z": 1
267
  }
268
  }
269
  ]
 
272
  "id": "bias_no_state_tiled_large_kernel",
273
  "priority": 10,
274
  "when": ["biasNoStateContract", "kernelSize >= 32", "dim(shapes.inputT, 2) >= 256", "dim(shapes.inputT, 2) % 8 == 0", "tunables.tiledWorkgroupSize >= 1", "floor(tunables.tiledWorkgroupSize) == tunables.tiledWorkgroupSize", "tunables.tiledWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.tiledWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "(tunables.tiledWorkgroupSize * 8 + 2 * kernelSizePadded - 1) * 4 <= device.limits.maxComputeWorkgroupStorageSize"],
275
+ "derive": {
276
  "hasBias": true,
277
  "hasState": false,
278
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
 
280
  "outputScalar": "dtypes.T",
281
  "workgroupSize": "tunables.tiledWorkgroupSize",
282
  "tileSize": "tunables.tiledWorkgroupSize * 8",
 
 
283
  "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1",
284
  "hasStateWindow": "windowed",
285
  "usesF16": "tensorDtypes.inputT == \"float16\""
 
288
  {
289
  "id": "main",
290
  "name": "CausalConvWithState.TiledLargeKernel",
291
+ "shader": "causal-conv-with-state-tiled.wgsl.jinja",
292
+ "derive": { "materializeConvBeforeActivation": false },
293
+ "bindings": ["input_2", "weight_2", "bias", "output_2", "present_state", "params"],
 
 
294
  "dispatch": {
295
+ "x": "min(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)",
296
+ "y": "ceilDiv(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)",
297
+ "z": 1
298
  }
299
  }
300
  ]
 
303
  "id": "state_no_bias_tiled_large_kernel",
304
  "priority": 10,
305
  "when": ["stateNoBiasContract", "kernelSize >= 32", "dim(shapes.inputT, 2) >= 256", "dim(shapes.inputT, 2) % 8 == 0", "tunables.tiledWorkgroupSize >= 1", "floor(tunables.tiledWorkgroupSize) == tunables.tiledWorkgroupSize", "tunables.tiledWorkgroupSize <= device.limits.maxComputeInvocationsPerWorkgroup", "tunables.tiledWorkgroupSize <= device.limits.maxComputeWorkgroupSizeX", "(tunables.tiledWorkgroupSize * 8 + 2 * kernelSizePadded - 1) * 4 <= device.limits.maxComputeWorkgroupStorageSize"],
306
+ "derive": {
307
  "hasBias": false,
308
  "hasState": true,
309
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
 
311
  "outputScalar": "dtypes.T",
312
  "workgroupSize": "tunables.tiledWorkgroupSize",
313
  "tileSize": "tunables.tiledWorkgroupSize * 8",
 
 
314
  "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1",
315
  "hasStateWindow": "windowed",
316
  "usesF16": "tensorDtypes.inputT == \"float16\""
 
319
  {
320
  "id": "main",
321
  "name": "CausalConvWithState.TiledLargeKernel",
322
+ "shader": "causal-conv-with-state-tiled.wgsl.jinja",
323
+ "derive": { "materializeConvBeforeActivation": false },
324
+ "bindings": ["input_2", "weight_2", "past_state", "output_2", "present_state", "params"],
 
 
325
  "dispatch": {
326
+ "x": "min(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)",
327
+ "y": "ceilDiv(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)",
328
+ "z": 1
329
  }
330
  }
331
  ]
 
334
  "id": "zero_state",
335
  "priority": 0,
336
  "when": ["zeroStateContract"],
337
+ "derive": {
338
  "hasBias": false,
339
  "hasState": false,
340
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
 
348
  {
349
  "id": "main",
350
  "name": "CausalConvWithState",
351
+ "shader": "causal-conv-with-state.wgsl.jinja",
352
+ "derive": { "materializeConvBeforeActivation": false },
353
+ "bindings": ["input_2", "weight_2", "output_2", "present_state", "params_2"],
 
 
354
  "dispatch": {
355
+ "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)",
356
+ "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)",
357
+ "z": 1
358
  }
359
  }
360
  ]
 
363
  "id": "state_bias",
364
  "priority": 0,
365
  "when": ["stateBiasContract"],
366
+ "derive": {
367
  "hasBias": true,
368
  "hasState": true,
369
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
 
377
  {
378
  "id": "main",
379
  "name": "CausalConvWithState",
380
+ "shader": "causal-conv-with-state.wgsl.jinja",
381
+ "derive": { "materializeConvBeforeActivation": false },
382
+ "bindings": ["input_2", "weight_2", "bias", "past_state", "output_2", "present_state", "params_2"],
 
 
383
  "dispatch": {
384
+ "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)",
385
+ "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)",
386
+ "z": 1
387
  }
388
  }
389
  ]
 
392
  "id": "bias_no_state",
393
  "priority": 0,
394
  "when": ["biasNoStateContract"],
395
+ "derive": {
396
  "hasBias": true,
397
  "hasState": false,
398
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
 
406
  {
407
  "id": "main",
408
  "name": "CausalConvWithState",
409
+ "shader": "causal-conv-with-state.wgsl.jinja",
410
+ "derive": { "materializeConvBeforeActivation": false },
411
+ "bindings": ["input_2", "weight_2", "bias", "output_2", "present_state", "params_2"],
 
 
412
  "dispatch": {
413
+ "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)",
414
+ "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)",
415
+ "z": 1
416
  }
417
  }
418
  ]
 
421
  "id": "state_no_bias",
422
  "priority": 0,
423
  "when": ["stateNoBiasContract"],
424
+ "derive": {
425
  "hasBias": false,
426
  "hasState": true,
427
  "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
 
435
  {
436
  "id": "main",
437
  "name": "CausalConvWithState",
438
+ "shader": "causal-conv-with-state.wgsl.jinja",
439
+ "derive": { "materializeConvBeforeActivation": false },
440
+ "bindings": ["input_2", "weight_2", "past_state", "output_2", "present_state", "params_2"],
 
 
441
  "dispatch": {
442
+ "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)",
443
+ "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)",
444
+ "z": 1
445
  }
446
  }
447
  ]
build/webgpu/metadata.json CHANGED
@@ -1,20 +1,36 @@
1
  {
2
  "name": "com.microsoft.CausalConvWithState",
3
- "id": "_com_microsoft_causalconvwithstate_webgpu_9e6e59f",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
- "bench.json": "JZU8nd2+4ByxWEyHozqDzF2xGDNwA4gaZ1jk71H/kF0=",
11
- "causal-conv-with-state-tiled.wgsl.jinja": "3AcQURBK9/4LFE/TCfWxDMIAMsjLIVcfsZ9dK/qlzGk=",
12
- "causal-conv-with-state-vec4.wgsl.jinja": "jjKAFL4nTGXXrRLF+TyBiJO4fNy0bEF1gK0KfIGeRoc=",
13
- "causal-conv-with-state.wgsl.jinja": "B03ROsgmw6YXaTQ6373ibBbKyt7lxTBThsR7ThfyXVU=",
14
- "manifest.json": "sR2aWi0R4vk7io6t9uwLdcVyP8jdQ5PIA0QHMHKQkaU=",
15
- "test.json": "oC4aXSiNHpXHdk6nE+nXhAZvkGulY8eJOcfiXZ0jXdw="
16
  }
17
  },
18
- "provenance": { "kernel": { "sha": "2e7068faf55e7f43df740015f6d1ee49391a41c5", "dirty": false } },
19
- "webgpu": { "manifestSpec": "1.0", "specialized": true, "opPath": "ops/com.microsoft.CausalConvWithState" }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
20
  }
 
1
  {
2
  "name": "com.microsoft.CausalConvWithState",
3
+ "id": "_com_microsoft_causalconvwithstate_webgpu_2f94cdc",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "backend": { "type": "webgpu" },
7
  "digest": {
8
  "algorithm": "sha256",
9
  "files": {
10
+ "bench.json": "xexf85uhEj1seBHH80A/gzgwsempJwEy1IPsqs2Q81w=",
11
+ "causal-conv-with-state-tiled.wgsl.jinja": "VcSN9BQ2I4NO+PfsdA4qMne9kISCvAZeturYtZLWpmI=",
12
+ "causal-conv-with-state-vec4.wgsl.jinja": "rxqAeKAR/c0Qpy1FL2kyiNO1Ki+dusqnxkQcYtnhV4E=",
13
+ "causal-conv-with-state.wgsl.jinja": "QrQcWLijzv+pENPVNjTA/UNNzCCNXrjiEySPL/rOuMQ=",
14
+ "manifest.json": "nUvFIfdnmA2uQhwLhzjd141hkTwVnTFEppipUB6Fk+E=",
15
+ "test.json": "ITwDaH/kaBYxzYqj2uGRdZXBX7uN8/++FZ/mKSQvnho="
16
  }
17
  },
18
+ "provenance": { "kernel": { "sha": "91d990483a174128daf7673f3f37a7c890493ae1", "dirty": false } },
19
+ "webgpu": {
20
+ "manifestSpec": "2.0",
21
+ "variants": {
22
+ "zero_state_vec4": ["causal-conv-with-state-vec4.wgsl.jinja"],
23
+ "state_bias_vec4": ["causal-conv-with-state-vec4.wgsl.jinja"],
24
+ "bias_no_state_vec4": ["causal-conv-with-state-vec4.wgsl.jinja"],
25
+ "state_no_bias_vec4": ["causal-conv-with-state-vec4.wgsl.jinja"],
26
+ "zero_state_tiled_large_kernel": ["causal-conv-with-state-tiled.wgsl.jinja"],
27
+ "state_bias_tiled_large_kernel": ["causal-conv-with-state-tiled.wgsl.jinja"],
28
+ "bias_no_state_tiled_large_kernel": ["causal-conv-with-state-tiled.wgsl.jinja"],
29
+ "state_no_bias_tiled_large_kernel": ["causal-conv-with-state-tiled.wgsl.jinja"],
30
+ "zero_state": ["causal-conv-with-state.wgsl.jinja"],
31
+ "state_bias": ["causal-conv-with-state.wgsl.jinja"],
32
+ "bias_no_state": ["causal-conv-with-state.wgsl.jinja"],
33
+ "state_no_bias": ["causal-conv-with-state.wgsl.jinja"]
34
+ }
35
+ }
36
  }
build/webgpu/test.json CHANGED
@@ -1,5 +1,4 @@
1
  {
2
- "op": "com.microsoft.CausalConvWithState",
3
  "cases": [
4
  {
5
  "name": "ort_kernel1_zero_size_state",
@@ -810,7 +809,7 @@
810
  {
811
  "name": "state_window2_pinned",
812
  "provenance": {
813
- "notes": "Hand-computed from the ONNX Runtime state_window contract (onnxruntime/core/graph/contrib_ops/bert_defs.cc, the CausalConvWithState schema). Upstream's own state_window cases are CUDA-only and compare against a replayed reference rather than pinned numbers, so the expected values here were worked out by hand instead of ported. Slot 0 is the carry state after position 1 and slot 1 after position 2, so slot 1 repeats what the unwindowed op writes."
814
  },
815
  "attrs": { "activation": "none", "state_window": 2 },
816
  "inputs": {
@@ -839,7 +838,7 @@
839
  {
840
  "name": "state_window4_longer_than_sequence",
841
  "provenance": {
842
- "notes": "Hand-computed from the ONNX Runtime state_window contract (onnxruntime/core/graph/contrib_ops/bert_defs.cc, the CausalConvWithState schema). Upstream's own state_window cases are CUDA-only and compare against a replayed reference rather than pinned numbers, so the expected values here were worked out by hand instead of ported. W exceeds the sequence length, so the leading W - T slots must be zero rather than uninitialized."
843
  },
844
  "attrs": { "activation": "none", "state_window": 4 },
845
  "inputs": {
@@ -868,7 +867,7 @@
868
  {
869
  "name": "state_window2_past_slot_pinned",
870
  "provenance": {
871
- "notes": "Hand-computed from the ONNX Runtime state_window contract (onnxruntime/core/graph/contrib_ops/bert_defs.cc, the CausalConvWithState schema). Upstream's own state_window cases are CUDA-only and compare against a replayed reference rather than pinned numbers, so the expected values here were worked out by hand instead of ported. past_state slot 0 is poisoned with large negatives that no correct read touches; only slot W-1 carries the previous call's state."
872
  },
873
  "attrs": { "activation": "none", "state_window": 2 },
874
  "inputs": {
@@ -925,7 +924,7 @@
925
  {
926
  "name": "large_kernel_tiled_zero_state_window2",
927
  "provenance": {
928
- "notes": "Windowed present_state on the large-kernel tiled path with no past state; the state-writing tile now strides over a (slot, element) grid instead of a single slot."
929
  },
930
  "attrs": { "activation": "none", "state_window": 2 },
931
  "inputs": {
@@ -1029,7 +1028,7 @@
1029
  {
1030
  "name": "vec4_state_window6_longer_than_sequence",
1031
  "provenance": {
1032
- "notes": "W = 6 exceeds the four-position input, so the vec4 path's two leading window slots hold no position from this call and must be zero. Only this variant can reach that branch with a window: the tiled path demands at least 256 positions, which no legal window exceeds."
1033
  },
1034
  "attrs": { "activation": "none", "state_window": 6 },
1035
  "inputs": {
@@ -1052,7 +1051,7 @@
1052
  {
1053
  "name": "vec4_state_window_past_state_prefix",
1054
  "provenance": {
1055
- "notes": "A windowed state whose window reaches back further than this call is long, WITH a past state: the early slots carry positions from before this call, so they have to come from past_state rather than from the input row."
1056
  },
1057
  "attrs": { "activation": "silu", "state_window": 6 },
1058
  "inputs": {
@@ -1080,7 +1079,7 @@
1080
  {
1081
  "name": "f16_scalar_state_bias_silu",
1082
  "provenance": {
1083
- "notes": "float16 tensors on the scalar kernel. ONNX Runtime registers this operator for the whole supported float set; this port pinned float32. Every tap and accumulation still runs in f32 and only the store narrows, which is what the kernel already did for float32."
1084
  },
1085
  "attrs": { "activation": "silu" },
1086
  "inputs": {
@@ -1108,9 +1107,7 @@
1108
  },
1109
  {
1110
  "name": "f16_k4_vec4_zero_state_silu",
1111
- "provenance": {
1112
- "notes": "float16 on the four-tap vectorized kernel, which read the bound element type directly and so was the only one of the three actually pinned to float32."
1113
- },
1114
  "attrs": { "activation": "silu" },
1115
  "inputs": {
1116
  "inputT": {
@@ -1182,7 +1179,7 @@
1182
  {
1183
  "name": "large_kernel_tiled_unaligned_k33_weight_tile_pad",
1184
  "provenance": {
1185
- "notes": "Kernel length 1 mod 4. The tap loop consumes four weights per iteration, so this shape reaches the tiled path only via the zero-padded weight tile; before that it fell to the untiled kernel."
1186
  },
1187
  "attrs": { "activation": "none" },
1188
  "inputs": {
@@ -1204,7 +1201,9 @@
1204
  },
1205
  {
1206
  "name": "large_kernel_tiled_unaligned_k34_weight_tile_pad",
1207
- "provenance": { "notes": "Kernel length 2 mod 4 -- the other half of the padded-tail arithmetic." },
 
 
1208
  "attrs": { "activation": "none" },
1209
  "inputs": {
1210
  "inputT": {
@@ -1278,6 +1277,67 @@
1278
  "outputT": { "dtype": "float32", "shape": [1, 2, 256], "tolerance": 0.00005 },
1279
  "presentStateT": { "dtype": "float32", "shape": [1, 2, 36], "tolerance": 0.000001 }
1280
  }
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1281
  }
1282
  ]
1283
  }
 
1
  {
 
2
  "cases": [
3
  {
4
  "name": "ort_kernel1_zero_size_state",
 
809
  {
810
  "name": "state_window2_pinned",
811
  "provenance": {
812
+ "notes": "Expected values are derived directly from ONNX Runtime's `state_window` schema. Slot 0 is the carry state after position 1 and slot 1 is the carry state after position 2."
813
  },
814
  "attrs": { "activation": "none", "state_window": 2 },
815
  "inputs": {
 
838
  {
839
  "name": "state_window4_longer_than_sequence",
840
  "provenance": {
841
+ "notes": "Expected values are derived directly from ONNX Runtime's `state_window` schema. The window exceeds the sequence length, so its leading `W - T` slots must be zero."
842
  },
843
  "attrs": { "activation": "none", "state_window": 4 },
844
  "inputs": {
 
867
  {
868
  "name": "state_window2_past_slot_pinned",
869
  "provenance": {
870
+ "notes": "Expected values are derived directly from ONNX Runtime's `state_window` schema. Large negative sentinels in past-state slot 0 must remain unread; only slot `W - 1` carries the preceding state."
871
  },
872
  "attrs": { "activation": "none", "state_window": 2 },
873
  "inputs": {
 
924
  {
925
  "name": "large_kernel_tiled_zero_state_window2",
926
  "provenance": {
927
+ "notes": "A windowed present state on the large-kernel tiled path uses a two-dimensional `(slot, element)` output grid with no past state."
928
  },
929
  "attrs": { "activation": "none", "state_window": 2 },
930
  "inputs": {
 
1028
  {
1029
  "name": "vec4_state_window6_longer_than_sequence",
1030
  "provenance": {
1031
+ "notes": "With `W = 6` and four input positions, the vec4 path must write zero to the two leading window slots that have no position in this call."
1032
  },
1033
  "attrs": { "activation": "none", "state_window": 6 },
1034
  "inputs": {
 
1051
  {
1052
  "name": "vec4_state_window_past_state_prefix",
1053
  "provenance": {
1054
+ "notes": "The window reaches before the start of this call, so its early slots must come from `past_state` rather than the current input row."
1055
  },
1056
  "attrs": { "activation": "silu", "state_window": 6 },
1057
  "inputs": {
 
1079
  {
1080
  "name": "f16_scalar_state_bias_silu",
1081
  "provenance": {
1082
+ "notes": "Float16 tensors exercise the scalar kernel; every tap and accumulation uses float32 and only the store narrows."
1083
  },
1084
  "attrs": { "activation": "silu" },
1085
  "inputs": {
 
1107
  },
1108
  {
1109
  "name": "f16_k4_vec4_zero_state_silu",
1110
+ "provenance": { "notes": "Float16 tensors exercise the four-tap vectorized kernel and its typed input loads." },
 
 
1111
  "attrs": { "activation": "silu" },
1112
  "inputs": {
1113
  "inputT": {
 
1179
  {
1180
  "name": "large_kernel_tiled_unaligned_k33_weight_tile_pad",
1181
  "provenance": {
1182
+ "notes": "Kernel length 33 leaves one live tap in the final four-wide iteration, requiring zero padding in the tiled weight buffer."
1183
  },
1184
  "attrs": { "activation": "none" },
1185
  "inputs": {
 
1201
  },
1202
  {
1203
  "name": "large_kernel_tiled_unaligned_k34_weight_tile_pad",
1204
+ "provenance": {
1205
+ "notes": "Kernel length 34 leaves two live taps in the final four-wide iteration, exercising the tiled weight-buffer tail."
1206
+ },
1207
  "attrs": { "activation": "none" },
1208
  "inputs": {
1209
  "inputT": {
 
1277
  "outputT": { "dtype": "float32", "shape": [1, 2, 256], "tolerance": 0.00005 },
1278
  "presentStateT": { "dtype": "float32", "shape": [1, 2, 36], "tolerance": 0.000001 }
1279
  }
1280
+ },
1281
+ {
1282
+ "name": "vec4_state_window_bias_no_past_state",
1283
+ "provenance": {
1284
+ "notes": "A windowed present_state on the aligned K=4 vec4 arm that also carries a bias and starts from the zero prefix, so the window slot stride is exercised without a past_state input."
1285
+ },
1286
+ "attrs": { "activation": "silu", "state_window": 3 },
1287
+ "inputs": {
1288
+ "inputT": {
1289
+ "dtype": "float32",
1290
+ "shape": [1, 2, 8],
1291
+ "data": { "kind": "fillFloat32", "scale": 0.4, "sinStep": 0.23, "cosStep": 0.19 }
1292
+ },
1293
+ "weightT": {
1294
+ "dtype": "float32",
1295
+ "shape": [2, 1, 4],
1296
+ "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.13, "cosStep": 0.29 }
1297
+ },
1298
+ "biasT": {
1299
+ "dtype": "float32",
1300
+ "shape": [2],
1301
+ "data": { "kind": "fillFloat32", "scale": 0.3, "sinStep": 0.07, "cosStep": 0.41 }
1302
+ }
1303
+ },
1304
+ "outputs": {
1305
+ "outputT": { "dtype": "float32", "shape": [1, 2, 8], "tolerance": 0.00002 },
1306
+ "presentStateT": { "dtype": "float32", "shape": [3, 1, 2, 3], "tolerance": 0.000001 }
1307
+ }
1308
+ },
1309
+ {
1310
+ "name": "vec4_state_window_bias_past_state",
1311
+ "provenance": {
1312
+ "notes": "The windowed vec4 arm with both a bias and a past_state: the pinned past slot and the window slot stride are read in the same render."
1313
+ },
1314
+ "attrs": { "activation": "silu", "state_window": 3 },
1315
+ "inputs": {
1316
+ "inputT": {
1317
+ "dtype": "float32",
1318
+ "shape": [1, 2, 8],
1319
+ "data": { "kind": "fillFloat32", "scale": 0.4, "sinStep": 0.31, "cosStep": 0.17 }
1320
+ },
1321
+ "weightT": {
1322
+ "dtype": "float32",
1323
+ "shape": [2, 1, 4],
1324
+ "data": { "kind": "fillFloat32", "scale": 0.2, "sinStep": 0.11, "cosStep": 0.37 }
1325
+ },
1326
+ "biasT": {
1327
+ "dtype": "float32",
1328
+ "shape": [2],
1329
+ "data": { "kind": "fillFloat32", "scale": 0.3, "sinStep": 0.05, "cosStep": 0.43 }
1330
+ },
1331
+ "pastStateT": {
1332
+ "dtype": "float32",
1333
+ "shape": [3, 1, 2, 3],
1334
+ "data": { "kind": "fillFloat32", "scale": 0.5, "sinStep": 0.29, "cosStep": 0.13 }
1335
+ }
1336
+ },
1337
+ "outputs": {
1338
+ "outputT": { "dtype": "float32", "shape": [1, 2, 8], "tolerance": 0.00002 },
1339
+ "presentStateT": { "dtype": "float32", "shape": [3, 1, 2, 3], "tolerance": 0.000001 }
1340
+ }
1341
  }
1342
  ]
1343
  }