flashinfer.cake_mhc_train.cake_mhc_train_forward_finish

flashinfer.cake_mhc_train.cake_mhc_train_forward_finish(X: Tensor, B: Tensor, stats: Tensor, *, out: Tensor | None = None) → Tensor

Residual mixing: Y[t, j, d] = BF16(post[t, j] B[t, d] + sum_i C[t, i, j] X[t, i, d]).

Parameters:

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

  • B (torch.Tensor) – BF16 [T, 4096] contiguous output of the external branch applied to U.

  • stats (torch.Tensor) – FP32 [T, 64] record returned by cake_mhc_train_forward_prepare() (or cake_mhc_train_forward_stats()).

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

Return type:

BF16 Y of shape [T, 4, 4096]; the FP32 sum is rounded once.