flashinfer.comm.cake_fp4_allreduce

flashinfer.comm.cake_fp4_allreduce(x: Tensor, workspace: CakeFp4AllReduceWorkspace, *, out: Tensor | None = None, launch_with_pdl: bool = True) → Tensor

Sum x over the workspace’s group through the NVFP4 transport.

Parameters:

  • x (torch.Tensor) – Contiguous bfloat16 [rows, 6144] partial of this rank; rows is at most workspace.capacity, a multiple of 4 for world size 4, and the same on every rank of the call.

  • workspace (CakeFp4AllReduceWorkspace) – The collective workspace of the group (family "allreduce", input format "bf16").

  • out (torch.Tensor, optional) – Contiguous bfloat16 [rows, 6144] output; allocated when omitted.

  • launch_with_pdl (bool) – Launch with the programmatic-stream-serialization attribute (the kernel waits for its predecessor before touching the windows).

Returns:

out: the rank-order sum of the NVFP4-quantised partials, accumulated in float32 and rounded to bfloat16, bitwise identical on every rank. The call allocates nothing when out is given and performs no host synchronisation, so it can be captured into a CUDA graph.

Return type:

torch.Tensor