CausalConvWithState

CausalConvWithState - 27

Version

This version of the operator has been available since version 27.

Summary

Stateful causal 1D depthwise convolution.

Used by Gated DeltaNet (Qwen3.5) and Mamba (Jamba, FalconMamba) as a preprocessing step. Replaces the 3-op pattern (Concat + Conv + Slice) with a single fused operation.

The convolution is causal (looks only at current and past positions) and depthwise (each channel is convolved independently with its own kernel).

The input, weight, past_state, output, and present_state tensors are rank-3 with shape (batch_size, channels, length). The optional bias input is rank-1 with shape (channels). For higher-dimensional data, use Reshape nodes before and after this operator to pack extra dimensions into the batch or channel axis.

Weight layout: (channels, 1, k) for depthwise convolution. The carry state stores the last (k-1) positions for incremental decode.

The optional activation attribute supports fused SiLU/Swish activation.

Attributes

  • activation - STRING (default is none):

    Fused activation function. One of: ‘silu’, ‘swish’, ‘none’. Default is ‘none’.

Inputs

Between 2 and 4 inputs.

  • input (heterogeneous) - T:

    Input tensor with shape (batch_size, channels, length). Channels-first layout.

  • weight (heterogeneous) - T:

    Depthwise convolution kernel with shape (channels, 1, k) where k is the kernel size. The middle dim of size 1 follows the ONNX Conv weight layout (M, C/group, k1, ..., kn): since this op is always depthwise, group = channels, so C/group = 1. Keeping this layout makes the weight tensor a drop-in for a depthwise Conv(group=channels) weight, so Conv <-> CausalConvWithState rewrites require no reshape.

  • bias (optional, heterogeneous) - T:

    Optional per-channel bias with shape (channels).

  • past_state (optional, heterogeneous) - T:

    Carry state from previous step with shape (batch_size, channels, k - 1). If not provided, padding is zero.

Outputs

  • output (heterogeneous) - T:

    Convolution output with same shape as input.

  • present_state (heterogeneous) - T:

    Updated carry state with shape (batch_size, channels, k - 1). Contains the last (k - 1) values of the effective padded/concatenated sequence along the causal axis, including any values from past_state or zero-padding when the current input is shorter than k - 1.

Type Constraints

  • T in ( tensor(bfloat16), tensor(float), tensor(float16) ):

    Constrain input and output types to float tensors.

Examples

_basic

import numpy as np
import onnx

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)

output, present_state = _compute(input_, weight)

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_basic",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

_with_bias

import numpy as np
import onnx

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight", "bias"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)
bias = np.random.randn(channels).astype(np.float32)

output, present_state = _compute(input_, weight, bias=bias)

expect(
    node,
    inputs=[input_, weight, bias],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_with_bias",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

_with_past_state

import numpy as np
import onnx

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight", "", "past_state"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)
past_state = np.random.randn(batch_size, channels, k - 1).astype(np.float32)

output, present_state = _compute(input_, weight, past_state=past_state)

expect(
    node,
    inputs=[input_, weight, past_state],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_with_past_state",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

_silu

import numpy as np
import onnx

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
    activation="silu",
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)

output, present_state = _compute(input_, weight, activation="silu")

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_silu",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

_swish_alias

import numpy as np
import onnx

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
    activation="swish",
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)

output, present_state = _compute(input_, weight, activation="swish")

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_swish_alias",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

_decode_step

import numpy as np
import onnx

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight", "bias", "past_state"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 1, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)
bias = np.random.randn(channels).astype(np.float32)
past_state = np.random.randn(batch_size, channels, k - 1).astype(np.float32)

output, present_state = _compute(
    input_, weight, bias=bias, past_state=past_state
)

expect(
    node,
    inputs=[input_, weight, bias, past_state],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_decode_step",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

_kernel_size_one

import numpy as np
import onnx

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 8, 1
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)

output, present_state = _compute(input_, weight)

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_kernel_size_one",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

_with_bias_and_past_state

import numpy as np
import onnx

# Multi-token (T>1) path through Concat(past, input) -> Conv(+bias).
node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight", "bias", "past_state"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)
bias = np.random.randn(channels).astype(np.float32)
past_state = np.random.randn(batch_size, channels, k - 1).astype(np.float32)

output, present_state = _compute(
    input_, weight, bias=bias, past_state=past_state
)

expect(
    node,
    inputs=[input_, weight, bias, past_state],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_with_bias_and_past_state",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

_silu_with_past_state

import numpy as np
import onnx

# Fused activation combined with concat-from-past variant of PaddedInput.
node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight", "", "past_state"],
    outputs=["output", "present_state"],
    activation="silu",
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)
past_state = np.random.randn(batch_size, channels, k - 1).astype(np.float32)

output, present_state = _compute(
    input_, weight, past_state=past_state, activation="silu"
)

expect(
    node,
    inputs=[input_, weight, past_state],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_silu_with_past_state",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

_b1_c1_degenerate

import numpy as np
import onnx

# Mamba/GDN inner-head edge case: B=1, C=1.
node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 1, 1, 6, 4
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)

output, present_state = _compute(input_, weight)

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_b1_c1_degenerate",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

_short_input_no_past_state

import numpy as np
import onnx

# L < k-1 with no past_state: zero-pad is wider than the input.
node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 2, 5
input_ = np.random.randn(batch_size, channels, length).astype(np.float32)
weight = np.random.randn(channels, 1, k).astype(np.float32)

output, present_state = _compute(input_, weight)

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_short_input_no_past_state",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

_fp16

import numpy as np
import onnx

node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.rand(batch_size, channels, length).astype(np.float16)
weight = np.random.rand(channels, 1, k).astype(np.float16)

output, present_state = _compute(input_, weight)

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_fp16",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)

_silu_fp16

import numpy as np
import onnx

# fp16 + SiLU: the reference upcasts Sigmoid/Mul to float32, so the
# function-body expansion must do the same to stay numerically faithful.
node = onnx.helper.make_node(
    "CausalConvWithState",
    inputs=["input", "weight"],
    outputs=["output", "present_state"],
    activation="silu",
)

batch_size, channels, length, k = 2, 4, 8, 4
input_ = np.random.rand(batch_size, channels, length).astype(np.float16)
weight = np.random.rand(channels, 1, k).astype(np.float16)

output, present_state = _compute(input_, weight, activation="silu")

expect(
    node,
    inputs=[input_, weight],
    outputs=[output, present_state],
    name="test_causal_conv_with_state_silu_fp16",
    opset_imports=[onnx.helper.make_opsetid("", 27)],
)