[TLERaw][Mega] Add Triton TLE NVSHMEM MegaMoE correctness case

[Zhang-kg]

Summary

This PR adds a Triton TLE + raw NVSHMEM MegaMoE integration case under:

python/test/tle/integration/megamoe/

The goal of this PR is to demonstrate that the current FlagTree/Triton TLE stack can run a small end-to-end MegaMoE-style operator with:

• raw NVSHMEM dispatch / receiver
• TLE warp specialization
• FP8 L1/L2 TensorCore tl.dot compute
• remote combine staging
• local y reduce
• repeated launch with workspace cleanup

This is a correctness/capability integration case, not yet a production-performance replacement for UserHopperMegaMoE.

Environment

Validated locally on:

GPU: 8x NVIDIA H100 80GB HBM3
CUDA: 12.8
Python: 3.10
NVSHMEM: 3.4.5 (nvidia-nvshmem-cu12==3.4.5)
MPI launcher: Open MPI 4.1.2 (mpirun)
Triton: FlagTree PR682-based Triton with TLE raw NVSHMEM support

The runnable instructions are documented in:

python/test/tle/integration/megamoe/RUNBOOK_CN.md

Implementation Status

Added files include:

python/test/tle/integration/megamoe/
  megamoe_operator/
    triton_tle_megamoe_operator.py
    triton_tle_megamoe_runtime.py
    ws_userhopper_dispatch_receiver_device.cu
    ws_userhopper_dispatch_receiver_extern_call.py
    ws_userhopper_dispatch_receiver_host.cu

tests/
  megamoe_local_harness.py
  run_isolated_operator.py

perf/
test_logs/
RUNBOOK_CN.md

Current main kernel path: _single_kernel_dispatch_receiver_l1_l2_tile_split_multi_cta_tldot_kernel

The implementation currently validates a merged dispatch/receiver/compute/combine path for small controlled shapes. The 8-rank H256 path uses a tile-split multi-CTA compute structure.

A file lock and atomic temporary output path were added around local host helper compilation to avoid concurrent mpirun ranks corrupting the generated .so.

Verified Cases

Case 1: 2-rank H128

Config:

world_size = 2
H = 128
I = 128
experts = 2
topk = 1
tokens/rank = 2
repeats = 3
cleanup = 1

Result:

PASS printed on both ranks
finalize still hangs; command is expected to be collected by timeout

Latest local timing from test_logs/local_2rank_h128_tokens2_repeats3.log:

rank 0 steady_avg_us = 4920.624, steady_max_us = 6031.936
rank 1 steady_avg_us = 3580.512, steady_max_us = 3926.080

Case 2: 8-rank H256 masked tile-split

Config:

world_size = 8
H = 256
I = 128
experts = 16
topk = 8
route_mode = masked
tokens/rank = 1
repeats = 2
cleanup = 1
compute_order = expert_wave_multi_cta_l1_l2_tile_split

Result:

PASS on all 8 ranks
process exits normally
checked = 6
counts = [3, 3]

Latest local timing from test_logs/local_8rank_h256_topk8_masked_tokens1_tile_split_repeats2.log:

steady_min_us across ranks ~= 16207.552
steady_max_us across ranks ~= 16471.872

Current Gap vs UserHopperMegaMoE

This PR does not claim production equivalence with UserHopperMegaMoE yet.

Known gaps:

• Only small correctness/capability shapes are covered.
• Production MoE shapes are not supported yet.
• The current H256 path relies on tile-split multi-CTA structure.
• It does not yet implement UserHopper’s full persistent scheduler.
• It does not yet implement the mature UserHopper TMA pipeline / mbarrier structure.
• It does not yet implement production-grade expert-wave scheduling and combine path.
• The 2-rank H128 case still has a finalize hang after PASS.
• Performance numbers are only for sanity/capability tracking, not for speedup claims.

Next Steps

Planned follow-up work:

1. Remove the 2-rank finalize hang.

2. Convert the current manual runner into a more standard gated integration test.

3. Expand supported shapes beyond the current small H128/H256 cases.

4. Move closer to UserHopper’s production structure:

   • persistent scheduling
   • expert-wave execution
   • TMA-style pipeline
   • finer synchronization and arrival tracking
   • production combine/reduce path

5. Add comparable benchmark cases against existing MegaMoE implementations once shapes and semantics are aligned.

[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.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

[Zhang-kg]

Summary

This PR adds production-shape Triton TLE + raw NVSHMEM MegaMoE candidates under:

python/test/tle/integration/megamoe/

Two validated candidates are included:

• v25 single-rank: BM64, one math warp-group, descriptorless raw TMA1D D8 token pull, and raw wide L2 scatter.
• v33 multi-rank: BM128, two math warp-groups, SMEM expert counting, parallel NVLink signaling, one D8 TMA1D pull stream, and raw wide L2 scatter.

Both candidates execute dispatch, L1/L2 GEMM, combine, and writeback in one persistent Triton/TLE kernel launch per rank.

The raw CUDA boundary is intentionally narrow:

• descriptorless TMA1D for the FP8 token payload;
• L2 wide scatter for gathered remote destinations.

Routing, TLE pipes, WGMMA, scale handling, SwiGLU, metadata, arrival signaling, and combine remain in Triton/TLE.

Compiler baseline

The branch has been rebased onto the current flagos-ai/main.

The kernels require safe inlining of helpers containing tle.local_pointers inside warp-specialized workers. The original implementation in PR #837 was reverted by PR #845; the narrower TLE-specific implementation was subsequently merged through PR #859 and is present in the current main branch.

No reverted PR #837 compiler implementation is carried by this PR.

Workload

GPU: 8x NVIDIA H100 80GB HBM3
CUDA: 12.8
NVSHMEM: 3.4.5

H=4096
I=1536
experts=128
topk=8
tokens/rank=512
drop=0
stages=4

Correctness

v25 single-rank

1/1 PASS
routed partial rows checked: 4096
scatter_bad=0
errors=0

v33 8-rank

8/8 PASS
routed partial rows checked per rank: 4096
scatter_bad=0
errors=0

The current correctness harness checks the integrated dispatch, L1/L2 compute, cross-rank scatter, and final output path. This is integration evidence, not a claim of complete UserHopper stage equivalence.

Performance

Timing protocol:

warmup=10
iterations=30
reduction=arithmetic mean

v25 single-rank

1272.5 us

This value was collected with the PR837-compatible validation build after forward-porting v25. Historical measurements from a different compiler build are not mixed into this result.

v33 8-rank

rank 0: 464.8 us
rank 1: 468.6 us
rank 2: 475.9 us
rank 3: 470.5 us
rank 4: 475.5 us
rank 5: 476.9 us
rank 6: 476.4 us
rank 7: 479.7 us

rank mean: 473.54 us
rank population stddev: 4.70 us

Reference CUDA event baseline:

335.8 us

Under this timing protocol, v33 reaches approximately:

335.8 / 473.54 = 70.9% of CUDA performance

Main files

megamoe_operator/production/
  production_runtime.py
  nvshmem_host.cu

  v25/
    run.py
    kernel.py
    raw_d8_tma1d_pull.cu
    raw_l2_wide_scatter.cu

  v33/
    run.py
    raw_d8_tma1d_pull.cu
    raw_l2_wide_scatter.cu

Detailed commands and environment variables are documented in:

python/test/tle/integration/megamoe/RUNBOOK_CN.md