flashinfer.comm.prepare_fp8_rowwise_gemm_allreduce

flashinfer.comm.prepare_fp8_rowwise_gemm_allreduce(a: Tensor, a_scale: Tensor, w: Tensor, w_scale: Tensor, group: ProcessGroup, *, backend: str = 'auto', max_rows: int | None = None)

Bind group, N, K_l and a row capacity once; return a capturable launcher.

The sample operands fix N and K_l; max_rows (default: the rows of a) sizes the symmetric workspace in this single collective and must be identical on every rank: the preparation checks it in that collective and raises ValueError on every rank when it differs. The returned launcher(a, a_scale, w, w_scale, *, out=None) serves any operands of the bound N and K_l with at most max_rows rows without another collective and is CUDA-graph capturable. Both backend="auto" and backend="cake" select the Cake route (SM100, four-rank NCCL groups, torch’s default CUDA or NVSHMEM symmetric-memory backend). Unsupported inputs raise during preparation instead of falling back to another implementation.

Parameters:

  • a (torch.Tensor) – Sample contiguous FP8 (e4m3) activation shard of shape [M, K_l]. a, a_scale, w, and w_scale must be CUDA tensors on one shared device.

  • a_scale (torch.Tensor) – Sample contiguous FP32 row scales of shape [M].

  • w (torch.Tensor) – Sample contiguous FP8 (e4m3) weight shard of shape [N, K_l].

  • w_scale (torch.Tensor) – Sample contiguous FP32 column scales of shape [N].

  • group (torch.distributed.ProcessGroup) – Four-rank NCCL process group on one SM100 node.

  • backend (str, optional) – "auto" and "cake" both select the Cake backend.

  • max_rows (int, optional) – Maximum row count accepted by the returned launcher. Must be at least the sample input’s row count. Defaults to the sample input’s row count and must match on every rank.

Returns:

A CUDA-graph-capturable launcher with signature launcher(a, a_scale, w, w_scale, *, out=None).

Return type:

Callable