flashinfer.comm.mnnvl_cutedsl_moe_finalize.moe_finalize_rmsnorm_allreduce

flashinfer.comm.mnnvl_cutedsl_moe_finalize.moe_finalize_rmsnorm_allreduce(routed_rows: Tensor, route_weights: Tensor, permuted_idx: Tensor, norm_weight: Tensor, output: Tensor, workspace: MNNVLCuteDSLMoEFinalizeAllReduceWorkspace, *, eps: float, gain_center: float = 0.0, weight_before_norm: bool = True) → None

Per-expert RMSNorm finalize of local routes, all-reduced into output.

For token t and route k with local row x = routed_rows[p] (p = permuted_idx[t, k]; routes with p < 0 or p >= routed_rows.shape[0] belong to other ranks):

s = x * (norm_weight + gain_center)    # if weight_before_norm
partial[t] = sum_k route_weights[t, k] * s * rsqrt(mean(s**2) + eps)
output[t] = sum over ranks of bf16(partial[t])

Without weight_before_norm the weight multiplies after the norm.

Parameters:

  • routed_rows – [rows, >=hidden] BF16 GEMM2 output in permuted order.

  • route_weights – [M, top_k] BF16 or FP32 routing weights.

  • permuted_idx – [M, top_k] int32 row of each route on this rank.

  • norm_weight – [hidden] BF16 RMSNorm weight.

  • output – [M, hidden] BF16 destination, identical on every rank.

  • workspace – Workspace created on every rank of the group.

  • eps – RMSNorm epsilon.

  • gain_center – Constant added to norm_weight.

  • weight_before_norm – Whether the weight scales the row before the norm.