flashinfer.cake_mhc_train.cake_mhc_train_forward_prepare

flashinfer.cake_mhc_train.cake_mhc_train_forward_prepare(X: Tensor, fn: Tensor, scale: Tensor, base: Tensor, *, rms_eps: float = 1e-05, hc_eps: float = 1e-06, sinkhorn_iters: int = 20, out: Tensor | None = None, stats_out: Tensor | None = None, workspace: Tensor | None = None) → tuple[Tensor, Tensor]

Forward mapping, coefficients and collapse: U and the per-token statistics.

Parameters:

  • X (torch.Tensor) – BF16 [T, 4, 4096] contiguous residual streams (T is a runtime scalar; T == 0 returns empty outputs without launching).

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

  • scale (torch.Tensor) – FP32 [3] group scales of the pre / post / Sinkhorn logits.

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

  • rms_eps (float) – Normalization epsilon (default 1e-5): a launch argument of the exported kernels, so any finite positive value runs on the same compiled programs.

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

  • sinkhorn_iters (int) – Sinkhorn iteration count, 1 to 64 (default 20). The default runs the compiled programs, whose chain is unrolled for it; every other count runs the run-time-count programs of the coefficient and token-owner passes. A count above 64 raises NotImplementedError.

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

  • stats_out (torch.Tensor, optional) – Preallocated FP32 [T, 64] contiguous statistics record.

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

Returns:

  • ``(U, stats)`` (BF16 U[t, d] = BF16(sum_i pre[t, i] X[t, i, d]) and the)

  • FP32 [T, 64] per-token record r | z[24] | pre[4] | post[4] | C[16] | pad[15]

  • (256 bytes per token) that the finish and backward entry points consume.