{ "domain": "com.microsoft", "name": "CausalConvWithState", "sinceVersion": 1, "inputs": { "inputT": { "onnx": "input", "dtype": "T", "rank": 3 }, "weightT": { "onnx": "weight", "dtype": "T", "rank": 3 }, "biasT": { "onnx": "bias", "dtype": "T", "rank": 1, "optional": true }, "pastStateT": { "onnx": "past_state", "dtype": "T", "rank": "3 if attrs.state_window == 0 else 4", "optional": true } }, "outputs": { "outputT": { "onnx": "output", "dtype": "T", "rank": 3, "shape": "shapes.inputT" }, "presentStateT": { "onnx": "present_state", "dtype": "T", "rank": "3 if attrs.state_window == 0 else 4", "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]" } }, "attributes": { "activation": { "default": "none" }, "ndim": { "default": 1 }, "state_window": { "default": 0 } }, "attributeConstraints": { "activation": { "values": ["none", "silu", "swish"] }, "ndim": { "values": [1] } }, "typeConstraints": { "T": ["float32", "float16"] }, "tunables": { "workgroupSize": { "default": 256 }, "tiledWorkgroupSize": { "default": 128 } }, "derive": { "stateWindow": "attrs.state_window", "windowed": "stateWindow > 0", "stateWindowOk": "stateWindow >= 0 and stateWindow <= 8", "kernelSize": "dim(shapes.weightT, ranks.weightT - 1)", "kernelSizePadded": "ceilDiv(kernelSize, 4) * 4", "weightRankOk": "ranks.weightT == 3 and dim(shapes.weightT, 1) == 1", "stateLength": "kernelSize - 1", "stateSlotStride": "dim(shapes.inputT, 0) * dim(shapes.inputT, 1) * stateLength", "windowedLengthOk": "not windowed or dim(shapes.inputT, 2) > 0", "presentStateOk": "(ranks.presentStateT == 3 and dim(shapes.presentStateT, 0) == dim(shapes.inputT, 0) and dim(shapes.presentStateT, 1) == dim(shapes.inputT, 1) and dim(shapes.presentStateT, 2) == stateLength) if not windowed else (ranks.presentStateT == 4 and dim(shapes.presentStateT, 0) == stateWindow and dim(shapes.presentStateT, 1) == dim(shapes.inputT, 0) and dim(shapes.presentStateT, 2) == dim(shapes.inputT, 1) and dim(shapes.presentStateT, 3) == stateLength)", "pastStateShapeOk": "present.pastStateT and ((ranks.pastStateT == 3 and dim(shapes.pastStateT, 0) == dim(shapes.inputT, 0) and dim(shapes.pastStateT, 1) == dim(shapes.inputT, 1) and dim(shapes.pastStateT, 2) == stateLength) if not windowed else (ranks.pastStateT == 4 and dim(shapes.pastStateT, 0) == stateWindow and dim(shapes.pastStateT, 1) == dim(shapes.inputT, 0) and dim(shapes.pastStateT, 2) == dim(shapes.inputT, 1) and dim(shapes.pastStateT, 3) == stateLength))", "commonContract": "ranks.inputT == 3 and weightRankOk and ranks.outputT == 3 and (tensorDtypes.inputT == \"float32\" or tensorDtypes.inputT == \"float16\") and tensorDtypes.weightT == tensorDtypes.inputT and tensorDtypes.outputT == tensorDtypes.inputT and tensorDtypes.presentStateT == tensorDtypes.inputT and f16Ok(dtypes.T) and dim(shapes.inputT, 1) == dim(shapes.weightT, 0) and dim(shapes.outputT, 0) == dim(shapes.inputT, 0) and dim(shapes.outputT, 1) == dim(shapes.inputT, 1) and dim(shapes.outputT, 2) == dim(shapes.inputT, 2) and stateWindowOk and windowedLengthOk and presentStateOk", "zeroStateContract": "commonContract and not present.pastStateT and not present.biasT", "biasNoStateContract": "commonContract and not present.pastStateT and present.biasT and ranks.biasT == 1 and tensorDtypes.biasT == tensorDtypes.inputT and dim(shapes.biasT, 0) == dim(shapes.inputT, 1)", "stateNoBiasContract": "commonContract and present.pastStateT and not present.biasT and tensorDtypes.pastStateT == tensorDtypes.inputT and pastStateShapeOk", "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)" }, "bindings": { "input": { "arg": "inputT", "buffer": "read-only-storage", "elementType": "$inputVec4" }, "weight": { "arg": "weightT", "buffer": "read-only-storage", "elementType": "$weightElem" }, "output": { "arg": "outputT", "buffer": "storage", "elementType": "$outputVec4" }, "present_state": { "arg": "presentStateT", "buffer": "storage", "elementType": "$outputScalar" }, "params": { "buffer": "uniform", "struct": [ { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" }, { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" }, { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" }, { "name": "stateWindow", "type": "u32", "value": "stateWindow" }, { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" } ] }, "bias": { "arg": "biasT", "buffer": "read-only-storage", "elementType": "$inputScalar" }, "past_state": { "arg": "pastStateT", "buffer": "read-only-storage", "elementType": "$inputScalar" }, "input_2": { "arg": "inputT", "name": "input", "buffer": "read-only-storage", "elementType": "$inputScalar" }, "weight_2": { "arg": "weightT", "name": "weight", "buffer": "read-only-storage", "elementType": "$inputScalar" }, "output_2": { "arg": "outputT", "name": "output", "buffer": "storage", "elementType": "$outputScalar" }, "params_2": { "name": "params", "buffer": "uniform", "struct": [ { "name": "batchSize", "type": "u32", "value": "dim(shapes.inputT, 0)" }, { "name": "channels", "type": "u32", "value": "dim(shapes.inputT, 1)" }, { "name": "length", "type": "u32", "value": "dim(shapes.inputT, 2)" }, { "name": "kernelSize", "type": "u32", "value": "kernelSize" }, { "name": "stateWindow", "type": "u32", "value": "stateWindow" }, { "name": "stateSlotStride", "type": "u32", "value": "stateSlotStride" } ] } }, "variants": [ { "id": "zero_state_vec4", "priority": 20, "when": ["zeroStateContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"], "derive": { "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", "workgroupSize": 256, "hasStateWindow": "windowed", "usesF16": "tensorDtypes.inputT == \"float16\"", "inputScalar": "dtypes.T", "outputScalar": "dtypes.T", "inputVec4": "\"vec4\" if tensorDtypes.inputT == \"float16\" else \"vec4\"", "outputVec4": "\"vec4\" if tensorDtypes.inputT == \"float16\" else \"vec4\"", "hasBias": false, "hasState": false, "weightElem": "(\"vec4\" if tensorDtypes.inputT == \"float16\" else \"vec4\") if kernelSize == 4 else dtypes.T" }, "passes": [ { "id": "main", "name": "CausalConvWithState.Vec4", "shader": "causal-conv-with-state-vec4.wgsl.jinja", "derive": { "materializeConvBeforeActivation": false }, "bindings": ["input", "weight", "output", "present_state", "params"], "dispatch": { "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)", "z": 1 } } ] }, { "id": "state_bias_vec4", "priority": 20, "when": ["stateBiasContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"], "derive": { "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", "workgroupSize": 256, "hasStateWindow": "windowed", "usesF16": "tensorDtypes.inputT == \"float16\"", "inputScalar": "dtypes.T", "outputScalar": "dtypes.T", "inputVec4": "\"vec4\" if tensorDtypes.inputT == \"float16\" else \"vec4\"", "outputVec4": "\"vec4\" if tensorDtypes.inputT == \"float16\" else \"vec4\"", "hasBias": true, "hasState": true, "weightElem": "(\"vec4\" if tensorDtypes.inputT == \"float16\" else \"vec4\") if kernelSize == 4 else dtypes.T" }, "passes": [ { "id": "main", "name": "CausalConvWithState.Vec4", "shader": "causal-conv-with-state-vec4.wgsl.jinja", "derive": { "materializeConvBeforeActivation": false }, "bindings": ["input", "weight", "bias", "past_state", "output", "present_state", "params"], "dispatch": { "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)", "z": 1 } } ] }, { "id": "bias_no_state_vec4", "priority": 20, "when": ["biasNoStateContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"], "derive": { "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", "workgroupSize": 256, "hasStateWindow": "windowed", "usesF16": "tensorDtypes.inputT == \"float16\"", "inputScalar": "dtypes.T", "outputScalar": "dtypes.T", "inputVec4": "\"vec4\" if tensorDtypes.inputT == \"float16\" else \"vec4\"", "outputVec4": "\"vec4\" if tensorDtypes.inputT == \"float16\" else \"vec4\"", "hasBias": true, "hasState": false, "weightElem": "(\"vec4\" if tensorDtypes.inputT == \"float16\" else \"vec4\") if kernelSize == 4 else dtypes.T" }, "passes": [ { "id": "main", "name": "CausalConvWithState.Vec4", "shader": "causal-conv-with-state-vec4.wgsl.jinja", "derive": { "materializeConvBeforeActivation": false }, "bindings": ["input", "weight", "bias", "output", "present_state", "params"], "dispatch": { "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)", "z": 1 } } ] }, { "id": "state_no_bias_vec4", "priority": 20, "when": ["stateNoBiasContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"], "derive": { "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", "workgroupSize": 256, "hasStateWindow": "windowed", "usesF16": "tensorDtypes.inputT == \"float16\"", "inputScalar": "dtypes.T", "outputScalar": "dtypes.T", "inputVec4": "\"vec4\" if tensorDtypes.inputT == \"float16\" else \"vec4\"", "outputVec4": "\"vec4\" if tensorDtypes.inputT == \"float16\" else \"vec4\"", "hasBias": false, "hasState": true, "weightElem": "(\"vec4\" if tensorDtypes.inputT == \"float16\" else \"vec4\") if kernelSize == 4 else dtypes.T" }, "passes": [ { "id": "main", "name": "CausalConvWithState.Vec4", "shader": "causal-conv-with-state-vec4.wgsl.jinja", "derive": { "materializeConvBeforeActivation": false }, "bindings": ["input", "weight", "past_state", "output", "present_state", "params"], "dispatch": { "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)), (workgroupSize)), 65535)", "z": 1 } } ] }, { "id": "zero_state_tiled_large_kernel", "priority": 10, "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"], "derive": { "hasBias": false, "hasState": false, "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", "inputScalar": "dtypes.T", "outputScalar": "dtypes.T", "workgroupSize": "tunables.tiledWorkgroupSize", "tileSize": "tunables.tiledWorkgroupSize * 8", "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1", "hasStateWindow": "windowed", "usesF16": "tensorDtypes.inputT == \"float16\"" }, "passes": [ { "id": "main", "name": "CausalConvWithState.TiledLargeKernel", "shader": "causal-conv-with-state-tiled.wgsl.jinja", "derive": { "materializeConvBeforeActivation": false }, "bindings": ["input_2", "weight_2", "output_2", "present_state", "params"], "dispatch": { "x": "min(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)", "y": "ceilDiv(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)", "z": 1 } } ] }, { "id": "state_bias_tiled_large_kernel", "priority": 10, "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"], "derive": { "hasBias": true, "hasState": true, "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", "inputScalar": "dtypes.T", "outputScalar": "dtypes.T", "workgroupSize": "tunables.tiledWorkgroupSize", "tileSize": "tunables.tiledWorkgroupSize * 8", "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1", "hasStateWindow": "windowed", "usesF16": "tensorDtypes.inputT == \"float16\"" }, "passes": [ { "id": "main", "name": "CausalConvWithState.TiledLargeKernel", "shader": "causal-conv-with-state-tiled.wgsl.jinja", "derive": { "materializeConvBeforeActivation": false }, "bindings": ["input_2", "weight_2", "bias", "past_state", "output_2", "present_state", "params"], "dispatch": { "x": "min(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)", "y": "ceilDiv(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)", "z": 1 } } ] }, { "id": "bias_no_state_tiled_large_kernel", "priority": 10, "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"], "derive": { "hasBias": true, "hasState": false, "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", "inputScalar": "dtypes.T", "outputScalar": "dtypes.T", "workgroupSize": "tunables.tiledWorkgroupSize", "tileSize": "tunables.tiledWorkgroupSize * 8", "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1", "hasStateWindow": "windowed", "usesF16": "tensorDtypes.inputT == \"float16\"" }, "passes": [ { "id": "main", "name": "CausalConvWithState.TiledLargeKernel", "shader": "causal-conv-with-state-tiled.wgsl.jinja", "derive": { "materializeConvBeforeActivation": false }, "bindings": ["input_2", "weight_2", "bias", "output_2", "present_state", "params"], "dispatch": { "x": "min(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)", "y": "ceilDiv(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)", "z": 1 } } ] }, { "id": "state_no_bias_tiled_large_kernel", "priority": 10, "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"], "derive": { "hasBias": false, "hasState": true, "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", "inputScalar": "dtypes.T", "outputScalar": "dtypes.T", "workgroupSize": "tunables.tiledWorkgroupSize", "tileSize": "tunables.tiledWorkgroupSize * 8", "inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1", "hasStateWindow": "windowed", "usesF16": "tensorDtypes.inputT == \"float16\"" }, "passes": [ { "id": "main", "name": "CausalConvWithState.TiledLargeKernel", "shader": "causal-conv-with-state-tiled.wgsl.jinja", "derive": { "materializeConvBeforeActivation": false }, "bindings": ["input_2", "weight_2", "past_state", "output_2", "present_state", "params"], "dispatch": { "x": "min(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)", "y": "ceilDiv(dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), tileSize), 65535)", "z": 1 } } ] }, { "id": "zero_state", "priority": 0, "when": ["zeroStateContract"], "derive": { "hasBias": false, "hasState": false, "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", "inputScalar": "dtypes.T", "outputScalar": "dtypes.T", "workgroupSize": "tunables.workgroupSize", "hasStateWindow": "windowed", "usesF16": "tensorDtypes.inputT == \"float16\"" }, "passes": [ { "id": "main", "name": "CausalConvWithState", "shader": "causal-conv-with-state.wgsl.jinja", "derive": { "materializeConvBeforeActivation": false }, "bindings": ["input_2", "weight_2", "output_2", "present_state", "params_2"], "dispatch": { "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)", "z": 1 } } ] }, { "id": "state_bias", "priority": 0, "when": ["stateBiasContract"], "derive": { "hasBias": true, "hasState": true, "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", "inputScalar": "dtypes.T", "outputScalar": "dtypes.T", "workgroupSize": "tunables.workgroupSize", "hasStateWindow": "windowed", "usesF16": "tensorDtypes.inputT == \"float16\"" }, "passes": [ { "id": "main", "name": "CausalConvWithState", "shader": "causal-conv-with-state.wgsl.jinja", "derive": { "materializeConvBeforeActivation": false }, "bindings": ["input_2", "weight_2", "bias", "past_state", "output_2", "present_state", "params_2"], "dispatch": { "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)", "z": 1 } } ] }, { "id": "bias_no_state", "priority": 0, "when": ["biasNoStateContract"], "derive": { "hasBias": true, "hasState": false, "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", "inputScalar": "dtypes.T", "outputScalar": "dtypes.T", "workgroupSize": "tunables.workgroupSize", "hasStateWindow": "windowed", "usesF16": "tensorDtypes.inputT == \"float16\"" }, "passes": [ { "id": "main", "name": "CausalConvWithState", "shader": "causal-conv-with-state.wgsl.jinja", "derive": { "materializeConvBeforeActivation": false }, "bindings": ["input_2", "weight_2", "bias", "output_2", "present_state", "params_2"], "dispatch": { "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)", "z": 1 } } ] }, { "id": "state_no_bias", "priority": 0, "when": ["stateNoBiasContract"], "derive": { "hasBias": false, "hasState": true, "useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"", "inputScalar": "dtypes.T", "outputScalar": "dtypes.T", "workgroupSize": "tunables.workgroupSize", "hasStateWindow": "windowed", "usesF16": "tensorDtypes.inputT == \"float16\"" }, "passes": [ { "id": "main", "name": "CausalConvWithState", "shader": "causal-conv-with-state.wgsl.jinja", "derive": { "materializeConvBeforeActivation": false }, "bindings": ["input_2", "weight_2", "past_state", "output_2", "present_state", "params_2"], "dispatch": { "x": "min(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)", "y": "ceilDiv(ceilDiv((dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))), (workgroupSize)), 65535)", "z": 1 } } ] } ] }