[TLERaw]Add AllReduce_RMSNorm with NVSHMEM

[zhendongmiao]

Kernel fusion of AllReduce and RMSNorm on Hopper, eliminating the intermediate HBM round‑trip of separate NCCL AllReduce + FlashInfer RMSNorm.

Two modes:
norm‑only – outputs only the normalized result.
dual‑output – writes both residual_out and norm_out (matches production fused‑add‑RMSNorm).

run

export NVSHMEM_HOME=/xxx/lib/python3.12/site-packages/nvidia/nvshmem

# 04-fuse_ar_rmsnorm
torchrun --nproc_per_node=8 python/tutorials/tle/raw/nvshmem/04-fuse_ar_rmsnorm/fuse_ar_rmsnorm.py

performance

config

• 1 node
• 8 GPUs, NVIDIA H20
• dtype: torch.float16

FlagTree

Tokens	04 norm	04 dual	NCCL+FI
8	52.688	44.768	67.952
32	54.752	52.448	66.656
128	53.408	47.808	67.600
512	62.688	66.656	85.552
4096	171.744	251.040	252.832
8192	304.160	470.320	419.056
16384	557.760	908.912	741.952
32768	1048.000	1743.856	1393.648