flashinfer.cake_mhc_train.cake_mhc_train_backward_prepare

flashinfer.cake_mhc_train.cake_mhc_train_backward_prepare(dY: Tensor, B: Tensor, X: Tensor, stats: Tensor, *, dB: Tensor | None = None, grad_stats_out: Tensor | None = None) → tuple[Tensor, Tensor]

Branch-output gradient and the per-token gradient statistics (before the branch backward).

Parameters:

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

  • B (torch.Tensor) – BF16 [T, 4096] contiguous branch output of the forward.

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

  • stats (torch.Tensor) – FP32 [T, 64] record of the forward.

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

  • grad_stats_out (torch.Tensor, optional) – Preallocated FP32 [T, 32] contiguous output.

Returns:

  • ``(dB, grad_stats)`` (dB[t, d] = BF16(sum_j post[t, j] dY[t, j, d]) (one)

  • rounding of the full FP32 sum) and the FP32 [T, 32] record

  • dpost[4] | dC[16] | zeros[12] that cake_mhc_train_backward_finish() consumes.