flashinfer.cake_mhc_train.cake_mhc_train_forward_stats

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

Recompute the per-token statistics of forward_prepare (checkpoint replay).

Parameters:

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

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

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

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

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

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

  • workspace (torch.Tensor, optional) – Caller-owned forward workspace (cake_mhc_train_forward_workspace()).

Returns:

  • FP32 [T, 64] stats, bitwise identical to the record

  • cake_mhc_train_forward_prepare() returns for the same inputs on the

  • same device (the same work partition without the U write).