flashinfer.cake_causal_conv1d_train.cake_causal_conv1d_train_backward

flashinfer.cake_causal_conv1d_train.cake_causal_conv1d_train_backward(x: Tensor, weight: Tensor, cu_seqlens: Tensor, dy: Tensor, *, bias: Tensor | None = None, silu_grad_preactivation: str = 'bf16', workspace: Tensor | None = None, dx: Tensor | None = None, dweight: Tensor | None = None, dbias: Tensor | None = None) → tuple[Tensor, ...]

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

Recomputes the FP32 preactivation from x, evaluates the SiLU derivative at it – rounded to BF16 (silu_grad_preactivation="bf16", the default) or unrounded ("fp32") – and produces dx, the weight gradient and, with a bias, the bias gradient in two ordinary launches on the current stream: the main program writes dx and per-CTA partial rows of the gradients, and a reduce program launched programmatically dependent on it sums the rows in a fixed order (no cooperative grid, so kernels resident on other streams never stall the call).

The call’s arguments select the route (bias, BF16 among weight / dweight, the SiLU’ policy) and the channel count selects the route’s program through a host-side rule of the backend’s registry (never the data), so a given (S, C, device) and route always run the same program. Before capturing a CUDA Graph, build every library of the device’s architecture (the backend’s cake_jit.specs_for_arch) so that no JIT build happens inside the capture; the captured graph holds the program selected for its shape and route.

Parameters:

  • x – The forward’s inputs (same contracts as cake_causal_conv1d_train_forward()); nothing else is saved.

  • weight – The forward’s inputs (same contracts as cake_causal_conv1d_train_forward()); nothing else is saved.

  • cu_seqlens – The forward’s inputs (same contracts as cake_causal_conv1d_train_forward()); nothing else is saved.

  • bias – The forward’s inputs (same contracts as cake_causal_conv1d_train_forward()); nothing else is saved.

  • dy (torch.Tensor) – BF16 gradient of y, x’s shape, contiguous.

  • silu_grad_preactivation ({"bf16", "fp32"}) – "bf16": SiLU’ at the BF16-rounded recomputed preactivation (the boundary flash-linear-attention’s backward uses; dweight within about 3e-5 rel-L2 of the exact FP64 gradient). "fp32": SiLU’ at the unrounded FP32 preactivation with an FP64 final reduction (about 2e-7; a slower program). dx is within one BF16 ulp of exact under both.

  • workspace (torch.Tensor, optional) – Caller-owned uint8 buffer from cake_causal_conv1d_train_backward_workspace() (the per-CTA FP32 dweight / dbias partials; fully written by the kernel, so it needs no zeroing); allocated when omitted. Its size (cake_causal_conv1d_train_backward_workspace_bytes()) is the maximum over every backward program of every route, so one buffer serves whichever program a call selects.

  • dx (torch.Tensor, optional) – Preallocated outputs: BF16 of x’s shape, [C, 4] and [C] (each FP32 or BF16; the parameter’s dtype when omitted), contiguous. dbias needs bias.

  • dweight (torch.Tensor, optional) – Preallocated outputs: BF16 of x’s shape, [C, 4] and [C] (each FP32 or BF16; the parameter’s dtype when omitted), contiguous. dbias needs bias.

  • dbias (torch.Tensor, optional) – Preallocated outputs: BF16 of x’s shape, [C, 4] and [C] (each FP32 or BF16; the parameter’s dtype when omitted), contiguous. dbias needs bias.

Returns:

  • (dx, dweight), with a third element dbias only when bias is

  • given (BF16 dx of x’s shape, ``dx[t, c] = bf16(sum_k weight[c, k] *)

  • dpre[t + 3 - k, c])`` over the tokens of the same document, the [C, 4]

  • dweight[c, k] = sum_t dpre[t, c] * x[t - 3 + k, c] and the [C]

  • dbias[c] = sum_t dpre[t, c], each reduced in a fixed order, so they are

  • bitwise reproducible for a given (S, C, device) and route. S == 0

  • returns the empty dx and zero gradients without launching.