Xenova's picture
Xenova HF Staff
sync 2e7068faf55e
53a7bdc verified
Raw
History Blame
46.6 kB
{
"domain": "com.microsoft",
"name": "CausalConvWithState",
"sinceVersion": 1,
"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.",
"inputs": [
{
"role": "input",
"dtype": "T",
"rank": 3,
"description": "Channels-first input tensor with shape `(batch_size, channels, sequence_length)` for the supported 1-D mode."
},
{
"role": "weight",
"dtype": "T",
"rank": 3,
"description": "Depthwise convolution kernel with shape `(channels, 1, kernel_size)` for the supported 1-D mode."
},
{
"role": "bias",
"dtype": "T",
"rank": 1,
"optional": true,
"description": "Optional per-channel bias with shape `(channels,)`."
},
{
"role": "past_state",
"dtype": "T",
"rank": "3 if attrs.state_window == 0 else 4",
"optional": true,
"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."
}
],
"outputs": [
{
"role": "output",
"dtype": "T",
"rank": 3,
"shape": "shapes.input",
"description": "Convolution output with the same shape as `input`."
},
{
"role": "present_state",
"dtype": "T",
"rank": "3 if attrs.state_window == 0 else 4",
"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]",
"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`."
}
],
"attributes": { "activation": "none", "ndim": 1, "state_window": 0 },
"attributeConstraints": { "activation": { "values": ["none", "silu", "swish"] }, "ndim": { "values": [1] } },
"attributeDescriptions": {
"activation": "Activation applied after convolution and bias. Defaults to `none`; `swish` is an alias of SiLU.",
"ndim": "Number of spatial dimensions. This implementation supports the contrib 1D mode (`ndim = 1`).",
"state_window": "Contrib extension selecting the number of rollback state slots to retain, in the range 0 through 8. Defaults to 0."
},
"typeConstraints": { "T": ["float32", "float16"] },
"args": {
"inputT": { "kind": "tensor", "semantic": "input", "role": "input" },
"weightT": { "kind": "tensor", "semantic": "weight", "role": "input" },
"biasT": { "kind": "tensor", "semantic": "bias", "role": "input", "required": false },
"pastStateT": { "kind": "tensor", "semantic": "past_state", "role": "input", "required": false },
"outputT": { "kind": "tensor", "semantic": "output", "role": "output" },
"presentStateT": { "kind": "tensor", "semantic": "present_state", "role": "output" }
},
"tunables": { "workgroupSize": 256, "tiledWorkgroupSize": 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)"
},
"bindingSets": {
"zeroScalar": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": {
"name": "Params",
"fields": [
{ "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" }
]
}
}
],
"zeroVec4": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputVec4"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$weightElem"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputVec4"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": {
"name": "Params",
"fields": [
{ "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" }
]
}
}
],
"biasNoState": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "bias",
"arg": "biasT",
"semantic": "bias",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": {
"name": "Params",
"fields": [
{ "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" }
]
}
}
],
"stateNoBias": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "past_state",
"arg": "pastStateT",
"semantic": "past_state",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": {
"name": "Params",
"fields": [
{ "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" }
]
}
}
],
"stateBias": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "bias",
"arg": "biasT",
"semantic": "bias",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "past_state",
"arg": "pastStateT",
"semantic": "past_state",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": {
"name": "Params",
"fields": [
{ "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" }
]
}
}
],
"zeroScalarIo": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
}
],
"biasNoStateIo": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "bias",
"arg": "biasT",
"semantic": "bias",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
}
],
"stateNoBiasIo": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "past_state",
"arg": "pastStateT",
"semantic": "past_state",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
}
],
"stateBiasIo": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "bias",
"arg": "biasT",
"semantic": "bias",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "past_state",
"arg": "pastStateT",
"semantic": "past_state",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
}
],
"zeroTiled": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": {
"name": "Params",
"fields": [
{ "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" }
]
}
}
],
"biasNoStateTiled": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "bias",
"arg": "biasT",
"semantic": "bias",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": {
"name": "Params",
"fields": [
{ "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" }
]
}
}
],
"stateNoBiasTiled": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "past_state",
"arg": "pastStateT",
"semantic": "past_state",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": {
"name": "Params",
"fields": [
{ "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" }
]
}
}
],
"stateBiasTiled": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "bias",
"arg": "biasT",
"semantic": "bias",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "past_state",
"arg": "pastStateT",
"semantic": "past_state",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": {
"name": "Params",
"fields": [
{ "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" }
]
}
}
],
"biasNoStateVec4": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputVec4"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$weightElem"
},
{
"name": "bias",
"arg": "biasT",
"semantic": "bias",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputVec4"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": {
"name": "Params",
"fields": [
{ "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" }
]
}
}
],
"stateNoBiasVec4": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputVec4"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$weightElem"
},
{
"name": "past_state",
"arg": "pastStateT",
"semantic": "past_state",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputVec4"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": {
"name": "Params",
"fields": [
{ "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" }
]
}
}
],
"stateBiasVec4": [
{
"name": "input",
"arg": "inputT",
"semantic": "input",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputVec4"
},
{
"name": "weight",
"arg": "weightT",
"semantic": "weight",
"buffer": { "type": "read-only-storage" },
"elementType": "$weightElem"
},
{
"name": "bias",
"arg": "biasT",
"semantic": "bias",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "past_state",
"arg": "pastStateT",
"semantic": "past_state",
"buffer": { "type": "read-only-storage" },
"elementType": "$inputScalar"
},
{
"name": "output",
"arg": "outputT",
"semantic": "output",
"buffer": { "type": "storage" },
"elementType": "$outputVec4"
},
{
"name": "present_state",
"arg": "presentStateT",
"semantic": "present_state",
"buffer": { "type": "storage" },
"elementType": "$outputScalar"
},
{
"name": "params",
"semantic": "kernel.params",
"buffer": { "type": "uniform" },
"struct": {
"name": "Params",
"fields": [
{ "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" }
]
}
}
]
},
"variants": [
{
"id": "zero_state_vec4",
"priority": 20,
"when": ["zeroStateContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"],
"constants": {
"useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
"workgroupSize": 256,
"hasStateWindow": "windowed",
"usesF16": "tensorDtypes.inputT == \"float16\"",
"inputScalar": "dtypes.T",
"outputScalar": "dtypes.T",
"inputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
"outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
"hasBias": false,
"hasState": false,
"kernelSize": "kernelSize",
"kernelSizePadded": "kernelSizePadded",
"weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T"
},
"passes": [
{
"id": "main",
"name": "CausalConvWithState.Vec4",
"source": {
"shader": "causal-conv-with-state-vec4.wgsl.jinja",
"inputs": { "materializeConvBeforeActivation": false }
},
"bindings": "zeroVec4",
"dispatch": {
"threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)",
"workgroupSize": "constants.workgroupSize"
}
}
]
},
{
"id": "state_bias_vec4",
"priority": 20,
"when": ["stateBiasContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"],
"constants": {
"useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
"workgroupSize": 256,
"hasStateWindow": "windowed",
"usesF16": "tensorDtypes.inputT == \"float16\"",
"inputScalar": "dtypes.T",
"outputScalar": "dtypes.T",
"inputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
"outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
"hasBias": true,
"hasState": true,
"kernelSize": "kernelSize",
"kernelSizePadded": "kernelSizePadded",
"weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T"
},
"passes": [
{
"id": "main",
"name": "CausalConvWithState.Vec4",
"source": {
"shader": "causal-conv-with-state-vec4.wgsl.jinja",
"inputs": { "materializeConvBeforeActivation": false }
},
"bindings": "stateBiasVec4",
"dispatch": {
"threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)",
"workgroupSize": "constants.workgroupSize"
}
}
]
},
{
"id": "bias_no_state_vec4",
"priority": 20,
"when": ["biasNoStateContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"],
"constants": {
"useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
"workgroupSize": 256,
"hasStateWindow": "windowed",
"usesF16": "tensorDtypes.inputT == \"float16\"",
"inputScalar": "dtypes.T",
"outputScalar": "dtypes.T",
"inputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
"outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
"hasBias": true,
"hasState": false,
"kernelSize": "kernelSize",
"kernelSizePadded": "kernelSizePadded",
"weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T"
},
"passes": [
{
"id": "main",
"name": "CausalConvWithState.Vec4",
"source": {
"shader": "causal-conv-with-state-vec4.wgsl.jinja",
"inputs": { "materializeConvBeforeActivation": false }
},
"bindings": "biasNoStateVec4",
"dispatch": {
"threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)",
"workgroupSize": "constants.workgroupSize"
}
}
]
},
{
"id": "state_no_bias_vec4",
"priority": 20,
"when": ["stateNoBiasContract", "kernelSize >= 2", "kernelSize <= 4", "dim(shapes.inputT, 2) >= 4", "dim(shapes.inputT, 2) % 4 == 0"],
"constants": {
"useSilu": "attrs.activation == \"silu\" or attrs.activation == \"swish\"",
"workgroupSize": 256,
"hasStateWindow": "windowed",
"usesF16": "tensorDtypes.inputT == \"float16\"",
"inputScalar": "dtypes.T",
"outputScalar": "dtypes.T",
"inputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
"outputVec4": "\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\"",
"hasBias": false,
"hasState": true,
"kernelSize": "kernelSize",
"kernelSizePadded": "kernelSizePadded",
"weightElem": "(\"vec4<f16>\" if tensorDtypes.inputT == \"float16\" else \"vec4<f32>\") if kernelSize == 4 else dtypes.T"
},
"passes": [
{
"id": "main",
"name": "CausalConvWithState.Vec4",
"source": {
"shader": "causal-conv-with-state-vec4.wgsl.jinja",
"inputs": { "materializeConvBeforeActivation": false }
},
"bindings": "stateNoBiasVec4",
"dispatch": {
"threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * (dim(shapes.outputT, 2) / 4)",
"workgroupSize": "constants.workgroupSize"
}
}
]
},
{
"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"],
"constants": {
"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",
"kernelSize": "kernelSize",
"kernelSizePadded": "kernelSizePadded",
"inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1",
"hasStateWindow": "windowed",
"usesF16": "tensorDtypes.inputT == \"float16\""
},
"passes": [
{
"id": "main",
"name": "CausalConvWithState.TiledLargeKernel",
"source": {
"shader": "causal-conv-with-state-tiled.wgsl.jinja",
"inputs": { "materializeConvBeforeActivation": false }
},
"bindings": "zeroTiled",
"dispatch": {
"workgroups": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), constants.tileSize)"
}
}
]
},
{
"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"],
"constants": {
"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",
"kernelSize": "kernelSize",
"kernelSizePadded": "kernelSizePadded",
"inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1",
"hasStateWindow": "windowed",
"usesF16": "tensorDtypes.inputT == \"float16\""
},
"passes": [
{
"id": "main",
"name": "CausalConvWithState.TiledLargeKernel",
"source": {
"shader": "causal-conv-with-state-tiled.wgsl.jinja",
"inputs": { "materializeConvBeforeActivation": false }
},
"bindings": "stateBiasTiled",
"dispatch": {
"workgroups": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), constants.tileSize)"
}
}
]
},
{
"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"],
"constants": {
"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",
"kernelSize": "kernelSize",
"kernelSizePadded": "kernelSizePadded",
"inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1",
"hasStateWindow": "windowed",
"usesF16": "tensorDtypes.inputT == \"float16\""
},
"passes": [
{
"id": "main",
"name": "CausalConvWithState.TiledLargeKernel",
"source": {
"shader": "causal-conv-with-state-tiled.wgsl.jinja",
"inputs": { "materializeConvBeforeActivation": false }
},
"bindings": "biasNoStateTiled",
"dispatch": {
"workgroups": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), constants.tileSize)"
}
}
]
},
{
"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"],
"constants": {
"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",
"kernelSize": "kernelSize",
"kernelSizePadded": "kernelSizePadded",
"inputTileSize": "tunables.tiledWorkgroupSize * 8 + kernelSizePadded - 1",
"hasStateWindow": "windowed",
"usesF16": "tensorDtypes.inputT == \"float16\""
},
"passes": [
{
"id": "main",
"name": "CausalConvWithState.TiledLargeKernel",
"source": {
"shader": "causal-conv-with-state-tiled.wgsl.jinja",
"inputs": { "materializeConvBeforeActivation": false }
},
"bindings": "stateNoBiasTiled",
"dispatch": {
"workgroups": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * ceilDiv(dim(shapes.outputT, 2), constants.tileSize)"
}
}
]
},
{
"id": "zero_state",
"priority": 0,
"when": ["zeroStateContract"],
"constants": {
"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",
"source": {
"shader": "causal-conv-with-state.wgsl.jinja",
"inputs": { "materializeConvBeforeActivation": false }
},
"bindings": "zeroScalar",
"dispatch": {
"threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))",
"workgroupSize": "constants.workgroupSize"
}
}
]
},
{
"id": "state_bias",
"priority": 0,
"when": ["stateBiasContract"],
"constants": {
"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",
"source": {
"shader": "causal-conv-with-state.wgsl.jinja",
"inputs": { "materializeConvBeforeActivation": false }
},
"bindings": "stateBias",
"dispatch": {
"threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))",
"workgroupSize": "constants.workgroupSize"
}
}
]
},
{
"id": "bias_no_state",
"priority": 0,
"when": ["biasNoStateContract"],
"constants": {
"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",
"source": {
"shader": "causal-conv-with-state.wgsl.jinja",
"inputs": { "materializeConvBeforeActivation": false }
},
"bindings": "biasNoState",
"dispatch": {
"threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))",
"workgroupSize": "constants.workgroupSize"
}
}
]
},
{
"id": "state_no_bias",
"priority": 0,
"when": ["stateNoBiasContract"],
"constants": {
"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",
"source": {
"shader": "causal-conv-with-state.wgsl.jinja",
"inputs": { "materializeConvBeforeActivation": false }
},
"bindings": "stateNoBias",
"dispatch": {
"threads": "dim(shapes.outputT, 0) * dim(shapes.outputT, 1) * max(1, dim(shapes.outputT, 2))",
"workgroupSize": "constants.workgroupSize"
}
}
]
}
]
}