flashinfer.cake_causal_conv1d_det_train_kcp.cake_causal_conv1d_det_multi_kcp

flashinfer.cake_causal_conv1d_det_train_kcp.cake_causal_conv1d_det_multi_kcp(xs: Sequence[Tensor], weights: Sequence[Tensor], ctx: Any, *, recompute_handle: Any | None = None, cooperative: bool | None = None) → tuple[Tensor, ...]

Autograd-enabled context-parallel convolutions of several tensors of one layer (for example its q, k and v) with one halo exchange per direction.

The result is bitwise tuple(cake_causal_conv1d_det_kcp(x, w, ctx) for x, w in zip(xs, weights)): outputs and every gradient, dweight included (each tensor keeps its own kernels and fixed-order reduction). The forward issues one halo collective instead of one per tensor (the history rows of all tensors travel side by side and are copied out after the wait), and so does the backward, so a layer has fewer global synchronization points.

Parameters:

  • xs (sequence of torch.Tensor) – n >= 1 BF16 [S_local, C_i] or [1, S_local, C_i] contiguous packed activations of this rank’s fragments, in the partition’s order; each C_i is a multiple of 8 (they may differ).

  • weights (sequence of torch.Tensor) – FP32 [C_i, 4] contiguous taps, one per tensor.

  • ctx (ContextDetKcp) – As in cake_causal_conv1d_det_kcp(). Every rank of the group calls this function once per layer, in the same order, with the same tensor count and channel counts. The ranks are not compared (that would cost a collective): a different count, channel count or order hangs the exchange or silently corrupts the history rows, as a mismatch of the single entry’s per-rank contract does.

  • recompute_handle (RecomputeHandleDetKcp, optional) – Keyword-only, default None (unchanged behavior). One handle for the whole call: the original forward stores the n history pools under "conv" and torch’s recompute inside the backward reads them back, so the recompute issues no collective; the backward still exchanges once. The rules of cake_causal_conv1d_det_kcp() apply (one handle per call and micro-step, created outside the outermost checkpoint region; every rank passes one or none). The same handle may carry the layer’s "kda" entry; its "conv" entry serves one multi call or one single call, never both.

  • cooperative (bool, optional) – As in cake_causal_conv1d_det_kcp(), one form for every tensor of the call; the result is bitwise the single calls with the same cooperative.

Returns:

The n BF16 outputs, each of its input’s shape.

Return type:

tuple of torch.Tensor