flashinfer.comm.tp_moe_finalize_rmsnorm

flashinfer.comm.tp_moe_finalize_rmsnorm(gemm2_permuted: Tensor, expert_weights: Tensor, expanded_idx_to_permuted_idx: Tensor, norm_weight: Tensor, output: Tensor, *, eps: float, gain_center: float, weight_before_norm: bool) → None

TP MoE finalize plus per-expert RMSNorm, for hidden size 4096 and top-k 8.

gemm2_permuted contains the raw, unweighted GEMM2 rows in physical expert order. expanded_idx_to_permuted_idx[t, k] selects the physical row for logical token t and expert slot k; a negative or out-of-range index contributes zero. The input may have a padded trailing dimension, but its first 4096 columns and every row start must be contiguous and 16-byte aligned. expert_weights may be BF16 or FP32; either dtype is converted to FP32 registers and remains FP32 through the routing reduction.

Initialize the workspace collectively with init_tp_moe_finalize_rmsnorm() before the first invocation. The operation writes its BF16 result to output. All invocations share one process-wide Lamport workspace and must not overlap, including CUDA graph replays. Every rank must issue the same call sequence. When switching streams, the caller must order their work with synchronization or event dependencies; this operation does not insert synchronization between streams. Warm up every specialization outside CUDA graph capture before capturing it.

Parameters:

  • gemm2_permuted (torch.Tensor) – BF16 GEMM2 output of shape (rows, >= 4096) in physical expert order. Rows may be padded, but the last dimension must have unit stride and the row stride must be a multiple of 8 elements.

  • expert_weights (torch.Tensor) – BF16 or FP32 routing weights of shape (num_tokens, 8).

  • expanded_idx_to_permuted_idx (torch.Tensor) – Int32 tensor of shape (num_tokens, 8) mapping each (token, top-k slot) to a row of gemm2_permuted.

  • norm_weight (torch.Tensor) – BF16 RMSNorm weight of shape (4096,).

  • output (torch.Tensor) – BF16 destination of shape (num_tokens, 4096).

  • eps (float) – Epsilon added to the RMS denominator.

  • gain_center (float) – Constant added to norm_weight in FP32.

  • weight_before_norm (bool) – Whether norm_weight scales the activations before the RMS reduction instead of after normalization.

  • gemm2_permuted) (All tensors must be contiguous (except the padded)

:param : :param 16-byte aligned: :param and on the device the workspace was initialized on.: