flashinfer.comm.cake_fp4_allreduce_nvfp4

flashinfer.comm.cake_fp4_allreduce_nvfp4(records: Tensor, global_scales: Tensor, workspace: CakeFp4AllReduceWorkspace, *, out: Tensor | None = None, launch_with_pdl: bool = True) → Tensor

Sum pre-quantised NVFP4 records over the workspace’s group.

Parameters:

  • records (torch.Tensor) – Contiguous uint8 [rows, 3456] records of this rank: 3072 packed E2M1 value bytes (element 2i in the low nibble) followed by 384 E4M3 block scales (one per 16 values). Record r decodes as e2m1(value) * e4m3(scale) * global_scales[r].

  • global_scales (torch.Tensor) – Contiguous float32 tensor of at least world_size elements holding the static scale of every rank’s records; identical on every rank.

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

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

  • launch_with_pdl (bool) – Launch with the programmatic-stream-serialization attribute.

Returns:

out: the rank-order float32 sum of the decoded records rounded to bfloat16, bitwise identical on every rank.

Return type:

torch.Tensor