From 48cce165a4fe4fa3ec9b1a9e46bb2021b41aaf84 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 22 Jun 2026 12:43:33 -0700 Subject: [PATCH] [None][feat] WideEP FT: add AlltoAll watchdog --- tensorrt_llm/_torch/alltoall_watchdog.py | 367 ++++++++++++++++++ .../_torch/distributed/moe_alltoall.py | 87 ++++- .../communication/nvlink_one_sided.py | 92 ++++- .../_torch/modules/test_alltoall_watchdog.py | 236 +++++++++++ 4 files changed, 778 insertions(+), 4 deletions(-) create mode 100644 tensorrt_llm/_torch/alltoall_watchdog.py create mode 100644 tests/unittest/_torch/modules/test_alltoall_watchdog.py diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py new file mode 100644 index 000000000000..e793650443bc --- /dev/null +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -0,0 +1,367 @@ +# 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. +""" + +from __future__ import annotations + +import threading +import time +from collections import deque +from dataclasses import dataclass +from typing import Callable, Deque, Mapping, Optional, Protocol, Sequence + +import torch + +from tensorrt_llm.logger import logger as tllm_logger + + +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 EPGroupHealthLike(Protocol): + """Subset of EPGroupHealth used by the watchdog.""" + + def get_mask(self) -> int: + """Return the active-rank bitmask.""" + + def mark_failed(self, rank: int) -> bool: + """Mark ``rank`` failed and return whether state changed.""" + + +@dataclass(frozen=True) +class AlltoAllWatchdogTimeout: + """Details emitted when an AlltoAll phase times out.""" + + phase: str + expected_flag: int + observed_flags: tuple[int, ...] + missing_ranks: tuple[int, ...] + marked_failed_ranks: tuple[int, ...] + elapsed_s: float + + +@dataclass(frozen=True) +class _CollectiveWatch: + phase: str + expected_flag: int + active_mask: int + start_s: float + + +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, + ) -> 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), + } + + 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 != "cpu": + flags = flags.detach().cpu() + return tuple(int(v) for v in flags.tolist()) + + +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. + """ + + VALID_PHASES = frozenset({"dispatch", "combine"}) + + def __init__( + self, + *, + ep_size: int, + ep_rank: int, + completion_reader: CompletionFlagReader, + timeout_s: float, + poll_interval_s: float = 0.05, + health: Optional[EPGroupHealthLike] = None, + on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], 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._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, + poll_interval_s: float = 0.05, + health: Optional[EPGroupHealthLike] = None, + on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], 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, + ) + 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._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._stopping = True + self._queue.clear() + self._cv.notify_all() + thread = self._thread + if thread is not None: + 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 expected_flag < 0: + raise ValueError( + f"expected_flag must be non-negative, 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._stopping: + 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(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 observed_flags[rank] != watch.expected_flag) + + def _handle_timeout(self, watch: _CollectiveWatch, + observed_flags: tuple[int, ...]) -> None: + elapsed_s = time.monotonic() - watch.start_s + missing_ranks = self._missing_ranks(watch, observed_flags) + marked_failed: list[int] = [] + if self._health is not None: + for rank in missing_ranks: + if rank == self._ep_rank: + continue + if self._health.mark_failed(rank): + marked_failed.append(rank) + + event = AlltoAllWatchdogTimeout( + phase=watch.phase, + expected_flag=watch.expected_flag, + observed_flags=observed_flags, + missing_ranks=missing_ranks, + marked_failed_ranks=tuple(marked_failed), + elapsed_s=elapsed_s, + ) + 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 _run(self) -> None: + 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( + 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}") + except BaseException as exc: # noqa: BLE001 - keep watchdog failures visible. + with self._cv: + self._last_error = exc + self._queue.clear() + self._cv.notify_all() + 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() + continue + + if time.monotonic() - watch.start_s >= self._timeout_s: + self._handle_timeout(watch, observed_flags) + 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() + 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 d99db01a32ca..df51a347800b 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -9,11 +9,13 @@ import os 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 (AlltoAllWatchdog, + AlltoAllWatchdogTimeout) from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -126,6 +128,11 @@ def __init__( num_slots: int, workspace_size_per_rank: int, num_experts: Optional[int] = None, + ep_group_health=None, + alltoall_watchdog_timeout_s: Optional[float] = None, + alltoall_watchdog_poll_interval_s: float = 0.05, + alltoall_watchdog_on_timeout: Optional[Callable[ + [AlltoAllWatchdogTimeout], None]] = None, ): """ Initialize MoeAlltoAll with workspace allocation. @@ -138,6 +145,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 EPGroupHealth-compatible object. When present, its mask is passed to the + CUDA kernels and used by the watchdog. + 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 +227,65 @@ def __init__( self.metainfo = self._WORKSPACE["metainfo"] # Internal state self._state: _A2AState = _A2AState() + self.ep_group_health = ep_group_health + self._watchdog_flag_generation = 0 + self._alltoall_watchdog: AlltoAllWatchdog | None = None + if alltoall_watchdog_timeout_s is not None: + self._watchdog_flag_generation = self._read_current_flag_val() + self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( + workspace=self.workspace, + metainfo=self.metainfo, + metainfo_index=self._METAINFO_INDEX, + ep_rank=self.ep_rank, + ep_size=self.ep_size, + timeout_s=alltoall_watchdog_timeout_s, + poll_interval_s=alltoall_watchdog_poll_interval_s, + health=self.ep_group_health, + on_timeout=alltoall_watchdog_on_timeout, + ) + + def _read_current_flag_val(self) -> int: + flag_val_offset = 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 int(flag_val.item()) + + def _get_active_rank_mask_tensor( + self, + active_rank_mask: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + if active_rank_mask is not None: + return active_rank_mask + if self.ep_group_health is None: + return None + return torch.tensor(self.ep_group_health.get_mask_words(), + dtype=torch.uint64, + device="cpu") + + def _active_mask_int( + self, active_rank_mask: Optional[torch.Tensor]) -> Optional[int]: + 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.ep_group_health is not None: + return self.ep_group_health.get_mask() + return None + + def _watch_collective(self, phase: str, + active_rank_mask: Optional[torch.Tensor]) -> None: + if self._alltoall_watchdog is None: + return + self._watchdog_flag_generation += 1 + self._alltoall_watchdog.watch( + phase=phase, + expected_flag=self._watchdog_flag_generation, + active_mask=self._active_mask_int(active_rank_mask), + ) def dispatch(self, token_selected_experts: torch.Tensor, @@ -221,7 +293,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 +305,7 @@ 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 ep_group_health for this dispatch. Returns: recv_tensors: List of tensors received, each has shape [ep_size, max_tokens_per_rank, payload_num_elements_per_token] @@ -246,6 +320,7 @@ def dispatch(self, 0 ) == self.eplb_stats_num_experts, "eplb_local_stats size must match eplb_stats_num_experts" + active_rank_mask = self._get_active_rank_mask_tensor(active_rank_mask) recv_tensors, combine_payload_offset, eplb_gathered_stats = torch.ops.trtllm.moe_a2a_dispatch( token_selected_experts, input_payloads, @@ -257,7 +332,9 @@ def dispatch(self, self.top_k, self.num_experts, eplb_local_stats, + active_rank_mask, ) + self._watch_collective("dispatch", active_rank_mask) if eplb_gathered_stats.numel() == 0: eplb_gathered_stats = None @@ -287,6 +364,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 +374,7 @@ 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 overriding ep_group_health for this combine. Returns: combined_output: [local_num_tokens, num_elements_per_token] tensor of combined results @@ -303,11 +382,13 @@ 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 = self._get_active_rank_mask_tensor(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._watch_collective("combine", active_rank_mask) # Reset state for next round self._state = _A2AState() 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 9b6a8c820f24..49da27e07997 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,13 @@ """ 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 (AlltoAllWatchdog, + AlltoAllWatchdogTimeout) 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,6 +153,11 @@ def __init__( dtype: Optional[torch.dtype] = None, num_experts: Optional[int] = None, use_low_precision_combine: bool = False, + ep_group_health=None, + alltoall_watchdog_timeout_s: Optional[float] = None, + alltoall_watchdog_poll_interval_s: float = 0.05, + alltoall_watchdog_on_timeout: Optional[Callable[ + [AlltoAllWatchdogTimeout], None]] = None, ): """ Initialize NVLinkOneSided with workspace allocation. @@ -169,6 +176,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 EPGroupHealth-compatible object. When present, its mask is passed to the + CUDA kernels and used by the watchdog. + 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) @@ -300,6 +313,29 @@ def __init__( 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_flag_generation = 0 + self._alltoall_watchdog: AlltoAllWatchdog | None = None + if alltoall_watchdog_timeout_s is not None: + self._watchdog_flag_generation = self._read_current_flag_val() + self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( + 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, + ep_size=self.ep_size, + timeout_s=alltoall_watchdog_timeout_s, + poll_interval_s=alltoall_watchdog_poll_interval_s, + health=self.ep_group_health, + on_timeout=alltoall_watchdog_on_timeout, + ) # Initialize dispatch state self._dispatch_state = {"phase": "idle"} @@ -307,6 +343,49 @@ def __init__( # Invalid token expert ID (default to -1), the kernels in TRTLLM-gen is hard-code to support -1 only. self.invalid_token_expert_id: int = -1 + def _read_current_flag_val(self) -> int: + flag_val_offset = self.moe_a2a_metainfo[ + self.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 int(flag_val.item()) + + def _get_active_rank_mask_tensor( + self, + active_rank_mask: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + if active_rank_mask is not None: + return active_rank_mask + if self.ep_group_health is None: + return None + return torch.tensor(self.ep_group_health.get_mask_words(), + dtype=torch.uint64, + device="cpu") + + def _active_mask_int( + self, active_rank_mask: Optional[torch.Tensor]) -> Optional[int]: + 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.ep_group_health is not None: + return self.ep_group_health.get_mask() + return None + + def _watch_collective(self, phase: str, + active_rank_mask: Optional[torch.Tensor]) -> None: + if self._alltoall_watchdog is None: + return + self._watchdog_flag_generation += 1 + self._alltoall_watchdog.watch( + phase=phase, + expected_flag=self._watchdog_flag_generation, + active_mask=self._active_mask_int(active_rank_mask), + ) + @staticmethod def is_platform_supported() -> bool: """ @@ -326,6 +405,9 @@ def destroy(self): return self._destroyed = True + if self._alltoall_watchdog is not None: + self._alltoall_watchdog.stop(timeout_s=1.0) + self._alltoall_watchdog = None workspace_key = getattr(self, "_workspace_key", None) if workspace_key is None: return @@ -409,6 +491,8 @@ 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 = self._get_active_rank_mask_tensor( + kwargs.get("active_rank_mask")) recv_buffers, combine_payload_offset, eplb_gathered_stats = ( torch.ops.trtllm.moe_a2a_dispatch( @@ -422,8 +506,10 @@ def dispatch( self.top_k, self.num_experts, eplb_local_stats, + active_rank_mask, ) ) + self._watch_collective("dispatch", active_rank_mask) if eplb_gathered_stats.numel() == 0: eplb_gathered_stats = None self._dispatch_state["eplb_gathered_stats"] = eplb_gathered_stats @@ -526,6 +612,8 @@ def combine( raise ValueError( f"final_hidden_states must be 2D or 3D, got {final_hidden_states.dim()}D" ) + active_rank_mask = self._get_active_rank_mask_tensor( + kwargs.get("active_rank_mask")) output = torch.ops.trtllm.moe_a2a_combine( final_hidden_states, int(local_num_tokens), @@ -538,7 +626,9 @@ def combine( int(combine_payload_offset), bool(self.payload_in_workspace), bool(self.use_low_precision_combine), + active_rank_mask, ) + self._watch_collective("combine", active_rank_mask) # Reset state for next round self._dispatch_state = {"phase": "idle"} 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..4ff2446899f5 --- /dev/null +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -0,0 +1,236 @@ +# 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 + +import pytest +import torch + +from tensorrt_llm._torch.alltoall_watchdog import (AlltoAllWatchdog, + AlltoAllWatchdogTimeout) +from tensorrt_llm._torch.modules.fused_moe.ep_group_health import EPGroupHealth + + +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]) + + +def _wait_for(predicate, 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_timeout_reports_and_marks_missing_remote_ranks() -> None: + health = EPGroupHealth(4) + 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 event.marked_failed_ranks == (1, 3) + assert health.get_failed_ranks() == frozenset({1, 3}) + + +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_but_does_not_mark_local_failed( +) -> 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 event.marked_failed_ranks == () + assert health.get_failed_ranks() == frozenset() + + +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({1}) + + +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 events[0].marked_failed_ranks == (0, ) + assert health.get_failed_ranks() == frozenset({0}) + + +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)