flashinfer.cake_mhc_train.cake_mhc_train_head_forward

flashinfer.cake_mhc_train.cake_mhc_train_head_forward(X: Tensor, *, out: Tensor | None = None) → Tensor

The model’s head: the mean of the four residual streams, Y = mean_i X[:, i, :].

Y[t, d] = BF16(((X[t, 0, d] + X[t, 1, d]) + X[t, 2, d] + X[t, 3, d]) * 0.25) with the sum in FP32 from zero in stream order and one rounding: bit for bit torch’s X.mean(dim=1) of a BF16 X, signed zeros and the NaN encoding included.

Parameters:

  • X (torch.Tensor) – BF16 [T, 4, H] contiguous residual streams, H in 4096, 5120 or 7168 (is_cake_mhc_train_head_supported()), 32-byte aligned. T is a runtime scalar; T == 0 returns an empty Y without launching.

  • out (torch.Tensor, optional) – Preallocated BF16 [T, H] contiguous Y (default None: a new tensor); it must not overlap X.

Returns:

  • BF16 Y of shape [T, H].

  • The launch runs on the current torch stream of X’s device and

  • allocates nothing beyond the default out, so it can be captured in a

  • CUDA graph once its program is loaded (one eager call first).