From 372243aea5e217633eeb89f1b1ac987c1b872458 Mon Sep 17 00:00:00 2001 From: "Jiahan Chang (Cyrus)" <173873397+jiahanc@users.noreply.github.com> Date: Mon, 13 Jul 2026 14:04:19 +0800 Subject: [PATCH] =?UTF-8?q?Revert=20"feat(comm):=20preserve=20all-reduce?= =?UTF-8?q?=20graph=20VAs=20across=20checkpoint=20restore=20(=E2=80=A6"?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This reverts commit 68ebdbd6ef095e2a9ac93ae9ee222fde152cab6e. --- docs/api/comm.rst | 25 --- flashinfer/comm/allreduce.py | 59 ------- flashinfer/comm/mnnvl.py | 155 ++++++------------ flashinfer/comm/trtllm_ar.py | 84 ++-------- flashinfer/comm/trtllm_mnnvl_ar.py | 97 ++++------- .../comm/test_trtllm_allreduce_checkpoint.py | 120 -------------- tests/comm/test_trtllm_mnnvl_allreduce.py | 2 +- 7 files changed, 100 insertions(+), 442 deletions(-) delete mode 100644 tests/comm/test_trtllm_allreduce_checkpoint.py diff --git a/docs/api/comm.rst b/docs/api/comm.rst index ea6cff9c5af..fd398488046 100644 --- a/docs/api/comm.rst +++ b/docs/api/comm.rst @@ -92,31 +92,6 @@ Unified AllReduce Fusion API TRTLLMAllReduceFusionWorkspace MNNVLAllReduceFusionWorkspace -All-reduce workspaces backed by ``SymmDeviceMemory`` preserve their CUDA -virtual addresses across process checkpoint/restore. After quiescing all -work, release the physical handles and restore them with a fresh communication -backend before replaying a captured CUDA graph: - -.. code-block:: python - - workspace.checkpoint_prepare() - workspace.checkpoint_restore(comm_backend) - -Both methods are collective. Every rank must call them in the same order, and -``comm_backend`` must reproduce the original rank and world size. Repeated -calls are no-ops after the workspace reaches the requested state. If an -exception occurs after detach or reattach begins, do not retry or reuse the -workspace; restart the affected rank. Workspaces backed by torch symmetric -memory do not support this lifecycle. - -.. autosummary:: - :toctree: ../generated - - TRTLLMAllReduceFusionWorkspace.checkpoint_prepare - TRTLLMAllReduceFusionWorkspace.checkpoint_restore - MNNVLAllReduceFusionWorkspace.checkpoint_prepare - MNNVLAllReduceFusionWorkspace.checkpoint_restore - vLLM AllReduce -------------- diff --git a/flashinfer/comm/allreduce.py b/flashinfer/comm/allreduce.py index 87295b870e2..c9fe5400f11 100644 --- a/flashinfer/comm/allreduce.py +++ b/flashinfer/comm/allreduce.py @@ -63,7 +63,6 @@ from .trtllm_ar import trtllm_allreduce_fusion from .trtllm_ar import trtllm_create_ipc_workspace_for_all_reduce_fusion -from .trtllm_ar import _initialize_allreduce_fusion_protocol from .trtllm_ar import check_trtllm_allreduce_fusion_workspace_metadata from .trtllm_ar import trtllm_moe_allreduce_fusion from .trtllm_ar import trtllm_moe_finalize_allreduce_fusion @@ -179,59 +178,6 @@ def is_buffer_size_sufficient( logger.warning("Workspace is insufficient for problem size. %s", e) return False - @flashinfer_api - def checkpoint_prepare(self) -> None: - """Detach physical backing; repeated successful calls are no-ops.""" - if not self.mem_handles or not all( - isinstance(handle, SymmDeviceMemory) for handle in self.mem_handles - ): - raise NotImplementedError( - "Stable-VA checkpointing is unavailable for workspaces backed " - "by torch symmetric memory" - ) - - mapped = [handle.mapped for handle in self.mem_handles] - if not any(mapped): - return - if not all(mapped): - raise RuntimeError("TRT-LLM symmetric-memory handle state is inconsistent") - - for handle in self.mem_handles: - handle._unmap_and_release_handles() - # Do not return until every rank has released all workspace handles. - self.mem_handles[0].comm_backend.barrier() - - @flashinfer_api - def checkpoint_restore(self, comm_backend: CommBackend) -> None: - """Restore physical backing; repeated successful calls are no-ops.""" - if not self.mem_handles or not all( - isinstance(handle, SymmDeviceMemory) for handle in self.mem_handles - ): - raise NotImplementedError( - "Stable-VA checkpointing is unavailable for workspaces backed " - "by torch symmetric memory" - ) - - mapped = [handle.mapped for handle in self.mem_handles] - if all(mapped): - return - if any(mapped): - raise RuntimeError("TRT-LLM symmetric-memory handle state is inconsistent") - for handle in self.mem_handles: - handle._create_and_map_handles(comm_backend) - - _initialize_allreduce_fusion_protocol( - ipc_handles=self.ipc_handles, - tp_rank=self.rank, - flag_size=self.metadata["flag_size"], - lamport_buffer_size=self.metadata["lamport_buffer_size"], - lamport_comm_size=self.metadata["lamport_comm_size"], - use_fp32_lamport=self.metadata["use_fp32_lamport"], - control_flag_ptr=self.metadata["control_flag_ptr"], - ) - torch.cuda.synchronize() - comm_backend.barrier() - def destroy(self) -> None: """Destroy workspace and free resources.""" if getattr(self, "_destroyed", False): @@ -730,11 +676,6 @@ def allreduce_fusion( # Dispatch based on workspace type if isinstance(workspace, TRTLLMAllReduceFusionWorkspace): # TensorRT-LLM backend implementation - if any( - isinstance(handle, SymmDeviceMemory) and not handle.mapped - for handle in workspace.mem_handles - ): - raise RuntimeError("TRT-LLM symmetric-memory handles are not attached") # ---- MOE Reduction pattern ---- if pattern == AllReduceFusionPattern.kMoEReductionARResidualRMSNorm: diff --git a/flashinfer/comm/mnnvl.py b/flashinfer/comm/mnnvl.py index e127b813ced..5dec41b6f54 100644 --- a/flashinfer/comm/mnnvl.py +++ b/flashinfer/comm/mnnvl.py @@ -1004,7 +1004,6 @@ def __init__( self.signal_pad_offset = 0 self.allocation_size = 0 self.comm_backend = comm_backend_for_handle_transfer or MPIBackend() - self._enable_multicast = enable_multicast # CUDA memory handles and pointers self.mc_ptr = 0 # CUdeviceptr mMcPtr @@ -1016,24 +1015,22 @@ def __init__( self.uc_handles: List[ int ] = [] # std::vector mUcHandles - self._mapped = False # Signal pad constants self.SIGNAL_PAD_ALIGNMENT = 16 self.SIGNAL_PAD_SIZE = SIGNAL_PAD_SIZE # Check if device supports multicasting - if self._enable_multicast: - multicast_supported = checkCudaErrors( - cuda.cuDeviceGetAttribute( - cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, - device_idx, - ) + multicast_supported = checkCudaErrors( + cuda.cuDeviceGetAttribute( + cuda.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED, + device_idx, + ) + ) + if multicast_supported == 0: + raise RuntimeError( + "[SymmDeviceMemory] Device does not support multicasting." ) - if multicast_supported == 0: - raise RuntimeError( - "[SymmDeviceMemory] Device does not support multicasting." - ) # Calculate signal pad offset with alignment (matching C++ exactly) self.signal_pad_offset = round_up(buf_size, self.SIGNAL_PAD_ALIGNMENT) @@ -1044,8 +1041,16 @@ def __init__( f"Signal pad offset: {self.signal_pad_offset}" ) - self._exchanger: Optional[HandleExchanger] = None - self._create_and_map_handles(self.comm_backend) + # Create handle exchanger + if is_mnnvl_fabric_supported(device_idx): + self._exchanger: HandleExchanger = FabricHandleExchanger( + self.comm_backend, self.group_rank, self.group_size + ) + else: + self._exchanger = PosixFDHandleExchanger( + self.comm_backend, self.group_rank, self.group_size + ) + self._alloc_mn_mcast_mem(buf_size, enable_multicast) if allocate_signal_pads: # Initialize signal pads @@ -1063,7 +1068,7 @@ def __init__( def __del__(self): """Destructor - cleanup allocated memory""" - if hasattr(self, "_exchanger") and self._exchanger is not None: + if hasattr(self, "_exchanger"): self._exchanger.close() # Skip cleanup during Python finalization to avoid segfaults @@ -1121,8 +1126,6 @@ def __del__(self): checkCudaErrors(cuda.cuMemRelease(self.mc_handle)) except Exception as e: logger.warning("Destructor: Failed to release MC handle: %s", e) - elif hasattr(self, "mc_ptr") and self.mc_ptr: - checkCudaErrors(cuda.cuMemAddressFree(self.mc_ptr, self.allocation_size)) def get_signal_pad_ptrs_host(self) -> List[int]: """Get the raw array of signal pad pointers to all ranks (including self)""" @@ -1172,65 +1175,20 @@ def get_usable_buffer_size(self) -> int: """Get the usable buffer size (excluding signal pad)""" return self.allocation_size - self.SIGNAL_PAD_SIZE - @property - def mapped(self) -> bool: - return self._mapped - - def _create_and_map_handles(self, comm: CommBackend) -> None: - """Create physical backing and map it at the reserved addresses.""" - # Create handle exchanger - if is_mnnvl_fabric_supported(self.device_idx): - self._exchanger = FabricHandleExchanger( - comm, self.group_rank, self.group_size - ) - else: - self._exchanger = PosixFDHandleExchanger( - comm, self.group_rank, self.group_size - ) - + def _alloc_mn_mcast_mem(self, buf_size: int, enable_multicast: bool): + """Allocate multi-node multicast memory using MNNVL""" self._verify_cuda_context() # Compute allocation size and get allocation properties - allocation_prop, mc_prop = self._get_allocation_prop(self.buf_size) + allocation_prop, mc_prop = self._get_allocation_prop(buf_size) # Allocate, exchange, and map unicast buffers self._allocate_unicast_buffers(allocation_prop) # Setup multicast object, exchange handles, map and bind memory - if self._enable_multicast: + if enable_multicast: self._setup_multicast(mc_prop) - self.comm_backend = comm - self._mapped = True - - def _unmap_and_release_handles(self) -> None: - """Unmap and release physical backing while retaining reserved addresses.""" - # Drain local work, then align ranks before changing shared mappings. - cuda.cuCtxSynchronize() - self.comm_backend.barrier() - - if self._enable_multicast: - checkCudaErrors( - cuda.cuMulticastUnbind( - self.mc_handle, self.device_idx, 0, self.allocation_size - ) - ) - checkCudaErrors(cuda.cuMemUnmap(self.mc_ptr, self.allocation_size)) - - for ptr in self.uc_ptrs: - checkCudaErrors(cuda.cuMemUnmap(ptr, self.allocation_size)) - - if self._enable_multicast: - checkCudaErrors(cuda.cuMemRelease(self.mc_handle)) - self.mc_handle = 0 - for handle in self.uc_handles: - checkCudaErrors(cuda.cuMemRelease(handle)) - - self._exchanger.close() - self.uc_handles = [0] * self.group_size - self._exchanger = None - self._mapped = False - def _verify_cuda_context(self): """Verify CUDA context is set to the correct device.""" try: @@ -1268,23 +1226,20 @@ def _get_allocation_prop(self, buf_size: int): buf_size + self.SIGNAL_PAD_SIZE, alloc_granularity ) - self._mc_granularity = alloc_granularity - mc_prop = None - if self._enable_multicast: - # Set up multicast properties - mc_prop = cuda.CUmulticastObjectProp() - mc_prop.numDevices = self.group_size - mc_prop.size = self.allocation_size - mc_prop.handleTypes = self._exchanger.handle_type - - # Get multicast granularity and adjust allocation size - self._mc_granularity = checkCudaErrors( - cuda.cuMulticastGetGranularity( - mc_prop, - cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED, - ) + # Set up multicast properties + mc_prop = cuda.CUmulticastObjectProp() + mc_prop.numDevices = self.group_size + mc_prop.size = self.allocation_size + mc_prop.handleTypes = self._exchanger.handle_type + + # Get multicast granularity and adjust allocation size + self._mc_granularity = checkCudaErrors( + cuda.cuMulticastGetGranularity( + mc_prop, + cuda.CUmulticastGranularity_flags.CU_MULTICAST_GRANULARITY_RECOMMENDED, ) - self.allocation_size = round_up(self.allocation_size, self._mc_granularity) + ) + self.allocation_size = round_up(self.allocation_size, self._mc_granularity) return allocation_prop, mc_prop @@ -1320,24 +1275,21 @@ def _allocate_unicast_buffers(self, allocation_prop): self._exchanger.handle_type, ) ) - self._exchanger.cleanup(all_shareable_uc_handles[p]) - self._exchanger.cleanup(local_shareable_uc_handle) + self._exchanger.cleanup(all_shareable_uc_handles[p]) # Reserve address space for UC pointers - if not self.uc_ptrs: - self.uc_ptrs = [0] * self.group_size - total_uc_size = self.allocation_size * self.group_size - self.total_uc_size = total_uc_size - uc_base_ptr = checkCudaErrors( - cuda.cuMemAddressReserve(total_uc_size, self._mc_granularity, 0, 0) - ) - self.uc_base_ptr = uc_base_ptr - for i in range(self.group_size): - offset = self.allocation_size * i - self.uc_ptrs[i] = int(uc_base_ptr) + offset + self.uc_ptrs = [0] * self.group_size + total_uc_size = self.allocation_size * self.group_size + self.total_uc_size = total_uc_size + uc_base_ptr = checkCudaErrors( + cuda.cuMemAddressReserve(total_uc_size, self._mc_granularity, 0, 0) + ) + self.uc_base_ptr = uc_base_ptr # Map UC memory for i in range(self.group_size): + offset = self.allocation_size * i + self.uc_ptrs[i] = int(uc_base_ptr) + offset checkCudaErrors( cuda.cuMemMap( self.uc_ptrs[i], self.allocation_size, 0, self.uc_handles[i], 0 @@ -1347,7 +1299,7 @@ def _allocate_unicast_buffers(self, allocation_prop): # Set memory access permissions for UC access_desc = self._get_mem_access_desc() checkCudaErrors( - cuda.cuMemSetAccess(self.uc_base_ptr, self.total_uc_size, [access_desc], 1) + cuda.cuMemSetAccess(uc_base_ptr, total_uc_size, [access_desc], 1) ) def _setup_multicast(self, mc_prop): @@ -1377,18 +1329,15 @@ def _setup_multicast(self, mc_prop): self._exchanger.handle_type, ) ) - self._exchanger.cleanup(shareable_mc_handle) + self._exchanger.cleanup(shareable_mc_handle) # Add device to multicast checkCudaErrors(cuda.cuMulticastAddDevice(self.mc_handle, self.device_idx)) # Reserve and map MC pointer - if not self.mc_ptr: - self.mc_ptr = checkCudaErrors( - cuda.cuMemAddressReserve( - self.allocation_size, self._mc_granularity, 0, 0 - ) - ) + self.mc_ptr = checkCudaErrors( + cuda.cuMemAddressReserve(self.allocation_size, self._mc_granularity, 0, 0) + ) checkCudaErrors( cuda.cuMemMap(self.mc_ptr, self.allocation_size, 0, self.mc_handle, 0) ) diff --git a/flashinfer/comm/trtllm_ar.py b/flashinfer/comm/trtllm_ar.py index d8a3666a6c0..77fba4049ab 100644 --- a/flashinfer/comm/trtllm_ar.py +++ b/flashinfer/comm/trtllm_ar.py @@ -16,7 +16,7 @@ import functools import logging -from ctypes import c_void_p, cast, create_string_buffer +from ctypes import c_void_p, cast from types import SimpleNamespace from typing import List, Optional, Tuple, Union from typing_extensions import deprecated @@ -431,7 +431,7 @@ def trtllm_moe_finalize_allreduce_fusion( MAX_ALL_REDUCE_BLOCKS = 24 LamportTokenNumThreshold = 16 -_symm_workspace_refs: dict[int, list[object]] = {} +_symm_workspace_refs: dict[int, list[torch.Tensor]] = {} @deprecated( @@ -555,34 +555,6 @@ def trtllm_destroy_ipc_workspace_for_all_reduce( MAX_COMM_SIZE = 2147483647 & ~((1 << 21) - 1) # MAX_INT32 rounded down to 2MB -def _initialize_allreduce_fusion_protocol( - ipc_handles: List[List[int]], - tp_rank: int, - flag_size: int, - lamport_buffer_size: int, - lamport_comm_size: int, - use_fp32_lamport: bool, - control_flag_ptr: int, -) -> None: - cudart.cudaMemset(c_void_p(ipc_handles[1][tp_rank]), 0, flag_size) - - lamport_dtype = torch.float32 if use_fp32_lamport else torch.float16 - aligned_size = round_up(lamport_buffer_size, 16) - trtllm_lamport_initialize( - ipc_handles[2][tp_rank], - aligned_size // (4 if use_fp32_lamport else 2), - lamport_dtype, - ) - - cudart.cudaMemset(c_void_p(control_flag_ptr), 0, 5 * 4) - lamport_comm_size_bytes = create_string_buffer( - lamport_comm_size.to_bytes(4, byteorder="little"), 4 - ) - cudart.cudaMemcpy( - c_void_p(control_flag_ptr + 3 * 4), cast(lamport_comm_size_bytes, c_void_p), 4 - ) - - @deprecated( "use the unified API allreduce.py instead. It will internally call trtllm_create_ipc_workspace_for_all_reduce_fusion." ) @@ -620,8 +592,7 @@ def trtllm_create_ipc_workspace_for_all_reduce_fusion( use_fp32_lamport, buffer_size, flag_size, lamport_comm_size, lamport_buffer_size - If create_metadata=True: and use_symm_dev_mem=True: (ipc_handles, workspace_tensor, mem_handles,metadata) where metadata contains: tp_rank, tp_size, max_token_num, hidden_dim, - use_fp32_lamport, buffer_size, flag_size, lamport_comm_size, - lamport_buffer_size, control_flag_ptr + use_fp32_lamport, buffer_size, flag_size, lamport_comm_size, lamport_buffer_size and mem_handles is a list of SymmDeviceMemory objects. Note: The optional parameters make the API clunky at this time. This will be refactored in the future, at the cost of backward compatibility, where the default behavior will be @@ -668,7 +639,12 @@ def trtllm_create_ipc_workspace_for_all_reduce_fusion( lamport_buffer_size = lamport_comm_size * 3 device = torch.device(f"cuda:{torch.cuda.current_device()}") - symm_refs: list[object] = [] + group_name = ( + group.group_name + if group is not None + else torch.distributed.group.WORLD.group_name + ) + symm_refs: list[torch.Tensor] = [] # we should init 3 buffers for all reduce fusion: # [buffer_size, flag_size, lamport_buffer_size] @@ -683,35 +659,16 @@ def trtllm_create_ipc_workspace_for_all_reduce_fusion( ]: aligned_size = round_up(size, 16) - if use_symm_dev_mem: - assert comm_backend is not None - handle = SymmDeviceMemory( - buf_size=aligned_size, - group_size=tp_size, - group_rank=tp_rank, - device_idx=device.index, - comm_backend_for_handle_transfer=comm_backend, - enable_multicast=False, - allocate_signal_pads=False, - ) - ptrs = handle.get_buffer_ptrs_host() - symm_refs.append(handle) - mem_handles.append(handle) - else: - group_name = ( - group.group_name - if group is not None - else torch.distributed.group.WORLD.group_name - ) - ptrs, tensor, handle = _alloc_symm_buffer_bytes( - aligned_size, - tp_size, - dtype, - device, - group_name, - ) - symm_refs.append((tensor, handle)) + ptrs, tensor, handle = _alloc_symm_buffer_bytes( + aligned_size, + tp_size, + dtype, + device, + group_name, + ) + symm_refs.append((tensor, handle)) ipc_handles.append(ptrs) + mem_handles.append(handle) logger.debug( "rank %s allocated ipc_handles: %s", @@ -721,9 +678,6 @@ def trtllm_create_ipc_workspace_for_all_reduce_fusion( _symm_workspace_refs[id(ipc_handles)] = symm_refs - if use_symm_dev_mem: - cudart.cudaMemset(c_void_p(ipc_handles[1][tp_rank]), 0, flag_size) - # Initialize lamport buffer aligned_lamport_buffer_size = round_up(lamport_buffer_size, 16) if use_fp32_lamport: @@ -774,7 +728,6 @@ def trtllm_create_ipc_workspace_for_all_reduce_fusion( ) if use_symm_dev_mem: - torch.cuda.synchronize() comm_backend.barrier() # must sync after create_workspace else: dist.barrier(group=group) @@ -792,7 +745,6 @@ def trtllm_create_ipc_workspace_for_all_reduce_fusion( "lamport_buffer_size": lamport_buffer_size, } if use_symm_dev_mem: - metadata["control_flag_ptr"] = flag_ptr.value return ipc_handles, workspace_tensor, mem_handles, metadata else: return ipc_handles, workspace_tensor, metadata diff --git a/flashinfer/comm/trtllm_mnnvl_ar.py b/flashinfer/comm/trtllm_mnnvl_ar.py index 0a76914d20a..bc48915c7b2 100644 --- a/flashinfer/comm/trtllm_mnnvl_ar.py +++ b/flashinfer/comm/trtllm_mnnvl_ar.py @@ -14,14 +14,15 @@ from typing_extensions import deprecated from flashinfer.comm.mapping import Mapping -from flashinfer.api_logging import flashinfer_api +from flashinfer.comm.mnnvl import TorchDistBackend from ..jit import gen_trtllm_mnnvl_comm_module from ..utils import register_custom_op from ..fp4_quantization import _compute_swizzled_layout_sf_size -from .mnnvl import CommBackend, McastGPUBuffer, MPIBackend, SymmDeviceMemory +from .mnnvl import CommBackend, MPIBackend from .trtllm_ar import QuantizationSFLayout from .workspace_base import AllReduceFusionWorkspace +from .torch_symmetric_memory import _alloc_symm_buffer_bytes def mpi_barrier(): @@ -146,16 +147,26 @@ def __init__( # support base_gpu_id != 0 scenarios where the actual CUDA device # index differs from the TP rank / local_rank. device = torch.device("cuda", torch.cuda.current_device()) - self.handle = McastGPUBuffer( - buf_size=requested_workspace_size, - group_size=mapping.tp_size, - group_rank=mapping.tp_rank, - device=device, - comm_backend_for_handle_transfer=comm_backend, + if isinstance(comm_backend, TorchDistBackend): + group = ( + comm_backend._group + if comm_backend._group is not None + else torch.distributed.group.WORLD + ) + group_name = group.group_name + else: + group_name = torch.distributed.group.WORLD.group_name + self.ptrs, self.tensor, self.handle = _alloc_symm_buffer_bytes( + requested_workspace_size, + mapping.tp_size, + torch.float32, + device, + group_name, ) - self.ptrs = self.handle.mcast_device_memory.get_buffer_ptrs_host() - allocated_size = self.handle.buf_size + # handle.buffer_size is the usable data size. torch symmetric memory + # allocator places signal_pad on top of it, not carved from within. + allocated_size = self.handle.buffer_size # We want the buffer size to be aligned to 16B which is the granularity for buffer management. self.buffer_size_bytes = ( math.floor(allocated_size / self.NUM_LAMPORT_BUFFERS) // 16 * 16 @@ -167,7 +178,9 @@ def __init__( f"[MNNVL Allreduce] Actual allocated size: {allocated_size} bytes, Actual buffer size per lamport buffer: {self.buffer_size_bytes} bytes, total workspace: {self.workspace_size_bytes} bytes." ) - self.handle.lamport_initialize(self.rank, torch.float32) + # lamport initialize tensor to negative zero. + self.tensor.fill_(-0.0) + # Wait until the initialization is done torch.cuda.synchronize() comm_backend.barrier() @@ -182,9 +195,9 @@ def __init__( device=torch.device("cuda", torch.cuda.current_device()), ) - self.uc_ptrs_dev = self.handle.get_buffer_ptrs_dev() - self.uc_ptr_local = self.handle.get_unicast_ptr(self.rank) - self.mc_ptr = self.handle.get_multicast_ptr() + self.uc_ptrs_dev = self.handle.buffer_ptrs_dev + self.uc_ptr_local = self.handle.buffer_ptrs[self.rank] + self.mc_ptr = self.handle.multicast_ptr @functools.cache def is_buffer_size_sufficient( @@ -239,53 +252,6 @@ def get_required_buffer_size_bytes( def backend(self) -> str: return "mnnvl" - def _require_handles_attached(self) -> None: - memory = getattr(self.handle, "mcast_device_memory", None) - if isinstance(memory, SymmDeviceMemory) and not memory.mapped: - raise RuntimeError("MNNVL handles are not attached") - - def _initialize_protocol(self) -> None: - self.handle.lamport_initialize(self.rank, torch.float32) - num_bytes_to_clear = [0] * 4 - self.buffer_flags.copy_( - torch.tensor( - [0, 2, self.buffer_size_bytes, 0, *num_bytes_to_clear, 0], - dtype=torch.uint32, - device=self.buffer_flags.device, - ) - ) - torch.cuda.synchronize() - - @flashinfer_api - def checkpoint_prepare(self) -> None: - """Detach physical backing; repeated successful calls are no-ops.""" - memory = getattr(self.handle, "mcast_device_memory", None) - if not isinstance(memory, SymmDeviceMemory): - raise NotImplementedError( - "Stable-VA checkpointing is unavailable for workspaces backed " - "by torch symmetric memory" - ) - if not memory.mapped: - return - memory._unmap_and_release_handles() - # Do not return until every rank has released its workspace handles. - memory.comm_backend.barrier() - - @flashinfer_api - def checkpoint_restore(self, comm_backend: CommBackend) -> None: - """Restore physical backing; repeated successful calls are no-ops.""" - memory = getattr(self.handle, "mcast_device_memory", None) - if not isinstance(memory, SymmDeviceMemory): - raise NotImplementedError( - "Stable-VA checkpointing is unavailable for workspaces backed " - "by torch symmetric memory" - ) - if memory.mapped: - return - memory._create_and_map_handles(comm_backend) - self._initialize_protocol() - comm_backend.barrier() - def destroy(self) -> None: """Destroy workspace and free resources.""" if getattr(self, "_destroyed", False): @@ -295,6 +261,7 @@ def destroy(self) -> None: del self.uc_ptrs_dev del self.uc_ptr_local del self.mc_ptr + del self.tensor del self.handle del self.ptrs self._destroyed = True @@ -466,8 +433,6 @@ def trtllm_mnnvl_allreduce( f"The output tensor must be 2D, got {len(output.shape)}D. The shape is {output.shape}." ) - workspace._require_handles_attached() - module = get_trtllm_mnnvl_comm_module() if strategy == MNNVLAllreduceFusionStrategy.AUTO: @@ -568,8 +533,6 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm( f"The residual output tensor must be 2D, got {len(residual_out.shape)}D. The shape is {residual_out.shape}." ) - workspace._require_handles_attached() - module = get_trtllm_mnnvl_comm_module() if strategy == MNNVLAllreduceFusionStrategy.AUTO: @@ -809,8 +772,6 @@ def trtllm_mnnvl_fused_allreduce_add_rmsnorm_quant( else: raise ValueError(f"Unsupported MNNVL quant_type: {quant_type}") - workspace._require_handles_attached() - if strategy == MNNVLAllreduceFusionStrategy.AUTO: strategy = MNNVLAllreduceFusionStrategy.select_strategy( workspace.tp_size, token_num, hidden_dim, input.dtype @@ -884,7 +845,7 @@ def get_allreduce_mnnvl_workspace( Returns: Tuple containing: - - MNNVLAllReduceFusionWorkspace: The CUDA VMM-backed workspace object + - MNNVLAllReduceFusionWorkspace: The workspace object backed by torch symmetric memory - torch.Tensor: Buffer flags tensor tracking state - int: Maximum number of elements that can fit in buffer """ diff --git a/tests/comm/test_trtllm_allreduce_checkpoint.py b/tests/comm/test_trtllm_allreduce_checkpoint.py deleted file mode 100644 index b19a0a3dac4..00000000000 --- a/tests/comm/test_trtllm_allreduce_checkpoint.py +++ /dev/null @@ -1,120 +0,0 @@ -import socket - -import pytest -import torch -import torch.distributed as dist -import torch.multiprocessing as mp - - -def test_checkpoint_lifecycle_rejects_torch_symmetric_memory_backing() -> None: - from flashinfer.comm.allreduce import ( - MNNVLAllReduceFusionWorkspace, - TRTLLMAllReduceFusionWorkspace, - ) - from flashinfer.comm.workspace_base import AllReduceFusionWorkspace - - trt_workspace = object.__new__(TRTLLMAllReduceFusionWorkspace) - AllReduceFusionWorkspace.__init__(trt_workspace, world_size=1, rank=0) - trt_workspace.mem_handles = [] - trt_workspace._destroyed = True - - mnnvl_workspace = object.__new__(MNNVLAllReduceFusionWorkspace) - AllReduceFusionWorkspace.__init__(mnnvl_workspace, world_size=1, rank=0) - mnnvl_workspace.handle = object() - mnnvl_workspace._destroyed = True - - for workspace in (trt_workspace, mnnvl_workspace): - with pytest.raises(NotImplementedError, match="torch symmetric memory"): - workspace.checkpoint_prepare() - with pytest.raises(NotImplementedError, match="torch symmetric memory"): - workspace.checkpoint_restore(None) - - -def _free_port() -> int: - with socket.socket() as sock: - sock.bind(("127.0.0.1", 0)) - return sock.getsockname()[1] - - -def _run_worker( - rank: int, world_size: int, port: int, num_tokens: int, use_oneshot: bool -) -> None: - torch.cuda.set_device(rank) - dist.init_process_group( - "gloo", - init_method=f"tcp://127.0.0.1:{port}", - rank=rank, - world_size=world_size, - ) - workspace = None - try: - import flashinfer.comm as comm - from flashinfer.comm.mnnvl import TorchDistBackend - - workspace = comm.create_allreduce_fusion_workspace( - backend="trtllm", - world_size=world_size, - rank=rank, - max_token_num=num_tokens, - hidden_dim=4096, - dtype=torch.bfloat16, - comm_backend=TorchDistBackend(), - ) - input_ = torch.full( - (num_tokens, 4096), rank + 1, dtype=torch.bfloat16, device="cuda" - ) - output = torch.empty_like(input_) - - def all_reduce() -> None: - comm.allreduce_fusion( - input=input_, - workspace=workspace, - output=output, - pattern=comm.AllReduceFusionPattern.kAllReduce, - use_oneshot=use_oneshot, - ) - - all_reduce() - torch.cuda.synchronize() - torch.testing.assert_close(output, torch.full_like(output, 3)) - - graph = torch.cuda.CUDAGraph() - dist.barrier() - with torch.cuda.graph(graph): - all_reduce() - - dist.barrier() - workspace.checkpoint_prepare() - workspace.checkpoint_prepare() - with pytest.raises(RuntimeError, match="not attached"): - all_reduce() - dist.barrier() - fresh_backend = TorchDistBackend() - workspace.checkpoint_restore(fresh_backend) - workspace.checkpoint_restore(fresh_backend) - dist.barrier() - - input_.fill_(rank + 2) - graph.replay() - torch.cuda.synchronize() - torch.testing.assert_close(output, torch.full_like(output, 5)) - finally: - if workspace is not None: - workspace.destroy() - dist.destroy_process_group() - - -@pytest.mark.parametrize("num_tokens,use_oneshot", [(1, True), (4, False)]) -@pytest.mark.skipif( - not torch.cuda.is_available() or torch.cuda.device_count() < 2, - reason="checkpointable TRT-LLM all-reduce requires two CUDA devices", -) -def test_graph_replay_after_symmetric_memory_remap( - num_tokens: int, use_oneshot: bool -) -> None: - mp.spawn( - _run_worker, - args=(2, _free_port(), num_tokens, use_oneshot), - nprocs=2, - join=True, - ) diff --git a/tests/comm/test_trtllm_mnnvl_allreduce.py b/tests/comm/test_trtllm_mnnvl_allreduce.py index a7dd10fc916..1e10e36df5f 100644 --- a/tests/comm/test_trtllm_mnnvl_allreduce.py +++ b/tests/comm/test_trtllm_mnnvl_allreduce.py @@ -458,7 +458,7 @@ def run_mnnvl_ar_full( multicast_ptr = legacy_workspace.mc_ptr buffer_ptrs_dev = legacy_workspace.uc_ptrs_dev - unicast_ptr = legacy_workspace.uc_ptr_local + unicast_ptr = legacy_workspace.handle.buffer_ptrs[mapping.tp_rank] else: workspace = trtllm_mnnvl_ar.MNNVLAllReduceFusionWorkspace(