flashinfer.comm.fp8_rowwise_gemm_allreduce

flashinfer.comm.fp8_rowwise_gemm_allreduce(a: Tensor, a_scale: Tensor, w: Tensor, w_scale: Tensor, group: ProcessGroup, *, backend: str = 'auto', out: Tensor | None = None) → Tensor

Fused row-wise FP8 GEMM + all-reduce over group.

backend="cake" selects the exact source-built SM100 fused kernel for four-rank NCCL groups; backend="auto" selects the same route (no other implementation exists). Unsupported inputs raise instead of falling back.

Parameters:

  • a (torch.Tensor) – This rank’s 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) – Contiguous FP32 row scales of shape [M].

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

  • w_scale (torch.Tensor) – 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.

  • out (torch.Tensor, optional) – Contiguous BF16 output buffer of shape [M, N]. A new tensor is allocated when omitted.

Returns:

The reduced BF16 product of shape [M, N] on every rank.

Return type:

torch.Tensor