flashinfer.comm.prepare_all_gather_matmul

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

Prepare the packed-QKV all-gather matmul launcher.

The returned callable binds w and group and accepts a new input tensor with the same shape, dtype, and device as inp. Both backend="auto" and backend="cake" select the source-built prepared launcher. Unsupported inputs raise during preparation instead of falling back to another implementation.