flashinfer.comm.prepare_all_gather_matmul

flashinfer.comm.prepare_all_gather_matmul(inp: Tensor, w: Tensor, group: ProcessGroup, *, backend: str = 'auto', max_rows: int | None = None, verbose: bool = False) → Callable[[Tensor], Tensor]

Prepare a capacity-bound Cake all-gather matmul launcher.

The returned callable binds w and group and accepts any contiguous input with the dtype, device and K of inp and at most max_rows rows (default: inp.shape[0]); the symmetric scratch is sized for that capacity in this one collective, so later calls with fewer rows perform no collective. w is the logical [8192, N] weight, contiguous or the transposed view of a contiguous [N, 8192] parameter, with N a positive multiple of 256. Both backend="auto" and backend="cake" select the Cake route (SM100 or SM103, bfloat16 or float16, two-, four- or eight-rank NCCL groups). Unsupported inputs raise during preparation instead of falling back to another implementation.