From adedf216ead334e448d1f88e80d0219ba72755c0 Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Mon, 20 Jul 2026 11:40:20 -0400 Subject: [PATCH 01/14] [None][feat] preserve native MoE A2A virtual addrs across restore Signed-off-by: Hannah Zhang --- tensorrt_llm/_mnnvl_utils.py | 373 +++++++++++------- .../_torch/distributed/moe_alltoall.py | 58 +++ .../communication/nvlink_one_sided.py | 57 +++ .../communication/nvlink_two_sided.py | 17 +- 4 files changed, 359 insertions(+), 146 deletions(-) diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index d116f8ae2925..a48d1e04abbc 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/tensorrt_llm/_mnnvl_utils.py @@ -18,7 +18,7 @@ import platform import sys from dataclasses import dataclass -from typing import List, Optional, Union +from typing import Any, List, Optional, Union import pynvml import torch @@ -51,6 +51,19 @@ def _check_cu_result(cu_func_ret): return None +@dataclass +class _MnnvlAllocationRecord: + comm: Any + comm_size: int + comm_rank: int + aligned_size: int + mem_handles: List[Any] + start_address: int + rank_stride: int + address_offset: int + mapped: bool = True + + class MnnvlMemory: """MNNVL memory management for tensor parallel (TP) operations.""" @@ -101,6 +114,11 @@ def __del__(self): if hasattr(self, "ptr"): type(self).close_mnnvl_memory(self.ptr) + @property + def mapped(self) -> bool: + """Whether physical handles are mapped into this VA reservation.""" + return type(self).allocated_map[self.ptr].mapped + def as_torch_strided_tensor(self, dtype): num_segments = type(self).comm.Get_size() return pack_strided_memory( @@ -110,7 +128,7 @@ def as_torch_strided_tensor(self, dtype): @property def local_mem_handle(self) -> int: """Return the local rank's CUmemGenericAllocationHandle.""" - _, _, mem_handles, _, _, _ = type(self).allocated_map[self.ptr] + mem_handles = type(self).allocated_map[self.ptr].mem_handles comm_rank = type(self).comm.Get_rank() return int(mem_handles[comm_rank]) @@ -195,6 +213,105 @@ def new_mnnvl_memory_address(cls, mapping: Mapping, size: int): cls.current_rank_stride = current_rank_stride cls.current_mem_offset = 0 + @classmethod + def _create_and_map_handles( + cls, + comm, + aligned_size: int, + start_address: int, + rank_stride: int, + address_offset: int, + ) -> List[Any]: + dev_id = int(_check_cu_result(cuda.cuCtxGetDevice())) + assert dev_id == MnnvlMemory.dev_id, ( + f"Different dev_id found dev_id={dev_id} but MnnvlMemory.dev_id={MnnvlMemory.dev_id}" + ) + allocation_prop = MnnvlMemory.get_allocation_prop(dev_id) + local_handle = _check_cu_result(cuda.cuMemCreate(aligned_size, allocation_prop, flags=0)) + exported_handle = None + pidfds = [] + remote_fds = [] + mem_handles = [None] * comm.Get_size() + mapped_rank_ptrs = [] + is_fabric = ( + allocation_prop.requestedHandleTypes + == cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC + ) + try: + exported_handle = _check_cu_result( + cuda.cuMemExportToShareableHandle( + local_handle, allocation_prop.requestedHandleTypes, 0 + ) + ) + if is_fabric: + all_handles_data = comm.allgather(exported_handle.data) + else: + all_exported_fds = comm.allgather(int(exported_handle)) + all_pids = comm.allgather(os.getpid()) + syscall = ctypes.CDLL(None, use_errno=True).syscall + for pid in all_pids: + pidfd = syscall(434, pid, 0) + if pidfd < 0: + err = ctypes.get_errno() + raise RuntimeError( + f"pidfd_open({pid}) failed with errno {err}: {os.strerror(err)}" + ) + pidfds.append(pidfd) + for pidfd, fd in zip(pidfds, all_exported_fds): + remote_fd = syscall(438, pidfd, fd, 0) + if remote_fd < 0: + err = ctypes.get_errno() + raise RuntimeError( + f"pidfd_getfd(pidfd={pidfd}, fd={fd}) failed with errno " + f"{err}: {os.strerror(err)}" + ) + remote_fds.append(remote_fd) + comm.barrier() + all_handles_data = remote_fds + + access_desc = cuda.CUmemAccessDesc() + access_desc.location = allocation_prop.location + access_desc.flags = cuda.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + for rank, handle_data in enumerate(all_handles_data): + rank_ptr = start_address + rank_stride * rank + address_offset + handle = ( + local_handle + if rank == comm.Get_rank() + else _check_cu_result( + cuda.cuMemImportFromShareableHandle( + handle_data, allocation_prop.requestedHandleTypes + ) + ) + ) + mem_handles[rank] = handle + _check_cu_result(cuda.cuMemMap(rank_ptr, aligned_size, 0, handle, 0)) + mapped_rank_ptrs.append(rank_ptr) + _check_cu_result(cuda.cuMemSetAccess(rank_ptr, aligned_size, [access_desc], 1)) + return mem_handles + except Exception: + for rank_ptr in reversed(mapped_rank_ptrs): + try: + _check_cu_result(cuda.cuMemUnmap(rank_ptr, aligned_size)) + except RuntimeError as error: + logger.warning("Failed to unmap incomplete MNNVL allocation: %s", error) + + handles_to_release = [handle for handle in mem_handles if handle is not None] + if mem_handles[comm.Get_rank()] is None: + handles_to_release.append(local_handle) + for handle in handles_to_release: + try: + _check_cu_result(cuda.cuMemRelease(handle)) + except RuntimeError as error: + logger.warning("Failed to release incomplete MNNVL allocation: %s", error) + raise + finally: + for pidfd in pidfds: + os.close(pidfd) + for remote_fd in remote_fds: + os.close(remote_fd) + if not is_fabric and exported_handle is not None: + os.close(int(exported_handle)) + @classmethod def open_mnnvl_memory(cls, mapping: Mapping, size: int): # Ensure MnnvlMemory is initialized (for dev_id and allocation_granularity) @@ -227,133 +344,23 @@ def open_mnnvl_memory(cls, mapping: Mapping, size: int): assert cls.current_mem_offset + aligned_size <= cls.current_rank_stride - allocation_prop = cls.get_allocation_prop(dev_id) - allocated_mem_handle = _check_cu_result( - cuda.cuMemCreate(aligned_size, allocation_prop, flags=0) - ) - exported_fabric_handle = _check_cu_result( - cuda.cuMemExportToShareableHandle( - allocated_mem_handle, allocation_prop.requestedHandleTypes, 0 - ) - ) - pidfds = [] - remote_fds = [] - try: - if ( - allocation_prop.requestedHandleTypes - == cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC - ): - all_handles_data = comm.allgather(exported_fabric_handle.data) - else: - all_handles_data = comm.allgather(exported_fabric_handle) - all_pids = comm.allgather(os.getpid()) - libc = ctypes.CDLL(None, use_errno=True) - syscall = libc.syscall - SYS_pidfd_open = 434 - SYS_pidfd_getfd = 438 - for i, pid in enumerate(all_pids): - pidfd = syscall(SYS_pidfd_open, pid, 0) - if pidfd < 0: - err = ctypes.get_errno() - raise RuntimeError( - f"pidfd_open({pid}) failed with errno {err}: {os.strerror(err)}" - ) - pidfds.append(pidfd) - - for i, (pidfd, fd) in enumerate(zip(pidfds, all_handles_data)): - remote_fd = syscall(SYS_pidfd_getfd, pidfd, fd, 0) - if remote_fd < 0: - err = ctypes.get_errno() - error_msg = f"pidfd_getfd(pidfd={pidfd}, fd={fd}) failed with errno {err}: {os.strerror(err)}." - if err == 1: # EPERM - error_msg += ( - " Permission denied. If running in a container, try adding --cap-add=SYS_PTRACE " - "to your docker run command." - ) - else: - error_msg += " This may be due to kernel version (requires Linux 5.6+)." - raise RuntimeError(error_msg) - remote_fds.append(remote_fd) - - all_handles_data = remote_fds - except Exception: - # Release resources on failure path to avoid leaks; then re-raise. - if isinstance(exported_fabric_handle, int): - try: - os.close(exported_fabric_handle) - except OSError as e: - logger.warning( - "Failed to close exported shareable handle on error: %s", - e, - ) - try: - _check_cu_result(cuda.cuMemRelease(allocated_mem_handle)) - except RuntimeError as e: - logger.warning( - "cuMemRelease failed during error cleanup (original error will be raised): %s", - e, - ) - for _pidfd in pidfds: - try: - os.close(_pidfd) - except OSError: - pass - for _rfd in remote_fds: - try: - os.close(_rfd) - except OSError: - pass - raise - # all_handles_data like b'\x00\x00\x00 \x00\x00\x00\x00\x8f\xec\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\t\x00\x00\x00\x00\x00\x1d\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00' # noqa: E501 - # can use buf = memoryview(data) to import if using plain buffer for data. - - madesc = cuda.CUmemAccessDesc() - madesc.location = allocation_prop.location - madesc.flags = cuda.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE - - mem_handles = [None] * comm_size - mem_handles[comm_rank] = allocated_mem_handle - mapped_rank_ptrs = [] - try: - for i, remote_handle_data in enumerate(all_handles_data): - rank_ptr = ( - cls.current_start_address + cls.current_rank_stride * i + cls.current_mem_offset - ) - if i != comm_rank: - # Fabric memory mapping - mem_handles[i] = _check_cu_result( - cuda.cuMemImportFromShareableHandle( - remote_handle_data, allocation_prop.requestedHandleTypes - ) - ) - - _check_cu_result(cuda.cuMemMap(rank_ptr, aligned_size, 0, mem_handles[i], 0)) - mapped_rank_ptrs.append(rank_ptr) - _check_cu_result(cuda.cuMemSetAccess(rank_ptr, aligned_size, [madesc], 1)) + mem_handles = cls._create_and_map_handles( + comm, + aligned_size, + cls.current_start_address, + cls.current_rank_stride, + cls.current_mem_offset, + ) except Exception: - # Clean up partially imported and mapped memory before propagating - # the failure to the caller. - for rank_ptr in reversed(mapped_rank_ptrs): - try: - _check_cu_result(cuda.cuMemUnmap(rank_ptr, aligned_size)) - except RuntimeError as e: - logger.warning("cuMemUnmap failed during error cleanup: %s", e) - for mem_handle in mem_handles: - if mem_handle is None: - continue - try: - _check_cu_result(cuda.cuMemRelease(mem_handle)) - except RuntimeError as e: - logger.warning("cuMemRelease failed during error cleanup: %s", e) if reserved_new_address: try: device_ptr = cuda.CUdeviceptr(cls.current_start_address) _check_cu_result( cuda.cuMemAddressFree(device_ptr, comm_size * cls.current_rank_stride) ) - except RuntimeError as e: - logger.warning("cuMemAddressFree failed during error cleanup: %s", e) + except RuntimeError as error: + logger.warning("cuMemAddressFree failed during error cleanup: %s", error) else: ( cls.current_start_address, @@ -361,16 +368,18 @@ def open_mnnvl_memory(cls, mapping: Mapping, size: int): cls.current_mem_offset, ) = previous_address_state raise - + # all_handles_data like b'\x00\x00\x00 \x00\x00\x00\x00\x8f\xec\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\t\x00\x00\x00\x00\x00\x1d\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00' # noqa: E501 ptr = cls.current_start_address + cls.current_mem_offset stride = cls.current_rank_stride - cls.allocated_map[ptr] = ( - mapping, - aligned_size, - mem_handles, - cls.current_start_address, - cls.current_rank_stride, - cls.current_mem_offset, + cls.allocated_map[ptr] = _MnnvlAllocationRecord( + comm=comm, + comm_size=comm_size, + comm_rank=comm_rank, + aligned_size=aligned_size, + mem_handles=mem_handles, + start_address=cls.current_start_address, + rank_stride=cls.current_rank_stride, + address_offset=cls.current_mem_offset, ) cls.address_refcnt[cls.current_start_address] = ( cls.address_refcnt.get(cls.current_start_address, 0) + 1 @@ -381,26 +390,69 @@ def open_mnnvl_memory(cls, mapping: Mapping, size: int): @classmethod def close_mnnvl_memory(cls, ptr: int): - mapping, aligned_size, mem_handles, start_address, rank_stride, address_offset = ( - cls.allocated_map.pop(ptr) - ) - comm = cls.get_comm(mapping) - comm_size = comm.Get_size() - for i in range(comm_size): - rank_ptr = start_address + i * rank_stride + address_offset - _check_cu_result(cuda.cuMemUnmap(rank_ptr, aligned_size)) - _check_cu_result(cuda.cuMemRelease(mem_handles[i])) - cls.address_refcnt[start_address] -= 1 - - if cls.address_refcnt[start_address] == 0: - cls.address_refcnt.pop(start_address) - device_ptr = cuda.CUdeviceptr(start_address) - _check_cu_result(cuda.cuMemAddressFree(device_ptr, comm_size * rank_stride)) - if start_address == cls.current_start_address: + record = cls.allocated_map.pop(ptr) + if record.mapped: + cls._unmap_and_release_handles(record) + cls.address_refcnt[record.start_address] -= 1 + + if cls.address_refcnt[record.start_address] == 0: + cls.address_refcnt.pop(record.start_address) + device_ptr = cuda.CUdeviceptr(record.start_address) + _check_cu_result( + cuda.cuMemAddressFree(device_ptr, record.comm_size * record.rank_stride) + ) + if record.start_address == cls.current_start_address: cls.current_start_address = 0 cls.current_rank_stride = 0 cls.current_mem_offset = 0 + @classmethod + def _unmap_and_release_handles(cls, record: _MnnvlAllocationRecord) -> None: + for rank in range(record.comm_size): + rank_ptr = record.start_address + rank * record.rank_stride + record.address_offset + _check_cu_result(cuda.cuMemUnmap(rank_ptr, record.aligned_size)) + _check_cu_result(cuda.cuMemRelease(record.mem_handles[rank])) + + def checkpoint_prepare(self) -> None: + """Collectively detach backing handles while retaining graph-visible VA.""" + cls = type(self) + record = cls.allocated_map[self.ptr] + if not record.mapped: + return + torch.cuda.synchronize() + record.comm.barrier() + cls._unmap_and_release_handles(record) + record.mem_handles = [None] * record.comm_size + record.mapped = False + record.comm.barrier() + + def checkpoint_restore(self, comm) -> None: + """Collectively remap fresh handles at the original virtual addresses.""" + cls = type(self) + record = cls.allocated_map[self.ptr] + if record.mapped: + return + comm_size = comm.Get_size() + comm_rank = comm.Get_rank() + if comm_size != record.comm_size or comm_rank != record.comm_rank: + raise RuntimeError( + "Cannot restore MNNVL memory with a communicator that differs from " + "the graph-visible allocation layout: " + f"rank/size {comm_rank}/{comm_size} != " + f"{record.comm_rank}/{record.comm_size}" + ) + torch.cuda.synchronize() + record.mem_handles = cls._create_and_map_handles( + comm, + record.aligned_size, + record.start_address, + record.rank_stride, + record.address_offset, + ) + record.comm = comm + cls.comm = comm + record.mapped = True + @staticmethod @functools.cache def support_nvlink(dev_id: int, need_all_up: bool = True): @@ -596,6 +648,37 @@ def get_moe_prepare_workspace(mapping: Mapping): ) return MnnvlMoe.moe_prepare_workspace_tensor + @staticmethod + def checkpoint_prepare() -> None: + """Detach TRT-native two-sided MoE workspaces for checkpointing.""" + for workspace in (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace): + if workspace is not None: + workspace.checkpoint_prepare() + + @staticmethod + def checkpoint_restore(comm) -> None: + """Restore TRT-native two-sided MoE workspaces at their original virtual addresses.""" + workspaces = (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) + for workspace in workspaces: + if workspace is not None: + workspace.checkpoint_restore(comm) + if MnnvlMoe.moe_workspace_tensor is not None: + assert MnnvlMoe.moe_mapping is not None + torch.ops.trtllm.moe_initialize_workspace( + MnnvlMoe.moe_workspace_tensor, + MnnvlMoe.moe_mapping.moe_ep_rank, + MnnvlMoe.moe_mapping.moe_ep_size, + ) + torch.cuda.synchronize() + comm.barrier() + + @staticmethod + def require_mapped() -> None: + """Reject kernel access while either native MoE workspace is detached.""" + for workspace in (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace): + if workspace is not None and not workspace.mapped: + raise RuntimeError("Native MoE All-to-All workspace handles are unmapped") + @staticmethod def compute_target_rank_id( token_selected_experts: torch.Tensor, expert_count: int, ep_size: int diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index 8d554921ef37..9219acdcd08a 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -446,6 +446,61 @@ def cft_initialize(self) -> None: f"CFT LE initialized (workspace-bound): ep_rank={self.ep_rank}, ep_size={self.ep_size}" ) + def _require_mapped(self) -> None: + if not self.mnnvl_mem.mapped: + raise RuntimeError( + "Native MoE All-to-All workspace handles are unmapped") + + def checkpoint_prepare(self) -> None: + """Collectively detach handles after the caller globally quiesces all owners. + + The local phase and watchdog checks are defense-in-depth only; they do + not prove that every wrapper sharing this workspace is quiescent. + """ + if self._state.phase != "idle": + raise RuntimeError( + "Cannot checkpoint during an active MoE All-to-All phase") + if self._alltoall_watchdog is not None: + raise RuntimeError( + "Checkpointing native MoE All-to-All with the watchdog enabled " + "is not supported") + self.mnnvl_mem.checkpoint_prepare() + + def checkpoint_restore(self, comm) -> None: + """Collectively restore handles after caller-proven global quiescence. + + The low-level remap is idempotent, so every shared owner may call this + method to reset its own protocol state. + """ + self.mnnvl_mem.checkpoint_restore(comm) + refreshed_metainfo = torch.ops.trtllm.moe_a2a_initialize( + self.workspace, + self.ep_rank, + self.ep_size, + self.max_num_tokens, + self.eplb_stats_num_experts, + ) + if not torch.equal(refreshed_metainfo, self.metainfo): + raise RuntimeError( + "MoE All-to-All metainfo changed during MNNVL restore; " + "captured CUDA graphs cannot be replayed safely") + self.metainfo = refreshed_metainfo + assert self._WORKSPACE is not None + self._WORKSPACE["metainfo"] = refreshed_metainfo + metainfo_index = self._METAINFO_INDEX + assert metainfo_index is not None + self._watchdog_coordinator = AlltoAllWatchdogCoordinator( + workspace_state=self._WORKSPACE, + workspace=self.workspace, + metainfo=refreshed_metainfo, + metainfo_index=metainfo_index, + ep_rank=self.ep_rank, + health=self.ep_group_health, + ) + torch.cuda.synchronize() + comm.barrier() + self.reset_state() + def dispatch(self, token_selected_experts: torch.Tensor, input_payloads: list[torch.Tensor], @@ -472,6 +527,7 @@ def dispatch(self, Returns: recv_tensors: List of tensors received, each has shape [ep_size, max_tokens_per_rank, payload_num_elements_per_token] """ + self._require_mapped() assert self._state.phase == "idle", "dispatch called twice without an intervening combine" reject_rank_mask_cuda_graph_capture(self._rank_mask_enabled) assert runtime_max_tokens_per_rank <= self.max_num_tokens, "runtime_max_tokens_per_rank must not exceed max_num_tokens" @@ -570,6 +626,7 @@ def combine( Returns: combined_output: [local_num_tokens, num_elements_per_token] tensor of combined results """ + self._require_mapped() assert self._state.phase == "dispatched", "combine called before a successful dispatch" reject_rank_mask_cuda_graph_capture(self._rank_mask_enabled) assert runtime_max_tokens_per_rank <= self.max_num_tokens, "runtime_max_tokens_per_rank must not exceed max_num_tokens" @@ -617,6 +674,7 @@ def get_combine_payload_tensor_in_workspace( Return the combine payload tensor in the workspace, which could be used as the output of MoE kernel to avoid extra copy. Passing the returned tensor to combine lets the C++ op detect workspace ownership. """ + self._require_mapped() if self._state.phase != "dispatched": raise RuntimeError( "get_combine_payload_tensor_in_workspace called before a successful dispatch" 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 8bc2754a58c1..ad8a49912300 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 @@ -574,6 +574,61 @@ def use_cft_for_combine(self, runtime_max_tokens_per_rank: int) -> bool: runtime_max_tokens_per_rank, ) + def _require_mapped(self) -> None: + if not self.mnnvl_mem.mapped: + raise RuntimeError("Native MoE All-to-All workspace handles are unmapped") + + def checkpoint_prepare(self) -> None: + """Collectively detach handles after the caller globally quiesces all owners. + + The local phase and watchdog checks are defense-in-depth only; they do + not prove that every wrapper sharing this workspace is quiescent. + """ + if self._dispatch_state.get("phase") != "idle": + raise RuntimeError("Cannot checkpoint during an active MoE All-to-All phase") + if self._alltoall_watchdog is not None: + raise RuntimeError( + "Checkpointing native MoE All-to-All with the watchdog enabled is not supported" + ) + self.mnnvl_mem.checkpoint_prepare() + + def checkpoint_restore(self, comm) -> None: + """Collectively restore handles after caller-proven global quiescence. + + The low-level remap is idempotent, so every shared owner may call this + method to reset its own protocol state. + """ + self.mnnvl_mem.checkpoint_restore(comm) + refreshed_metainfo = torch.ops.trtllm.moe_a2a_initialize( + self.workspace, + self.ep_rank, + self.ep_size, + self.max_num_tokens_per_rank, + self.eplb_stats_num_experts, + ) + if not torch.equal(refreshed_metainfo, self.moe_a2a_metainfo): + raise RuntimeError( + "MoE All-to-All metainfo changed during MNNVL restore; " + "captured CUDA graphs cannot be replayed safely" + ) + self.moe_a2a_metainfo = refreshed_metainfo + self._workspace_state["metainfo"] = refreshed_metainfo + self._watchdog_coordinator = AlltoAllWatchdogCoordinator( + workspace_state=self._workspace_state, + workspace=self.workspace, + metainfo=refreshed_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, + ) + torch.cuda.synchronize() + comm.barrier() + self._dispatch_state = {"phase": "idle"} + def dispatch( self, hidden_states: torch.Tensor, @@ -604,6 +659,7 @@ def dispatch( Tuple of (hidden_states, hidden_states_sf, token_selected_slots, token_final_scales) Each tensor has shape [ep_size, max_tokens_per_rank, ...] """ + self._require_mapped() if self._dispatch_state.get("phase") == "dispatched": raise RuntimeError("dispatch called twice without an intervening combine") reject_rank_mask_cuda_graph_capture(self._rank_mask_enabled) @@ -744,6 +800,7 @@ def combine( Combined output tensor [local_num_tokens, hidden_size] """ + self._require_mapped() if self._dispatch_state.get("phase") != "dispatched": raise RuntimeError("combine called before a successful dispatch") reject_rank_mask_cuda_graph_capture(self._rank_mask_enabled) diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py index 61d03b3a973b..dab6204e48aa 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -99,6 +99,17 @@ def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) """ return True + def checkpoint_prepare(self) -> None: + """Detach TRT-native two-sided workspaces after global quiescence.""" + if self._dispatch_state: + raise RuntimeError("Cannot checkpoint during an active MoE All-to-All phase") + MnnvlMoe.checkpoint_prepare() + + def checkpoint_restore(self, comm) -> None: + """Restore TRT-native two-sided workspaces and protocol state.""" + MnnvlMoe.checkpoint_restore(comm) + self._dispatch_state = {} + def prepare_dispatch( self, token_selected_slots: torch.Tensor, @@ -108,6 +119,7 @@ def prepare_dispatch( """ NVLINK two-sided comm prepare dispatch: gather EPLB statistics and prepare alltoall_info. """ + MnnvlMoe.require_mapped() all_rank_max_num_tokens = max(all_rank_num_tokens) top_k = token_selected_slots.shape[1] @@ -144,6 +156,7 @@ def dispatch( """ NVLINK two-sided comm dispatch (post-quant, uses alltoall_info from prepare_dispatch). """ + MnnvlMoe.require_mapped() # Read alltoall_info from dispatch_state (set by prepare_dispatch) alltoall_info = self._dispatch_state.get("alltoall_info") if alltoall_info is None: @@ -189,6 +202,7 @@ def combine( """ NVLINK two-sided comm combine - reads from self._dispatch_state. """ + MnnvlMoe.require_mapped() if isinstance(final_hidden_states, list): final_hidden_states = final_hidden_states[0] @@ -204,4 +218,5 @@ def combine( do_reduce=self.alltoall_result_do_sum, ) + self._dispatch_state = {} return final_hidden_states From eda2d04fd77517ccb0b51369a7a605cf964294f5 Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Mon, 20 Jul 2026 11:40:51 -0400 Subject: [PATCH 02/14] [None][test] cover native MoE A2A checkpoint lifecycle Signed-off-by: Hannah Zhang --- .../_torch/test_mnnvl_memory_lifecycle.py | 246 ++++++++++++++++++ 1 file changed, 246 insertions(+) create mode 100644 tests/unittest/_torch/test_mnnvl_memory_lifecycle.py diff --git a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py new file mode 100644 index 000000000000..91bb82f64055 --- /dev/null +++ b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py @@ -0,0 +1,246 @@ +# 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. + +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest + +import tensorrt_llm._mnnvl_utils as mnnvl +from tensorrt_llm._torch.distributed.moe_alltoall import MoeAlltoAll, _A2AState + + +class _FakeComm: + def __init__(self, rank=0, size=2): + self.rank = rank + self.size = size + self.barrier_count = 0 + + def Get_rank(self): + return self.rank + + def Get_size(self): + return self.size + + def barrier(self): + self.barrier_count += 1 + + +class _TestMnnvlMemory(mnnvl.MnnvlMemory): + pass + + +@pytest.fixture +def memory(monkeypatch): + comm = _FakeComm() + record = mnnvl._MnnvlAllocationRecord( + comm=comm, + comm_size=2, + comm_rank=0, + aligned_size=64, + mem_handles=[11, 22], + start_address=1000, + rank_stride=256, + address_offset=32, + ) + obj = _TestMnnvlMemory.__new__(_TestMnnvlMemory) + obj.ptr = 1032 + _TestMnnvlMemory.allocated_map = {obj.ptr: record} + _TestMnnvlMemory.address_refcnt = {record.start_address: 1} + _TestMnnvlMemory.current_start_address = record.start_address + _TestMnnvlMemory.current_rank_stride = record.rank_stride + _TestMnnvlMemory.current_mem_offset = record.address_offset + record.aligned_size + + monkeypatch.setattr(mnnvl.torch.cuda, "synchronize", Mock()) + monkeypatch.setattr(mnnvl, "_check_cu_result", lambda result: result) + monkeypatch.setattr(mnnvl.cuda, "cuMemUnmap", Mock(return_value=None)) + monkeypatch.setattr(mnnvl.cuda, "cuMemRelease", Mock(return_value=None)) + yield obj, record + _TestMnnvlMemory.allocated_map = {} + _TestMnnvlMemory.address_refcnt = {} + if hasattr(obj, "ptr"): + del obj.ptr + + +def test_checkpoint_prepare_preserves_va_and_is_idempotent(memory): + obj, record = memory + + obj.checkpoint_prepare() + + assert not obj.mapped + assert record.start_address == 1000 + assert record.rank_stride == 256 + assert record.address_offset == 32 + assert record.mem_handles == [None, None] + assert record.comm.barrier_count == 2 + assert [call.args[0] for call in mnnvl.cuda.cuMemUnmap.call_args_list] == [1032, 1288] + + obj.checkpoint_prepare() + assert record.comm.barrier_count == 2 + + +def test_checkpoint_restore_reuses_layout_with_fresh_handles(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + restored_comm = _FakeComm() + create_and_map = Mock(return_value=[33, 44]) + monkeypatch.setattr(_TestMnnvlMemory, "_create_and_map_handles", create_and_map) + + obj.checkpoint_restore(restored_comm) + + create_and_map.assert_called_once_with(restored_comm, 64, 1000, 256, 32) + assert obj.ptr == 1032 + assert obj.mapped + assert record.mem_handles == [33, 44] + assert record.comm is restored_comm + assert _TestMnnvlMemory.comm is restored_comm + + +def test_checkpoint_restore_rejects_changed_rank_layout(memory): + obj, _ = memory + obj.checkpoint_prepare() + + with pytest.raises(RuntimeError, match="rank/size 1/2 != 0/2"): + obj.checkpoint_restore(_FakeComm(rank=1)) + + assert not obj.mapped + + +def test_close_detached_memory_only_frees_va(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + address_free = Mock(return_value=None) + monkeypatch.setattr(mnnvl.cuda, "CUdeviceptr", lambda value: value) + monkeypatch.setattr(mnnvl.cuda, "cuMemAddressFree", address_free) + + _TestMnnvlMemory.close_mnnvl_memory(obj.ptr) + + address_free.assert_called_once_with( + record.start_address, record.comm_size * record.rank_stride + ) + assert obj.ptr not in _TestMnnvlMemory.allocated_map + assert record.start_address not in _TestMnnvlMemory.address_refcnt + del obj.ptr + + +def test_create_and_map_handles_cleans_partial_allocation(monkeypatch): + comm = Mock() + comm.Get_rank.return_value = 0 + comm.Get_size.return_value = 2 + comm.allgather.return_value = [b"local", b"remote"] + allocation_prop = SimpleNamespace( + requestedHandleTypes=mnnvl.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC, + location=object(), + ) + access_desc = SimpleNamespace(location=None, flags=None) + + monkeypatch.setattr(mnnvl.MnnvlMemory, "dev_id", 0) + monkeypatch.setattr( + mnnvl.MnnvlMemory, + "get_allocation_prop", + Mock(return_value=allocation_prop), + ) + monkeypatch.setattr(mnnvl, "_check_cu_result", lambda result: result) + monkeypatch.setattr(mnnvl.cuda, "cuCtxGetDevice", Mock(return_value=0)) + monkeypatch.setattr(mnnvl.cuda, "cuMemCreate", Mock(return_value=11)) + monkeypatch.setattr( + mnnvl.cuda, + "cuMemExportToShareableHandle", + Mock(return_value=SimpleNamespace(data=b"local")), + ) + monkeypatch.setattr(mnnvl.cuda, "CUmemAccessDesc", Mock(return_value=access_desc)) + monkeypatch.setattr(mnnvl.cuda, "cuMemImportFromShareableHandle", Mock(return_value=22)) + map_memory = Mock(side_effect=[None, RuntimeError("map failed")]) + unmap_memory = Mock(return_value=None) + release_memory = Mock(return_value=None) + monkeypatch.setattr(mnnvl.cuda, "cuMemMap", map_memory) + monkeypatch.setattr(mnnvl.cuda, "cuMemSetAccess", Mock(return_value=None)) + monkeypatch.setattr(mnnvl.cuda, "cuMemUnmap", unmap_memory) + monkeypatch.setattr(mnnvl.cuda, "cuMemRelease", release_memory) + + with pytest.raises(RuntimeError, match="map failed"): + _TestMnnvlMemory._create_and_map_handles(comm, 64, 1000, 256, 32) + + unmap_memory.assert_called_once_with(1032, 64) + assert [call.args[0] for call in release_memory.call_args_list] == [11, 22] + + +def test_create_and_map_handles_releases_local_handle_on_export_failure(monkeypatch): + comm = Mock() + comm.Get_rank.return_value = 0 + comm.Get_size.return_value = 2 + allocation_prop = SimpleNamespace( + requestedHandleTypes=mnnvl.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC, + location=object(), + ) + release_memory = Mock(return_value=None) + + monkeypatch.setattr(mnnvl.MnnvlMemory, "dev_id", 0) + monkeypatch.setattr( + mnnvl.MnnvlMemory, + "get_allocation_prop", + Mock(return_value=allocation_prop), + ) + monkeypatch.setattr(mnnvl, "_check_cu_result", lambda result: result) + monkeypatch.setattr(mnnvl.cuda, "cuCtxGetDevice", Mock(return_value=0)) + monkeypatch.setattr(mnnvl.cuda, "cuMemCreate", Mock(return_value=11)) + monkeypatch.setattr( + mnnvl.cuda, + "cuMemExportToShareableHandle", + Mock(side_effect=RuntimeError("export failed")), + ) + monkeypatch.setattr(mnnvl.cuda, "cuMemRelease", release_memory) + + with pytest.raises(RuntimeError, match="export failed"): + _TestMnnvlMemory._create_and_map_handles(comm, 64, 1000, 256, 32) + + release_memory.assert_called_once_with(11) + + +def _make_moe_alltoall_for_lifecycle(): + obj = MoeAlltoAll.__new__(MoeAlltoAll) + obj._destroyed = True + obj._state = _A2AState() + obj._alltoall_watchdog = None + obj.mnnvl_mem = Mock(mapped=True) + return obj + + +def test_moe_alltoall_checkpoint_prepare_delegates_with_shared_owners(): + first = _make_moe_alltoall_for_lifecycle() + second = _make_moe_alltoall_for_lifecycle() + shared_memory = Mock(mapped=True) + first.mnnvl_mem = shared_memory + second.mnnvl_mem = shared_memory + + first.checkpoint_prepare() + second.checkpoint_prepare() + + assert shared_memory.checkpoint_prepare.call_count == 2 + + +def test_moe_alltoall_checkpoint_prepare_rejects_active_phase(): + obj = _make_moe_alltoall_for_lifecycle() + obj._state.phase = "dispatched" + with pytest.raises(RuntimeError, match="active MoE All-to-All phase"): + obj.checkpoint_prepare() + + +def test_moe_alltoall_rejects_unmapped_workspace(): + obj = _make_moe_alltoall_for_lifecycle() + obj.mnnvl_mem.mapped = False + + with pytest.raises(RuntimeError, match="workspace handles are unmapped"): + obj._require_mapped() From 61107414cccdf122923efdd79c31b8248476c684 Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Thu, 23 Jul 2026 14:16:24 -0400 Subject: [PATCH 03/14] fix: add require_mapped Signed-off-by: Hannah Zhang --- .../_torch/modules/fused_moe/communication/nvlink_one_sided.py | 1 + 1 file changed, 1 insertion(+) 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 ad8a49912300..1511789cf8cc 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 @@ -897,6 +897,7 @@ def get_combine_payload_tensor_in_workspace( Returns: Tensor view into workspace [ep_size, max_tokens_per_rank, hidden_size] """ + self._require_mapped() if self._dispatch_state.get("phase") != "dispatched": raise RuntimeError( "get_combine_payload_tensor_in_workspace called before a successful dispatch" From 33911f802c6b45f8f785e094454ca52cde82faba Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Tue, 28 Jul 2026 11:12:03 -0400 Subject: [PATCH 04/14] feat(moe): share MNNVL workspace checkpoint lifecycle Signed-off-by: Hannah Zhang --- tensorrt_llm/_mnnvl_utils.py | 132 +++- .../_torch/distributed/moe_alltoall.py | 125 ++-- .../_torch/mnnvl_alltoall_workspace.py | 330 +++++++++ .../communication/nvlink_one_sided.py | 154 ++-- .../communication/nvlink_two_sided.py | 35 +- .../_torch/test_mnnvl_alltoall_workspace.py | 698 ++++++++++++++++++ .../_torch/test_mnnvl_memory_lifecycle.py | 110 ++- 7 files changed, 1381 insertions(+), 203 deletions(-) create mode 100644 tensorrt_llm/_torch/mnnvl_alltoall_workspace.py create mode 100644 tests/unittest/_torch/test_mnnvl_alltoall_workspace.py diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index a48d1e04abbc..515f049de691 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/tensorrt_llm/_mnnvl_utils.py @@ -18,6 +18,7 @@ import platform import sys from dataclasses import dataclass +from enum import Enum from typing import Any, List, Optional, Union import pynvml @@ -51,6 +52,14 @@ def _check_cu_result(cu_func_ret): return None +class _MnnvlAllocationState(Enum): + MAPPED = "mapped" + PREPARING = "preparing" + UNMAPPED = "unmapped" + RESTORING = "restoring" + BROKEN = "broken" + + @dataclass class _MnnvlAllocationRecord: comm: Any @@ -61,7 +70,7 @@ class _MnnvlAllocationRecord: start_address: int rank_stride: int address_offset: int - mapped: bool = True + state: _MnnvlAllocationState = _MnnvlAllocationState.MAPPED class MnnvlMemory: @@ -116,8 +125,8 @@ def __del__(self): @property def mapped(self) -> bool: - """Whether physical handles are mapped into this VA reservation.""" - return type(self).allocated_map[self.ptr].mapped + """Whether the allocation is mapped and ready for data-path access.""" + return type(self).allocated_map[self.ptr].state is _MnnvlAllocationState.MAPPED def as_torch_strided_tensor(self, dtype): num_segments = type(self).comm.Get_size() @@ -390,8 +399,18 @@ def open_mnnvl_memory(cls, mapping: Mapping, size: int): @classmethod def close_mnnvl_memory(cls, ptr: int): - record = cls.allocated_map.pop(ptr) - if record.mapped: + record = cls.allocated_map[ptr] + if record.state not in ( + _MnnvlAllocationState.MAPPED, + _MnnvlAllocationState.UNMAPPED, + ): + logger.warning( + "Skipping cleanup of MNNVL allocation in terminal state %s", + record.state.value, + ) + return + cls.allocated_map.pop(ptr) + if record.state is _MnnvlAllocationState.MAPPED: cls._unmap_and_release_handles(record) cls.address_refcnt[record.start_address] -= 1 @@ -417,21 +436,30 @@ def checkpoint_prepare(self) -> None: """Collectively detach backing handles while retaining graph-visible VA.""" cls = type(self) record = cls.allocated_map[self.ptr] - if not record.mapped: + if record.state is _MnnvlAllocationState.UNMAPPED: return - torch.cuda.synchronize() - record.comm.barrier() - cls._unmap_and_release_handles(record) - record.mem_handles = [None] * record.comm_size - record.mapped = False - record.comm.barrier() - - def checkpoint_restore(self, comm) -> None: - """Collectively remap fresh handles at the original virtual addresses.""" + if record.state is not _MnnvlAllocationState.MAPPED: + raise RuntimeError(f"Cannot prepare MNNVL allocation in {record.state.value} state") + record.state = _MnnvlAllocationState.PREPARING + try: + torch.cuda.synchronize() + record.comm.barrier() + cls._unmap_and_release_handles(record) + record.mem_handles = [None] * record.comm_size + record.comm.barrier() + except Exception: + record.state = _MnnvlAllocationState.BROKEN + raise + record.state = _MnnvlAllocationState.UNMAPPED + + def checkpoint_restore(self, comm) -> bool: + """Remap fresh handles while keeping data-path access disabled.""" cls = type(self) record = cls.allocated_map[self.ptr] - if record.mapped: - return + if record.state is _MnnvlAllocationState.MAPPED: + return False + if record.state is not _MnnvlAllocationState.UNMAPPED: + raise RuntimeError(f"Cannot restore MNNVL allocation in {record.state.value} state") comm_size = comm.Get_size() comm_rank = comm.Get_rank() if comm_size != record.comm_size or comm_rank != record.comm_rank: @@ -441,17 +469,35 @@ def checkpoint_restore(self, comm) -> None: f"rank/size {comm_rank}/{comm_size} != " f"{record.comm_rank}/{record.comm_size}" ) - torch.cuda.synchronize() - record.mem_handles = cls._create_and_map_handles( - comm, - record.aligned_size, - record.start_address, - record.rank_stride, - record.address_offset, - ) + record.state = _MnnvlAllocationState.RESTORING + try: + torch.cuda.synchronize() + record.mem_handles = cls._create_and_map_handles( + comm, + record.aligned_size, + record.start_address, + record.rank_stride, + record.address_offset, + ) + except Exception: + record.state = _MnnvlAllocationState.BROKEN + raise record.comm = comm cls.comm = comm - record.mapped = True + return True + + def _checkpoint_restore_complete(self) -> None: + """Publish a restored allocation after frontend protocol readiness.""" + record = type(self).allocated_map[self.ptr] + if record.state is not _MnnvlAllocationState.RESTORING: + raise RuntimeError(f"Cannot complete MNNVL restore in {record.state.value} state") + record.state = _MnnvlAllocationState.MAPPED + + def _checkpoint_restore_failed(self) -> None: + """Make a failed frontend restore terminal and fail closed.""" + record = type(self).allocated_map[self.ptr] + if record.state is _MnnvlAllocationState.RESTORING: + record.state = _MnnvlAllocationState.BROKEN @staticmethod @functools.cache @@ -659,18 +705,28 @@ def checkpoint_prepare() -> None: def checkpoint_restore(comm) -> None: """Restore TRT-native two-sided MoE workspaces at their original virtual addresses.""" workspaces = (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) - for workspace in workspaces: - if workspace is not None: - workspace.checkpoint_restore(comm) - if MnnvlMoe.moe_workspace_tensor is not None: - assert MnnvlMoe.moe_mapping is not None - torch.ops.trtllm.moe_initialize_workspace( - MnnvlMoe.moe_workspace_tensor, - MnnvlMoe.moe_mapping.moe_ep_rank, - MnnvlMoe.moe_mapping.moe_ep_size, - ) - torch.cuda.synchronize() - comm.barrier() + restored_workspaces = [] + try: + for workspace in workspaces: + if workspace is not None and workspace.checkpoint_restore(comm): + restored_workspaces.append(workspace) + if not restored_workspaces: + return + if MnnvlMoe.moe_workspace_tensor is not None: + assert MnnvlMoe.moe_mapping is not None + torch.ops.trtllm.moe_initialize_workspace( + MnnvlMoe.moe_workspace_tensor, + MnnvlMoe.moe_mapping.moe_ep_rank, + MnnvlMoe.moe_mapping.moe_ep_size, + ) + torch.cuda.synchronize() + comm.barrier() + except Exception: + for workspace in restored_workspaces: + workspace._checkpoint_restore_failed() + raise + for workspace in restored_workspaces: + workspace._checkpoint_restore_complete() @staticmethod def require_mapped() -> None: diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index 9219acdcd08a..44d573337ba3 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -34,6 +34,8 @@ DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, ActiveRankMaskSnapshot, AlltoAllWatchdog, AlltoAllWatchdogCoordinator, AlltoAllWatchdogTimeout, EPGroupHealthLike, reject_rank_mask_cuda_graph_capture) +from tensorrt_llm._torch.mnnvl_alltoall_workspace import \ + _MnnvlAlltoAllWorkspaceLifecycle from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -370,48 +372,59 @@ def __init__( assert workspace_entry[ "can_use_cft_counted_writes"] == self.can_use_cft_counted_writes, "reuse workspace with different CFT mode" + workspace_state = workspace_entry self.mnnvl_mem = workspace_entry["mnnvl_mem"] self.workspace = workspace_entry["workspace"] - self.metainfo = workspace_entry["metainfo"] # Internal state self._state: _A2AState = _A2AState() self.ep_group_health = ep_group_health # Keep the kernel specialization stable for this communicator's lifetime. self._rank_mask_enabled = ep_group_health is not None - workspace_state = workspace_entry self._workspace_state = workspace_state - metainfo_index = self._METAINFO_INDEX - assert metainfo_index is not None - self._watchdog_coordinator = AlltoAllWatchdogCoordinator( - workspace_state=workspace_state, - workspace=self.workspace, - metainfo=self.metainfo, - metainfo_index=metainfo_index, - ep_rank=self.ep_rank, - health=self.ep_group_health, - ) - self._destroyed = False - self._alltoall_watchdog: AlltoAllWatchdog | None = None if (alltoall_watchdog_timeout_s is None and self.ep_group_health is not None): alltoall_watchdog_timeout_s = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S - if alltoall_watchdog_timeout_s is not None: - self._alltoall_watchdog = self._watchdog_coordinator.acquire_watchdog( + metainfo_index = self._METAINFO_INDEX + assert metainfo_index is not None + self._workspace_lifecycle = ( + _MnnvlAlltoAllWorkspaceLifecycle.get_or_create( + workspace_state=workspace_state, + memory=self.mnnvl_mem, + workspace=self.workspace, + metainfo=workspace_state["metainfo"], + metainfo_index=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, - on_timeout=alltoall_watchdog_on_timeout, - ) + health=self.ep_group_health, + )) + self._workspace_lifecycle.register( + self, + watchdog_timeout_s=alltoall_watchdog_timeout_s, + watchdog_poll_interval_s=alltoall_watchdog_poll_interval_s, + watchdog_on_timeout=alltoall_watchdog_on_timeout, + ) + self._destroyed = False + + @property + def metainfo(self) -> torch.Tensor: + return self._workspace_lifecycle.metainfo + + @property + def _watchdog_coordinator(self) -> AlltoAllWatchdogCoordinator: + return self._workspace_lifecycle.coordinator + + @property + def _alltoall_watchdog(self) -> AlltoAllWatchdog | None: + return self._workspace_lifecycle.watchdog_for(self) def destroy(self) -> None: """Stop background watchdog resources owned by this wrapper.""" if getattr(self, "_destroyed", False): return self._destroyed = True - watchdog = getattr(self, "_alltoall_watchdog", None) - if watchdog is not None: - self._watchdog_coordinator.release_watchdog(watchdog) - self._alltoall_watchdog = None + lifecycle = getattr(self, "_workspace_lifecycle", None) + if lifecycle is not None: + lifecycle.unregister(self) def __del__(self) -> None: if not sys.is_finalizing(): @@ -452,53 +465,31 @@ def _require_mapped(self) -> None: "Native MoE All-to-All workspace handles are unmapped") def checkpoint_prepare(self) -> None: - """Collectively detach handles after the caller globally quiesces all owners. - - The local phase and watchdog checks are defense-in-depth only; they do - not prove that every wrapper sharing this workspace is quiescent. - """ - if self._state.phase != "idle": + """Collectively detach handles after every shared owner is idle.""" + if self.can_use_cft_counted_writes: raise RuntimeError( - "Cannot checkpoint during an active MoE All-to-All phase") - if self._alltoall_watchdog is not None: - raise RuntimeError( - "Checkpointing native MoE All-to-All with the watchdog enabled " - "is not supported") - self.mnnvl_mem.checkpoint_prepare() + "Checkpointing a CFT-backed MoE All-to-All workspace is not supported" + ) + self._workspace_lifecycle.checkpoint_prepare() def checkpoint_restore(self, comm) -> None: - """Collectively restore handles after caller-proven global quiescence. - - The low-level remap is idempotent, so every shared owner may call this - method to reset its own protocol state. - """ - self.mnnvl_mem.checkpoint_restore(comm) - refreshed_metainfo = torch.ops.trtllm.moe_a2a_initialize( - self.workspace, - self.ep_rank, - self.ep_size, - self.max_num_tokens, - self.eplb_stats_num_experts, - ) - if not torch.equal(refreshed_metainfo, self.metainfo): - raise RuntimeError( - "MoE All-to-All metainfo changed during MNNVL restore; " - "captured CUDA graphs cannot be replayed safely") - self.metainfo = refreshed_metainfo - assert self._WORKSPACE is not None - self._WORKSPACE["metainfo"] = refreshed_metainfo - metainfo_index = self._METAINFO_INDEX - assert metainfo_index is not None - self._watchdog_coordinator = AlltoAllWatchdogCoordinator( - workspace_state=self._WORKSPACE, - workspace=self.workspace, - metainfo=refreshed_metainfo, - metainfo_index=metainfo_index, - ep_rank=self.ep_rank, - health=self.ep_group_health, + """Collectively restore handles and all shared frontend state.""" + self._workspace_lifecycle.checkpoint_restore( + comm, + lambda: torch.ops.trtllm.moe_a2a_initialize( + self.workspace, + self.ep_rank, + self.ep_size, + self.max_num_tokens, + self.eplb_stats_num_experts, + self.can_use_cft_counted_writes, + ), ) - torch.cuda.synchronize() - comm.barrier() + + def _mnnvl_checkpoint_is_idle(self) -> bool: + return self._state.phase == "idle" + + def _mnnvl_checkpoint_reset(self) -> None: self.reset_state() def dispatch(self, diff --git a/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py b/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py new file mode 100644 index 000000000000..e800c95cc7cb --- /dev/null +++ b/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py @@ -0,0 +1,330 @@ +# 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. + +from dataclasses import dataclass +from typing import Callable, Mapping, MutableMapping, Protocol, cast +from weakref import WeakSet + +import torch + +from tensorrt_llm._mnnvl_utils import MnnvlMemory +from tensorrt_llm._torch.alltoall_watchdog import ( + AlltoAllWatchdog, + AlltoAllWatchdogCoordinator, + AlltoAllWatchdogTimeout, + EPGroupHealthLike, +) + +_WORKSPACE_LIFECYCLE_KEY = "mnnvl_alltoall_workspace_lifecycle" + + +class _CollectiveCommunicator(Protocol): + def barrier(self) -> None: + """Synchronize every rank participating in the workspace.""" + + def allgather(self, value: bool) -> list[bool]: + """Gather one readiness value from each participating rank.""" + ... + + +def _collect_active_ranks( + comm: _CollectiveCommunicator, + *, + local_clients_idle: bool, + expected_size: int, +) -> list[int]: + """Collectively return ranks whose local workspace clients are active.""" + clients_idle_by_rank = comm.allgather(local_clients_idle) + if len(clients_idle_by_rank) != expected_size: + raise RuntimeError( + "MNNVL workspace communicator size does not match the MoE EP group: " + f"{len(clients_idle_by_rank)} != {expected_size}" + ) + return [rank for rank, clients_idle in enumerate(clients_idle_by_rank) if not clients_idle] + + +class _WorkspaceClient(Protocol): + def _mnnvl_checkpoint_is_idle(self) -> bool: + """Return whether this client can safely enter a checkpoint.""" + ... + + def _mnnvl_checkpoint_reset(self) -> None: + """Reset frontend protocol state after a successful restore.""" + + +@dataclass(frozen=True) +class _WatchdogConfig: + ep_size: int + timeout_s: float + poll_interval_s: float + on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None + + +class _MnnvlAlltoAllWorkspaceLifecycle: + """Own checkpoint and watchdog transitions for one shared MoE workspace. + + The checkpoint coordinator must prevent new dispatches from starting while + a transition is in progress. + """ + + def __init__( + self, + *, + workspace_state: MutableMapping[str, object], + memory: MnnvlMemory, + workspace: torch.Tensor, + metainfo: torch.Tensor, + metainfo_index: Mapping[str, int], + ep_rank: int, + ep_size: int, + health: EPGroupHealthLike | None, + ) -> None: + self._workspace_state = workspace_state + self._memory = memory + self._workspace = workspace + self._metainfo = metainfo + self._metainfo_index = dict(metainfo_index) + self._ep_rank = int(ep_rank) + self._ep_size = int(ep_size) + self._health = health + self._clients: WeakSet[_WorkspaceClient] = WeakSet() + self._watchdog_clients: WeakSet[_WorkspaceClient] = WeakSet() + self._watchdog_config: _WatchdogConfig | None = None + self._watchdog: AlltoAllWatchdog | None = None + self._coordinator = self._create_coordinator() + + @classmethod + def get_or_create( + cls, + *, + workspace_state: MutableMapping[str, object], + memory: MnnvlMemory, + workspace: torch.Tensor, + metainfo: torch.Tensor, + metainfo_index: Mapping[str, int], + ep_rank: int, + ep_size: int, + health: EPGroupHealthLike | None, + ) -> "_MnnvlAlltoAllWorkspaceLifecycle": + lifecycle = workspace_state.get(_WORKSPACE_LIFECYCLE_KEY) + if lifecycle is None: + lifecycle = cls( + workspace_state=workspace_state, + memory=memory, + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=ep_rank, + ep_size=ep_size, + health=health, + ) + workspace_state[_WORKSPACE_LIFECYCLE_KEY] = lifecycle + return lifecycle + if not isinstance(lifecycle, cls): + raise TypeError("invalid MNNVL All-to-All workspace lifecycle state") + lifecycle._validate_shared_context( + memory=memory, + workspace=workspace, + metainfo=metainfo, + metainfo_index=metainfo_index, + ep_rank=ep_rank, + ep_size=ep_size, + health=health, + ) + return lifecycle + + @property + def metainfo(self) -> torch.Tensor: + return self._metainfo + + @property + def coordinator(self) -> AlltoAllWatchdogCoordinator: + return self._coordinator + + def watchdog_for(self, client: _WorkspaceClient) -> AlltoAllWatchdog | None: + if client not in self._watchdog_clients: + return None + return self._watchdog + + def register( + self, + client: _WorkspaceClient, + *, + watchdog_timeout_s: float | None, + watchdog_poll_interval_s: float, + watchdog_on_timeout: Callable[[AlltoAllWatchdogTimeout], None] | None, + ) -> None: + if client in self._clients: + return + if watchdog_timeout_s is None: + self._clients.add(client) + return + + config = _WatchdogConfig( + ep_size=self._ep_size, + timeout_s=float(watchdog_timeout_s), + poll_interval_s=float(watchdog_poll_interval_s), + on_timeout=watchdog_on_timeout, + ) + self._validate_watchdog_config(config) + self._clients.add(client) + self._watchdog_clients.add(client) + try: + if self._memory.mapped: + self._start_watchdog() + except Exception: + self._watchdog_clients.discard(client) + self._clients.discard(client) + if not self._watchdog_clients: + self._watchdog_config = None + raise + + def unregister(self, client: _WorkspaceClient) -> None: + self._clients.discard(client) + self._watchdog_clients.discard(client) + if not self._watchdog_clients: + self._stop_watchdog() + self._watchdog_config = None + + def checkpoint_prepare(self) -> None: + """Stop shared readers, then collectively detach the backing handles.""" + if not self._memory.mapped: + self._stop_watchdog() + self._memory.checkpoint_prepare() + return + local_clients_idle = all( + client._mnnvl_checkpoint_is_idle() for client in list(self._clients) + ) + comm = cast(_CollectiveCommunicator | None, self._memory.comm) + if comm is None: + raise RuntimeError("MNNVL workspace communicator is not initialized") + active_ranks = _collect_active_ranks( + comm, + local_clients_idle=local_clients_idle, + expected_size=self._ep_size, + ) + if active_ranks: + raise RuntimeError( + f"Cannot checkpoint during an active MoE All-to-All phase on ranks {active_ranks}" + ) + self._stop_watchdog() + self._memory.checkpoint_prepare() + + def checkpoint_restore( + self, + comm: _CollectiveCommunicator, + initialize_frontend: Callable[[], torch.Tensor], + ) -> None: + """Restore backing handles and publish the workspace after frontend readiness.""" + if self._memory.mapped: + return + restored = self._memory.checkpoint_restore(comm) + if restored is False: + return + try: + refreshed_metainfo = initialize_frontend() + if not torch.equal(refreshed_metainfo, self._metainfo): + raise RuntimeError( + "MoE All-to-All metainfo changed during MNNVL restore; " + "captured CUDA graphs cannot be replayed safely" + ) + self._metainfo = refreshed_metainfo + self._workspace_state["metainfo"] = refreshed_metainfo + self._coordinator = self._create_coordinator() + torch.cuda.synchronize() + comm.barrier() + self._start_watchdog() + for client in list(self._clients): + client._mnnvl_checkpoint_reset() + self._memory._checkpoint_restore_complete() + except Exception: + self._memory._checkpoint_restore_failed() + self._stop_watchdog() + raise + + def _create_coordinator(self) -> AlltoAllWatchdogCoordinator: + return AlltoAllWatchdogCoordinator( + workspace_state=self._workspace_state, + workspace=self._workspace, + metainfo=self._metainfo, + metainfo_index=self._metainfo_index, + ep_rank=self._ep_rank, + health=self._health, + ) + + def _start_watchdog(self) -> None: + config = self._watchdog_config + if config is None or not self._watchdog_clients or self._watchdog is not None: + return + self._watchdog = self._coordinator.acquire_watchdog( + ep_size=config.ep_size, + timeout_s=config.timeout_s, + poll_interval_s=config.poll_interval_s, + on_timeout=config.on_timeout, + ) + + def _stop_watchdog(self) -> None: + if self._watchdog is None: + return + self._coordinator.release_watchdog(self._watchdog) + self._watchdog = None + + def _validate_shared_context( + self, + *, + memory: MnnvlMemory, + workspace: torch.Tensor, + metainfo: torch.Tensor, + metainfo_index: Mapping[str, int], + ep_rank: int, + ep_size: int, + health: EPGroupHealthLike | None, + ) -> None: + if ( + self._memory is not memory + or self._workspace is not workspace + or self._metainfo is not metainfo + or self._metainfo_index != dict(metainfo_index) + or self._ep_rank != ep_rank + or self._ep_size != ep_size + ): + raise ValueError( + "MNNVL All-to-All wrappers sharing a workspace must use the " + "same allocation, metadata layout, and rank layout" + ) + if self._health is health: + return + if self._clients or self._watchdog is not None or self._watchdog_config is not None: + raise ValueError( + "MNNVL All-to-All wrappers sharing a workspace must use the same EP health object" + ) + self._health = health + self._coordinator = self._create_coordinator() + + def _validate_watchdog_config(self, requested: _WatchdogConfig) -> None: + existing = self._watchdog_config + if existing is None: + self._watchdog_config = requested + return + if ( + existing.ep_size != requested.ep_size + or existing.timeout_s != requested.timeout_s + or existing.poll_interval_s != requested.poll_interval_s + or existing.on_timeout is not requested.on_timeout + ): + raise ValueError( + "MNNVL All-to-All wrappers sharing a workspace must use the " + "same watchdog configuration" + ) 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 1511789cf8cc..909a21b245fe 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 sys from typing import Callable, Dict, List, Optional, Tuple import torch @@ -40,6 +41,7 @@ EPGroupHealthLike, reject_rank_mask_cuda_graph_capture, ) +from tensorrt_llm._torch.mnnvl_alltoall_workspace import _MnnvlAlltoAllWorkspaceLifecycle from tensorrt_llm.bindings import internal as _tllm_internal from tensorrt_llm.logger import logger as tllm_logger from tensorrt_llm.mapping import Mapping @@ -396,7 +398,7 @@ def __init__( workspace_state = NVLinkOneSided._WORKSPACES.get(self._workspace_key) memory_cls = CftMnnvlMemory if self.can_use_cft_counted_writes else MnnvlMemory - + workspace_created = workspace_state is None if workspace_state is None: tllm_logger.info( f"NVLinkOneSided: Allocating workspace with size {self.workspace_size_per_rank} bytes." @@ -429,7 +431,6 @@ def __init__( "metainfo": metainfo, "cft_initialized": False, } - NVLinkOneSided._WORKSPACES[self._workspace_key] = workspace_state else: expected_workspace_state = { "workspace_size_per_rank": self.workspace_size_per_rank, @@ -449,41 +450,53 @@ def __init__( f"reuse workspace with different {key}" ) - NVLinkOneSided._WORKSPACE = workspace_state - NVLinkOneSided._WORKSPACE_REFCOUNTS[self._workspace_key] = ( - NVLinkOneSided._WORKSPACE_REFCOUNTS.get(self._workspace_key, 0) + 1 - ) self._destroyed = False + self._workspace_registered = False self._workspace_state = workspace_state self.mnnvl_mem = workspace_state["mnnvl_mem"] self.workspace = workspace_state["workspace"] - self.moe_a2a_metainfo = workspace_state["metainfo"] self.max_num_tokens_per_rank = workspace_state["max_num_tokens_per_rank"] self.ep_group_health = ep_group_health # Keep the kernel specialization stable for this communicator's lifetime. self._rank_mask_enabled = ep_group_health is not None - self._watchdog_coordinator = AlltoAllWatchdogCoordinator( + if alltoall_watchdog_timeout_s is None and self.ep_group_health is not None: + alltoall_watchdog_timeout_s = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S + flag_val_offset_index = self.FLAG_VAL_OFFSET_INDEX + dispatch_flags_offset_index = self.DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX + combine_flags_offset_index = self.COMBINE_COMPLETION_FLAGS_OFFSET_INDEX + if ( + flag_val_offset_index is None + or dispatch_flags_offset_index is None + or combine_flags_offset_index is None + ): + raise RuntimeError("MoE All-to-All metadata indices are not initialized") + self._workspace_lifecycle = _MnnvlAlltoAllWorkspaceLifecycle.get_or_create( workspace_state=workspace_state, + memory=self.mnnvl_mem, workspace=self.workspace, - metainfo=self.moe_a2a_metainfo, + metainfo=workspace_state["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, + "FLAG_VAL_OFFSET_INDEX": flag_val_offset_index, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": dispatch_flags_offset_index, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": combine_flags_offset_index, }, ep_rank=self.ep_rank, + ep_size=self.ep_size, health=self.ep_group_health, ) - self._alltoall_watchdog: AlltoAllWatchdog | None = None - if alltoall_watchdog_timeout_s is None and self.ep_group_health is not None: - alltoall_watchdog_timeout_s = DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S - if alltoall_watchdog_timeout_s is not None: - self._alltoall_watchdog = self._watchdog_coordinator.acquire_watchdog( - ep_size=self.ep_size, - timeout_s=alltoall_watchdog_timeout_s, - poll_interval_s=alltoall_watchdog_poll_interval_s, - on_timeout=alltoall_watchdog_on_timeout, - ) + self._workspace_lifecycle.register( + self, + watchdog_timeout_s=alltoall_watchdog_timeout_s, + watchdog_poll_interval_s=alltoall_watchdog_poll_interval_s, + watchdog_on_timeout=alltoall_watchdog_on_timeout, + ) + if workspace_created: + NVLinkOneSided._WORKSPACES[self._workspace_key] = workspace_state + NVLinkOneSided._WORKSPACE = workspace_state + NVLinkOneSided._WORKSPACE_REFCOUNTS[self._workspace_key] = ( + NVLinkOneSided._WORKSPACE_REFCOUNTS.get(self._workspace_key, 0) + 1 + ) + self._workspace_registered = True # Initialize CFT Logical Endpoints by binding the LE to the workspace. # The LE IS the workspace — no separate allocation or payload layout needed. @@ -503,6 +516,24 @@ 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 + @property + def moe_a2a_metainfo(self) -> torch.Tensor: + return self._require_workspace_lifecycle().metainfo + + @property + def _watchdog_coordinator(self) -> AlltoAllWatchdogCoordinator: + return self._require_workspace_lifecycle().coordinator + + @property + def _alltoall_watchdog(self) -> AlltoAllWatchdog | None: + return self._require_workspace_lifecycle().watchdog_for(self) + + def _require_workspace_lifecycle(self) -> _MnnvlAlltoAllWorkspaceLifecycle: + lifecycle = self._workspace_lifecycle + if lifecycle is None: + raise RuntimeError("NVLinkOneSided workspace has been destroyed") + return lifecycle + @staticmethod def is_platform_supported() -> bool: """ @@ -522,12 +553,14 @@ def destroy(self): return self._destroyed = True - if self._alltoall_watchdog is not None: - self._watchdog_coordinator.release_watchdog(self._alltoall_watchdog) - self._alltoall_watchdog = None + lifecycle = getattr(self, "_workspace_lifecycle", None) + if lifecycle is not None: + lifecycle.unregister(self) workspace_key = getattr(self, "_workspace_key", None) - if workspace_key is None: + if workspace_key is None or not getattr(self, "_workspace_registered", False): + self._workspace_lifecycle = None return + self._workspace_registered = False if torch.cuda.is_available(): torch.cuda.synchronize() @@ -545,10 +578,14 @@ def destroy(self): self.mnnvl_mem = None self.workspace = None - self.moe_a2a_metainfo = None self._workspace_state = None + self._workspace_lifecycle = None self._dispatch_state = {"phase": "destroyed"} + def __del__(self) -> None: + if not sys.is_finalizing(): + self.destroy() + def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) -> bool: """ Check if NVLINK one-sided comm is feasible for the given workload at runtime. @@ -579,54 +616,31 @@ def _require_mapped(self) -> None: raise RuntimeError("Native MoE All-to-All workspace handles are unmapped") def checkpoint_prepare(self) -> None: - """Collectively detach handles after the caller globally quiesces all owners. - - The local phase and watchdog checks are defense-in-depth only; they do - not prove that every wrapper sharing this workspace is quiescent. - """ - if self._dispatch_state.get("phase") != "idle": - raise RuntimeError("Cannot checkpoint during an active MoE All-to-All phase") - if self._alltoall_watchdog is not None: + """Collectively detach handles after every shared owner is idle.""" + if self.can_use_cft_counted_writes: raise RuntimeError( - "Checkpointing native MoE All-to-All with the watchdog enabled is not supported" + "Checkpointing a CFT-backed MoE All-to-All workspace is not supported" ) - self.mnnvl_mem.checkpoint_prepare() + self._require_workspace_lifecycle().checkpoint_prepare() def checkpoint_restore(self, comm) -> None: - """Collectively restore handles after caller-proven global quiescence. - - The low-level remap is idempotent, so every shared owner may call this - method to reset its own protocol state. - """ - self.mnnvl_mem.checkpoint_restore(comm) - refreshed_metainfo = torch.ops.trtllm.moe_a2a_initialize( - self.workspace, - self.ep_rank, - self.ep_size, - self.max_num_tokens_per_rank, - self.eplb_stats_num_experts, - ) - if not torch.equal(refreshed_metainfo, self.moe_a2a_metainfo): - raise RuntimeError( - "MoE All-to-All metainfo changed during MNNVL restore; " - "captured CUDA graphs cannot be replayed safely" - ) - self.moe_a2a_metainfo = refreshed_metainfo - self._workspace_state["metainfo"] = refreshed_metainfo - self._watchdog_coordinator = AlltoAllWatchdogCoordinator( - workspace_state=self._workspace_state, - workspace=self.workspace, - metainfo=refreshed_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, + """Collectively restore handles and all shared frontend state.""" + self._require_workspace_lifecycle().checkpoint_restore( + comm, + lambda: torch.ops.trtllm.moe_a2a_initialize( + self.workspace, + self.ep_rank, + self.ep_size, + self.max_num_tokens_per_rank, + self.eplb_stats_num_experts, + self.can_use_cft_counted_writes, + ), ) - torch.cuda.synchronize() - comm.barrier() + + def _mnnvl_checkpoint_is_idle(self) -> bool: + return self._dispatch_state.get("phase") == "idle" + + def _mnnvl_checkpoint_reset(self) -> None: self._dispatch_state = {"phase": "idle"} def dispatch( diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py index dab6204e48aa..5eb0f0cd3812 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py @@ -23,10 +23,12 @@ import os from typing import List, Optional, Tuple +from weakref import WeakSet import torch from tensorrt_llm._mnnvl_utils import MnnvlMemory, MnnvlMoe +from tensorrt_llm._torch.mnnvl_alltoall_workspace import _collect_active_ranks from tensorrt_llm.mapping import Mapping from .base import Communication @@ -42,6 +44,8 @@ class NVLinkTwoSided(Communication): The required symmetric memory size is proportional to the communication channels opened. """ + _INSTANCES: WeakSet = WeakSet() + def __init__( self, mapping: Mapping, @@ -76,6 +80,7 @@ def __init__( # Initialize dispatch state self._dispatch_state = {} + self._INSTANCES.add(self) @staticmethod def is_platform_supported() -> bool: @@ -101,14 +106,38 @@ def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) def checkpoint_prepare(self) -> None: """Detach TRT-native two-sided workspaces after global quiescence.""" - if self._dispatch_state: - raise RuntimeError("Cannot checkpoint during an active MoE All-to-All phase") + workspaces = (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) + if all(workspace is None or not workspace.mapped for workspace in workspaces): + MnnvlMoe.checkpoint_prepare() + return + local_clients_idle = not any(instance._dispatch_state for instance in self._INSTANCES) + workspace = MnnvlMoe.moe_workspace + assert workspace is not None + comm = workspace.comm + if comm is None: + raise RuntimeError("MNNVL workspace communicator is not initialized") + active_ranks = _collect_active_ranks( + comm, + local_clients_idle=local_clients_idle, + expected_size=self.ep_size, + ) + if active_ranks: + raise RuntimeError( + f"Cannot checkpoint during an active MoE All-to-All phase on ranks {active_ranks}" + ) MnnvlMoe.checkpoint_prepare() def checkpoint_restore(self, comm) -> None: """Restore TRT-native two-sided workspaces and protocol state.""" + restore_required = any( + workspace is not None and not workspace.mapped + for workspace in (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) + ) MnnvlMoe.checkpoint_restore(comm) - self._dispatch_state = {} + if not restore_required: + return + for instance in self._INSTANCES: + instance._dispatch_state = {} def prepare_dispatch( self, diff --git a/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py b/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py new file mode 100644 index 000000000000..9325e126f31c --- /dev/null +++ b/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py @@ -0,0 +1,698 @@ +# 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. + +import gc +from types import SimpleNamespace +from unittest.mock import Mock +from weakref import WeakSet + +import pytest +import torch + +import tensorrt_llm._mnnvl_utils as mnnvl +import tensorrt_llm._torch.modules.fused_moe.communication.nvlink_one_sided as one_sided_module +from tensorrt_llm._torch.distributed.moe_alltoall import MoeAlltoAll +from tensorrt_llm._torch.mnnvl_alltoall_workspace import _MnnvlAlltoAllWorkspaceLifecycle +from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_one_sided import NVLinkOneSided +from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_two_sided import NVLinkTwoSided + + +class _Client: + def __init__(self, *, idle: bool = True) -> None: + self.idle = idle + self.reset_count = 0 + + def _mnnvl_checkpoint_is_idle(self) -> bool: + return self.idle + + def _mnnvl_checkpoint_reset(self) -> None: + self.reset_count += 1 + + +class _FailingResetClient(_Client): + def _mnnvl_checkpoint_reset(self) -> None: + raise RuntimeError("frontend reset failed") + + +class _FakeComm: + def __init__(self, clients_idle_by_rank: list[bool] | None = None) -> None: + self.barrier_count = 0 + self.allgather_count = 0 + self.clients_idle_by_rank = clients_idle_by_rank + + def barrier(self) -> None: + self.barrier_count += 1 + + def allgather(self, local_clients_idle: bool) -> list[bool]: + self.allgather_count += 1 + if self.clients_idle_by_rank is not None: + return self.clients_idle_by_rank + return [local_clients_idle, local_clients_idle] + + +def _make_lifecycle() -> tuple[_MnnvlAlltoAllWorkspaceLifecycle, Mock, torch.Tensor]: + workspace_state = {} + memory = Mock(mapped=True) + memory.comm = _FakeComm() + workspace = torch.zeros(1, dtype=torch.uint8) + metainfo = torch.tensor([1]) + lifecycle = _MnnvlAlltoAllWorkspaceLifecycle.get_or_create( + workspace_state=workspace_state, + memory=memory, + workspace=workspace, + metainfo=metainfo, + metainfo_index={ + "FLAG_VAL_OFFSET_INDEX": 0, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 0, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 0, + }, + ep_rank=0, + ep_size=2, + health=None, + ) + return lifecycle, memory, metainfo + + +def _register_without_watchdog( + lifecycle: _MnnvlAlltoAllWorkspaceLifecycle, + client: _Client, +) -> None: + lifecycle.register( + client, + watchdog_timeout_s=None, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + + +def test_checkpoint_prepare_rejects_any_active_shared_client() -> None: + lifecycle, memory, _ = _make_lifecycle() + idle = _Client() + _register_without_watchdog(lifecycle, idle) + active = _Client(idle=False) + _register_without_watchdog(lifecycle, active) + + with pytest.raises(RuntimeError, match="active MoE All-to-All phase"): + lifecycle.checkpoint_prepare() + + memory.checkpoint_prepare.assert_not_called() + + +def test_repeated_checkpoint_prepare_skips_shared_preflight() -> None: + lifecycle, memory, _ = _make_lifecycle() + memory.mapped = False + + lifecycle.checkpoint_prepare() + + assert memory.comm.allgather_count == 0 + memory.checkpoint_prepare.assert_called_once_with() + + +def test_detached_checkpoint_prepare_stops_stale_watchdog() -> None: + lifecycle, memory, _ = _make_lifecycle() + coordinator = Mock() + watchdog = Mock() + coordinator.acquire_watchdog.return_value = watchdog + lifecycle._coordinator = coordinator + lifecycle.register( + _Client(), + watchdog_timeout_s=5.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + memory.mapped = False + + lifecycle.checkpoint_prepare() + + coordinator.release_watchdog.assert_called_once_with(watchdog) + assert memory.comm.allgather_count == 0 + memory.checkpoint_prepare.assert_called_once_with() + + +def test_checkpoint_prepare_rejects_uninitialized_communicator() -> None: + lifecycle, memory, _ = _make_lifecycle() + memory.comm = None + + with pytest.raises(RuntimeError, match="communicator is not initialized"): + lifecycle.checkpoint_prepare() + + memory.checkpoint_prepare.assert_not_called() + + +def test_checkpoint_prepare_rejects_communicator_size_mismatch() -> None: + lifecycle, memory, _ = _make_lifecycle() + memory.comm = _FakeComm(clients_idle_by_rank=[True]) + + with pytest.raises(RuntimeError, match="communicator size does not match"): + lifecycle.checkpoint_prepare() + + memory.checkpoint_prepare.assert_not_called() + + +def test_checkpoint_prepare_rejects_remote_active_client_before_watchdog_stop() -> None: + lifecycle, memory, _ = _make_lifecycle() + memory.comm = _FakeComm(clients_idle_by_rank=[True, False]) + coordinator = Mock() + watchdog = Mock() + coordinator.acquire_watchdog.return_value = watchdog + lifecycle._coordinator = coordinator + client = _Client() + lifecycle.register( + client, + watchdog_timeout_s=5.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + + with pytest.raises(RuntimeError, match=r"active MoE All-to-All phase on ranks \[1\]"): + lifecycle.checkpoint_prepare() + + coordinator.release_watchdog.assert_not_called() + memory.checkpoint_prepare.assert_not_called() + + +def test_checkpoint_suspends_and_recreates_one_shared_watchdog( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle, memory, metainfo = _make_lifecycle() + old_coordinator = Mock() + old_watchdog = Mock() + old_coordinator.acquire_watchdog.return_value = old_watchdog + lifecycle._coordinator = old_coordinator + first = _Client() + second = _Client() + + for client in (first, second): + lifecycle.register( + client, + watchdog_timeout_s=5.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + + old_coordinator.acquire_watchdog.assert_called_once_with( + ep_size=2, + timeout_s=5.0, + poll_interval_s=0.1, + on_timeout=None, + ) + assert lifecycle.watchdog_for(first) is old_watchdog + assert lifecycle.watchdog_for(second) is old_watchdog + + lifecycle.checkpoint_prepare() + + old_coordinator.release_watchdog.assert_called_once_with(old_watchdog) + memory.checkpoint_prepare.assert_called_once_with() + + memory.mapped = False + memory.checkpoint_restore.return_value = True + new_coordinator = Mock() + new_watchdog = Mock() + new_coordinator.acquire_watchdog.return_value = new_watchdog + monkeypatch.setattr( + lifecycle, + "_create_coordinator", + Mock(return_value=new_coordinator), + ) + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + comm = _FakeComm() + + lifecycle.checkpoint_restore(comm, Mock(return_value=metainfo)) + + memory._checkpoint_restore_complete.assert_called_once_with() + new_coordinator.acquire_watchdog.assert_called_once_with( + ep_size=2, + timeout_s=5.0, + poll_interval_s=0.1, + on_timeout=None, + ) + assert lifecycle.watchdog_for(first) is new_watchdog + assert lifecycle.watchdog_for(second) is new_watchdog + assert first.reset_count == 1 + assert second.reset_count == 1 + assert comm.barrier_count == 1 + + +def test_unregister_stops_watchdog_after_last_enabled_client() -> None: + lifecycle, _, _ = _make_lifecycle() + coordinator = Mock() + watchdog = Mock() + coordinator.acquire_watchdog.return_value = watchdog + lifecycle._coordinator = coordinator + first = _Client() + second = _Client() + for client in (first, second): + lifecycle.register( + client, + watchdog_timeout_s=5.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + + lifecycle.unregister(first) + coordinator.release_watchdog.assert_not_called() + + lifecycle.unregister(second) + coordinator.release_watchdog.assert_called_once_with(watchdog) + + +def test_shared_watchdog_configuration_mismatch_rejects_new_client() -> None: + lifecycle, _, _ = _make_lifecycle() + coordinator = Mock() + coordinator.acquire_watchdog.return_value = Mock() + lifecycle._coordinator = coordinator + first = _Client() + second = _Client() + lifecycle.register( + first, + watchdog_timeout_s=5.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + + with pytest.raises(ValueError, match="same watchdog configuration"): + lifecycle.register( + second, + watchdog_timeout_s=10.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + + assert lifecycle.watchdog_for(second) is None + coordinator.acquire_watchdog.assert_called_once() + + +def test_watchdog_registration_is_deferred_while_workspace_is_unmapped( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle, memory, metainfo = _make_lifecycle() + memory.mapped = False + memory.checkpoint_restore.return_value = True + old_coordinator = Mock() + lifecycle._coordinator = old_coordinator + client = _Client() + + lifecycle.register( + client, + watchdog_timeout_s=5.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + + old_coordinator.acquire_watchdog.assert_not_called() + assert lifecycle.watchdog_for(client) is None + + new_coordinator = Mock() + new_watchdog = Mock() + new_coordinator.acquire_watchdog.return_value = new_watchdog + monkeypatch.setattr( + lifecycle, + "_create_coordinator", + Mock(return_value=new_coordinator), + ) + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + + lifecycle.checkpoint_restore(_FakeComm(), Mock(return_value=metainfo)) + + assert lifecycle.watchdog_for(client) is new_watchdog + + +def test_checkpoint_restore_failure_before_watchdog_start_stays_unpublished( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle, memory, _ = _make_lifecycle() + client = _Client() + _register_without_watchdog(lifecycle, client) + memory.mapped = False + memory.checkpoint_restore.return_value = True + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + + with pytest.raises(RuntimeError, match="frontend restore failed"): + lifecycle.checkpoint_restore( + _FakeComm(), + Mock(side_effect=RuntimeError("frontend restore failed")), + ) + + memory._checkpoint_restore_failed.assert_called_once_with() + memory._checkpoint_restore_complete.assert_not_called() + assert client.reset_count == 0 + + +def test_checkpoint_restore_rejects_changed_metainfo_and_fails_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle, memory, _ = _make_lifecycle() + client = _Client() + _register_without_watchdog(lifecycle, client) + memory.mapped = False + memory.checkpoint_restore.return_value = True + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + + with pytest.raises(RuntimeError, match="metainfo changed"): + lifecycle.checkpoint_restore( + _FakeComm(), + Mock(return_value=torch.tensor([2])), + ) + + memory._checkpoint_restore_failed.assert_called_once_with() + memory._checkpoint_restore_complete.assert_not_called() + assert client.reset_count == 0 + + +def test_checkpoint_restore_failure_after_watchdog_start_stops_it( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle, memory, metainfo = _make_lifecycle() + old_coordinator = Mock() + old_watchdog = Mock() + old_coordinator.acquire_watchdog.return_value = old_watchdog + lifecycle._coordinator = old_coordinator + client = _FailingResetClient() + lifecycle.register( + client, + watchdog_timeout_s=5.0, + watchdog_poll_interval_s=0.1, + watchdog_on_timeout=None, + ) + lifecycle.checkpoint_prepare() + memory.mapped = False + memory.checkpoint_restore.return_value = True + new_coordinator = Mock() + new_watchdog = Mock() + new_coordinator.acquire_watchdog.return_value = new_watchdog + monkeypatch.setattr( + lifecycle, + "_create_coordinator", + Mock(return_value=new_coordinator), + ) + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + + with pytest.raises(RuntimeError, match="frontend reset failed"): + lifecycle.checkpoint_restore(_FakeComm(), Mock(return_value=metainfo)) + + new_coordinator.release_watchdog.assert_called_once_with(new_watchdog) + memory._checkpoint_restore_failed.assert_called_once_with() + memory._checkpoint_restore_complete.assert_not_called() + + +def test_two_sided_checkpoint_prepare_rejects_active_shared_owner( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instances = WeakSet() + monkeypatch.setattr(NVLinkTwoSided, "_INSTANCES", instances) + checkpoint_prepare = Mock() + monkeypatch.setattr(mnnvl.MnnvlMoe, "checkpoint_prepare", checkpoint_prepare) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_workspace", + Mock(comm=_FakeComm()), + ) + idle = NVLinkTwoSided.__new__(NVLinkTwoSided) + idle.ep_size = 2 + idle._dispatch_state = {} + active = NVLinkTwoSided.__new__(NVLinkTwoSided) + active._dispatch_state = {"alltoall_info": object()} + instances.update((idle, active)) + + with pytest.raises( + RuntimeError, + match=r"active MoE All-to-All phase on ranks \[0, 1\]", + ): + idle.checkpoint_prepare() + + checkpoint_prepare.assert_not_called() + + +def test_two_sided_repeated_checkpoint_prepare_skips_shared_preflight( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instances = WeakSet() + monkeypatch.setattr(NVLinkTwoSided, "_INSTANCES", instances) + checkpoint_prepare = Mock() + monkeypatch.setattr(mnnvl.MnnvlMoe, "checkpoint_prepare", checkpoint_prepare) + comm = _FakeComm() + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_workspace", + Mock(mapped=False, comm=comm), + ) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_prepare_workspace", + Mock(mapped=False), + ) + owner = NVLinkTwoSided.__new__(NVLinkTwoSided) + owner.ep_size = 2 + owner._dispatch_state = {} + instances.add(owner) + + owner.checkpoint_prepare() + + assert comm.allgather_count == 0 + checkpoint_prepare.assert_called_once_with() + + +def test_two_sided_checkpoint_prepare_rejects_uninitialized_communicator( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instances = WeakSet() + monkeypatch.setattr(NVLinkTwoSided, "_INSTANCES", instances) + checkpoint_prepare = Mock() + monkeypatch.setattr(mnnvl.MnnvlMoe, "checkpoint_prepare", checkpoint_prepare) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_workspace", + Mock(mapped=True, comm=None), + ) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_prepare_workspace", + Mock(mapped=True), + ) + owner = NVLinkTwoSided.__new__(NVLinkTwoSided) + owner.ep_size = 2 + owner._dispatch_state = {} + instances.add(owner) + + with pytest.raises(RuntimeError, match="communicator is not initialized"): + owner.checkpoint_prepare() + + checkpoint_prepare.assert_not_called() + + +def test_two_sided_checkpoint_restore_resets_all_shared_owners( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instances = WeakSet() + monkeypatch.setattr(NVLinkTwoSided, "_INSTANCES", instances) + checkpoint_restore = Mock() + monkeypatch.setattr(mnnvl.MnnvlMoe, "checkpoint_restore", checkpoint_restore) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_workspace", + Mock(mapped=False), + ) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_prepare_workspace", + Mock(mapped=False), + ) + first = NVLinkTwoSided.__new__(NVLinkTwoSided) + first._dispatch_state = {"alltoall_info": object()} + second = NVLinkTwoSided.__new__(NVLinkTwoSided) + second._dispatch_state = {"alltoall_info": object()} + instances.update((first, second)) + comm = Mock() + + first.checkpoint_restore(comm) + + checkpoint_restore.assert_called_once_with(comm) + assert first._dispatch_state == {} + assert second._dispatch_state == {} + + +def test_two_sided_checkpoint_restore_noop_preserves_shared_owner_state( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instances = WeakSet() + monkeypatch.setattr(NVLinkTwoSided, "_INSTANCES", instances) + checkpoint_restore = Mock() + monkeypatch.setattr(mnnvl.MnnvlMoe, "checkpoint_restore", checkpoint_restore) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_workspace", + Mock(mapped=True), + ) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_prepare_workspace", + Mock(mapped=True), + ) + first = NVLinkTwoSided.__new__(NVLinkTwoSided) + first._dispatch_state = {"alltoall_info": object()} + second = NVLinkTwoSided.__new__(NVLinkTwoSided) + second._dispatch_state = {"alltoall_info": object()} + instances.update((first, second)) + comm = Mock() + + first.checkpoint_restore(comm) + + checkpoint_restore.assert_called_once_with(comm) + assert first._dispatch_state + assert second._dispatch_state + + +@pytest.mark.parametrize("wrapper_type", [MoeAlltoAll, NVLinkOneSided]) +def test_frontend_checkpoint_delegates_to_shared_lifecycle( + wrapper_type: type[MoeAlltoAll] | type[NVLinkOneSided], +) -> None: + wrapper = wrapper_type.__new__(wrapper_type) + wrapper._workspace_lifecycle = Mock() + comm = Mock() + + wrapper.checkpoint_prepare() + wrapper.checkpoint_restore(comm) + + wrapper._workspace_lifecycle.checkpoint_prepare.assert_called_once_with() + wrapper._workspace_lifecycle.checkpoint_restore.assert_called_once() + assert wrapper._workspace_lifecycle.checkpoint_restore.call_args.args[0] is comm + + +@pytest.mark.parametrize("wrapper_type", [MoeAlltoAll, NVLinkOneSided]) +def test_frontend_destroy_unregisters_from_shared_lifecycle( + wrapper_type: type[MoeAlltoAll] | type[NVLinkOneSided], +) -> None: + wrapper = wrapper_type.__new__(wrapper_type) + wrapper._destroyed = False + lifecycle = Mock() + wrapper._workspace_lifecycle = lifecycle + if wrapper_type is NVLinkOneSided: + wrapper._workspace_key = None + + wrapper.destroy() + wrapper.destroy() + + lifecycle.unregister.assert_called_once_with(wrapper) + + +def test_one_sided_checkpoint_rejects_destroyed_workspace( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACES", {}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE_REFCOUNTS", {}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE", None) + wrapper = NVLinkOneSided.__new__(NVLinkOneSided) + wrapper._destroyed = False + wrapper._workspace_lifecycle = Mock() + wrapper._workspace_key = ("test",) + wrapper._workspace_registered = True + wrapper.destroy() + + with pytest.raises(RuntimeError, match="workspace has been destroyed"): + wrapper.checkpoint_prepare() + + +def test_one_sided_finalizer_unregisters_from_shared_lifecycle() -> None: + class _Lifecycle: + def __init__(self) -> None: + self.unregister_count = 0 + + def unregister(self, client: object) -> None: + self.unregister_count += 1 + + lifecycle = _Lifecycle() + wrapper = NVLinkOneSided.__new__(NVLinkOneSided) + wrapper._destroyed = False + wrapper._workspace_lifecycle = lifecycle + wrapper._workspace_key = None + + del wrapper + gc.collect() + + assert lifecycle.unregister_count == 1 + + +def test_one_sided_aborted_construction_does_not_release_sibling_workspace( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace_key = ("shared",) + workspace_state = {} + refcounts = {workspace_key: 2} + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACES", {workspace_key: workspace_state}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE_REFCOUNTS", refcounts) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE", workspace_state) + wrapper = NVLinkOneSided.__new__(NVLinkOneSided) + wrapper._destroyed = False + wrapper._workspace_lifecycle = Mock() + wrapper._workspace_key = workspace_key + wrapper._workspace_registered = False + + wrapper.destroy() + + assert refcounts[workspace_key] == 2 + assert NVLinkOneSided._WORKSPACES[workspace_key] is workspace_state + assert NVLinkOneSided._WORKSPACE is workspace_state + + +def test_one_sided_failed_registration_does_not_publish_new_workspace( + monkeypatch: pytest.MonkeyPatch, +) -> None: + class _Memory: + mapped = True + + @staticmethod + def initialize() -> None: + pass + + def __init__(self, mapping: object, size: int) -> None: + self.comm = _FakeComm() + + def as_torch_strided_tensor(self, dtype: torch.dtype) -> torch.Tensor: + return torch.zeros(1, dtype=torch.uint8) + + lifecycle = Mock() + lifecycle.register.side_effect = RuntimeError("registration failed") + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACES", {}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE_REFCOUNTS", {}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE", None) + monkeypatch.setattr(NVLinkOneSided, "is_platform_supported", Mock(return_value=True)) + monkeypatch.setattr(NVLinkOneSided, "_init_constants", Mock()) + monkeypatch.setattr(NVLinkOneSided, "FLAG_VAL_OFFSET_INDEX", 0) + monkeypatch.setattr(NVLinkOneSided, "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX", 0) + monkeypatch.setattr(NVLinkOneSided, "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX", 0) + monkeypatch.setattr(one_sided_module, "MnnvlMemory", _Memory) + monkeypatch.setattr( + _MnnvlAlltoAllWorkspaceLifecycle, + "get_or_create", + Mock(return_value=lifecycle), + ) + monkeypatch.setattr( + torch.ops.trtllm, + "moe_a2a_initialize", + Mock(return_value=torch.tensor([1])), + ) + mapping = SimpleNamespace(world_size=2, moe_ep_size=2, moe_ep_rank=0) + + with pytest.raises(RuntimeError, match="registration failed"): + NVLinkOneSided( + mapping=mapping, + num_slots=2, + top_k=1, + max_num_tokens_per_rank=1, + ) + + assert NVLinkOneSided._WORKSPACES == {} + assert NVLinkOneSided._WORKSPACE_REFCOUNTS == {} + assert NVLinkOneSided._WORKSPACE is None diff --git a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py index 91bb82f64055..d62dea1b8bfa 100644 --- a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py +++ b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py @@ -19,7 +19,7 @@ import pytest import tensorrt_llm._mnnvl_utils as mnnvl -from tensorrt_llm._torch.distributed.moe_alltoall import MoeAlltoAll, _A2AState +from tensorrt_llm._torch.distributed.moe_alltoall import MoeAlltoAll class _FakeComm: @@ -98,14 +98,48 @@ def test_checkpoint_restore_reuses_layout_with_fresh_handles(memory, monkeypatch create_and_map = Mock(return_value=[33, 44]) monkeypatch.setattr(_TestMnnvlMemory, "_create_and_map_handles", create_and_map) - obj.checkpoint_restore(restored_comm) + assert obj.checkpoint_restore(restored_comm) create_and_map.assert_called_once_with(restored_comm, 64, 1000, 256, 32) assert obj.ptr == 1032 - assert obj.mapped + assert not obj.mapped + assert record.state is mnnvl._MnnvlAllocationState.RESTORING assert record.mem_handles == [33, 44] assert record.comm is restored_comm assert _TestMnnvlMemory.comm is restored_comm + obj._checkpoint_restore_complete() + assert obj.mapped + + +def test_checkpoint_prepare_failure_is_terminal_and_fails_closed(memory): + obj, record = memory + mnnvl.cuda.cuMemUnmap.side_effect = [None, RuntimeError("unmap failed")] + + with pytest.raises(RuntimeError, match="unmap failed"): + obj.checkpoint_prepare() + + assert not obj.mapped + assert record.state is mnnvl._MnnvlAllocationState.BROKEN + with pytest.raises(RuntimeError, match="broken state"): + obj.checkpoint_prepare() + + +def test_checkpoint_restore_failure_is_terminal_and_fails_closed(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + monkeypatch.setattr( + _TestMnnvlMemory, + "_create_and_map_handles", + Mock(side_effect=RuntimeError("restore failed")), + ) + + with pytest.raises(RuntimeError, match="restore failed"): + obj.checkpoint_restore(_FakeComm()) + + assert not obj.mapped + assert record.state is mnnvl._MnnvlAllocationState.BROKEN + with pytest.raises(RuntimeError, match="broken state"): + obj.checkpoint_restore(_FakeComm()) def test_checkpoint_restore_rejects_changed_rank_layout(memory): @@ -118,6 +152,54 @@ def test_checkpoint_restore_rejects_changed_rank_layout(memory): assert not obj.mapped +def test_mnnvl_moe_restore_publishes_all_workspaces_after_frontend_ready(monkeypatch): + first = Mock() + first.checkpoint_restore.return_value = True + second = Mock() + second.checkpoint_restore.return_value = True + workspace_tensor = Mock() + mapping = SimpleNamespace(moe_ep_rank=0, moe_ep_size=2) + initialize_workspace = Mock() + synchronize = Mock() + comm = _FakeComm() + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace", first) + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_prepare_workspace", second) + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace_tensor", workspace_tensor) + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_mapping", mapping) + monkeypatch.setattr( + mnnvl.torch.ops.trtllm, + "moe_initialize_workspace", + initialize_workspace, + ) + monkeypatch.setattr(mnnvl.torch.cuda, "synchronize", synchronize) + + mnnvl.MnnvlMoe.checkpoint_restore(comm) + + first.checkpoint_restore.assert_called_once_with(comm) + second.checkpoint_restore.assert_called_once_with(comm) + initialize_workspace.assert_called_once_with(workspace_tensor, 0, 2) + synchronize.assert_called_once_with() + assert comm.barrier_count == 1 + first._checkpoint_restore_complete.assert_called_once_with() + second._checkpoint_restore_complete.assert_called_once_with() + + +def test_mnnvl_moe_restore_failure_marks_earlier_workspace_broken(monkeypatch): + first = Mock() + first.checkpoint_restore.return_value = True + second = Mock() + second.checkpoint_restore.side_effect = RuntimeError("second restore failed") + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace", first) + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_prepare_workspace", second) + + with pytest.raises(RuntimeError, match="second restore failed"): + mnnvl.MnnvlMoe.checkpoint_restore(_FakeComm()) + + first._checkpoint_restore_failed.assert_called_once_with() + first._checkpoint_restore_complete.assert_not_called() + second._checkpoint_restore_complete.assert_not_called() + + def test_close_detached_memory_only_frees_va(memory, monkeypatch): obj, record = memory obj.checkpoint_prepare() @@ -212,32 +294,10 @@ def test_create_and_map_handles_releases_local_handle_on_export_failure(monkeypa def _make_moe_alltoall_for_lifecycle(): obj = MoeAlltoAll.__new__(MoeAlltoAll) obj._destroyed = True - obj._state = _A2AState() - obj._alltoall_watchdog = None obj.mnnvl_mem = Mock(mapped=True) return obj -def test_moe_alltoall_checkpoint_prepare_delegates_with_shared_owners(): - first = _make_moe_alltoall_for_lifecycle() - second = _make_moe_alltoall_for_lifecycle() - shared_memory = Mock(mapped=True) - first.mnnvl_mem = shared_memory - second.mnnvl_mem = shared_memory - - first.checkpoint_prepare() - second.checkpoint_prepare() - - assert shared_memory.checkpoint_prepare.call_count == 2 - - -def test_moe_alltoall_checkpoint_prepare_rejects_active_phase(): - obj = _make_moe_alltoall_for_lifecycle() - obj._state.phase = "dispatched" - with pytest.raises(RuntimeError, match="active MoE All-to-All phase"): - obj.checkpoint_prepare() - - def test_moe_alltoall_rejects_unmapped_workspace(): obj = _make_moe_alltoall_for_lifecycle() obj.mnnvl_mem.mapped = False From 60e9fea3a196c0c5122ff3a7c6083acceb7d6ab1 Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Fri, 31 Jul 2026 11:18:20 -0400 Subject: [PATCH 05/14] feat: address mr comments Signed-off-by: Hannah Zhang --- tensorrt_llm/_mnnvl_utils.py | 25 ++++++- .../_torch/mnnvl_alltoall_workspace.py | 65 +++++++++++++---- .../communication/nvlink_one_sided.py | 18 ++++- .../_torch/test_mnnvl_alltoall_workspace.py | 69 ++++++++++++++++++- .../_torch/test_mnnvl_memory_lifecycle.py | 25 ++++++- 5 files changed, 177 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index 515f049de691..4a0bdcace285 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/tensorrt_llm/_mnnvl_utils.py @@ -65,12 +65,14 @@ class _MnnvlAllocationRecord: comm: Any comm_size: int comm_rank: int + comm_membership: tuple[int, ...] aligned_size: int mem_handles: List[Any] start_address: int rank_stride: int address_offset: int state: _MnnvlAllocationState = _MnnvlAllocationState.MAPPED + pending_comm: Any = None class MnnvlMemory: @@ -336,6 +338,12 @@ def open_mnnvl_memory(cls, mapping: Mapping, size: int): comm = cls.get_comm(mapping) comm_rank = comm.Get_rank() comm_size = comm.Get_size() + comm_membership = tuple(int(rank) for rank in comm.allgather(mapping.rank)) + if len(comm_membership) != comm_size: + raise RuntimeError( + "MNNVL communicator membership size does not match its rank count: " + f"{len(comm_membership)} != {comm_size}" + ) all_rank_allocate_sizes = comm.allgather(size) assert len(all_rank_allocate_sizes) == comm_size assert all(x == size for x in all_rank_allocate_sizes), "Not all rank allocating same size." @@ -384,6 +392,7 @@ def open_mnnvl_memory(cls, mapping: Mapping, size: int): comm=comm, comm_size=comm_size, comm_rank=comm_rank, + comm_membership=comm_membership, aligned_size=aligned_size, mem_handles=mem_handles, start_address=cls.current_start_address, @@ -469,6 +478,13 @@ def checkpoint_restore(self, comm) -> bool: f"rank/size {comm_rank}/{comm_size} != " f"{record.comm_rank}/{record.comm_size}" ) + comm_membership = tuple(int(rank) for rank in comm.allgather(self.mapping.rank)) + if comm_membership != record.comm_membership: + raise RuntimeError( + "Cannot restore MNNVL memory with a communicator whose ordered " + "membership differs from the graph-visible allocation layout: " + f"{comm_membership} != {record.comm_membership}" + ) record.state = _MnnvlAllocationState.RESTORING try: torch.cuda.synchronize() @@ -482,8 +498,7 @@ def checkpoint_restore(self, comm) -> bool: except Exception: record.state = _MnnvlAllocationState.BROKEN raise - record.comm = comm - cls.comm = comm + record.pending_comm = comm return True def _checkpoint_restore_complete(self) -> None: @@ -491,6 +506,11 @@ def _checkpoint_restore_complete(self) -> None: record = type(self).allocated_map[self.ptr] if record.state is not _MnnvlAllocationState.RESTORING: raise RuntimeError(f"Cannot complete MNNVL restore in {record.state.value} state") + if record.pending_comm is None: + raise RuntimeError("Cannot complete MNNVL restore without a replacement communicator") + record.comm = record.pending_comm + type(self).comm = record.pending_comm + record.pending_comm = None record.state = _MnnvlAllocationState.MAPPED def _checkpoint_restore_failed(self) -> None: @@ -498,6 +518,7 @@ def _checkpoint_restore_failed(self) -> None: record = type(self).allocated_map[self.ptr] if record.state is _MnnvlAllocationState.RESTORING: record.state = _MnnvlAllocationState.BROKEN + record.pending_comm = None @staticmethod @functools.cache diff --git a/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py b/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py index e800c95cc7cb..c48212d64cfe 100644 --- a/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py +++ b/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py @@ -55,6 +55,22 @@ def _collect_active_ranks( return [rank for rank, clients_idle in enumerate(clients_idle_by_rank) if not clients_idle] +def _collect_unready_ranks( + comm: _CollectiveCommunicator, + *, + local_ready: bool, + expected_size: int, +) -> list[int]: + """Collectively return ranks that could not finish local restore work.""" + ready_by_rank = comm.allgather(local_ready) + if len(ready_by_rank) != expected_size: + raise RuntimeError( + "MNNVL workspace communicator size does not match the MoE EP group: " + f"{len(ready_by_rank)} != {expected_size}" + ) + return [rank for rank, ready in enumerate(ready_by_rank) if not ready] + + class _WorkspaceClient(Protocol): def _mnnvl_checkpoint_is_idle(self) -> bool: """Return whether this client can safely enter a checkpoint.""" @@ -75,8 +91,10 @@ class _WatchdogConfig: class _MnnvlAlltoAllWorkspaceLifecycle: """Own checkpoint and watchdog transitions for one shared MoE workspace. - The checkpoint coordinator must prevent new dispatches from starting while - a transition is in progress. + A top-level engine checkpoint coordinator must atomically stop admission and + drain or abort in-flight work before invoking this resource hook. The local + client and rank-wide idle checks are preflight validation only; they do not + prevent a new dispatch from starting after the idle vote. """ def __init__( @@ -199,7 +217,7 @@ def unregister(self, client: _WorkspaceClient) -> None: self._watchdog_config = None def checkpoint_prepare(self) -> None: - """Stop shared readers, then collectively detach the backing handles.""" + """Preflight shared readers, then collectively detach backing handles.""" if not self._memory.mapped: self._stop_watchdog() self._memory.checkpoint_prepare() @@ -233,21 +251,38 @@ def checkpoint_restore( restored = self._memory.checkpoint_restore(comm) if restored is False: return + local_error: Exception | None = None try: - refreshed_metainfo = initialize_frontend() - if not torch.equal(refreshed_metainfo, self._metainfo): + try: + refreshed_metainfo = initialize_frontend() + if not torch.equal(refreshed_metainfo, self._metainfo): + raise RuntimeError( + "MoE All-to-All metainfo changed during MNNVL restore; " + "captured CUDA graphs cannot be replayed safely" + ) + self._metainfo = refreshed_metainfo + self._workspace_state["metainfo"] = refreshed_metainfo + self._coordinator = self._create_coordinator() + torch.cuda.synchronize() + self._start_watchdog() + for client in list(self._clients): + client._mnnvl_checkpoint_reset() + except Exception as error: + local_error = error + + unready_ranks = _collect_unready_ranks( + comm, + local_ready=local_error is None, + expected_size=self._ep_size, + ) + if unready_ranks: + self._stop_watchdog() + if local_error is not None: + raise local_error raise RuntimeError( - "MoE All-to-All metainfo changed during MNNVL restore; " - "captured CUDA graphs cannot be replayed safely" + "MNNVL workspace restore failed on ranks " + f"{unready_ranks}; refusing to publish the restored workspace" ) - self._metainfo = refreshed_metainfo - self._workspace_state["metainfo"] = refreshed_metainfo - self._coordinator = self._create_coordinator() - torch.cuda.synchronize() - comm.barrier() - self._start_watchdog() - for client in list(self._clients): - client._mnnvl_checkpoint_reset() self._memory._checkpoint_restore_complete() except Exception: self._memory._checkpoint_restore_failed() 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 909a21b245fe..c4e3ce49cbcb 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 @@ -548,7 +548,7 @@ def supports_post_quant_dispatch(self) -> bool: return True def destroy(self): - """Release this instance's reference to the shared symmetric workspace.""" + """Release shared state during explicit, rank-coordinated teardown.""" if getattr(self, "_destroyed", False): return @@ -583,8 +583,20 @@ def destroy(self): self._dispatch_state = {"phase": "destroyed"} def __del__(self) -> None: - if not sys.is_finalizing(): - self.destroy() + if sys.is_finalizing(): + return + lifecycle = getattr(self, "_workspace_lifecycle", None) + if lifecycle is not None: + lifecycle.unregister(self) + workspace_key = getattr(self, "_workspace_key", None) + if workspace_key is not None and getattr(self, "_workspace_registered", False): + refcount = NVLinkOneSided._WORKSPACE_REFCOUNTS.get(workspace_key, 0) - 1 + if refcount > 0: + NVLinkOneSided._WORKSPACE_REFCOUNTS[workspace_key] = refcount + else: + NVLinkOneSided._WORKSPACE_REFCOUNTS.pop(workspace_key, None) + self._workspace_registered = False + self._workspace_lifecycle = None def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) -> bool: """ diff --git a/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py b/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py index 9325e126f31c..801769c66c48 100644 --- a/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py +++ b/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py @@ -47,16 +47,23 @@ def _mnnvl_checkpoint_reset(self) -> None: class _FakeComm: - def __init__(self, clients_idle_by_rank: list[bool] | None = None) -> None: + def __init__( + self, + clients_idle_by_rank: list[bool] | None = None, + gathered_values: list[list[bool]] | None = None, + ) -> None: self.barrier_count = 0 self.allgather_count = 0 self.clients_idle_by_rank = clients_idle_by_rank + self.gathered_values = list(gathered_values or []) def barrier(self) -> None: self.barrier_count += 1 def allgather(self, local_clients_idle: bool) -> list[bool]: self.allgather_count += 1 + if self.gathered_values: + return self.gathered_values.pop(0) if self.clients_idle_by_rank is not None: return self.clients_idle_by_rank return [local_clients_idle, local_clients_idle] @@ -242,7 +249,7 @@ def test_checkpoint_suspends_and_recreates_one_shared_watchdog( assert lifecycle.watchdog_for(second) is new_watchdog assert first.reset_count == 1 assert second.reset_count == 1 - assert comm.barrier_count == 1 + assert comm.allgather_count == 1 def test_unregister_stops_watchdog_after_last_enabled_client() -> None: @@ -407,6 +414,43 @@ def test_checkpoint_restore_failure_after_watchdog_start_stops_it( memory._checkpoint_restore_complete.assert_not_called() +def test_checkpoint_restore_remote_failure_fails_closed_on_every_rank( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle, memory, metainfo = _make_lifecycle() + client = _Client() + _register_without_watchdog(lifecycle, client) + memory.mapped = False + memory.checkpoint_restore.return_value = True + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + comm = _FakeComm(gathered_values=[[True, False]]) + + with pytest.raises(RuntimeError, match=r"restore failed on ranks \[1\]"): + lifecycle.checkpoint_restore(comm, Mock(return_value=metainfo)) + + memory._checkpoint_restore_failed.assert_called_once_with() + memory._checkpoint_restore_complete.assert_not_called() + + +def test_checkpoint_restore_reports_local_failure_to_every_rank( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle, memory, _ = _make_lifecycle() + memory.mapped = False + memory.checkpoint_restore.return_value = True + monkeypatch.setattr(torch.cuda, "synchronize", Mock()) + comm = _FakeComm(gathered_values=[[False, True]]) + + with pytest.raises(RuntimeError, match="frontend restore failed"): + lifecycle.checkpoint_restore( + comm, + Mock(side_effect=RuntimeError("frontend restore failed")), + ) + + assert comm.allgather_count == 1 + memory._checkpoint_restore_failed.assert_called_once_with() + + def test_two_sided_checkpoint_prepare_rejects_active_shared_owner( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -624,6 +668,27 @@ def unregister(self, client: object) -> None: assert lifecycle.unregister_count == 1 +def test_one_sided_finalizer_preserves_collective_workspace_cache( + monkeypatch: pytest.MonkeyPatch, +) -> None: + workspace_key = ("shared",) + workspace_state = {"memory": object()} + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACES", {workspace_key: workspace_state}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE_REFCOUNTS", {workspace_key: 1}) + monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE", workspace_state) + wrapper = NVLinkOneSided.__new__(NVLinkOneSided) + wrapper._workspace_lifecycle = Mock() + wrapper._workspace_key = workspace_key + wrapper._workspace_registered = True + + del wrapper + gc.collect() + + assert NVLinkOneSided._WORKSPACES[workspace_key] is workspace_state + assert NVLinkOneSided._WORKSPACE is workspace_state + assert workspace_key not in NVLinkOneSided._WORKSPACE_REFCOUNTS + + def test_one_sided_aborted_construction_does_not_release_sibling_workspace( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py index d62dea1b8bfa..6608bbc17a6c 100644 --- a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py +++ b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py @@ -23,9 +23,10 @@ class _FakeComm: - def __init__(self, rank=0, size=2): + def __init__(self, rank=0, size=2, membership=(0, 1)): self.rank = rank self.size = size + self.membership = membership self.barrier_count = 0 def Get_rank(self): @@ -37,6 +38,9 @@ def Get_size(self): def barrier(self): self.barrier_count += 1 + def allgather(self, value): + return list(self.membership) + class _TestMnnvlMemory(mnnvl.MnnvlMemory): pass @@ -49,6 +53,7 @@ def memory(monkeypatch): comm=comm, comm_size=2, comm_rank=0, + comm_membership=(0, 1), aligned_size=64, mem_handles=[11, 22], start_address=1000, @@ -57,6 +62,7 @@ def memory(monkeypatch): ) obj = _TestMnnvlMemory.__new__(_TestMnnvlMemory) obj.ptr = 1032 + obj.mapping = SimpleNamespace(rank=0) _TestMnnvlMemory.allocated_map = {obj.ptr: record} _TestMnnvlMemory.address_refcnt = {record.start_address: 1} _TestMnnvlMemory.current_start_address = record.start_address @@ -105,10 +111,23 @@ def test_checkpoint_restore_reuses_layout_with_fresh_handles(memory, monkeypatch assert not obj.mapped assert record.state is mnnvl._MnnvlAllocationState.RESTORING assert record.mem_handles == [33, 44] - assert record.comm is restored_comm - assert _TestMnnvlMemory.comm is restored_comm + assert record.comm is not restored_comm + assert record.pending_comm is restored_comm obj._checkpoint_restore_complete() assert obj.mapped + assert record.comm is restored_comm + assert _TestMnnvlMemory.comm is restored_comm + + +def test_checkpoint_restore_rejects_changed_ordered_membership(memory): + obj, record = memory + obj.checkpoint_prepare() + + with pytest.raises(RuntimeError, match="ordered membership differs"): + obj.checkpoint_restore(_FakeComm(membership=(1, 0))) + + assert record.state is mnnvl._MnnvlAllocationState.UNMAPPED + assert record.pending_comm is None def test_checkpoint_prepare_failure_is_terminal_and_fails_closed(memory): From acc4d5e2ab5668ed03a77f34ab2e418674ad924d Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Wed, 12 Aug 2026 16:31:07 -0600 Subject: [PATCH 06/14] fix(moe): harden shared MNNVL restore cleanup Signed-off-by: Hannah Zhang --- tensorrt_llm/_mnnvl_utils.py | 22 ++++++++++ .../integration/test_lists/test-db/l0_a10.yml | 1 + .../_torch/test_mnnvl_memory_lifecycle.py | 43 +++++++++++++++++++ 3 files changed, 66 insertions(+) diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index 4a0bdcace285..21efff8f7836 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/tensorrt_llm/_mnnvl_utils.py @@ -517,6 +517,28 @@ def _checkpoint_restore_failed(self) -> None: """Make a failed frontend restore terminal and fail closed.""" record = type(self).allocated_map[self.ptr] if record.state is _MnnvlAllocationState.RESTORING: + for rank, handle in enumerate(record.mem_handles): + rank_ptr = record.start_address + rank * record.rank_stride + record.address_offset + try: + _check_cu_result(cuda.cuMemUnmap(rank_ptr, record.aligned_size)) + except RuntimeError as error: + logger.warning( + "Failed to unmap unpublished MNNVL restore for rank %d: %s", + rank, + error, + ) + if handle is None: + continue + try: + _check_cu_result(cuda.cuMemRelease(handle)) + except RuntimeError as error: + logger.warning( + "Failed to release unpublished MNNVL restore handle for rank %d: %s", + rank, + error, + ) + else: + record.mem_handles[rank] = None record.state = _MnnvlAllocationState.BROKEN record.pending_comm = None diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 5b39eea23efa..323a6823881c 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -18,6 +18,7 @@ l0_a10: - unittest/_torch/sampler/test_penalties.py - unittest/_torch/test_tensor_lru_cache.py - unittest/_torch/test_torch_multi_arange.py + - unittest/_torch/test_mnnvl_alltoall_workspace.py - unittest/utils/test_util.py - unittest/_torch/modeling/test_modeling_mistral.py - unittest/_torch/modeling/test_modeling_pixtral.py diff --git a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py index 6608bbc17a6c..c044fbdb979a 100644 --- a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py +++ b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py @@ -161,6 +161,49 @@ def test_checkpoint_restore_failure_is_terminal_and_fails_closed(memory, monkeyp obj.checkpoint_restore(_FakeComm()) +def test_failed_frontend_restore_releases_unpublished_handles(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + monkeypatch.setattr( + _TestMnnvlMemory, + "_create_and_map_handles", + Mock(return_value=[33, 44]), + ) + + assert obj.checkpoint_restore(_FakeComm()) + obj._checkpoint_restore_failed() + + assert record.state is mnnvl._MnnvlAllocationState.BROKEN + assert record.mem_handles == [None, None] + assert record.pending_comm is None + assert [call.args[0] for call in mnnvl.cuda.cuMemUnmap.call_args_list[-2:]] == [1032, 1288] + assert [call.args[0] for call in mnnvl.cuda.cuMemRelease.call_args_list[-2:]] == [33, 44] + + +def test_failed_frontend_restore_cleanup_continues_after_cuda_errors(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + monkeypatch.setattr( + _TestMnnvlMemory, + "_create_and_map_handles", + Mock(return_value=[33, 44]), + ) + + assert obj.checkpoint_restore(_FakeComm()) + mnnvl.cuda.cuMemUnmap.reset_mock() + mnnvl.cuda.cuMemRelease.reset_mock() + mnnvl.cuda.cuMemUnmap.side_effect = [RuntimeError("unmap failed"), None] + mnnvl.cuda.cuMemRelease.side_effect = [RuntimeError("release failed"), None] + + obj._checkpoint_restore_failed() + + assert record.state is mnnvl._MnnvlAllocationState.BROKEN + assert record.mem_handles == [33, None] + assert record.pending_comm is None + assert [call.args[0] for call in mnnvl.cuda.cuMemUnmap.call_args_list] == [1032, 1288] + assert [call.args[0] for call in mnnvl.cuda.cuMemRelease.call_args_list] == [33, 44] + + def test_checkpoint_restore_rejects_changed_rank_layout(memory): obj, _ = memory obj.checkpoint_prepare() From 71600c15c805bcf608021c13320dc1970da2b6a7 Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Wed, 12 Aug 2026 16:02:47 -0600 Subject: [PATCH 07/14] fix(moe): address MNNVL lifecycle review feedback Signed-off-by: Hannah Zhang --- tensorrt_llm/_mnnvl_utils.py | 29 +++- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../test-db/l0_gb200_multi_gpus.yml | 1 + .../_torch/modules/moe/test_moe_comm.py | 138 +++++++++++++++++- .../_torch/test_mnnvl_memory_lifecycle.py | 103 +++++++++++++ 5 files changed, 265 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index 21efff8f7836..91dc635751f7 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/tensorrt_llm/_mnnvl_utils.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import ctypes +import errno import functools import os import platform @@ -272,10 +273,18 @@ def _create_and_map_handles( remote_fd = syscall(438, pidfd, fd, 0) if remote_fd < 0: err = ctypes.get_errno() - raise RuntimeError( + error_msg = ( f"pidfd_getfd(pidfd={pidfd}, fd={fd}) failed with errno " - f"{err}: {os.strerror(err)}" + f"{err}: {os.strerror(err)}." ) + if err == errno.EPERM: + error_msg += ( + " Permission denied. If running in a container, try adding " + "--cap-add=SYS_PTRACE to your docker run command." + ) + elif err == errno.ENOSYS: + error_msg += " This may be due to kernel version (requires Linux 5.6+)." + raise RuntimeError(error_msg) remote_fds.append(remote_fd) comm.barrier() all_handles_data = remote_fds @@ -317,11 +326,20 @@ def _create_and_map_handles( raise finally: for pidfd in pidfds: - os.close(pidfd) + try: + os.close(pidfd) + except OSError as error: + logger.warning("Failed to close MNNVL pidfd: %s", error) for remote_fd in remote_fds: - os.close(remote_fd) + try: + os.close(remote_fd) + except OSError as error: + logger.warning("Failed to close imported MNNVL file descriptor: %s", error) if not is_fabric and exported_handle is not None: - os.close(int(exported_handle)) + try: + os.close(int(exported_handle)) + except OSError as error: + logger.warning("Failed to close exported MNNVL file descriptor: %s", error) @classmethod def open_mnnvl_memory(cls, mapping: Mapping, size: int): @@ -385,7 +403,6 @@ def open_mnnvl_memory(cls, mapping: Mapping, size: int): cls.current_mem_offset, ) = previous_address_state raise - # all_handles_data like b'\x00\x00\x00 \x00\x00\x00\x00\x8f\xec\x02\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\t\x00\x00\x00\x00\x00\x1d\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00' # noqa: E501 ptr = cls.current_start_address + cls.current_mem_offset stride = cls.current_rank_stride cls.allocated_map[ptr] = _MnnvlAllocationRecord( diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 323a6823881c..c2ed366ccf46 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -19,6 +19,7 @@ l0_a10: - unittest/_torch/test_tensor_lru_cache.py - unittest/_torch/test_torch_multi_arange.py - unittest/_torch/test_mnnvl_alltoall_workspace.py + - unittest/_torch/test_mnnvl_memory_lifecycle.py - unittest/utils/test_util.py - unittest/_torch/modeling/test_modeling_mistral.py - unittest/_torch/modeling/test_modeling_pixtral.py diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index dceecb3761be..6882aee2741c 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -55,6 +55,7 @@ l0_gb200_multi_gpus: - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_nccl_ep_cuda_graph_replay_uses_updated_routing - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_postquant + - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_preserves_moe_graph_addresses - disaggregated/test_disaggregated.py::test_disaggregated_overlap_transceiver_runtime_python_fabric_memory[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_overlap_transceiver_runtime_python_bounce[TinyLlama-1.1B-Chat-v1.0] diff --git a/tests/unittest/_torch/modules/moe/test_moe_comm.py b/tests/unittest/_torch/modules/moe/test_moe_comm.py index 4c3bbb84ce79..85b9c67aa68d 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_comm.py +++ b/tests/unittest/_torch/modules/moe/test_moe_comm.py @@ -62,7 +62,7 @@ from mpi4py import MPI import tensorrt_llm as tllm -from tensorrt_llm._mnnvl_utils import MnnvlMemory +from tensorrt_llm._mnnvl_utils import MnnvlMemory, MnnvlMoe from tensorrt_llm._torch.modules.fused_moe.communication.allgather_reducescatter import ( AllGatherReduceScatter, ) @@ -2559,6 +2559,113 @@ def _run_rank_mask_one_rank_masked_test( assert saw_dead, f"dead rank {dead_rank} did not appear in results" +def _worker_mnnvl_checkpoint_graph_replay(config: CommTestConfig) -> bool: + """Exercise stable-VA checkpoint cycles through a captured MoE communication graph.""" + rank = tllm.mpi_rank() + torch.cuda.set_device(rank) + mapping = Mapping( + rank=rank, + tp_size=config.ep_size, + moe_ep_size=config.ep_size, + world_size=config.ep_size, + ) + communication = create_comm_object(config.comm_type, mapping, config) + worker_inputs = _prepare_worker_inputs(rank, config) + _, local_slot_start, local_slot_end = _compute_ep_partition( + config.num_experts, config.ep_size, rank + ) + + def forward() -> torch.Tensor: + dispatch_outputs = _run_worker_dispatch(communication, worker_inputs, config) + received_hidden_states = _to_bf16( + dispatch_outputs.recv_hs, + dispatch_outputs.recv_sf, + worker_inputs.global_scale, + config.quant_mode, + ) + moe_output = simple_moe( + received_hidden_states, + dispatch_outputs.recv_slots, + dispatch_outputs.recv_scales, + local_slot_start, + local_slot_end, + ) + return communication.combine( + moe_output, + all_rank_max_num_tokens=max(config.all_num_tokens), + ) + + with torch.inference_mode(): + forward() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + graph_output = forward() + graph.replay() + torch.cuda.synchronize() + + if config.comm_type == COMM_NVLINK_ONE_SIDED: + stable_addresses = (communication.mnnvl_mem.ptr, communication.workspace.data_ptr()) + + def allocations_are_mapped() -> bool: + return communication.mnnvl_mem.mapped + + else: + stable_addresses = ( + communication.alltoall_workspace.data_ptr(), + communication.alltoall_prepare_workspace.data_ptr(), + ) + + def allocations_are_mapped() -> bool: + return all( + workspace is None or workspace.mapped + for workspace in (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) + ) + + def any_allocation_is_mapped() -> bool: + return any( + workspace is not None and workspace.mapped + for workspace in (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) + ) + + if config.comm_type == COMM_NVLINK_ONE_SIDED: + + def any_allocation_is_mapped() -> bool: + return communication.mnnvl_mem.mapped + + for cycle in range(3): + free_before_prepare = torch.cuda.mem_get_info()[0] + communication.checkpoint_prepare() + assert not any_allocation_is_mapped() + fresh_comm = MPI.COMM_WORLD.Split(0, rank) + communication.checkpoint_restore(fresh_comm) + assert allocations_are_mapped() + memory_drift_bytes = free_before_prepare - torch.cuda.mem_get_info()[0] + assert memory_drift_bytes <= 64 * 1024**2 + + if config.comm_type == COMM_NVLINK_ONE_SIDED: + restored_addresses = ( + communication.mnnvl_mem.ptr, + communication.workspace.data_ptr(), + ) + else: + restored_addresses = ( + communication.alltoall_workspace.data_ptr(), + communication.alltoall_prepare_workspace.data_ptr(), + ) + assert restored_addresses == stable_addresses + + with torch.inference_mode(): + worker_inputs.hs.add_(cycle + 1) + expected = forward().clone() + graph.replay() + torch.cuda.synchronize() + torch.testing.assert_close(graph_output, expected) + + if hasattr(communication, "destroy"): + communication.destroy() + return True + + # ============================================================================ # Test Class # ============================================================================ @@ -2624,6 +2731,35 @@ def test_nccl_ep_cuda_graph_replay_uses_updated_routing(self, mpi_pool_executor) """Verify LL CUDA graph replay reads routing written after capture.""" _run_nccl_ep_cuda_graph_replay_test(mpi_pool_executor) + @pytest.mark.threadleak(enabled=False) + @pytest.mark.parametrize("mpi_pool_executor", [4], indirect=True) + @pytest.mark.parametrize( + "comm_type", + [COMM_NVLINK_ONE_SIDED, COMM_NVLINK_TWO_SIDED], + ) + def test_mnnvl_checkpoint_preserves_moe_graph_addresses( + self, + mpi_pool_executor, + comm_type: str, + ) -> None: + """Verify repeated checkpoint cycles preserve MoE graph pointers and outputs.""" + config = CommTestConfig( + comm_type=comm_type, + ep_size=mpi_pool_executor.num_workers, + num_experts=32, + top_k=2, + hidden_size=128, + all_num_tokens=[8] * mpi_pool_executor.num_workers, + ) + skip_reason = _get_skip_reason(config) + if skip_reason: + pytest.skip(skip_reason) + results = mpi_pool_executor.map( + _worker_mnnvl_checkpoint_graph_replay, + [config] * config.ep_size, + ) + assert all(results) + @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize( "mpi_pool_executor,local_num_tokens,top_k", diff --git a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py index c044fbdb979a..d7ebc7a9a82a 100644 --- a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py +++ b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py @@ -17,9 +17,11 @@ from unittest.mock import Mock import pytest +import torch import tensorrt_llm._mnnvl_utils as mnnvl from tensorrt_llm._torch.distributed.moe_alltoall import MoeAlltoAll +from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_two_sided import NVLinkTwoSided class _FakeComm: @@ -353,6 +355,82 @@ def test_create_and_map_handles_releases_local_handle_on_export_failure(monkeypa release_memory.assert_called_once_with(11) +@pytest.mark.parametrize( + ("errno_value", "expected_hint"), + [(mnnvl.errno.EPERM, "--cap-add=SYS_PTRACE"), (mnnvl.errno.ENOSYS, "requires Linux 5.6+")], +) +def test_create_and_map_handles_preserves_pidfd_error_hint(monkeypatch, errno_value, expected_hint): + comm = Mock() + comm.Get_rank.return_value = 0 + comm.Get_size.return_value = 2 + comm.allgather.side_effect = [[100, 101], [1000, 1001]] + allocation_prop = SimpleNamespace( + requestedHandleTypes=( + mnnvl.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + ), + location=object(), + ) + syscall = Mock(side_effect=[10, 11, -1]) + + monkeypatch.setattr(mnnvl.MnnvlMemory, "dev_id", 0) + monkeypatch.setattr( + mnnvl.MnnvlMemory, + "get_allocation_prop", + Mock(return_value=allocation_prop), + ) + monkeypatch.setattr(mnnvl, "_check_cu_result", lambda result: result) + monkeypatch.setattr(mnnvl.cuda, "cuCtxGetDevice", Mock(return_value=0)) + monkeypatch.setattr(mnnvl.cuda, "cuMemCreate", Mock(return_value=11)) + monkeypatch.setattr(mnnvl.cuda, "cuMemExportToShareableHandle", Mock(return_value=100)) + monkeypatch.setattr(mnnvl.cuda, "cuMemRelease", Mock(return_value=None)) + monkeypatch.setattr(mnnvl.ctypes, "CDLL", Mock(return_value=SimpleNamespace(syscall=syscall))) + monkeypatch.setattr(mnnvl.ctypes, "get_errno", Mock(return_value=errno_value)) + close_fd = Mock() + monkeypatch.setattr(mnnvl.os, "close", close_fd) + + with pytest.raises(RuntimeError, match=expected_hint): + _TestMnnvlMemory._create_and_map_handles(comm, 64, 1000, 256, 32) + + assert [call.args[0] for call in close_fd.call_args_list] == [10, 11, 100] + + +def test_create_and_map_handles_close_failure_does_not_mask_original_error(monkeypatch): + comm = Mock() + comm.Get_rank.return_value = 0 + comm.Get_size.return_value = 2 + comm.allgather.side_effect = [[100, 101], [1000, 1001]] + allocation_prop = SimpleNamespace( + requestedHandleTypes=( + mnnvl.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR + ), + location=object(), + ) + access_desc = SimpleNamespace(location=None, flags=None) + syscall = Mock(side_effect=[10, 11, 20, 21]) + + monkeypatch.setattr(mnnvl.MnnvlMemory, "dev_id", 0) + monkeypatch.setattr( + mnnvl.MnnvlMemory, + "get_allocation_prop", + Mock(return_value=allocation_prop), + ) + monkeypatch.setattr(mnnvl, "_check_cu_result", lambda result: result) + monkeypatch.setattr(mnnvl.cuda, "cuCtxGetDevice", Mock(return_value=0)) + monkeypatch.setattr(mnnvl.cuda, "cuMemCreate", Mock(return_value=11)) + monkeypatch.setattr(mnnvl.cuda, "cuMemExportToShareableHandle", Mock(return_value=100)) + monkeypatch.setattr(mnnvl.cuda, "CUmemAccessDesc", Mock(return_value=access_desc)) + monkeypatch.setattr(mnnvl.cuda, "cuMemMap", Mock(side_effect=RuntimeError("map failed"))) + monkeypatch.setattr(mnnvl.cuda, "cuMemRelease", Mock(return_value=None)) + monkeypatch.setattr(mnnvl.ctypes, "CDLL", Mock(return_value=SimpleNamespace(syscall=syscall))) + close_fd = Mock(side_effect=[OSError("close failed"), None, None, None, None]) + monkeypatch.setattr(mnnvl.os, "close", close_fd) + + with pytest.raises(RuntimeError, match="map failed"): + _TestMnnvlMemory._create_and_map_handles(comm, 64, 1000, 256, 32) + + assert [call.args[0] for call in close_fd.call_args_list] == [10, 11, 20, 21, 100] + + def _make_moe_alltoall_for_lifecycle(): obj = MoeAlltoAll.__new__(MoeAlltoAll) obj._destroyed = True @@ -366,3 +444,28 @@ def test_moe_alltoall_rejects_unmapped_workspace(): with pytest.raises(RuntimeError, match="workspace handles are unmapped"): obj._require_mapped() + + +def test_two_sided_combine_requires_new_prepare_before_next_dispatch(monkeypatch): + communication = NVLinkTwoSided.__new__(NVLinkTwoSided) + communication._dispatch_state = { + "alltoall_info": object(), + "original_token_count": 1, + } + communication.alltoall_workspace = object() + communication.ep_rank = 0 + communication.ep_size = 1 + communication.top_k = 1 + communication.use_low_precision_combine = False + communication.alltoall_result_do_sum = True + monkeypatch.setattr(mnnvl.MnnvlMoe, "require_mapped", Mock()) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "mnnvl_moe_alltoallv_combine", + Mock(return_value=torch.ones(1, 1)), + ) + + communication.combine(torch.ones(1, 1)) + + with pytest.raises(ValueError, match=r"requires prepare_dispatch\(\) to be called first"): + communication.dispatch(None, None, None, None, [1]) From 8e0c1b9286470dc78c8708f35b70790a1784f3d3 Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Mon, 17 Aug 2026 11:46:22 -0400 Subject: [PATCH 08/14] fix(moe): coordinate MNNVL checkpoint lifecycle Signed-off-by: Hannah Zhang --- tensorrt_llm/_mnnvl_utils.py | 308 +++++++++++++---- .../_torch/distributed/moe_alltoall.py | 25 +- .../_torch/mnnvl_alltoall_workspace.py | 39 ++- .../communication/communication_factory.py | 6 + .../communication/nvlink_one_sided.py | 27 +- .../communication/nvlink_two_sided.py | 38 +- .../nvlink_two_sided_flashinfer.py | 5 + tensorrt_llm/_torch/pyexecutor/py_executor.py | 62 +++- tensorrt_llm/executor/base_worker.py | 113 +++--- .../integration/test_lists/test-db/l0_a10.yml | 1 + .../test-db/l0_gb200_multi_gpus.yml | 2 + .../modules/moe/test_communication_factory.py | 117 +++++++ .../modules/moe/test_kimi_k3_situ_moe.py | 1 + .../_torch/modules/moe/test_moe_comm.py | 325 +++++++++++++++++- .../_torch/test_mnnvl_alltoall_workspace.py | 100 +++++- .../_torch/test_mnnvl_memory_lifecycle.py | 186 +++++++++- .../test_sleep_collective_rpc_guards.py | 173 ++++++++++ 17 files changed, 1367 insertions(+), 161 deletions(-) diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index 91dc635751f7..ffa354770d66 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/tensorrt_llm/_mnnvl_utils.py @@ -18,6 +18,7 @@ import os import platform import sys +import time from dataclasses import dataclass from enum import Enum from typing import Any, List, Optional, Union @@ -35,6 +36,30 @@ from .logger import logger from .mapping import Mapping +_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S = 20.0 +_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S = 0.01 + + +def _checkpoint_allgather(comm, value, *, operation: str): + """Run a bounded object allgather when the communicator supports it.""" + iallgather = getattr(comm, "iallgather", None) + if iallgather is None: + return comm.allgather(value) + + request = iallgather(value) + deadline = time.monotonic() + _MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S + while True: + test_result = request.test() + if isinstance(test_result, tuple): + ready, result = test_result + if ready: + return result + elif test_result: + return request.wait() + if time.monotonic() >= deadline: + raise TimeoutError(f"Timed out waiting for MNNVL checkpoint {operation} allgather") + time.sleep(_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S) + def _check_cu_result(cu_func_ret): if isinstance(cu_func_ret, tuple): @@ -234,59 +259,117 @@ def _create_and_map_handles( rank_stride: int, address_offset: int, ) -> List[Any]: - dev_id = int(_check_cu_result(cuda.cuCtxGetDevice())) - assert dev_id == MnnvlMemory.dev_id, ( - f"Different dev_id found dev_id={dev_id} but MnnvlMemory.dev_id={MnnvlMemory.dev_id}" - ) - allocation_prop = MnnvlMemory.get_allocation_prop(dev_id) - local_handle = _check_cu_result(cuda.cuMemCreate(aligned_size, allocation_prop, flags=0)) + local_handle = None exported_handle = None pidfds = [] remote_fds = [] mem_handles = [None] * comm.Get_size() mapped_rank_ptrs = [] - is_fabric = ( - allocation_prop.requestedHandleTypes - == cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC - ) + is_fabric = False try: - exported_handle = _check_cu_result( - cuda.cuMemExportToShareableHandle( - local_handle, allocation_prop.requestedHandleTypes, 0 + local_error = None + local_handle_data = None + local_pid = None + try: + dev_id = int(_check_cu_result(cuda.cuCtxGetDevice())) + assert dev_id == MnnvlMemory.dev_id, ( + f"Different dev_id found dev_id={dev_id} but " + f"MnnvlMemory.dev_id={MnnvlMemory.dev_id}" + ) + allocation_prop = MnnvlMemory.get_allocation_prop(dev_id) + is_fabric = ( + allocation_prop.requestedHandleTypes + == cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC + ) + local_handle = _check_cu_result( + cuda.cuMemCreate(aligned_size, allocation_prop, flags=0) ) + exported_handle = _check_cu_result( + cuda.cuMemExportToShareableHandle( + local_handle, allocation_prop.requestedHandleTypes, 0 + ) + ) + local_handle_data = exported_handle.data if is_fabric else int(exported_handle) + local_pid = os.getpid() + except Exception as error: + local_error = f"{type(error).__name__}: {error}" + + exported_by_rank = _checkpoint_allgather( + comm, + { + "error": local_error, + "handle": local_handle_data, + "is_fabric": is_fabric, + "pid": local_pid, + }, + operation="handle export", ) + export_errors = [ + f"rank {rank}: {payload['error']}" + for rank, payload in enumerate(exported_by_rank) + if payload["error"] is not None + ] + if export_errors: + raise RuntimeError( + "MNNVL handle export failed before mapping:\n" + "\n".join(export_errors) + ) + fabric_modes = {payload["is_fabric"] for payload in exported_by_rank} + if len(fabric_modes) != 1: + raise RuntimeError("MNNVL ranks selected inconsistent shareable handle types") + is_fabric = fabric_modes.pop() if is_fabric: - all_handles_data = comm.allgather(exported_handle.data) + all_handles_data = [payload["handle"] for payload in exported_by_rank] else: - all_exported_fds = comm.allgather(int(exported_handle)) - all_pids = comm.allgather(os.getpid()) + all_exported_fds = [payload["handle"] for payload in exported_by_rank] + all_pids = [payload["pid"] for payload in exported_by_rank] syscall = ctypes.CDLL(None, use_errno=True).syscall - for pid in all_pids: - pidfd = syscall(434, pid, 0) - if pidfd < 0: - err = ctypes.get_errno() - raise RuntimeError( - f"pidfd_open({pid}) failed with errno {err}: {os.strerror(err)}" - ) - pidfds.append(pidfd) - for pidfd, fd in zip(pidfds, all_exported_fds): - remote_fd = syscall(438, pidfd, fd, 0) - if remote_fd < 0: - err = ctypes.get_errno() - error_msg = ( - f"pidfd_getfd(pidfd={pidfd}, fd={fd}) failed with errno " - f"{err}: {os.strerror(err)}." - ) - if err == errno.EPERM: - error_msg += ( - " Permission denied. If running in a container, try adding " - "--cap-add=SYS_PTRACE to your docker run command." + fd_import_error = None + try: + for pid in all_pids: + pidfd = syscall(434, pid, 0) + if pidfd < 0: + err = ctypes.get_errno() + raise RuntimeError( + f"pidfd_open({pid}) failed with errno {err}: {os.strerror(err)}" + ) + pidfds.append(pidfd) + for pidfd, fd in zip(pidfds, all_exported_fds): + remote_fd = syscall(438, pidfd, fd, 0) + if remote_fd < 0: + err = ctypes.get_errno() + error_msg = ( + f"pidfd_getfd(pidfd={pidfd}, fd={fd}) failed with errno " + f"{err}: {os.strerror(err)}." ) - elif err == errno.ENOSYS: - error_msg += " This may be due to kernel version (requires Linux 5.6+)." - raise RuntimeError(error_msg) - remote_fds.append(remote_fd) - comm.barrier() + if err == errno.EPERM: + error_msg += ( + " Permission denied. If running in a container, try adding " + "--cap-add=SYS_PTRACE to your docker run command." + ) + elif err == errno.ENOSYS: + error_msg += ( + " This may be due to kernel version (requires Linux 5.6+)." + ) + raise RuntimeError(error_msg) + remote_fds.append(remote_fd) + except Exception as error: + fd_import_error = f"{type(error).__name__}: {error}" + + fd_import_errors = _checkpoint_allgather( + comm, + fd_import_error, + operation="POSIX file descriptor import readiness", + ) + failed_ranks = [ + f"rank {rank}: {error}" + for rank, error in enumerate(fd_import_errors) + if error is not None + ] + if failed_ranks: + raise RuntimeError( + "MNNVL POSIX file descriptor import failed on one or more ranks:\n" + + "\n".join(failed_ranks) + ) all_handles_data = remote_fds access_desc = cuda.CUmemAccessDesc() @@ -313,33 +396,33 @@ def _create_and_map_handles( try: _check_cu_result(cuda.cuMemUnmap(rank_ptr, aligned_size)) except RuntimeError as error: - logger.warning("Failed to unmap incomplete MNNVL allocation: %s", error) + logger.warning(f"Failed to unmap incomplete MNNVL allocation: {error}") handles_to_release = [handle for handle in mem_handles if handle is not None] - if mem_handles[comm.Get_rank()] is None: + if local_handle is not None and mem_handles[comm.Get_rank()] is None: handles_to_release.append(local_handle) for handle in handles_to_release: try: _check_cu_result(cuda.cuMemRelease(handle)) except RuntimeError as error: - logger.warning("Failed to release incomplete MNNVL allocation: %s", error) + logger.warning(f"Failed to release incomplete MNNVL allocation: {error}") raise finally: for pidfd in pidfds: try: os.close(pidfd) except OSError as error: - logger.warning("Failed to close MNNVL pidfd: %s", error) + logger.warning(f"Failed to close MNNVL pidfd: {error}") for remote_fd in remote_fds: try: os.close(remote_fd) except OSError as error: - logger.warning("Failed to close imported MNNVL file descriptor: %s", error) + logger.warning(f"Failed to close imported MNNVL file descriptor: {error}") if not is_fabric and exported_handle is not None: try: os.close(int(exported_handle)) except OSError as error: - logger.warning("Failed to close exported MNNVL file descriptor: %s", error) + logger.warning(f"Failed to close exported MNNVL file descriptor: {error}") @classmethod def open_mnnvl_memory(cls, mapping: Mapping, size: int): @@ -395,7 +478,7 @@ def open_mnnvl_memory(cls, mapping: Mapping, size: int): cuda.cuMemAddressFree(device_ptr, comm_size * cls.current_rank_stride) ) except RuntimeError as error: - logger.warning("cuMemAddressFree failed during error cleanup: %s", error) + logger.warning(f"cuMemAddressFree failed during error cleanup: {error}") else: ( cls.current_start_address, @@ -431,8 +514,7 @@ def close_mnnvl_memory(cls, ptr: int): _MnnvlAllocationState.UNMAPPED, ): logger.warning( - "Skipping cleanup of MNNVL allocation in terminal state %s", - record.state.value, + f"Skipping cleanup of MNNVL allocation in terminal state {record.state.value}" ) return cls.allocated_map.pop(ptr) @@ -453,13 +535,32 @@ def close_mnnvl_memory(cls, ptr: int): @classmethod def _unmap_and_release_handles(cls, record: _MnnvlAllocationRecord) -> None: + first_error = None for rank in range(record.comm_size): rank_ptr = record.start_address + rank * record.rank_stride + record.address_offset - _check_cu_result(cuda.cuMemUnmap(rank_ptr, record.aligned_size)) - _check_cu_result(cuda.cuMemRelease(record.mem_handles[rank])) + try: + _check_cu_result(cuda.cuMemUnmap(rank_ptr, record.aligned_size)) + except RuntimeError as error: + if first_error is None: + first_error = error + continue + try: + _check_cu_result(cuda.cuMemRelease(record.mem_handles[rank])) + except RuntimeError as error: + if first_error is None: + first_error = error + else: + record.mem_handles[rank] = None + if first_error is not None: + raise first_error def checkpoint_prepare(self) -> None: - """Collectively detach backing handles while retaining graph-visible VA.""" + """Detach local backing handles while retaining graph-visible VA. + + The engine checkpoint coordinator is responsible for quiescing and + invoking this operation on every rank. No collective follows the CUDA + mutation, so a local failure cannot strand peers in a trailing barrier. + """ cls = type(self) record = cls.allocated_map[self.ptr] if record.state is _MnnvlAllocationState.UNMAPPED: @@ -469,15 +570,24 @@ def checkpoint_prepare(self) -> None: record.state = _MnnvlAllocationState.PREPARING try: torch.cuda.synchronize() - record.comm.barrier() cls._unmap_and_release_handles(record) record.mem_handles = [None] * record.comm_size - record.comm.barrier() except Exception: record.state = _MnnvlAllocationState.BROKEN raise record.state = _MnnvlAllocationState.UNMAPPED + def checkpoint_fail_closed(self) -> None: + """Make a timed-out checkpoint allocation terminal. + + A timed-out nonblocking MPI collective may still be outstanding. The + allocation must not be reused by a later checkpoint attempt because a + new collective could then be matched against the abandoned request. + """ + record = type(self).allocated_map[self.ptr] + record.state = _MnnvlAllocationState.BROKEN + record.pending_comm = None + def checkpoint_restore(self, comm) -> bool: """Remap fresh handles while keeping data-path access disabled.""" cls = type(self) @@ -495,7 +605,18 @@ def checkpoint_restore(self, comm) -> bool: f"rank/size {comm_rank}/{comm_size} != " f"{record.comm_rank}/{record.comm_size}" ) - comm_membership = tuple(int(rank) for rank in comm.allgather(self.mapping.rank)) + try: + comm_membership = tuple( + int(rank) + for rank in _checkpoint_allgather( + comm, + self.mapping.rank, + operation="communicator membership", + ) + ) + except Exception: + self.checkpoint_fail_closed() + raise if comm_membership != record.comm_membership: raise RuntimeError( "Cannot restore MNNVL memory with a communicator whose ordered " @@ -503,6 +624,7 @@ def checkpoint_restore(self, comm) -> bool: f"{comm_membership} != {record.comm_membership}" ) record.state = _MnnvlAllocationState.RESTORING + local_error = None try: torch.cuda.synchronize() record.mem_handles = cls._create_and_map_handles( @@ -512,9 +634,31 @@ def checkpoint_restore(self, comm) -> bool: record.rank_stride, record.address_offset, ) - except Exception: + except TimeoutError: record.state = _MnnvlAllocationState.BROKEN raise + except Exception as error: + local_error = f"{type(error).__name__}: {error}" + + try: + restore_errors = _checkpoint_allgather( + comm, + local_error, + operation="mapping readiness", + ) + except Exception: + self._checkpoint_restore_failed() + raise + failed_ranks = [ + f"rank {rank}: {error}" + for rank, error in enumerate(restore_errors) + if error is not None + ] + if failed_ranks: + self._checkpoint_restore_failed() + raise RuntimeError( + "MNNVL checkpoint restore failed on one or more ranks:\n" + "\n".join(failed_ranks) + ) record.pending_comm = comm return True @@ -535,24 +679,21 @@ def _checkpoint_restore_failed(self) -> None: record = type(self).allocated_map[self.ptr] if record.state is _MnnvlAllocationState.RESTORING: for rank, handle in enumerate(record.mem_handles): + if handle is None: + continue rank_ptr = record.start_address + rank * record.rank_stride + record.address_offset try: _check_cu_result(cuda.cuMemUnmap(rank_ptr, record.aligned_size)) except RuntimeError as error: logger.warning( - "Failed to unmap unpublished MNNVL restore for rank %d: %s", - rank, - error, + f"Failed to unmap unpublished MNNVL restore for rank {rank}: {error}" ) - if handle is None: - continue try: _check_cu_result(cuda.cuMemRelease(handle)) except RuntimeError as error: logger.warning( - "Failed to release unpublished MNNVL restore handle for rank %d: %s", - rank, - error, + "Failed to release unpublished MNNVL restore handle " + f"for rank {rank}: {error}" ) else: record.mem_handles[rank] = None @@ -772,15 +913,36 @@ def checkpoint_restore(comm) -> None: restored_workspaces.append(workspace) if not restored_workspaces: return - if MnnvlMoe.moe_workspace_tensor is not None: - assert MnnvlMoe.moe_mapping is not None - torch.ops.trtllm.moe_initialize_workspace( - MnnvlMoe.moe_workspace_tensor, - MnnvlMoe.moe_mapping.moe_ep_rank, - MnnvlMoe.moe_mapping.moe_ep_size, + restored_main_workspace = any( + workspace is MnnvlMoe.moe_workspace for workspace in restored_workspaces + ) + local_error = None + try: + if restored_main_workspace and MnnvlMoe.moe_workspace_tensor is not None: + assert MnnvlMoe.moe_mapping is not None + torch.ops.trtllm.moe_initialize_workspace( + MnnvlMoe.moe_workspace_tensor, + MnnvlMoe.moe_mapping.moe_ep_rank, + MnnvlMoe.moe_mapping.moe_ep_size, + ) + torch.cuda.synchronize() + except Exception as error: + local_error = f"{type(error).__name__}: {error}" + readiness_errors = _checkpoint_allgather( + comm, + local_error, + operation="two-sided frontend readiness", + ) + failed_ranks = [ + f"rank {rank}: {error}" + for rank, error in enumerate(readiness_errors) + if error is not None + ] + if failed_ranks: + raise RuntimeError( + "Native two-sided MoE restore failed on one or more ranks:\n" + + "\n".join(failed_ranks) ) - torch.cuda.synchronize() - comm.barrier() except Exception: for workspace in restored_workspaces: workspace._checkpoint_restore_failed() diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index 44d573337ba3..f3328c08c990 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -397,13 +397,15 @@ def __init__( ep_size=self.ep_size, health=self.ep_group_health, )) + self._destroyed = False + self._workspace_registered = False self._workspace_lifecycle.register( self, watchdog_timeout_s=alltoall_watchdog_timeout_s, watchdog_poll_interval_s=alltoall_watchdog_poll_interval_s, watchdog_on_timeout=alltoall_watchdog_on_timeout, ) - self._destroyed = False + self._workspace_registered = True @property def metainfo(self) -> torch.Tensor: @@ -423,8 +425,10 @@ def destroy(self) -> None: return self._destroyed = True lifecycle = getattr(self, "_workspace_lifecycle", None) - if lifecycle is not None: + if lifecycle is not None and getattr(self, "_workspace_registered", + False): lifecycle.unregister(self) + self._workspace_registered = False def __del__(self) -> None: if not sys.is_finalizing(): @@ -472,8 +476,21 @@ def checkpoint_prepare(self) -> None: ) self._workspace_lifecycle.checkpoint_prepare() - def checkpoint_restore(self, comm) -> None: - """Collectively restore handles and all shared frontend state.""" + def checkpoint_restore(self, comm=None) -> None: + """Collectively restore handles and all shared frontend state. + + Args: + comm: An mpi4py-like communicator exposing ``Get_rank()``, + ``Get_size()``, ``allgather()``, and ``barrier()``. Its local + rank and size must match the communicator used for the + original allocation. Every rank must call this method + symmetrically. + """ + if comm is None: + comm = self.mnnvl_mem.comm + if comm is None: + raise RuntimeError( + "MNNVL workspace communicator is not initialized") self._workspace_lifecycle.checkpoint_restore( comm, lambda: torch.ops.trtllm.moe_a2a_initialize( diff --git a/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py b/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py index c48212d64cfe..a5e1e025673a 100644 --- a/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py +++ b/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py @@ -19,7 +19,7 @@ import torch -from tensorrt_llm._mnnvl_utils import MnnvlMemory +from tensorrt_llm._mnnvl_utils import MnnvlMemory, _checkpoint_allgather from tensorrt_llm._torch.alltoall_watchdog import ( AlltoAllWatchdog, AlltoAllWatchdogCoordinator, @@ -46,7 +46,11 @@ def _collect_active_ranks( expected_size: int, ) -> list[int]: """Collectively return ranks whose local workspace clients are active.""" - clients_idle_by_rank = comm.allgather(local_clients_idle) + clients_idle_by_rank = _checkpoint_allgather( + comm, + local_clients_idle, + operation="workspace idle readiness", + ) if len(clients_idle_by_rank) != expected_size: raise RuntimeError( "MNNVL workspace communicator size does not match the MoE EP group: " @@ -62,7 +66,11 @@ def _collect_unready_ranks( expected_size: int, ) -> list[int]: """Collectively return ranks that could not finish local restore work.""" - ready_by_rank = comm.allgather(local_ready) + ready_by_rank = _checkpoint_allgather( + comm, + local_ready, + operation="frontend restore readiness", + ) if len(ready_by_rank) != expected_size: raise RuntimeError( "MNNVL workspace communicator size does not match the MoE EP group: " @@ -91,10 +99,12 @@ class _WatchdogConfig: class _MnnvlAlltoAllWorkspaceLifecycle: """Own checkpoint and watchdog transitions for one shared MoE workspace. - A top-level engine checkpoint coordinator must atomically stop admission and - drain or abort in-flight work before invoking this resource hook. The local - client and rank-wide idle checks are preflight validation only; they do not - prevent a new dispatch from starting after the idle vote. + ``PyExecutor`` invokes this resource hook from the engine sleep/wakeup + PREPARE/COMMIT/ABORT control path. That coordinator stops admission, drains + in-flight work, aggregates bounded per-rank results, and reopens admission + only after every rank has completed COMMIT. The local client and subgroup + idle checks here remain defense-in-depth validation while the executor is + parked. """ def __init__( @@ -228,11 +238,16 @@ def checkpoint_prepare(self) -> None: comm = cast(_CollectiveCommunicator | None, self._memory.comm) if comm is None: raise RuntimeError("MNNVL workspace communicator is not initialized") - active_ranks = _collect_active_ranks( - comm, - local_clients_idle=local_clients_idle, - expected_size=self._ep_size, - ) + try: + active_ranks = _collect_active_ranks( + comm, + local_clients_idle=local_clients_idle, + expected_size=self._ep_size, + ) + except TimeoutError: + self._memory.checkpoint_fail_closed() + self._stop_watchdog() + raise if active_ranks: raise RuntimeError( f"Cannot checkpoint during an active MoE All-to-All phase on ranks {active_ranks}" 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 0afe4697cf25..c7628cc1c7cf 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.py @@ -127,6 +127,12 @@ def create_strategy( """ # Extract parameters from model_config mapping = model_config.mapping + if mapping.has_cp_helix(): + raise ValueError( + "MoE communication requires Helix context-parallel ranks to be " + "repurposed through Mapping.repurpose_helix_cp_to_tp() before " + "strategy selection" + ) if hidden_size is None: hidden_size = model_config.pretrained_config.hidden_size act_dtype = model_config.torch_dtype 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 c4e3ce49cbcb..65055d76a205 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 @@ -298,6 +298,11 @@ def __init__( alltoall_watchdog_on_timeout: Optional callback invoked when the watchdog reports suspects. """ super().__init__(mapping) + if mapping.has_cp_helix(): + raise ValueError( + "NVLinkOneSided does not support Helix context parallelism because " + "its MNNVL communicator covers only the tensor-parallel group" + ) if self.mapping.world_size != self.ep_size: raise RuntimeError("Currently NVLinkOneSided only supports pure EP for MoE.") @@ -554,7 +559,7 @@ def destroy(self): self._destroyed = True lifecycle = getattr(self, "_workspace_lifecycle", None) - if lifecycle is not None: + if lifecycle is not None and getattr(self, "_workspace_registered", False): lifecycle.unregister(self) workspace_key = getattr(self, "_workspace_key", None) if workspace_key is None or not getattr(self, "_workspace_registered", False): @@ -585,8 +590,10 @@ def destroy(self): def __del__(self) -> None: if sys.is_finalizing(): return + # Finalizers cannot safely perform collective cache eviction because + # their order is nondeterministic across ranks. lifecycle = getattr(self, "_workspace_lifecycle", None) - if lifecycle is not None: + if lifecycle is not None and getattr(self, "_workspace_registered", False): lifecycle.unregister(self) workspace_key = getattr(self, "_workspace_key", None) if workspace_key is not None and getattr(self, "_workspace_registered", False): @@ -635,8 +642,20 @@ def checkpoint_prepare(self) -> None: ) self._require_workspace_lifecycle().checkpoint_prepare() - def checkpoint_restore(self, comm) -> None: - """Collectively restore handles and all shared frontend state.""" + def checkpoint_restore(self, comm=None) -> None: + """Collectively restore handles and all shared frontend state. + + Args: + comm: An mpi4py-like communicator exposing ``Get_rank()``, + ``Get_size()``, ``allgather()``, and ``barrier()``. Its local + rank and size must match the communicator used for the + original allocation. Every rank must call this method + symmetrically. + """ + if comm is None: + comm = self.mnnvl_mem.comm + if comm is None: + raise RuntimeError("MNNVL workspace communicator is not initialized") self._require_workspace_lifecycle().checkpoint_restore( comm, lambda: torch.ops.trtllm.moe_a2a_initialize( diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py index 5eb0f0cd3812..d249ab0b2ce9 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py @@ -56,6 +56,11 @@ def __init__( alltoall_result_do_sum: bool = False, ): super().__init__(mapping) + if mapping.has_cp_helix(): + raise ValueError( + "NVLinkTwoSided does not support Helix context parallelism because " + "its MNNVL communicator covers only the tensor-parallel group" + ) # Store needed parameters self.num_experts = num_experts @@ -116,19 +121,38 @@ def checkpoint_prepare(self) -> None: comm = workspace.comm if comm is None: raise RuntimeError("MNNVL workspace communicator is not initialized") - active_ranks = _collect_active_ranks( - comm, - local_clients_idle=local_clients_idle, - expected_size=self.ep_size, - ) + try: + active_ranks = _collect_active_ranks( + comm, + local_clients_idle=local_clients_idle, + expected_size=self.ep_size, + ) + except TimeoutError: + for candidate in workspaces: + if candidate is not None: + candidate.checkpoint_fail_closed() + raise if active_ranks: raise RuntimeError( f"Cannot checkpoint during an active MoE All-to-All phase on ranks {active_ranks}" ) MnnvlMoe.checkpoint_prepare() - def checkpoint_restore(self, comm) -> None: - """Restore TRT-native two-sided workspaces and protocol state.""" + def checkpoint_restore(self, comm=None) -> None: + """Restore TRT-native two-sided workspaces and protocol state. + + Args: + comm: An mpi4py-like communicator exposing ``Get_rank()``, + ``Get_size()``, ``allgather()``, and ``barrier()``. Its local + rank and size must match the communicator used for the + original allocations. Every rank must call this method + symmetrically. + """ + workspace = MnnvlMoe.moe_workspace or MnnvlMoe.moe_prepare_workspace + if comm is None and workspace is not None: + comm = workspace.comm + if comm is None: + raise RuntimeError("MNNVL workspace communicator is not initialized") restore_required = any( workspace is not None and not workspace.mapped for workspace in (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py index 6963a046dd68..b3aaa2f8ad0e 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided_flashinfer.py @@ -84,6 +84,11 @@ def __init__( alltoall_result_do_sum: bool = False, ): super().__init__(mapping) + if mapping.has_cp_helix(): + raise ValueError( + "NVLinkTwoSidedFlashinfer does not support Helix context parallelism " + "because its MNNVL communicator covers only the tensor-parallel group" + ) # Store needed parameters self.num_experts = num_experts diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 19ab0cffc873..70a337335d17 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3041,6 +3041,7 @@ def _sleep_wakeup_listener_loop(self) -> None: op_id = msg.get("op_id") error_msg = None release_control_request = True + has_mnnvl_resources = False try: # Decode tags inside the try so KeyError / ValueError from # a malformed message still results in an error ACK being @@ -3092,8 +3093,12 @@ def _sleep_wakeup_listener_loop(self) -> None: if action == _SleepWakeupAction.PREPARE: # Prepared means this rank is quiesced and # ready to commit, but VMM state is unchanged. + has_mnnvl_resources = ( + self._has_mnnvl_checkpoint_resources(tags)) release_control_request = False elif target_action == _SleepWakeupAction.SLEEP: + self._run_mnnvl_checkpoint_resources( + target_action, tags) release_with_tag(*tags) torch.cuda.synchronize() gc.collect() @@ -3101,6 +3106,8 @@ def _sleep_wakeup_listener_loop(self) -> None: elif target_action == _SleepWakeupAction.WAKEUP: materialize_with_tag(*tags) torch.cuda.synchronize() + self._run_mnnvl_checkpoint_resources( + target_action, tags) else: error_msg = ( f"unknown target action '{target_action}'") @@ -3112,7 +3119,7 @@ def _sleep_wakeup_listener_loop(self) -> None: logger.warning("Sleep/wakeup listener: %s, ignoring.", error_msg) except (KeyError, TypeError, ValueError, RuntimeError, - torch.OutOfMemoryError) as exc: + TimeoutError, torch.OutOfMemoryError) as exc: error_msg = (f"rank {self.dist.rank} '{action}' failed: " f"{exc}\n{traceback.format_exc()}") logger.error("Sleep/wakeup listener: error executing '%s':", @@ -3156,6 +3163,7 @@ def _sleep_wakeup_listener_loop(self) -> None: "error": error_msg, "op_id": op_id, "phase": action, + "has_mnnvl_resources": has_mnnvl_resources, }, dest=0, tag=_SleepWakeupTag.ACK, @@ -3163,6 +3171,58 @@ def _sleep_wakeup_listener_loop(self) -> None: finally: set_thread_local_mpi_comm(None) + def _mnnvl_checkpoint_resources(self, tags) -> list: + """Return unique native MNNVL MoE resources selected by engine tags.""" + from tensorrt_llm._torch.distributed.moe_alltoall import MoeAlltoAll + from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_one_sided import \ + NVLinkOneSided + from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_two_sided import \ + NVLinkTwoSided + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + selected_engines = [] + if ExecutorMemoryType.MODEL_ENGINE_MAIN in tags: + selected_engines.append(self.model_engine) + if ExecutorMemoryType.MODEL_ENGINE_DRAFT in tags and self.draft_model_engine is not None: + selected_engines.append(self.draft_model_engine) + + resource_types = (MoeAlltoAll, NVLinkOneSided, NVLinkTwoSided) + resources = [] + seen = set() + for engine in selected_engines: + model = getattr(engine, "model", None) + if model is None: + continue + for module in model.modules(): + resource = getattr(module, "comm", None) + if not isinstance(resource, resource_types): + continue + lifecycle = getattr(resource, "_workspace_lifecycle", None) + key = (("workspace", + id(lifecycle)) if lifecycle is not None else + ("resource_type", type(resource))) + if key in seen: + continue + seen.add(key) + resources.append(resource) + return resources + + def _has_mnnvl_checkpoint_resources(self, tags) -> bool: + return bool(self._mnnvl_checkpoint_resources(tags)) + + def _run_mnnvl_checkpoint_resources(self, action, tags) -> None: + """Execute process-local native MNNVL hooks while the engine is parked.""" + resources = self._mnnvl_checkpoint_resources(tags) + if action == _SleepWakeupAction.SLEEP: + for resource in resources: + resource.checkpoint_prepare() + return + if action == _SleepWakeupAction.WAKEUP: + for resource in resources: + resource.checkpoint_restore() + return + raise ValueError(f"unknown MNNVL checkpoint action '{action}'") + def _ring_broadcast_sample_state( self, executed_batch: Optional[BatchStatePP], diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 098d4674603d..e56bffd9e0b9 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -647,11 +647,12 @@ def _multi_rank_sleep_wakeup( 2. Send PREPARE to every non-rank-0 rank via the dedicated ``_sleep_wakeup_comm`` communicator. Peers quiesce and ACK without changing VMM state. - 3. Execute the VMM operation (``release_with_tag`` or - ``materialize_with_tag``) locally on rank-0. - 4. Send COMMIT to prepared peers and collect ACKs after their local VMM - operations. If PREPARE or rank-0 local execution fails, send ABORT - so peers leave the control barrier without changing VMM state. + 3. When native MNNVL MoE resources are selected, send COMMIT to every + prepared peer before rank-0 enters the subgroup checkpoint + collectives. Otherwise execute the local VMM operation first. + 4. Execute the selected VMM and MNNVL operations and collect bounded + COMMIT ACKs. If PREPARE or pre-commit local execution fails, send + ABORT so peers leave the control barrier without changing VMM state. 5. Exit ``control_action()``, resuming rank-0's event loop. Args: @@ -683,6 +684,9 @@ def _multi_rank_sleep_wakeup( tag_strings = [t.value for t in tags] op_id = uuid.uuid4().hex target_action = _SleepWakeupAction(action) + has_mnnvl_resources = getattr(self.engine, + "_has_mnnvl_checkpoint_resources", + lambda _tags: False)(tags) prepare_msg = { "action": _SleepWakeupAction.PREPARE, "target_action": target_action, @@ -707,6 +711,8 @@ def _multi_rank_sleep_wakeup( errors = [] local_error = None abort_sent = False + commit_ranks = [] + commit_sent_early = False def send_abort(reason: str, ranks: Optional[list[int]] = None) -> list[int]: @@ -739,7 +745,9 @@ def send_abort(reason: str, ) return abort_ranks - def drain_acks(ranks: list[int], phase: _SleepWakeupAction) -> None: + def drain_acks(ranks: list[int], + phase: _SleepWakeupAction) -> list[dict]: + received_acks = [] ack_deadline = time.monotonic() + _SLEEP_WAKEUP_ACK_TIMEOUT_S for src in ranks: try: @@ -760,10 +768,40 @@ def drain_acks(ranks: list[int], phase: _SleepWakeupAction) -> None: exc_info=True, ) continue + received_acks.append(ack) if ack.get("status") != "ok": errors.append( ack.get("error") or f"rank {src} returned unknown {phase} ACK") + return received_acks + + def send_commits() -> list[int]: + sent_ranks = [] + failed_ranks = [] + for dest in prepared_ranks: + try: + sleep_wakeup_comm.send( + commit_msg, + dest=dest, + tag=_SleepWakeupTag.ACTION, + ) + sent_ranks.append(dest) + except Exception as exc: + commit_error = ( + f"rank 0 failed to send '{action}' commit to " + f"rank {dest}: {exc}") + errors.append(commit_error) + failed_ranks.append(dest) + logger.error( + "_multi_rank_sleep_wakeup: %s", + commit_error, + exc_info=True, + ) + if failed_ranks: + abort_ranks = send_abort("\n".join(errors), + ranks=failed_ranks) + drain_acks(abort_ranks, _SleepWakeupAction.ABORT) + return sent_ranks try: # Phase 1: prepare peers. A prepared peer has reached the @@ -793,15 +831,27 @@ def drain_acks(ranks: list[int], phase: _SleepWakeupAction) -> None: break if not errors: - drain_acks(prepared_ranks, _SleepWakeupAction.PREPARE) + prepare_acks = drain_acks(prepared_ranks, + _SleepWakeupAction.PREPARE) + has_mnnvl_resources = has_mnnvl_resources or any( + ack.get("has_mnnvl_resources", False) + for ack in prepare_acks) if not errors: - # Execute locally on rank-0. Only CUDA/VMM errors are - # captured as local_error. Peers are still prepared but - # uncommitted, so local failure can abort them without - # changing their VMM state. + # MNNVL resource hooks contain subgroup handle exchange, + # so peers must enter COMMIT before rank 0 executes them. + if has_mnnvl_resources: + commit_ranks = send_commits() + commit_sent_early = True + + if not errors or commit_sent_early: torch.cuda.synchronize() if action == _SleepWakeupAction.SLEEP: + run_mnnvl = (getattr( + self.engine, "_run_mnnvl_checkpoint_resources", + None) if has_mnnvl_resources else None) + if run_mnnvl is not None: + run_mnnvl(target_action, tags) release_with_tag(*tags) torch.cuda.synchronize() gc.collect() @@ -809,7 +859,12 @@ def drain_acks(ranks: list[int], phase: _SleepWakeupAction) -> None: else: materialize_with_tag(*tags) torch.cuda.synchronize() - except (RuntimeError, torch.OutOfMemoryError) as exc: + run_mnnvl = (getattr( + self.engine, "_run_mnnvl_checkpoint_resources", + None) if has_mnnvl_resources else None) + if run_mnnvl is not None: + run_mnnvl(target_action, tags) + except (RuntimeError, TimeoutError, torch.OutOfMemoryError) as exc: local_error = (f"rank 0 '{action}' failed: {exc}\n" f"{traceback.format_exc()}") logger.error( @@ -821,35 +876,15 @@ def drain_acks(ranks: list[int], phase: _SleepWakeupAction) -> None: if local_error: errors.append(local_error) - if errors and prepared_ranks and not abort_sent: + if commit_sent_early: + drain_acks(commit_ranks, _SleepWakeupAction.COMMIT) + + if (errors and prepared_ranks and not abort_sent + and not commit_sent_early): abort_ranks = send_abort("\n".join(errors)) drain_acks(abort_ranks, _SleepWakeupAction.ABORT) - elif not errors: - commit_ranks = [] - commit_failed_ranks = [] - for dest in prepared_ranks: - try: - sleep_wakeup_comm.send( - commit_msg, - dest=dest, - tag=_SleepWakeupTag.ACTION, - ) - commit_ranks.append(dest) - except Exception as exc: - commit_error = ( - f"rank 0 failed to send '{action}' commit to " - f"rank {dest}: {exc}") - errors.append(commit_error) - commit_failed_ranks.append(dest) - logger.error( - "_multi_rank_sleep_wakeup: %s", - commit_error, - exc_info=True, - ) - if commit_failed_ranks: - abort_ranks = send_abort("\n".join(errors), - ranks=commit_failed_ranks) - drain_acks(abort_ranks, _SleepWakeupAction.ABORT) + elif not errors and not commit_sent_early: + commit_ranks = send_commits() drain_acks(commit_ranks, _SleepWakeupAction.COMMIT) if errors: diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index c2ed366ccf46..41db4b16b4fb 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -20,6 +20,7 @@ l0_a10: - unittest/_torch/test_torch_multi_arange.py - unittest/_torch/test_mnnvl_alltoall_workspace.py - unittest/_torch/test_mnnvl_memory_lifecycle.py + - unittest/executor/test_sleep_collective_rpc_guards.py - unittest/utils/test_util.py - unittest/_torch/modeling/test_modeling_mistral.py - unittest/_torch/modeling/test_modeling_pixtral.py diff --git a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml index 6882aee2741c..f3597fbe7a9a 100644 --- a/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml +++ b/tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml @@ -56,6 +56,8 @@ l0_gb200_multi_gpus: - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_nccl_ep_cuda_graph_replay_uses_updated_routing - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_moe_comm_postquant - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_preserves_moe_graph_addresses + - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_engine_checkpoint_coordination + - unittest/_torch/modules/moe/test_moe_comm.py::TestMoEComm::test_mnnvl_checkpoint_failure_is_collective_and_bounded - disaggregated/test_disaggregated.py::test_disaggregated_overlap_transceiver_runtime_python_fabric_memory[TinyLlama-1.1B-Chat-v1.0] - disaggregated/test_disaggregated.py::test_disaggregated_overlap_transceiver_runtime_python_bounce[TinyLlama-1.1B-Chat-v1.0] diff --git a/tests/unittest/_torch/modules/moe/test_communication_factory.py b/tests/unittest/_torch/modules/moe/test_communication_factory.py index 3919ea233af4..92bb2544ed6a 100644 --- a/tests/unittest/_torch/modules/moe/test_communication_factory.py +++ b/tests/unittest/_torch/modules/moe/test_communication_factory.py @@ -17,12 +17,18 @@ import sys from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch from tensorrt_llm._torch.modules.fused_moe import nccl_ep_utils from tensorrt_llm._torch.modules.fused_moe.communication import communication_factory +from tensorrt_llm._torch.modules.fused_moe.communication import nvlink_one_sided as one_sided_module +from tensorrt_llm._torch.modules.fused_moe.communication import nvlink_two_sided as two_sided_module +from tensorrt_llm._torch.modules.fused_moe.communication import ( + nvlink_two_sided_flashinfer as flashinfer_module, +) from tensorrt_llm._torch.modules.fused_moe.communication.allgather_reducescatter import ( AllGatherReduceScatter, ) @@ -39,6 +45,7 @@ def _make_model_config( moe_tp_size=1, moe_ep_size=2, moe_ep_rank=0, + has_cp_helix=Mock(return_value=False), ) return SimpleNamespace( mapping=mapping, @@ -57,6 +64,116 @@ def _strategy_unavailable(*args, **kwargs): raise RuntimeError("strategy unavailable") +@pytest.mark.parametrize("use_flashinfer", [False, True]) +def test_forced_two_sided_rejects_helix_before_workspace_allocation( + monkeypatch: pytest.MonkeyPatch, + use_flashinfer: bool, +) -> None: + model_config = _make_model_config() + model_config.mapping.has_cp_helix = Mock(return_value=True) + monkeypatch.setattr( + communication_factory.NVLinkTwoSided, + "is_platform_supported", + Mock(return_value=True), + ) + monkeypatch.setattr( + communication_factory.NVLinkTwoSidedFlashinfer, + "is_platform_supported", + Mock(return_value=True), + ) + native_initialize = Mock(side_effect=AssertionError("native allocation reached")) + flashinfer_symbols = Mock(side_effect=AssertionError("FlashInfer allocation reached")) + monkeypatch.setattr(two_sided_module.MnnvlMemory, "initialize", native_initialize) + monkeypatch.setattr(flashinfer_module, "_flashinfer_mnnvl", flashinfer_symbols) + + with pytest.raises(ValueError, match="does not support Helix context parallelism"): + communication_factory.CommunicationFactory._create_forced_method( + "NVLINK_TWO_SIDED", + model_config, + num_experts=32, + num_slots=32, + top_k=8, + expert_size_per_partition=16, + payload_in_workspace=False, + alltoall_result_do_sum=True, + use_flashinfer=use_flashinfer, + hidden_size=4096, + ) + + native_initialize.assert_not_called() + flashinfer_symbols.assert_not_called() + + +def test_forced_one_sided_rejects_helix_before_workspace_allocation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_config = _make_model_config() + model_config.mapping.world_size = model_config.mapping.moe_ep_size + model_config.mapping.has_cp_helix = Mock(return_value=True) + monkeypatch.setattr( + communication_factory.NVLinkOneSided, + "is_platform_supported", + Mock(return_value=True), + ) + native_initialize = Mock(side_effect=AssertionError("native allocation reached")) + monkeypatch.setattr(one_sided_module.MnnvlMemory, "initialize", native_initialize) + + with pytest.raises(ValueError, match="does not support Helix context parallelism"): + communication_factory.CommunicationFactory._create_forced_method( + "NVLINK_ONE_SIDED", + model_config, + num_experts=32, + num_slots=32, + top_k=8, + expert_size_per_partition=16, + payload_in_workspace=False, + alltoall_result_do_sum=True, + use_flashinfer=False, + hidden_size=4096, + ) + + native_initialize.assert_not_called() + + +def test_auto_selection_rejects_unrepurposed_helix( + monkeypatch: pytest.MonkeyPatch, +) -> None: + model_config = _make_model_config() + model_config.mapping.world_size = model_config.mapping.moe_ep_size + model_config.mapping.has_cp_helix = Mock(return_value=True) + monkeypatch.setattr( + communication_factory.NVLinkOneSided, + "is_platform_supported", + Mock(return_value=True), + ) + monkeypatch.setattr( + communication_factory.NVLinkTwoSided, + "is_platform_supported", + Mock(return_value=True), + ) + monkeypatch.setattr( + communication_factory.NVLinkTwoSidedFlashinfer, + "is_platform_supported", + Mock(return_value=True), + ) + native_initialize = Mock(side_effect=AssertionError("native allocation reached")) + flashinfer_symbols = Mock(side_effect=AssertionError("FlashInfer allocation reached")) + monkeypatch.setattr(one_sided_module.MnnvlMemory, "initialize", native_initialize) + monkeypatch.setattr(flashinfer_module, "_flashinfer_mnnvl", flashinfer_symbols) + with pytest.raises(ValueError, match="repurpose_helix_cp_to_tp"): + communication_factory.CommunicationFactory.create_strategy( + model_config, + num_experts=32, + num_slots=32, + top_k=8, + expert_size_per_partition=16, + hidden_size=4096, + ) + + native_initialize.assert_not_called() + flashinfer_symbols.assert_not_called() + + def _install_failing_nccl_module(monkeypatch: pytest.MonkeyPatch, error: BaseException): def fail_get_version(): raise error diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py index 8b564563a92e..757d06bfe90a 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py @@ -206,6 +206,7 @@ def test_communication_factory_accepts_model_selected_method(monkeypatch): dp_size=16, moe_tp_size=1, moe_ep_size=16, + has_cp_helix=lambda: False, ) model_config = SimpleNamespace( mapping=mapping, diff --git a/tests/unittest/_torch/modules/moe/test_moe_comm.py b/tests/unittest/_torch/modules/moe/test_moe_comm.py index 85b9c67aa68d..3ec329ca75cf 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_comm.py +++ b/tests/unittest/_torch/modules/moe/test_moe_comm.py @@ -50,9 +50,11 @@ import os import pickle import sys +import threading import traceback from dataclasses import dataclass from functools import lru_cache +from types import SimpleNamespace from typing import Dict, List, Optional, Set, Tuple from unittest.mock import MagicMock @@ -62,6 +64,7 @@ from mpi4py import MPI import tensorrt_llm as tllm +import tensorrt_llm._mnnvl_utils as mnnvl from tensorrt_llm._mnnvl_utils import MnnvlMemory, MnnvlMoe from tensorrt_llm._torch.modules.fused_moe.communication.allgather_reducescatter import ( AllGatherReduceScatter, @@ -2559,17 +2562,9 @@ def _run_rank_mask_one_rank_masked_test( assert saw_dead, f"dead rank {dead_rank} did not appear in results" -def _worker_mnnvl_checkpoint_graph_replay(config: CommTestConfig) -> bool: - """Exercise stable-VA checkpoint cycles through a captured MoE communication graph.""" - rank = tllm.mpi_rank() - torch.cuda.set_device(rank) - mapping = Mapping( - rank=rank, - tp_size=config.ep_size, - moe_ep_size=config.ep_size, - world_size=config.ep_size, - ) - communication = create_comm_object(config.comm_type, mapping, config) +def _exercise_mnnvl_checkpoint_graph_replay( + config: CommTestConfig, rank: int, communication +) -> bool: worker_inputs = _prepare_worker_inputs(rank, config) _, local_slot_start, local_slot_end = _compute_ep_partition( config.num_experts, config.ep_size, rank @@ -2661,8 +2656,270 @@ def any_allocation_is_mapped() -> bool: torch.cuda.synchronize() torch.testing.assert_close(graph_output, expected) - if hasattr(communication, "destroy"): - communication.destroy() + return True + + +def _worker_mnnvl_checkpoint_graph_replay(config: CommTestConfig) -> bool: + """Exercise stable-VA checkpoint cycles through a captured MoE communication graph.""" + rank = tllm.mpi_rank() + torch.cuda.set_device(rank) + mapping = Mapping( + rank=rank, + tp_size=config.ep_size, + moe_ep_size=config.ep_size, + world_size=config.ep_size, + ) + communication = create_comm_object(config.comm_type, mapping, config) + try: + return _exercise_mnnvl_checkpoint_graph_replay(config, rank, communication) + finally: + if hasattr(communication, "destroy"): + communication.destroy() + + +def _mnnvl_workspace_is_mapped(communication) -> bool: + if isinstance(communication, NVLinkOneSided): + return communication.mnnvl_mem.mapped + workspaces = (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) + assert any(workspace is not None for workspace in workspaces) + return all(workspace is None or workspace.mapped for workspace in workspaces) + + +def _worker_mnnvl_engine_checkpoint_coordination(config: CommTestConfig) -> bool: + """Join the production engine coordinator, listener, and native resource hooks.""" + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor + from tensorrt_llm.executor.base_worker import BaseWorker + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + rank = tllm.mpi_rank() + world_size = config.ep_size + torch.cuda.set_device(rank) + assert MPI.Query_thread() >= MPI.THREAD_MULTIPLE + + mapping = Mapping( + rank=rank, + tp_size=world_size, + moe_ep_size=world_size, + world_size=world_size, + ) + communication = create_comm_object(config.comm_type, mapping, config) + sleep_wakeup_comm = MPI.COMM_WORLD.Dup() + control_comm = MPI.COMM_WORLD.Dup() + + class _Model: + def modules(self): + return [SimpleNamespace(comm=communication)] + + executor = object.__new__(PyExecutor) + executor._sleep_wakeup_comm = sleep_wakeup_comm + executor._sleep_wakeup_lock = threading.Lock() + executor._sleep_wakeup_listener_thread = None + executor.control_request_barrier = threading.Event() + executor.control_action_done = threading.Event() + executor._active_control_id = None + executor.device_id = rank + executor.dist = SimpleNamespace(rank=rank, world_size=world_size) + executor.model_engine = SimpleNamespace(model=_Model()) + executor.draft_model_engine = None + + class _ControlRequestQueue: + def enqueue_control_request(self, *, drain, control_id): + assert drain + control_comm.bcast(control_id, root=0) + executor._active_control_id = control_id + executor.control_action_done.clear() + executor.control_request_barrier.set() + + executor.executor_request_queue = _ControlRequestQueue() + retained_objects = getattr(sys, "_trtllm_mnnvl_coordination_test_objects", None) + if retained_objects is None: + retained_objects = [] + setattr(sys, "_trtllm_mnnvl_coordination_test_objects", retained_objects) + retained_objects.append((communication, sleep_wakeup_comm, control_comm, executor)) + worker = object.__new__(BaseWorker) + worker.doing_shutdown = True + worker._backend = "pytorch" + worker.rank = rank + worker.engine = executor + worker.llm_args = SimpleNamespace( + parallel_config=SimpleNamespace(world_size=world_size), + sleep_config=object(), + ) + + listener = None + listener_shutdown = False + try: + initial_error = None + try: + assert executor._has_mnnvl_checkpoint_resources([ExecutorMemoryType.MODEL_ENGINE_MAIN]) + assert _mnnvl_workspace_is_mapped(communication) + except Exception as error: + initial_error = f"rank {rank}: {error}" + initial_errors = MPI.COMM_WORLD.allgather(initial_error) + assert not any(initial_errors), initial_errors + + if rank != 0: + listener = threading.Thread( + target=executor._sleep_wakeup_listener_loop, + name=f"mnnvl-checkpoint-listener-{rank}", + daemon=True, + ) + executor._sleep_wakeup_listener_thread = listener + listener.start() + + MPI.COMM_WORLD.Barrier() + expected_mapped_states = (False, True) + if rank == 0: + for action, expected_mapped in zip(("sleep", "wakeup"), expected_mapped_states): + action_error = None + try: + worker._multi_rank_sleep_wakeup(action, [ExecutorMemoryType.MODEL_ENGINE_MAIN]) + except Exception as error: + action_error = error + control_comm.bcast(action_error is None, root=0) + if action_error is not None: + raise action_error + + local_state_error = None + try: + assert _mnnvl_workspace_is_mapped(communication) is expected_mapped + except Exception as error: + local_state_error = f"rank {rank}: {error}" + state_errors = control_comm.allgather(local_state_error) + assert not any(state_errors), state_errors + executor._shutdown_sleep_wakeup_listeners() + listener_shutdown = True + else: + # This is the only emulated production piece: the regular request + # broadcaster delivers the control sentinel to the executor loop. + # The real listener and its dedicated communicator remain under test. + for expected_mapped in expected_mapped_states: + control_id = control_comm.bcast(None, root=0) + executor._active_control_id = control_id + executor.control_action_done.clear() + executor.control_request_barrier.set() + action_completed = executor.control_action_done.wait(timeout=60.0) + executor.control_action_done.clear() + executor._active_control_id = None + action_succeeded = control_comm.bcast(None, root=0) + if not action_succeeded: + break + + local_state_error = None + try: + assert action_completed, "listener did not release the control request" + assert _mnnvl_workspace_is_mapped(communication) is expected_mapped + except Exception as error: + local_state_error = f"rank {rank}: {error}" + state_errors = control_comm.allgather(local_state_error) + assert not any(state_errors), state_errors + assert listener is not None + listener.join(timeout=60.0) + assert not listener.is_alive() + return True + finally: + if rank == 0 and not listener_shutdown: + executor._shutdown_sleep_wakeup_listeners() + if listener is not None and listener.is_alive(): + listener.join(timeout=1.0) + # The process-global retention list keeps the duplicated communicators, + # executor shell, and native resource alive until MPI_Finalize. Freeing + # or destroying any of them while a failed peer listener still owns a + # request can hide the original failure with a hang/crash. + + +class _FailureInjectionMnnvlMemory(MnnvlMemory): + pass + + +def _worker_mnnvl_checkpoint_failure_injection(_unused=None) -> bool: + """Prove local CUDA failures do not strand peers in checkpoint collectives.""" + from types import SimpleNamespace + from unittest.mock import Mock, patch + + rank = tllm.mpi_rank() + comm = MPI.COMM_WORLD.Split(0, rank) + comm_size = comm.Get_size() + obj = _FailureInjectionMnnvlMemory.__new__(_FailureInjectionMnnvlMemory) + obj.ptr = 1032 + obj.mapping = SimpleNamespace(rank=rank) + + def install_record(state, handles): + record = mnnvl._MnnvlAllocationRecord( + comm=comm, + comm_size=comm_size, + comm_rank=rank, + comm_membership=tuple(range(comm_size)), + aligned_size=64, + mem_handles=handles, + start_address=1000, + rank_stride=256, + address_offset=32, + state=state, + ) + _FailureInjectionMnnvlMemory.allocated_map = {obj.ptr: record} + return record + + prepare_record = install_record( + mnnvl._MnnvlAllocationState.MAPPED, + list(range(1, comm_size + 1)), + ) + unmap = Mock() + if rank == 0: + unmap.side_effect = [ + None, + RuntimeError("injected partial unmap failure"), + *([None] * (comm_size - 2)), + ] + with ( + patch.object(mnnvl.torch.cuda, "synchronize"), + patch.object(mnnvl.cuda, "cuMemUnmap", unmap), + patch.object(mnnvl.cuda, "cuMemRelease", return_value=None), + patch.object(mnnvl, "_check_cu_result", side_effect=lambda result: result), + ): + prepare_error = None + try: + obj.checkpoint_prepare() + except RuntimeError as error: + prepare_error = str(error) + + prepare_errors = comm.allgather(prepare_error) + assert "injected partial unmap failure" in prepare_errors[0] + assert all(error is None for error in prepare_errors[1:]) + assert prepare_record.state is ( + mnnvl._MnnvlAllocationState.BROKEN if rank == 0 else mnnvl._MnnvlAllocationState.UNMAPPED + ) + + restore_record = install_record( + mnnvl._MnnvlAllocationState.UNMAPPED, + [None] * comm_size, + ) + create_and_map = Mock( + side_effect=(RuntimeError("injected remap failure") if rank == 0 else None), + return_value=list(range(11, 11 + comm_size)), + ) + with ( + patch.object(mnnvl.torch.cuda, "synchronize"), + patch.object(mnnvl.cuda, "cuMemUnmap", return_value=None), + patch.object(mnnvl.cuda, "cuMemRelease", return_value=None), + patch.object(mnnvl, "_check_cu_result", side_effect=lambda result: result), + patch.object( + _FailureInjectionMnnvlMemory, + "_create_and_map_handles", + create_and_map, + ), + ): + restore_error = None + try: + obj.checkpoint_restore(comm) + except RuntimeError as error: + restore_error = str(error) + + restore_errors = comm.allgather(restore_error) + assert all("injected remap failure" in error for error in restore_errors) + assert restore_record.state is mnnvl._MnnvlAllocationState.BROKEN + _FailureInjectionMnnvlMemory.allocated_map = {} + del obj.ptr return True @@ -2760,6 +3017,48 @@ def test_mnnvl_checkpoint_preserves_moe_graph_addresses( ) assert all(results) + @pytest.mark.threadleak(enabled=False) + @pytest.mark.parametrize("mpi_pool_executor", [4], indirect=True) + @pytest.mark.parametrize( + "comm_type", + [COMM_NVLINK_ONE_SIDED, COMM_NVLINK_TWO_SIDED], + ) + def test_mnnvl_engine_checkpoint_coordination( + self, + mpi_pool_executor, + comm_type: str, + ) -> None: + """Exercise BaseWorker and PyExecutor around native MNNVL checkpoint hooks.""" + config = CommTestConfig( + comm_type=comm_type, + ep_size=mpi_pool_executor.num_workers, + num_experts=32, + top_k=2, + hidden_size=128, + all_num_tokens=[8] * mpi_pool_executor.num_workers, + ) + skip_reason = _get_skip_reason(config) + if skip_reason: + pytest.skip(skip_reason) + results = mpi_pool_executor.map( + _worker_mnnvl_engine_checkpoint_coordination, + [config] * config.ep_size, + ) + assert all(results) + + @pytest.mark.threadleak(enabled=False) + @pytest.mark.parametrize("mpi_pool_executor", [4], indirect=True) + def test_mnnvl_checkpoint_failure_is_collective_and_bounded( + self, + mpi_pool_executor, + ) -> None: + """A single-rank unmap or remap failure must return on every rank.""" + results = mpi_pool_executor.map( + _worker_mnnvl_checkpoint_failure_injection, + [None] * mpi_pool_executor.num_workers, + ) + assert all(results) + @pytest.mark.threadleak(enabled=False) @pytest.mark.parametrize( "mpi_pool_executor,local_num_tokens,top_k", diff --git a/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py b/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py index 801769c66c48..0c3edfadee6f 100644 --- a/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py +++ b/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py @@ -168,6 +168,20 @@ def test_checkpoint_prepare_rejects_communicator_size_mismatch() -> None: memory.checkpoint_prepare.assert_not_called() +def test_checkpoint_prepare_timeout_fails_closed(monkeypatch) -> None: + lifecycle, memory, _ = _make_lifecycle() + request = SimpleNamespace(test=lambda: (False, None)) + memory.comm = SimpleNamespace(iallgather=lambda value: request) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S", 0.0) + + with pytest.raises(TimeoutError, match="workspace idle readiness"): + lifecycle.checkpoint_prepare() + + memory.checkpoint_fail_closed.assert_called_once_with() + memory.checkpoint_prepare.assert_not_called() + + def test_checkpoint_prepare_rejects_remote_active_client_before_watchdog_stop() -> None: lifecycle, memory, _ = _make_lifecycle() memory.comm = _FakeComm(clients_idle_by_rank=[True, False]) @@ -536,6 +550,35 @@ def test_two_sided_checkpoint_prepare_rejects_uninitialized_communicator( checkpoint_prepare.assert_not_called() +def test_two_sided_checkpoint_prepare_timeout_fails_closed( + monkeypatch: pytest.MonkeyPatch, +) -> None: + instances = WeakSet() + monkeypatch.setattr(NVLinkTwoSided, "_INSTANCES", instances) + request = SimpleNamespace(test=lambda: (False, None)) + comm = SimpleNamespace(iallgather=lambda value: request) + main_workspace = Mock(mapped=True, comm=comm) + prepare_workspace = Mock(mapped=True) + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace", main_workspace) + monkeypatch.setattr( + mnnvl.MnnvlMoe, + "moe_prepare_workspace", + prepare_workspace, + ) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S", 0.0) + owner = NVLinkTwoSided.__new__(NVLinkTwoSided) + owner.ep_size = 2 + owner._dispatch_state = {} + instances.add(owner) + + with pytest.raises(TimeoutError, match="workspace idle readiness"): + owner.checkpoint_prepare() + + main_workspace.checkpoint_fail_closed.assert_called_once_with() + prepare_workspace.checkpoint_fail_closed.assert_called_once_with() + + def test_two_sided_checkpoint_restore_resets_all_shared_owners( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -620,6 +663,7 @@ def test_frontend_destroy_unregisters_from_shared_lifecycle( ) -> None: wrapper = wrapper_type.__new__(wrapper_type) wrapper._destroyed = False + wrapper._workspace_registered = True lifecycle = Mock() wrapper._workspace_lifecycle = lifecycle if wrapper_type is NVLinkOneSided: @@ -631,6 +675,53 @@ def test_frontend_destroy_unregisters_from_shared_lifecycle( lifecycle.unregister.assert_called_once_with(wrapper) +def test_moe_alltoall_aborted_registration_does_not_unregister( + monkeypatch: pytest.MonkeyPatch, +) -> None: + lifecycle = Mock() + lifecycle.register.side_effect = RuntimeError("registration failed") + monkeypatch.setattr(MoeAlltoAll, "_WORKSPACE", None) + monkeypatch.setattr(MoeAlltoAll, "_init_constants", Mock()) + monkeypatch.setattr( + MoeAlltoAll, + "_METAINFO_INDEX", + { + "FLAG_VAL_OFFSET_INDEX": 0, + "DISPATCH_COMPLETION_FLAGS_OFFSET_INDEX": 0, + "COMBINE_COMPLETION_FLAGS_OFFSET_INDEX": 0, + }, + ) + monkeypatch.setattr(mnnvl.MnnvlMemory, "initialize", Mock()) + memory = Mock() + memory.as_torch_strided_tensor.return_value = torch.zeros(1, dtype=torch.uint8) + monkeypatch.setattr( + "tensorrt_llm._torch.distributed.moe_alltoall.MnnvlMemory", + Mock(return_value=memory), + ) + monkeypatch.setattr( + _MnnvlAlltoAllWorkspaceLifecycle, + "get_or_create", + Mock(return_value=lifecycle), + ) + monkeypatch.setattr( + torch.ops.trtllm, + "moe_a2a_initialize", + Mock(return_value=torch.tensor([1])), + ) + mapping = SimpleNamespace(moe_ep_size=2, moe_ep_rank=0) + + with pytest.raises(RuntimeError, match="registration failed"): + MoeAlltoAll( + mapping=mapping, + max_num_tokens=1, + top_k=1, + num_slots=2, + workspace_size_per_rank=1, + ) + + lifecycle.unregister.assert_not_called() + + def test_one_sided_checkpoint_rejects_destroyed_workspace( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -661,6 +752,7 @@ def unregister(self, client: object) -> None: wrapper._destroyed = False wrapper._workspace_lifecycle = lifecycle wrapper._workspace_key = None + wrapper._workspace_registered = True del wrapper gc.collect() @@ -748,7 +840,12 @@ def as_torch_strided_tensor(self, dtype: torch.dtype) -> torch.Tensor: "moe_a2a_initialize", Mock(return_value=torch.tensor([1])), ) - mapping = SimpleNamespace(world_size=2, moe_ep_size=2, moe_ep_rank=0) + mapping = SimpleNamespace( + world_size=2, + moe_ep_size=2, + moe_ep_rank=0, + has_cp_helix=Mock(return_value=False), + ) with pytest.raises(RuntimeError, match="registration failed"): NVLinkOneSided( @@ -761,3 +858,4 @@ def as_torch_strided_tensor(self, dtype: torch.dtype) -> torch.Tensor: assert NVLinkOneSided._WORKSPACES == {} assert NVLinkOneSided._WORKSPACE_REFCOUNTS == {} assert NVLinkOneSided._WORKSPACE is None + lifecycle.unregister.assert_not_called() diff --git a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py index d7ebc7a9a82a..6e9fcb1f3a62 100644 --- a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py +++ b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py @@ -41,13 +41,29 @@ def barrier(self): self.barrier_count += 1 def allgather(self, value): - return list(self.membership) + if isinstance(value, int) and not isinstance(value, bool): + return list(self.membership) + return [value] * self.size class _TestMnnvlMemory(mnnvl.MnnvlMemory): pass +def test_checkpoint_allgather_timeout_is_bounded(monkeypatch): + request = SimpleNamespace(test=lambda: (False, None)) + comm = SimpleNamespace(iallgather=lambda value: request) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S", 0.0) + + with pytest.raises(TimeoutError, match="mapping readiness"): + mnnvl._checkpoint_allgather( + comm, + None, + operation="mapping readiness", + ) + + @pytest.fixture def memory(monkeypatch): comm = _FakeComm() @@ -92,11 +108,11 @@ def test_checkpoint_prepare_preserves_va_and_is_idempotent(memory): assert record.rank_stride == 256 assert record.address_offset == 32 assert record.mem_handles == [None, None] - assert record.comm.barrier_count == 2 + assert record.comm.barrier_count == 0 assert [call.args[0] for call in mnnvl.cuda.cuMemUnmap.call_args_list] == [1032, 1288] obj.checkpoint_prepare() - assert record.comm.barrier_count == 2 + assert record.comm.barrier_count == 0 def test_checkpoint_restore_reuses_layout_with_fresh_handles(memory, monkeypatch): @@ -163,6 +179,87 @@ def test_checkpoint_restore_failure_is_terminal_and_fails_closed(memory, monkeyp obj.checkpoint_restore(_FakeComm()) +def test_checkpoint_restore_membership_timeout_is_terminal(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S", 0.0) + + request = SimpleNamespace(test=lambda: (False, None)) + comm = SimpleNamespace( + Get_rank=lambda: 0, + Get_size=lambda: 2, + iallgather=lambda value: request, + ) + with pytest.raises(TimeoutError, match="communicator membership"): + obj.checkpoint_restore(comm) + + assert record.state is mnnvl._MnnvlAllocationState.BROKEN + assert record.pending_comm is None + with pytest.raises(RuntimeError, match="broken state"): + obj.checkpoint_restore(comm) + + +def test_checkpoint_restore_timeout_cleans_locally_mapped_handles(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + monkeypatch.setattr( + _TestMnnvlMemory, + "_create_and_map_handles", + Mock(return_value=[33, 44]), + ) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S", 0.0) + + class Request: + def __init__(self, result=None): + self.result = result + + def test(self): + return (self.result is not None, self.result) + + class TimeoutComm(_FakeComm): + def __init__(self): + super().__init__() + self.requests = [Request([0, 1]), Request()] + + def iallgather(self, value): + return self.requests.pop(0) + + with pytest.raises(TimeoutError, match="mapping readiness"): + obj.checkpoint_restore(TimeoutComm()) + + assert record.state is mnnvl._MnnvlAllocationState.BROKEN + assert record.mem_handles == [None, None] + assert [call.args[0] for call in mnnvl.cuda.cuMemUnmap.call_args_list[-2:]] == [1032, 1288] + + +def test_handle_exchange_timeout_does_not_start_second_collective(memory, monkeypatch): + obj, record = memory + obj.checkpoint_prepare() + monkeypatch.setattr( + _TestMnnvlMemory, + "_create_and_map_handles", + Mock(side_effect=TimeoutError("handle exchange timed out")), + ) + + class CountingComm(_FakeComm): + def __init__(self): + super().__init__() + self.allgather_count = 0 + + def allgather(self, value): + self.allgather_count += 1 + return super().allgather(value) + + comm = CountingComm() + with pytest.raises(TimeoutError, match="handle exchange timed out"): + obj.checkpoint_restore(comm) + + assert comm.allgather_count == 1 + assert record.state is mnnvl._MnnvlAllocationState.BROKEN + + def test_failed_frontend_restore_releases_unpublished_handles(memory, monkeypatch): obj, record = memory obj.checkpoint_prepare() @@ -243,11 +340,39 @@ def test_mnnvl_moe_restore_publishes_all_workspaces_after_frontend_ready(monkeyp second.checkpoint_restore.assert_called_once_with(comm) initialize_workspace.assert_called_once_with(workspace_tensor, 0, 2) synchronize.assert_called_once_with() - assert comm.barrier_count == 1 + assert comm.barrier_count == 0 first._checkpoint_restore_complete.assert_called_once_with() second._checkpoint_restore_complete.assert_called_once_with() +def test_mnnvl_moe_restore_prepare_only_skips_main_workspace_initialization(monkeypatch): + main_workspace = Mock() + main_workspace.checkpoint_restore.return_value = False + prepare_workspace = Mock() + prepare_workspace.checkpoint_restore.return_value = True + workspace_tensor = Mock() + initialize_workspace = Mock() + synchronize = Mock() + comm = _FakeComm() + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace", main_workspace) + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_prepare_workspace", prepare_workspace) + monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace_tensor", workspace_tensor) + monkeypatch.setattr( + mnnvl.torch.ops.trtllm, + "moe_initialize_workspace", + initialize_workspace, + ) + monkeypatch.setattr(mnnvl.torch.cuda, "synchronize", synchronize) + + mnnvl.MnnvlMoe.checkpoint_restore(comm) + + initialize_workspace.assert_not_called() + synchronize.assert_called_once_with() + assert comm.barrier_count == 0 + main_workspace._checkpoint_restore_complete.assert_not_called() + prepare_workspace._checkpoint_restore_complete.assert_called_once_with() + + def test_mnnvl_moe_restore_failure_marks_earlier_workspace_broken(monkeypatch): first = Mock() first.checkpoint_restore.return_value = True @@ -283,9 +408,18 @@ def test_close_detached_memory_only_frees_va(memory, monkeypatch): def test_create_and_map_handles_cleans_partial_allocation(monkeypatch): comm = Mock() + comm.iallgather = None comm.Get_rank.return_value = 0 comm.Get_size.return_value = 2 - comm.allgather.return_value = [b"local", b"remote"] + comm.allgather.side_effect = lambda value: [ + value, + { + "error": None, + "handle": b"remote", + "is_fabric": True, + "pid": 1001, + }, + ] allocation_prop = SimpleNamespace( requestedHandleTypes=mnnvl.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC, location=object(), @@ -325,8 +459,18 @@ def test_create_and_map_handles_cleans_partial_allocation(monkeypatch): def test_create_and_map_handles_releases_local_handle_on_export_failure(monkeypatch): comm = Mock() + comm.iallgather = None comm.Get_rank.return_value = 0 comm.Get_size.return_value = 2 + comm.allgather.side_effect = lambda value: [ + value, + { + "error": None, + "handle": b"remote", + "is_fabric": True, + "pid": 1001, + }, + ] allocation_prop = SimpleNamespace( requestedHandleTypes=mnnvl.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC, location=object(), @@ -361,9 +505,22 @@ def test_create_and_map_handles_releases_local_handle_on_export_failure(monkeypa ) def test_create_and_map_handles_preserves_pidfd_error_hint(monkeypatch, errno_value, expected_hint): comm = Mock() + comm.iallgather = None comm.Get_rank.return_value = 0 comm.Get_size.return_value = 2 - comm.allgather.side_effect = [[100, 101], [1000, 1001]] + comm.allgather.side_effect = lambda value: ( + [ + value, + { + "error": None, + "handle": 101, + "is_fabric": False, + "pid": 1001, + }, + ] + if isinstance(value, dict) + else [value, value] + ) allocation_prop = SimpleNamespace( requestedHandleTypes=( mnnvl.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR @@ -391,14 +548,28 @@ def test_create_and_map_handles_preserves_pidfd_error_hint(monkeypatch, errno_va with pytest.raises(RuntimeError, match=expected_hint): _TestMnnvlMemory._create_and_map_handles(comm, 64, 1000, 256, 32) + assert comm.allgather.call_count == 2 assert [call.args[0] for call in close_fd.call_args_list] == [10, 11, 100] def test_create_and_map_handles_close_failure_does_not_mask_original_error(monkeypatch): comm = Mock() + comm.iallgather = None comm.Get_rank.return_value = 0 comm.Get_size.return_value = 2 - comm.allgather.side_effect = [[100, 101], [1000, 1001]] + comm.allgather.side_effect = lambda value: ( + [ + value, + { + "error": None, + "handle": 101, + "is_fabric": False, + "pid": 1001, + }, + ] + if isinstance(value, dict) + else [value, value] + ) allocation_prop = SimpleNamespace( requestedHandleTypes=( mnnvl.cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR @@ -428,6 +599,7 @@ def test_create_and_map_handles_close_failure_does_not_mask_original_error(monke with pytest.raises(RuntimeError, match="map failed"): _TestMnnvlMemory._create_and_map_handles(comm, 64, 1000, 256, 32) + assert comm.allgather.call_count == 2 assert [call.args[0] for call in close_fd.call_args_list] == [10, 11, 20, 21, 100] diff --git a/tests/unittest/executor/test_sleep_collective_rpc_guards.py b/tests/unittest/executor/test_sleep_collective_rpc_guards.py index 4576f88ede6c..90add390febd 100644 --- a/tests/unittest/executor/test_sleep_collective_rpc_guards.py +++ b/tests/unittest/executor/test_sleep_collective_rpc_guards.py @@ -426,6 +426,109 @@ def test_multiple_peer_errors_aggregated(self): assert "rank 2 OOM" in msg +class TestMnnvlSleepWakeupCoordination: + """Native MNNVL hooks enter COMMIT collectively and remain tag-scoped.""" + + def test_pyexecutor_discovers_shared_native_resource_once(self): + from types import SimpleNamespace + from unittest.mock import Mock + + from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_one_sided import ( + NVLinkOneSided, + ) + from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor, _SleepWakeupAction + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + lifecycle = object() + resources = [] + modules = [] + for _ in range(2): + resource = NVLinkOneSided.__new__(NVLinkOneSided) + resource._workspace_lifecycle = lifecycle + resource._workspace_registered = False + resource.checkpoint_prepare = Mock() + resource.checkpoint_restore = Mock() + resources.append(resource) + modules.append(SimpleNamespace(comm=resource)) + + executor = object.__new__(PyExecutor) + executor.model_engine = SimpleNamespace(model=SimpleNamespace(modules=lambda: modules)) + executor.draft_model_engine = None + + tags = [ExecutorMemoryType.MODEL_ENGINE_MAIN] + assert executor._mnnvl_checkpoint_resources(tags) == [resources[0]] + executor._run_mnnvl_checkpoint_resources(_SleepWakeupAction.SLEEP, tags) + resources[0].checkpoint_prepare.assert_called_once_with() + resources[1].checkpoint_prepare.assert_not_called() + assert not executor._has_mnnvl_checkpoint_resources([ExecutorMemoryType.KV_CACHE]) + + def test_mnnvl_commit_reaches_peer_before_rank_zero_hook(self): + from unittest.mock import patch + + from tensorrt_llm._torch.pyexecutor.py_executor import _SleepWakeupAction + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + worker, _ = _make_proto_worker( + [ + {"status": "ok", "has_mnnvl_resources": True}, + {"status": "ok"}, + ], + world_size=2, + ) + events = [] + original_send = worker.engine._sleep_wakeup_comm.send + + def record_send(payload, *args, **kwargs): + events.append(("send", payload["action"])) + return original_send(payload, *args, **kwargs) + + worker.engine._sleep_wakeup_comm.send = record_send + # Model partitions can differ: rank 0 may own no MoE layer while a + # peer does. The PREPARE ACK must still select peer-first COMMIT. + worker.engine._has_mnnvl_checkpoint_resources = lambda tags: False + worker.engine._run_mnnvl_checkpoint_resources = lambda action, tags: events.append( + ("local_mnnvl", action) + ) + + with ( + patch( + "tensorrt_llm._torch.virtual_memory.release_with_tag", + side_effect=lambda *tags: events.append(("local_vmm", "sleep")), + ), + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize"), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.MODEL_ENGINE_MAIN]) + + commit_index = events.index(("send", _SleepWakeupAction.COMMIT)) + hook_index = events.index(("local_mnnvl", _SleepWakeupAction.SLEEP)) + vmm_index = events.index(("local_vmm", "sleep")) + assert commit_index < hook_index < vmm_index + + def test_kv_only_sleep_does_not_run_mnnvl_hook(self): + from unittest.mock import Mock, patch + + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + worker, _ = _make_proto_worker([{"status": "ok"}, {"status": "ok"}], world_size=2) + run_mnnvl = Mock() + worker.engine._has_mnnvl_checkpoint_resources = lambda tags: False + worker.engine._run_mnnvl_checkpoint_resources = run_mnnvl + + with ( + patch("tensorrt_llm._torch.virtual_memory.release_with_tag"), + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize"), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.KV_CACHE]) + + run_mnnvl.assert_not_called() + + class TestMultiRankSendFailureRecovery: """Partial rank-0 broadcast failures must not leave peer ACKs undrained.""" @@ -667,6 +770,75 @@ def _noop_control_action(**kwargs): (1, _SleepWakeupTag.ACK), ] + def test_mnnvl_partial_commit_still_runs_rank_zero_bounded_phase(self): + """Once a peer sees MNNVL COMMIT, rank 0 must enter the local phase.""" + import threading + from contextlib import contextmanager + from unittest.mock import Mock, patch + + from tensorrt_llm._torch.pyexecutor.py_executor import _SleepWakeupAction + from tensorrt_llm.executor.base_worker import BaseWorker + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + class FakeComm: + def __init__(self): + self.op_id = None + self.ack_phases = {1: [], 2: []} + + def send(self, payload, dest, tag): + self.op_id = payload.get("op_id", self.op_id) + if payload["action"] == _SleepWakeupAction.COMMIT and dest == 2: + raise RuntimeError("injected MNNVL commit send failure") + self.ack_phases[dest].append(payload["action"]) + + def iprobe(self, source, tag): + return True + + def recv(self, source, tag): + return { + "status": "ok", + "op_id": self.op_id, + "phase": self.ack_phases[source].pop(0), + "has_mnnvl_resources": True, + } + + @contextmanager + def control_action(**kwargs): + yield None + + worker = object.__new__(BaseWorker) + worker._backend = "pytorch" + worker.rank = 0 + worker.llm_args = SimpleNamespace( + backend="pytorch", + parallel_config=SimpleNamespace(world_size=3), + sleep_config=object(), + ) + run_mnnvl = Mock() + worker.engine = SimpleNamespace( + _sleep_wakeup_lock=threading.Lock(), + _sleep_wakeup_comm=FakeComm(), + control_action=control_action, + _has_mnnvl_checkpoint_resources=lambda tags: True, + _run_mnnvl_checkpoint_resources=run_mnnvl, + ) + + with ( + patch("tensorrt_llm._torch.virtual_memory.release_with_tag") as release, + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize"), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + with pytest.raises(RuntimeError, match="commit send failure"): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.MODEL_ENGINE_MAIN]) + + run_mnnvl.assert_called_once_with( + _SleepWakeupAction.SLEEP, + [ExecutorMemoryType.MODEL_ENGINE_MAIN], + ) + release.assert_called_once_with(ExecutorMemoryType.MODEL_ENGINE_MAIN) + class TestListenerUncaughtExceptionSendsErrorAck: """An exception that bypasses the narrow except clause must still produce an error ACK. @@ -803,6 +975,7 @@ def send(self, payload, dest, tag): "error": None, "op_id": op_id, "phase": _SleepWakeupAction.PREPARE, + "has_mnnvl_resources": False, } ] From ef33cf087dc711bb02d21bf65087d2b492962cf5 Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Tue, 18 Aug 2026 08:43:52 -0400 Subject: [PATCH 09/14] test(moe): fix MNNVL lifecycle CI coverage Signed-off-by: Hannah Zhang --- tests/integration/test_lists/test-db/l0_a10.yml | 1 - tests/integration/test_lists/test-db/l0_cpu.yml | 1 + tests/unittest/_torch/test_mnnvl_memory_lifecycle.py | 6 +++++- 3 files changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 41db4b16b4fb..c2ed366ccf46 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -20,7 +20,6 @@ l0_a10: - unittest/_torch/test_torch_multi_arange.py - unittest/_torch/test_mnnvl_alltoall_workspace.py - unittest/_torch/test_mnnvl_memory_lifecycle.py - - unittest/executor/test_sleep_collective_rpc_guards.py - unittest/utils/test_util.py - unittest/_torch/modeling/test_modeling_mistral.py - unittest/_torch/modeling/test_modeling_pixtral.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index c605b8e36d6e..effeab272165 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -70,6 +70,7 @@ l0_cpu: - unittest/executor/test_multi_frontend_routing.py - unittest/executor/test_event_loop_error_broadcast.py - unittest/executor/test_rpc_worker_mixin.py + - unittest/executor/test_sleep_collective_rpc_guards.py - unittest/executor/test_stats_serializer.py - unittest/executor/test_spec_dec_perf_metrics.py - unittest/executor/test_ray_stub.py diff --git a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py index 6e9fcb1f3a62..be1df3fc19ec 100644 --- a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py +++ b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py @@ -594,7 +594,11 @@ def test_create_and_map_handles_close_failure_does_not_mask_original_error(monke monkeypatch.setattr(mnnvl.cuda, "cuMemRelease", Mock(return_value=None)) monkeypatch.setattr(mnnvl.ctypes, "CDLL", Mock(return_value=SimpleNamespace(syscall=syscall))) close_fd = Mock(side_effect=[OSError("close failed"), None, None, None, None]) - monkeypatch.setattr(mnnvl.os, "close", close_fd) + monkeypatch.setattr( + mnnvl, + "os", + SimpleNamespace(close=close_fd, getpid=mnnvl.os.getpid, strerror=mnnvl.os.strerror), + ) with pytest.raises(RuntimeError, match="map failed"): _TestMnnvlMemory._create_and_map_handles(comm, 64, 1000, 256, 32) From 8bd72d75c3ec7f1b703df7e293f245bff0e1ef0e Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Tue, 18 Aug 2026 10:05:53 -0400 Subject: [PATCH 10/14] fix(moe): harden MNNVL checkpoint coordination Signed-off-by: Hannah Zhang --- tensorrt_llm/_mnnvl_utils.py | 116 +++++++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 30 ++- tensorrt_llm/executor/base_worker.py | 48 +++- .../_torch/test_mnnvl_alltoall_workspace.py | 26 +- .../_torch/test_mnnvl_memory_lifecycle.py | 138 +++++++++- .../test_sleep_collective_rpc_guards.py | 243 +++++++++++++++++- 6 files changed, 539 insertions(+), 62 deletions(-) diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index ffa354770d66..f1d7a8479bdf 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/tensorrt_llm/_mnnvl_utils.py @@ -38,27 +38,104 @@ _MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S = 20.0 _MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S = 0.01 +_MNNVL_CHECKPOINT_REQUEST_CLEANUP_TIMEOUT_S = 0.1 +_MNNVL_CHECKPOINT_ALLGATHER_TAG = 31415 +_MNNVL_CHECKPOINT_ORPHANED_REQUESTS: list[Any] = [] -def _checkpoint_allgather(comm, value, *, operation: str): - """Run a bounded object allgather when the communicator supports it.""" - iallgather = getattr(comm, "iallgather", None) - if iallgather is None: +def _cancel_checkpoint_requests(requests: list[Any]) -> None: + """Bound cancellation cleanup so timeout handling cannot hang in MPI.""" + for request in requests: + try: + request.Cancel() + except Exception as error: + logger.warning(f"Failed to cancel MNNVL checkpoint request: {error}") + + pending_requests = list(requests) + cleanup_timeout_s = min( + _MNNVL_CHECKPOINT_REQUEST_CLEANUP_TIMEOUT_S, + _MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S, + ) + deadline = time.monotonic() + cleanup_timeout_s + while pending_requests: + incomplete_requests = [] + for request in pending_requests: + try: + ready, _ = request.test() + except Exception as error: + logger.warning(f"Failed to poll MNNVL checkpoint request cleanup: {error}") + incomplete_requests.append(request) + continue + if not ready: + incomplete_requests.append(request) + pending_requests = incomplete_requests + if not pending_requests or time.monotonic() >= deadline: + break + time.sleep(_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S) + + if pending_requests: + # An active receive must not be freed: mpi4py owns its receive buffer, + # and MPI may still write into it. Keep the wrappers alive until the + # enclosing fail-closed path terminates the worker. + _MNNVL_CHECKPOINT_ORPHANED_REQUESTS.extend(pending_requests) + logger.error( + "Retaining %d active MNNVL checkpoint requests until worker termination", + len(pending_requests), + ) + + +def _checkpoint_allgather(comm: Any, value: Any, *, operation: str) -> list[Any]: + """Run a bounded object allgather over nonblocking point-to-point requests. + + Production mpi4py communicators provide ``isend`` and ``irecv`` but no + nonblocking object allgather. The blocking fallback is retained only for + lightweight test communicators that do not implement point-to-point APIs. + """ + isend = getattr(comm, "isend", None) + irecv = getattr(comm, "irecv", None) + if isend is None or irecv is None: return comm.allgather(value) - request = iallgather(value) - deadline = time.monotonic() + _MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S - while True: - test_result = request.test() - if isinstance(test_result, tuple): - ready, result = test_result - if ready: - return result - elif test_result: - return request.wait() - if time.monotonic() >= deadline: - raise TimeoutError(f"Timed out waiting for MNNVL checkpoint {operation} allgather") - time.sleep(_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S) + rank = comm.Get_rank() + size = comm.Get_size() + results: list[Any] = [None] * size + results[rank] = value + receive_requests: dict[int, Any] = {} + send_requests: list[Any] = [] + try: + for peer in range(size): + if peer != rank: + receive_requests[peer] = irecv( + source=peer, + tag=_MNNVL_CHECKPOINT_ALLGATHER_TAG, + ) + for peer in range(size): + if peer != rank: + send_requests.append( + isend( + value, + dest=peer, + tag=_MNNVL_CHECKPOINT_ALLGATHER_TAG, + ) + ) + if not receive_requests: + return results + deadline = time.monotonic() + _MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S + while receive_requests or send_requests: + for peer, request in list(receive_requests.items()): + ready, result = request.test() + if ready: + results[peer] = result + del receive_requests[peer] + send_requests = [request for request in send_requests if not request.test()[0]] + if not receive_requests and not send_requests: + return results + if time.monotonic() >= deadline: + raise TimeoutError(f"Timed out waiting for MNNVL checkpoint {operation} allgather") + time.sleep(_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S) + except Exception: + _cancel_checkpoint_requests([*receive_requests.values(), *send_requests]) + raise def _check_cu_result(cu_func_ret): @@ -580,9 +657,8 @@ def checkpoint_prepare(self) -> None: def checkpoint_fail_closed(self) -> None: """Make a timed-out checkpoint allocation terminal. - A timed-out nonblocking MPI collective may still be outstanding. The - allocation must not be reused by a later checkpoint attempt because a - new collective could then be matched against the abandoned request. + A timeout means ranks may have reached different checkpoint phases. + The allocation must not be reused by a later checkpoint attempt. """ record = type(self).allocated_map[self.ptr] record.state = _MnnvlAllocationState.BROKEN diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 70a337335d17..edc0dd141773 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -3040,6 +3040,7 @@ def _sleep_wakeup_listener_loop(self) -> None: action = msg.get("action", "") op_id = msg.get("op_id") error_msg = None + abort_acknowledged = False release_control_request = True has_mnnvl_resources = False try: @@ -3069,10 +3070,13 @@ def _sleep_wakeup_listener_loop(self) -> None: None) if op_id is None: release_control_request = False - elif (not self.control_request_barrier.is_set() - or active_control_id != op_id): - self._record_sleep_wakeup_abort(op_id, error_msg) - release_control_request = False + else: + if (not self.control_request_barrier.is_set() + or active_control_id != op_id): + self._record_sleep_wakeup_abort( + op_id, error_msg) + release_control_request = False + abort_acknowledged = True logger.warning("Sleep/wakeup listener: %s", error_msg) elif action in (_SleepWakeupAction.PREPARE, _SleepWakeupAction.COMMIT, @@ -3157,14 +3161,18 @@ def _sleep_wakeup_listener_loop(self) -> None: # exiting control_action() and broadcasting new requests # before our executor loop has cleared its control barrier # and is ready to participate in the next collective. + ack_error = None if abort_acknowledged else error_msg + ack = { + "status": "ok" if ack_error is None else "error", + "error": ack_error, + "op_id": op_id, + "phase": action, + "has_mnnvl_resources": has_mnnvl_resources, + } + if abort_acknowledged: + ack["reason"] = error_msg self._sleep_wakeup_comm.send( - { - "status": "ok" if error_msg is None else "error", - "error": error_msg, - "op_id": op_id, - "phase": action, - "has_mnnvl_resources": has_mnnvl_resources, - }, + ack, dest=0, tag=_SleepWakeupTag.ACK, ) diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index e56bffd9e0b9..5b3da7f0c19b 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -651,9 +651,12 @@ def _multi_rank_sleep_wakeup( prepared peer before rank-0 enters the subgroup checkpoint collectives. Otherwise execute the local VMM operation first. 4. Execute the selected VMM and MNNVL operations and collect bounded - COMMIT ACKs. If PREPARE or pre-commit local execution fails, send - ABORT so peers leave the control barrier without changing VMM state. - 5. Exit ``control_action()``, resuming rank-0's event loop. + COMMIT ACKs. If PREPARE or the initial rank-0 synchronization fails, + send ABORT before any rank changes VMM state. Once COMMIT or local + mutation starts, any error fail-stops the distributed worker because + rank state can no longer be reconciled safely. + 5. Exit ``control_action()`` and resume rank-0 only after a successful + operation or a recoverable pre-COMMIT abort. Args: action: ``"sleep"`` or ``"wakeup"``. @@ -713,9 +716,12 @@ def _multi_rank_sleep_wakeup( abort_sent = False commit_ranks = [] commit_sent_early = False + local_commit_started = False + abort_incomplete = False def send_abort(reason: str, ranks: Optional[list[int]] = None) -> list[int]: + nonlocal abort_incomplete abort_ranks = [] abort_dests = ranks if ranks is not None else range( 1, world_size) @@ -734,6 +740,7 @@ def send_abort(reason: str, ) abort_ranks.append(abort_dest) except Exception as abort_exc: + abort_incomplete = True abort_error = ( "rank 0 failed to send sleep/wakeup abort " f"to rank {abort_dest}: {abort_exc}") @@ -747,6 +754,7 @@ def send_abort(reason: str, def drain_acks(ranks: list[int], phase: _SleepWakeupAction) -> list[dict]: + nonlocal abort_incomplete received_acks = [] ack_deadline = time.monotonic() + _SLEEP_WAKEUP_ACK_TIMEOUT_S for src in ranks: @@ -757,6 +765,8 @@ def drain_acks(ranks: list[int], expected_op_id=op_id, expected_phase=phase) except Exception as exc: + if phase == _SleepWakeupAction.ABORT: + abort_incomplete = True errors.append( f"rank 0 failed to receive {phase} ACK from " f"rank {src}: {exc}") @@ -770,6 +780,8 @@ def drain_acks(ranks: list[int], continue received_acks.append(ack) if ack.get("status") != "ok": + if phase == _SleepWakeupAction.ABORT: + abort_incomplete = True errors.append( ack.get("error") or f"rank {src} returned unknown {phase} ACK") @@ -841,11 +853,14 @@ def send_commits() -> list[int]: # MNNVL resource hooks contain subgroup handle exchange, # so peers must enter COMMIT before rank 0 executes them. if has_mnnvl_resources: + # A failed MPI send has uncertain delivery, so any + # attempted COMMIT requires the bounded local phase. + commit_sent_early = bool(prepared_ranks) commit_ranks = send_commits() - commit_sent_early = True if not errors or commit_sent_early: torch.cuda.synchronize() + local_commit_started = True if action == _SleepWakeupAction.SLEEP: run_mnnvl = (getattr( self.engine, "_run_mnnvl_checkpoint_resources", @@ -864,7 +879,7 @@ def send_commits() -> list[int]: None) if has_mnnvl_resources else None) if run_mnnvl is not None: run_mnnvl(target_action, tags) - except (RuntimeError, TimeoutError, torch.OutOfMemoryError) as exc: + except Exception as exc: local_error = (f"rank 0 '{action}' failed: {exc}\n" f"{traceback.format_exc()}") logger.error( @@ -888,9 +903,30 @@ def send_commits() -> list[int]: drain_acks(commit_ranks, _SleepWakeupAction.COMMIT) if errors: - raise RuntimeError( + operation_error = RuntimeError( f"{action}() failed on {len(errors)} rank(s):\n" + "\n".join(errors)) + if commit_sent_early or local_commit_started or abort_incomplete: + self._fail_stop_divergent_sleep_wakeup(operation_error) + raise operation_error + + def _fail_stop_divergent_sleep_wakeup(self, error: RuntimeError) -> None: + """Stop every rank after a sleep/wakeup operation may have diverged.""" + from tensorrt_llm._torch.pyexecutor.hang_detector import \ + propagate_hard_kill + + self._set_fatal_error(error) + if self.engine is not None: + if getattr(self.engine, "_fatal_error", None) is None: + self.engine._fatal_error = error + self.engine.is_shutdown = True + try: + logger.critical( + "Distributed sleep/wakeup state may have diverged; hard-killing all ranks: %s", + error, + ) + finally: + propagate_hard_kill() def sleep(self, sleep_tags: list[str]) -> None: """Release GPU virtual memory for the specified memory type tags. diff --git a/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py b/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py index 0c3edfadee6f..aa61439cc41c 100644 --- a/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py +++ b/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py @@ -170,8 +170,17 @@ def test_checkpoint_prepare_rejects_communicator_size_mismatch() -> None: def test_checkpoint_prepare_timeout_fails_closed(monkeypatch) -> None: lifecycle, memory, _ = _make_lifecycle() - request = SimpleNamespace(test=lambda: (False, None)) - memory.comm = SimpleNamespace(iallgather=lambda value: request) + request = SimpleNamespace( + test=lambda: (False, None), + Cancel=Mock(), + Free=Mock(), + ) + memory.comm = SimpleNamespace( + Get_rank=lambda: 0, + Get_size=lambda: 2, + irecv=lambda source, tag: request, + isend=lambda value, dest, tag: SimpleNamespace(test=lambda: (True, None)), + ) monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S", 0.0) @@ -555,8 +564,17 @@ def test_two_sided_checkpoint_prepare_timeout_fails_closed( ) -> None: instances = WeakSet() monkeypatch.setattr(NVLinkTwoSided, "_INSTANCES", instances) - request = SimpleNamespace(test=lambda: (False, None)) - comm = SimpleNamespace(iallgather=lambda value: request) + request = SimpleNamespace( + test=lambda: (False, None), + Cancel=Mock(), + Free=Mock(), + ) + comm = SimpleNamespace( + Get_rank=lambda: 0, + Get_size=lambda: 2, + irecv=lambda source, tag: request, + isend=lambda value, dest, tag: SimpleNamespace(test=lambda: (True, None)), + ) main_workspace = Mock(mapped=True, comm=comm) prepare_workspace = Mock(mapped=True) monkeypatch.setattr(mnnvl.MnnvlMoe, "moe_workspace", main_workspace) diff --git a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py index be1df3fc19ec..5a20dde00f6a 100644 --- a/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py +++ b/tests/unittest/_torch/test_mnnvl_memory_lifecycle.py @@ -51,8 +51,24 @@ class _TestMnnvlMemory(mnnvl.MnnvlMemory): def test_checkpoint_allgather_timeout_is_bounded(monkeypatch): - request = SimpleNamespace(test=lambda: (False, None)) - comm = SimpleNamespace(iallgather=lambda value: request) + orphaned_requests = [] + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_ORPHANED_REQUESTS", orphaned_requests) + receive_request = SimpleNamespace( + test=lambda: (False, None), + Cancel=Mock(), + Free=Mock(), + ) + send_request = SimpleNamespace( + test=lambda: (False, None), + Cancel=Mock(), + Free=Mock(), + ) + comm = SimpleNamespace( + Get_rank=lambda: 0, + Get_size=lambda: 2, + irecv=lambda **kwargs: receive_request, + isend=lambda value, **kwargs: send_request, + ) monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S", 0.0) @@ -62,6 +78,82 @@ def test_checkpoint_allgather_timeout_is_bounded(monkeypatch): None, operation="mapping readiness", ) + receive_request.Cancel.assert_called_once_with() + receive_request.Free.assert_not_called() + send_request.Cancel.assert_called_once_with() + send_request.Free.assert_not_called() + assert orphaned_requests == [receive_request, send_request] + + +def test_checkpoint_request_cleanup_retains_request_when_cancel_fails(monkeypatch): + orphaned_requests = [] + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_ORPHANED_REQUESTS", orphaned_requests) + request = SimpleNamespace( + test=Mock(return_value=(False, None)), + Cancel=Mock(side_effect=RuntimeError("cancel failed")), + Free=Mock(), + ) + monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) + + mnnvl._cancel_checkpoint_requests([request]) + + request.Cancel.assert_called_once_with() + request.test.assert_called_once_with() + request.Free.assert_not_called() + assert orphaned_requests == [request] + + +def test_checkpoint_allgather_uses_bounded_object_requests(): + class _ReceiveRequest: + def __init__(self, value): + self.value = value + + def test(self): + return True, self.value + + comm = SimpleNamespace( + Get_rank=lambda: 1, + Get_size=lambda: 3, + irecv=lambda source, tag: _ReceiveRequest(f"rank-{source}"), + isend=lambda value, dest, tag: SimpleNamespace(test=lambda: (True, None)), + allgather=Mock(side_effect=AssertionError("blocking fallback used")), + ) + + result = mnnvl._checkpoint_allgather(comm, "rank-1", operation="test") + + assert result == ["rank-0", "rank-1", "rank-2"] + comm.allgather.assert_not_called() + + +def test_checkpoint_allgather_post_failure_retires_partial_requests(): + requests = [ + SimpleNamespace(test=Mock(return_value=(True, None)), Cancel=Mock(), Free=Mock()), + SimpleNamespace(test=Mock(return_value=(True, None)), Cancel=Mock(), Free=Mock()), + SimpleNamespace(test=Mock(return_value=(True, None)), Cancel=Mock(), Free=Mock()), + ] + receive_requests = iter(requests[:2]) + send_requests = iter([requests[2], RuntimeError("post failed")]) + + def isend(value, dest, tag): + result = next(send_requests) + if isinstance(result, Exception): + raise result + return result + + comm = SimpleNamespace( + Get_rank=lambda: 0, + Get_size=lambda: 3, + irecv=lambda source, tag: next(receive_requests), + isend=isend, + ) + + with pytest.raises(RuntimeError, match="post failed"): + mnnvl._checkpoint_allgather(comm, None, operation="test") + + for request in requests: + request.Cancel.assert_called_once_with() + request.test.assert_called_once_with() + request.Free.assert_not_called() @pytest.fixture @@ -185,11 +277,16 @@ def test_checkpoint_restore_membership_timeout_is_terminal(memory, monkeypatch): monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S", 0.0) - request = SimpleNamespace(test=lambda: (False, None)) + receive_request = SimpleNamespace( + test=lambda: (False, None), + Cancel=Mock(), + Free=Mock(), + ) comm = SimpleNamespace( Get_rank=lambda: 0, Get_size=lambda: 2, - iallgather=lambda value: request, + irecv=lambda source, tag: receive_request, + isend=lambda value, dest, tag: SimpleNamespace(test=lambda: (True, None)), ) with pytest.raises(TimeoutError, match="communicator membership"): obj.checkpoint_restore(comm) @@ -211,20 +308,31 @@ def test_checkpoint_restore_timeout_cleans_locally_mapped_handles(memory, monkey monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_TIMEOUT_S", 0.0) monkeypatch.setattr(mnnvl, "_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S", 0.0) + pending = object() + class Request: - def __init__(self, result=None): + def __init__(self, result=pending): self.result = result def test(self): - return (self.result is not None, self.result) + return (self.result is not pending, None if self.result is pending else self.result) + + def Cancel(self): + pass + + def Free(self): + pass class TimeoutComm(_FakeComm): def __init__(self): super().__init__() - self.requests = [Request([0, 1]), Request()] + self.receive_requests = iter([Request(1), Request()]) + + def irecv(self, source, tag): + return next(self.receive_requests) - def iallgather(self, value): - return self.requests.pop(0) + def isend(self, value, dest, tag): + return Request(None) with pytest.raises(TimeoutError, match="mapping readiness"): obj.checkpoint_restore(TimeoutComm()) @@ -408,7 +516,8 @@ def test_close_detached_memory_only_frees_va(memory, monkeypatch): def test_create_and_map_handles_cleans_partial_allocation(monkeypatch): comm = Mock() - comm.iallgather = None + comm.isend = None + comm.irecv = None comm.Get_rank.return_value = 0 comm.Get_size.return_value = 2 comm.allgather.side_effect = lambda value: [ @@ -459,7 +568,8 @@ def test_create_and_map_handles_cleans_partial_allocation(monkeypatch): def test_create_and_map_handles_releases_local_handle_on_export_failure(monkeypatch): comm = Mock() - comm.iallgather = None + comm.isend = None + comm.irecv = None comm.Get_rank.return_value = 0 comm.Get_size.return_value = 2 comm.allgather.side_effect = lambda value: [ @@ -505,7 +615,8 @@ def test_create_and_map_handles_releases_local_handle_on_export_failure(monkeypa ) def test_create_and_map_handles_preserves_pidfd_error_hint(monkeypatch, errno_value, expected_hint): comm = Mock() - comm.iallgather = None + comm.isend = None + comm.irecv = None comm.Get_rank.return_value = 0 comm.Get_size.return_value = 2 comm.allgather.side_effect = lambda value: ( @@ -554,7 +665,8 @@ def test_create_and_map_handles_preserves_pidfd_error_hint(monkeypatch, errno_va def test_create_and_map_handles_close_failure_does_not_mask_original_error(monkeypatch): comm = Mock() - comm.iallgather = None + comm.isend = None + comm.irecv = None comm.Get_rank.return_value = 0 comm.Get_size.return_value = 2 comm.allgather.side_effect = lambda value: ( diff --git a/tests/unittest/executor/test_sleep_collective_rpc_guards.py b/tests/unittest/executor/test_sleep_collective_rpc_guards.py index 90add390febd..f39d96001e31 100644 --- a/tests/unittest/executor/test_sleep_collective_rpc_guards.py +++ b/tests/unittest/executor/test_sleep_collective_rpc_guards.py @@ -567,8 +567,9 @@ def recv(self, source, tag): if source == 1: return {"status": "ok", "op_id": self.op_id} return { - "status": "error", - "error": "rank 0 aborted sleep/wakeup before local execution", + "status": "ok", + "error": None, + "reason": "rank 0 aborted sleep/wakeup before local execution", "op_id": self.op_id, } @@ -589,7 +590,6 @@ def _noop_control_action(**kwargs): _sleep_wakeup_comm=FakeComm(), control_action=_noop_control_action, ) - with ( patch("tensorrt_llm._torch.virtual_memory.release_with_tag") as release, patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), @@ -602,7 +602,6 @@ def _noop_control_action(**kwargs): msg = str(exc_info.value) assert "simulated mid-broadcast send failure" in msg - assert "rank 0 aborted sleep/wakeup" in msg assert release.call_count == 0 assert send_calls == [ (_SleepWakeupAction.PREPARE, 1, _SleepWakeupTag.ACTION), @@ -676,6 +675,8 @@ def _noop_control_action(**kwargs): _sleep_wakeup_comm=FakeComm(), control_action=_noop_control_action, ) + fail_stop = MagicMock() + w._fail_stop_divergent_sleep_wakeup = fail_stop monkeypatch.setattr(py_executor, "_SLEEP_WAKEUP_ACK_TIMEOUT_S", 0.0) monkeypatch.setattr(py_executor, "_SLEEP_WAKEUP_ACK_POLL_INTERVAL_S", 0.0) @@ -690,6 +691,72 @@ def _noop_control_action(**kwargs): with pytest.raises(RuntimeError, match="timed out waiting"): w._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.KV_CACHE]) + fail_stop.assert_called_once() + + def test_abort_send_failure_fail_stops_prepared_worker(self): + """Rank 0 must not resume if a prepared peer cannot be aborted.""" + from unittest.mock import patch + + from tensorrt_llm._torch.pyexecutor.py_executor import _SleepWakeupAction + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + worker, _ = _make_proto_worker( + [{"status": "error", "error": "prepare failed"}], + world_size=2, + ) + original_send = worker.engine._sleep_wakeup_comm.send + + def fail_abort(payload, *args, **kwargs): + original_send(payload, *args, **kwargs) + if payload["action"] == _SleepWakeupAction.ABORT: + raise RuntimeError("abort send failed") + + worker.engine._sleep_wakeup_comm.send = fail_abort + fail_stop = MagicMock() + worker._fail_stop_divergent_sleep_wakeup = fail_stop + + with ( + patch("tensorrt_llm._torch.virtual_memory.release_with_tag") as release, + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize"), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + with pytest.raises(RuntimeError, match="abort send failed"): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.KV_CACHE]) + + release.assert_not_called() + fail_stop.assert_called_once() + + def test_abort_error_ack_fail_stops_prepared_worker(self): + """A genuine ABORT handler error leaves peer recovery inconclusive.""" + from unittest.mock import patch + + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + worker, _ = _make_proto_worker( + [ + {"status": "error", "error": "prepare failed"}, + {"status": "error", "error": "abort handler failed"}, + ], + world_size=2, + ) + fail_stop = MagicMock() + worker._fail_stop_divergent_sleep_wakeup = fail_stop + + with ( + patch("tensorrt_llm._torch.virtual_memory.release_with_tag") as release, + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize"), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + with pytest.raises(RuntimeError, match="abort handler failed"): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.KV_CACHE]) + + release.assert_not_called() + fail_stop.assert_called_once() + def test_commit_send_failure_aborts_uncommitted_rank(self): """A prepared rank that misses COMMIT must receive ABORT to unblock.""" import threading @@ -743,6 +810,8 @@ def _noop_control_action(**kwargs): _sleep_wakeup_comm=FakeComm(), control_action=_noop_control_action, ) + fail_stop = MagicMock() + w._fail_stop_divergent_sleep_wakeup = fail_stop with ( patch("tensorrt_llm._torch.virtual_memory.release_with_tag") as release, @@ -755,6 +824,7 @@ def _noop_control_action(**kwargs): w._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.KV_CACHE]) release.assert_called_once() + fail_stop.assert_called_once() assert send_calls == [ (_SleepWakeupAction.PREPARE, 1, _SleepWakeupTag.ACTION), (_SleepWakeupAction.PREPARE, 2, _SleepWakeupTag.ACTION), @@ -822,6 +892,8 @@ def control_action(**kwargs): _has_mnnvl_checkpoint_resources=lambda tags: True, _run_mnnvl_checkpoint_resources=run_mnnvl, ) + fail_stop = Mock() + worker._fail_stop_divergent_sleep_wakeup = fail_stop with ( patch("tensorrt_llm._torch.virtual_memory.release_with_tag") as release, @@ -838,6 +910,152 @@ def control_action(**kwargs): [ExecutorMemoryType.MODEL_ENGINE_MAIN], ) release.assert_called_once_with(ExecutorMemoryType.MODEL_ENGINE_MAIN) + fail_stop.assert_called_once() + + def test_mnnvl_all_commit_send_errors_are_treated_as_uncertain_delivery(self): + """A send error cannot prove that no peer received MNNVL COMMIT.""" + import threading + from contextlib import contextmanager + from unittest.mock import Mock, patch + + from tensorrt_llm._torch.pyexecutor.py_executor import _SleepWakeupAction + from tensorrt_llm.executor.base_worker import BaseWorker + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + class FakeComm: + def __init__(self): + self.op_id = None + self.ack_phases = {1: []} + + def send(self, payload, dest, tag): + self.op_id = payload.get("op_id", self.op_id) + if payload["action"] == _SleepWakeupAction.COMMIT: + raise RuntimeError("injected MNNVL commit send failure") + self.ack_phases[dest].append(payload["action"]) + + def iprobe(self, source, tag): + return True + + def recv(self, source, tag): + return { + "status": "ok", + "op_id": self.op_id, + "phase": self.ack_phases[source].pop(0), + "has_mnnvl_resources": True, + } + + @contextmanager + def control_action(**kwargs): + yield None + + worker = object.__new__(BaseWorker) + worker._backend = "pytorch" + worker.rank = 0 + worker.llm_args = SimpleNamespace( + backend="pytorch", + parallel_config=SimpleNamespace(world_size=2), + sleep_config=object(), + ) + run_mnnvl = Mock() + worker.engine = SimpleNamespace( + _sleep_wakeup_lock=threading.Lock(), + _sleep_wakeup_comm=FakeComm(), + control_action=control_action, + _has_mnnvl_checkpoint_resources=lambda tags: True, + _run_mnnvl_checkpoint_resources=run_mnnvl, + ) + fail_stop = Mock() + worker._fail_stop_divergent_sleep_wakeup = fail_stop + + with ( + patch("tensorrt_llm._torch.virtual_memory.release_with_tag") as release, + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize"), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + with pytest.raises(RuntimeError, match="commit send failure"): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.MODEL_ENGINE_MAIN]) + + run_mnnvl.assert_called_once_with( + _SleepWakeupAction.SLEEP, + [ExecutorMemoryType.MODEL_ENGINE_MAIN], + ) + release.assert_called_once_with(ExecutorMemoryType.MODEL_ENGINE_MAIN) + fail_stop.assert_called_once() + + def test_postcommit_non_runtime_error_fail_stops_worker(self): + """Every ordinary exception after COMMIT must take the fail-stop path.""" + from unittest.mock import Mock, patch + + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + worker, _ = _make_proto_worker( + [{"status": "ok"}, {"status": "ok"}], + world_size=2, + ) + worker.engine._has_mnnvl_checkpoint_resources = lambda tags: True + worker.engine._run_mnnvl_checkpoint_resources = Mock( + side_effect=ValueError("invalid checkpoint state") + ) + fail_stop = Mock() + worker._fail_stop_divergent_sleep_wakeup = fail_stop + + with ( + patch("tensorrt_llm._torch.virtual_memory.release_with_tag"), + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize"), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + with pytest.raises(RuntimeError, match="invalid checkpoint state"): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.MODEL_ENGINE_MAIN]) + + fail_stop.assert_called_once() + + def test_divergent_commit_fail_stops_worker(self): + """A potentially divergent operation poisons the worker before returning.""" + from unittest.mock import patch + + from tensorrt_llm.executor.base_worker import BaseWorker + + worker = object.__new__(BaseWorker) + worker._fatal_error = None + worker.engine = SimpleNamespace(_fatal_error=None, is_shutdown=False) + error = RuntimeError("divergent sleep state") + + with patch("tensorrt_llm._torch.pyexecutor.hang_detector.propagate_hard_kill") as hard_kill: + worker._fail_stop_divergent_sleep_wakeup(error) + + assert worker._fatal_error is error + assert worker.engine._fatal_error is error + assert worker.engine.is_shutdown + hard_kill.assert_called_once_with() + + def test_precommit_local_failure_remains_recoverable(self): + """A failure before local mutation aborts peers without fail-stopping.""" + from unittest.mock import Mock, patch + + from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + + worker, _ = _make_proto_worker( + [{"status": "ok"}, {"status": "ok"}], + world_size=2, + ) + fail_stop = Mock() + worker._fail_stop_divergent_sleep_wakeup = fail_stop + + with ( + patch("tensorrt_llm._torch.virtual_memory.release_with_tag"), + patch("tensorrt_llm._torch.virtual_memory.materialize_with_tag"), + patch("torch.cuda.synchronize", side_effect=RuntimeError("sync failed")), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + ): + with pytest.raises(RuntimeError, match="sync failed"): + worker._multi_rank_sleep_wakeup("sleep", [ExecutorMemoryType.KV_CACHE]) + + fail_stop.assert_not_called() class TestListenerUncaughtExceptionSendsErrorAck: @@ -979,7 +1197,7 @@ def send(self, payload, dest, tag): } ] - def test_abort_unblocks_control_request_and_sends_error_ack(self): + def test_abort_unblocks_control_request_and_sends_success_ack(self): """Abort messages release the non-rank executor control barrier.""" import threading from unittest.mock import patch @@ -1031,8 +1249,9 @@ def send(self, payload, dest, tag): assert executor.control_action_done.is_set() assert not executor.control_request_barrier.is_set() assert sent_acks - assert sent_acks[0]["status"] == "error" - assert "rank 0 send failed" in sent_acks[0]["error"] + assert sent_acks[0]["status"] == "ok" + assert sent_acks[0]["error"] is None + assert "rank 0 send failed" in sent_acks[0]["reason"] def test_abort_before_control_barrier_unblocks_later_control_request(self): """An early ABORT is recorded and later consumed by matching control.""" @@ -1095,7 +1314,9 @@ def send(self, payload, dest, tag): assert executor.control_requests == [] assert not executor.control_request_barrier.is_set() assert sent_acks - assert sent_acks[0]["status"] == "error" + assert sent_acks[0]["status"] == "ok" + assert sent_acks[0]["error"] is None + assert "rank 0 send failed" in sent_acks[0]["reason"] def test_shutdown_sends_ack_before_listener_exits(self): """Shutdown messages are acknowledged so rank-0 can drain them.""" @@ -1188,6 +1409,8 @@ def test_local_failure_drains_all_peer_acks(self): # Two peers prepare ok, then abort ok after rank-0 local failure. responses = [{"status": "ok"}, {"status": "ok"}] w, recv_calls = _make_proto_worker(responses, world_size=3) + fail_stop = MagicMock() + w._fail_stop_divergent_sleep_wakeup = fail_stop with ( patch( @@ -1209,6 +1432,7 @@ def test_local_failure_drains_all_peer_acks(self): f"{len(recv_calls)}; stale ACKs would corrupt the next " "sleep/wakeup call." ) + fail_stop.assert_called_once() def test_local_failure_plus_peer_error_both_reported(self): """When rank-0 fails locally and a peer also errors, both messages must appear. @@ -1226,6 +1450,8 @@ def test_local_failure_plus_peer_error_both_reported(self): {"status": "error", "error": "rank 2 also failed"}, ] w, _ = _make_proto_worker(responses) + fail_stop = MagicMock() + w._fail_stop_divergent_sleep_wakeup = fail_stop with ( patch( @@ -1243,6 +1469,7 @@ def test_local_failure_plus_peer_error_both_reported(self): msg = str(exc_info.value) assert "rank 0 VMM fault" in msg assert "rank 2 also failed" in msg + fail_stop.assert_called_once() class TestSingleRankLockAcquired: From ec5139e2891c2d15c08b09bd24025dae6a86c4e3 Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Tue, 18 Aug 2026 11:53:55 -0400 Subject: [PATCH 11/14] fix(moe): correct checkpoint failure logging Signed-off-by: Hannah Zhang --- tensorrt_llm/_mnnvl_utils.py | 4 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 76 +++++++------------ tensorrt_llm/executor/base_worker.py | 38 +++------- 3 files changed, 41 insertions(+), 77 deletions(-) diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index f1d7a8479bdf..afd3d8807e72 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/tensorrt_llm/_mnnvl_utils.py @@ -79,8 +79,8 @@ def _cancel_checkpoint_requests(requests: list[Any]) -> None: # enclosing fail-closed path terminates the worker. _MNNVL_CHECKPOINT_ORPHANED_REQUESTS.extend(pending_requests) logger.error( - "Retaining %d active MNNVL checkpoint requests until worker termination", - len(pending_requests), + f"Retaining {len(pending_requests)} active MNNVL checkpoint requests " + "until worker termination" ) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index edc0dd141773..b46d350da389 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -233,23 +233,16 @@ def _recv_sleep_wakeup_ack_until(comm, if (expected_op_id is not None and ack.get("op_id") != expected_op_id): logger.warning( - "Ignoring stale sleep/wakeup ACK from rank %d for op_id=%s " - "(expected op_id=%s).", - source, - ack.get("op_id"), - expected_op_id, + f"Ignoring stale sleep/wakeup ACK from rank {source} for " + f"op_id={ack.get('op_id')} (expected op_id={expected_op_id})." ) continue ack_phase = ack.get("phase") if (expected_phase is not None and ack_phase is not None and ack_phase != expected_phase): logger.warning( - "Ignoring stale sleep/wakeup ACK from rank %d for phase=%s " - "(expected phase=%s).", - source, - ack_phase, - expected_phase, - ) + f"Ignoring stale sleep/wakeup ACK from rank {source} for " + f"phase={ack_phase} (expected phase={expected_phase}).") continue return ack if time.monotonic() >= deadline: @@ -1343,10 +1336,8 @@ def _shutdown_sleep_wakeup_listeners(self) -> None: return if self.dist.rank == 0: - logger.info( - "Sending shutdown to %d sleep/wakeup listener thread(s).", - self.dist.world_size - 1, - ) + logger.info(f"Sending shutdown to {self.dist.world_size - 1} " + "sleep/wakeup listener thread(s).") shutdown_ack_deadline = (time.monotonic() + _SLEEP_WAKEUP_ACK_TIMEOUT_S) shutdown_errors = [] @@ -1367,7 +1358,7 @@ def _shutdown_sleep_wakeup_listeners(self) -> None: f"rank {dest} shutdown send failed: {exc}") logger.warning( "Failed to send sleep/wakeup listener shutdown to " - "rank %d: %s", dest, exc) + f"rank {dest}: {exc}") for src in shutdown_ranks: try: ack = _recv_sleep_wakeup_ack_until(self._sleep_wakeup_comm, @@ -1378,7 +1369,7 @@ def _shutdown_sleep_wakeup_listeners(self) -> None: f"rank {src} shutdown ACK recv failed: {exc}") logger.warning( "Failed to receive sleep/wakeup listener shutdown ACK " - "from rank %d: %s", src, exc) + f"from rank {src}: {exc}") continue if ack.get("status") != "ok": shutdown_errors.append( @@ -1386,18 +1377,16 @@ def _shutdown_sleep_wakeup_listeners(self) -> None: or f"rank {src} returned unknown shutdown ACK") if shutdown_errors: logger.warning( - "Sleep/wakeup listener shutdown completed with errors: %s", - "; ".join(shutdown_errors)) + "Sleep/wakeup listener shutdown completed with errors: " + f"{'; '.join(shutdown_errors)}") elif self._sleep_wakeup_listener_thread is not None: self._sleep_wakeup_listener_thread.join( timeout=_SLEEP_WAKEUP_LISTENER_JOIN_TIMEOUT_S) if self._sleep_wakeup_listener_thread.is_alive(): logger.warning( - "Sleep/wakeup listener thread did not exit within %.1f " - "seconds on rank %d.", - _SLEEP_WAKEUP_LISTENER_JOIN_TIMEOUT_S, - self.dist.rank, - ) + "Sleep/wakeup listener thread did not exit within " + f"{_SLEEP_WAKEUP_LISTENER_JOIN_TIMEOUT_S:.1f} seconds on " + f"rank {self.dist.rank}.") def _record_sleep_wakeup_abort(self, control_id: str, error_msg: str) -> None: @@ -3023,8 +3012,8 @@ def _sleep_wakeup_listener_loop(self) -> None: tag=_SleepWakeupTag.ACTION) if msg.get("action") == _SleepWakeupAction.SHUTDOWN: logger.debug( - "Sleep/wakeup listener (rank %d): received shutdown, " - "exiting.", self.dist.rank) + f"Sleep/wakeup listener (rank {self.dist.rank}): " + "received shutdown, exiting.") self._sleep_wakeup_comm.send( { "status": "ok", @@ -3077,7 +3066,7 @@ def _sleep_wakeup_listener_loop(self) -> None: op_id, error_msg) release_control_request = False abort_acknowledged = True - logger.warning("Sleep/wakeup listener: %s", error_msg) + logger.warning(f"Sleep/wakeup listener: {error_msg}") elif action in (_SleepWakeupAction.PREPARE, _SleepWakeupAction.COMMIT, _SleepWakeupAction.SLEEP, @@ -3090,8 +3079,8 @@ def _sleep_wakeup_listener_loop(self) -> None: f"stale control message for op_id={op_id}; " f"active control_id={active_control_id}") release_control_request = False - logger.warning("Sleep/wakeup listener: %s", - error_msg) + logger.warning( + f"Sleep/wakeup listener: {error_msg}") else: torch.cuda.synchronize() if action == _SleepWakeupAction.PREPARE: @@ -3116,19 +3105,17 @@ def _sleep_wakeup_listener_loop(self) -> None: error_msg = ( f"unknown target action '{target_action}'") logger.warning( - "Sleep/wakeup listener: %s, ignoring.", - error_msg) + f"Sleep/wakeup listener: {error_msg}, ignoring." + ) else: error_msg = f"unknown action '{action}'" - logger.warning("Sleep/wakeup listener: %s, ignoring.", - error_msg) + logger.warning( + f"Sleep/wakeup listener: {error_msg}, ignoring.") except (KeyError, TypeError, ValueError, RuntimeError, TimeoutError, torch.OutOfMemoryError) as exc: error_msg = (f"rank {self.dist.rank} '{action}' failed: " f"{exc}\n{traceback.format_exc()}") - logger.error("Sleep/wakeup listener: error executing '%s':", - action, - exc_info=True) + logger.error(f"Sleep/wakeup listener: {error_msg}") finally: # Always ACK so rank-0 does not deadlock; carry error # details so rank-0 can raise after all ranks respond. @@ -3140,14 +3127,9 @@ def _sleep_wakeup_listener_loop(self) -> None: exc = sys.exc_info()[1] error_msg = ( f"rank {self.dist.rank} '{action}' failed with " - f"uncaught {type(exc).__name__}: {exc!r}") - logger.error( - "Sleep/wakeup listener: uncaught exception on " - "rank %d during '%s':", - self.dist.rank, - action, - exc_info=True, - ) + f"uncaught {type(exc).__name__}: {exc!r}\n" + f"{traceback.format_exc()}") + logger.error(f"Sleep/wakeup listener: {error_msg}") # Unblock the executor loop that is waiting in # _handle_control_request(). Clear control_request_barrier # first so that it is clean for the next sleep/wakeup cycle @@ -4558,10 +4540,8 @@ def _handle_control_request(self): pending_abort = self._pop_sleep_wakeup_abort(control_id) if pending_abort is not None: logger.warning( - "[control_action] skipping aborted control request %s: %s", - control_id, - pending_abort, - ) + f"[control_action] skipping aborted control request {control_id}: " + f"{pending_abort}") self.control_request_barrier.set() self.control_request_barrier.clear() self._active_control_id = None diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 5b3da7f0c19b..c9d2a27ded40 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -746,10 +746,8 @@ def send_abort(reason: str, f"to rank {abort_dest}: {abort_exc}") errors.append(abort_error) logger.error( - "_multi_rank_sleep_wakeup: %s", - abort_error, - exc_info=True, - ) + f"_multi_rank_sleep_wakeup: {abort_error}\n" + f"{traceback.format_exc()}") return abort_ranks def drain_acks(ranks: list[int], @@ -771,11 +769,8 @@ def drain_acks(ranks: list[int], f"rank 0 failed to receive {phase} ACK from " f"rank {src}: {exc}") logger.error( - "_multi_rank_sleep_wakeup: failed to receive %s " - "ACK from rank %d", - phase, - src, - exc_info=True, + "_multi_rank_sleep_wakeup: failed to receive " + f"{phase} ACK from rank {src}\n{traceback.format_exc()}" ) continue received_acks.append(ack) @@ -805,10 +800,8 @@ def send_commits() -> list[int]: errors.append(commit_error) failed_ranks.append(dest) logger.error( - "_multi_rank_sleep_wakeup: %s", - commit_error, - exc_info=True, - ) + f"_multi_rank_sleep_wakeup: {commit_error}\n" + f"{traceback.format_exc()}") if failed_ranks: abort_ranks = send_abort("\n".join(errors), ranks=failed_ranks) @@ -831,11 +824,8 @@ def send_commits() -> list[int]: f"rank 0 failed to send '{action}' prepare to rank " f"{dest}: {exc}") errors.append(send_error) - logger.error( - "_multi_rank_sleep_wakeup: %s", - send_error, - exc_info=True, - ) + logger.error(f"_multi_rank_sleep_wakeup: {send_error}\n" + f"{traceback.format_exc()}") abort_ranks = send_abort(send_error) abort_sent = True drain_acks(prepared_ranks, _SleepWakeupAction.PREPARE) @@ -882,11 +872,7 @@ def send_commits() -> list[int]: except Exception as exc: local_error = (f"rank 0 '{action}' failed: {exc}\n" f"{traceback.format_exc()}") - logger.error( - "_multi_rank_sleep_wakeup: rank-0 local %s failed:", - action, - exc_info=True, - ) + logger.error(f"_multi_rank_sleep_wakeup: {local_error}") finally: if local_error: errors.append(local_error) @@ -921,10 +907,8 @@ def _fail_stop_divergent_sleep_wakeup(self, error: RuntimeError) -> None: self.engine._fatal_error = error self.engine.is_shutdown = True try: - logger.critical( - "Distributed sleep/wakeup state may have diverged; hard-killing all ranks: %s", - error, - ) + logger.critical("Distributed sleep/wakeup state may have diverged; " + f"hard-killing all ranks: {error}") finally: propagate_hard_kill() From a6cb14a12a3f2599c1445dc29e9ee3af08eead9a Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Tue, 18 Aug 2026 15:36:46 -0400 Subject: [PATCH 12/14] fix(runtime): enforce sleep wakeup admission state Signed-off-by: Hannah Zhang --- tensorrt_llm/_mnnvl_utils.py | 42 ++++++- .../_torch/distributed/moe_alltoall.py | 12 +- .../_torch/mnnvl_alltoall_workspace.py | 23 ++-- .../modules/fused_moe/MOE_DEVELOPER_GUIDE.md | 12 +- .../fused_moe/communication/__init__.py | 3 +- .../modules/fused_moe/communication/base.py | 26 ++++- .../communication/nvlink_one_sided.py | 11 +- .../communication/nvlink_two_sided.py | 11 +- .../pyexecutor/executor_request_queue.py | 110 ++++++++++++++++++ tensorrt_llm/_torch/pyexecutor/py_executor.py | 72 +++++++++--- tensorrt_llm/executor/base_worker.py | 101 +++++++++++----- .../executor/test_executor_request_queue.py | 67 ++++++++++- .../test_sleep_collective_rpc_guards.py | 102 ++++++++++++++-- 13 files changed, 507 insertions(+), 85 deletions(-) diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index afd3d8807e72..6c2f5153a4ea 100644 --- a/tensorrt_llm/_mnnvl_utils.py +++ b/tensorrt_llm/_mnnvl_utils.py @@ -21,7 +21,7 @@ import time from dataclasses import dataclass from enum import Enum -from typing import Any, List, Optional, Union +from typing import Any, List, Optional, Protocol, Union import pynvml import torch @@ -43,6 +43,34 @@ _MNNVL_CHECKPOINT_ORPHANED_REQUESTS: list[Any] = [] +class MnnvlCheckpointCommunicator(Protocol): + """Structural contract required by bounded MNNVL checkpoint collectives.""" + + def Get_rank(self) -> int: + """Return this process's rank in the communicator.""" + ... + + def Get_size(self) -> int: + """Return the communicator size.""" + ... + + def barrier(self) -> None: + """Synchronize all communicator members.""" + ... + + def allgather(self, value: Any) -> list[Any]: + """Gather a Python object from every communicator member.""" + ... + + def isend(self, value: Any, *, dest: int, tag: int) -> Any: + """Start a nonblocking Python-object send.""" + ... + + def irecv(self, *, source: int, tag: int) -> Any: + """Start a nonblocking Python-object receive.""" + ... + + def _cancel_checkpoint_requests(requests: list[Any]) -> None: """Bound cancellation cleanup so timeout handling cannot hang in MPI.""" for request in requests: @@ -84,7 +112,12 @@ def _cancel_checkpoint_requests(requests: list[Any]) -> None: ) -def _checkpoint_allgather(comm: Any, value: Any, *, operation: str) -> list[Any]: +def _checkpoint_allgather( + comm: MnnvlCheckpointCommunicator, + value: Any, + *, + operation: str, +) -> list[Any]: """Run a bounded object allgather over nonblocking point-to-point requests. Production mpi4py communicators provide ``isend`` and ``irecv`` but no @@ -133,6 +166,7 @@ def _checkpoint_allgather(comm: Any, value: Any, *, operation: str) -> list[Any] if time.monotonic() >= deadline: raise TimeoutError(f"Timed out waiting for MNNVL checkpoint {operation} allgather") time.sleep(_MNNVL_CHECKPOINT_COLLECTIVE_POLL_INTERVAL_S) + return results except Exception: _cancel_checkpoint_requests([*receive_requests.values(), *send_requests]) raise @@ -664,7 +698,7 @@ def checkpoint_fail_closed(self) -> None: record.state = _MnnvlAllocationState.BROKEN record.pending_comm = None - def checkpoint_restore(self, comm) -> bool: + def checkpoint_restore(self, comm: MnnvlCheckpointCommunicator) -> bool: """Remap fresh handles while keeping data-path access disabled.""" cls = type(self) record = cls.allocated_map[self.ptr] @@ -979,7 +1013,7 @@ def checkpoint_prepare() -> None: workspace.checkpoint_prepare() @staticmethod - def checkpoint_restore(comm) -> None: + def checkpoint_restore(comm: MnnvlCheckpointCommunicator) -> None: """Restore TRT-native two-sided MoE workspaces at their original virtual addresses.""" workspaces = (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) restored_workspaces = [] diff --git a/tensorrt_llm/_torch/distributed/moe_alltoall.py b/tensorrt_llm/_torch/distributed/moe_alltoall.py index f3328c08c990..12e80d4591e3 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -28,7 +28,8 @@ import torch -from tensorrt_llm._mnnvl_utils import CftMnnvlMemory, MnnvlMemory +from tensorrt_llm._mnnvl_utils import (CftMnnvlMemory, + MnnvlCheckpointCommunicator, MnnvlMemory) from tensorrt_llm._torch.alltoall_watchdog import ( DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, ActiveRankMaskSnapshot, @@ -419,6 +420,10 @@ def _watchdog_coordinator(self) -> AlltoAllWatchdogCoordinator: def _alltoall_watchdog(self) -> AlltoAllWatchdog | None: return self._workspace_lifecycle.watchdog_for(self) + def checkpoint_resource_key(self) -> int: + """Identify wrappers sharing the same MNNVL workspace lifecycle.""" + return id(self._workspace_lifecycle) + def destroy(self) -> None: """Stop background watchdog resources owned by this wrapper.""" if getattr(self, "_destroyed", False): @@ -476,7 +481,10 @@ def checkpoint_prepare(self) -> None: ) self._workspace_lifecycle.checkpoint_prepare() - def checkpoint_restore(self, comm=None) -> None: + def checkpoint_restore( + self, + comm: MnnvlCheckpointCommunicator | None = None, + ) -> None: """Collectively restore handles and all shared frontend state. Args: diff --git a/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py b/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py index a5e1e025673a..07dde82ee935 100644 --- a/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py +++ b/tensorrt_llm/_torch/mnnvl_alltoall_workspace.py @@ -19,7 +19,11 @@ import torch -from tensorrt_llm._mnnvl_utils import MnnvlMemory, _checkpoint_allgather +from tensorrt_llm._mnnvl_utils import ( + MnnvlCheckpointCommunicator, + MnnvlMemory, + _checkpoint_allgather, +) from tensorrt_llm._torch.alltoall_watchdog import ( AlltoAllWatchdog, AlltoAllWatchdogCoordinator, @@ -30,17 +34,8 @@ _WORKSPACE_LIFECYCLE_KEY = "mnnvl_alltoall_workspace_lifecycle" -class _CollectiveCommunicator(Protocol): - def barrier(self) -> None: - """Synchronize every rank participating in the workspace.""" - - def allgather(self, value: bool) -> list[bool]: - """Gather one readiness value from each participating rank.""" - ... - - def _collect_active_ranks( - comm: _CollectiveCommunicator, + comm: MnnvlCheckpointCommunicator, *, local_clients_idle: bool, expected_size: int, @@ -60,7 +55,7 @@ def _collect_active_ranks( def _collect_unready_ranks( - comm: _CollectiveCommunicator, + comm: MnnvlCheckpointCommunicator, *, local_ready: bool, expected_size: int, @@ -235,7 +230,7 @@ def checkpoint_prepare(self) -> None: local_clients_idle = all( client._mnnvl_checkpoint_is_idle() for client in list(self._clients) ) - comm = cast(_CollectiveCommunicator | None, self._memory.comm) + comm = cast(MnnvlCheckpointCommunicator | None, self._memory.comm) if comm is None: raise RuntimeError("MNNVL workspace communicator is not initialized") try: @@ -257,7 +252,7 @@ def checkpoint_prepare(self) -> None: def checkpoint_restore( self, - comm: _CollectiveCommunicator, + comm: MnnvlCheckpointCommunicator, initialize_frontend: Callable[[], torch.Tensor], ) -> None: """Restore backing handles and publish the workspace after frontend readiness.""" diff --git a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md index b03b26a9bb49..6930810ee2da 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md +++ b/tensorrt_llm/_torch/modules/fused_moe/MOE_DEVELOPER_GUIDE.md @@ -164,6 +164,16 @@ Still on old path (standalone, with embedded communication): Communication strategies are auto-selected at runtime by `CommunicationFactory` based on hardware and configuration. Skipped for `FUSED_COMM` backends. See `communication_factory.py` for selection logic and `base.py` for the `Communication` ABC. +Communication strategies whose native workspace must participate in executor +sleep/wakeup implement the runtime-checkable `CheckpointableCommunication` +protocol from `communication/base.py`. The public +`checkpoint_resource_key()` identifies wrappers that share one underlying +workspace so the executor invokes `checkpoint_prepare()` and +`checkpoint_restore()` once per resource. Discovery must use this protocol; +do not add concrete strategy checks or inspect private workspace attributes. +The executor request queue owns the persistent admission state around the +checkpoint operation and reopens admission only after a complete wakeup. + ### MegaMoE (`fused_moe/mega_moe/`) | File | Role | @@ -428,7 +438,7 @@ When adding new components, use these reference implementations: | New `EXTERNAL_COMM` Backend | `fused_moe_cutlass.py` (`CutlassFusedMoE`) | `capabilities`, `can_implement`, `run_moe`, `create_weights`, `load_weights`; then add the class to `moe_resolution.BACKEND_CANDIDATES`. Add a fixed `descriptor.identity` only for a one-implementation leaf class | | New `FUSED_COMM` Backend | `mega_moe/mega_moe_deepgemm.py` (`MegaMoEDeepGemm`), `mega_moe/mega_moe_cute_dsl.py` (`MegaMoECuteDsl`) | Same as above + override `scheduler_kind = MoESchedulerKind.FUSED_COMM` and `validate_configurable_moe` for backend-specific constraints. For NVFP4 CuteDSL specifically, mirror the `MegaMoECuteDsl` pattern: capability probe for the CUDA 13 Cutlass DSL runtime, JSON-friendly tactic dict, lazy kernel import via `cute_dsl_kernels/mega_moe_nvfp4/import_kernel()`, and `quantize_input` that short-circuits zero-token input. | | New Quantization Method | `quantization.py` → `FP8QDQFusedMoEMethod` | Subclass `FusedMoEMethod`, implement quant/dequant ops | -| New Communication Strategy | `communication/nvlink_one_sided.py` (`NVLinkOneSided`) | Subclass `Communication`, implement `prepare_dispatch`, `dispatch`, `combine` | +| New Communication Strategy | `communication/nvlink_one_sided.py` (`NVLinkOneSided`) | Subclass `Communication`, implement `prepare_dispatch`, `dispatch`, `combine`; also implement `CheckpointableCommunication` when native workspace mappings must follow executor sleep/wakeup | | New Scheduler | `moe_scheduler.py` (`ExternalCommMoEScheduler` / `FusedCommMoEScheduler`) | Subclass `MoEScheduler`, implement `forward`; add new `MoESchedulerKind` value and wire into `create_moe_scheduler` factory | | Backend Tests | `test_moe_backend.py` | Follow existing parametrize patterns | | Integration Tests | `test_moe_module.py` | Test Backend × Communication × EPLB combinations | diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/__init__.py b/tensorrt_llm/_torch/modules/fused_moe/communication/__init__.py index 9858693c5f37..6e8780f35115 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/__init__.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/__init__.py @@ -30,7 +30,7 @@ """ from .allgather_reducescatter import AllGatherReduceScatter -from .base import Communication +from .base import CheckpointableCommunication, Communication from .communication_factory import CommunicationFactory from .deep_ep import DeepEP from .deep_ep_low_latency import DeepEPLowLatency @@ -40,6 +40,7 @@ __all__ = [ # Base classes and types + "CheckpointableCommunication", "Communication", # Communication strategies "AllGatherReduceScatter", diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/base.py b/tensorrt_llm/_torch/modules/fused_moe/communication/base.py index 0cfe8d5e5d7b..688e6849a564 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/base.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/base.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -24,13 +24,35 @@ """ from abc import ABC, abstractmethod -from typing import List, Optional, Tuple +from collections.abc import Hashable +from typing import List, Optional, Protocol, Tuple, runtime_checkable import torch from tensorrt_llm.mapping import Mapping +@runtime_checkable +class CheckpointableCommunication(Protocol): + """MoE communication resource participating in engine sleep/wakeup. + + Implementations own their local checkpoint details and expose a stable key + so the executor invokes each shared workspace exactly once. + """ + + def checkpoint_resource_key(self) -> Hashable: + """Return the identity shared by wrappers using one checkpoint resource.""" + ... + + def checkpoint_prepare(self) -> None: + """Detach checkpoint-backed resources after global quiescence.""" + ... + + def checkpoint_restore(self) -> None: + """Restore checkpoint-backed resources before admission reopens.""" + ... + + class Communication(ABC): """ Abstract base class for MoE communication methods 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 65055d76a205..be6e75c40dc2 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,7 @@ import torch -from tensorrt_llm._mnnvl_utils import CftMnnvlMemory, MnnvlMemory +from tensorrt_llm._mnnvl_utils import CftMnnvlMemory, MnnvlCheckpointCommunicator, MnnvlMemory from tensorrt_llm._torch.alltoall_watchdog import ( DEFAULT_ALLTOALL_WATCHDOG_POLL_INTERVAL_S, DEFAULT_ALLTOALL_WATCHDOG_TIMEOUT_S, @@ -634,6 +634,10 @@ def _require_mapped(self) -> None: if not self.mnnvl_mem.mapped: raise RuntimeError("Native MoE All-to-All workspace handles are unmapped") + def checkpoint_resource_key(self) -> int: + """Identify wrappers sharing the same MNNVL workspace lifecycle.""" + return id(self._require_workspace_lifecycle()) + def checkpoint_prepare(self) -> None: """Collectively detach handles after every shared owner is idle.""" if self.can_use_cft_counted_writes: @@ -642,7 +646,10 @@ def checkpoint_prepare(self) -> None: ) self._require_workspace_lifecycle().checkpoint_prepare() - def checkpoint_restore(self, comm=None) -> None: + def checkpoint_restore( + self, + comm: MnnvlCheckpointCommunicator | None = None, + ) -> None: """Collectively restore handles and all shared frontend state. Args: diff --git a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py index d249ab0b2ce9..2582a965490f 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py +++ b/tensorrt_llm/_torch/modules/fused_moe/communication/nvlink_two_sided.py @@ -27,7 +27,7 @@ import torch -from tensorrt_llm._mnnvl_utils import MnnvlMemory, MnnvlMoe +from tensorrt_llm._mnnvl_utils import MnnvlCheckpointCommunicator, MnnvlMemory, MnnvlMoe from tensorrt_llm._torch.mnnvl_alltoall_workspace import _collect_active_ranks from tensorrt_llm.mapping import Mapping @@ -109,6 +109,10 @@ def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) """ return True + def checkpoint_resource_key(self) -> int: + """Identify the process-global TRT-native two-sided workspaces.""" + return id(MnnvlMoe) + def checkpoint_prepare(self) -> None: """Detach TRT-native two-sided workspaces after global quiescence.""" workspaces = (MnnvlMoe.moe_workspace, MnnvlMoe.moe_prepare_workspace) @@ -138,7 +142,10 @@ def checkpoint_prepare(self) -> None: ) MnnvlMoe.checkpoint_prepare() - def checkpoint_restore(self, comm=None) -> None: + def checkpoint_restore( + self, + comm: MnnvlCheckpointCommunicator | None = None, + ) -> None: """Restore TRT-native two-sided workspaces and protocol state. Args: diff --git a/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py b/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py index 8dcb325c2775..7d1b860c5cd4 100644 --- a/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py +++ b/tensorrt_llm/_torch/pyexecutor/executor_request_queue.py @@ -1,5 +1,6 @@ import dataclasses import datetime +import enum import queue import threading import time @@ -15,6 +16,16 @@ CONTROL_REQUEST_ID = -2 +class RequestAdmissionState(enum.Enum): + """Persistent request-admission state across engine sleep and wakeup.""" + + RUNNING = "running" + PARKING = "parking" + PARKED = "parked" + WAKING = "waking" + FAILED = "failed" + + @dataclasses.dataclass class RequestQueueItem: id: int @@ -61,8 +72,97 @@ def __init__( self.enable_iter_perf_stats = enable_iter_perf_stats self.start_times = {} self.active = True + self.admission_state = RequestAdmissionState.RUNNING + self._pending_sleep_tags: frozenset[str] = frozenset() + self._parked_tags: frozenset[str] = frozenset() + self._pending_wakeup_tags: frozenset[str] = frozenset() self.batch_wait_timeout_ms = batch_wait_timeout_ms + def _transition_admission_state( + self, + expected: RequestAdmissionState, + target: RequestAdmissionState, + ) -> None: + if self.admission_state is not expected: + raise RuntimeError( + "Invalid executor admission transition: " + f"{self.admission_state.value} -> {target.value}; " + f"expected {expected.value}") + self.admission_state = target + + def begin_sleep_transition(self, tags: Iterable[str]) -> None: + """Close admission before the executor starts draining for sleep.""" + with self.enqueue_lock: + sleep_tags = frozenset(tags) + self._transition_admission_state( + RequestAdmissionState.RUNNING, + RequestAdmissionState.PARKING, + ) + self._pending_sleep_tags = sleep_tags + + def complete_sleep_transition(self) -> None: + """Publish PARKED only after every rank has completed sleep.""" + with self.enqueue_lock: + self._transition_admission_state( + RequestAdmissionState.PARKING, + RequestAdmissionState.PARKED, + ) + self._parked_tags = self._pending_sleep_tags + self._pending_sleep_tags = frozenset() + + def abort_sleep_transition(self) -> None: + """Reopen admission after a recoverable pre-mutation sleep failure.""" + with self.enqueue_lock: + if self.admission_state is RequestAdmissionState.FAILED: + return + self._transition_admission_state( + RequestAdmissionState.PARKING, + RequestAdmissionState.RUNNING, + ) + self._pending_sleep_tags = frozenset() + + def begin_wakeup_transition(self, tags: Iterable[str]) -> None: + """Keep admission closed while parked resources are restored.""" + with self.enqueue_lock: + wakeup_tags = frozenset(tags) + self._transition_admission_state( + RequestAdmissionState.PARKED, + RequestAdmissionState.WAKING, + ) + self._pending_wakeup_tags = wakeup_tags + + def complete_wakeup_transition(self) -> None: + """Reopen admission only after all originally parked tags are restored.""" + with self.enqueue_lock: + remaining_tags = self._parked_tags - self._pending_wakeup_tags + target = RequestAdmissionState.PARKED + if not remaining_tags: + target = RequestAdmissionState.RUNNING + self._transition_admission_state(RequestAdmissionState.WAKING, + target) + self._parked_tags = remaining_tags + self._pending_wakeup_tags = frozenset() + + def abort_wakeup_transition(self) -> None: + """Return to PARKED after a recoverable pre-mutation wakeup failure.""" + with self.enqueue_lock: + if self.admission_state is RequestAdmissionState.FAILED: + return + self._transition_admission_state( + RequestAdmissionState.WAKING, + RequestAdmissionState.PARKED, + ) + self._pending_wakeup_tags = frozenset() + + def fail_sleep_wakeup_transition(self) -> None: + """Permanently close admission after sleep/wakeup may have mutated state.""" + with self.enqueue_lock: + self.admission_state = RequestAdmissionState.FAILED + + def get_admission_state(self) -> RequestAdmissionState: + with self.enqueue_lock: + return self.admission_state + def _get_request_id(self, request: Optional[ExecutorRequest] = None): # if request has a disagg_request_id, use it as request id so that # corresponding context and generation requests have the same request id @@ -93,6 +193,10 @@ def _enqueue_impl(self, requests: Iterable[ExecutorRequest]) -> List[int]: req_ids = [] with self.enqueue_lock: assert self.active, "PyExecutor has already been shutdown." + if self.admission_state is not RequestAdmissionState.RUNNING: + raise RuntimeError( + "Cannot enqueue requests while executor admission is " + f"{self.admission_state.value}") start_time = time.time() for request in requests: req_id = self._get_request_id(request) @@ -140,6 +244,12 @@ def enqueue_shutdown_request(self): self.active = False def can_enqueue_request(self) -> bool: + with self.enqueue_lock: + return (self.active and self.dist.rank == 0 + and self.admission_state is RequestAdmissionState.RUNNING) + + def can_enqueue_control_request(self) -> bool: + """Return whether rank zero can enqueue shutdown/control sentinels.""" with self.enqueue_lock: return self.active and self.dist.rank == 0 diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index b46d350da389..c50faf9b0776 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -42,7 +42,8 @@ from tensorrt_llm.executor.request import TruncateKVCacheRequest from tensorrt_llm.inputs.multimodal import strip_mm_data_for_generation from tensorrt_llm.inputs.registry import get_multimodal_encoder_item_metadata -from tensorrt_llm.llmapi.llm_args import PeftCacheConfig, WaitingQueuePolicy +from tensorrt_llm.llmapi.llm_args import (ExecutorMemoryType, PeftCacheConfig, + WaitingQueuePolicy) from tensorrt_llm.logger import logger from tensorrt_llm.mapping import CpType from tensorrt_llm.runtime.kv_cache_manager_v2 import OutOfPagesError @@ -68,7 +69,8 @@ from .connectors.kv_cache_connector import KvCacheConnectorManager from .dwdp import DwdpManager from .error_classification import ErrorBudget -from .executor_request_queue import ExecutorRequestQueue, RequestQueueItem +from .executor_request_queue import (ExecutorRequestQueue, + RequestAdmissionState, RequestQueueItem) from .guided_decoder import GuidedDecoder from .handle_additional_outputs import HandleAdditionalOutputs from .handle_logits import HandleLogits @@ -104,6 +106,9 @@ if TYPE_CHECKING: from ray.actor import ActorHandle + from ..modules.fused_moe.communication.base import \ + CheckpointableCommunication + _UNBOUNDED_STATS_MAX_LEN = -1 @@ -1534,6 +1539,35 @@ def can_enqueue_requests(self) -> bool: """ return self.executor_request_queue.can_enqueue_request() + def begin_sleep_transition(self, tags: list[ExecutorMemoryType]) -> None: + self.executor_request_queue.begin_sleep_transition(tag.value + for tag in tags) + + def complete_sleep_transition(self) -> None: + self.executor_request_queue.complete_sleep_transition() + + def abort_sleep_transition(self) -> None: + self.executor_request_queue.abort_sleep_transition() + + def begin_wakeup_transition(self, tags: list[ExecutorMemoryType]) -> None: + self.executor_request_queue.begin_wakeup_transition(tag.value + for tag in tags) + + def complete_wakeup_transition(self) -> None: + self.executor_request_queue.complete_wakeup_transition() + + def abort_wakeup_transition(self) -> None: + self.executor_request_queue.abort_wakeup_transition() + + def fail_sleep_wakeup_transition(self) -> None: + self.executor_request_queue.fail_sleep_wakeup_transition() + + def get_request_admission_state(self) -> RequestAdmissionState: + return self.executor_request_queue.get_admission_state() + + def can_shutdown(self) -> bool: + return self.executor_request_queue.can_enqueue_control_request() + def get_latest_iteration_stats(self): """ Returns the per-iterations statistics computed since last call to this method. @@ -3161,14 +3195,13 @@ def _sleep_wakeup_listener_loop(self) -> None: finally: set_thread_local_mpi_comm(None) - def _mnnvl_checkpoint_resources(self, tags) -> list: + def _mnnvl_checkpoint_resources( + self, + tags: list[ExecutorMemoryType], + ) -> list["CheckpointableCommunication"]: """Return unique native MNNVL MoE resources selected by engine tags.""" - from tensorrt_llm._torch.distributed.moe_alltoall import MoeAlltoAll - from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_one_sided import \ - NVLinkOneSided - from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_two_sided import \ - NVLinkTwoSided - from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType + from tensorrt_llm._torch.modules.fused_moe.communication.base import \ + CheckpointableCommunication selected_engines = [] if ExecutorMemoryType.MODEL_ENGINE_MAIN in tags: @@ -3176,8 +3209,7 @@ def _mnnvl_checkpoint_resources(self, tags) -> list: if ExecutorMemoryType.MODEL_ENGINE_DRAFT in tags and self.draft_model_engine is not None: selected_engines.append(self.draft_model_engine) - resource_types = (MoeAlltoAll, NVLinkOneSided, NVLinkTwoSided) - resources = [] + resources: list[CheckpointableCommunication] = [] seen = set() for engine in selected_engines: model = getattr(engine, "model", None) @@ -3185,22 +3217,26 @@ def _mnnvl_checkpoint_resources(self, tags) -> list: continue for module in model.modules(): resource = getattr(module, "comm", None) - if not isinstance(resource, resource_types): + if not isinstance(resource, CheckpointableCommunication): continue - lifecycle = getattr(resource, "_workspace_lifecycle", None) - key = (("workspace", - id(lifecycle)) if lifecycle is not None else - ("resource_type", type(resource))) + key = resource.checkpoint_resource_key() if key in seen: continue seen.add(key) resources.append(resource) return resources - def _has_mnnvl_checkpoint_resources(self, tags) -> bool: + def _has_mnnvl_checkpoint_resources( + self, + tags: list[ExecutorMemoryType], + ) -> bool: return bool(self._mnnvl_checkpoint_resources(tags)) - def _run_mnnvl_checkpoint_resources(self, action, tags) -> None: + def _run_mnnvl_checkpoint_resources( + self, + action: _SleepWakeupAction, + tags: list[ExecutorMemoryType], + ) -> None: """Execute process-local native MNNVL hooks while the engine is parked.""" resources = self._mnnvl_checkpoint_resources(tags) if action == _SleepWakeupAction.SLEEP: diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index c9d2a27ded40..bc63be7d2e56 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -22,7 +22,7 @@ import weakref from pathlib import Path from queue import Empty, Queue -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union import torch @@ -633,7 +633,7 @@ def _check_sleep_wakeup_preconditions(self, method: str) -> None: def _multi_rank_sleep_wakeup( self, - action: str, + action: Literal["sleep", "wakeup"], tags: list[ExecutorMemoryType], ) -> None: """Coordinate a sleep or wakeup operation across all MPI ranks. @@ -655,8 +655,9 @@ def _multi_rank_sleep_wakeup( send ABORT before any rank changes VMM state. Once COMMIT or local mutation starts, any error fail-stops the distributed worker because rank state can no longer be reconciled safely. - 5. Exit ``control_action()`` and resume rank-0 only after a successful - operation or a recoverable pre-COMMIT abort. + 5. Exit ``control_action()`` only after a successful operation or a + recoverable pre-COMMIT abort. The public ``sleep()``/``wakeup()`` + wrapper owns the persistent admission transition around this helper. Args: action: ``"sleep"`` or ``"wakeup"``. @@ -903,6 +904,13 @@ def _fail_stop_divergent_sleep_wakeup(self, error: RuntimeError) -> None: self._set_fatal_error(error) if self.engine is not None: + fail_transition = getattr( + self.engine, + "fail_sleep_wakeup_transition", + None, + ) + if fail_transition is not None: + fail_transition() if getattr(self.engine, "_fatal_error", None) is None: self.engine._fatal_error = error self.engine.is_shutdown = True @@ -937,9 +945,9 @@ def sleep(self, sleep_tags: list[str]) -> None: value strings (e.g. ``["kv_cache"]``). Returns: - None. The call is synchronous; when it returns all requested - VMM-tagged allocations have been released on every rank and the - event loop has been resumed. + None. The call is synchronous; when it returns all requested + VMM-tagged allocations have been released on every rank and request + admission remains closed until ``wakeup()`` succeeds. Raises: ValueError: If the backend is not ``"pytorch"`` or @@ -951,15 +959,30 @@ def sleep(self, sleep_tags: list[str]) -> None: tags = [ExecutorMemoryType(tag) for tag in sleep_tags] logger.info(f"Sleep: {tags}") - if self.llm_args.parallel_config.world_size > 1: - self._multi_rank_sleep_wakeup("sleep", tags) - else: - with self.engine._sleep_wakeup_lock, self.engine.control_action(): - torch.cuda.synchronize() - release_with_tag(*tags) - torch.cuda.synchronize() - gc.collect() - torch.cuda.empty_cache() + self.engine.begin_sleep_transition(tags) + local_mutation_started = False + try: + if self.llm_args.parallel_config.world_size > 1: + self._multi_rank_sleep_wakeup("sleep", tags) + else: + with self.engine._sleep_wakeup_lock, self.engine.control_action( + ): + torch.cuda.synchronize() + local_mutation_started = True + release_with_tag(*tags) + torch.cuda.synchronize() + gc.collect() + torch.cuda.empty_cache() + except Exception: + if local_mutation_started: + self.engine.fail_sleep_wakeup_transition() + self.engine.abort_sleep_transition() + raise + try: + self.engine.complete_sleep_transition() + except Exception: + self.engine.fail_sleep_wakeup_transition() + raise def wakeup(self, wakeup_tags: list[str]) -> None: """Materialize GPU virtual memory for the specified memory type tags. @@ -975,9 +998,10 @@ def wakeup(self, wakeup_tags: list[str]) -> None: value strings (e.g. ``["kv_cache"]``). Returns: - None. The call is synchronous; when it returns all requested - VMM-tagged allocations have been materialized on every rank and the - event loop has been resumed. + None. The call is synchronous; when it returns all requested + VMM-tagged allocations have been materialized on every rank. + Request admission reopens once every tag from the corresponding + sleep transition has been restored. Raises: ValueError: If the backend is not ``"pytorch"`` or @@ -989,13 +1013,28 @@ def wakeup(self, wakeup_tags: list[str]) -> None: tags = [ExecutorMemoryType(tag) for tag in wakeup_tags] logger.info(f"Wakeup: {tags}") - if self.llm_args.parallel_config.world_size > 1: - self._multi_rank_sleep_wakeup("wakeup", tags) - else: - with self.engine._sleep_wakeup_lock, self.engine.control_action(): - torch.cuda.synchronize() - materialize_with_tag(*tags) - torch.cuda.synchronize() + self.engine.begin_wakeup_transition(tags) + local_mutation_started = False + try: + if self.llm_args.parallel_config.world_size > 1: + self._multi_rank_sleep_wakeup("wakeup", tags) + else: + with self.engine._sleep_wakeup_lock, self.engine.control_action( + ): + torch.cuda.synchronize() + local_mutation_started = True + materialize_with_tag(*tags) + torch.cuda.synchronize() + except Exception: + if local_mutation_started: + self.engine.fail_sleep_wakeup_transition() + self.engine.abort_wakeup_transition() + raise + try: + self.engine.complete_wakeup_transition() + except Exception: + self.engine.fail_sleep_wakeup_transition() + raise def shutdown(self): if self.doing_shutdown: @@ -1003,9 +1042,13 @@ def shutdown(self): else: self.doing_shutdown = True - if self.engine is not None and self.engine.can_enqueue_requests(): - self.engine.shutdown() - self.engine = None + if self.engine is not None: + can_shutdown = getattr(self.engine, "can_shutdown", None) + if can_shutdown is None: + can_shutdown = self.engine.can_enqueue_requests + if can_shutdown(): + self.engine.shutdown() + self.engine = None def get_disaggregated_params(self) -> dict: if self.engine is None or self.engine.kv_cache_transceiver is None: diff --git a/tests/unittest/_torch/executor/test_executor_request_queue.py b/tests/unittest/_torch/executor/test_executor_request_queue.py index 1478cb4ab990..d4b46e4502d8 100644 --- a/tests/unittest/_torch/executor/test_executor_request_queue.py +++ b/tests/unittest/_torch/executor/test_executor_request_queue.py @@ -16,7 +16,8 @@ import pytest from tensorrt_llm._torch.pyexecutor.executor_request_queue import ( - SHUTDOWN_REQUEST_ID, ExecutorRequestQueue, RequestQueueItem) + SHUTDOWN_REQUEST_ID, ExecutorRequestQueue, RequestAdmissionState, + RequestQueueItem) pytestmark = pytest.mark.cpu_only @@ -65,6 +66,7 @@ def test_executor_queue_init(executor_queue, mock_dist): assert executor_queue.next_request_id == 8 assert executor_queue.enable_iter_perf_stats assert executor_queue.active + assert executor_queue.get_admission_state() is RequestAdmissionState.RUNNING assert isinstance(executor_queue.request_queue, queue.Queue) assert isinstance(executor_queue.enqueue_lock, type(threading.Lock())) @@ -155,6 +157,69 @@ def test_enqueue_request_after_shutdown(executor_queue): executor_queue.enqueue_request(Mock()) +def test_sleep_wakeup_admission_state_rejects_direct_enqueue(executor_queue): + """Admission stays closed from PARKING until the matching wakeup completes.""" + executor_queue.begin_sleep_transition(["model", "kv_cache"]) + assert executor_queue.get_admission_state() is RequestAdmissionState.PARKING + assert not executor_queue.can_enqueue_request() + with pytest.raises(RuntimeError, match="admission is parking"): + executor_queue.enqueue_request(Mock()) + + executor_queue.complete_sleep_transition() + assert executor_queue.get_admission_state() is RequestAdmissionState.PARKED + assert executor_queue.can_enqueue_control_request() + with pytest.raises(RuntimeError, match="admission is parked"): + executor_queue.enqueue_request(Mock()) + + executor_queue.begin_wakeup_transition(["kv_cache"]) + assert executor_queue.get_admission_state() is RequestAdmissionState.WAKING + with pytest.raises(RuntimeError, match="admission is waking"): + executor_queue.enqueue_request(Mock()) + + executor_queue.complete_wakeup_transition() + assert executor_queue.get_admission_state() is RequestAdmissionState.PARKED + with pytest.raises(RuntimeError, match="admission is parked"): + executor_queue.enqueue_request(Mock()) + + executor_queue.begin_wakeup_transition(["model"]) + executor_queue.complete_wakeup_transition() + assert executor_queue.get_admission_state() is RequestAdmissionState.RUNNING + assert executor_queue.can_enqueue_request() + + +def test_wakeup_allows_extra_tags_and_reopens_after_parked_tags(executor_queue): + executor_queue.begin_sleep_transition(["model"]) + executor_queue.complete_sleep_transition() + + executor_queue.begin_wakeup_transition(["model", "kv_cache"]) + executor_queue.complete_wakeup_transition() + + assert executor_queue.get_admission_state() is RequestAdmissionState.RUNNING + + +def test_recoverable_sleep_wakeup_failures_restore_prior_state(executor_queue): + executor_queue.begin_sleep_transition(["model"]) + executor_queue.abort_sleep_transition() + assert executor_queue.get_admission_state() is RequestAdmissionState.RUNNING + + executor_queue.begin_sleep_transition(["model"]) + executor_queue.complete_sleep_transition() + executor_queue.begin_wakeup_transition(["model"]) + executor_queue.abort_wakeup_transition() + assert executor_queue.get_admission_state() is RequestAdmissionState.PARKED + + +def test_post_mutation_failure_permanently_closes_admission(executor_queue): + executor_queue.begin_sleep_transition(["model"]) + executor_queue.fail_sleep_wakeup_transition() + executor_queue.abort_sleep_transition() + + assert executor_queue.get_admission_state() is RequestAdmissionState.FAILED + assert not executor_queue.can_enqueue_request() + with pytest.raises(RuntimeError, match="admission is failed"): + executor_queue.enqueue_request(Mock()) + + @pytest.mark.parametrize( "rank,active,expected", [ diff --git a/tests/unittest/executor/test_sleep_collective_rpc_guards.py b/tests/unittest/executor/test_sleep_collective_rpc_guards.py index f39d96001e31..9231d0000247 100644 --- a/tests/unittest/executor/test_sleep_collective_rpc_guards.py +++ b/tests/unittest/executor/test_sleep_collective_rpc_guards.py @@ -50,6 +50,15 @@ def _make_worker(backend="pytorch", world_size=1, sleep_config=_SLEEP_CONFIG_DEF parallel_config=SimpleNamespace(world_size=world_size), sleep_config=sleep_config, ) + w.engine = SimpleNamespace( + begin_sleep_transition=MagicMock(), + complete_sleep_transition=MagicMock(), + abort_sleep_transition=MagicMock(), + begin_wakeup_transition=MagicMock(), + complete_wakeup_transition=MagicMock(), + abort_wakeup_transition=MagicMock(), + fail_sleep_wakeup_transition=MagicMock(), + ) return w @@ -111,6 +120,28 @@ def test_multirank_dispatches_to_helper(self, method): call_action, call_tags = mock_helper.call_args[0] assert call_action == method assert len(call_tags) == 1 + getattr(w.engine, f"begin_{method}_transition").assert_called_once_with(call_tags) + getattr(w.engine, f"complete_{method}_transition").assert_called_once_with() + getattr(w.engine, f"abort_{method}_transition").assert_not_called() + + def test_multirank_recoverable_failure_restores_admission(self, method): + """A helper failure before mutation restores the prior admission state.""" + from unittest.mock import patch + + w = _make_worker(world_size=2) + with ( + patch.object( + w, + "_multi_rank_sleep_wakeup", + side_effect=RuntimeError("prepare failed"), + ), + pytest.raises(RuntimeError, match="prepare failed"), + ): + getattr(w, method)(["kv_cache"]) + + getattr(w.engine, f"abort_{method}_transition").assert_called_once_with() + getattr(w.engine, f"complete_{method}_transition").assert_not_called() + w.engine.fail_sleep_wakeup_transition.assert_not_called() def test_backend_checked_before_sleep_config(self, method): """Backend check fires even when sleep_config is also absent.""" @@ -433,21 +464,27 @@ def test_pyexecutor_discovers_shared_native_resource_once(self): from types import SimpleNamespace from unittest.mock import Mock - from tensorrt_llm._torch.modules.fused_moe.communication.nvlink_one_sided import ( - NVLinkOneSided, + from tensorrt_llm._torch.modules.fused_moe.communication.base import ( + CheckpointableCommunication, ) from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor, _SleepWakeupAction from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType - lifecycle = object() + class FutureCheckpointCommunication: + def __init__(self, resource_key): + self.resource_key = resource_key + self.checkpoint_prepare = Mock() + self.checkpoint_restore = Mock() + + def checkpoint_resource_key(self): + return self.resource_key + + shared_resource_key = object() resources = [] modules = [] for _ in range(2): - resource = NVLinkOneSided.__new__(NVLinkOneSided) - resource._workspace_lifecycle = lifecycle - resource._workspace_registered = False - resource.checkpoint_prepare = Mock() - resource.checkpoint_restore = Mock() + resource = FutureCheckpointCommunication(shared_resource_key) + assert isinstance(resource, CheckpointableCommunication) resources.append(resource) modules.append(SimpleNamespace(comm=resource)) @@ -1021,7 +1058,11 @@ def test_divergent_commit_fail_stops_worker(self): worker = object.__new__(BaseWorker) worker._fatal_error = None - worker.engine = SimpleNamespace(_fatal_error=None, is_shutdown=False) + worker.engine = SimpleNamespace( + _fatal_error=None, + is_shutdown=False, + fail_sleep_wakeup_transition=MagicMock(), + ) error = RuntimeError("divergent sleep state") with patch("tensorrt_llm._torch.pyexecutor.hang_detector.propagate_hard_kill") as hard_kill: @@ -1030,6 +1071,7 @@ def test_divergent_commit_fail_stops_worker(self): assert worker._fatal_error is error assert worker.engine._fatal_error is error assert worker.engine.is_shutdown + worker.engine.fail_sleep_wakeup_transition.assert_called_once_with() hard_kill.assert_called_once_with() def test_precommit_local_failure_remains_recoverable(self): @@ -1513,6 +1555,13 @@ def _noop_control_action(): w.engine = SimpleNamespace( _sleep_wakeup_lock=SpyLock(), control_action=_noop_control_action, + begin_sleep_transition=MagicMock(), + complete_sleep_transition=MagicMock(), + abort_sleep_transition=MagicMock(), + begin_wakeup_transition=MagicMock(), + complete_wakeup_transition=MagicMock(), + abort_wakeup_transition=MagicMock(), + fail_sleep_wakeup_transition=MagicMock(), ) with ( @@ -1526,6 +1575,41 @@ def _noop_control_action(): assert lock_entered, f"{method}() with world_size=1 did not acquire _sleep_wakeup_lock" + @pytest.mark.parametrize( + ("method", "mutation"), + [("sleep", "release_with_tag"), ("wakeup", "materialize_with_tag")], + ) + def test_post_mutation_failure_permanently_closes_admission(self, method, mutation): + """A single-rank error after VMM mutation starts is fail-closed.""" + import threading + from contextlib import contextmanager + from unittest.mock import patch + + w = _make_worker() + + @contextmanager + def _noop_control_action(): + yield None + + w.engine._sleep_wakeup_lock = threading.Lock() + w.engine.control_action = _noop_control_action + + with ( + patch(f"tensorrt_llm._torch.virtual_memory.{mutation}"), + patch( + "torch.cuda.synchronize", + side_effect=[None, RuntimeError("post-mutation sync failed")], + ), + patch("gc.collect"), + patch("torch.cuda.empty_cache"), + pytest.raises(RuntimeError, match="post-mutation sync failed"), + ): + getattr(w, method)(["kv_cache"]) + + w.engine.fail_sleep_wakeup_transition.assert_called_once_with() + getattr(w.engine, f"abort_{method}_transition").assert_called_once_with() + getattr(w.engine, f"complete_{method}_transition").assert_not_called() + # --------------------------------------------------------------------------- # GenerationExecutorProxy / GenerationExecutorRpcProxy collective_rpc() From 711bad55adcbc3c4011e9cfc1eac5f9c46cef18a Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Wed, 19 Aug 2026 10:54:07 -0400 Subject: [PATCH 13/14] test(runtime): initialize sleep listener control state Signed-off-by: Hannah Zhang --- .../test_sleep_collective_rpc_guards.py | 24 +++++++------------ 1 file changed, 8 insertions(+), 16 deletions(-) diff --git a/tests/unittest/executor/test_sleep_collective_rpc_guards.py b/tests/unittest/executor/test_sleep_collective_rpc_guards.py index 9231d0000247..c93849de08ab 100644 --- a/tests/unittest/executor/test_sleep_collective_rpc_guards.py +++ b/tests/unittest/executor/test_sleep_collective_rpc_guards.py @@ -1115,6 +1115,7 @@ def test_uncaught_exception_sends_error_ack(self): because error_msg is still None when the exception bypasses the except clause — leaving rank-0 with an inconsistent view of the operation. """ + import threading from types import SimpleNamespace from unittest.mock import patch @@ -1134,6 +1135,10 @@ def send(self, payload, dest, tag): executor._sleep_wakeup_comm = FakeComm() executor.device_id = 0 executor.dist = SimpleNamespace(rank=1) + executor.control_request_barrier = threading.Event() + executor.control_request_barrier.set() + executor.control_action_done = threading.Event() + executor._active_control_id = None with ( patch("torch.cuda.set_device"), @@ -1146,23 +1151,8 @@ def send(self, payload, dest, tag): ), patch("torch.cuda.synchronize"), ): - # The loop receives one message then raises MemoryError which is - # not in the narrow except. We patch recv to raise StopIteration - # on the second call so the loop terminates cleanly in the test. - call_count = [0] - - def _recv_once(self_inner, source, tag): - call_count[0] += 1 - if call_count[0] > 1: - raise StopIteration - return {"action": "sleep", "tags": ["kv_cache"]} - - executor._sleep_wakeup_comm.recv = lambda source, tag: _recv_once(None, source, tag) - - try: + with pytest.raises(MemoryError, match="simulated OOM outside except list"): executor._sleep_wakeup_listener_loop() - except StopIteration: - pass assert sent_acks, "finally block must send an ACK even for uncaught exceptions" assert sent_acks[0]["status"] == "error", ( @@ -1171,6 +1161,8 @@ def _recv_once(self_inner, source, tag): ) assert sent_acks[0]["error"] is not None assert "MemoryError" in sent_acks[0]["error"] + assert not executor.control_request_barrier.is_set() + assert executor.control_action_done.is_set() class TestListenerAbortAndShutdown: From ec78fe2b0530d984510c566dcc8135db8e9dba05 Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Tue, 25 Aug 2026 10:30:57 -0400 Subject: [PATCH 14/14] test(moe): update MNNVL lifecycle fixtures Signed-off-by: Hannah Zhang --- tests/unittest/_torch/test_mnnvl_alltoall_workspace.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py b/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py index aa61439cc41c..fbd210fa9b87 100644 --- a/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py +++ b/tests/unittest/_torch/test_mnnvl_alltoall_workspace.py @@ -664,6 +664,7 @@ def test_frontend_checkpoint_delegates_to_shared_lifecycle( wrapper_type: type[MoeAlltoAll] | type[NVLinkOneSided], ) -> None: wrapper = wrapper_type.__new__(wrapper_type) + wrapper.can_use_cft_counted_writes = False wrapper._workspace_lifecycle = Mock() comm = Mock() @@ -698,7 +699,7 @@ def test_moe_alltoall_aborted_registration_does_not_unregister( ) -> None: lifecycle = Mock() lifecycle.register.side_effect = RuntimeError("registration failed") - monkeypatch.setattr(MoeAlltoAll, "_WORKSPACE", None) + monkeypatch.setattr(MoeAlltoAll, "_WORKSPACES", {}) monkeypatch.setattr(MoeAlltoAll, "_init_constants", Mock()) monkeypatch.setattr( MoeAlltoAll, @@ -748,6 +749,7 @@ def test_one_sided_checkpoint_rejects_destroyed_workspace( monkeypatch.setattr(NVLinkOneSided, "_WORKSPACE", None) wrapper = NVLinkOneSided.__new__(NVLinkOneSided) wrapper._destroyed = False + wrapper.can_use_cft_counted_writes = False wrapper._workspace_lifecycle = Mock() wrapper._workspace_key = ("test",) wrapper._workspace_registered = True