flashinfer.cake_causal_conv1d_train.cake_causal_conv1d_train_forward

flashinfer.cake_causal_conv1d_train.cake_causal_conv1d_train_forward(x: Tensor, weight: Tensor, cu_seqlens: Tensor, *, bias: Tensor | None = None, out: Tensor | None = None, workspace: Tensor | None = None) → Tensor

Packed width-4 depthwise causal convolution + SiLU, training forward.

Parameters:

  • x (torch.Tensor) – BF16 [S, C] or [1, S, C] contiguous packed activations of every document; C is a multiple of 8.

  • weight (torch.Tensor) – FP32 or BF16 [C, 4] contiguous taps; weight[c, 3] multiplies the current token, weight[c, 0] the token three positions back. Any other width is rejected.

  • cu_seqlens (torch.Tensor) – int32 [num_docs + 1] contiguous device prefix sums of the document lengths (cu_seqlens[0] == 0, cu_seqlens[-1] == S; a repeated value is an empty document). The host reads its metadata only: no element is copied to the host and no stream is synchronized.

  • bias (torch.Tensor, optional) – FP32 or BF16 [C] contiguous per-channel bias, added after the four taps; its dtype is independent of weight’s.

  • out (torch.Tensor, optional) – Preallocated BF16 output of x’s shape, contiguous.

  • workspace (torch.Tensor, optional) – Caller-owned int32 buffer from cake_causal_conv1d_train_forward_workspace() (the token-chunk counters: zero before its first launch, left zero by every launch, so one buffer serves every later call of one stream, CUDA Graph replays included; two launches that may run at once need two buffers). Its size (cake_causal_conv1d_train_forward_workspace_bytes()) is the most over every forward route. When omitted the package uses a zeroed buffer it keeps per (device, current stream), or a zeroed buffer of the call’s own inside a CUDA Graph capture.

Returns:

BF16 y of x’s shape: y[t, c] = bf16(silu(sum_k weight[c, k] * x[t - 3 + k, c] + bias[c])) with zero history before each document’s first token. S == 0 returns the empty output without launching.

Return type:

torch.Tensor