From ad9590a48829777c760396b8e108536e879fcd98 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 1/8] [None][feat] WideEP FT: add AlltoAll watchdog Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/alltoall_watchdog.py | 366 ++++++++++++++++++ .../_torch/distributed/moe_alltoall.py | 87 ++++- .../communication/nvlink_one_sided.py | 78 +++- .../_torch/modules/test_alltoall_watchdog.py | 229 +++++++++++ 4 files changed, 756 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..7c35669e29ee --- /dev/null +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -0,0 +1,366 @@ +# 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 ca3a50dcfd12..8284b2d1087a 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.reset_state() 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..7b5ed16e8256 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,12 @@ """ 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 +152,10 @@ 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 +174,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 +311,26 @@ 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 +338,42 @@ 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 +393,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 +479,7 @@ 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 +493,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 +599,7 @@ 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 +612,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.reset_state() 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..1168bb0f00fe --- /dev/null +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -0,0 +1,229 @@ +# 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) From 80b8335db69c45443cea60e0ae206b5720a63519 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 22 Jun 2026 15:04:38 -0700 Subject: [PATCH 2/8] [None][fix] Address AlltoAll watchdog review findings Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/alltoall_watchdog.py | 116 +++++++++++++++--- .../_torch/distributed/moe_alltoall.py | 28 ++++- .../communication/communication_factory.py | 13 ++ .../communication/nvlink_one_sided.py | 11 +- .../modules/fused_moe/fused_moe_cutlass.py | 7 ++ .../modules/fused_moe/fused_moe_trtllm_gen.py | 9 +- .../_torch/modules/fused_moe/wide_ep_ft.py | 82 +++++++++++++ .../_torch/modules/test_alltoall_watchdog.py | 102 ++++++++++++++- 8 files changed, 345 insertions(+), 23 deletions(-) create mode 100644 tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py index 7c35669e29ee..79a6347b58e2 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -31,8 +31,13 @@ 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) + class CompletionFlagReader(Protocol): """Reads one phase's rank-local completion flag row.""" @@ -51,6 +56,10 @@ def mark_failed(self, rank: int) -> bool: """Mark ``rank`` failed and return whether state changed.""" +class CompletionFlagReadTimeout(TimeoutError): + """Raised when the host watchdog cannot read completion flags in time.""" + + @dataclass(frozen=True) class AlltoAllWatchdogTimeout: """Details emitted when an AlltoAll phase times out.""" @@ -61,6 +70,7 @@ class AlltoAllWatchdogTimeout: missing_ranks: tuple[int, ...] marked_failed_ranks: tuple[int, ...] elapsed_s: float + poll_timed_out: bool = False @dataclass(frozen=True) @@ -81,6 +91,7 @@ def __init__( 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]") @@ -97,11 +108,50 @@ def __init__( "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._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() + + host_flags = torch.empty( + (self._ep_size,), + dtype=torch.int32, + device="cpu", + pin_memory=prefer_pinned(), + ) + event = torch.cuda.Event(blocking=False) + 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)) + 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()) @@ -123,8 +173,8 @@ def __init__( ep_size: int, ep_rank: int, completion_reader: CompletionFlagReader, - timeout_s: float, - poll_interval_s: float = 0.05, + timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, health: Optional[EPGroupHealthLike] = None, on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, ) -> None: @@ -160,8 +210,8 @@ def from_workspace( metainfo_index: Mapping[str, int], ep_rank: int, ep_size: int, - timeout_s: float, - poll_interval_s: float = 0.05, + timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, health: Optional[EPGroupHealthLike] = None, on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, ) -> "AlltoAllWatchdog": @@ -178,6 +228,7 @@ def from_workspace( 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, @@ -288,11 +339,18 @@ def _missing_ranks( if observed_flags[rank] != watch.expected_flag ) - def _handle_timeout(self, watch: _CollectiveWatch, observed_flags: tuple[int, ...]) -> None: + 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) marked_failed: list[int] = [] - if self._health is not None: + has_known_flags = UNKNOWN_COMPLETION_FLAG not in observed_flags + if self._health is not None and (has_known_flags or not poll_timed_out): for rank in missing_ranks: if rank == self._ep_rank: continue @@ -306,20 +364,37 @@ def _handle_timeout(self, watch: _CollectiveWatch, observed_flags: tuple[int, .. missing_ranks=missing_ranks, marked_failed_ranks=tuple(marked_failed), elapsed_s=elapsed_s, + poll_timed_out=poll_timed_out, ) - 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 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, marked ranks %s", + self._ep_rank, + watch.phase, + elapsed_s, + watch.expected_flag, + list(self._active_ranks(watch.active_mask)), + list(observed_flags), + list(marked_failed), + ) + 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 _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: @@ -337,6 +412,11 @@ def _run(self) -> None: 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 BaseException as exc: # noqa: BLE001 - keep watchdog failures visible. with self._cv: self._last_error = exc @@ -350,16 +430,20 @@ def _run(self) -> None: 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: - self._handle_timeout(watch, observed_flags) + self._handle_timeout(watch, observed_flags, poll_timed_out=poll_timed_out) 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: diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index 8284b2d1087a..fe2c1020d201 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -8,14 +8,17 @@ # ruff: noqa: E501 import os +import sys from dataclasses import dataclass 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._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, 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 @@ -130,7 +133,8 @@ def __init__( 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_poll_interval_s: + float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, alltoall_watchdog_on_timeout: Optional[Callable[ [AlltoAllWatchdogTimeout], None]] = None, ): @@ -228,8 +232,12 @@ def __init__( # Internal state self._state: _A2AState = _A2AState() self.ep_group_health = ep_group_health + self._destroyed = False self._watchdog_flag_generation = 0 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._watchdog_flag_generation = self._read_current_flag_val() self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( @@ -244,6 +252,20 @@ def __init__( 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: + watchdog.stop(timeout_s=1.0) + self._alltoall_watchdog = None + + def __del__(self) -> None: + if not sys.is_finalizing(): + self.destroy() + def _read_current_flag_val(self) -> int: flag_val_offset = self.metainfo[ self._METAINFO_INDEX["FLAG_VAL_OFFSET_INDEX"]].item() 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 7b5ed16e8256..9b3d72306775 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 @@ -30,7 +30,12 @@ import torch from tensorrt_llm._mnnvl_utils import MnnvlMemory -from tensorrt_llm._torch.alltoall_watchdog import AlltoAllWatchdog, AlltoAllWatchdogTimeout +from tensorrt_llm._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + 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 @@ -154,7 +159,7 @@ def __init__( 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_poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, alltoall_watchdog_on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, ): """ @@ -314,6 +319,8 @@ def __init__( self.ep_group_health = ep_group_health self._watchdog_flag_generation = 0 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._watchdog_flag_generation = self._read_current_flag_val() self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( 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..69e5c53f061e --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py @@ -0,0 +1,82 @@ +# 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 = "TRTLLM_ENABLE_WIDE_EP_FT" +_TIMEOUT_ENV = "TRTLLM_ALLTOALL_WATCHDOG_TIMEOUT_S" +_POLL_INTERVAL_ENV = "TRTLLM_ALLTOALL_WATCHDOG_POLL_INTERVAL_S" + +_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, "0").lower() in {"1", "true", "yes", "on"} + + +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 the shared EP health object 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 ``TRTLLM_ENABLE_WIDE_EP_FT=1`` to create + one process-local health object shared by all MoE communication layers. + """ + + 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/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index 1168bb0f00fe..ba4b6774f412 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -16,12 +16,21 @@ import threading import time +from types import SimpleNamespace import pytest import torch -from tensorrt_llm._torch.alltoall_watchdog import AlltoAllWatchdog, AlltoAllWatchdogTimeout +from tensorrt_llm._torch.alltoall_watchdog import ( + DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, + DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + UNKNOWN_COMPLETION_FLAG, + AlltoAllWatchdog, + 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: @@ -43,6 +52,23 @@ def read_completion_flags(self, phase: str) -> tuple[int, ...]: 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, timeout_s: float = 1.0) -> None: deadline = time.monotonic() + timeout_s while time.monotonic() < deadline: @@ -74,6 +100,32 @@ def test_watchdog_completes_when_all_active_flags_arrive() -> None: assert health.all_active() is True +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_wide_ep_ft_options_create_shared_health_when_enabled(monkeypatch) -> None: + monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "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_watchdog_timeout_reports_and_marks_missing_remote_ranks() -> None: health = EPGroupHealth(4) reader = FakeCompletionFlagReader(ep_size=4) @@ -149,6 +201,54 @@ def test_watchdog_reports_local_missing_but_does_not_mark_local_failed() -> None 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 event.marked_failed_ranks == () + assert health.all_active() is True + + +def test_watchdog_poll_timeout_with_prior_snapshot_marks_known_missing_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 event.marked_failed_ranks == (1,) + assert health.get_failed_ranks() == frozenset({1}) + + def test_watchdog_preserves_fifo_order_and_clears_followups_after_timeout() -> None: health = EPGroupHealth(3) reader = FakeCompletionFlagReader(ep_size=3) From 15d429376f851cf5ca293331985afb4d62b36a63 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 22 Jun 2026 15:20:17 -0700 Subject: [PATCH 3/8] fix: make AlltoAll watchdog stop terminal Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/alltoall_watchdog.py | 6 +++++- .../_torch/modules/test_alltoall_watchdog.py | 18 ++++++++++++++++++ 2 files changed, 23 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py index 79a6347b58e2..2688cfa0e2d7 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -197,6 +197,7 @@ def __init__( 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 @@ -249,6 +250,8 @@ def last_error(self) -> BaseException | None: 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 @@ -262,6 +265,7 @@ def start(self) -> None: 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() @@ -291,7 +295,7 @@ def watch( self.start() with self._cv: - if self._stopping: + if self._closed: raise RuntimeError("cannot queue a stopped AlltoAllWatchdog") self._queue.append( _CollectiveWatch( diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index ba4b6774f412..cda4d813d146 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -108,6 +108,24 @@ def test_watchdog_defaults_match_design_doc() -> None: 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) -> None: monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "1") model_config = SimpleNamespace( From 0f4d98f38807cddf0503b45284252124070e6ba8 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 23 Jun 2026 10:09:02 -0700 Subject: [PATCH 4/8] fix: address AlltoAll watchdog review comments Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- .../moeAlltoAllKernels.cu | 8 +- .../communicationKernels/moeAlltoAllKernels.h | 4 +- tensorrt_llm/_torch/alltoall_watchdog.py | 37 ++- .../_torch/distributed/moe_alltoall.py | 32 ++- .../communication/nvlink_one_sided.py | 30 +- .../_torch/modules/moe/test_moe_comm.py | 267 +++++++++++++++++- .../_torch/modules/test_alltoall_watchdog.py | 7 +- 7 files changed, 344 insertions(+), 41 deletions(-) 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 index 2688cfa0e2d7..c41a3fb8fdb3 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -26,8 +26,9 @@ import threading import time from collections import deque +from collections.abc import Callable, Mapping, Sequence from dataclasses import dataclass -from typing import Callable, Deque, Mapping, Optional, Protocol, Sequence +from typing import Protocol import torch @@ -110,6 +111,8 @@ def __init__( } 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) @@ -123,13 +126,17 @@ def _read_cuda_flags(self, flags: torch.Tensor) -> tuple[int, ...]: assert self._copy_stream is not None self._prune_retired_copies() - host_flags = torch.empty( - (self._ep_size,), - dtype=torch.int32, - device="cpu", - pin_memory=prefer_pinned(), - ) - event = torch.cuda.Event(blocking=False) + 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) @@ -139,6 +146,8 @@ def _read_cuda_flags(self, flags: torch.Tensor) -> tuple[int, ...]: 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" ) @@ -175,8 +184,8 @@ def __init__( completion_reader: CompletionFlagReader, timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, - health: Optional[EPGroupHealthLike] = None, - on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, + 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}") @@ -196,7 +205,7 @@ def __init__( self._on_timeout = on_timeout self._cv = threading.Condition() - self._queue: Deque[_CollectiveWatch] = deque() + self._queue: deque[_CollectiveWatch] = deque() self._closed = False self._stopping = False self._thread: threading.Thread | None = None @@ -213,8 +222,8 @@ def from_workspace( ep_size: int, timeout_s: float = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, poll_interval_s: float = DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, - health: Optional[EPGroupHealthLike] = None, - on_timeout: Optional[Callable[[AlltoAllWatchdogTimeout], None]] = None, + 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( @@ -421,7 +430,7 @@ def _run(self) -> None: except CompletionFlagReadTimeout: observed_flags = last_observed_flags poll_timed_out = True - except BaseException as exc: # noqa: BLE001 - keep watchdog failures visible. + except Exception as exc: # noqa: BLE001 - keep watchdog failures visible. with self._cv: self._last_error = exc self._queue.clear() diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index fe2c1020d201..b16c2a25bc6d 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -9,6 +9,7 @@ import os import sys +import threading from dataclasses import dataclass from typing import Callable, Dict, Optional @@ -212,6 +213,8 @@ def __init__( "mnnvl_mem": mnnvl_mem, "workspace": workspace, "metainfo": metainfo, + "watchdog_flag_generation": 0, + "watchdog_flag_generation_lock": threading.Lock(), } else: assert self._WORKSPACE[ @@ -229,17 +232,20 @@ def __init__( self.mnnvl_mem = self._WORKSPACE["mnnvl_mem"] self.workspace = self._WORKSPACE["workspace"] self.metainfo = self._WORKSPACE["metainfo"] + if "watchdog_flag_generation_lock" not in self._WORKSPACE: + self._WORKSPACE["watchdog_flag_generation_lock"] = threading.Lock() + self._WORKSPACE[ + "watchdog_flag_generation"] = self._read_current_flag_val() # Internal state self._state: _A2AState = _A2AState() self.ep_group_health = ep_group_health self._destroyed = False - self._watchdog_flag_generation = 0 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._watchdog_flag_generation = self._read_current_flag_val() + self._sync_watchdog_flag_generation() self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( workspace=self.workspace, metainfo=self.metainfo, @@ -276,6 +282,25 @@ def _read_current_flag_val(self) -> int: flag_val = flag_val.detach().cpu() return int(flag_val.item()) + def _sync_watchdog_flag_generation(self) -> None: + workspace_state = self._WORKSPACE + assert workspace_state is not None + lock = workspace_state["watchdog_flag_generation_lock"] + with lock: + workspace_state["watchdog_flag_generation"] = max( + int(workspace_state["watchdog_flag_generation"]), + self._read_current_flag_val(), + ) + + def _next_watchdog_flag_generation(self) -> int: + workspace_state = self._WORKSPACE + assert workspace_state is not None + lock = workspace_state["watchdog_flag_generation_lock"] + with lock: + workspace_state["watchdog_flag_generation"] = ( + int(workspace_state["watchdog_flag_generation"]) + 1) + return int(workspace_state["watchdog_flag_generation"]) + def _get_active_rank_mask_tensor( self, active_rank_mask: Optional[torch.Tensor]) -> Optional[torch.Tensor]: @@ -302,10 +327,9 @@ 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, + expected_flag=self._next_watchdog_flag_generation(), active_mask=self._active_mask_int(active_rank_mask), ) 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 9b3d72306775..db5966615356 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,6 +25,7 @@ """ import os +import threading from typing import Callable, Dict, List, Optional, Tuple import torch @@ -287,6 +288,8 @@ def __init__( "mnnvl_mem": mnnvl_mem, "workspace": workspace, "metainfo": metainfo, + "watchdog_flag_generation": 0, + "watchdog_flag_generation_lock": threading.Lock(), } NVLinkOneSided._WORKSPACES[self._workspace_key] = workspace_state else: @@ -312,17 +315,20 @@ 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"] + if "watchdog_flag_generation_lock" not in workspace_state: + workspace_state["watchdog_flag_generation_lock"] = threading.Lock() + workspace_state["watchdog_flag_generation"] = self._read_current_flag_val() self.ep_group_health = ep_group_health - self._watchdog_flag_generation = 0 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._watchdog_flag_generation = self._read_current_flag_val() + self._sync_watchdog_flag_generation() self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( workspace=self.workspace, metainfo=self.moe_a2a_metainfo, @@ -354,6 +360,22 @@ def _read_current_flag_val(self) -> int: flag_val = flag_val.detach().cpu() return int(flag_val.item()) + def _sync_watchdog_flag_generation(self) -> None: + lock = self._workspace_state["watchdog_flag_generation_lock"] + with lock: + self._workspace_state["watchdog_flag_generation"] = max( + int(self._workspace_state["watchdog_flag_generation"]), + self._read_current_flag_val(), + ) + + def _next_watchdog_flag_generation(self) -> int: + lock = self._workspace_state["watchdog_flag_generation_lock"] + with lock: + self._workspace_state["watchdog_flag_generation"] = ( + int(self._workspace_state["watchdog_flag_generation"]) + 1 + ) + return int(self._workspace_state["watchdog_flag_generation"]) + def _get_active_rank_mask_tensor( self, active_rank_mask: Optional[torch.Tensor] ) -> Optional[torch.Tensor]: @@ -374,10 +396,9 @@ def _active_mask_int(self, active_rank_mask: Optional[torch.Tensor]) -> Optional 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, + expected_flag=self._next_watchdog_flag_generation(), active_mask=self._active_mask_int(active_rank_mask), ) @@ -424,6 +445,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: 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 index cda4d813d146..b44140dbd21f 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -16,6 +16,7 @@ import threading import time +from collections.abc import Callable from types import SimpleNamespace import pytest @@ -69,7 +70,7 @@ def read_completion_flags(self, phase: str) -> tuple[int, ...]: raise CompletionFlagReadTimeout("blocked") -def _wait_for(predicate, timeout_s: float = 1.0) -> None: +def _wait_for(predicate: Callable[[], bool], timeout_s: float = 1.0) -> None: deadline = time.monotonic() + timeout_s while time.monotonic() < deadline: if predicate(): @@ -126,7 +127,9 @@ def test_watchdog_stop_is_terminal() -> None: watchdog.watch(phase="dispatch", expected_flag=1) -def test_wide_ep_ft_options_create_shared_health_when_enabled(monkeypatch) -> None: +def test_wide_ep_ft_options_create_shared_health_when_enabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "1") model_config = SimpleNamespace( extra_attrs={}, From da348ac1fce23c79211d20f0ccb20b6e9f966dd2 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 29 Jun 2026 11:40:25 -0700 Subject: [PATCH 5/8] fix: accept advanced watchdog completion flags Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/alltoall_watchdog.py | 4 ++-- .../_torch/modules/test_alltoall_watchdog.py | 22 +++++++++++++++++++ 2 files changed, 24 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py index c41a3fb8fdb3..81d13438f10a 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -339,7 +339,7 @@ def _active_ranks(self, active_mask: int) -> tuple[int, ...]: def _phase_complete(self, watch: _CollectiveWatch, observed_flags: tuple[int, ...]) -> bool: return all( - observed_flags[rank] == watch.expected_flag + observed_flags[rank] >= watch.expected_flag for rank in self._active_ranks(watch.active_mask) ) @@ -349,7 +349,7 @@ def _missing_ranks( return tuple( rank for rank in self._active_ranks(watch.active_mask) - if observed_flags[rank] != watch.expected_flag + if observed_flags[rank] < watch.expected_flag ) def _handle_timeout( diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index b44140dbd21f..5d8a7826c70e 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -101,6 +101,28 @@ def test_watchdog_completes_when_all_active_flags_arrive() -> None: 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_defaults_match_design_doc() -> None: reader = FakeCompletionFlagReader(ep_size=1) watchdog = AlltoAllWatchdog(ep_size=1, ep_rank=0, completion_reader=reader) From 3a734e4d143b1d86fd7cff66bb0bb0b7a3f32588 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Mon, 29 Jun 2026 12:38:33 -0700 Subject: [PATCH 6/8] fix: harden shared AlltoAll watchdog state Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/alltoall_watchdog.py | 247 +++++++++++++++++- .../_torch/distributed/moe_alltoall.py | 109 ++------ .../communication/nvlink_one_sided.py | 105 +++----- .../_torch/modules/test_alltoall_watchdog.py | 178 ++++++++++++- 4 files changed, 464 insertions(+), 175 deletions(-) diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py index 81d13438f10a..1cb3fe88aae7 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -26,7 +26,7 @@ import threading import time from collections import deque -from collections.abc import Callable, Mapping, Sequence +from collections.abc import Callable, Mapping, MutableMapping, Sequence from dataclasses import dataclass from typing import Protocol @@ -38,6 +38,23 @@ 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 +_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): @@ -53,6 +70,9 @@ class EPGroupHealthLike(Protocol): 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 mark_failed(self, rank: int) -> bool: """Mark ``rank`` failed and return whether state changed.""" @@ -82,6 +102,25 @@ class _CollectiveWatch: 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.""" @@ -166,6 +205,168 @@ def read_completion_flags(self, phase: str) -> tuple[int, ...]: return tuple(int(v) for v in flags.tolist()) +class AlltoAllWatchdogCoordinator: + """Shared watchdog plumbing for MoE AlltoAll frontends.""" + + 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 active_rank_mask_tensor(self, active_rank_mask: torch.Tensor | None) -> torch.Tensor | None: + if active_rank_mask is not None: + return active_rank_mask + if self._health is None: + return None + return torch.tensor(self._health.get_mask_words(), dtype=torch.uint64, device="cpu") + + 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. @@ -279,7 +480,7 @@ def stop(self, timeout_s: float | None = None) -> None: self._queue.clear() self._cv.notify_all() thread = self._thread - if thread is not None: + if thread is not None and thread is not threading.current_thread(): thread.join(timeout=timeout_s) def watch( @@ -292,8 +493,10 @@ def watch( """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 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() @@ -339,7 +542,7 @@ def _active_ranks(self, active_mask: int) -> tuple[int, ...]: def _phase_complete(self, watch: _CollectiveWatch, observed_flags: tuple[int, ...]) -> bool: return all( - observed_flags[rank] >= watch.expected_flag + _completion_flag_reached(observed_flags[rank], watch.expected_flag) for rank in self._active_ranks(watch.active_mask) ) @@ -349,7 +552,7 @@ def _missing_ranks( return tuple( rank for rank in self._active_ranks(watch.active_mask) - if observed_flags[rank] < watch.expected_flag + if not _completion_flag_reached(observed_flags[rank], watch.expected_flag) ) def _handle_timeout( @@ -362,8 +565,11 @@ def _handle_timeout( elapsed_s = time.monotonic() - watch.start_s missing_ranks = self._missing_ranks(watch, observed_flags) marked_failed: list[int] = [] - has_known_flags = UNKNOWN_COMPLETION_FLAG not in observed_flags - if self._health is not None and (has_known_flags or not poll_timed_out): + if ( + self._health is not None + and not poll_timed_out + and UNKNOWN_COMPLETION_FLAG not in observed_flags + ): for rank in missing_ranks: if rank == self._ep_rank: continue @@ -405,6 +611,14 @@ def _handle_timeout( 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 @@ -418,7 +632,8 @@ def _run(self) -> None: try: observed_flags = tuple( - int(v) for v in self._completion_reader.read_completion_flags(watch.phase) + _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( @@ -431,10 +646,7 @@ def _run(self) -> None: observed_flags = last_observed_flags poll_timed_out = True except Exception as exc: # noqa: BLE001 - keep watchdog failures visible. - with self._cv: - self._last_error = exc - self._queue.clear() - self._cv.notify_all() + self._stop_after_error(exc) tllm_logger.error("AlltoAll watchdog stopped after polling error: %s", exc) return @@ -448,7 +660,14 @@ def _run(self) -> None: continue if time.monotonic() - watch.start_s >= self._timeout_s: - self._handle_timeout(watch, observed_flags, poll_timed_out=poll_timed_out) + 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 diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index b16c2a25bc6d..6e44510a8573 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -9,7 +9,6 @@ import os import sys -import threading from dataclasses import dataclass from typing import Callable, Dict, Optional @@ -19,7 +18,7 @@ from tensorrt_llm._torch.alltoall_watchdog import ( DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, AlltoAllWatchdog, - AlltoAllWatchdogTimeout) + 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 @@ -132,13 +131,13 @@ def __init__( num_slots: int, workspace_size_per_rank: int, num_experts: Optional[int] = None, - ep_group_health=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. @@ -213,8 +212,6 @@ def __init__( "mnnvl_mem": mnnvl_mem, "workspace": workspace, "metainfo": metainfo, - "watchdog_flag_generation": 0, - "watchdog_flag_generation_lock": threading.Lock(), } else: assert self._WORKSPACE[ @@ -232,29 +229,31 @@ def __init__( self.mnnvl_mem = self._WORKSPACE["mnnvl_mem"] self.workspace = self._WORKSPACE["workspace"] self.metainfo = self._WORKSPACE["metainfo"] - if "watchdog_flag_generation_lock" not in self._WORKSPACE: - self._WORKSPACE["watchdog_flag_generation_lock"] = threading.Lock() - self._WORKSPACE[ - "watchdog_flag_generation"] = self._read_current_flag_val() # 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._sync_watchdog_flag_generation() - self._alltoall_watchdog = AlltoAllWatchdog.from_workspace( - workspace=self.workspace, - metainfo=self.metainfo, - metainfo_index=self._METAINFO_INDEX, - ep_rank=self.ep_rank, + 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, - health=self.ep_group_health, on_timeout=alltoall_watchdog_on_timeout, ) @@ -265,74 +264,13 @@ def destroy(self) -> None: self._destroyed = True watchdog = getattr(self, "_alltoall_watchdog", None) if watchdog is not None: - watchdog.stop(timeout_s=1.0) + self._watchdog_coordinator.release_watchdog(watchdog) self._alltoall_watchdog = None def __del__(self) -> None: if not sys.is_finalizing(): self.destroy() - 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 _sync_watchdog_flag_generation(self) -> None: - workspace_state = self._WORKSPACE - assert workspace_state is not None - lock = workspace_state["watchdog_flag_generation_lock"] - with lock: - workspace_state["watchdog_flag_generation"] = max( - int(workspace_state["watchdog_flag_generation"]), - self._read_current_flag_val(), - ) - - def _next_watchdog_flag_generation(self) -> int: - workspace_state = self._WORKSPACE - assert workspace_state is not None - lock = workspace_state["watchdog_flag_generation_lock"] - with lock: - workspace_state["watchdog_flag_generation"] = ( - int(workspace_state["watchdog_flag_generation"]) + 1) - return int(workspace_state["watchdog_flag_generation"]) - - 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._alltoall_watchdog.watch( - phase=phase, - expected_flag=self._next_watchdog_flag_generation(), - active_mask=self._active_mask_int(active_rank_mask), - ) - def dispatch(self, token_selected_experts: torch.Tensor, input_payloads: list[torch.Tensor], @@ -366,7 +304,8 @@ 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) + active_rank_mask = self._watchdog_coordinator.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, @@ -380,7 +319,9 @@ def dispatch(self, eplb_local_stats, active_rank_mask, ) - self._watch_collective("dispatch", 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 @@ -428,13 +369,15 @@ 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) + active_rank_mask = self._watchdog_coordinator.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, active_rank_mask) - self._watch_collective("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/nvlink_one_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_one_sided.py index db5966615356..0811abc9b5d0 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,7 +25,6 @@ """ import os -import threading from typing import Callable, Dict, List, Optional, Tuple import torch @@ -35,7 +34,9 @@ DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, AlltoAllWatchdog, + AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, + EPGroupHealthLike, ) from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger @@ -158,11 +159,11 @@ def __init__( dtype: Optional[torch.dtype] = None, num_experts: Optional[int] = None, use_low_precision_combine: bool = False, - ep_group_health=None, + 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. @@ -288,8 +289,6 @@ def __init__( "mnnvl_mem": mnnvl_mem, "workspace": workspace, "metainfo": metainfo, - "watchdog_flag_generation": 0, - "watchdog_flag_generation_lock": threading.Lock(), } NVLinkOneSided._WORKSPACES[self._workspace_key] = workspace_state else: @@ -320,28 +319,27 @@ 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"] - if "watchdog_flag_generation_lock" not in workspace_state: - workspace_state["watchdog_flag_generation_lock"] = threading.Lock() - workspace_state["watchdog_flag_generation"] = self._read_current_flag_val() 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._sync_watchdog_flag_generation() - 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, + 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, - health=self.ep_group_health, on_timeout=alltoall_watchdog_on_timeout, ) @@ -351,57 +349,6 @@ 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 _sync_watchdog_flag_generation(self) -> None: - lock = self._workspace_state["watchdog_flag_generation_lock"] - with lock: - self._workspace_state["watchdog_flag_generation"] = max( - int(self._workspace_state["watchdog_flag_generation"]), - self._read_current_flag_val(), - ) - - def _next_watchdog_flag_generation(self) -> int: - lock = self._workspace_state["watchdog_flag_generation_lock"] - with lock: - self._workspace_state["watchdog_flag_generation"] = ( - int(self._workspace_state["watchdog_flag_generation"]) + 1 - ) - return int(self._workspace_state["watchdog_flag_generation"]) - - 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._alltoall_watchdog.watch( - phase=phase, - expected_flag=self._next_watchdog_flag_generation(), - active_mask=self._active_mask_int(active_rank_mask), - ) - @staticmethod def is_platform_supported() -> bool: """ @@ -422,7 +369,7 @@ def destroy(self): self._destroyed = True if self._alltoall_watchdog is not None: - self._alltoall_watchdog.stop(timeout_s=1.0) + self._watchdog_coordinator.release_watchdog(self._alltoall_watchdog) self._alltoall_watchdog = None workspace_key = getattr(self, "_workspace_key", None) if workspace_key is None: @@ -508,7 +455,9 @@ 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")) + active_rank_mask = self._watchdog_coordinator.active_rank_mask_tensor( + kwargs.get("active_rank_mask") + ) recv_buffers, combine_payload_offset, eplb_gathered_stats = ( torch.ops.trtllm.moe_a2a_dispatch( @@ -525,7 +474,9 @@ def dispatch( active_rank_mask, ) ) - self._watch_collective("dispatch", 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 @@ -628,7 +579,9 @@ 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")) + active_rank_mask = self._watchdog_coordinator.active_rank_mask_tensor( + kwargs.get("active_rank_mask") + ) output = torch.ops.trtllm.moe_a2a_combine( final_hidden_states, int(local_num_tokens), @@ -643,7 +596,9 @@ def combine( bool(self.use_low_precision_combine), active_rank_mask, ) - self._watch_collective("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/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index 5d8a7826c70e..0028f8a67685 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -27,6 +27,7 @@ DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, UNKNOWN_COMPLETION_FLAG, AlltoAllWatchdog, + AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, CompletionFlagReadTimeout, ) @@ -123,6 +124,33 @@ def test_watchdog_completes_when_flags_advance_past_expected_generation() -> Non 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) @@ -268,7 +296,7 @@ def test_watchdog_poll_timeout_without_snapshot_fails_closed() -> None: assert health.all_active() is True -def test_watchdog_poll_timeout_with_prior_snapshot_marks_known_missing_rank() -> None: +def test_watchdog_poll_timeout_with_prior_snapshot_does_not_mark_failed_rank() -> None: health = EPGroupHealth(3) events: list[AlltoAllWatchdogTimeout] = [] @@ -288,8 +316,33 @@ def test_watchdog_poll_timeout_with_prior_snapshot_marks_known_missing_rank() -> assert event.poll_timed_out is True assert event.observed_flags == (1, 0, 1) assert event.missing_ranks == (1,) - assert event.marked_failed_ranks == (1,) - assert health.get_failed_ranks() == frozenset({1}) + assert event.marked_failed_ranks == () + 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: @@ -359,6 +412,125 @@ def test_watchdog_from_workspace_reads_phase_specific_offsets() -> None: assert health.get_failed_ranks() == frozenset({0}) +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({1}) + 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( From 05db03d105d6e0af6f23b5121d0c58379b4aa1c8 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 30 Jun 2026 12:43:40 -0700 Subject: [PATCH 7/8] fix: make AlltoAll watchdog detection-only Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/alltoall_watchdog.py | 72 ++++++++++++------- .../_torch/distributed/moe_alltoall.py | 16 +++-- .../communication/nvlink_one_sided.py | 15 ++-- .../modules/fused_moe/ep_group_health.py | 13 ++-- .../_torch/modules/fused_moe/wide_ep_ft.py | 15 ++-- .../_torch/modules/test_alltoall_watchdog.py | 62 ++++++++++++---- 6 files changed, 135 insertions(+), 58 deletions(-) diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py index 1cb3fe88aae7..166d4507e926 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -19,6 +19,12 @@ 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 @@ -65,7 +71,7 @@ def read_completion_flags(self, phase: str) -> Sequence[int]: class EPGroupHealthLike(Protocol): - """Subset of EPGroupHealth used by the watchdog.""" + """Read-only committed EP membership used by AlltoAll frontends.""" def get_mask(self) -> int: """Return the active-rank bitmask.""" @@ -73,9 +79,6 @@ def get_mask(self) -> int: def get_mask_words(self) -> tuple[int, ...]: """Return the active-rank bitmask split into uint64 words.""" - def mark_failed(self, rank: int) -> bool: - """Mark ``rank`` failed and return whether state changed.""" - class CompletionFlagReadTimeout(TimeoutError): """Raised when the host watchdog cannot read completion flags in time.""" @@ -83,13 +86,17 @@ class CompletionFlagReadTimeout(TimeoutError): @dataclass(frozen=True) class AlltoAllWatchdogTimeout: - """Details emitted when an AlltoAll phase times out.""" + """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, ...] - marked_failed_ranks: tuple[int, ...] elapsed_s: float poll_timed_out: bool = False @@ -206,7 +213,11 @@ def read_completion_flags(self, phase: str) -> tuple[int, ...]: class AlltoAllWatchdogCoordinator: - """Shared watchdog plumbing for MoE AlltoAll frontends.""" + """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, @@ -256,12 +267,38 @@ def read_current_flag_val(self) -> int: return _normalize_completion_flag(int(flag_val.item())) def active_rank_mask_tensor(self, active_rank_mask: torch.Tensor | None) -> torch.Tensor | None: + """Capture a caller override or the current committed mask. + + The returned tensor does not alias a caller-owned override, so it can + safely represent the mask used by one dispatch/combine pair. + """ if active_rank_mask is not None: - return active_rank_mask + return active_rank_mask.detach().clone() if self._health is None: return None return torch.tensor(self._health.get_mask_words(), dtype=torch.uint64, device="cpu") + def active_rank_mask_for_combine( + self, + dispatch_active_rank_mask: torch.Tensor | None, + requested_active_rank_mask: torch.Tensor | None, + ) -> torch.Tensor | None: + """Reuse the dispatch mask and reject a conflicting combine override. + + This guarantees one local mask snapshot across a dispatch/combine pair. + Atomic membership across layers and iterations remains the recovery + coordinator's responsibility. + """ + 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() @@ -372,7 +409,8 @@ class AlltoAllWatchdog: 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. + hide a still-spinning dispatch. A timeout is reported through + ``on_timeout`` without mutating the committed EP membership. """ VALID_PHASES = frozenset({"dispatch", "combine"}) @@ -564,24 +602,11 @@ def _handle_timeout( ) -> 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 - and not poll_timed_out - and UNKNOWN_COMPLETION_FLAG not in observed_flags - ): - 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, poll_timed_out=poll_timed_out, ) @@ -589,14 +614,13 @@ def _handle_timeout( 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, marked ranks %s", + "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), - list(marked_failed), ) else: tllm_logger.warning( diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index 6e44510a8573..f1692d2473fc 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -31,6 +31,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: torch.Tensor | None = None class MoeAlltoAll: @@ -149,8 +150,8 @@ 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. + 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. @@ -289,7 +290,8 @@ 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. + active_rank_mask: Optional uint64 CPU tensor overriding committed membership for this dispatch. The + captured value is reused by combine. Returns: recv_tensors: List of tensors received, each has shape [ep_size, max_tokens_per_rank, payload_num_elements_per_token] @@ -329,6 +331,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 = active_rank_mask self._state.phase = "dispatched" if invalid_token_expert_id is not None: @@ -361,7 +364,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 overriding ep_group_health for this combine. + active_rank_mask: Optional uint64 CPU tensor. If supplied, it must match the mask captured by dispatch + for this collective. Returns: combined_output: [local_num_tokens, num_elements_per_token] tensor of combined results @@ -369,8 +373,8 @@ 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._watchdog_coordinator.active_rank_mask_tensor( - active_rank_mask) + active_rank_mask = self._watchdog_coordinator.active_rank_mask_for_combine( + self._state.active_rank_mask, 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, 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 0811abc9b5d0..076e1f5e6b81 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 @@ -181,8 +181,8 @@ 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. + 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. @@ -424,7 +424,8 @@ 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; the captured value is reused by combine. Returns: Tuple of (hidden_states, hidden_states_sf, token_selected_slots, token_final_scales) @@ -483,6 +484,7 @@ def dispatch( 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"] = active_rank_mask self._dispatch_state["phase"] = "dispatched" # Extract results from recv_buffers @@ -545,6 +547,8 @@ 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. Returns: Combined output tensor [local_num_tokens, hidden_size] @@ -579,8 +583,9 @@ def combine( raise ValueError( f"final_hidden_states must be 2D or 3D, got {final_hidden_states.dim()}D" ) - active_rank_mask = self._watchdog_coordinator.active_rank_mask_tensor( - kwargs.get("active_rank_mask") + active_rank_mask = self._watchdog_coordinator.active_rank_mask_for_combine( + self._dispatch_state.get("active_rank_mask"), + kwargs.get("active_rank_mask"), ) output = torch.ops.trtllm.moe_a2a_combine( final_hidden_states, 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..b1aba97c818c 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py +++ b/tensorrt_llm/_torch/modules/fused_moe/ep_group_health.py @@ -16,17 +16,18 @@ """EP group health tracking 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 the committed active membership of an Expert Parallel +(EP) group. It is the single source of truth for the data-plane rank mask within +one process and 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 +Failure detection does not mutate this object directly. Cross-process consensus, +expert-placement preparation, and atomic publication of a new committed mask are +the responsibility of higher-layer recovery coordination components and are not performed here. """ diff --git a/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py b/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py index 69e5c53f061e..fb1586ae1054 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py +++ b/tensorrt_llm/_torch/modules/fused_moe/wide_ep_ft.py @@ -26,17 +26,19 @@ from .ep_group_health import EPGroupHealth -_ENABLE_ENV = "TRTLLM_ENABLE_WIDE_EP_FT" +_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, "0").lower() in {"1", "true", "yes", "on"} + return os.environ.get(_ENABLE_ENV) == "1" def _float_option(extra_attrs: dict, key: str, env_name: str, default: float) -> float: @@ -50,12 +52,15 @@ def _float_option(extra_attrs: dict, key: str, env_name: str, default: float) -> def get_wide_ep_ft_options( model_config: Any, ) -> tuple[Optional[EPGroupHealth], Optional[float], float]: - """Return the shared EP health object and watchdog timing for a model. + """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 ``TRTLLM_ENABLE_WIDE_EP_FT=1`` to create - one process-local health object shared by all MoE communication layers. + ``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", {}) diff --git a/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index 0028f8a67685..931b0dd5e5a6 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -180,7 +180,7 @@ def test_watchdog_stop_is_terminal() -> None: def test_wide_ep_ft_options_create_shared_health_when_enabled( monkeypatch: pytest.MonkeyPatch, ) -> None: - monkeypatch.setenv("TRTLLM_ENABLE_WIDE_EP_FT", "1") + monkeypatch.setenv("TLLM_FAULT_TOLERANCE_MODE", "1") model_config = SimpleNamespace( extra_attrs={}, mapping=SimpleNamespace(moe_ep_size=4), @@ -197,8 +197,50 @@ def test_wide_ep_ft_options_create_shared_health_when_enabled( assert poll_again_s == poll_interval_s -def test_watchdog_timeout_reports_and_marks_missing_remote_ranks() -> None: +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_dispatch_mask_for_combine() -> 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_mask = coordinator.active_rank_mask_tensor(None) + assert dispatch_mask is not None + + # Simulate a higher-layer commit at an invalid mid-collective point. The + # local dispatch/combine pair must keep using its dispatch snapshot. + health.mark_failed(2) + combine_mask = coordinator.active_rank_mask_for_combine(dispatch_mask, None) + + assert combine_mask is dispatch_mask + assert combine_mask.tolist() == [0b1111, 0] + with pytest.raises(ValueError, match="mask captured at dispatch"): + coordinator.active_rank_mask_for_combine( + dispatch_mask, + 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] = [] @@ -221,8 +263,8 @@ def test_watchdog_timeout_reports_and_marks_missing_remote_ranks() -> None: 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}) + assert not hasattr(event, "marked_failed_ranks") + assert health.snapshot() == committed_before def test_watchdog_ignores_ranks_already_failed_in_health_mask() -> None: @@ -248,7 +290,7 @@ def test_watchdog_ignores_ranks_already_failed_in_health_mask() -> None: assert health.get_failed_ranks() == frozenset({2}) -def test_watchdog_reports_local_missing_but_does_not_mark_local_failed() -> None: +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]) @@ -268,7 +310,6 @@ def test_watchdog_reports_local_missing_but_does_not_mark_local_failed() -> None event = events[0] assert event.missing_ranks == (0,) - assert event.marked_failed_ranks == () assert health.get_failed_ranks() == frozenset() @@ -292,7 +333,6 @@ def test_watchdog_poll_timeout_without_snapshot_fails_closed() -> None: assert event.poll_timed_out is True assert event.observed_flags == (UNKNOWN_COMPLETION_FLAG,) * 3 assert event.missing_ranks == (0, 1, 2) - assert event.marked_failed_ranks == () assert health.all_active() is True @@ -316,7 +356,6 @@ def test_watchdog_poll_timeout_with_prior_snapshot_does_not_mark_failed_rank() - assert event.poll_timed_out is True assert event.observed_flags == (1, 0, 1) assert event.missing_ranks == (1,) - assert event.marked_failed_ranks == () assert health.all_active() is True @@ -370,7 +409,7 @@ def test_watchdog_preserves_fifo_order_and_clears_followups_after_timeout() -> N assert len(events) == 1 assert events[0].phase == "dispatch" assert events[0].missing_ranks == (1,) - assert health.get_failed_ranks() == frozenset({1}) + assert health.get_failed_ranks() == frozenset() def test_watchdog_from_workspace_reads_phase_specific_offsets() -> None: @@ -408,8 +447,7 @@ def test_watchdog_from_workspace_reads_phase_specific_offsets() -> None: 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}) + assert health.get_failed_ranks() == frozenset() def test_workspace_coordinators_share_fifo_watchdog() -> None: @@ -459,7 +497,7 @@ def test_workspace_coordinators_share_fifo_watchdog() -> None: assert len(events) == 1 assert events[0].phase == "dispatch" assert events[0].missing_ranks == (1,) - assert health.get_failed_ranks() == frozenset({1}) + assert health.get_failed_ranks() == frozenset() finally: for coordinator, watchdog in zip(coordinators, watchdogs): coordinator.release_watchdog(watchdog) From 7e8ed418447f70ac4fb8b86a1b4a46f5f29f3604 Mon Sep 17 00:00:00 2001 From: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> Date: Tue, 30 Jun 2026 13:27:57 -0700 Subject: [PATCH 8/8] fix: bind AlltoAll mask to committed generation Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com> --- tensorrt_llm/_torch/alltoall_watchdog.py | 99 ++++++++++++++++--- .../_torch/distributed/moe_alltoall.py | 23 +++-- .../communication/nvlink_one_sided.py | 17 +++- .../modules/fused_moe/ep_group_health.py | 14 ++- .../_torch/modules/test_alltoall_watchdog.py | 79 +++++++++++++-- 5 files changed, 190 insertions(+), 42 deletions(-) diff --git a/tensorrt_llm/_torch/alltoall_watchdog.py b/tensorrt_llm/_torch/alltoall_watchdog.py index 166d4507e926..2b71aaefa257 100644 --- a/tensorrt_llm/_torch/alltoall_watchdog.py +++ b/tensorrt_llm/_torch/alltoall_watchdog.py @@ -46,6 +46,9 @@ 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() @@ -70,6 +73,18 @@ 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.""" @@ -79,6 +94,22 @@ def get_mask(self) -> int: 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.""" @@ -266,29 +297,75 @@ def read_current_flag_val(self) -> int: flag_val = flag_val.detach().cpu() return _normalize_completion_flag(int(flag_val.item())) - def active_rank_mask_tensor(self, active_rank_mask: torch.Tensor | None) -> torch.Tensor | None: - """Capture a caller override or the current committed mask. + 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. + 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 active_rank_mask.detach().clone() + return ActiveRankMaskSnapshot( + active_rank_mask=active_rank_mask.detach().clone(), + committed_generation=None, + ) if self._health is None: - return None - return torch.tensor(self._health.get_mask_words(), dtype=torch.uint64, device="cpu") + 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_active_rank_mask: torch.Tensor | None, + dispatch_snapshot: ActiveRankMaskSnapshot, requested_active_rank_mask: torch.Tensor | None, ) -> torch.Tensor | None: - """Reuse the dispatch mask and reject a conflicting combine override. + """Validate the dispatch epoch and return its captured rank mask. - This guarantees one local mask snapshot across a dispatch/combine pair. - Atomic membership across layers and iterations remains the recovery - coordinator's responsibility. + 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) diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index f1692d2473fc..2bf31ef680a4 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -17,8 +17,9 @@ 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, AlltoAllWatchdog, - AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, EPGroupHealthLike) + 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 @@ -31,7 +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: torch.Tensor | None = None + active_rank_mask_snapshot: ActiveRankMaskSnapshot | None = None class MoeAlltoAll: @@ -290,8 +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. The - captured value is reused by combine. + 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] @@ -306,8 +308,9 @@ def dispatch(self, 0 ) == self.eplb_stats_num_experts, "eplb_local_stats size must match eplb_stats_num_experts" - active_rank_mask = self._watchdog_coordinator.active_rank_mask_tensor( + 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, @@ -331,7 +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 = active_rank_mask + self._state.active_rank_mask_snapshot = active_rank_mask_snapshot self._state.phase = "dispatched" if invalid_token_expert_id is not None: @@ -365,7 +368,7 @@ def combine( 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. + 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 @@ -373,8 +376,10 @@ 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( - self._state.active_rank_mask, active_rank_mask) + 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, 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 076e1f5e6b81..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 @@ -33,6 +33,7 @@ from tensorrt_llm._torch.alltoall_watchdog import ( DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, + ActiveRankMaskSnapshot, AlltoAllWatchdog, AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, @@ -425,7 +426,8 @@ def dispatch( all_rank_num_tokens: Token counts per rank [ep_size] use_dp_padding: Whether to use DP padding (optional) **kwargs: Strategy-specific arguments. ``active_rank_mask`` may override the committed membership - for dispatch; the captured value is reused by combine. + 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) @@ -456,9 +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 = self._watchdog_coordinator.active_rank_mask_tensor( + 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( @@ -484,7 +487,7 @@ def dispatch( 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"] = active_rank_mask + self._dispatch_state["active_rank_mask_snapshot"] = active_rank_mask_snapshot self._dispatch_state["phase"] = "dispatched" # Extract results from recv_buffers @@ -548,7 +551,8 @@ def combine( 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. + 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] @@ -583,8 +587,11 @@ 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( - self._dispatch_state.get("active_rank_mask"), + active_rank_mask_snapshot, kwargs.get("active_rank_mask"), ) output = torch.ops.trtllm.moe_a2a_combine( 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 b1aba97c818c..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,22 +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 the committed active membership of an Expert Parallel -(EP) group. It is the single source of truth for the data-plane rank mask 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 (read-only expected-peer snapshot) * The MoE load balancer (emergency-mask reconfiguration) * The model engine and PyExecutor (degraded health reporting) -Failure detection does not mutate this object directly. Cross-process consensus, -expert-placement preparation, and atomic publication of a new committed mask are -the responsibility of higher-layer recovery coordination components and are 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/tests/unittest/_torch/modules/test_alltoall_watchdog.py b/tests/unittest/_torch/modules/test_alltoall_watchdog.py index 931b0dd5e5a6..cdff07cbcea4 100644 --- a/tests/unittest/_torch/modules/test_alltoall_watchdog.py +++ b/tests/unittest/_torch/modules/test_alltoall_watchdog.py @@ -211,7 +211,7 @@ def test_wide_ep_ft_options_ignore_legacy_enable_flag(monkeypatch: pytest.Monkey assert timeout_s is None -def test_watchdog_coordinator_reuses_dispatch_mask_for_combine() -> None: +def test_watchdog_coordinator_reuses_committed_mask_when_generation_is_unchanged() -> None: health = EPGroupHealth(4) coordinator = AlltoAllWatchdogCoordinator( workspace_state={}, @@ -221,19 +221,80 @@ def test_watchdog_coordinator_reuses_dispatch_mask_for_combine() -> None: ep_rank=0, health=health, ) - dispatch_mask = coordinator.active_rank_mask_tensor(None) - assert dispatch_mask is not None + dispatch_snapshot = coordinator.capture_active_rank_mask(None) + assert dispatch_snapshot.active_rank_mask is not None + assert dispatch_snapshot.committed_generation == 0 - # Simulate a higher-layer commit at an invalid mid-collective point. The - # local dispatch/combine pair must keep using its dispatch snapshot. - health.mark_failed(2) - combine_mask = coordinator.active_rank_mask_for_combine(dispatch_mask, None) + combine_mask = coordinator.active_rank_mask_for_combine(dispatch_snapshot, None) - assert combine_mask is dispatch_mask + 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_mask, + dispatch_snapshot, torch.tensor(health.get_mask_words(), dtype=torch.uint64), )