flashinfer.cake_causal_conv1d_det_train_kcp.cake_causal_conv1d_det_kcp

flashinfer.cake_causal_conv1d_det_train_kcp.cake_causal_conv1d_det_kcp(x_local: Tensor, weight: Tensor, ctx: Any, *, recompute_handle: Any | None = None, cooperative: bool | 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 flashinfer.cake_causal_conv1d_det_train.cake_causal_conv1d_det() 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 (ContextDetKcp) – From flashinfer.cake_kda_det_train_kcp.build_context_det_kcp (the same context the KDA step of the rank uses).

  • recompute_handle (RecomputeHandleDetKcp, optional) – Keyword-only, default None (unchanged behavior). A handle (any object with an entries dict, such as flashinfer.cake_causal_conv1d_det_train_kcp.RecomputeHandleDetKcp) 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.

  • cooperative (bool, optional) – Keyword-only, default None (the package default, flashinfer.cake_causal_conv1d_det_train_backend.kcp.COOPERATIVE_DEFAULT, currently True). The launch form of the backward’s main kernel; the forward is the same either way. True: one cooperative launch, which starts only when all its CTAs fit on the GPU at once, so it waits for any kernel of another stream (for example a communication kernel) that occupies an SM. False: an ordinary launch one SM short of the co-resident CTA count, then a fixed-order reduce kernel on the same stream; it starts beside a small kernel of another stream. dx is bitwise the same in both forms. dweight is reproducible and within the same FP64-oracle bound in each form, but its bits differ between the forms, so keep one form for runs whose dweight bits must match. Each rank may choose its form independently (no collective depends on it).

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