flashinfer.cake_mhc_train.cake_mhc_train_prepare

flashinfer.cake_mhc_train.cake_mhc_train_prepare(X: Tensor, fn: Tensor, scale: Tensor, base: Tensor, *, rms_eps: float = 1e-05, hc_eps: float = 1e-06, sinkhorn_iters: int = 20, save_stats: bool = True, workspace: Tensor | None = None, backward_workspace: Tensor | None = None)

Autograd-enabled forward_prepare: (U, handle); pass handle to cake_mhc_train_finish().

Parameters:

  • X (torch.Tensor) – BF16 [T, 4, 4096] contiguous residual streams (may require gradients).

  • fn (torch.Tensor) – FP32 [24, 16384] contiguous mapping matrix (may require gradients).

  • scale (torch.Tensor) – FP32 [3] group scales (may require gradients).

  • base (torch.Tensor) – FP32 [24] logit offsets (may require gradients).

  • 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()).

  • save_stats (bool) – True (default): the per-token statistics are kept for the backward. False: they are dropped after the forward and recomputed from X in the backward (checkpoint replay); the gradients are bitwise the same.

  • workspace (torch.Tensor, optional) – Caller-owned forward workspace; a per-device default when omitted.

  • backward_workspace (torch.Tensor, optional) – Caller-owned backward workspace; a per-device default when omitted.

Returns:

  • ``(U, handle)`` (BF16 U of shape [T, 4096] and the handle of this)

  • prepare node for one cake_mhc_train_finish(). The backward of the

  • prepare node returns the single-cast dX and FP32 dfn, dscale,

  • dbase.