flashinfer.comm.cake_fp4_allreduce_rmsnorm

flashinfer.comm.cake_fp4_allreduce_rmsnorm(x: 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

All-reduce x through the NVFP4 transport, add residual and apply the RMSNorm of spec.

The sum is the rank-order float32 sum of the NVFP4-quantised partials (as cake_fp4_allreduce()): z = sum + residual in float32, residual_out = bf16(z), rinv = rsqrt(mean(z^2) + eps) from the float32 sum of squares, y = bf16(z) * rinv * (weight_bias + weight) (the bfloat16-rounded z enters the product). residual_out and the quantised outputs are bitwise identical on every rank.

Parameters:

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

  • residual (torch.Tensor) – Contiguous bfloat16 [rows, 6144] residual stream of this rank (the same tensor on every rank gives rank-identical outputs).

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

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

  • out (torch.Tensor, optional) – Bfloat16 output buffer [rows, 6144] (bfloat16 output only).

  • values (torch.Tensor, optional) – Buffers of a quantised output (see CakeFp4AllReduceFusedNormResult); scales may be larger than needed.

  • scales (torch.Tensor, optional) – Buffers of a quantised output (see CakeFp4AllReduceFusedNormResult); scales may be larger than needed.

  • residual_out (torch.Tensor, optional) – Bfloat16 [rows, 6144] buffer for the summed residual; must not alias residual.

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

Returns:

The output tensors (the given buffers, or freshly allocated ones). The call allocates nothing when every buffer is given and performs no host synchronisation, so it can be captured into a CUDA graph.

Return type:

CakeFp4AllReduceFusedNormResult