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 for up to rows tokens.

The workspace holds the uint32 chunk-claim and completion counters in its first cake_jit.COUNTERS_BYTES bytes, then the FP32 per-chunk dw partial rows: one row per chunk of the backward chunk table, whose count is bounded by min(rows, MAX_ROUNDS x slots) of the route’s largest persistent grid (cake_jit.chunk_capacity). The size grows with rows and a buffer sized for rows serves every call with at most that many tokens. The 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 and CUDA Graph replays without host writes.