flashinfer.cake_rmsnorm_train.cake_rmsnorm_train_backward_workspace_bytes

flashinfer.cake_rmsnorm_train.cake_rmsnorm_train_backward_workspace_bytes(rows: int, hidden: int, device: device, *, residual: bool = False) → int

Bytes of the caller-owned backward workspace of width hidden on device.

The workspace holds the FP32 per-chunk dw partials (one row per CTA of the persistent backward grid, sized for the device’s largest chunk count so it is independent of rows) followed by the uint32 completion counters. Counters must be zero on first use (allocate with torch.zeros or use cake_rmsnorm_train_backward_workspace()); the kernels leave them zero, so the same buffer is reusable across calls, any token count and CUDA Graph replays without host writes.