flashinfer.cake_mhc_train.cake_mhc_train_finish

flashinfer.cake_mhc_train.cake_mhc_train_finish(X: Tensor, B: Tensor, handle) → Tensor

Autograd-enabled forward_finish: Y from X, the branch output B and the prepare handle.

Parameters:

• X (torch.Tensor) – The BF16 [T, 4, 4096] tensor passed to cake_mhc_train_prepare().

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

• handle (object) – The handle returned by cake_mhc_train_prepare(), for one finish. A reentrant checkpoint region (use_reentrant=True) holds both calls or starts after the prepare; it does not end between them.

Returns:

• BF16 Y of shape [T, 4, 4096]. Its backward returns dB for the

• branch and no gradient for X (the residual-mixing contribution is part

• of the single-cast dX of the prepare node).