flashinfer.cake_mhc_train.cake_mhc_train_backward_finish

flashinfer.cake_mhc_train.cake_mhc_train_backward_finish(X: Tensor, dY: Tensor, dU: Tensor, B: Tensor | None, fn: Tensor, scale: Tensor, base: Tensor, stats: Tensor, grad_stats: Tensor, *, rms_eps: float = 1e-05, hc_eps: float = 1e-06, sinkhorn_iters: int = 20, accumulate: bool = False, workspace: Tensor | None = None, dX: Tensor | None = None, dfn: Tensor | None = None, dscale: Tensor | None = None, dbase: Tensor | None = None) → tuple[Tensor, Tensor, Tensor, Tensor]

Input gradient and parameter gradients (after the branch backward produced dU).

Parameters:

  • X (torch.Tensor) – BF16 [T, 4, 4096] contiguous residual streams of the forward.

  • dY (torch.Tensor) – BF16 [T, 4, 4096] contiguous gradient of Y.

  • dU (torch.Tensor) – BF16 [T, 4096] contiguous gradient of U from the branch backward.

  • B (torch.Tensor, optional) – BF16 [T, 4096] branch output of the forward; accepted for the prepare / finish symmetry and validated when given, not read.

  • fn (torch.Tensor) – FP32 [24, 16384] contiguous mapping matrix of the forward.

  • scale (torch.Tensor) – FP32 [3] group scales of the forward.

  • base (torch.Tensor) – FP32 [24] logit offsets of the forward.

  • stats (torch.Tensor) – FP32 [T, 64] record of the forward (saved, or recomputed with cake_mhc_train_forward_stats(); both give bitwise the same gradients).

  • grad_stats (torch.Tensor) – FP32 [T, 32] record returned by cake_mhc_train_backward_prepare().

  • rms_eps (float) – Normalization epsilon (default 1e-5; a launch argument).

  • hc_eps (float) – Coefficient epsilon (default 1e-6; a launch argument).

  • sinkhorn_iters (int) – Sinkhorn iteration count, 1 to 64 (default 20; as in cake_mhc_train_forward_prepare()).

  • accumulate (bool) – False (default): dfn, dscale and dbase are written. True: the caller passes all three and the fixed-order sums are added to their current values (reproducible across consecutive accumulations).

  • workspace (torch.Tensor, optional) – Caller-owned uint8 buffer from cake_mhc_train_backward_workspace(); a per-device default buffer is used when omitted.

  • dX (torch.Tensor, optional) – Preallocated BF16 [T, 4, 4096] contiguous output.

  • dfn (torch.Tensor, optional) – Preallocated FP32 [24, 16384] contiguous output (required with accumulate=True).

  • dscale (torch.Tensor, optional) – Preallocated FP32 [3] output (required with accumulate=True).

  • dbase (torch.Tensor, optional) – Preallocated FP32 [24] output (required with accumulate=True).

Returns:

  • ``(dX, dfn, dscale, dbase)`` (the BF16 dX sums the mapping, collapse and)

  • residual-mixing contributions in FP32 and rounds once; the FP32 parameter

  • gradients are fixed-order reductions. With T == 0 nothing is launched

  • dX is empty and the parameter gradients are zeroed (accumulate=False)

  • or left untouched (accumulate=True).