flashinfer.comm.cake_fp4_allreduce_rmsnorm_nvfp4

flashinfer.comm.cake_fp4_allreduce_rmsnorm_nvfp4(records: Tensor, global_scales: Tensor, residual: Tensor, spec: CakeFp4AllReduceFusedNormSpec, workspace: CakeFp4AllReduceWorkspace, *, out: Tensor | None = None, values: Tensor | None = None, scales: Tensor | None = None, residual_out: Tensor | None = None, launch_with_pdl: bool = True) → CakeFp4AllReduceFusedNormResult

Sum pre-quantised NVFP4 records, add residual and apply the RMSNorm of spec.

The sum is the rank-order float32 sum of the decoded records (as cake_fp4_allreduce_nvfp4()); the epilogue is that of cake_fp4_allreduce_rmsnorm(), with every output precision available.

Parameters:

  • records (torch.Tensor) – Contiguous uint8 [rows, 3456] records of this rank, encoded as for cake_fp4_allreduce_nvfp4(); rows is at most workspace.capacity and the same on every rank.

  • 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.

  • residual (torch.Tensor) – Contiguous bfloat16 [rows, 6144] residual stream of this rank.

  • spec (CakeFp4AllReduceFusedNormSpec) – Weight, eps, weight bias, output precision and (NVFP4) global scale.

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

  • out (torch.Tensor, optional) – Output buffers as for cake_fp4_allreduce_rmsnorm().

  • values (torch.Tensor, optional) – Output buffers as for cake_fp4_allreduce_rmsnorm().

  • scales (torch.Tensor, optional) – Output buffers as for cake_fp4_allreduce_rmsnorm().

  • residual_out (torch.Tensor, optional) – Output buffers as for cake_fp4_allreduce_rmsnorm().

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

Returns:

The output tensors (the given buffers, or freshly allocated ones).

Return type:

CakeFp4AllReduceFusedNormResult