[TLERaw] Add NVSHMEM allreduce RMSNorm tutorial

[Fucov]

Summary

This PR adds a TLE-Raw NVSHMEM tutorial for a fused distributed operator with the following contract:

reduced      = AllReduce_SUM(x)
residual_out = reduced + residual
norm_out     = RMSNorm(residual_out, weight, eps)

The tutorial demonstrates how to integrate a CUDA/NVSHMEM device-side collective into FlagTree through TLE-Raw, followed by a Triton residual + RMSNorm epilogue.

The implementation is added under:

python/tutorials/tle/raw/nvshmem/05-allreduce-rmsnorm/
├── allreduce-device.cu
├── allreduce-host.cu
├── allreduce.py
└── benchmark.py

Motivation

The existing NVSHMEM TLE-Raw tutorials demonstrate lower-level communication primitives such as CUDA-IPC AllReduce.

This tutorial extends that path to a representative LLM distributed fusion pattern:

tensor-parallel AllReduce
        ↓
residual add
        ↓
RMSNorm

The goal is to provide a compact example showing how TLE-Raw can be used to combine:

• CUDA/NVSHMEM device functions,

• symmetric communication workspace,

• distributed synchronization,

• Triton-side tensor computation,

• correctness validation, and

• reproducible operator benchmarking.

This PR is intended as a TLE-Raw integration and optimization example rather than as a universal replacement for an existing production collective implementation.

Implementation

1. TLE-Raw NVSHMEM AllReduce

allreduce-device.cu implements the device-side BF16 collective path used by the tutorial.

The native CUDA functions are registered from Python through the current TLE-Raw @dialect(...) interface and invoked from the Triton path with tle_raw.call(...).

The implementation uses rank-visible communication buffers and explicit synchronization required by the NVSHMEM communication path.

2. Row-parallel execution

The collective stage maps token rows to independent CTAs.

Instead of serializing multiple token rows through one CTA, the implementation uses approximately:

one CTA per token row

so the available parallelism scales with the token dimension.

This is especially important when moving from decode-sized inputs to larger token counts.

3. Residual + RMSNorm epilogue

After the distributed reduction completes, the epilogue computes:

residual_out = reduced + residual
norm_out     = RMSNorm(residual_out, weight, eps)

with FP32 accumulation where required for normalization.

4. Host-side NVSHMEM lifecycle

allreduce-host.cu contains the host/runtime integration required by the tutorial, including:

• NVSHMEM initialization,

• distributed bootstrap,

• symmetric workspace management,

• host-side C ABI, and

• resource cleanup.

5. Python API and correctness test

allreduce.py provides:

• TLE-Raw CUDA function registration,

• the reusable operator/communicator wrapper,

• distributed initialization,

• a PyTorch reference implementation, and

• standalone correctness validation through torchrun.

6. Benchmark harness

benchmark.py provides a standalone distributed benchmark with:

• configurable token shapes,

• warmup iterations,

• repeated timing,

• per-rank synchronization,

• rank-max latency aggregation,

• percentile statistics, and

• JSON result output.

The benchmark implementation reuses the operator from allreduce.py rather than maintaining a second implementation.

Correctness

Validated on NVIDIA H20 GPUs with:

world size: 4
hidden size: 2048
dtype: BF16

Token counts:

Tokens | Max absolute error | Status -- | -- | -- 1 | 0.03125 | PASS 16 | 0.03125 | PASS 64 | 0.06250 | PASS 128 | 0.06250 | PASS 256 | 0.06250 | PASS 1024 | 0.06250 | PASS 4096 | 0.06250 | PASS

In these serving workloads, the current NVSHMEM route remains behind the mature FlashInfer production path, corresponding to roughly a 12–15% decode-throughput gap.

This framework result is included as an engineering boundary rather than as a performance-win claim. The tutorial contribution itself focuses on making the TLE-Raw NVSHMEM fusion path correct, reusable, and reproducibly benchmarkable.

Files changed

Only the new tutorial directory is added:

python/tutorials/tle/raw/nvshmem/05-allreduce-rmsnorm/

with four source files:

• allreduce-device.cu

• allreduce-host.cu

• allreduce.py

• benchmark.py

No framework-specific integration code, generated benchmark artifacts, or machine-specific paths are included in the PR.

Validation

Local validation performed:

• TLE-Raw import/runtime validation

• NVSHMEM distributed correctness

• TP4 / H2048 / BF16 correctness across T=1...4096

• standalone benchmark execution

• git diff --check

• copyright notice audit

• third-party source provenance audit

• no generated benchmark/profile artifacts included

• no machine-specific absolute paths included

Repository CI is expected to provide the final formatting/build validation.

License / copyright

All newly added source files contain the FlagOS copyright notice.

The final tutorial code was written against the current FlagTree TLE-Raw APIs. No GPL/LGPL/AGPL-derived implementation code or third-party source text is included in the added files.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.

---

root seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}