[https://nvbugs/6430702][perf] Restore NVLink one-sided A2A fast path

[chienchunhung]

Summary

• Compile active-rank mask checks out of non-FT NVLink one-sided AllToAll kernels.
• Restore the pre-regression combine payload loop for both masked and unmasked specializations.
• Keep an FT-only dispatch safety guard until post-commit routing is enforced end to end, plus mask checks around peer counters, EPLB statistics, and completion synchronization.
• Pass the construction-time rank-mask mode explicitly through the real and fake operator schemas while treating the mask as per-epoch data.
• Reject CUDA graphs in rank-mask mode until generation-scoped invalidation and recapture are available.
• Align multi-GPU rank-mask coverage with route-safety, mode-consistency, and failed-epoch semantics.

Root cause

The reported regression is between post-merge jobs 2825 and 2826:

• Good commit: 7c8dde830bac813e23605d47a1d27c92d5437a92 (12.88 ms)
• Bad commit: 045705139d125dbfd0614b369096f7bb5ddcbebf (14.43 ms)
• Metric: Mean Gen Worker Per-Iter Device Step Time, lower is better

The relevant kernel diff added two is_rank_active checks inside vectorized_combine_impl: one in the payload-load pass and another in the conversion pass. These execute in the token/top-k/vectorized payload hot path even when every rank is active.

Design alignment

The active-rank mask is a next-launch primitive. The full post-commit routing guarantee is not yet enforced end to end, so the masked dispatch specialization defensively rejects inactive route targets before remote workspace access. That sentinel is an internal fail-closed artifact, not valid model output. A failed or generation-mismatched epoch must be discarded.

The compile-time flag is named ENABLE_RANK_MASK rather than ENABLE_FAULT_TOLERANCE so it remains independent from the running-kernel execution-abort mechanism. Its value is fixed when the communicator is constructed; the mask tensor supplies one epoch's committed membership. The masked specialization skips inactive peers in route, completion, counter, and EPLB-stat paths, but not in combine payload accumulation.

The mask remains captured by value in kernel arguments. Rank-mask mode therefore rejects configured CUDA graphs and direct stream capture until the 1a.11 generation-scoped graph lifecycle lands.

Related: #15524, #16025

NVBUG: https://nvbugs/6430702

Validation

• clang-format --dry-run --Werror on the changed C++/CUDA files
• ruff check and ruff format --check on the changed non-legacy Python files
• scripts/legacy_utils.py lint-precommit tensorrt_llm/_torch/distributed/moe_alltoall.py
• scripts/legacy_utils.py lint-precommit tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
• yapf -d tensorrt_llm/_torch/distributed/moe_alltoall.py
• yapf -d tensorrt_llm/_torch/custom_ops/cpp_custom_ops.py
• git diff --check
• Source comparison confirms the combine payload loop matches good commit 7c8dde83

Not run locally: CUDA build, MPI multi-GPU tests, or performance benchmarks because this worktree has no TRT-LLM build environment, nvcc, or pytest. These require CI/GPU validation.

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58505 [ run ] triggered by Bot. Commit: d20af81 Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds compile-time active-rank masking to MoE All-to-All CUDA dispatch/combine kernels via a new ENABLE_RANK_MASK template parameter and enable_rank_mask struct field, updates op-layer mask resolution to require an explicit mask, tightens Python communicator validation of active_rank_mask, and updates tests for the new dispatch routing metadata contract.

Changes

Active-Rank Masking Gating

Layer / File(s)	Summary
Kernel header contract cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h	Adds enable_rank_mask{false} to MoeA2ADispatchParams/MoeA2ACombineParams and rewords active_rank_mask documentation for both kernel pointer structs.
Dispatch kernel routing/sync cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu	Adds compile-time is_rank_active specialization, extends moeA2ADispatchKernel with ENABLE_RANK_MASK, removes dead-rank filtering from routing, and templates completion/wait sync loops.
Dispatch launch ...moeAlltoAllKernels.cu	moe_a2a_dispatch_launch conditionally validates ep_rank activity and switches kernel instantiation via SWITCH_BOOL(enable_rank_mask, ...).
Combine kernel accumulation/sync ...moeAlltoAllKernels.cu	vectorized_combine_impl now skips only on negative dst_idx/topk_send_indices; moeA2ACombineKernel gains ENABLE_RANK_MASK and templated sync loops.
Combine launch ...moeAlltoAllKernels.cu	moe_a2a_combine_launch conditionally validates ep_rank activity and switches kernel instantiation via SWITCH_BOOL(enable_rank_mask, ...).
Op-layer mask resolution cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp	Adds hasActiveRankMask; resolveActiveRankMask now requires a defined mask; moeA2ADispatchOp/moeA2ACombineOp set enable_rank_mask and conditionally resolve the mask.
Python MoeAlltoAll gating tensorrt_llm/_torch/distributed/moe_alltoall.py	Adds _rank_mask_enabled; dispatch()/combine() raise ValueError on active_rank_mask without rank-mask mode, otherwise capture/derive the mask via coordinator APIs.
NVLinkOneSided gating tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py	Adds _rank_mask_enabled; dispatch()/combine() validate active_rank_mask overrides and pass requested masks to capture/derive helpers.
Test updates for new routing metadata tests/unittest/_torch/modules/moe/test_moe_comm.py	Test helpers now return/use topk_send_indices; removes dead_ranks param; one-rank-masked test excludes dead rank from routing and asserts exact expected output; removes obsolete inactive-before-combine test.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant PythonCaller as MoeAlltoAll/NVLinkOneSided
  participant Op as moeAlltoAllOp.cpp
  participant Launch as moe_a2a_dispatch/combine_launch
  participant Kernel as moeA2ADispatch/CombineKernel

  PythonCaller->>Op: dispatch/combine(active_rank_mask?)
  alt rank-mask mode disabled and mask provided
    PythonCaller->>PythonCaller: raise ValueError
  else mask provided or absent
    Op->>Op: hasActiveRankMask(activeRankMask)
    Op->>Op: set params.enable_rank_mask
    opt mask present
      Op->>Op: resolveActiveRankMask(validate mask)
    end
    Op->>Launch: launch(params)
    alt enable_rank_mask true
      Launch->>Launch: validate local ep_rank active
      Launch->>Kernel: instantiate with ENABLE_RANK_MASK=true
    else
      Launch->>Kernel: instantiate with ENABLE_RANK_MASK=false
    end
    Kernel->>Kernel: is_rank_active<ENABLE_RANK_MASK> in sync loops
    Kernel-->>PythonCaller: routed/combined output
  end

Loading

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 83.33% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title is concise, specific, and matches the main change: restoring the NVLink one-sided A2A fast path.
Description check	✅ Passed	The description covers summary, root cause, design alignment, related issues, and validation, matching the template well enough.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[tensorrt-cicd]

PR_Github #58505 [ run ] completed with state FAILURE. Commit: d20af81
/LLM/main/L0_MergeRequest_PR pipeline #47112 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58689 [ run ] triggered by Bot. Commit: c8a0ab7 Link to invocation

[tensorrt-cicd]

PR_Github #58689 [ run ] completed with state FAILURE. Commit: c8a0ab7
/LLM/main/L0_MergeRequest_PR pipeline #47275 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58880 [ run ] triggered by Bot. Commit: 79c9cbc Link to invocation

[tensorrt-cicd]

PR_Github #58880 [ run ] completed with state SUCCESS. Commit: 79c9cbc
/LLM/main/L0_MergeRequest_PR pipeline #47423 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast --stage-list "DGX_B200-4_GPUs-PyTorch-2,DGX_B200-8_GPUs-PyTorch-3"

[tensorrt-cicd]

PR_Github #59008 [ run ] triggered by Bot. Commit: 79c9cbc Link to invocation

[tensorrt-cicd]

PR_Github #59008 [ run ] completed with state SUCCESS. Commit: 79c9cbc
/LLM/main/L0_MergeRequest_PR pipeline #47537 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59031 [ run ] triggered by Bot. Commit: 79c9cbc Link to invocation

[tensorrt-cicd]

PR_Github #59031 [ run ] completed with state SUCCESS. Commit: 79c9cbc
/LLM/main/L0_MergeRequest_PR pipeline #47557 completed with status: 'SUCCESS'

CI Report

Link to invocation