feat(moe_ep): fault-tolerance rank mask (NCCL-EP + NIXL-EP)

[Anerudhan]

Why

NCCL-EP gained fault tolerance (nccl_ep.h enable_mask + ncclEpMask*): a peer that times out during dispatch/combine is masked and skipped instead of tripping a GPU trap(). NIXL-EP has the equivalent (update/query/clean_mask_buffer). flashinfer.moe_ep exposed neither — it never set enable_mask, never called a mask API, and gave callers no way to learn a rank had died. One slow or dead EP rank killed the whole job.

This adds the FT surface over both transports. moe_ep is pure Python, so there is no CUDA/C++ change here.

---

How the FT API is used

FT is opt-in per Fleet via a knob, not new FleetParams fields — FleetParams is the frozen sizing dataclass, while optional transport features already live in the knob namespace (cf. FleetAlgoKnobTopologyCapacity). The runtime API lands on the Fleet ABC as concrete raising defaults, not @abstractmethod, so no existing or out-of-tree Fleet breaks.

1. Probe, then enable

from flashinfer.moe_ep import (
    FleetAlgoKnobFaultTolerance, FleetAlgoKnobTopologyCapacity,
    MoEEpLayer, supports_fault_tolerance,
)

# Needs more than the backend being built: nccl_ep also needs an nccl4py with
# GroupConfig.enable_mask AND a libnccl_ep exporting ncclEpMask*. Never raises.
assert supports_fault_tolerance("nccl_ep")

knobs = [FleetAlgoKnobFaultTolerance(timeout_ms=5000)]   # 0 = transport default
# nixl only: size the capacity for the largest world you will ever reach
knobs.append(FleetAlgoKnobTopologyCapacity(n=32))

layer = MoEEpLayer(bootstrap, fleet_params, weights, fleet_knobs=knobs,
                   backend=SplitConfig(comm=NCCLEPConfig(), kernel=IdentityConfig()))

FleetAlgoKnobFaultTolerance(enabled=True, timeout_ms=0, reconcile_timeout_s=30.0, coordinator_takeover_s=10.0). LOW_LATENCY only on both transports — validate_fleet_params rejects FT + HIGH_THROUGHPUT at construction, because nccl leaves the mask buffer NULL under HT and the mask APIs then abort the process.

2. Serve, poll, recover

fleet = layer._ensure_fleet()

for step in ...:
    layer.forward(t)

    # Between iterations only — both transports read the mask LIVE from the
    # dispatch/combine kernels, so mutating it mid-collective is a race.
    if fleet.query_fault():                       # free on nccl (pinned host flag)
        agreed = fleet.reconcile_active_mask()    # store-collective, death-tolerant
        fleet.clear_faults(readmit=False)         # re-arm detection; keep serving DEGRADED

        # ... later, if the peer comes back (collective over survivors, blocking):
        # fleet.clear_faults(readmit=True)

Method	Collective?	Blocks host?	Stream-ordered?
supports_fault_tolerance	—	no	no
query_fault()	local	nccl no / nixl yes (small D2H)	nccl no / nixl yes
query_active_mask(out=None)	local	no	yes
set_active_mask(mask)	local (must be applied identically everywhere)	no	yes
reconcile_active_mask()	store-collective, tolerates dead ranks	yes (≤ timeout)	yes
clear_faults(readmit=False)	local	no	no
clear_faults(readmit=True)	collective over survivors	yes	yes
active_mask_epoch	—	no	no

Canonical mask: int32[world_size], 1 = active, CUDA tensor — matching ncclEpMaskQuery and vLLM's query_active_mask() naming.

3. Rules that will bite you

1. The steady state is read-only. The transport discovers the fault and masks the peer; the application's job is to notice (query_fault / query_active_mask). set_active_mask is the exceptional reconciliation path, not a per-step call — most callers only reach it via reconcile_active_mask(). If you do write the mask, write it only between iterations (kernels read it live).
2. All survivors must reconcile in the same iteration slot — they must pass the same active_mask_epoch.
3. clear_faults(readmit=True) and update_topology() are alternatives, not a sequence. The former re-admits a merely-delayed rank on the same communicator; the latter destroys the group and builds a new ncclComm_t, which is the only way to add or replace a process. Re-admitting after a rebuild is meaningless; before one it is wasted work. (If you do both, MaskClean must come first — it needs a live handle on the current group.)
4. No FT call during CUDA-graph capture — but note dispatch and combine themselves ARE safe to capture: neither transport compacts the surviving ranks' layout, so they re-read the mask on every replay and a rank that fails later is still skipped. What must not be captured is a decision about fault state. query_fault() is a host read, not stream work, so it can't be captured at all — it returns the capture-time answer and freezes the branch taken on it into the graph forever. All four FT entry points raise on capture with a per-operation reason. (Thanks @lrbison — an earlier draft of this said "stale offsets", which was wrong.)
5. A rank can be told it is dead. reconcile_active_mask() raises MoEEpRankEvictedError when the survivors agreed this rank is gone. It can't apply that (a rank may not mask itself) and mustn't keep serving (peers stopped sending it tokens).
6. Dropped tokens are not renormalized, and experts are not re-homed. y_degraded[t] == y_healthy[t] * Σ_{alive k} topk_weights[t][k]. Both omissions are deliberate: implicit renormalization would add a kernel to every forward, hide a serving-quality event, and divide by zero when a token's whole top-k died; re-homing the dead rank's experts is an EPLB-style job that belongs to the framework. Keep serving on a partial mask with reconcile_active_mask() → clear_faults(readmit=False) — no update_topology, no new communicator. Runbook has the opt-in renormalization snippet.

---

Call stacks

Group/Buffer creation (where masking is switched on)

nccl_ep                                        nixl_ep
-------                                        -------
NcclEpFleet.__init__                           NixlEpFleet.__init__
 ├ _index_knobs -> self._ft                     ├ _index_knobs -> self._ft
 ├ validate_fleet_params(fault_tolerance=)      ├ validate_fleet_params(fault_tolerance=)
 ├ _check_ft_supported()                        ├ nixl_ep.Buffer(..., timeout_ms=)   [TypeError -> actionable]
 │   ├ dataclasses.fields(GroupConfig)          ├ update_memory_buffers(cap, ...)
 │   │   -> needs "enable_mask"                 │   -> allocates mask_buffer[capacity], 0xFF-memset
 │   └ mask_ffi().available                     └ connect_ranks([0, world))
 │       -> needs the 5 ncclEpMask* symbols         -> unmasks [0, world); tail stays masked
 └ _build_group_config()
     ├ kwargs["enable_mask"] = True
     ├ kwargs["timeout_ns"] = timeout_ms * 1e6      (omitted when 0)
     └ nccl.ep.Group.create -> ncclEpCreateGroup
         -> allocates mask_buffer[nRanks] + pinned async-error flag

query_fault()

nccl_ep                                        nixl_ep
NcclEpFleet.query_fault                        NixlEpFleet.query_fault
 ├ reject_graph_capture("query_fault")          ├ self.query_active_mask()      (kernel + D2H)
 │   ^ NOT because it touches a stream --        │   ^ which itself rejects capture
 │     because it does NOT: a host read          └ compare against self._ft_applied
 │     cannot be captured, so it would
 │     return the capture-time answer and       (no host error flag on this transport,
 │     freeze the branch into the graph          so the applied mask IS the state)
 └ mask_ffi().get_async_error(group)
    ├ group.get_async_error()  [native first]
    └ ctypes ncclEpGetAsyncError(group.ptr, &o)
        -> libnccl_ep.so -> reads PINNED HOST flag: no stream, no sync, free

query_active_mask()

nccl_ep                                        nixl_ep
 ├ _reject_graph_capture()                      ├ _reject_graph_capture()
 ├ _ft_bufs() -> device int32[world]            ├ _ft_raw() -> device int32[CAPACITY]
 └ mask_ffi().mask_query(group, dev_ptr, str)   ├ buffer.query_mask_buffer(raw)
     -> ncclEpMaskQuery                         │   -> nixl_ep_cpp -> kernel copy
     -> D2D copy of mask_buffer                 │   (asserts numel == max_num_ranks)
     -> ALREADY 1 = active, identity            └ (raw[:world] == 0).to(int32)
                                                    ^ NOT 1-raw: buffer is 0xFF-memset so
                                                      untouched entries read back as -1

set_active_mask(mask)

nccl_ep                                        nixl_ep
 ├ _normalize_mask -> list[int], rejects        ├ _normalize_mask (same guard)
 │   masking the local rank                     └ for each CHANGED rank r != self:
 ├ host.copy_(...)   PINNED staging buffer          buffer.update_mask_buffer(r, mask=(a==MASKED))
 │   ^ pinned because Update is stream-ordered:      -> atomicExch kernel, one launch EACH
 │     a pageable buffer mutated next call is        -> so we push only the DIFF; a blind
 │     a use-after-write race                          range(world) loop would inject `world`
 └ mask_ffi().mask_update(group, host_ptr, str)        launches into the steady state
     -> ncclEpMaskUpdate  (HOST ptr, unlike Query)

reconcile_active_mask() — shared, transport-free

FaultToleranceMixin.reconcile_active_mask          (both backends)
 ├ local = self.query_active_mask().cpu().tolist()
 ├ store = self._ft_store()          -> resolve_rendezvous_store(subsystem="ft")
 └ reconcile_masks_via_store(store, rank, world, local, epoch=active_mask_epoch, ...)
     ├ 1. store.set("ft/gen{E}/local/{rank}", bytes(view))
     ├ 2. poll ONLY ranks we still believe alive, until timeout_s
     ├ 3. elementwise-AND what arrived; mask believed-alive ranks that never reported
     ├ 4. coord = min(active); coord publishes via ATOMIC compare_set("ft/gen{E}/decision")
     │      everyone else adopts whatever that key holds  <- kills split brain
     │      (coordinator itself dead -> mask it, re-elect, bounded by world_size)
     └ 5. _adopt(): raises MoEEpRankEvictedError if the decision masks US
 └ self.set_active_mask(agreed)   -> backend-specific path above

Deliberately not a torch.distributed allreduce: that would hang on exactly the rank being masked out. With a store, a missing key is the death signal.

clear_faults()

nccl_ep                                        nixl_ep
readmit=False:                                 readmit=False:
 └ mask_ffi().error_clear -> ncclEpErrorClear    └ return    (no sticky flag to re-arm)

readmit=True:                                  readmit=True:
 ├ guard: a handle must exist                   ├ buffer.clean_mask_buffer()
 │   (MaskClean asserts on the LL staging       │   -> zeroes ALL `capacity` entries,
 │    buffer -> would SIGABRT from C)           │      marking the never-connected tail ACTIVE
 ├ warn under EXPERT_MAJOR                      ├ _ft_applied = [ACTIVE] * world
 │   (MaskClean computes reset offsets          └ re-mask [world, capacity)   <- fixes that bug
 │    assuming RANK_MAJOR)
 ├ mask_clean -> ncclEpMaskClean
 │   COLLECTIVE over survivors; internally cudaStreamSynchronize's
 └ error_clear -> ncclEpErrorClear
     ^ always paired: MaskClean does NOT clear the flag, which is exactly why
       `readmit` is a flag on one method rather than two a caller can mis-sequence

---

Three transport facts that shaped the code

1. NIXL's polarity is "nonzero = masked", not "1 = masked". The buffer is 0xFF-memset at allocation (an untouched entry reads back as -1) and the kernels test != 0. The normalization must be (raw == 0); the obvious 1 - raw yields 2 for never-connected capacity-tail ranks and silently poisons every downstream sum()/bool().
2. clean_mask_buffer zeroes all capacity entries, marking the never-connected tail active — a live bug on any fleet sized above its world.
3. disconnect_ranks is suffix-only, so NIXL cannot evict a middle rank; masked-and-degraded is the terminal state there.

Each has a regression test.

nccl4py binding gap

nccl4py binds GroupConfig.enable_mask/timeout_ns but its Group stops at create/create_handle/destroy/.ptr — the five mask functions are unreachable from Python. This adds a ctypes shim on Group.ptr that tries a native Group method first, so it retires itself with no call-site churn once those bindings land. Note the symbols live in libnccl_ep.so, not libnccl.so.2; we bind the process-global namespace, which is the only resolution guaranteed to be the same library the caller's group came from.

Reconciliation, and a hole the tests found

Each transport masks locally — NCCL-EP's header calls mask consistency "a framework-level concern" — so survivors can disagree, and disagreeing masks deadlock the next dispatch. The decision is published via a single atomic compare_set, which is what prevents a split brain when a straggler's key lands between two survivors' polls; there is a dedicated regression test for that timing.

Writing those tests surfaced a case the design missed: a rank that is alive but that some peer already timed out on gets ANDed out of the group. It can't apply that decision and can't ignore it, so it now raises MoEEpRankEvictedError.

Also fixed (incidental)

• validate_fleet_params now rejects a nixl topology capacity below world size — it previously sailed through and went out of bounds inside the transport.
• update_topology now rejects growing past the capacity, for the same reason. test_update_topology_diffs_ranks was exercising that invalid config (4→6 on capacity 4); it now passes capacity=8, keeping its intent.

Testing

254 passed in tests/moe_ep/; 166 FT-related tests pass. The 6 failures in the full run are pre-existing and environmental (nvfp4/CuTeDSL ninja JIT) — confirmed failing identically on base 6258e522.

The protocol, shim and both backends' wiring are covered host-only (HashStore + threads; a fake ctypes library; fake transports), so they run in CI with no GPU.

Not yet run: the 4-GPU tiers — test_moe_ep_fault_tolerance_multirank.py (stalls a middle rank and walks detect → reconcile → degrade → re-admit, asserting the degraded output exactly equals the surviving-weight scaling) and smoke_ft_ep.py (hard SIGTERM kill). This host has 1 GPU and no transport built. bash tests/moe_ep/run_tests.sh ft runs both.

vLLM

No vLLM code here, by design. vLLM already owns the contract (support_fault_tolerance / query_active_mask / query_fault + the per-step hook in gpu_model_runner.py); the FlashInfer-EP managers just hardcode False, and since the manager is transport-parameterized, one base-class change covers both LL backends. docs/design_docs/vllm_moe_ep_integration.md §8 records the target shape and the four things that commit must get right — including that vLLM's existing nixl/deepep query_active_mask() return raw buffers with the opposite polarity, and that self._fleets holds several fleets sharing one EP group so FT state must be hoisted to a primary fleet.

Review follow-up (commit 9)

@lrbison's review caught a real defect: query_fault() on nccl_ep had no capture guard, because I'd reasoned "host read, therefore capture-safe" — exactly backwards. Being a host read is why it can't be captured, so it silently returned the capture-time answer and froze the branch on it into the graph. NIXL's already raised (it goes via query_active_mask), so the same call behaved differently per backend. Both now raise, with a regression test each. The "stale offsets" rationale was also wrong and is replaced with per-operation reasons.

Commits

1. FleetAlgoKnobFaultTolerance + Fleet FT API surface (no backend touched)
2. promote nixl_ep's store resolver to core.bootstrap_utils (pure move)
3. TCPStore-based active-mask reconciliation
4. ctypes shim for the ncclEpMask* API
5. wire nccl_ep enable_mask/timeout_ns + FT methods
6. wire nixl_ep timeout_ms + FT methods (carries the polarity/capacity/tail fixes)
7. multirank fault injection + FT smoke + run_tests.sh ft
8. docs + vLLM design note
9. fix: guard query_fault against graph capture + correct the capture rationale (review follow-up)

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features
  • Added opt-in fault tolerance for MoE expert-parallel fleets, including timeouts and coordinator takeover.
  • Introduced public APIs to query faults, manage active-rank masks, reconcile fleet state, clear faults, and track mask epochs.
  • Added backend capability detection to ensure fault tolerance only activates when supported.
• Documentation
  • Added design docs and an operational runbook describing recovery behavior, ordering rules, and transport-specific constraints.
  • Added a vLLM integration design note for future wiring into the fault-tolerance flow.
• Tests
  • Added host-only unit tests, backend-specific mocks, NCCL ctypes-shim tests, multi-GPU end-to-end coverage, and FT smoke tests.

[gemini-code-assist]

Caution

The consumer version of Gemini Code Assist on GitHub has been sunset. All code review activity has officially ceased.

[coderabbitai]

📝 Walkthrough

Walkthrough

Adds opt-in MoE expert-parallel fault tolerance with active-mask management, backend capability detection, NCCL-EP and NIXL-EP implementations, store-based reconciliation, rank eviction handling, validation, documentation, and distributed tests.

Changes

MoE EP fault tolerance

Layer / File(s)	Summary
API contracts and validation flashinfer/moe_ep/{__init__.py,algo_knobs.py,errors.py}, flashinfer/moe_ep/core/comm/fleet.py, flashinfer/moe_ep/core/validation/common.py, flashinfer/moe_ep/modes/split_layer.py, docs/design_docs/*, tests/moe_ep/test_fault_tolerance_api.py, tests/moe_ep/test_constraints.py	Adds the fault-tolerance knob, public capability and Fleet APIs, related errors, LOW_LATENCY validation, split-layer wiring, and API-level tests.
Store-based mask reconciliation flashinfer/moe_ep/core/bootstrap_utils.py, flashinfer/moe_ep/core/comm/fault_tolerance.py, tests/moe_ep/test_fault_tolerance_reconcile.py	Adds namespaced rendezvous stores and coordinated active-mask reconciliation with takeover and rank-eviction behavior.
NCCL-EP mask transport flashinfer/moe_ep/backends/split/comm/nccl_ep/*, tests/moe_ep/nccl_ep/*	Adds NCCL-EP mask FFI bindings, group configuration, mask buffers, fault queries, clearing, and mock coverage.
NIXL-EP mask transport flashinfer/moe_ep/backends/split/comm/nixl_ep/fleet.py, tests/moe_ep/nixl_ep/test_fleet_mock.py	Adds NIXL timeout handling, mask normalization, capacity checks, incremental mask updates, and readmission behavior.
Operational verification and integration notes docs/design_docs/{moe_ep_runbook.md,vllm_moe_ep_integration.md}, tests/moe_ep/{run_tests.sh,smoke_ft_ep.py,test_moe_ep_fault_tolerance_multirank.py}	Documents FT operation and future vLLM wiring, and adds host-only, smoke, and multirank recovery tests.

Estimated code review effort: 5 (Critical) | ~120 minutes

Suggested labels: op: moe

Suggested reviewers: aleozlx, yzh119, saltyminty, bkryu, sricketts

Sequence Diagram(s)

sequenceDiagram
  participant MoEEpLayer
  participant Fleet
  participant Transport
  participant Store
  MoEEpLayer->>Fleet: dispatch and combine with FT enabled
  Transport-->>Fleet: report timeout or fault
  Fleet->>Transport: query and update active mask
  Fleet->>Store: reconcile rank views
  Store-->>Fleet: agreed mask or eviction
  Fleet->>MoEEpLayer: continue degraded or readmitted execution

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 30.73% 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
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 clearly names the new fault-tolerance rank-mask work across NCCL-EP and NIXL-EP.
Description check	✅ Passed	The description is detailed and covers the change, usage, and tests, though it omits some template sections.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 6

🧹 Nitpick comments (4)

flashinfer/moe_ep/backends/split/comm/nixl_ep/fleet.py (1)

142-157: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Narrow the TypeError translation so unrelated ctor errors aren't misreported.

Any TypeError raised inside Buffer.__init__ (bad arg type, internal call) is currently reported as "build does not accept timeout_ms" whenever buf_kwargs is non-empty. Gate on the argument actually being rejected.

♻️ Suggested tightening

         except TypeError as e:
-            if not buf_kwargs:
+            if not buf_kwargs or "timeout_ms" not in str(e):
                 raise

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/moe_ep/backends/split/comm/nixl_ep/fleet.py` around lines 142 -
157, In the Buffer construction error handling, narrow the TypeError translation
in the surrounding initialization method so it only raises
MoEEpFaultToleranceUnsupportedError when the exception indicates the timeout_ms
keyword is unsupported. Re-raise unrelated TypeErrors from
nixl_ep.Buffer.__init__ unchanged, while preserving the existing behavior for
the genuine staged-build argument rejection.

tests/moe_ep/run_tests.sh (1)

221-242: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Silent pass when neither backend supports FT.

If supports_fault_tolerance is false for both nccl_ep and nixl_ep, the loop body never sets rc=1, so run_ft returns 0 and run_tests.sh ft reports "PASS" while exercising zero FT tests. Consider tracking whether at least one backend ran and surfacing that distinctly (warn or fail) so CI doesn't silently certify FT as validated when it never ran.

♻️ Proposed fix to surface a "no backend ran" state

 run_ft() {
   local rc=0
   local expected_ok=$(( NPROC_SMOKE - 1 ))
+  local ran_any=0

   for backend in nccl_ep nixl_ep; do
     if ! "${PY}" -c "from flashinfer.moe_ep import supports_fault_tolerance as s; raise SystemExit(0 if s('${backend}') else 1)"; then
       echo "${backend} cannot serve the FT API here; skipping its FT tests"
       continue
     fi
+    ran_any=1
     ...
   done

+  if [ "${ran_any}" -eq 0 ]; then
+    echo "no backend supports fault tolerance on this host; FT was not exercised" >&2
+    rc=1
+  fi
   return "${rc}"
 }

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/moe_ep/run_tests.sh` around lines 221 - 242, Update the FT backend loop
in run_ft to track whether at least one backend passes supports_fault_tolerance
and actually runs tests. After the loop, detect when none ran and return a
distinct non-success outcome or warning that prevents run_tests.sh ft from
reporting a misleading PASS, while preserving existing rc handling when a
backend does run.

docs/design_docs/moe_ep_runbook.md (1)

323-330: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Add a language hint to fenced diagrams (MD040).

Both blocks are plain-text diagrams; markdownlint flags the missing language identifier.

📝 Suggested fix

-```
+```text
 HEALTHY ──query_fault()──> FAULT_DETECTED      (local; quiesce in-flight combine/complete)

and similarly for the block at line 371.

Also applies to: 371-373

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/design_docs/moe_ep_runbook.md` around lines 323 - 330, Add the text
language identifier to the fenced code blocks containing the plain-text state
diagrams, including the blocks around the HEALTHY/DEGRADED flow and the
corresponding block near the later referenced section, without changing their
diagram contents.

Source: Linters/SAST tools

flashinfer/moe_ep/backends/split/comm/nccl_ep/_mask_ffi.py (1)

235-235: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

__all__ ordering flagged by Ruff (RUF022).

-__all__ = ["mask_ffi", "_MaskFfi"]
+__all__ = ["_MaskFfi", "mask_ffi"]

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/moe_ep/backends/split/comm/nccl_ep/_mask_ffi.py` at line 235,
Reorder the entries in __all__ to satisfy Ruff RUF022, placing the public symbol
mask_ffi before the private _MaskFfi symbol.

Source: Linters/SAST tools

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@flashinfer/moe_ep/backends/split/comm/nccl_ep/_mask_ffi.py`:
- Around line 229-232: Update mask_ffi() and its caching behavior so an
unsuccessful _MaskFfi binding is not memoized, allowing later calls to retry
after nccl.ep initialization; retain memoization only once the shim has
successfully bound the library. Preserve the existing process-wide reuse for
successful bindings and ensure supports_fault_tolerance("nccl_ep") can recover
from an early probe.

In `@flashinfer/moe_ep/backends/split/comm/nccl_ep/fleet.py`:
- Around line 179-181: Reset the topology-scoped FT state in update_topology
alongside the existing cache clearing: set _any_handle_created to False and
_ft_store_obj to None after the NCCL group is recreated. Preserve these resets
for every topology update so MaskClean cannot use stale handle state or an FT
store from the previous rank set.

In `@flashinfer/moe_ep/backends/split/comm/nixl_ep/fleet.py`:
- Around line 234-235: Update the topology-reset logic around _ft_applied in
update_topology to re-derive the applied-state mirror from the current transport
mask after resizing, rather than initializing every rank to ACTIVE. Preserve the
transport’s existing masked state for ranks retained in the new world so a later
set_active_mask([1] * world_size) detects and clears those masks correctly.

In `@flashinfer/moe_ep/core/bootstrap_utils.py`:
- Line 55: Synchronize access to the _STORE_GENS counter in
resolve_rendezvous_store so concurrent calls for the same (subsystem, ranks) key
cannot reuse a generation number. Add a module-level lock and hold it across the
get-and-increment operation, while leaving store creation and unrelated logic
unchanged.

In `@tests/moe_ep/test_moe_ep_fault_tolerance_multirank.py`:
- Around line 212-221: Update the alive-mask construction near surviving_w to
use the explicit EXPERT_MAJOR rank-to-expert ownership mapping supported by
nixl_ep, rather than indexing expected with e // experts_per_rank. Preserve the
existing dead-rank masking and assertion behavior while ensuring each expert
resolves to its owning rank under expert-major layout.
- Around line 162-226: Fix the barrier mismatch in the reconciliation and
teardown flow of the fault-tolerance test: create and use a survivors-only
process group for the two post-eviction synchronization barriers, including the
degraded-forward synchronization and final cleanup, while keeping the victim’s
existing WORLD-group barriers before eviction. Ensure the survivors-only group
is destroyed or released appropriately after use.

---

Nitpick comments:
In `@docs/design_docs/moe_ep_runbook.md`:
- Around line 323-330: Add the text language identifier to the fenced code
blocks containing the plain-text state diagrams, including the blocks around the
HEALTHY/DEGRADED flow and the corresponding block near the later referenced
section, without changing their diagram contents.

In `@flashinfer/moe_ep/backends/split/comm/nccl_ep/_mask_ffi.py`:
- Line 235: Reorder the entries in __all__ to satisfy Ruff RUF022, placing the
public symbol mask_ffi before the private _MaskFfi symbol.

In `@flashinfer/moe_ep/backends/split/comm/nixl_ep/fleet.py`:
- Around line 142-157: In the Buffer construction error handling, narrow the
TypeError translation in the surrounding initialization method so it only raises
MoEEpFaultToleranceUnsupportedError when the exception indicates the timeout_ms
keyword is unsupported. Re-raise unrelated TypeErrors from
nixl_ep.Buffer.__init__ unchanged, while preserving the existing behavior for
the genuine staged-build argument rejection.

In `@tests/moe_ep/run_tests.sh`:
- Around line 221-242: Update the FT backend loop in run_ft to track whether at
least one backend passes supports_fault_tolerance and actually runs tests. After
the loop, detect when none ran and return a distinct non-success outcome or
warning that prevents run_tests.sh ft from reporting a misleading PASS, while
preserving existing rc handling when a backend does run.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: ddaf2258-1286-4310-86f8-e80a3572dd5b

📥 Commits

Reviewing files that changed from the base of the PR and between d536efd and e5bba68.

📒 Files selected for processing (24)

• docs/design_docs/MoE_EP_impl.md
• docs/design_docs/moe_ep_architecture.md
• docs/design_docs/moe_ep_runbook.md
• docs/design_docs/vllm_moe_ep_integration.md
• flashinfer/moe_ep/__init__.py
• flashinfer/moe_ep/algo_knobs.py
• flashinfer/moe_ep/backends/split/comm/nccl_ep/_mask_ffi.py
• flashinfer/moe_ep/backends/split/comm/nccl_ep/fleet.py
• flashinfer/moe_ep/backends/split/comm/nixl_ep/fleet.py
• flashinfer/moe_ep/core/bootstrap_utils.py
• flashinfer/moe_ep/core/comm/fault_tolerance.py
• flashinfer/moe_ep/core/comm/fleet.py
• flashinfer/moe_ep/core/validation/common.py
• flashinfer/moe_ep/errors.py
• flashinfer/moe_ep/modes/split_layer.py
• tests/moe_ep/nccl_ep/test_fleet_mock.py
• tests/moe_ep/nccl_ep/test_mask_ffi.py
• tests/moe_ep/nixl_ep/test_fleet_mock.py
• tests/moe_ep/run_tests.sh
• tests/moe_ep/smoke_ft_ep.py
• tests/moe_ep/test_constraints.py
• tests/moe_ep/test_fault_tolerance_api.py
• tests/moe_ep/test_fault_tolerance_reconcile.py
• tests/moe_ep/test_moe_ep_fault_tolerance_multirank.py

[coderabbitai]

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

A failed early probe is cached for the life of the process.

mask_ffi() is functools.cached, so if the first call happens before nccl.ep has initialized and _resolve_libnccl_ep() returns None (no package path, no CONDA_PREFIX, SONAME not on the loader path), available=False is memoized permanently. A later NcclEpFleet(..., FleetAlgoKnobFaultTolerance()) then fails _check_ft_supported() even though libnccl_ep.so is loaded by that point. supports_fault_tolerance("nccl_ep") is exactly the API users are told to call first, so this ordering is likely in practice.

Suggest only memoizing a successful binding.

♻️ Cache only when the shim actually bound

-@functools.cache
-def mask_ffi() -> _MaskFfi:
-    """Process-wide shim instance (symbol binding is done once)."""
-    return _MaskFfi()
+_CACHED: "_MaskFfi | None" = None
+
+
+def mask_ffi() -> _MaskFfi:
+    """Process-wide shim instance (successful symbol binding is done once).
+
+    An unavailable shim is NOT cached: the first call may land before
+    ``nccl.ep`` has dlopened libnccl_ep, and caching that miss would make FT
+    look unsupported for the rest of the process.
+    """
+    global _CACHED
+    if _CACHED is None or not _CACHED.available:
+        _CACHED = _MaskFfi()
+    return _CACHED

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	@functools.cache
	def mask_ffi() -> _MaskFfi:
	"""Process-wide shim instance (symbol binding is done once)."""
	return _MaskFfi()
	_CACHED: "_MaskFfi | None" = None
	def mask_ffi() -> _MaskFfi:
	"""Process-wide shim instance (successful symbol binding is done once).
	An unavailable shim is NOT cached: the first call may land before
	``nccl.ep`` has dlopened libnccl_ep, and caching that miss would make FT
	look unsupported for the rest of the process.
	"""
	global _CACHED
	if _CACHED is None or not _CACHED.available:
	_CACHED = _MaskFfi()
	return _CACHED

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/moe_ep/backends/split/comm/nccl_ep/_mask_ffi.py` around lines 229
- 232, Update mask_ffi() and its caching behavior so an unsuccessful _MaskFfi
binding is not memoized, allowing later calls to retry after nccl.ep
initialization; retain memoization only once the shim has successfully bound the
library. Preserve the existing process-wide reuse for successful bindings and
ensure supports_fault_tolerance("nccl_ep") can recover from an early probe.

[coderabbitai]

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

_any_handle_created (and the FT store) survive update_topology(), defeating the MaskClean guard.

update_topology() destroys and re-creates the nccl.ep.Group, so the LL staging buffer that ncclEpMaskClean asserts on is gone until a new handle is created — but _any_handle_created stays True from the pre-resize group, so clear_faults(readmit=True) will pass the Python guard and SIGABRT from C. _ft_store_obj has the same staleness problem: it's cached against the old rank set, and reusing it across a topology change puts survivors on a prefix that no longer matches the group identity.

Reset both in update_topology():

🐛 Reset FT state on topology change

# in update_topology(), alongside self._hot_cache.clear()
self._any_handle_created = False
self._ft_store_obj = None

Also applies to: 290-290

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/moe_ep/backends/split/comm/nccl_ep/fleet.py` around lines 179 -
181, Reset the topology-scoped FT state in update_topology alongside the
existing cache clearing: set _any_handle_created to False and _ft_store_obj to
None after the NCCL group is recreated. Preserve these resets for every topology
update so MaskClean cannot use stale handle state or an FT store from the
previous rank set.

[coderabbitai]

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

_ft_applied is reset to all-ACTIVE without re-syncing the transport mask.

If a rank was masked before update_topology and is still within the new world, the transport mask keeps it masked while _ft_applied now claims ACTIVE. A subsequent set_active_mask([1]*ws) then diffs to a no-op and never un-masks it, so the fleet stays silently degraded (and query_fault() reports a permanent fault). Re-derive the mirror from the transport instead of assuming all-active.

🐛 Suggested re-sync

-        self._ft_applied = [ACTIVE] * bootstrap.world_size
         self._hot_ft_bufs = None
+        if self._ft is not None:
+            # The transport mask survives connect/disconnect, so re-derive the
+            # mirror rather than assuming a clean slate.
+            self._ft_applied = [
+                int(v) for v in self.query_active_mask().cpu().tolist()
+            ]
+        else:
+            self._ft_applied = [ACTIVE] * bootstrap.world_size

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	self._ft_applied = [ACTIVE] * bootstrap.world_size
	self._hot_ft_bufs = None
	self._hot_ft_bufs = None
	if self._ft is not None:
	# The transport mask survives connect/disconnect, so re-derive the
	# mirror rather than assuming a clean slate.
	self._ft_applied = [
	int(v) for v in self.query_active_mask().cpu().tolist()
	]
	else:
	self._ft_applied = [ACTIVE] * bootstrap.world_size

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/moe_ep/backends/split/comm/nixl_ep/fleet.py` around lines 234 -
235, Update the topology-reset logic around _ft_applied in update_topology to
re-derive the applied-state mirror from the current transport mask after
resizing, rather than initializing every rank to ACTIVE. Preserve the
transport’s existing masked state for ranks retained in the new world so a later
set_active_mask([1] * world_size) detects and clears those masks correctly.

[coderabbitai]

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

Unsynchronized read-modify-write on _STORE_GENS.

gen = _STORE_GENS.get(key, 0) followed by _STORE_GENS[key] = gen + 1 is not atomic. Concurrent calls to resolve_rendezvous_store from multiple threads in the same process for the same (subsystem, ranks) key can race and hand out a duplicate/lost generation number, which is exactly the collision this counter was introduced to prevent — the resulting store-prefix mismatch across ranks would silently break rendezvous rather than raise.

🔒 Suggested fix

+import threading
+
+_STORE_GENS_LOCK = threading.Lock()
 _STORE_GENS: dict = {}
...
-    gen = _STORE_GENS.get(key, 0)
-    _STORE_GENS[key] = gen + 1
+    with _STORE_GENS_LOCK:
+        gen = _STORE_GENS.get(key, 0)
+        _STORE_GENS[key] = gen + 1

Also applies to: 105-115

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/moe_ep/core/bootstrap_utils.py` at line 55, Synchronize access to
the _STORE_GENS counter in resolve_rendezvous_store so concurrent calls for the
same (subsystem, ranks) key cannot reuse a generation number. Add a module-level
lock and hold it across the get-and-increment operation, while leaving store
creation and unrelated logic unchanged.

[coderabbitai]

🩺 Stability & Availability | 🔴 Critical | 🏗️ Heavy lift

Mismatched dist.barrier() collective counts between the victim and survivors will hang the survivors.

Victim calls dist.barrier() on the default (WORLD) process group 3 times in this test (lines 166, 179, 196) before returning at line 197. Survivors call it 4 times (166, 179, 223, 226). Since dist.barrier() collectives on a shared process group are matched by issue-order per rank (not by source line), this 1-call deficit persists for the rest of the module: test_readmit_restores_full_strength adds 4 symmetric calls per rank afterward, so the deficit surfaces at the very last barrier (line 303), which survivors will block on until _PG_TIMEOUT (60 min) elapses with no victim participant left to match it. This directly contradicts the file's own docstring claim that "every process stays alive, so the torch process group stays usable for the test's own barriers" (lines 8-15) — the victim is alive, it just stops calling barrier() two calls early.

This is very likely why the PR reports the 4-GPU multirank tests were never actually executed (dev host has 1 GPU) — the bug would only manifest under a real 4-rank run.

🔒 Proposed fix: give the two post-eviction syncs their own survivors-only group

     survivors = [r for r in range(world_size) if r != _VICTIM]
     expected = [1 if r in survivors else 0 for r in range(world_size)]
+    survivors_group = dist.new_group(survivors)
     ...
     if rank == _VICTIM:
         with pytest.raises(MoEEpRankEvictedError):
             fleet.reconcile_active_mask()
         print(f"rank {rank}: correctly evicted")
         layer.destroy()
         dist.barrier()
         return

     agreed = fleet.reconcile_active_mask().cpu().tolist()
     assert agreed == expected, f"rank {rank} agreed {agreed}, expected {expected}"
     fleet.clear_faults(readmit=False)
     assert fleet.query_fault() is False, "clear_faults must re-arm detection"
     ...
     print(f"rank {rank}: degraded forward matches surviving-weight scaling")
-    dist.barrier()
+    dist.barrier(group=survivors_group)

     layer.destroy()
-    dist.barrier()
+    dist.barrier(group=survivors_group)

(Alternatively, simply add two matching dist.barrier() calls in the victim branch before return, but a survivors-only group is more robust against future edits changing the call count on either side.)

📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change

	# --- HEALTHY ---------------------------------------------------------
	y = layer.forward(t)
	torch.cuda.synchronize()
	torch.testing.assert_close(y, x, atol=5e-2, rtol=5e-2)
	dist.barrier()
	fleet = layer._ensure_fleet()
	assert fleet.supports_fault_tolerance is True
	assert fleet.query_active_mask().cpu().tolist() == [1] * world_size
	assert fleet.query_fault() is False
	# --- FAULT: the victim stalls past the timeout ------------------------
	if rank == _VICTIM:
	time.sleep(_STALL_S)
	else:
	layer.forward(t) # survivors' kernels time out waiting on the victim
	torch.cuda.synchronize()
	dist.barrier() # the victim is alive, so the torch PG is still healthy
	if rank in survivors:
	assert fleet.query_fault() is True, "survivors must observe the fault"
	mask = fleet.query_active_mask().cpu().tolist()
	assert mask[_VICTIM] == 0, f"rank {_VICTIM} should be masked, got {mask}"
	# --- RECONCILE -------------------------------------------------------
	# Every rank reconciles in the same iteration slot. The victim learns it
	# was evicted (it stalled long enough for the survivors to give up).
	from flashinfer.moe_ep.errors import MoEEpRankEvictedError
	if rank == _VICTIM:
	with pytest.raises(MoEEpRankEvictedError):
	fleet.reconcile_active_mask()
	print(f"rank {rank}: correctly evicted")
	layer.destroy()
	dist.barrier()
	return
	agreed = fleet.reconcile_active_mask().cpu().tolist()
	assert agreed == expected, f"rank {rank} agreed {agreed}, expected {expected}"
	fleet.clear_faults(readmit=False)
	assert fleet.query_fault() is False, "clear_faults must re-arm detection"
	# --- DEGRADED --------------------------------------------------------
	# The dead rank's experts are gone and combine does NOT renormalize, so
	# each token comes out scaled by the surviving weight fraction rather
	# than 1. Assert exactly that, not "close to x".
	y_deg = layer.forward(t)
	torch.cuda.synchronize()
	assert torch.isfinite(y_deg.float()).all(), "degraded output must not be NaN/Inf"
	experts_per_rank = num_experts // world_size
	alive = torch.tensor(
	[expected[e // experts_per_rank] for e in range(num_experts)],
	device="cuda",
	dtype=torch.bool,
	)
	surviving_w = (topk_weights * alive[topk_ids]).sum(-1) # [num_tokens]
	torch.testing.assert_close(
	y_deg.float(), x.float() * surviving_w.unsqueeze(-1), atol=8e-2, rtol=8e-2
	)
	print(f"rank {rank}: degraded forward matches surviving-weight scaling")
	dist.barrier()
	layer.destroy()
	dist.barrier()
	survivors_group = dist.new_group(survivors)
	# --- HEALTHY ---------------------------------------------------------
	y = layer.forward(t)
	torch.cuda.synchronize()
	torch.testing.assert_close(y, x, atol=5e-2, rtol=5e-2)
	dist.barrier()
	fleet = layer._ensure_fleet()
	assert fleet.supports_fault_tolerance is True
	assert fleet.query_active_mask().cpu().tolist() == [1] * world_size
	assert fleet.query_fault() is False
	# --- FAULT: the victim stalls past the timeout ------------------------
	if rank == _VICTIM:
	time.sleep(_STALL_S)
	else:
	layer.forward(t) # survivors' kernels time out waiting on the victim
	torch.cuda.synchronize()
	dist.barrier() # the victim is alive, so the torch PG is still healthy
	if rank in survivors:
	assert fleet.query_fault() is True, "survivors must observe the fault"
	mask = fleet.query_active_mask().cpu().tolist()
	assert mask[_VICTIM] == 0, f"rank {_VICTIM} should be masked, got {mask}"
	# --- RECONCILE -------------------------------------------------------
	# Every rank reconciles in the same iteration slot. The victim learns it
	# was evicted (it stalled long enough for the survivors to give up).
	from flashinfer.moe_ep.errors import MoEEpRankEvictedError
	if rank == _VICTIM:
	with pytest.raises(MoEEpRankEvictedError):
	fleet.reconcile_active_mask()
	print(f"rank {rank}: correctly evicted")
	layer.destroy()
	dist.barrier()
	return
	agreed = fleet.reconcile_active_mask().cpu().tolist()
	assert agreed == expected, f"rank {rank} agreed {agreed}, expected {expected}"
	fleet.clear_faults(readmit=False)
	assert fleet.query_fault() is False, "clear_faults must re-arm detection"
	# --- DEGRADED --------------------------------------------------------
	# The dead rank's experts are gone and combine does NOT renormalize, so
	# each token comes out scaled by the surviving weight fraction rather
	# than 1. Assert exactly that, not "close to x".
	y_deg = layer.forward(t)
	torch.cuda.synchronize()
	assert torch.isfinite(y_deg.float()).all(), "degraded output must not be NaN/Inf"
	experts_per_rank = num_experts // world_size
	alive = torch.tensor(
	[expected[e // experts_per_rank] for e in range(num_experts)],
	device="cuda",
	dtype=torch.bool,
	)
	surviving_w = (topk_weights * alive[topk_ids]).sum(-1) # [num_tokens]
	torch.testing.assert_close(
	y_deg.float(), x.float() * surviving_w.unsqueeze(-1), atol=8e-2, rtol=8e-2
	)
	print(f"rank {rank}: degraded forward matches surviving-weight scaling")
	dist.barrier(group=survivors_group)
	layer.destroy()
	dist.barrier(group=survivors_group)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/moe_ep/test_moe_ep_fault_tolerance_multirank.py` around lines 162 -
226, Fix the barrier mismatch in the reconciliation and teardown flow of the
fault-tolerance test: create and use a survivors-only process group for the two
post-eviction synchronization barriers, including the degraded-forward
synchronization and final cleanup, while keeping the victim’s existing
WORLD-group barriers before eviction. Ensure the survivors-only group is
destroyed or released appropriately after use.

[lrbison]

from the NCCL perspective it looks about right to me. I scanned the documentation and the code. I have a few comments about these "gotchyas":

1. Mutate the mask only between iterations (kernels read it live).

True, however the expectation is the kernel will discover the problem first, so torch application should be typically checking the mask/error rather than trying to update it.

3. clear_faults(readmit=True) before update_topology, never after: ncclEpMaskClean needs a live handle on the current group, which update_topology destroys.

As I understand it, update_topology() creates a new NCCL communicator? (it is the only option for NCCL to admit a new process). In that case the graph has changed as well.

4. No FT call during CUDA-graph capture — a captured query replays stale offsets, a captured set freezes one mask into every replay. All stream-ordered entry points raise on capture.

I think I agree with this point just not the wording. I agree that querying if mask has changed or if there has been an error during a graph capture is going to capture a graph that ignores errors by assuming whatever the capture saw will be true forever.

But regarding stale offsets, NCCL-EP doesn't "squeeze" the remaining ranks in data layout. Instead the data arriving from those ranks is left in an unknown state, and an error is flagged. The application may end up invoking the dispatch and combine again and again before the CUDA-graph is complete, and each time we respect the mask and don't wait on masked ranks. In fact, the application may choose to re-route experts away from the failed rank with something like EPLB, in which case the rank remains failed, but the application can clear the error state and continue with a partial mask (without update_topology). I think you call this clear_faults(readmit=False) right?

[Anerudhan]

True, however the expectation is the kernel will discover the problem first, so torch application should be typically checking the mask/error rather than trying to update it.

In steady state, query_fault() and query_active_mask() are read only.

set_active_mask() is not to be called, but if a kernel is timed out only the local rank will set the mask.

If there are two rank which have different mask we need to combine them with bitwise and, to have agreed survivors ?

Done via reconcile_active_mask().

Hope that makes sense. Let me know if my understanding is incorrect.

As I understand it, update_topology() creates a new NCCL communicator? (it is the only option for NCCL to admit a new process). In that case the graph has changed as well.

Yes. update_topology creates new communicator, and then creates a new group over it.

• clear_faults(readmit=True) --> ncclEpMaskClean on the communicator: re-admits a merely-delayed rank.
• update_topology() → new communicator + new group: the only way to add or replace a process.

No FT call during CUDA-graph capture — a captured query replays stale offsets, a captured set freezes one mask into every replay. All stream-ordered entry points raise on capture.

Correct, will update the doc.

Question to you @lrbison
ncclEpMaskUpdate has no callerd, so I have kept the reconciliation "push" step separable. Is the framework pushing a
reconciled mask back via ncclEpMaskUpdate the use it was for ?

[coderabbitai]

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (2)

docs/design_docs/moe_ep_runbook.md (2)

439-440: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Do not use a clamp as the zero-survivor fallback.

When surviving_w == 0, clamping to 1e-6 does not drop the token; it can amplify residual output instead. Explicitly mask/drop zero-survivor tokens and divide only rows with positive surviving weight.

Proposed handling

- out.div_(surviving_w.clamp_min(1e-6).unsqueeze(-1))
+ valid = surviving_w > 0
+ out[valid] /= surviving_w[valid].unsqueeze(-1)
+ out[~valid] = 0

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/design_docs/moe_ep_runbook.md` around lines 439 - 440, Update the
normalization logic around surviving_w so zero-survivor tokens are explicitly
masked or dropped rather than handled with clamp_min. Divide only rows where
surviving_w is positive, while preserving the existing normalization for valid
rows and ensuring zero-survivor outputs are not amplified.

---

436-440: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Guard the expert-owner mapping against invalid topology shapes.

This formula assumes num_experts is evenly divisible by world_size; otherwise the final experts can map to owner world_size and alive[...] raises an out-of-bounds error. It also divides by zero when there are fewer than world_size experts. State and validate the precondition, or use the framework’s actual expert-ownership map.

Proposed validation

+if num_experts <= 0 or num_experts % world_size != 0:
+    raise ValueError("num_experts must be a positive multiple of world_size")
 owner = torch.arange(num_experts, device=d) // (num_experts // world_size)

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docs/design_docs/moe_ep_runbook.md` around lines 436 - 440, Update the
expert-owner mapping around owner and alive indexing to validate that
num_experts is at least world_size and evenly divisible by world_size before
applying the formula; otherwise fail with a clear precondition error. Prefer the
framework’s authoritative expert-ownership map if available, while preserving
the existing masking and normalization behavior for valid topology shapes.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docs/design_docs/moe_ep_runbook.md`:
- Line 424: Update the fenced code block in the runbook section to include a
language identifier, using text or another appropriate language after the
opening fence, while preserving the block’s contents.
- Around line 366-387: Update the CUDA-graph capture contract section to state
one definitive behavior for every guarded fault-tolerance API, preferably that
each call is rejected immediately during capture. Explicitly enumerate all four
APIs, including query_fault(), set_active_mask(), query_active_mask(), and the
remaining guarded operation identified elsewhere in the document, and remove
contradictory claims that query_fault() returns a capture-time value or that
only three operations are listed.

---

Outside diff comments:
In `@docs/design_docs/moe_ep_runbook.md`:
- Around line 439-440: Update the normalization logic around surviving_w so
zero-survivor tokens are explicitly masked or dropped rather than handled with
clamp_min. Divide only rows where surviving_w is positive, while preserving the
existing normalization for valid rows and ensuring zero-survivor outputs are not
amplified.
- Around line 436-440: Update the expert-owner mapping around owner and alive
indexing to validate that num_experts is at least world_size and evenly
divisible by world_size before applying the formula; otherwise fail with a clear
precondition error. Prefer the framework’s authoritative expert-ownership map if
available, while preserving the existing masking and normalization behavior for
valid topology shapes.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: b8871016-62c4-41f5-8ddb-1939ce65f0f4

📥 Commits

Reviewing files that changed from the base of the PR and between e5bba68 and b7cd997.

📒 Files selected for processing (7)

• docs/design_docs/moe_ep_runbook.md
• flashinfer/moe_ep/backends/split/comm/nccl_ep/fleet.py
• flashinfer/moe_ep/backends/split/comm/nixl_ep/fleet.py
• flashinfer/moe_ep/core/comm/fault_tolerance.py
• flashinfer/moe_ep/core/comm/fleet.py
• tests/moe_ep/nccl_ep/test_fleet_mock.py
• tests/moe_ep/nixl_ep/test_fleet_mock.py

🚧 Files skipped from review as they are similar to previous changes (6)

• tests/moe_ep/nixl_ep/test_fleet_mock.py
• flashinfer/moe_ep/core/comm/fault_tolerance.py
• tests/moe_ep/nccl_ep/test_fleet_mock.py
• flashinfer/moe_ep/core/comm/fleet.py
• flashinfer/moe_ep/backends/split/comm/nixl_ep/fleet.py
• flashinfer/moe_ep/backends/split/comm/nccl_ep/fleet.py

[lrbison]

@Anerudhan

Yes -- I think we are in agreement about query_active_mask / set_active_mask.

If there are two rank which have different mask we need to combine them with bitwise and, to have agreed survivors

The approach we took was to assume that communication would naturally drive eventual consistency. However you are right to explicitly synchronize masks at the FlashInfer level to be sure.

ncclEpMaskUpdate: Is the framework pushing a reconciled mask back via ncclEpMaskUpdate the use it was for ?

Yes, the use case is framework-detected updates. This could be recovering from a false-positive timeout: Some rank decided to move on and mask out a peer, but after the forward pass the framework did its own synchronization and decided that actually all peers are healthy and should rejoin without any process changes. Or it could be some separate scale-down signal that application is handling by masking and exiting some of the ranks. For scale-up or replacement the API is not useful unless NCCL eventually provides a way to admit new processes into an existing group (which isn't possible today).

[bkryu]

/bot run tests/moe_ep

[flashinfer-bot]

GitLab MR !1066 has been created, and the CI pipeline #59999869 is currently running. I'll report back once the pipeline job completes.

[bkryu]

Hi @Anerudhan, CI is failing on test_deep_gemm_mega_kernel_vs_reference.py in cu130 configs. Plz take a look

[mhoqueanik]

Hi @bkryu! Filed a fix PR here: #4221

[Anerudhan]

#59999869

@mhoqueanik is going to be fixing this in a separate PR. Not introduced by this PR

Covered in #4221