unify memory allocation for single-node and multi-node nvlink all reduce(+fusion) using torch symmetric memory

flashinfer/comm/torch_symmetric_memory.py

@@ -0,0 +1,16 @@
+import torch
+import torch.distributed._symmetric_memory as symm_mem
+
+
+def _alloc_symm_buffer_bytes(
+    size_bytes: int, tp_size: int, dtype: torch.dtype, device: torch.device, group_name: str
+) -> tuple[list[int], torch.Tensor]:
+    numel = size_bytes // torch.tensor([], dtype=dtype).element_size()
+    tensor = symm_mem.empty(numel, dtype=dtype, device=device)
+    handle = symm_mem.rendezvous(tensor, group=group_name)
+    ptrs: list[int] = []
+    for peer in range(tp_size):
+        buf = handle.get_buffer(peer, (numel,), dtype, storage_offset=0)
+        ptrs.append(buf.data_ptr())
+    return ptrs, tensor, handle
+

flashinfer/comm/trtllm_ar.py

@@ -29,7 +29,7 @@
 from ..jit.comm import gen_trtllm_comm_module
 from ..utils import register_custom_op, round_up
 from .cuda_ipc import create_shared_buffer, cudart, free_shared_buffer
-
+from .torch_symmetric_memory import _alloc_symm_buffer_bytes
 
 class AllReduceStrategyType:
     # NOTE: for trtllm_custom_all_reduce
@@ -397,6 +397,8 @@ def trtllm_moe_finalize_allreduce_fusion(
 MAX_ALL_REDUCE_BLOCKS = 24
 LamportTokenNumThreshold = 16
 
+_symm_workspace_refs: dict[int, list[torch.Tensor]] = {}
+
 
 @deprecated(
     "trtllm_create_ipc_workspace_for_all_reduce and trtllm_custom_all_reduce are deprecated and will be removed in the next major bump, use allreduce.py instead."
@@ -449,25 +451,36 @@ def trtllm_create_ipc_workspace_for_all_reduce(
     flag_size = FLAG_SIZE * tp_size * 2
     lamport_buffer_size = tp_size * LamportTokenNumThreshold * tp_size * hidden_dim * 2
 
+    device = torch.device(f"cuda:{torch.cuda.current_device()}")
+    group_name = group.group_name
+    symm_refs: list[torch.Tensor] = []
     ipc_handles = list()
 
-    for size in [
-        buffer_size,
-        buffer_size,
-        flag_size,
-        flag_size,
-        lamport_buffer_size,
-        lamport_buffer_size,
-        lamport_buffer_size,
+    for size, dtype in [
+        (buffer_size, torch.float32),
+        (buffer_size, torch.float32),
+        (flag_size, torch.int32),
+        (flag_size, torch.int32),
+        (lamport_buffer_size, torch.float16),
+        (lamport_buffer_size, torch.float16),
+        (lamport_buffer_size, torch.float16),
     ]:
         # all sizes should be aligned to 1LU << 21 bytes (2MB)
         aligned_size = round_up(size, 1 << 21)
-        ipc_handles.append(create_shared_buffer(aligned_size, group))
 
+        ptrs, tensor, handle = _alloc_symm_buffer_bytes(
+            aligned_size,
+            tp_size,
+            dtype,
+            device,
+            group_name,
+        )
+        symm_refs.append((tensor, handle))
+        ipc_handles.append(ptrs)
     print(
         f"rank {rank} allocated ipc_handles: {[[hex(handle) for handle in sublist] for sublist in ipc_handles]}"
     )
-
+    _symm_workspace_refs[id(ipc_handles)] = symm_refs
     trtllm_lamport_initialize_all(
         ipc_handles[4][rank],
         ipc_handles[5][rank],
@@ -491,10 +504,7 @@ def trtllm_destroy_ipc_workspace_for_all_reduce(
     The workspace should be destroyed after calling trtllm_custom_all_reduce.
     The workspace can be reused for multiple all reduce calls under the same configuration.
     """
-
-    for ipc_handle in workspace:
-        free_shared_buffer(ipc_handle, group)
-
+    _ = _symm_workspace_refs.pop(id(workspace), None)
 
 BarrierFlagCount = 256
 
@@ -584,35 +594,40 @@ 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()}")
+    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]
 
     ipc_handles: List[List[int]] = list()
     mem_handles: List[SymmDeviceMemory] = list()
-    for size in [buffer_size, flag_size, lamport_buffer_size]:
+    #for size in [buffer_size, flag_size, lamport_buffer_size]:
+    lamport_buffer_dtype = torch.float16 if not use_fp32_lamport else torch.float32
+    for size, dtype in [
+        (buffer_size, torch.float32),
+        (flag_size, torch.int32),
+        (lamport_buffer_size, lamport_buffer_dtype),
+    ]:
         # todo(review): confirm we need this alignment
         # all sizes should be aligned to 1LU << 21 bytes (2MB)
         aligned_size = round_up(size, 1 << 21)
 
-        if not use_symm_dev_mem:
-            ipc_handles.append(create_shared_buffer(aligned_size, group))
-        else:
-            symm_mem = SymmDeviceMemory(
-                aligned_size,
-                tp_size,
-                tp_rank,
-                torch.device("cuda", tp_rank).index,
-                comm_backend,
-                enable_multicast=False,
-                allocate_signal_pads=False,
-            )
-            ipc_handles.append(symm_mem.uc_ptrs)
-            mem_handles.append(symm_mem)
+        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)
 
     print(
         f"rank {tp_rank} allocated ipc_handles: {[[hex(handle) for handle in sublist] for sublist in ipc_handles]}"
     )
-
+    _symm_workspace_refs[id(ipc_handles)] = symm_refs
     # Initialize lamport buffer
     aligned_lamport_buffer_size = round_up(lamport_buffer_size, 1 << 21)
     if use_fp32_lamport:
@@ -703,8 +718,7 @@ def trtllm_destroy_ipc_workspace_for_all_reduce_fusion(
     The workspace can be reused for multiple all reduce fusion calls under the same configuration.
     """
 
-    for ipc_handle in workspace:
-        free_shared_buffer(ipc_handle, group)
+    _ = _symm_workspace_refs.pop(id(workspace), None)
 
 
 # allReduce fused quant utils

flashinfer/comm/trtllm_mnnvl_ar.py

@@ -19,7 +19,21 @@
 from ..utils import register_custom_op
 from .mnnvl import McastGPUBuffer, CommBackend, MPIBackend
 from .workspace_base import AllReduceFusionWorkspace
-
+from .torch_symmetric_memory import _alloc_symm_buffer_bytes
+from ..cuda_utils import checkCudaErrors
+try:
+    # cuda-python >= 12.9 (has cuda.bindings.driver)
+    from cuda.bindings import driver as cuda
+except ImportError:
+    try:
+        # cuda-python < 12.9 (no cuda.bindings.driver, use cuda as driver)
+        # from cuda import cuda is not available in cuda-python >= 13.0
+        from cuda import cuda
+    except ImportError as e:
+        raise ImportError(
+            "Could not import the 'cuda' module. "
+            "Please install cuda-python that matches your CUDA version."
+        ) from e
 
 def mpi_barrier():
     from mpi4py import MPI
@@ -131,17 +145,19 @@ def __init__(
             f"[MNNVL Allreduce] TP size: {mapping.tp_size}, rank: {mapping.tp_rank}, Allocating workspace with requested size {buffer_size_in_bytes} bytes per buffer."
         )
 
-        # Allocate the workspace
-        self.mcast_buffer_handle = McastGPUBuffer(
+        device=torch.device("cuda", mapping.local_rank)
+        #TODO (asamani): fix group based on comm_backend?
+        group = torch.distributed.group.WORLD
+        group_name = group.group_name if group is not None else torch.distributed.group.WORLD.group_name
+        self.ptrs, self.tensor, self.handle = _alloc_symm_buffer_bytes(
             requested_workspace_size,
             mapping.tp_size,
-            mapping.tp_rank,
-            torch.device("cuda", mapping.local_rank),
-            comm_backend,
+            torch.float32, #TODO(asamani): should this always be f32 or dtype?
+            device,
+            group_name,
         )
-
         # Get the actual usable buffer size after allocation (buf_size is updated by McastGPUBuffer)
-        allocated_size = self.mcast_buffer_handle.buf_size
+        allocated_size = self.handle.buffer_size - self.handle.signal_pad_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
@@ -154,7 +170,8 @@ def __init__(
         )
 
         # We use FP32 for sentinel value regardless of the real dtype
-        self.mcast_buffer_handle.lamport_initialize(mapping.tp_rank, torch.float32)
+        self.lamport_initialize(mapping.tp_rank, torch.float32, allocated_size)
+
         # Wait until the initialization is done
         torch.cuda.synchronize()
         comm_backend.barrier()
@@ -170,9 +187,26 @@ def __init__(
             device=torch.device("cuda", mapping.local_rank),
         )
 
-        self.uc_ptrs_dev = self.mcast_buffer_handle.get_buffer_ptrs_dev()
-        self.uc_ptr_local = self.mcast_buffer_handle.get_unicast_ptr(self.rank)
-        self.mc_ptr = self.mcast_buffer_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
+
+    def lamport_initialize(self, rank: int, dtype: torch.dtype, allocated_size: int):
+        if dtype == torch.bfloat16 or dtype == torch.float16:
+            neg_zero = 0x8000
+            dsize = 2
+            memset_func = cuda.cuMemsetD16
+        elif dtype == torch.float32:
+            neg_zero = 0x80000000
+            dsize = 4
+            memset_func = cuda.cuMemsetD32
+        else:
+            raise ValueError(f"Unsupported dtype: {dtype}")
+
+        num_elements = (allocated_size) // dsize
+        checkCudaErrors(
+            memset_func(int(self.ptrs[rank]), neg_zero, num_elements)
+        )
 
     @functools.cache
     def is_buffer_size_sufficient(
@@ -232,11 +266,13 @@ def destroy(self) -> None:
         if getattr(self, "_destroyed", False):
             return  # Already destroyed, nothing to do
 
-        del self.mcast_buffer_handle
         del self.buffer_flags
         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
 
 

tests/comm/test_trtllm_allreduce_fusion.py

@@ -110,7 +110,7 @@ def _run_correctness_worker(
                                     dist.barrier(group=group)
                                     test_passed = True
                                     print(
-                                        f"test RANK {rank}: token{token_num}-hidden_dim{hidden_dim}-dtype{dtype}-pattern{pattern_code}-layout{swizzled_layout_code}-pdl{launch_with_pdl} start"
+                                        f"test RANK {rank}: token{token_num}-hidden_dim{hidden_dim}-dtype{dtype}-pattern{pattern_code}-layout{swizzled_layout_code}-pdl{launch_with_pdl}-isOneShot{use_oneshot}-isTriggerCompletionAtEnd{trigger_completion_at_end}-isFP32Acc{fp32_acc} start"
                                     )
                                     dist.barrier(group=group)
                                     torch.cuda.synchronize()
@@ -384,11 +384,11 @@ def _run_correctness_worker(
                                     dist.barrier(group=group)
                                     if test_passed:
                                         print(
-                                            f"test RANK {rank}: token{token_num}-hidden_dim{hidden_dim}-dtype{dtype}-pattern{pattern_code}-layout{swizzled_layout_code}-pdl{launch_with_pdl} passed"
+                                            f"test RANK {rank}: token{token_num}-hidden_dim{hidden_dim}-dtype{dtype}-pattern{pattern_code}-layout{swizzled_layout_code}-pdl{launch_with_pdl}-isOneShot{use_oneshot}-isTriggerCompletionAtEnd{trigger_completion_at_end}-isFP32Acc{fp32_acc} passed"
                                         )
                                     else:
                                         print(
-                                            f"test RANK {rank}: token{token_num}-hidden_dim{hidden_dim}-dtype{dtype}-pattern{pattern_code}-layout{swizzled_layout_code}-pdl{launch_with_pdl} failed"
+                                            f"test RANK {rank}: token{token_num}-hidden_dim{hidden_dim}-dtype{dtype}-pattern{pattern_code}-layout{swizzled_layout_code}-pdl{launch_with_pdl}-isOneShot{use_oneshot}-isTriggerCompletionAtEnd{trigger_completion_at_end}-isFP32Acc{fp32_acc} failed"
                                         )
     finally:
         dist.barrier(group=group)