From 8fc2cae50e0edf0ee5a1ec1b3b3762a3ade34493 Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Mon, 20 Jul 2026 11:40:20 -0400 Subject: [PATCH 1/3] [None][feat] preserve native MoE A2A virtual addrs across restore Signed-off-by: Hannah Zhang --- tensorrt_llm/_mnnvl_utils.py | 343 +++++++++++------- .../_torch/distributed/moe_alltoall.py | 72 ++++ .../communication/nvlink_one_sided.py | 57 +++ .../communication/nvlink_two_sided.py | 17 +- 4 files changed, 364 insertions(+), 125 deletions(-) diff --git a/tensorrt_llm/_mnnvl_utils.py b/tensorrt_llm/_mnnvl_utils.py index 6a2a2ad49e2e..7ac6bbb6e62d 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.""" @@ -98,6 +111,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( @@ -182,147 +200,150 @@ def new_mnnvl_memory_address(cls, mapping: Mapping, size: int): cls.current_mem_offset = 0 @classmethod - def open_mnnvl_memory(cls, mapping: Mapping, size: int): - # Ensure MnnvlMemory is initialized (for dev_id and allocation_granularity) - MnnvlMemory.initialize() - - dev = _check_cu_result(cuda.cuCtxGetDevice()) - dev_id = int(dev) - if MnnvlMemory.dev_id is None: - MnnvlMemory.dev_id = dev_id + 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}" ) - comm = cls.get_comm(mapping) - comm_rank = comm.Get_rank() - comm_size = comm.Get_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." - granularity = MnnvlMemory.get_allocation_granularity(dev_id) - aligned_size = (size + granularity - 1) // granularity * granularity - - if cls.current_mem_offset + aligned_size > cls.current_rank_stride: - cls.new_mnnvl_memory_address(mapping, aligned_size) - - assert cls.current_mem_offset + aligned_size <= cls.current_rank_stride - allocation_prop = MnnvlMemory.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 - ) - ) + 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: - if ( - allocation_prop.requestedHandleTypes - == cuda.CUmemAllocationHandleType.CU_MEM_HANDLE_TYPE_FABRIC - ): - all_handles_data = comm.allgather(exported_fabric_handle.data) + 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_handles_data = comm.allgather(exported_fabric_handle) + all_exported_fds = comm.allgather(int(exported_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) + 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 i, (pidfd, fd) in enumerate(zip(pidfds, all_handles_data)): - remote_fd = syscall(SYS_pidfd_getfd, pidfd, fd, 0) + 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 {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) + 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 - 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, + + 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 + ) ) - 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: + 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: - os.close(_pidfd) - except OSError: - pass - for _rfd in remote_fds: + _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: - os.close(_rfd) - except OSError: - pass + _check_cu_result(cuda.cuMemRelease(handle)) + except RuntimeError as error: + logger.warning("Failed to release incomplete MNNVL allocation: %s", error) 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. + 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)) - madesc = cuda.CUmemAccessDesc() - madesc.location = allocation_prop.location - madesc.flags = cuda.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE + @classmethod + def open_mnnvl_memory(cls, mapping: Mapping, size: int): + # Ensure MnnvlMemory is initialized (for dev_id and allocation_granularity) + MnnvlMemory.initialize() - mem_handles = [None] * comm_size + dev = _check_cu_result(cuda.cuCtxGetDevice()) + dev_id = int(dev) + if MnnvlMemory.dev_id is None: + MnnvlMemory.dev_id = dev_id + assert dev_id == MnnvlMemory.dev_id, ( + f"Different dev_id found dev_id={dev_id} but MnnvlMemory.dev_id={MnnvlMemory.dev_id}" + ) + comm = cls.get_comm(mapping) + comm_rank = comm.Get_rank() + comm_size = comm.Get_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." + granularity = MnnvlMemory.get_allocation_granularity(dev_id) + aligned_size = (size + granularity - 1) // granularity * granularity - 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: - # Local memory mapping - mem_handles[i] = allocated_mem_handle - _check_cu_result(cuda.cuMemMap(rank_ptr, aligned_size, 0, allocated_mem_handle, 0)) - else: - # Fabric memory mapping - imported_mem_handle = _check_cu_result( - cuda.cuMemImportFromShareableHandle( - remote_handle_data, allocation_prop.requestedHandleTypes - ) - ) - mem_handles[i] = imported_mem_handle - _check_cu_result(cuda.cuMemMap(rank_ptr, aligned_size, 0, imported_mem_handle, 0)) + if cls.current_mem_offset + aligned_size > cls.current_rank_stride: + cls.new_mnnvl_memory_address(mapping, aligned_size) - _check_cu_result(cuda.cuMemSetAccess(rank_ptr, aligned_size, [madesc], 1)) + assert cls.current_mem_offset + aligned_size <= cls.current_rank_stride - ptr = cls.current_start_address + cls.current_mem_offset - stride = cls.current_rank_stride - cls.allocated_map[ptr] = ( - mapping, + mem_handles = cls._create_and_map_handles( + comm, aligned_size, - mem_handles, cls.current_start_address, cls.current_rank_stride, cls.current_mem_offset, ) + # 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( + 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 ) @@ -332,26 +353,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): @@ -489,6 +553,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 24f98307c70b..85d5bab26064 100644 --- a/tensorrt_llm/_torch/distributed/moe_alltoall.py +++ b/tensorrt_llm/_torch/distributed/moe_alltoall.py @@ -1,3 +1,17 @@ +# 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. """ MoE All-to-All Operations @@ -276,6 +290,61 @@ def __del__(self) -> None: if not sys.is_finalizing(): self.destroy() + 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], @@ -302,6 +371,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" @@ -385,6 +455,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" @@ -429,6 +500,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. See "payload_in_workspace" in combine method. """ + 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 55b7b2297a7d..f213aed784b9 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 @@ -409,6 +409,61 @@ def is_workload_feasible(self, all_rank_num_tokens: List[int], num_chunks: int) """ return True + 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, @@ -439,6 +494,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) @@ -569,6 +625,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 297d098b127fcabc64420789a59cd7a1736a0c58 Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Mon, 20 Jul 2026 11:40:51 -0400 Subject: [PATCH 2/3] [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 db678080a7c14b1d88e1537fda63f9e3417ce7f0 Mon Sep 17 00:00:00 2001 From: Hannah Zhang Date: Thu, 23 Jul 2026 14:16:24 -0400 Subject: [PATCH 3/3] 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 f213aed784b9..af98e26e12d0 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 @@ -716,6 +716,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"