b2b_causal_conv1d#

b2b_causal_conv1d(x, weight_proj, weight_mixer, skip_bias)#

Back-to-back causal 1D convolution. Fused kernel performing projection convolution, pre-gating, mixer convolution, and post-gating. The operation is performed in a causal manner, meaning each position only attends to previous positions in the sequence. In code terms,

y_gated = b2b_causal_conv1d(x, weight_proj, weight_mixer, skip_bias)

is equivalent to,

y = conv1d_proj(x)
z = y[:,1::3, :] * y[:, 2::3, :]
y_gated = mixer(z) + mixer.skip_bias * z
y = y[:, ::3, :] * y_gated

Note

The input tensor is expected to be of shape (batch_size, 3*dim, seq_len), and its channels are read interleaved: channel 3*c gates the block output, channels 3*c + 1 and 3*c + 2 form the pre-gate product for channel c.

Both convolutions correlate: they apply the taps unflipped over a causally left-padded input, so output position t of channel c is sum_k w[c, k] * x[c, t - (K - 1) + k]. This is the convention the native kernel used as well – the two were measured bit-identical, forward and all four gradients, across the whole native kernel-width allowlist – so routing to cuDNN changed no caller’s numerics.

Because the operator correlates rather than convolves, mixer weights shared with an FFT-based convolution must be flipped along the last dimension for that other consumer. The flip is not applied by this operator and never was:

weight_mixer_for_fft = torch.flip(weight_mixer, [-1])

Note

Requires cuDNN >= 9.24 and nvidia-cudnn-frontend >= 1.27.0. An environment that cannot serve the kernel raises RuntimeError or NotImplementedError naming the GPU, the compute capability and both versions. FP64 and non-CUDA inputs raise ValueError; there is no CPU or native fallback.

Parameters:

  • x (torch.Tensor) – Input tensor of shape (batch_size, 3*dim, seq_len).

  • weight_proj (torch.Tensor) – Projection weight tensor of shape (3 * dim, kernel_size_proj), with kernel_size_proj in [2, 32].

  • weight_mixer (torch.Tensor) – Mixer weight tensor of shape (dim, kernel_size_mixer), with kernel_size_mixer in [2, 256].

  • skip_bias (torch.Tensor) – Skip bias tensor of shape (dim,).

Returns:

Output tensor of shape (batch_size, dim, seq_len).

Return type:

torch.Tensor

Raises:

  • ValueError – an input is FP64, is not on CUDA, has the wrong rank or shape, or asks for a kernel width outside cuDNN’s envelope.

  • NotImplementedError – the installed frontend does not export the B2B bindings.

  • RuntimeError – the frontend or the cuDNN runtime is missing or too old, or the kernel reported an execution fault.

Example

batch_size, dim, seq_len, kernel_size = 2, 4, 10, 3
x = torch.randn(batch_size, 3*dim, seq_len, device="cuda")
weight_proj = torch.randn(3*dim, kernel_size, device="cuda")
weight_mixer = torch.randn(dim, kernel_size, device="cuda")
skip_bias = torch.randn(dim, device="cuda")
y_gated = b2b_causal_conv1d(x, weight_proj, weight_mixer, skip_bias)
print(y_gated.shape)  # torch.Size([2, 4, 10])
On this page