[Feature] NcclEP backend hardcodes LOW_LATENCY + RANK_MAJOR; algorithm/layout should be selectable (HIGH_THROUGHPUT + FLAT measured working on AWS EFA)

Summary

The NcclEP MoE communication backend constructs its nccl.ep group with a hardcoded Algorithm.LOW_LATENCY and defaults the layout to RANK_MAJOR:

• tensorrt_llm/_torch/modules/fused_moe/nccl_ep_utils.py:200 → algorithm=Algorithm.LOW_LATENCY
• tensorrt_llm/_torch/modules/fused_moe/nccl_ep_utils.py:135 → self.layout = Layout.RANK_MAJOR if layout is None else ...

(verified against main @ 7a3b1bf5, 2026-08-14)

On AWS EFA the correct algorithm/layout is fabric- and path-dependent, which is exactly why a hardcode is the wrong default. We have measured HIGH_THROUGHPUT + FLAT running correctly on EFA both at 16 ranks and inside a real trtllm-serve; and on a device-initiated (IBGDA-class) LOW_LATENCY dispatch path on EFA we have separately observed a CUDA illegal memory access at the first dispatch. Because which tuple is safe depends on the substrate, it should be selectable rather than fixed.

Substrate caveat (measured both ways — important, we are not claiming LL is broken). The LOW_LATENCY CUDA IMA was observed on a standalone nccl_ep dispatch harness on a device-initiated path. On a different EFA substrate — NCCL runtime 2.30.4 + nccl_ep 0.1.0 + NCCL-GIN CPU-proxy path, driven inside a real trtllm-serve — LOW_LATENCY + RANK_MAJOR ran clean (serve up, HTTP 200, correct output, zero illegal-memory-access, 16/16 cross-node). So LOW_LATENCY is not universally broken on EFA; the failure is path-dependent (device-initiated vs CPU-proxy). The ask is selectability, not "change the default."

System Info

• nccl.ep from nccl4py (libnccl_ep 0.1.0), NCCL runtime 2.30.4 (the minimum nccl_ep_utils.is_nccl_ep_installed() requires)
• Fabric: AWS EFA on p5en (H200), libfabric efa provider, NCCL GIN CPU-proxy path, efa-direct confirmed on every rank. Non-IBGDA.
• 2 nodes × 8 GPU = 16 ranks. Shape: 128 experts, hidden 2048, top-8, 128 tokens/rank.

Actual behavior

The working HIGH_THROUGHPUT+FLAT tuple additionally needs (all already expressible via the existing config surface): num_qp_per_rank >= 32, recv_topk_idx dtype int64, and 2D [max_recv, hidden] recv buffers.

One extra fix the HT/FLAT path requires in a real serve (found by running it end-to-end). The nccl_ep 0.1 HT kernel ABI returns recv_topk_idx int64, but the downstream consumer torch.ops.trtllm.fused_moe requires int32 routing ids — a conflict invisible under the LOW_LATENCY/RANK_MAJOR default (buffer is already int32) and only reachable on the HT/FLAT path. Left unhandled it raises token_selected_experts dtype is Long, while Int is expected (torch_custom_ops.py) and the serve dies at startup on every rank. The fix is a one-line narrowing at the dispatch return boundary — recv_slots_global.to(torch.int32) — safe because expert ids (< num_experts) and the -1 unrouted sentinel both fit int32, and under LL it is a no-op. So the complete HT/FLAT enablement is algorithm gate + HT-aware FLAT default + int32 return narrowing, not the algorithm gate alone.

Expected behavior

A user on an EFA path where the hardcoded LOW_LATENCY tuple takes the CUDA IMA should be able to select the HIGH_THROUGHPUT + FLAT tuple that works on the same fabric — the hardcode is the only blocker, and the alternative tuple is already fully supported by the same nccl.ep build. Even where LOW_LATENCY runs (the CPU-proxy path above), selectability lets an operator pick the algorithm suiting their fabric rather than depending on which path the build happens to take.

(Notably, is_nccl_ep_installed() requires NCCL ≥ 2.30.4, while the shipping pip pins and the NGC containers we checked carry 2.28.9 — so this backend does not activate in those environments, and the EFA path appears not to have been widely exercised yet.)

Proposed fix

Make the algorithm/layout selectable rather than hardcoded — an env override TRTLLM_NCCL_EP_ALGO (LOW_LATENCY | HIGH_THROUGHPUT) plus an HT-aware layout default, and the int32 return-narrowing the HT/FLAT path requires. LOW_LATENCY stays the default; nothing changes for existing users.

We have this staged as a small patch (20 insertions / 2 deletions in nccl_ep_utils.py for the gate, plus the one-line narrowing) and will open a draft PR referencing this issue.

Scope note (honest): correctness-parity passes (16/16 ranks × 3 rounds, plus a real trtllm-serve producing coherent + arithmetically-correct output). Native nccl_ep at this shape is ~3.7× slower than our CPU-proxy reference, so this is filed as a correctness/enablement ask (make the tuple selectable), explicitly not a performance win.

Reproduction

Minimal torchrun 2×8 driving nccl.ep Group.create(GroupConfig(algorithm=..., ...)) → dispatch → combine with an identity-echo correctness check reproduces the algorithm/layout behavior; a full trtllm-serve (Qwen3-30B-A3B, EP8) additionally exercises the torch.ops.trtllm.fused_moe int32-consumer path the standalone harness does not. On the device-initiated path LOW_LATENCY IMAs; on the CPU-proxy path it runs clean; HIGH_THROUGHPUT + FLAT (with the int32 narrowing) passes on both the harness and a real serve. Harness available on request (all public NGC-rooted).

Start in tensorrt_llm/_torch/modules/fused_moe/nccl_ep_utils.py at the layout default and Group.create configuration around lines 135 and 200. Review the minimal torchrun nccl.ep dispatch/combine reproduction, then the trtllm-serve path that exercises fused_moe. Done means algorithm and layout can be selected without changing the LOW_LATENCY default, and the HT/FLAT path completes with the expected routing-id dtype.