flashinfer.cake_mhc_train.cake_mhc_train_head_backward

flashinfer.cake_mhc_train.cake_mhc_train_head_backward(dY: Tensor, *, dX: Tensor | None = None, accumulate: bool = False) → Tensor

Gradient of the head: dX[t, i, d] = BF16(dY[t, d] * 0.25) for each of the four streams.

Bit for bit torch’s MeanBackward of X.mean(dim=1) (the quarter is exact unless it is subnormal, and then rounds once). With accumulate=True the kernel adds instead, dX[t, i, d] = BF16(dX[t, i, d] + BF16(dY[t, d] * 0.25)), bit for bit torch’s BF16 dX.add_(...) of that gradient (autograd’s accumulation into an existing X.grad).

Parameters:

  • dY (torch.Tensor) – BF16 [T, H] contiguous gradient of Y, H in 4096, 5120 or 7168, 32-byte aligned; T == 0 launches nothing.

  • dX (torch.Tensor, optional) – BF16 [T, 4, H] contiguous gradient of X to write (default None: a new tensor) or, with accumulate=True, to add into (then required); it must not overlap dY.

  • accumulate (bool) – Add into dX instead of overwriting it (default False).

Returns:

  • BF16 dX of shape [T, 4, H] (the caller’s tensor when given).

  • The launch runs on the current torch stream of dY’s device.