fft_causal_conv1d_packed#

fft_causal_conv1d_packed(
x,
weight,
cu_seqlens,
max_seqlen,
min_seqlen=1,
)#

Causal depthwise FFT convolution over a packed (variable-length) batch.

Each sequence is convolved independently: for sequence s with off = cu_seqlens[s] and L = cu_seqlens[s + 1] - off, y[:, off : off + L] == linear_conv(x[:, off : off + L], weight)[:, :L]. Nothing is read or written outside a sequence’s own span, so there is no cross-sequence bleed. Sequences are grouped by the FFT size they need (bucket_fft_size() – the clamped rule, not next_fft_size(), which is not total over the supported domain), and short sequences do not pay for the longest sequence in the pack.

Parameters:

  • x (torch.Tensor) – Packed input of shape (hidden_dim, total_seq_len).

  • weight (torch.Tensor) – Weight tensor of shape (hidden_dim, filter_len), unpadded.

  • cu_seqlens (torch.Tensor) – int32 CUDA tensor of shape (num_seqs + 1,) holding the cumulative sequence offsets, cu_seqlens[0] == 0.

  • max_seqlen (int) – Maximum L_s in the pack. Required, flash-attention style: it is the only host-side handle for pruning the bucket launch list and for rejecting L_s > 16384 instead of silently truncating.

  • min_seqlen (int) – Minimum L_s in the pack. Optional; prunes the low buckets.

Returns:

Output of shape (hidden_dim, total_seq_len), same dtype as x.

Return type:

torch.Tensor

Note

Differentiable w.r.t. x and weight; cu_seqlens and the two length scalars get None. dweight accumulates over every sequence in the pack and over every bucket, deterministically – two identical runs give bit-identical dweight. Supported dtypes: float32, bfloat16, float16.

Note

Preconditions not checkable on the host (checking them would need a device sync): cu_seqlens must be non-decreasing, cu_seqlens[0] == 0, cu_seqlens[-1] == total_seq_len, and every L_s must lie in [min_seqlen, max_seqlen]. y is not zero-initialized – it is fully covered only when cu_seqlens partitions [0, total_seq_len). A caller violating this gets uninitialized output, not a fault.

Example

hidden_dim, filter_len = 128, 128
lengths = [512, 1000, 37]
cu_seqlens = torch.tensor([0, 512, 1512, 1549], dtype=torch.int32, device="cuda")
x = torch.randn(hidden_dim, 1549, device="cuda")
weight = torch.randn(hidden_dim, filter_len, device="cuda")
y = fft_causal_conv1d_packed(x, weight, cu_seqlens, max_seqlen=max(lengths), min_seqlen=min(lengths))
print(y.shape)  # torch.Size([128, 1549])