flashinfer.cake_causal_conv1d_train.cake_causal_conv1d

flashinfer.cake_causal_conv1d_train.cake_causal_conv1d(x: Tensor, weight: Tensor, cu_seqlens: Tensor, bias: Tensor | None = None, *, silu_grad_preactivation: str = 'bf16') → Tensor

Autograd-enabled packed width-4 depthwise causal convolution + SiLU.

Parameters:

  • x (torch.Tensor) – BF16 [S, C] or [1, S, C] contiguous packed activations on a supported CUDA device.

  • weight (torch.Tensor) – FP32 or BF16 [C, 4] contiguous taps.

  • cu_seqlens (torch.Tensor) – int32 [num_docs + 1] device prefix sums of the document lengths.

  • bias (torch.Tensor, optional) – FP32 or BF16 [C] contiguous per-channel bias.

  • silu_grad_preactivation ({"bf16", "fp32"}) – The weight gradient’s SiLU’ policy (see cake_causal_conv1d_train_backward()).

Returns:

BF16 y of x’s shape. The backward saves only x, weight, bias and cu_seqlens and launches the fused dx + deterministic dweight (+ dbias) kernel of cake_causal_conv1d_train_backward(); the parameters’ gradients come in the parameters’ dtypes.

Return type:

torch.Tensor