flashinfer.cake_causal_conv1d_train.cake_causal_conv1d_kcp

flashinfer.cake_causal_conv1d_train.cake_causal_conv1d_kcp(x_local: Tensor, weight: Tensor, ctx: Any, *, recompute_handle: Any | None = None) → Tensor

Autograd-enabled context-parallel packed width-4 depthwise causal convolution + SiLU of this rank’s fragments.

The first tokens of a fragment that continues a document read the history rows of the preceding fragment (exchanged through the context’s halo collective) instead of zeros, so the result equals the single-rank cake_causal_conv1d() of the whole batch on this rank’s tokens. The KCP path takes FP32 taps and no bias (the package’s pre-round-6 contract).

Parameters:

  • x_local (torch.Tensor) – BF16 [S_local, C] or [1, S_local, C] contiguous packed activations of this rank’s fragments, in the partition’s order; C is a multiple of 8.

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

  • ctx (KcpContext) – From flashinfer.kda_train.build_kcp_context (the same context the KDA step of the rank uses).

  • recompute_handle (KcpRecomputeHandle, optional) – Keyword-only, default None (unchanged behavior). A handle (any object with an entries dict, such as flashinfer.cake_causal_conv1d_train.KcpRecomputeHandle) makes the checkpoint recompute of this call run without the forward halo exchange: the original forward exchanges as always and stores its history rows (the saved set’s own tensor, key "conv") in the handle; torch’s recompute of the forward inside the backward reads them back, so y and the gradients are bitwise those of the same checkpoint without a handle. The backward still exchanges. Create one handle per convolution call and micro-step, outside the outermost checkpoint region, and pass it as an argument of the checkpointed function. A forward outside the backward that finds an entry raises (a handle reused across micro-steps). The forward halo exchange is a collective: every rank of ctx.group passes a handle for the same call or none, created at the same checkpoint nesting level. A mismatch pairs one rank’s recompute exchange with another rank’s backward exchange and hangs or corrupts data; no check is made. Never share one handle between two convolution calls (also not in nested regions), and never clear() or reuse a handle while a backward that needs it is pending.

Returns:

BF16 y of x_local’s shape. The backward runs the head program, the backward halo exchange and the fused dx + deterministic FP32 dweight kernel of this rank’s fragments.

Return type:

torch.Tensor