flashinfer.cake_rmsnorm_train.cake_rmsnorm_train_backward_workspace

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

Allocate a zero-initialized backward workspace (see ..._workspace_bytes).