flashinfer.cake_mhc_train.cake_mhc_train_head

flashinfer.cake_mhc_train.cake_mhc_train_head(X: Tensor) → Tensor

Autograd-enabled cake_mhc_train_head_forward(): Y = mean_i X[:, i, :].

Parameters:

X (torch.Tensor) – BF16 [T, 4, H] contiguous residual streams (may require gradients), H in 4096, 5120 or 7168.

Returns:

• BF16 Y of shape [T, H]. Its backward is

• cake_mhc_train_head_backward() (bit for bit torch’s gradient of

• X.mean(dim=1)); the node saves no tensor, so an activation checkpoint

• around it keeps nothing alive and its backward needs no recomputation.