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 BF16 launcher for TP8/N=1280 on SM100 or SM103, or TP4/N=2560 on SM103. SM100 TP8 uses asynchronous peer copies; the fused peer-copy specialization remains specific to SM103. Unsupported inputs raise during preparation instead of falling back to another implementation.