diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu index 0c4431b8eaef..0f2d453c363a 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu @@ -741,7 +741,7 @@ __device__ void vectorized_combine_impl(T* dst_typed_base, int size_per_token, i { int target_rank = ptrs.topk_target_ranks[local_token_idx * TOP_K + k]; int dst_idx = ptrs.topk_send_indices[local_token_idx * TOP_K + k]; - if (dst_idx < 0) + if (dst_idx < 0 || !is_rank_active(ptrs.active_rank_mask, target_rank)) { acc[k].fill(0.0f); continue; @@ -766,8 +766,12 @@ __device__ void vectorized_combine_impl(T* dst_typed_base, int size_per_token, i #pragma unroll for (int k = 0; k < TOP_K; ++k) { - if (ptrs.topk_send_indices[local_token_idx * TOP_K + k] < 0) + int target_rank = ptrs.topk_target_ranks[local_token_idx * TOP_K + k]; + int dst_idx = ptrs.topk_send_indices[local_token_idx * TOP_K + k]; + if (dst_idx < 0 || !is_rank_active(ptrs.active_rank_mask, target_rank)) + { continue; // acc[k] already holds 0.0f from fill() above + } #pragma unroll for (int j = elems_per_vec - 1; j >= 0; --j) acc[k][j] = static_cast(reinterpret_cast(&acc[k])[j]); diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h index 9a6f3904c501..138ca92e71a8 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h @@ -93,8 +93,8 @@ struct CombineKernelPointers int const* topk_send_indices; // dst index per k, -1 for duplicates // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Combine skips flag - // writes/waits to/from masked peers; per-token accumulation uses topk_send_indices[k] < 0 - // (set by dispatch) to skip dead-targeted slots, so no explicit mask check is needed there. + // writes/waits to/from masked peers and also skips per-token accumulation for ranks that + // become inactive between dispatch and combine. uint64_t active_rank_mask[kRankMaskWords]; }; diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py new file mode 100644 index 000000000000..2b71aaefa257 --- /dev/null +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -0,0 +1,783 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Host-side watchdog for MoE AlltoAll completion flags. + +The NVLinkOneSided kernels signal each collective by writing the current +``flag_val`` into the rank-local completion flag table. A dead peer in the +silent-spin failure mode never writes its slot, so this watchdog polls the same +table from a CPU thread and reports peers whose flags do not reach the expected +generation before a bounded timeout. + +Timeouts are detection events, not committed membership changes. The optional +EP health object supplies the already-committed active peers to watch, and is +read-only here. Higher-layer recovery coordination consumes ``on_timeout`` +events and commits a new membership only after the data plane and expert +placement are ready for the same generation. +""" + +from __future__ import annotations + +import threading +import time +from collections import deque +from collections.abc import Callable, Mapping, MutableMapping, Sequence +from dataclasses import dataclass +from typing import Protocol + +import torch + +from tensorrt_llm._utils import prefer_pinned +from tensorrt_llm.logger import logger as tllm_logger + +DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S = 5.0 +DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S = 0.1 +UNKNOWN_COMPLETION_FLAG = -(2**63) +_COMPLETION_FLAG_MASK = (1 << 32) - 1 +_COMPLETION_FLAG_HALF_RANGE = 1 << 31 +_ACTIVE_RANK_MASK_WORD_BITS = 64 +_ACTIVE_RANK_MASK_WORDS = 2 +_ACTIVE_RANK_MASK_WORD_MASK = (1 << _ACTIVE_RANK_MASK_WORD_BITS) - 1 +_WORKSPACE_WATCHDOG_STATE_KEY = "alltoall_watchdog_shared_state" +_WORKSPACE_WATCHDOG_STATE_INIT_LOCK = threading.Lock() + + +def _normalize_completion_flag(value: int) -> int: + return int(value) & _COMPLETION_FLAG_MASK + + +def _completion_flag_reached(observed: int, expected: int) -> bool: + if observed == UNKNOWN_COMPLETION_FLAG: + return False + # Counter values are ordered modulo uint32. A queued watch cannot lag by + # half the counter space within the bounded watchdog timeout. + distance = (_normalize_completion_flag(observed) - expected) & _COMPLETION_FLAG_MASK + return distance < _COMPLETION_FLAG_HALF_RANGE + + +class CompletionFlagReader(Protocol): + """Reads one phase's rank-local completion flag row.""" + + def read_completion_flags(self, phase: str) -> Sequence[int]: + """Return ``ep_size`` flag values for ``phase``.""" + + +class EPGroupHealthSnapshotLike(Protocol): + """Atomic fields required from a committed-membership snapshot.""" + + @property + def mask(self) -> int: + """Return the committed active-rank bitmask.""" + + @property + def generation(self) -> int: + """Return the matching committed-membership generation.""" + + +class EPGroupHealthLike(Protocol): + """Read-only committed EP membership used by AlltoAll frontends.""" + + def get_mask(self) -> int: + """Return the active-rank bitmask.""" + + def get_mask_words(self) -> tuple[int, ...]: + """Return the active-rank bitmask split into uint64 words.""" + + def snapshot(self) -> EPGroupHealthSnapshotLike: + """Return one atomic committed mask and generation snapshot.""" + + +@dataclass(frozen=True) +class ActiveRankMaskSnapshot: + """One dispatch epoch's rank mask and optional committed generation. + + ``committed_generation`` is populated only when ``active_rank_mask`` was + read from the coordinator-owned EP health object. Explicit caller masks do + not inherit a generation from unrelated health state. + """ + + active_rank_mask: torch.Tensor | None + committed_generation: int | None + + +class CompletionFlagReadTimeout(TimeoutError): + """Raised when the host watchdog cannot read completion flags in time.""" + + +@dataclass(frozen=True) +class AlltoAllWatchdogTimeout: + """Detection details emitted when an AlltoAll phase times out. + + ``missing_ranks`` contains suspects relative to the committed active mask + captured for this collective. Receiving this event does not publish a new + active mask. + """ + + phase: str + expected_flag: int + observed_flags: tuple[int, ...] + missing_ranks: tuple[int, ...] + elapsed_s: float + poll_timed_out: bool = False + + +@dataclass(frozen=True) +class _CollectiveWatch: + phase: str + expected_flag: int + active_mask: int + start_s: float + + +@dataclass(frozen=True) +class _SharedWatchdogConfig: + ep_size: int + ep_rank: int + timeout_s: float + poll_interval_s: float + health: EPGroupHealthLike | None + on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None + + +@dataclass +class _WorkspaceWatchdogState: + lock: threading.Lock + generation: int = 0 + watchdog: AlltoAllWatchdog | None = None + config: _SharedWatchdogConfig | None = None + ref_count: int = 0 + + +class _TorchCompletionFlagReader: + """Completion-flag reader backed by the MoE AlltoAll workspace tensor.""" + + def __init__( + self, + workspace: torch.Tensor, + ep_rank: int, + ep_size: int, + dispatch_completion_flags_offset: int, + combine_completion_flags_offset: int, + device_copy_timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + ) -> None: + if workspace.dim() != 2: + raise ValueError("workspace must be a 2D tensor [ep_size, size_per_rank]") + if not 0 <= ep_rank < ep_size: + raise ValueError(f"ep_rank must be in [0, {ep_size}), got {ep_rank}") + if workspace.size(0) != ep_size: + raise ValueError( + f"workspace first dimension must equal ep_size={ep_size}, got {workspace.size(0)}" + ) + self._workspace = workspace + self._ep_rank = ep_rank + self._ep_size = ep_size + self._offsets = { + "dispatch": int(dispatch_completion_flags_offset), + "combine": int(combine_completion_flags_offset), + } + self._device_copy_timeout_s = float(device_copy_timeout_s) + self._copy_stream: torch.cuda.Stream | None = None + self._host_flags: torch.Tensor | None = None + self._copy_event: torch.cuda.Event | None = None + self._retired_copies: list[tuple[torch.Tensor, torch.cuda.Event]] = [] + if workspace.device.type == "cuda": + self._copy_stream = torch.cuda.Stream(device=workspace.device) + + def _prune_retired_copies(self) -> None: + self._retired_copies = [ + (host_flags, event) for host_flags, event in self._retired_copies if not event.query() + ] + + def _read_cuda_flags(self, flags: torch.Tensor) -> tuple[int, ...]: + assert self._copy_stream is not None + self._prune_retired_copies() + + if self._host_flags is None: + self._host_flags = torch.empty( + (self._ep_size,), + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + if self._copy_event is None: + self._copy_event = torch.cuda.Event(blocking=False) + host_flags = self._host_flags + event = self._copy_event + with torch.cuda.device(flags.device), torch.cuda.stream(self._copy_stream): + host_flags.copy_(flags.detach(), non_blocking=True) + event.record(self._copy_stream) + + deadline_s = time.monotonic() + self._device_copy_timeout_s + while not event.query(): + remaining_s = deadline_s - time.monotonic() + if remaining_s <= 0: + self._retired_copies.append((host_flags, event)) + self._host_flags = None + self._copy_event = None + raise CompletionFlagReadTimeout( + "timed out copying AlltoAll completion flags to host" + ) + time.sleep(min(remaining_s, 0.001)) + + return tuple(int(v) for v in host_flags.tolist()) + + def read_completion_flags(self, phase: str) -> tuple[int, ...]: + offset = self._offsets[phase] + end = offset + self._ep_size * 4 + flags = self._workspace[self._ep_rank, offset:end].view(torch.int32) + if flags.device.type == "cuda": + return self._read_cuda_flags(flags) + if flags.device.type != "cpu": + flags = flags.detach().cpu() + return tuple(int(v) for v in flags.tolist()) + + +class AlltoAllWatchdogCoordinator: + """Shared watchdog plumbing for MoE AlltoAll frontends. + + ``health`` is the committed membership snapshot source. This class reads + it when a dispatch starts but never publishes detected failures into it. + """ + + def __init__( + self, + *, + workspace_state: MutableMapping[str, object], + workspace: torch.Tensor, + metainfo: torch.Tensor, + metainfo_index: Mapping[str, int], + ep_rank: int, + health: EPGroupHealthLike | None = None, + ) -> None: + self._workspace_state = workspace_state + self._workspace = workspace + self._metainfo = metainfo + self._metainfo_index = metainfo_index + self._ep_rank = ep_rank + self._health = health + + def _find_shared_state(self) -> _WorkspaceWatchdogState | None: + state = self._workspace_state.get(_WORKSPACE_WATCHDOG_STATE_KEY) + if state is None: + return None + if not isinstance(state, _WorkspaceWatchdogState): + raise TypeError("invalid shared AlltoAll watchdog workspace state") + return state + + def _get_shared_state(self) -> _WorkspaceWatchdogState: + state = self._find_shared_state() + if state is not None: + return state + with _WORKSPACE_WATCHDOG_STATE_INIT_LOCK: + state = self._workspace_state.get(_WORKSPACE_WATCHDOG_STATE_KEY) + if state is None: + state = _WorkspaceWatchdogState(lock=threading.Lock()) + self._workspace_state[_WORKSPACE_WATCHDOG_STATE_KEY] = state + if not isinstance(state, _WorkspaceWatchdogState): + raise TypeError("invalid shared AlltoAll watchdog workspace state") + return state + + def read_current_flag_val(self) -> int: + flag_val_offset = int(self._metainfo[self._metainfo_index["FLAG_VAL_OFFSET_INDEX"]].item()) + flag_val = self._workspace[self._ep_rank, flag_val_offset : flag_val_offset + 4].view( + torch.int32 + ) + if flag_val.device.type != "cpu": + flag_val = flag_val.detach().cpu() + return _normalize_completion_flag(int(flag_val.item())) + + def capture_active_rank_mask( + self, active_rank_mask: torch.Tensor | None + ) -> ActiveRankMaskSnapshot: + """Capture a caller override or one coherent committed-mask epoch. + + The returned tensor does not alias a caller-owned override, so it can + safely represent the mask used by one dispatch/combine pair. When the + mask comes from ``health``, derive both fixed-width ABI words from one + atomic mask/generation snapshot. + """ + if active_rank_mask is not None: + return ActiveRankMaskSnapshot( + active_rank_mask=active_rank_mask.detach().clone(), + committed_generation=None, + ) + if self._health is None: + return ActiveRankMaskSnapshot( + active_rank_mask=None, + committed_generation=None, + ) + + health_snapshot = self._health.snapshot() + mask_words = tuple( + (health_snapshot.mask >> (_ACTIVE_RANK_MASK_WORD_BITS * index)) + & _ACTIVE_RANK_MASK_WORD_MASK + for index in range(_ACTIVE_RANK_MASK_WORDS) + ) + return ActiveRankMaskSnapshot( + active_rank_mask=torch.tensor( + mask_words, + dtype=torch.uint64, + device="cpu", + ), + committed_generation=health_snapshot.generation, + ) + + def active_rank_mask_tensor(self, active_rank_mask: torch.Tensor | None) -> torch.Tensor | None: + """Return only the tensor from :meth:`capture_active_rank_mask`. + + Dispatch frontends must retain the full snapshot so combine can verify + the committed generation before launching its kernel. + """ + return self.capture_active_rank_mask(active_rank_mask).active_rank_mask + + def active_rank_mask_for_combine( + self, + dispatch_snapshot: ActiveRankMaskSnapshot, + requested_active_rank_mask: torch.Tensor | None, + ) -> torch.Tensor | None: + """Validate the dispatch epoch and return its captured rank mask. + + A coordinator-owned generation change between dispatch and combine is + an invalid asynchronous commit, so fail closed before launching the + combine kernel. Explicit caller masks are independent of health and + therefore carry no committed generation. + """ + dispatch_active_rank_mask = dispatch_snapshot.active_rank_mask + committed_generation = dispatch_snapshot.committed_generation + if committed_generation is not None: + if self._health is None: + raise RuntimeError("committed rank-mask snapshot has no EP health source") + current_generation = self._health.snapshot().generation + if current_generation != committed_generation: + raise RuntimeError( + "committed EP membership changed between dispatch and combine " + f"(generation {committed_generation} -> {current_generation}); " + "aborting collective epoch" + ) + + if requested_active_rank_mask is None: + return dispatch_active_rank_mask + requested_mask = self.active_rank_mask_tensor(requested_active_rank_mask) + assert requested_mask is not None + if dispatch_active_rank_mask is None or not torch.equal( + dispatch_active_rank_mask, requested_mask + ): + raise ValueError("active_rank_mask must match the mask captured at dispatch") + return dispatch_active_rank_mask + + def active_mask_int(self, active_rank_mask: torch.Tensor | None) -> int | None: + if active_rank_mask is not None: + mask_cpu = active_rank_mask.detach().cpu() + return sum(int(word) << (64 * idx) for idx, word in enumerate(mask_cpu.tolist())) + if self._health is not None: + return self._health.get_mask() + return None + + def acquire_watchdog( + self, + *, + ep_size: int, + timeout_s: float, + poll_interval_s: float, + on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None = None, + ) -> AlltoAllWatchdog: + config = _SharedWatchdogConfig( + ep_size=int(ep_size), + ep_rank=self._ep_rank, + timeout_s=float(timeout_s), + poll_interval_s=float(poll_interval_s), + health=self._health, + on_timeout=on_timeout, + ) + state = self._get_shared_state() + with state.lock: + if state.watchdog is None: + state.generation = self.read_current_flag_val() + state.watchdog = AlltoAllWatchdog.from_workspace( + workspace=self._workspace, + metainfo=self._metainfo, + metainfo_index=self._metainfo_index, + ep_rank=self._ep_rank, + ep_size=ep_size, + timeout_s=timeout_s, + poll_interval_s=poll_interval_s, + health=self._health, + on_timeout=on_timeout, + ) + state.config = config + elif not self._matching_config(state.config, config): + raise ValueError( + "AlltoAll wrappers sharing a workspace must use the same " + "watchdog configuration and EP health object" + ) + state.ref_count += 1 + watchdog = state.watchdog + assert watchdog is not None + return watchdog + + @staticmethod + def _matching_config( + existing: _SharedWatchdogConfig | None, + requested: _SharedWatchdogConfig, + ) -> bool: + return ( + existing is not None + and existing.ep_size == requested.ep_size + and existing.ep_rank == requested.ep_rank + and existing.timeout_s == requested.timeout_s + and existing.poll_interval_s == requested.poll_interval_s + and existing.health is requested.health + and existing.on_timeout is requested.on_timeout + ) + + def release_watchdog(self, watchdog: AlltoAllWatchdog) -> None: + state = self._get_shared_state() + watchdog_to_stop: AlltoAllWatchdog | None = None + with state.lock: + if state.watchdog is not watchdog or state.ref_count <= 0: + raise RuntimeError("attempted to release an unregistered AlltoAll watchdog") + state.ref_count -= 1 + if state.ref_count == 0: + watchdog_to_stop = state.watchdog + state.watchdog = None + state.config = None + if watchdog_to_stop is not None: + watchdog_to_stop.stop() + + def watch_collective( + self, + watchdog: AlltoAllWatchdog | None, + phase: str, + active_rank_mask: torch.Tensor | None, + ) -> None: + if watchdog is None: + state = self._find_shared_state() + if state is not None: + with state.lock: + if state.watchdog is not None: + state.generation = (state.generation + 1) & _COMPLETION_FLAG_MASK + return + active_mask = self.active_mask_int(active_rank_mask) + state = self._get_shared_state() + with state.lock: + if state.watchdog is not watchdog: + raise RuntimeError("AlltoAll watchdog is not registered for this workspace") + state.generation = (state.generation + 1) & _COMPLETION_FLAG_MASK + watchdog.watch( + phase=phase, + expected_flag=state.generation, + active_mask=active_mask, + ) + + +class AlltoAllWatchdog: + """Background host thread that watches AlltoAll completion flags. + + The watchdog is intentionally opt-in. Callers queue phases with + :meth:`watch`; the thread polls them in FIFO order so a queued combine cannot + hide a still-spinning dispatch. A timeout is reported through + ``on_timeout`` without mutating the committed EP membership. + """ + + VALID_PHASES = frozenset({"dispatch", "combine"}) + + def __init__( + self, + *, + ep_size: int, + ep_rank: int, + completion_reader: CompletionFlagReader, + timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + health: EPGroupHealthLike | None = None, + on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None = None, + ) -> None: + if ep_size <= 0: + raise ValueError(f"ep_size must be > 0, got {ep_size}") + if not 0 <= ep_rank < ep_size: + raise ValueError(f"ep_rank must be in [0, {ep_size}), got {ep_rank}") + if timeout_s <= 0: + raise ValueError(f"timeout_s must be > 0, got {timeout_s}") + if poll_interval_s <= 0: + raise ValueError(f"poll_interval_s must be > 0, got {poll_interval_s}") + + self._ep_size = int(ep_size) + self._ep_rank = int(ep_rank) + self._completion_reader = completion_reader + self._timeout_s = float(timeout_s) + self._poll_interval_s = float(poll_interval_s) + self._health = health + self._on_timeout = on_timeout + + self._cv = threading.Condition() + self._queue: deque[_CollectiveWatch] = deque() + self._closed = False + self._stopping = False + self._thread: threading.Thread | None = None + self._last_error: BaseException | None = None + + @classmethod + def from_workspace( + cls, + *, + workspace: torch.Tensor, + metainfo: torch.Tensor, + metainfo_index: Mapping[str, int], + ep_rank: int, + ep_size: int, + timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + health: EPGroupHealthLike | None = None, + on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None = None, + ) -> "AlltoAllWatchdog": + """Build a watchdog from the MoE AlltoAll workspace and metainfo.""" + dispatch_offset = int( + metainfo[metainfo_index["DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX"]].item() + ) + combine_offset = int( + metainfo[metainfo_index["COMBINE_COMPLETION_FLAGS_OFFSET_INDEX"]].item() + ) + reader = _TorchCompletionFlagReader( + workspace=workspace, + ep_rank=ep_rank, + ep_size=ep_size, + dispatch_completion_flags_offset=dispatch_offset, + combine_completion_flags_offset=combine_offset, + device_copy_timeout_s=poll_interval_s, + ) + return cls( + ep_size=ep_size, + ep_rank=ep_rank, + completion_reader=reader, + timeout_s=timeout_s, + poll_interval_s=poll_interval_s, + health=health, + on_timeout=on_timeout, + ) + + @property + def last_error(self) -> BaseException | None: + """Return the last polling-thread error, if any.""" + with self._cv: + return self._last_error + + def start(self) -> None: + """Start the background polling thread. Idempotent.""" + with self._cv: + if self._closed: + raise RuntimeError("cannot start a stopped AlltoAllWatchdog") + if self._thread is not None and self._thread.is_alive(): + return + self._stopping = False + self._thread = threading.Thread( + target=self._run, + name=f"AlltoAllWatchdog-rank{self._ep_rank}", + daemon=True, + ) + self._thread.start() + + def stop(self, timeout_s: float | None = None) -> None: + """Stop the polling thread and wait for it to exit.""" + with self._cv: + self._closed = True + self._stopping = True + self._queue.clear() + self._cv.notify_all() + thread = self._thread + if thread is not None and thread is not threading.current_thread(): + thread.join(timeout=timeout_s) + + def watch( + self, + *, + phase: str, + expected_flag: int, + active_mask: int | None = None, + ) -> None: + """Queue a just-launched AlltoAll phase for watchdog polling.""" + if phase not in self.VALID_PHASES: + raise ValueError(f"phase must be one of {sorted(self.VALID_PHASES)}, got {phase!r}") + if not 0 <= expected_flag <= _COMPLETION_FLAG_MASK: + raise ValueError( + f"expected_flag must be in [0, {_COMPLETION_FLAG_MASK}], got {expected_flag}" + ) + if active_mask is None: + if self._health is not None: + active_mask = self._health.get_mask() + else: + active_mask = (1 << self._ep_size) - 1 + if not (active_mask >> self._ep_rank) & 1: + raise ValueError("active_mask must include the local ep_rank") + + self.start() + with self._cv: + if self._closed: + raise RuntimeError("cannot queue a stopped AlltoAllWatchdog") + self._queue.append( + _CollectiveWatch( + phase=phase, + expected_flag=int(expected_flag), + active_mask=int(active_mask), + start_s=time.monotonic(), + ) + ) + self._cv.notify_all() + + def wait_until_idle(self, timeout_s: float) -> bool: + """Wait until all queued phases complete or timeout handling clears them.""" + deadline = time.monotonic() + timeout_s + with self._cv: + while self._queue: + remaining = deadline - time.monotonic() + if remaining <= 0: + return False + self._cv.wait(timeout=remaining) + return True + + def __enter__(self) -> "AlltoAllWatchdog": + self.start() + return self + + def __exit__(self, exc_type, exc, tb) -> None: + self.stop(timeout_s=1.0) + + def _active_ranks(self, active_mask: int) -> tuple[int, ...]: + return tuple(rank for rank in range(self._ep_size) if (active_mask >> rank) & 1) + + def _phase_complete(self, watch: _CollectiveWatch, observed_flags: tuple[int, ...]) -> bool: + return all( + _completion_flag_reached(observed_flags[rank], watch.expected_flag) + for rank in self._active_ranks(watch.active_mask) + ) + + def _missing_ranks( + self, watch: _CollectiveWatch, observed_flags: tuple[int, ...] + ) -> tuple[int, ...]: + return tuple( + rank + for rank in self._active_ranks(watch.active_mask) + if not _completion_flag_reached(observed_flags[rank], watch.expected_flag) + ) + + def _handle_timeout( + self, + watch: _CollectiveWatch, + observed_flags: tuple[int, ...], + *, + poll_timed_out: bool = False, + ) -> None: + elapsed_s = time.monotonic() - watch.start_s + missing_ranks = self._missing_ranks(watch, observed_flags) + event = AlltoAllWatchdogTimeout( + phase=watch.phase, + expected_flag=watch.expected_flag, + observed_flags=observed_flags, + missing_ranks=missing_ranks, + elapsed_s=elapsed_s, + poll_timed_out=poll_timed_out, + ) + if poll_timed_out: + tllm_logger.error( + "AlltoAll watchdog could not read completion flags on rank %d " + "during %s before timeout %.3fs; expected flag %d, active " + "ranks %s, observed flags %s; reporting detection event only", + self._ep_rank, + watch.phase, + elapsed_s, + watch.expected_flag, + list(self._active_ranks(watch.active_mask)), + list(observed_flags), + ) + else: + tllm_logger.warning( + "AlltoAll watchdog timeout on rank %d during %s: expected flag %d, " + "missing ranks %s, observed flags %s", + self._ep_rank, + watch.phase, + watch.expected_flag, + list(missing_ranks), + list(observed_flags), + ) + if self._on_timeout is not None: + self._on_timeout(event) + + def _stop_after_error(self, exc: Exception) -> None: + with self._cv: + self._last_error = exc + self._closed = True + self._stopping = True + self._queue.clear() + self._cv.notify_all() + + def _run(self) -> None: + last_observed_flags = tuple(UNKNOWN_COMPLETION_FLAG for _ in range(self._ep_size)) + poll_timed_out = False + while True: + with self._cv: + while not self._queue and not self._stopping: + self._cv.wait() + if self._stopping: + return + watch = self._queue[0] + + try: + observed_flags = tuple( + _normalize_completion_flag(int(v)) + for v in self._completion_reader.read_completion_flags(watch.phase) + ) + if len(observed_flags) != self._ep_size: + raise RuntimeError( + f"completion reader returned {len(observed_flags)} flags; " + f"expected ep_size={self._ep_size}" + ) + last_observed_flags = observed_flags + poll_timed_out = False + except CompletionFlagReadTimeout: + observed_flags = last_observed_flags + poll_timed_out = True + except Exception as exc: # noqa: BLE001 - keep watchdog failures visible. + self._stop_after_error(exc) + tllm_logger.error("AlltoAll watchdog stopped after polling error: %s", exc) + return + + if self._phase_complete(watch, observed_flags): + with self._cv: + if self._queue and self._queue[0] is watch: + self._queue.popleft() + self._cv.notify_all() + last_observed_flags = tuple(UNKNOWN_COMPLETION_FLAG for _ in range(self._ep_size)) + poll_timed_out = False + continue + + if time.monotonic() - watch.start_s >= self._timeout_s: + try: + self._handle_timeout(watch, observed_flags, poll_timed_out=poll_timed_out) + except Exception as exc: # noqa: BLE001 - keep watchdog failures visible. + self._stop_after_error(exc) + tllm_logger.error( + "AlltoAll watchdog stopped after timeout handling error: %s", exc + ) + return + with self._cv: + # The GPU stream is no longer trustworthy once a collective + # times out. Drop queued follow-on phases so they do not + # produce duplicate or misleading reports. + self._queue.clear() + self._cv.notify_all() + last_observed_flags = tuple(UNKNOWN_COMPLETION_FLAG for _ in range(self._ep_size)) + poll_timed_out = False + continue + + with self._cv: + self._cv.wait(timeout=self._poll_interval_s) diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index ca3a50dcfd12..2bf31ef680a4 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -8,12 +8,18 @@ # ruff: noqa: E501 import os +import sys from dataclasses import dataclass -from typing import Dict, Optional +from typing import Callable, Dict, Optional import torch from tensorrt_llm._mnnvl_utils import MnnvlMemory +from tensorrt_llm._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, ActiveRankMaskSnapshot, + AlltoAllWatchdog, AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, + EPGroupHealthLike) from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -26,6 +32,7 @@ class _A2AState: local_num_tokens: int | None = None combine_payload_offset: int | None = None eplb_gathered_stats: torch.Tensor | None = None + active_rank_mask_snapshot: ActiveRankMaskSnapshot | None = None class MoeAlltoAll: @@ -126,7 +133,13 @@ def __init__( num_slots: int, workspace_size_per_rank: int, num_experts: Optional[int] = None, - ): + ep_group_health: Optional[EPGroupHealthLike] = None, + alltoall_watchdog_timeout_s: Optional[float] = None, + alltoall_watchdog_poll_interval_s: + float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + alltoall_watchdog_on_timeout: Optional[Callable[ + [AlltoAllWatchdogTimeout], None]] = None, + ) -> None: """ Initialize MoeAlltoAll with workspace allocation. @@ -138,6 +151,12 @@ def __init__( Note: The terminology is mapped to `num_experts` in this class and the kernels. num_experts: (Optional) Number of experts for EPLB stats (must be <= num_slots). DO NOT provide this parameter if EPLB is not enabled. Note: The terminology is mapped to `eplb_stats_num_experts` in this class and the kernels. + ep_group_health: Optional read-only committed EP membership. When present, its mask is passed to the + CUDA kernels and defines the peers expected by the watchdog. Timeout detection never mutates it. + alltoall_watchdog_timeout_s: Optional timeout for the host-side AlltoAll watchdog. If None, the + watchdog is disabled. + alltoall_watchdog_poll_interval_s: Poll interval for the watchdog thread. + alltoall_watchdog_on_timeout: Optional callback invoked when the watchdog reports suspects. """ # Check for environment variable override workspace_mb_env = os.environ.get("TRTLLM_MOE_A2A_WORKSPACE_MB") @@ -214,6 +233,45 @@ def __init__( self.metainfo = self._WORKSPACE["metainfo"] # Internal state self._state: _A2AState = _A2AState() + self.ep_group_health = ep_group_health + workspace_state = self._WORKSPACE + assert workspace_state is not None + metainfo_index = self._METAINFO_INDEX + assert metainfo_index is not None + self._watchdog_coordinator = AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=self.workspace, + metainfo=self.metainfo, + metainfo_index=metainfo_index, + ep_rank=self.ep_rank, + health=self.ep_group_health, + ) + self._destroyed = False + self._alltoall_watchdog: AlltoAllWatchdog | None = None + if (alltoall_watchdog_timeout_s is None + and self.ep_group_health is not None): + alltoall_watchdog_timeout_s = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S + if alltoall_watchdog_timeout_s is not None: + self._alltoall_watchdog = self._watchdog_coordinator.acquire_watchdog( + ep_size=self.ep_size, + timeout_s=alltoall_watchdog_timeout_s, + poll_interval_s=alltoall_watchdog_poll_interval_s, + on_timeout=alltoall_watchdog_on_timeout, + ) + + def destroy(self) -> None: + """Stop background watchdog resources owned by this wrapper.""" + if getattr(self, "_destroyed", False): + return + self._destroyed = True + watchdog = getattr(self, "_alltoall_watchdog", None) + if watchdog is not None: + self._watchdog_coordinator.release_watchdog(watchdog) + self._alltoall_watchdog = None + + def __del__(self) -> None: + if not sys.is_finalizing(): + self.destroy() def dispatch(self, token_selected_experts: torch.Tensor, @@ -221,7 +279,8 @@ def dispatch(self, runtime_max_tokens_per_rank: int, invalid_token_expert_id: Optional[int] = None, expert_id_payload_index: Optional[int] = None, - eplb_local_stats: Optional[torch.Tensor] = None): + eplb_local_stats: Optional[torch.Tensor] = None, + active_rank_mask: Optional[torch.Tensor] = None): """ Perform MoE all-to-all dispatch operation. @@ -232,6 +291,9 @@ def dispatch(self, invalid_token_expert_id: If not None, set the token_selected_experts of the invalid tokens to this expert id. This is used to notify the MoE to skip these tokens for GroupGEMM. expert_id_payload_index: The index of token_selected_experts in the input_payloads. Must be provided if invalid_token_expert_id is not None. eplb_local_stats: (Optional) [num_experts] tensor containing local statistics for EPLB + active_rank_mask: Optional uint64 CPU tensor overriding committed membership for this dispatch. When + omitted, the committed mask and generation are captured together. Combine reuses that mask and + fails closed if the committed generation changes first. Returns: recv_tensors: List of tensors received, each has shape [ep_size, max_tokens_per_rank, payload_num_elements_per_token] @@ -246,6 +308,9 @@ def dispatch(self, 0 ) == self.eplb_stats_num_experts, "eplb_local_stats size must match eplb_stats_num_experts" + active_rank_mask_snapshot = self._watchdog_coordinator.capture_active_rank_mask( + active_rank_mask) + active_rank_mask = active_rank_mask_snapshot.active_rank_mask recv_tensors, combine_payload_offset, eplb_gathered_stats = torch.ops.trtllm.moe_a2a_dispatch( token_selected_experts, input_payloads, @@ -257,7 +322,11 @@ def dispatch(self, self.top_k, self.num_experts, eplb_local_stats, + active_rank_mask, ) + self._watchdog_coordinator.watch_collective(self._alltoall_watchdog, + "dispatch", + active_rank_mask) if eplb_gathered_stats.numel() == 0: eplb_gathered_stats = None @@ -265,6 +334,7 @@ def dispatch(self, self._state.local_num_tokens = token_selected_experts.size(0) self._state.combine_payload_offset = combine_payload_offset self._state.eplb_gathered_stats = eplb_gathered_stats + self._state.active_rank_mask_snapshot = active_rank_mask_snapshot self._state.phase = "dispatched" if invalid_token_expert_id is not None: @@ -287,6 +357,7 @@ def combine( runtime_max_tokens_per_rank: int, payload_in_workspace: bool = False, use_low_precision_combine: bool = False, + active_rank_mask: Optional[torch.Tensor] = None, ): """ Perform MoE all-to-all combine operation. @@ -296,6 +367,8 @@ def combine( runtime_max_tokens_per_rank: Maximum of the number of tokens of each DP rank's local batch. payload_in_workspace: If True, 'payload' is a view into 'workspace' at 'combine_payload_offset' and no staging copy is needed. If False, the op stages 'payload' into the workspace region before combining. use_low_precision_combine: If True, quantize the combine payload to FP8 for NVLink transfer (halves NVLink bandwidth usage, output precision is preserved). + active_rank_mask: Optional uint64 CPU tensor. If supplied, it must match the mask captured by dispatch + for this collective. A committed-generation change since dispatch aborts the collective epoch. Returns: combined_output: [local_num_tokens, num_elements_per_token] tensor of combined results @@ -303,11 +376,17 @@ def combine( assert self._state.phase == "dispatched", "combine called before a successful dispatch" assert runtime_max_tokens_per_rank <= self.max_num_tokens, "runtime_max_tokens_per_rank must not exceed max_num_tokens" + active_rank_mask_snapshot = self._state.active_rank_mask_snapshot + assert active_rank_mask_snapshot is not None + active_rank_mask = self._watchdog_coordinator.active_rank_mask_for_combine( + active_rank_mask_snapshot, active_rank_mask) output = torch.ops.trtllm.moe_a2a_combine( payload, self._state.local_num_tokens, self.workspace, self.metainfo, runtime_max_tokens_per_rank, self.ep_rank, self.ep_size, self.top_k, self._state.combine_payload_offset, - payload_in_workspace, use_low_precision_combine) + payload_in_workspace, use_low_precision_combine, active_rank_mask) + self._watchdog_coordinator.watch_collective(self._alltoall_watchdog, + "combine", active_rank_mask) # Reset state for next round self.reset_state() diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py index a4ec2ceefe44..17a67deee219 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py @@ -28,6 +28,7 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm.logger import logger +from ..wide_ep_ft import get_wide_ep_ft_options from .allgather_reducescatter import AllGatherReduceScatter from .base import Communication from .deep_ep import DeepEP @@ -133,6 +134,9 @@ def create_strategy( try: enable_eplb = model_config.moe_load_balancer is not None + ep_group_health, watchdog_timeout_s, watchdog_poll_interval_s = get_wide_ep_ft_options( + model_config + ) strategy = NVLinkOneSided( mapping, num_slots, @@ -143,6 +147,9 @@ def create_strategy( dtype=act_dtype, num_experts=num_experts if enable_eplb else None, use_low_precision_combine=use_low_precision_combine, + ep_group_health=ep_group_health, + alltoall_watchdog_timeout_s=watchdog_timeout_s, + alltoall_watchdog_poll_interval_s=watchdog_poll_interval_s, ) logger.info("Selected communication strategy: NVLinkOneSided") return strategy @@ -285,6 +292,9 @@ def _create_forced_method( ) elif method in ["NVLINK_ONE_SIDED"]: enable_eplb = model_config.moe_load_balancer is not None + ep_group_health, watchdog_timeout_s, watchdog_poll_interval_s = get_wide_ep_ft_options( + model_config + ) return NVLinkOneSided( mapping, num_slots, @@ -295,6 +305,9 @@ def _create_forced_method( dtype=act_dtype, num_experts=num_experts if enable_eplb else None, use_low_precision_combine=use_low_precision_combine, + ep_group_health=ep_group_health, + alltoall_watchdog_timeout_s=watchdog_timeout_s, + alltoall_watchdog_poll_interval_s=watchdog_poll_interval_s, ) elif method == "DEEPEP": return DeepEP( diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py index f9068f71d717..8f0b3c9486aa 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py @@ -25,11 +25,20 @@ """ import os -from typing import Dict, List, Optional, Tuple +from typing import Callable, Dict, List, Optional, Tuple import torch from tensorrt_llm._mnnvl_utils import MnnvlMemory +from tensorrt_llm._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + ActiveRankMaskSnapshot, + AlltoAllWatchdog, + AlltoAllWatchdogCoordinator, + AlltoAllWatchdogTimeout, + EPGroupHealthLike, +) from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -151,7 +160,11 @@ def __init__( dtype: Optional[torch.dtype] = None, num_experts: Optional[int] = None, use_low_precision_combine: bool = False, - ): + ep_group_health: EPGroupHealthLike | None = None, + alltoall_watchdog_timeout_s: Optional[float] = None, + alltoall_watchdog_poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + alltoall_watchdog_on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, + ) -> None: """ Initialize NVLinkOneSided with workspace allocation. @@ -169,6 +182,12 @@ def __init__( use_low_precision_combine: If True, quantize the combine payload to FP8 for NVLink transfer (halves NVLink bandwidth usage, output precision is preserved). Corresponds to model_config.use_low_precision_moe_combine. + ep_group_health: Optional read-only committed EP membership. When present, its mask is passed to + the CUDA kernels and defines the peers expected by the watchdog. Timeout detection never mutates it. + alltoall_watchdog_timeout_s: Optional timeout for the host-side AlltoAll watchdog. If None, the + watchdog is disabled. + alltoall_watchdog_poll_interval_s: Poll interval for the watchdog thread. + alltoall_watchdog_on_timeout: Optional callback invoked when the watchdog reports suspects. """ super().__init__(mapping) @@ -296,10 +315,34 @@ def __init__( NVLinkOneSided._WORKSPACE_REFCOUNTS.get(self._workspace_key, 0) + 1 ) self._destroyed = False + self._workspace_state = workspace_state self.mnnvl_mem = workspace_state["mnnvl_mem"] self.workspace = workspace_state["workspace"] self.moe_a2a_metainfo = workspace_state["metainfo"] self.max_num_tokens_per_rank = workspace_state["max_num_tokens_per_rank"] + self.ep_group_health = ep_group_health + self._watchdog_coordinator = AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=self.workspace, + metainfo=self.moe_a2a_metainfo, + metainfo_index={ + "FLAG_VAL_OFFSET_INDEX": self.FLAG_VAL_OFFSET_INDEX, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": self.DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": self.COMBINE_COMPLETION_FLAGS_OFFSET_INDEX, + }, + ep_rank=self.ep_rank, + health=self.ep_group_health, + ) + self._alltoall_watchdog: AlltoAllWatchdog | None = None + if alltoall_watchdog_timeout_s is None and self.ep_group_health is not None: + alltoall_watchdog_timeout_s = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S + if alltoall_watchdog_timeout_s is not None: + self._alltoall_watchdog = self._watchdog_coordinator.acquire_watchdog( + ep_size=self.ep_size, + timeout_s=alltoall_watchdog_timeout_s, + poll_interval_s=alltoall_watchdog_poll_interval_s, + on_timeout=alltoall_watchdog_on_timeout, + ) # Initialize dispatch state self._dispatch_state = {"phase": "idle"} @@ -326,6 +369,9 @@ def destroy(self): return self._destroyed = True + if self._alltoall_watchdog is not None: + self._watchdog_coordinator.release_watchdog(self._alltoall_watchdog) + self._alltoall_watchdog = None workspace_key = getattr(self, "_workspace_key", None) if workspace_key is None: return @@ -347,6 +393,7 @@ def destroy(self): self.mnnvl_mem = None self.workspace = None self.moe_a2a_metainfo = None + self._workspace_state = None self._dispatch_state = {"phase": "destroyed"} def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) -> bool: @@ -378,7 +425,9 @@ def dispatch( token_final_scales: Router weights [local_num_tokens, top_k] all_rank_num_tokens: Token counts per rank [ep_size] use_dp_padding: Whether to use DP padding (optional) - **kwargs: Strategy-specific arguments (unused) + **kwargs: Strategy-specific arguments. ``active_rank_mask`` may override the committed membership + for dispatch. Without an override, the committed mask and generation are captured together; + combine reuses that mask and fails closed if the generation changes first. Returns: Tuple of (hidden_states, hidden_states_sf, token_selected_slots, token_final_scales) @@ -409,6 +458,10 @@ def dispatch( assert eplb_local_stats.size(0) == self.eplb_stats_num_experts, ( "eplb_local_stats size must match eplb_stats_num_experts" ) + active_rank_mask_snapshot = self._watchdog_coordinator.capture_active_rank_mask( + kwargs.get("active_rank_mask") + ) + active_rank_mask = active_rank_mask_snapshot.active_rank_mask recv_buffers, combine_payload_offset, eplb_gathered_stats = ( torch.ops.trtllm.moe_a2a_dispatch( @@ -422,14 +475,19 @@ def dispatch( self.top_k, self.num_experts, eplb_local_stats, + active_rank_mask, ) ) + self._watchdog_coordinator.watch_collective( + self._alltoall_watchdog, "dispatch", active_rank_mask + ) if eplb_gathered_stats.numel() == 0: eplb_gathered_stats = None self._dispatch_state["eplb_gathered_stats"] = eplb_gathered_stats self._dispatch_state["combine_payload_offset"] = int(combine_payload_offset) self._dispatch_state["local_num_tokens"] = token_selected_slots.size(0) self._dispatch_state["runtime_max_tokens_per_rank"] = runtime_max_tokens_per_rank + self._dispatch_state["active_rank_mask_snapshot"] = active_rank_mask_snapshot self._dispatch_state["phase"] = "dispatched" # Extract results from recv_buffers @@ -492,6 +550,9 @@ def combine( final_hidden_states: Output from MoE computation Shape: [ep_size, max_tokens_per_rank, hidden_size] or [ep_size * max_tokens_per_rank, hidden_size] (will be reshaped) + **kwargs: Strategy-specific arguments. If ``active_rank_mask`` is supplied, it must match the mask + captured by dispatch for this collective. A committed-generation change since dispatch aborts + the collective epoch. Returns: Combined output tensor [local_num_tokens, hidden_size] @@ -526,6 +587,13 @@ def combine( raise ValueError( f"final_hidden_states must be 2D or 3D, got {final_hidden_states.dim()}D" ) + active_rank_mask_snapshot = self._dispatch_state.get("active_rank_mask_snapshot") + if not isinstance(active_rank_mask_snapshot, ActiveRankMaskSnapshot): + raise RuntimeError("combine called but dispatch rank-mask snapshot is missing") + active_rank_mask = self._watchdog_coordinator.active_rank_mask_for_combine( + active_rank_mask_snapshot, + kwargs.get("active_rank_mask"), + ) output = torch.ops.trtllm.moe_a2a_combine( final_hidden_states, int(local_num_tokens), @@ -538,6 +606,10 @@ def combine( int(combine_payload_offset), bool(self.payload_in_workspace), bool(self.use_low_precision_combine), + active_rank_mask, + ) + self._watchdog_coordinator.watch_collective( + self._alltoall_watchdog, "combine", active_rank_mask ) # Reset state for next round diff --git a/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py b/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py index 1864d4ba825d..dd4bcd14136b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py +++ b/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py @@ -13,21 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""EP group health tracking for WideEP fault tolerance. +"""Committed EP data-plane membership for WideEP fault tolerance. This module provides :class:`EPGroupHealth`, a process-local, thread-safe data -structure that records which ranks in an Expert Parallel (EP) group are currently -alive vs. failed. It is the single source of truth for EP rank health within one -process and is consumed by: +structure that records which Expert Parallel (EP) ranks the recovery coordinator +has committed as included vs. excluded from the data plane. It is consumed by: * AlltoAll communication backends (rank masking on dispatch / combine) - * The host-side AlltoAll watchdog (failure detection) + * The host-side AlltoAll watchdog (read-only expected-peer snapshot) * The MoE load balancer (emergency-mask reconfiguration) * The model engine and PyExecutor (degraded health reporting) -Cross-process consensus on which ranks are dead (failure broadcast across the EP -group) is the responsibility of higher-layer coordination components and is not -performed here. +Detected or suspected physical liveness is separate evidence. Higher-layer +coordination reconciles that evidence and commits membership; detectors and +telemetry consumers must not mutate this object to drive recovery. """ import threading diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py index 3a89ec2fb35d..0b619bb31916 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -20,6 +20,7 @@ Fp4QuantizedTensor) from .interface import AlltoallMethodType, MoE from .quantization import UnquantizedFusedMoEMethod +from .wide_ep_ft import get_wide_ep_ft_options # isort: off from .quantization import ( @@ -331,6 +332,8 @@ def __init__( dtype, self.num_experts if self.layer_load_balancer else None, ) + ep_group_health, watchdog_timeout_s, watchdog_poll_interval_s = ( + get_wide_ep_ft_options(model_config)) self.moe_a2a = MoeAlltoAll( mapping=self.mapping, @@ -340,6 +343,10 @@ def __init__( workspace_size_per_rank=workspace_size, num_experts=self.num_experts if self.layer_load_balancer else None, + ep_group_health=ep_group_health, + alltoall_watchdog_timeout_s=watchdog_timeout_s, + alltoall_watchdog_poll_interval_s= + watchdog_poll_interval_s, ) elif self.alltoall_method_type == AlltoallMethodType.DeepEP or self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: raise NotImplementedError( diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 820d92a95e4a..6f59870c374b 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -36,6 +36,7 @@ from ...utils import ActivationType, AuxStreamType, Fp4QuantizedTensor from .interface import AlltoallMethodType, MoE, MoEWeightLoadingMode from .moe_op_backend import MoEOpBackend, get_op_backend +from .wide_ep_ft import get_wide_ep_ft_options # isort: off from .quantization import ( @@ -291,6 +292,8 @@ def __init__( ep_size, self.routing_method.experts_per_token, max_num_tokens, hidden_size, dtype, self.num_experts if self.layer_load_balancer else None) + ep_group_health, watchdog_timeout_s, watchdog_poll_interval_s = ( + get_wide_ep_ft_options(model_config)) self.moe_a2a = MoeAlltoAll( mapping=self.mapping, @@ -299,7 +302,11 @@ def __init__( num_slots=self.num_slots, workspace_size_per_rank=workspace_size, num_experts=self.num_experts - if self.layer_load_balancer else None) + if self.layer_load_balancer else None, + ep_group_health=ep_group_health, + alltoall_watchdog_timeout_s=watchdog_timeout_s, + alltoall_watchdog_poll_interval_s= + watchdog_poll_interval_s) elif self.alltoall_method_type == AlltoallMethodType.DeepEP or self.alltoall_method_type == AlltoallMethodType.DeepEPLowLatency: raise NotImplementedError( "DeepEP and DeepEPLowLatency are not supported for TRTLLMGenFusedMoE yet" diff --git a/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py b/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py new file mode 100644 index 000000000000..fb1586ae1054 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py @@ -0,0 +1,87 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Shared WideEP fault-tolerance options for MoE communication paths.""" + +from __future__ import annotations + +import os +from typing import Any, Optional + +from tensorrt_llm._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, +) + +from .ep_group_health import EPGroupHealth + +_ENABLE_ENV = "TLLM_FAULT_TOLERANCE_MODE" +_TIMEOUT_ENV = "TRTLLM_ALLTOALL_WATCHDOG_TIMEOUT_S" +_POLL_INTERVAL_ENV = "TRTLLM_ALLTOALL_WATCHDOG_POLL_INTERVAL_S" + +# This object contains membership committed by higher-layer recovery +# coordination. Detection threads must treat it as read-only. +_HEALTH_KEY = "wide_ep_ft_ep_group_health" +_TIMEOUT_KEY = "alltoall_watchdog_timeout_s" +_POLL_INTERVAL_KEY = "alltoall_watchdog_poll_interval_s" + + +def _env_enabled() -> bool: + return os.environ.get(_ENABLE_ENV) == "1" + + +def _float_option(extra_attrs: dict, key: str, env_name: str, default: float) -> float: + if key in extra_attrs: + return float(extra_attrs[key]) + if env_name in os.environ: + return float(os.environ[env_name]) + return default + + +def get_wide_ep_ft_options( + model_config: Any, +) -> tuple[Optional[EPGroupHealth], Optional[float], float]: + """Return committed EP membership and watchdog timing for a model. + + WideEP FT remains opt-in until the integration PR wires a public model + option. Callers can either inject ``wide_ep_ft_ep_group_health`` through + ``ModelConfig.extra_attrs`` or set ``TLLM_FAULT_TOLERANCE_MODE=1`` to create + one process-local membership object shared by all MoE communication layers. + The AlltoAll watchdog reads this object to determine expected peers and + reports suspects through its ``on_timeout`` seam; it never mutates the + committed membership directly. + """ + + extra_attrs = getattr(model_config, "extra_attrs", {}) + health = extra_attrs.get(_HEALTH_KEY) or extra_attrs.get("ep_group_health") + if health is None and _env_enabled(): + health = EPGroupHealth(model_config.mapping.moe_ep_size) + extra_attrs[_HEALTH_KEY] = health + + poll_interval_s = _float_option( + extra_attrs, + _POLL_INTERVAL_KEY, + _POLL_INTERVAL_ENV, + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + ) + if health is None: + return None, None, poll_interval_s + + timeout_s = _float_option( + extra_attrs, + _TIMEOUT_KEY, + _TIMEOUT_ENV, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + ) + return health, timeout_s, poll_interval_s diff --git a/tests/unittest/_torch/modules/moe/test_moe_comm.py b/tests/unittest/_torch/modules/moe/test_moe_comm.py index 057f87a95f7a..ac75c895cf6d 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_comm.py +++ b/tests/unittest/_torch/modules/moe/test_moe_comm.py @@ -225,14 +225,31 @@ def _read_nvlink_topk_target_ranks( return raw.view(torch.int32).view(max_num_tokens, top_k).cpu() -def _run_nvlink_rank_mask_dispatch_combine( +def _read_nvlink_topk_send_indices( + comm: NVLinkOneSided, + max_num_tokens: int, + top_k: int, +) -> torch.Tensor: + """Read topk_send_indices[max_num_tokens, top_k] from NVLinkOneSided workspace.""" + from tensorrt_llm.bindings import internal as _tllm_internal + + offset_index = int(_tllm_internal.thop.MOE_A2A_TOPK_SEND_INDICES_OFFSET_INDEX) + offset = comm.moe_a2a_metainfo[offset_index].item() + raw = comm.workspace[ + comm.ep_rank, + offset : offset + max_num_tokens * top_k * 4, + ] + return raw.view(torch.int32).view(max_num_tokens, top_k).cpu() + + +def _run_nvlink_rank_mask_dispatch( comm: NVLinkOneSided, token_selected_experts: torch.Tensor, payload: torch.Tensor, runtime_max_tokens_per_rank: int, active_rank_mask: Optional[torch.Tensor], -) -> Tuple[torch.Tensor, torch.Tensor]: - """Run raw NVLink one-sided dispatch/combine with an optional active rank mask.""" +) -> Tuple[List[torch.Tensor], int, torch.Tensor, torch.Tensor]: + """Run raw NVLink one-sided dispatch with an optional active rank mask.""" recv_tensors, combine_payload_offset, _ = torch.ops.trtllm.moe_a2a_dispatch( token_selected_experts, [payload], @@ -252,24 +269,105 @@ def _run_nvlink_rank_mask_dispatch_combine( runtime_max_tokens_per_rank, comm.top_k, ) + topk_send_indices = _read_nvlink_topk_send_indices( + comm, + runtime_max_tokens_per_rank, + comm.top_k, + ) + return recv_tensors, int(combine_payload_offset), topk_target_ranks, topk_send_indices - combined = torch.ops.trtllm.moe_a2a_combine( - recv_tensors[0], - token_selected_experts.size(0), + +def _run_nvlink_rank_mask_combine( + comm: NVLinkOneSided, + combine_payload: torch.Tensor, + local_num_tokens: int, + runtime_max_tokens_per_rank: int, + combine_payload_offset: int, + active_rank_mask: Optional[torch.Tensor], +) -> torch.Tensor: + """Run raw NVLink one-sided combine with an optional active rank mask.""" + return torch.ops.trtllm.moe_a2a_combine( + combine_payload, + local_num_tokens, comm.workspace, comm.moe_a2a_metainfo, runtime_max_tokens_per_rank, comm.ep_rank, comm.ep_size, comm.top_k, - int(combine_payload_offset), + combine_payload_offset, False, # payload_in_workspace False, # use_low_precision active_rank_mask, ) + + +def _run_nvlink_rank_mask_dispatch_combine( + comm: NVLinkOneSided, + token_selected_experts: torch.Tensor, + payload: torch.Tensor, + runtime_max_tokens_per_rank: int, + active_rank_mask: Optional[torch.Tensor], +) -> Tuple[torch.Tensor, torch.Tensor]: + """Run raw NVLink one-sided dispatch/combine with an optional active rank mask.""" + recv_tensors, combine_payload_offset, topk_target_ranks, _ = _run_nvlink_rank_mask_dispatch( + comm, + token_selected_experts, + payload, + runtime_max_tokens_per_rank, + active_rank_mask, + ) + combined = _run_nvlink_rank_mask_combine( + comm, + recv_tensors[0], + token_selected_experts.size(0), + runtime_max_tokens_per_rank, + combine_payload_offset, + active_rank_mask, + ) return combined.cpu(), topk_target_ranks +def _expected_nvlink_rank_mask_combine_output( + comm: NVLinkOneSided, + payload: torch.Tensor, + topk_target_ranks: torch.Tensor, + topk_send_indices: torch.Tensor, + local_num_tokens: int, + runtime_max_tokens_per_rank: int, + dead_ranks: Set[int], +) -> torch.Tensor: + """Compute combine output from dispatched workspace while skipping dead ranks.""" + from tensorrt_llm.bindings import internal as _tllm_internal + + hidden_size = payload.shape[-1] + expected = torch.zeros( + (local_num_tokens, hidden_size), + dtype=payload.dtype, + device=payload.device, + ) + payload_offset_index = int(_tllm_internal.thop.MOE_A2A_PAYLOAD_DATA_OFFSET_INDEX) + payload_offset = comm.moe_a2a_metainfo[payload_offset_index].item() + bytes_per_rank = ( + comm.ep_size * runtime_max_tokens_per_rank * hidden_size * payload.element_size() + ) + + for token_idx in range(local_num_tokens): + for k in range(comm.top_k): + target_rank = int(topk_target_ranks[token_idx, k].item()) + dst_idx = int(topk_send_indices[token_idx, k].item()) + if dst_idx < 0 or target_rank in dead_ranks: + continue + raw = comm.workspace[target_rank, payload_offset : payload_offset + bytes_per_rank] + recv_payload = raw.view(payload.dtype).view( + comm.ep_size, + runtime_max_tokens_per_rank, + hidden_size, + ) + expected[token_idx] += recv_payload[comm.ep_rank, dst_idx] + return expected.cpu() + + # ============================================================================ # Source Encoding Utilities # ============================================================================ @@ -992,6 +1090,7 @@ def _worker_rank_mask_all_active_matches_no_mask(config: CommTestConfig) -> dict ) return { + "rank": rank, "output_eq": torch.equal(out_no_mask, out_all_active), "topk_eq": torch.equal(topk_no_mask, topk_all_active), } @@ -1040,7 +1139,7 @@ def _worker_rank_mask_one_rank_masked( if rank == dead_rank: MPI.COMM_WORLD.barrier() - return {"status": "dead"} + return {"rank": rank, "status": "dead"} local_num_tokens = config.all_num_tokens[rank] torch.manual_seed(0xA2A + rank) @@ -1069,6 +1168,7 @@ def _worker_rank_mask_one_rank_masked( MPI.COMM_WORLD.barrier() return { + "rank": rank, "status": "alive", "combined": combined, "topk_target_ranks": topk_target_ranks, @@ -1082,6 +1182,83 @@ def _worker_rank_mask_one_rank_masked( comm.destroy() +def _worker_rank_mask_inactive_before_combine( + config: CommTestConfig, + dead_rank: int, +) -> dict: + """Dispatch with all ranks active, then omit one rank from combine's active mask.""" + rank = tllm.mpi_rank() + torch.cuda.set_device(rank) + + comm = None + try: + mapping = Mapping( + rank=rank, + tp_size=config.ep_size, + moe_ep_size=config.ep_size, + world_size=config.ep_size, + ) + comm = create_comm_object(config.comm_type, mapping, config) + + local_num_tokens = config.all_num_tokens[rank] + torch.manual_seed(0xA2A + rank) + token_selected_experts = torch.randint( + 0, + config.num_experts, + (local_num_tokens, config.top_k), + dtype=torch.int32, + device="cuda", + ) + payload = _make_rank_mask_payload(local_num_tokens, config.hidden_size, rank) + + recv_tensors, combine_payload_offset, topk_target_ranks, topk_send_indices = ( + _run_nvlink_rank_mask_dispatch( + comm, + token_selected_experts, + payload, + local_num_tokens, + active_rank_mask=_ep_mask_words(config.ep_size, dead_ranks=set()), + ) + ) + + if rank == dead_rank: + MPI.COMM_WORLD.barrier() + return {"rank": rank, "status": "dead"} + + dead_ranks = {dead_rank} + expected = _expected_nvlink_rank_mask_combine_output( + comm, + payload, + topk_target_ranks, + topk_send_indices, + local_num_tokens, + local_num_tokens, + dead_ranks, + ) + combined = _run_nvlink_rank_mask_combine( + comm, + recv_tensors[0], + local_num_tokens, + local_num_tokens, + combine_payload_offset, + active_rank_mask=_ep_mask_words(config.ep_size, dead_ranks=dead_ranks), + ).cpu() + + MPI.COMM_WORLD.barrier() + return { + "rank": rank, + "status": "alive", + "combined": combined, + "expected": expected, + } + except Exception: + traceback.print_exc() + raise + finally: + if comm is not None and hasattr(comm, "destroy"): + comm.destroy() + + # ============================================================================ # Verification Functions # ============================================================================ @@ -1807,7 +1984,7 @@ def _make_postquant_test_params(): @pytest.fixture(autouse=True) -def setup_test(): +def setup_test() -> None: torch.manual_seed(0x1234) tllm.logger.set_level("error") @@ -1886,7 +2063,8 @@ def _run_rank_mask_all_active_test( ) ) - for rank, result in enumerate(results): + for result in results: + rank = result["rank"] assert result["output_eq"], ( f"rank {rank}: combine output differs between no-mask and all-active mask" ) @@ -1915,7 +2093,8 @@ def _run_rank_mask_one_rank_masked_test( ) saw_dead = False - for rank, result in enumerate(results): + for result in results: + rank = result["rank"] if result["status"] == "dead": assert rank == dead_rank saw_dead = True @@ -1952,6 +2131,45 @@ def _run_rank_mask_one_rank_masked_test( assert saw_dead, f"dead rank {dead_rank} did not appear in results" +def _run_rank_mask_inactive_before_combine_test( + mpi_pool_executor, + dead_rank: int, + local_num_tokens: int, + top_k: int, +) -> None: + ep_size = mpi_pool_executor.num_workers + config = _make_rank_mask_config(ep_size, local_num_tokens, top_k) + _skip_if_rank_mask_config_unsupported(config) + assert 0 <= dead_rank < ep_size + + worker_args = [(config, dead_rank)] * config.ep_size + results = list( + mpi_pool_executor.map( + _worker_rank_mask_inactive_before_combine, + *zip(*worker_args), + ) + ) + + saw_dead = False + for result in results: + rank = result["rank"] + if result["status"] == "dead": + assert rank == dead_rank + saw_dead = True + continue + + assert result["status"] == "alive" + combined = result["combined"] + expected = result["expected"] + assert combined is not None + assert expected is not None + assert torch.equal(combined, expected), ( + f"rank {rank}: combine output included a rank masked inactive before combine" + ) + + assert saw_dead, f"dead rank {dead_rank} did not appear in results" + + # ============================================================================ # Test Class # ============================================================================ @@ -2019,7 +2237,7 @@ def test_moe_comm_rank_mask_all_active_matches_no_mask( mpi_pool_executor, local_num_tokens: int, top_k: int, - ): + ) -> None: """Verify all-active active_rank_mask matches omitted mask for NVLinkOneSided.""" _run_rank_mask_all_active_test(mpi_pool_executor, local_num_tokens, top_k) @@ -2039,7 +2257,7 @@ def test_moe_comm_rank_mask_one_rank_masked_completes( dead_rank: int, local_num_tokens: int, top_k: int, - ): + ) -> None: """Verify masked-dead rank is skipped by raw NVLinkOneSided moe_a2a ops.""" _run_rank_mask_one_rank_masked_test( mpi_pool_executor, @@ -2047,3 +2265,26 @@ def test_moe_comm_rank_mask_one_rank_masked_completes( local_num_tokens, top_k, ) + + @pytest.mark.threadleak(enabled=False) + @pytest.mark.parametrize( + "mpi_pool_executor,dead_rank,local_num_tokens,top_k", + [ + (4, 2, 16, 2), + ], + indirect=["mpi_pool_executor"], + ) + def test_moe_comm_rank_mask_inactive_before_combine_skips_stale_dispatch_slots( + self, + mpi_pool_executor, + dead_rank: int, + local_num_tokens: int, + top_k: int, + ) -> None: + """Verify combine skips slots from a rank masked inactive after dispatch.""" + _run_rank_mask_inactive_before_combine_test( + mpi_pool_executor, + dead_rank, + local_num_tokens, + top_k, + ) diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py new file mode 100644 index 000000000000..cdff07cbcea4 --- /dev/null +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -0,0 +1,643 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for AlltoAllWatchdog (WideEP fault tolerance, PR 1a.4).""" + +import threading +import time +from collections.abc import Callable +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + UNKNOWN_COMPLETION_FLAG, + AlltoAllWatchdog, + AlltoAllWatchdogCoordinator, + AlltoAllWatchdogTimeout, + CompletionFlagReadTimeout, +) +from tensorrt_llm._torch.modules.fused_moe.ep_group_health import EPGroupHealth +from tensorrt_llm._torch.modules.fused_moe.wide_ep_ft import get_wide_ep_ft_options + + +class FakeCompletionFlagReader: + """Thread-safe completion flag reader for pure-Python watchdog tests.""" + + def __init__(self, ep_size: int) -> None: + self._lock = threading.Lock() + self._flags = { + "dispatch": [0 for _ in range(ep_size)], + "combine": [0 for _ in range(ep_size)], + } + + def set_flags(self, phase: str, flags: list[int]) -> None: + with self._lock: + self._flags[phase] = list(flags) + + def read_completion_flags(self, phase: str) -> tuple[int, ...]: + with self._lock: + return tuple(self._flags[phase]) + + +class TimeoutCompletionFlagReader: + def read_completion_flags(self, phase: str) -> tuple[int, ...]: + raise CompletionFlagReadTimeout("blocked") + + +class OneGoodReadThenTimeoutReader: + def __init__(self, flags: tuple[int, ...]) -> None: + self._flags = flags + self._read_count = 0 + + def read_completion_flags(self, phase: str) -> tuple[int, ...]: + self._read_count += 1 + if self._read_count == 1: + return self._flags + raise CompletionFlagReadTimeout("blocked") + + +def _wait_for(predicate: Callable[[], bool], timeout_s: float = 1.0) -> None: + deadline = time.monotonic() + timeout_s + while time.monotonic() < deadline: + if predicate(): + return + time.sleep(0.005) + raise AssertionError("condition was not reached before timeout") + + +def test_watchdog_completes_when_all_active_flags_arrive() -> None: + health = EPGroupHealth(4) + reader = FakeCompletionFlagReader(ep_size=4) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.2, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + reader.set_flags("dispatch", [1, 1, 1, 1]) + assert watchdog.wait_until_idle(timeout_s=1.0) + + assert events == [] + assert health.all_active() is True + + +def test_watchdog_completes_when_flags_advance_past_expected_generation() -> None: + health = EPGroupHealth(4) + reader = FakeCompletionFlagReader(ep_size=4) + reader.set_flags("dispatch", [2, 1, 5, 3]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + assert watchdog.wait_until_idle(timeout_s=1.0) + + assert events == [] + assert health.all_active() is True + + +def test_watchdog_handles_signed_uint32_generation_boundaries() -> None: + reader = FakeCompletionFlagReader(ep_size=2) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=2, + ep_rank=0, + completion_reader=reader, + timeout_s=0.05, + poll_interval_s=0.005, + on_timeout=events.append, + ) as watchdog: + reader.set_flags("dispatch", [-(1 << 31), -(1 << 31)]) + watchdog.watch(phase="dispatch", expected_flag=1 << 31) + assert watchdog.wait_until_idle(timeout_s=1.0) + + reader.set_flags("dispatch", [-1, -1]) + watchdog.watch(phase="dispatch", expected_flag=(1 << 32) - 1) + assert watchdog.wait_until_idle(timeout_s=1.0) + + reader.set_flags("dispatch", [0, 0]) + watchdog.watch(phase="dispatch", expected_flag=0) + assert watchdog.wait_until_idle(timeout_s=1.0) + + assert events == [] + + +def test_watchdog_defaults_match_design_doc() -> None: + reader = FakeCompletionFlagReader(ep_size=1) + watchdog = AlltoAllWatchdog(ep_size=1, ep_rank=0, completion_reader=reader) + + assert watchdog._timeout_s == DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S + assert watchdog._poll_interval_s == DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S + + +def test_watchdog_stop_is_terminal() -> None: + reader = FakeCompletionFlagReader(ep_size=1) + watchdog = AlltoAllWatchdog( + ep_size=1, + ep_rank=0, + completion_reader=reader, + timeout_s=0.2, + poll_interval_s=0.005, + ) + watchdog.start() + watchdog.stop(timeout_s=1.0) + + with pytest.raises(RuntimeError, match="stopped AlltoAllWatchdog"): + watchdog.start() + with pytest.raises(RuntimeError, match="stopped AlltoAllWatchdog"): + watchdog.watch(phase="dispatch", expected_flag=1) + + +def test_wide_ep_ft_options_create_shared_health_when_enabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("TLLM_FAULT_TOLERANCE_MODE", "1") + model_config = SimpleNamespace( + extra_attrs={}, + mapping=SimpleNamespace(moe_ep_size=4), + ) + + health, timeout_s, poll_interval_s = get_wide_ep_ft_options(model_config) + health_again, timeout_again_s, poll_again_s = get_wide_ep_ft_options(model_config) + + assert isinstance(health, EPGroupHealth) + assert health_again is health + assert timeout_s == DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S + assert timeout_again_s == timeout_s + assert poll_interval_s == DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S + assert poll_again_s == poll_interval_s + + +def test_wide_ep_ft_options_ignore_legacy_enable_flag(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("TLLM_FAULT_TOLERANCE_MODE", raising=False) + monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "1") + model_config = SimpleNamespace( + extra_attrs={}, + mapping=SimpleNamespace(moe_ep_size=4), + ) + + health, timeout_s, _ = get_wide_ep_ft_options(model_config) + + assert health is None + assert timeout_s is None + + +def test_watchdog_coordinator_reuses_committed_mask_when_generation_is_unchanged() -> None: + health = EPGroupHealth(4) + coordinator = AlltoAllWatchdogCoordinator( + workspace_state={}, + workspace=torch.zeros((4, 1), dtype=torch.uint8), + metainfo=torch.zeros((1,), dtype=torch.int64), + metainfo_index={}, + ep_rank=0, + health=health, + ) + dispatch_snapshot = coordinator.capture_active_rank_mask(None) + assert dispatch_snapshot.active_rank_mask is not None + assert dispatch_snapshot.committed_generation == 0 + + combine_mask = coordinator.active_rank_mask_for_combine(dispatch_snapshot, None) + + assert combine_mask is dispatch_snapshot.active_rank_mask + assert combine_mask.tolist() == [0b1111, 0] + + +def test_watchdog_coordinator_fails_closed_on_committed_generation_change() -> None: + health = EPGroupHealth(4) + coordinator = AlltoAllWatchdogCoordinator( + workspace_state={}, + workspace=torch.zeros((4, 1), dtype=torch.uint8), + metainfo=torch.zeros((1,), dtype=torch.int64), + metainfo_index={}, + ep_rank=0, + health=health, + ) + dispatch_snapshot = coordinator.capture_active_rank_mask(None) + + health.mark_failed(2) + health.mark_active(2) + assert health.get_mask() == 0b1111 + + with pytest.raises( + RuntimeError, + match="committed EP membership changed between dispatch and combine", + ): + coordinator.active_rank_mask_for_combine(dispatch_snapshot, None) + + +def test_watchdog_coordinator_converts_atomic_snapshot_to_two_mask_words() -> None: + health = EPGroupHealth(72) + health.mark_failed(70) + coordinator = AlltoAllWatchdogCoordinator( + workspace_state={}, + workspace=torch.zeros((72, 1), dtype=torch.uint8), + metainfo=torch.zeros((1,), dtype=torch.int64), + metainfo_index={}, + ep_rank=0, + health=health, + ) + + dispatch_snapshot = coordinator.capture_active_rank_mask(None) + + assert dispatch_snapshot.committed_generation == 1 + assert dispatch_snapshot.active_rank_mask is not None + assert dispatch_snapshot.active_rank_mask.tolist() == [(1 << 64) - 1, 0xBF] + + +def test_watchdog_coordinator_explicit_mask_is_not_bound_to_health_generation() -> None: + health = EPGroupHealth(4) + coordinator = AlltoAllWatchdogCoordinator( + workspace_state={}, + workspace=torch.zeros((4, 1), dtype=torch.uint8), + metainfo=torch.zeros((1,), dtype=torch.int64), + metainfo_index={}, + ep_rank=0, + health=health, + ) + explicit_mask = torch.tensor([0b1101, 0], dtype=torch.uint64) + dispatch_snapshot = coordinator.capture_active_rank_mask(explicit_mask) + assert dispatch_snapshot.committed_generation is None + + health.mark_failed(2) + combine_mask = coordinator.active_rank_mask_for_combine(dispatch_snapshot, None) + + assert combine_mask is dispatch_snapshot.active_rank_mask + assert combine_mask.tolist() == explicit_mask.tolist() + with pytest.raises(ValueError, match="mask captured at dispatch"): + coordinator.active_rank_mask_for_combine( + dispatch_snapshot, + torch.tensor(health.get_mask_words(), dtype=torch.uint64), + ) + + +def test_watchdog_no_detected_failure_publication_to_committed_health() -> None: + health = EPGroupHealth(4) + committed_before = health.snapshot() + reader = FakeCompletionFlagReader(ep_size=4) + reader.set_flags("dispatch", [1, 0, 1, 0]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + _wait_for(lambda: len(events) == 1) + assert watchdog.wait_until_idle(timeout_s=1.0) + + event = events[0] + assert event.phase == "dispatch" + assert event.expected_flag == 1 + assert event.observed_flags == (1, 0, 1, 0) + assert event.missing_ranks == (1, 3) + assert not hasattr(event, "marked_failed_ranks") + assert health.snapshot() == committed_before + + +def test_watchdog_ignores_ranks_already_failed_in_health_mask() -> None: + health = EPGroupHealth(4) + assert health.mark_failed(2) is True + reader = FakeCompletionFlagReader(ep_size=4) + reader.set_flags("dispatch", [1, 1, 0, 1]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.05, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + assert watchdog.wait_until_idle(timeout_s=1.0) + + assert events == [] + assert health.get_failed_ranks() == frozenset({2}) + + +def test_watchdog_reports_local_missing_without_changing_committed_health() -> None: + health = EPGroupHealth(4) + reader = FakeCompletionFlagReader(ep_size=4) + reader.set_flags("combine", [0, 2, 2, 2]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=4, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="combine", expected_flag=2) + _wait_for(lambda: len(events) == 1) + + event = events[0] + assert event.missing_ranks == (0,) + assert health.get_failed_ranks() == frozenset() + + +def test_watchdog_poll_timeout_without_snapshot_fails_closed() -> None: + health = EPGroupHealth(3) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=3, + ep_rank=0, + completion_reader=TimeoutCompletionFlagReader(), + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + _wait_for(lambda: len(events) == 1) + + event = events[0] + assert event.poll_timed_out is True + assert event.observed_flags == (UNKNOWN_COMPLETION_FLAG,) * 3 + assert event.missing_ranks == (0, 1, 2) + assert health.all_active() is True + + +def test_watchdog_poll_timeout_with_prior_snapshot_does_not_mark_failed_rank() -> None: + health = EPGroupHealth(3) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=3, + ep_rank=0, + completion_reader=OneGoodReadThenTimeoutReader((1, 0, 1)), + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + _wait_for(lambda: len(events) == 1) + + event = events[0] + assert event.poll_timed_out is True + assert event.observed_flags == (1, 0, 1) + assert event.missing_ranks == (1,) + assert health.all_active() is True + + +def test_watchdog_callback_error_stops_and_clears_queue() -> None: + health = EPGroupHealth(2) + reader = FakeCompletionFlagReader(ep_size=2) + reader.set_flags("dispatch", [1, 0]) + + def raise_from_callback(event: AlltoAllWatchdogTimeout) -> None: + raise RuntimeError(f"callback failed for {event.phase}") + + with AlltoAllWatchdog( + ep_size=2, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=raise_from_callback, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + _wait_for(lambda: watchdog.last_error is not None) + assert watchdog.wait_until_idle(timeout_s=1.0) + assert isinstance(watchdog.last_error, RuntimeError) + with pytest.raises(RuntimeError, match="stopped AlltoAllWatchdog"): + watchdog.watch(phase="dispatch", expected_flag=2) + + +def test_watchdog_preserves_fifo_order_and_clears_followups_after_timeout() -> None: + health = EPGroupHealth(3) + reader = FakeCompletionFlagReader(ep_size=3) + reader.set_flags("dispatch", [1, 0, 1]) + reader.set_flags("combine", [0, 0, 0]) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog( + ep_size=3, + ep_rank=0, + completion_reader=reader, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=1) + watchdog.watch(phase="combine", expected_flag=2) + _wait_for(lambda: len(events) == 1) + assert watchdog.wait_until_idle(timeout_s=1.0) + time.sleep(0.05) + + assert len(events) == 1 + assert events[0].phase == "dispatch" + assert events[0].missing_ranks == (1,) + assert health.get_failed_ranks() == frozenset() + + +def test_watchdog_from_workspace_reads_phase_specific_offsets() -> None: + ep_size = 3 + ep_rank = 1 + workspace = torch.zeros((ep_size, 64), dtype=torch.uint8) + metainfo = torch.zeros((10,), dtype=torch.int64) + metainfo_index = { + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 4, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 5, + } + metainfo[4] = 4 + metainfo[5] = 20 + workspace[ep_rank, 4:16].view(torch.int32).copy_(torch.tensor([7, 7, 7], dtype=torch.int32)) + workspace[ep_rank, 20:32].view(torch.int32).copy_(torch.tensor([0, 8, 8], dtype=torch.int32)) + health = EPGroupHealth(ep_size) + events: list[AlltoAllWatchdogTimeout] = [] + + with AlltoAllWatchdog.from_workspace( + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=ep_rank, + ep_size=ep_size, + timeout_s=0.02, + poll_interval_s=0.005, + health=health, + on_timeout=events.append, + ) as watchdog: + watchdog.watch(phase="dispatch", expected_flag=7) + assert watchdog.wait_until_idle(timeout_s=1.0) + + watchdog.watch(phase="combine", expected_flag=8) + _wait_for(lambda: len(events) == 1) + + assert events[0].phase == "combine" + assert events[0].missing_ranks == (0,) + assert health.get_failed_ranks() == frozenset() + + +def test_workspace_coordinators_share_fifo_watchdog() -> None: + ep_size = 3 + ep_rank = 0 + workspace_state: dict[str, object] = {} + workspace = torch.zeros((ep_size, 64), dtype=torch.uint8) + metainfo = torch.tensor([0, 4, 16], dtype=torch.int64) + metainfo_index = { + "FLAG_VAL_OFFSET_INDEX": 0, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 1, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 2, + } + workspace[ep_rank, 4:16].view(torch.int32).copy_(torch.tensor([1, 0, 1])) + health = EPGroupHealth(ep_size) + events: list[AlltoAllWatchdogTimeout] = [] + on_timeout = events.append + coordinators = [ + AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=ep_rank, + health=health, + ) + for _ in range(2) + ] + watchdogs = [ + coordinator.acquire_watchdog( + ep_size=ep_size, + timeout_s=0.02, + poll_interval_s=0.005, + on_timeout=on_timeout, + ) + for coordinator in coordinators + ] + + try: + assert watchdogs[0] is watchdogs[1] + coordinators[0].watch_collective(watchdogs[0], "dispatch", None) + coordinators[1].watch_collective(watchdogs[1], "combine", None) + _wait_for(lambda: len(events) == 1) + assert watchdogs[0].wait_until_idle(timeout_s=1.0) + time.sleep(0.05) + + assert len(events) == 1 + assert events[0].phase == "dispatch" + assert events[0].missing_ranks == (1,) + assert health.get_failed_ranks() == frozenset() + finally: + for coordinator, watchdog in zip(coordinators, watchdogs): + coordinator.release_watchdog(watchdog) + + +def test_workspace_coordinator_wraps_shared_generation() -> None: + workspace_state: dict[str, object] = {} + workspace = torch.zeros((1, 32), dtype=torch.uint8) + workspace[0, 0:4].view(torch.int32).fill_(-1) + metainfo = torch.tensor([0, 4, 8], dtype=torch.int64) + metainfo_index = { + "FLAG_VAL_OFFSET_INDEX": 0, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 1, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 2, + } + coordinator = AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=0, + ) + watchdog = coordinator.acquire_watchdog( + ep_size=1, + timeout_s=0.05, + poll_interval_s=0.005, + ) + + try: + coordinator.watch_collective(watchdog, "dispatch", None) + assert watchdog.wait_until_idle(timeout_s=1.0) + finally: + coordinator.release_watchdog(watchdog) + + +def test_unmonitored_coordinator_advances_shared_generation() -> None: + workspace_state: dict[str, object] = {} + workspace = torch.zeros((2, 32), dtype=torch.uint8) + metainfo = torch.tensor([0, 4, 12], dtype=torch.int64) + metainfo_index = { + "FLAG_VAL_OFFSET_INDEX": 0, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 1, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 2, + } + coordinators = [ + AlltoAllWatchdogCoordinator( + workspace_state=workspace_state, + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=0, + ) + for _ in range(2) + ] + events: list[AlltoAllWatchdogTimeout] = [] + watchdog = coordinators[0].acquire_watchdog( + ep_size=2, + timeout_s=0.02, + poll_interval_s=0.005, + on_timeout=events.append, + ) + + try: + coordinators[1].watch_collective(None, "dispatch", None) + coordinators[0].watch_collective(watchdog, "dispatch", None) + _wait_for(lambda: len(events) == 1) + assert events[0].expected_flag == 2 + finally: + coordinators[0].release_watchdog(watchdog) + + +def test_watchdog_rejects_active_mask_without_local_rank() -> None: + reader = FakeCompletionFlagReader(ep_size=4) + with AlltoAllWatchdog( + ep_size=4, + ep_rank=2, + completion_reader=reader, + timeout_s=0.1, + poll_interval_s=0.005, + ) as watchdog: + with pytest.raises(ValueError, match="local ep_rank"): + watchdog.watch(phase="dispatch", expected_flag=1, active_mask=0b1011)