flashinfer.comm.mnnvl_cutedsl_ag.all_gather_rmsnorm_fp4_quantize¶

flashinfer.comm.mnnvl_cutedsl_ag.all_gather_rmsnorm_fp4_quantize(shard: Tensor, weight: Tensor, global_scale: Tensor, workspace: MNNVLCuteDSLAllGatherNormQuantWorkspace, *, eps: float, gain_center: float = 0.0) → tuple[Tensor, Tensor]¶

Gather column shards across TP, then RMSNorm and NVFP4-quantize the rows.

With x = concat_r(shard[r], dim=1):

y = bf16(x * rsqrt(mean(x**2) + eps) * (weight + gain_center))
packed, scales = fp4_quantize(
    y, global_scale, sf_vec_size=16, is_sf_swizzled_layout=False
)

Parameters:

shard – [M, H/TP] BF16 columns [rank*H/TP, (rank+1)*H/TP).
weight – [H] BF16 RMSNorm weight.
global_scale – [1] FP32 NVFP4 global scale.
workspace – Workspace created on every rank of the TP group.
eps – RMSNorm epsilon.
gain_center – Constant added to weight.

Returns:

[M, H/2] uint8 E2M1 pairs and [M, H/16] uint8 E4M3 block scales in linear layout.

Return type:

(packed, scales)