4. Causal Depthwise 1D Convolution

Problem

Before the SSM scan runs, Mamba's block first mixes each channel with a short window of its own recent history via a causal depthwise 1D convolution (kernel size 4 in the reference implementation) — cheap local context that lets the subsequent per-channel recurrence see more than one input sample's worth of information per state update, at a fraction of the cost of attention or a full (non-depthwise) convolution.

Two properties matter: "depthwise" means each channel is convolved with its own kernel and channels never mix (unlike a standard convolution, which sums over input channels); "causal" means the output at time t must depend only on inputs at time <= t, which for a real-time voice model isn't optional — the model can't see the future. That means padding K-1 zeros on the left of the sequence only, never symmetric padding.

Source: src/4_causal_depthwise_conv1d.py

def causal_depthwise_conv1d(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
    ...

Examples:
>>> x = torch.tensor([[[1.0], [2.0], [3.0], [4.0]]])
>>> weight = torch.tensor([[1.0, 1.0]])  # causal moving sum of current + previous
>>> causal_depthwise_conv1d(x, weight)
tensor([[[1.], [3.], [5.], [7.]]])

Step-by-Step Approach

1. Note the shape mismatch: x comes in as (batch, length, dim) (channels-last, the natural layout for a sequence model), but F.conv1d expects channels-first: (batch, channels, length). Transpose first.
2. Left-pad the transposed input by kernel_size - 1 along the length axis using F.pad(x_t, (K-1, 0)) — the (K-1, 0) argument means "pad K-1 on the left, 0 on the right," which is what makes the convolution causal.
3. Reshape the per-channel weight from (dim, kernel_size) to the 3D shape F.conv1d expects for a depthwise kernel: (dim, 1, kernel_size), via weight.unsqueeze(1).
4. Call F.conv1d(x_pad, weight.unsqueeze(1), bias=bias, groups=dim) — groups=dim is what makes this depthwise: each of the dim input channels is convolved with its own single-channel kernel, with zero cross-channel mixing.
5. Transpose the output back to (batch, length, dim) to match the input layout.
6. Verify the output length matches the input length exactly (no shrinkage) — left-padding by exactly K-1 and using no additional stride/dilation guarantees this; if the lengths don't match, the padding amount is wrong.

The key insight is that "depthwise" and "causal" are two independent, orthogonal properties of the convolution — groups=dim controls the first (no channel mixing), and the padding argument controls the second (no future leakage) — and it's easy to get one right while silently breaking the other (e.g. defaulting to symmetric padding, which is fine for depthwise-ness but breaks causality).

Reference solution

def causal_depthwise_conv1d(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None) -> torch.Tensor:
    batch, length, dim = x.shape
    kernel_size = weight.shape[1]
    x_t = x.transpose(1, 2)  # (batch, dim, length) -- conv1d wants channels before length
    x_pad = F.pad(x_t, (kernel_size - 1, 0))  # left-pad only: causal, no peeking at future
    # groups=dim makes this depthwise: each channel convolved with its own (1, kernel_size) filter
    out = F.conv1d(x_pad, weight.unsqueeze(1), bias=bias, groups=dim)
    return out.transpose(1, 2)  # back to (batch, length, dim)

Key Functions & Tricks

• F.conv1d(input, weight, bias, groups=dim) — groups equal to the channel count makes it fully depthwise (one filter per channel)
• F.pad(x, (left, right)) — asymmetric 1D padding; (K-1, 0) is the causal-only-look-backward pattern
• tensor.transpose(1, 2) — swap length/channel axes between the model's channels-last convention and conv1d's channels-first requirement
• weight.unsqueeze(1) — reshape (dim, kernel_size) to the (out_channels, in_channels/groups, kernel_size) shape conv1d expects, with in_channels/groups=1 for depthwise
• Left-padding vs. symmetric padding — the single detail distinguishing a causal convolution from a standard one

How to Recognize This Pattern