flashinfer.comm.mnnvl_cutedsl_ag.all_gather_gated_residual_next_rmsnorm

flashinfer.comm.mnnvl_cutedsl_ag.all_gather_gated_residual_next_rmsnorm(branch_shard: Tensor, residual: Tensor, alpha: Tensor, beta: Tensor, next_weight: Tensor, workspace: MNNVLCuteDSLAllGatherDualNormWorkspace, *, post_eps: float, next_eps: float, next_gain_center: float = 0.0, post_weight: Tensor | None = None, post_gain_center: float = 0.0) → tuple[Tensor, Tensor]

Gather column shards across TP, then the gated residual and next RMSNorm.

With branch = concat_r(branch_shard[r], dim=1) and b = branch * (post_weight + post_gain_center) (or branch):

residual_out = alpha * residual + (beta * b) * rsqrt(mean(b**2) + post_eps)
ffn_input = (
    residual_out
    * rsqrt(mean(residual_out**2) + next_eps)
    * (next_weight + next_gain_center)
)
Parameters:

  • branch_shard – [M, H/TP] BF16 columns [rank*H/TP, (rank+1)*H/TP).

  • residual – [M, H] FP32 stream, identical on every rank.

  • alpha – [H] FP32 stream gate.

  • beta – [H] FP32 branch gate.

  • next_weight – [H] BF16 weight applied after the next normalization.

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

  • post_eps – Epsilon of the branch RMSNorm.

  • next_eps – Epsilon of the next RMSNorm.

  • next_gain_center – Constant added to next_weight.

  • post_weight – Optional [H] BF16 weight applied before the branch norm.

  • post_gain_center – Constant added to post_weight.

Returns:

(residual_out, ffn_input).