[TRTLLM-12199][feat] WideEP FT: add EPGroupHealth thread-safe rank mask (1a.1) by chienchunhung · Pull Request #13302 · NVIDIA/TensorRT-LLM

chienchunhung · 2026-04-22T00:14:32Z

Description

Adds EPGroupHealth, a process-local, thread-safe rank-health primitive for WideEP fault tolerance (PR 1a.1).

This PR has no runtime behavior change. It only adds the shared in-process state object that follow-up PRs will consume from the AlltoAll watchdog, NVLink rank-mask binding, EPLB reconfigure path, and PyExecutor health reporting.

EPGroupHealth provides:

Idempotent mark_failed(rank) / mark_active(rank) mutators.
is_active, get_active_count, get_failed_ranks, and all_active queries.
Active-rank mask export as both a Python int and little-endian uint64 words for the kernel ABI.
A monotonic generation counter for cheap iteration-boundary change detection.
Atomic immutable snapshot() reads for consumers that need coherent multi-field state.

Background

Phase 1 of WideEP FT needs one shared in-process answer to "which EP ranks are alive?" before later PRs can safely wire detection, kernel masking, EPLB remap, and health reporting. This PR intentionally limits scope to that primitive; cross-rank consensus, watchdog detection, kernel changes, and EPLB integration remain in follow-up PRs.

Test Coverage

Adds pure-Python unit coverage for:

construction and validation
idempotent fail/reactivate semantics
generation-counter behavior
mask-word layout, including NVL72 and multi-word cases
immutable failed-rank and snapshot views
concurrent update races

Run:
pytest tests/unittest/_torch/modules/test_ep_group_health.py -v

PR Checklist

Please review the following before submitting your PR:

PR description clearly explains what and why. If using CodeRabbit summary, please make sure it makes sense.
PR follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions).
Any new dependencies have been scanned for license and vulnerabilities.
CODEOWNERS updated if ownership changes.
Documentation updated as needed.
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

Summary by CodeRabbit

New Features
- Added a health monitoring system for tracking Expert Parallel rank status, enabling automatic detection and management of node failures in distributed training scenarios.
Tests
- Added comprehensive unit test coverage validating rank health tracking, state transitions, concurrent safety, and failure recovery behavior.

chienchunhung · 2026-04-24T02:22:14Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-24T02:28:05Z

PR_Github #45300 [ run ] triggered by Bot. Commit: 86afc95 Link to invocation

tensorrt-cicd · 2026-04-24T11:37:42Z

PR_Github #45300 [ run ] completed with state SUCCESS. Commit: 86afc95
/LLM/main/L0_MergeRequest_PR pipeline #35554 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

Link to invocation

Cross-link NVIDIA#13302 (the EPGroupHealth implementation) from the 1a.1 row of the implementation plan so future readers can trace each plan item to its filed PR. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Made-with: Cursor

chienchunhung · 2026-04-24T19:37:30Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-04-24T19:45:58Z

PR_Github #45436 [ run ] triggered by Bot. Commit: 6823d55 Link to invocation

…move v1 files Final batch of the v2 rewrite. §0 executive summary — problem statement, approach, key MPI-vs-Ray decision, two failure modes named, the four TRT-LLM uniqueness properties, headline timeline numbers, and an explicit "what v2 changes vs v1" list for readers coming from the prior version. README — complete rewrite as navigation for v2. Section table with one-line summaries per section. In-flight PR table (NVIDIA#13302 and NVIDIA#13404 against this design). Consolidated terminology table including the rank/process/slot distinction that a reviewer flagged as missing. Scope & non-goals stated once, not repeated throughout. Removing v1 files: - 01-background.md → content folded into 01-user-journey-and-stack.md (§1.3) and 00-executive-summary.md - 02-current-state.md → folded into 01-user-journey-and-stack.md (§1.2) and 03-failure-modes-and-gaps.md - 03-competitive-landscape.md → folded into 02-stack-comparison-and-positioning.md - 04-two-phase-recovery.md → superseded by 04-architecture-overview.md (now three-phase) - 05-rank-masking.md → folded into 05-phase-1-immediate-survival.md §5.1 - 06-eplb-adaptation.md → folded into 05-phase-1-immediate-survival.md §5.2 - 07-failure-detection.md → folded into 05-phase-1-immediate-survival.md §5.3 and §5.4 - 08-mx-gms-integration.md → split between 05-phase-1 (PR NVIDIA#12718 integration in §5.3) and 06-phase-2 (MX-GMS in §6.3) - 09-implementation-plan.md → superseded by 08-implementation-plan.md - 10-risks.md → superseded by 09-risks-and-open-questions.md - COMBINED.md → superseded; single-file view can be regenerated from the split files if needed New v2 file set: README + §0-§9 (11 files) + 3 workflow artifacts (redesign-outline, redesign-research-pass, redesign-research-pass-report). Section count held at 10 numbered sections (§0-§9) but with cleaner phase boundaries — one section per phase. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

tensorrt-cicd · 2026-04-25T15:06:23Z

PR_Github #45436 [ run ] completed with state SUCCESS. Commit: 6823d55
/LLM/main/L0_MergeRequest_PR pipeline #35667 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

Link to invocation

…ype scaffold Adds PR NVIDIA#14198 (DO NOT SUBMIT — preview only) to the in-flight PRs table in the README and updates mvp-prototype-plan.md to reflect that the prototype scaffolding is now shipped on branch WideEP-FT/mvp-prototype: * README.md In-flight PRs table: new row for NVIDIA#14198 (status flagged as "Draft (DO NOT SUBMIT — preview only)" to distinguish from the MVP component PRs). * README.md prototype paragraph: links NVIDIA#14198 alongside the cherry-picked NVIDIA#13302 and NVIDIA#14160; notes the discardable prototypes/wide_ep_ft_mvp/ directory. * mvp-prototype-plan.md Status header: bumped to reflect scaffolding-shipped state. * mvp-prototype-plan.md §6 sequencing: new "Current status (2026-05-15)" paragraph documenting that both 1a.1 (NVIDIA#13302) and 1d.0 (NVIDIA#14160) are private cherry-picks on the prototype branch pending their merges to main, and that the kernel-side 1a.2 + 1a.3 work is deferred per prototypes/wide_ep_ft_mvp/kernel/README.md. Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

chienchunhung · 2026-06-13T02:31:52Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-13T02:37:34Z

PR_Github #53991 [ run ] triggered by Bot. Commit: 09bdf4c Link to invocation

coderabbitai · 2026-06-13T02:37:39Z

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6643826d-5f13-4389-9233-2413ad497e63

📥 Commits

Reviewing files that changed from the base of the PR and between c2b7cd9 and 09bdf4c.

📒 Files selected for processing (2)

tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py
tests/unittest/_torch/modules/test_ep_group_health.py

📝 Walkthrough

Walkthrough

This PR introduces EPGroupHealth, a thread-safe process-local tracker for Expert Parallel rank availability, maintaining an active bitmask with idempotent failure transitions, generation bumping on state changes, and NVLink AlltoAll-compatible uint64 word-based mask export, validated by comprehensive unit and concurrency tests.

Changes

EP Group Health Implementation and Tests

Layer / File(s)	Summary
Data types and class initialization `tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py` (lines 1–102)	Module constant `EP_MASK_NUM_WORDS`, immutable `EPGroupHealthSnapshot` NamedTuple, and `EPGroupHealth` initialization with internal bitmask, lock, active count, failed-ranks set, and generation counter.
State transitions and core methods `tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py` (lines 103–179)	Lock-protected `ep_size` and `generation` properties; rank validation; idempotent `mark_failed` and `mark_active` that update bitmask and failed-ranks set with generation bumps only on effective state changes.
Query and accessor methods `tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py` (lines 180–273)	`is_active`, `get_mask`, `get_mask_words` (little-endian `uint64` chunking with validation), `get_active_count`, `get_failed_ranks`, `all_active`, `snapshot`, `__len__` (EP group size), and `__repr__` for coherent reads and diagnostics.
Validation and state transition unit tests `tests/unittest/_torch/modules/test_ep_group_health.py` (lines 1–273)	Tests constructor invariants, rank-range validation, `mark_failed`/`mark_active` idempotency and reversibility, failure accumulation, snapshot immutability of `get_failed_ranks()`, `get_mask_words()` layout correctness across word boundaries and NVLink cases, generation counter semantics, and `__repr__` content.
Concurrency, oscillation, and immutability tests `tests/unittest/_torch/modules/test_ep_group_health.py` (lines 280–460)	Multi-threaded stress tests (barriers/events for distinct and contending ranks, fail/reactivate scenarios), oscillation with monotonic generation, `snapshot()` atomic multi-field consistency and immutability, and kernel-ABI constant assertion for `EP_MASK_NUM_WORDS == 2`.

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 47.62% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The PR title clearly identifies the main change: adding EPGroupHealth, a thread-safe rank mask primitive for WideEP fault tolerance. It is concise and directly reflects the core contribution of the changeset.
Description check	✅ Passed	The PR description is comprehensive and follows the template structure. It explains what EPGroupHealth does, why it is needed (background context), lists key features and API, clarifies scope boundaries, and provides test coverage information with a test command. The checklist is completed appropriately.
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.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

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

tensorrt-cicd · 2026-06-13T09:16:39Z

PR_Github #53991 [ run ] completed with state SUCCESS. Commit: 09bdf4c
/LLM/main/L0_MergeRequest_PR pipeline #43077 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 · 2026-06-14T17:22:09Z

/bot run --disable-fail-fast --stage-list "DGX_B200-8_GPUs-PyTorch-1,GB300-PyTorch-1"

tensorrt-cicd · 2026-06-14T17:27:56Z

PR_Github #54147 [ run ] triggered by Bot. Commit: 09bdf4c Link to invocation

tensorrt-cicd · 2026-06-14T19:15:49Z

PR_Github #54147 [ run ] completed with state SUCCESS. Commit: 09bdf4c
/LLM/main/L0_MergeRequest_PR pipeline #43230 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

chienchunhung · 2026-06-15T18:59:59Z

/bot run

tensorrt-cicd · 2026-06-15T19:00:09Z

PR_Github #54347 [ run ] completed with state ABORTED. Commit: 4ecc284

Link to invocation

EPGroupHealth is the process-local, thread-safe primitive that tracks which ranks in an EP group are alive vs. failed. It is the foundational data structure for WideEP fault tolerance (PR 1a.1 of the implementation plan), shared across the AlltoAll communication backends, the host-side watchdog, the MoE load balancer, and PyExecutor health reporting. Internally backed by an arbitrary-precision Python int bitmask; the uint64[N] format expected by the NVLink AlltoAll kernels is exposed via get_mask_words() (default 2 words covers the NVL72 72-rank case). mark_failed / mark_active are idempotent and bump a monotonic generation counter only on effective state change, so consumers can compare a cached generation to detect mask changes without diffing the full mask. Tests cover the single-threaded API surface plus three concurrent race scenarios (distinct-rank fan-out across 128 threads, same-rank idempotency under contention, mixed fail/reactivate pairs). Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

tensorrt-cicd · 2026-06-15T19:05:32Z

PR_Github #54350 [ run ] triggered by Bot. Commit: 9b76997 Link to invocation

tensorrt-cicd · 2026-06-15T22:36:56Z

PR_Github #54350 [ run ] completed with state FAILURE. Commit: 9b76997
/LLM/main/L0_MergeRequest_PR pipeline #43421 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 · 2026-06-15T23:25:37Z

/bot run --disable-fail-fast --stage-list "H100_PCIe-PyTorch-Ray-1"

tensorrt-cicd · 2026-06-15T23:34:06Z

PR_Github #54373 [ run ] triggered by Bot. Commit: 9b76997 Link to invocation

tensorrt-cicd · 2026-06-16T00:08:21Z

PR_Github #54373 [ run ] completed with state SUCCESS. Commit: 9b76997
/LLM/main/L0_MergeRequest_PR pipeline #43443 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

chienchunhung · 2026-06-16T00:19:10Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-16T00:25:01Z

PR_Github #54385 [ run ] triggered by Bot. Commit: 9b76997 Link to invocation

tensorrt-cicd · 2026-06-16T07:12:17Z

PR_Github #54385 [ run ] completed with state FAILURE. Commit: 9b76997
/LLM/main/L0_MergeRequest_PR pipeline #43453 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

Link to invocation

chienchunhung · 2026-06-16T16:47:41Z

/bot run --disable-fail-fast --stage-list "DGX_B200-4_GPUs-PyTorch-Ray-1"

tensorrt-cicd · 2026-06-16T16:54:52Z

PR_Github #54637 [ run ] triggered by Bot. Commit: 9b76997 Link to invocation

tensorrt-cicd · 2026-06-16T17:45:53Z

PR_Github #54637 [ run ] completed with state SUCCESS. Commit: 9b76997
/LLM/main/L0_MergeRequest_PR pipeline #43669 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

chienchunhung · 2026-06-16T20:36:28Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-16T20:42:06Z

PR_Github #54666 [ run ] triggered by Bot. Commit: 9b76997 Link to invocation

tensorrt-cicd · 2026-06-16T21:57:23Z

PR_Github #54666 [ run ] completed with state FAILURE. Commit: 9b76997
/LLM/main/L0_MergeRequest_PR pipeline #43698 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 · 2026-06-16T22:36:23Z

/bot run --disable-fail-fast --stage-list "GB200-4_GPUs-PyTorch-1,GB200-4_GPUs-PyTorch-2,GB200-4_GPUs-PyTorch-3,GB200-4_GPUs-PyTorch-4,GB200-4_GPUs-PyTorch-5"

tensorrt-cicd · 2026-06-16T22:43:28Z

PR_Github #54682 [ run ] triggered by Bot. Commit: 9b76997 Link to invocation

tensorrt-cicd · 2026-06-17T00:20:04Z

PR_Github #54682 [ run ] completed with state FAILURE. Commit: 9b76997
/LLM/main/L0_MergeRequest_PR pipeline #43713 (Partly Tested) 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 · 2026-06-17T04:49:06Z

/bot run --disable-fail-fast

chienchunhung · 2026-06-17T05:27:33Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-17T05:33:01Z

PR_Github #54768 [ run ] triggered by Bot. Commit: 9b76997 Link to invocation

tensorrt-cicd · 2026-06-17T07:49:34Z

PR_Github #54768 [ run ] completed with state SUCCESS. Commit: 9b76997
/LLM/main/L0_MergeRequest_PR pipeline #43787 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

…sk (1a.1) (NVIDIA#13302) Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Signed-off-by: GitLab CI Bot <gitlab-ci@nvidia.com>

github-actions Bot assigned chienchunhung Apr 22, 2026

chienchunhung changed the title ~~[None][feat] WideEP FT: add EPGroupHealth thread-safe rank mask~~ [TRTLLM-12199][feat] WideEP FT: add EPGroupHealth thread-safe rank mask Apr 24, 2026

chienchunhung force-pushed the WideEP-FT/1a.1-ep-group-health branch 2 times, most recently from 1e04363 to 86afc95 Compare April 24, 2026 02:20

chienchunhung force-pushed the WideEP-FT/1a.1-ep-group-health branch from 86afc95 to 6823d55 Compare April 24, 2026 19:33

chienchunhung changed the title ~~[TRTLLM-12199][feat] WideEP FT: add EPGroupHealth thread-safe rank mask~~ [TRTLLM-12199][feat] WideEP FT: add EPGroupHealth thread-safe rank mask (1a.1) Apr 24, 2026

chienchunhung force-pushed the WideEP-FT/1a.1-ep-group-health branch from 6823d55 to 0323d2f Compare May 13, 2026 23:12

chienchunhung mentioned this pull request May 19, 2026

(DO NOT SUBMIT) WideEP FT MVP prorotype #14198

Draft

1 task

chienchunhung force-pushed the WideEP-FT/1a.1-ep-group-health branch from 0323d2f to 09bdf4c Compare June 13, 2026 00:07

chienchunhung requested review from dongxuy04 and xxi-nv June 13, 2026 02:31

chienchunhung marked this pull request as ready for review June 13, 2026 02:31

chienchunhung requested a review from a team as a code owner June 13, 2026 02:31

xxi-nv approved these changes Jun 15, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py Outdated

chienchunhung force-pushed the WideEP-FT/1a.1-ep-group-health branch from 3255353 to 14e5491 Compare June 15, 2026 18:58

chienchunhung force-pushed the WideEP-FT/1a.1-ep-group-health branch from 14e5491 to 9b76997 Compare June 15, 2026 19:04

chienchunhung requested a review from pcastonguay June 16, 2026 20:37

chienchunhung merged commit 9e69568 into NVIDIA:main Jun 18, 2026
7 checks passed

Uh oh!

Conversation

chienchunhung commented Apr 22, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Background

Test Coverage

PR Checklist

GitHub Bot Help

Summary by CodeRabbit

Uh oh!

chienchunhung commented Apr 24, 2026

Uh oh!

tensorrt-cicd commented Apr 24, 2026

Uh oh!

tensorrt-cicd commented Apr 24, 2026

Uh oh!

chienchunhung commented Apr 24, 2026

Uh oh!

tensorrt-cicd commented Apr 24, 2026

Uh oh!

tensorrt-cicd commented Apr 25, 2026

Uh oh!

chienchunhung commented Jun 13, 2026

Uh oh!

tensorrt-cicd commented Jun 13, 2026

Uh oh!

coderabbitai Bot commented Jun 13, 2026

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

tensorrt-cicd commented Jun 13, 2026

Uh oh!

chienchunhung commented Jun 14, 2026

Uh oh!

tensorrt-cicd commented Jun 14, 2026

Uh oh!

tensorrt-cicd commented Jun 14, 2026

Uh oh!

Uh oh!

chienchunhung commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

chienchunhung commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

chienchunhung commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

chienchunhung commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

chienchunhung commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

chienchunhung commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 17, 2026

Uh oh!

chienchunhung commented Jun 17, 2026

chienchunhung commented Apr 22, 2026 •

edited by coderabbitai Bot

Loading