[Feature][Example] Support return_peers in _allocate_tensor and add ag_gemm_ipc example

chengyupku · chengyupku · commit 9de9e1ba3a97 · 2025-09-25T21:30:18.000+08:00
diff --git a/examples/distributed/example_all_to_all.py b/examples/distributed/example_all_to_all.py
@@ -1,14 +1,12 @@
 import torch
 import pynvshmem
-import os
 import tilelang
 import tilelang.language as T
 from tilelang.profiler import TensorSupplyType
 from tilelang.distributed.utils import init_distributed
 import argparse
 import random
 
-
 tilelang.disable_cache()
 
 
diff --git a/examples/distributed/example_allgather_gemm_ipc.py b/examples/distributed/example_allgather_gemm_ipc.py
@@ -0,0 +1,199 @@
+import os
+import tilelang
+import tilelang.language as T
+import argparse
+import torch
+import torch.distributed as dist
+import torch.multiprocessing
+from tilelang.distributed.utils import init_dist
+from cuda import cudart
+from tilelang.distributed.utils import set_signal, wait_eq
+
+tilelang.disable_cache()
+os.environ['NCCL_DEBUG'] = 'WARN'  # silence NCCL log
+
+
+def gemm_kernel(M,
+                N,
+                K,
+                num_rank,
+                block_M,
+                block_N,
+                block_K,
+                threads,
+                dtype="float16",
+                accum_dtype="float"):
+
+    @T.prim_func
+    def main(
+            A: T.Tensor((M, K), dtype),
+            B: T.Tensor((K, N // num_rank), dtype),
+            C: T.Tensor((M, N // num_rank), dtype),
+    ):
+        with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=threads) as (bx, by):
+            A_shared = T.alloc_shared((block_M, block_K), dtype)
+            B_shared = T.alloc_shared((block_K, block_N), dtype)
+            C_local = T.alloc_fragment((block_M, block_N), accum_dtype)
+
+            T.clear(C_local)
+            for k in T.Pipelined(T.ceildiv(K, block_K), num_stages=3):
+                T.copy(A[by * block_M, k * block_K], A_shared)
+                T.copy(B[k * block_K, bx * block_N], B_shared)
+                T.gemm(A_shared, B_shared, C_local)
+            T.copy(C_local, C[by * block_M, bx * block_N])
+
+    return main
+
+
+def cp_engine_producer_all_gather_put(local_tensor, ag_buffer, signal_buffer, M_per_rank, N,
+                                      signal_target, rank, local_world_size, world_size,
+                                      intranode_ag_stream):
+    local_rank = rank % local_world_size
+    n_nodes = world_size // local_world_size
+    node_rank = rank // local_world_size
+
+    for i in range(1, local_world_size):
+        segment = rank * M_per_rank * N
+        local_dst_rank = (local_rank + local_world_size - i) % local_world_size
+        src_ptr = ag_buffer[local_rank].data_ptr() + segment * local_tensor.element_size()
+        dst_ptr = ag_buffer[local_dst_rank].data_ptr() + segment * local_tensor.element_size()
+        # Using copy engine to perform intranode transmission
+        # Sending rank-th local tensor to other ranks inside the node.
+        (err,) = cudart.cudaMemcpyAsync(
+            dst_ptr,
+            src_ptr,
+            M_per_rank * N * local_tensor.element_size(),
+            cudart.cudaMemcpyKind.cudaMemcpyDefault,
+            intranode_ag_stream.cuda_stream,
+        )
+        # Notify the peer that the transmission is done.
+        set_signal(signal_buffer[local_dst_rank][rank], signal_target, intranode_ag_stream)
+
+    for i in range(1, n_nodes):
+        recv_rank = local_rank + (node_rank + n_nodes - i) % n_nodes * local_world_size
+        recv_segment = recv_rank * M_per_rank * N
+        # Waiting for the internode data ready
+        wait_eq(signal_buffer[local_rank][recv_rank], signal_target, intranode_ag_stream)
+        src_ptr = ag_buffer[local_rank].data_ptr() + recv_segment * local_tensor.element_size()
+        for j in range(1, local_world_size):
+            local_dst_rank = (local_rank + local_world_size - j) % local_world_size
+            dst_ptr = ag_buffer[local_dst_rank].data_ptr(
+            ) + recv_segment * local_tensor.element_size()
+            # Sending (local_rank + j*local_world_size) % world_size -th local tensor to other ranks inside the node.
+            (err,) = cudart.cudaMemcpyAsync(
+                dst_ptr,
+                src_ptr,
+                M_per_rank * N * local_tensor.element_size(),
+                cudart.cudaMemcpyKind.cudaMemcpyDefault,
+                intranode_ag_stream.cuda_stream,
+            )
+            # Notify the peer that the transmission is done.
+            set_signal(signal_buffer[local_dst_rank][recv_rank], signal_target, intranode_ag_stream)
+
+
+def ag_gemm_op(A, B, C, ag_buffer, signal_buffer, M_per_rank, N, signal_target, rank, group,
+               local_world_size, world_size, gemm_kernel, ag_stream):
+
+    dist.barrier(group)
+
+    # all_gather A to ag_buffer
+    with torch.cuda.stream(ag_stream):
+        cp_engine_producer_all_gather_put(A, ag_buffer, signal_buffer, M_per_rank, N, signal_target,
+                                          rank, local_world_size, world_size, ag_stream)
+
+    current_stream = torch.cuda.current_stream()
+    current_stream.wait_stream(ag_stream)
+
+    dist.barrier(group)
+    torch.cuda.synchronize()
+
+    torch.cuda.synchronize()
+    torch.distributed.barrier(group)
+    gemm_kernel(ag_buffer[rank], B, C)
+    torch.cuda.synchronize()
+    torch.distributed.barrier(group)
+
+    return C
+
+
+def torch_ag_gemm(
+    pg: torch.distributed.ProcessGroup,
+    local_input: torch.Tensor,
+    local_weight: torch.Tensor,
+    ag_out: torch.Tensor,
+):
+    torch.distributed.all_gather_into_tensor(ag_out, local_input, pg)
+    ag_gemm_output = torch.matmul(ag_out, local_weight)
+    return ag_gemm_output
+
+
+def main(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
+    dtype = torch.float16
+    M = args.M if args else 8192
+    N = args.N if args else 8192
+    K = args.K if args else 8192
+    M_per_rank = M // num_local_ranks
+    N_per_rank = N // num_local_ranks
+
+    BLOCK_M = 128
+    BLOCK_N = 128
+    BLOCK_K = 64
+    threads = 256
+    assert num_local_ranks == 2, "this example only supports 2 ranks copying to each other"
+
+    rank, num_ranks, group = init_dist(local_rank, num_local_ranks)
+    allocator = tilelang.get_allocator(
+        size=2**30,
+        device="cuda",
+        is_distributed=True,
+        local_rank=local_rank,
+        num_local_ranks=num_local_ranks,
+        group=group)
+    kernel = tilelang.compile(gemm_kernel(M, N, K, num_ranks, BLOCK_M, BLOCK_N, BLOCK_K, threads))
+    kernel.initialize(allocator=allocator)
+    if local_rank == 0:
+        print(kernel.get_kernel_source())
+
+    A = tilelang.tensor((M_per_rank, K), dtype, allocator=allocator).normal_()
+    B = tilelang.tensor((K, N_per_rank), dtype, allocator=allocator).normal_()
+    C = tilelang.tensor((M, N_per_rank), dtype, allocator=allocator)
+    ag_buffer = tilelang.tensor((M, K), dtype, allocator=allocator, return_peers=True)
+    signal_buffer = tilelang.tensor((num_local_ranks,),
+                                    torch.int32,
+                                    allocator=allocator,
+                                    return_peers=True)
+    signal_buffer[rank].fill_(0)
+    ag_buffer[rank][rank * M_per_rank:(rank + 1) * M_per_rank, :].copy_(A)
+
+    dist.barrier(group)
+
+    ag_stream = torch.cuda.Stream()
+    signal_target = 1
+
+    tilelang_C = ag_gemm_op(A, B, C, ag_buffer, signal_buffer, M_per_rank, K, signal_target, rank,
+                            group, num_local_ranks, num_local_ranks, kernel, ag_stream)
+
+    torch_ag_buffer = torch.empty([M, K], dtype=dtype, device="cuda")
+    torch_C = torch_ag_gemm(group, A, B, torch_ag_buffer)
+
+    if torch.allclose(torch_C, tilelang_C, atol=1e-6, rtol=1e-6):
+        print(f"rank {local_rank} check passed.✅")
+    else:
+        print(f"rank {local_rank} check failed.❌")
+        print(f"torch_C: {torch_C}, tilelang_C: {tilelang_C}")
+        raise ValueError("Test failed")
+
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--num-processes', type=int, default=2, help='Number of processes to spawn (default: 2)')
+    parser.add_argument('--M', type=int, default=8192, help='M dimension')
+    parser.add_argument('--N', type=int, default=8192, help='N dimension')
+    parser.add_argument('--K', type=int, default=8192, help='K dimension')
+    args = parser.parse_args()
+    num_processes = args.num_processes
+
+    torch.multiprocessing.spawn(main, args=(num_processes, args), nprocs=num_processes)
diff --git a/tilelang/distributed/utils.py b/tilelang/distributed/utils.py
@@ -3,7 +3,7 @@
 import datetime
 import os
 import inspect
-from typing import List, Union, Tuple, Callable, Sequence
+from typing import List, Union, Tuple, Callable, Sequence, Optional
 from contextlib import contextmanager
 
 import importlib.metadata
@@ -263,3 +263,36 @@ def supports_p2p_native_atomic():
         cudart.cudaDeviceP2PAttr.cudaDevP2PAttrNativeAtomicSupported, 0, 1)
     CUDA_CHECK(err)
     return support == 1
+
+
+def set_signal(signal_tensor: torch.Tensor,
+               signal: int,
+               stream: Optional[torch.cuda.Stream] = None):
+    stream = stream or torch.cuda.current_stream()
+    if signal_tensor.dtype == torch.int32:
+        (err,) = cuda.cuStreamWriteValue32(
+            stream.cuda_stream,
+            signal_tensor.data_ptr(),
+            signal,
+            cuda.CUstreamWriteValue_flags.CU_STREAM_WRITE_VALUE_DEFAULT,
+        )
+        CUDA_CHECK(err)
+    else:
+        raise Exception(f"Unsupported signal dtype {signal_tensor.dtype}")
+
+
+def wait_eq(signal_tensor: torch.Tensor,
+            signal: int,
+            stream: Optional[torch.cuda.Stream] = None,
+            require_i64=False):
+    stream = stream or torch.cuda.current_stream()
+    if signal_tensor.dtype == torch.int32:
+        (err,) = cuda.cuStreamWaitValue32(
+            stream.cuda_stream,
+            signal_tensor.data_ptr(),
+            signal,
+            cuda.CUstreamWaitValue_flags.CU_STREAM_WAIT_VALUE_EQ,
+        )
+        CUDA_CHECK(err)
+    else:
+        raise Exception(f"Unsupported signal dtype {signal_tensor.dtype}")
diff --git a/tilelang/utils/allocator.py b/tilelang/utils/allocator.py
@@ -125,7 +125,7 @@ def _alloc(self):
         if rc != 0:
             msg = _libcudart.cudaGetErrorString(rc)
             raise RuntimeError(f"cudaMalloc failed: {rc} {msg.decode() if msg else ''}")
-        self._ptr = self._base_ptr
+        self._ptr.value = self._base_ptr.value
 
     def _free(self):
         if getattr(self, "_base_ptr", None) and self._base_ptr.value:
@@ -166,6 +166,7 @@ def initialized(self) -> bool:
     def _allocate_tensor(self,
                          shape: Tuple[int, ...],
                          dtype: torch.dtype,
+                         return_peers=False,
                          take_ownership: bool = False) -> torch.Tensor:
 
         numel = _prod_shape(shape)
@@ -198,13 +199,23 @@ def _allocate_tensor(self,
 
         t = tensor_from_ptr(cur_ptr_val, shape, dtype_str, self._device, take_ownership)
 
+        if return_peers:
+            peer_ts = []
+            for i in range(self._group.size()):
+                if i == self._local_rank:
+                    peer_ts.append(t)
+                else:
+                    peer_ptr_val = int(self._buffer_ptrs[i]) + current_offset
+                    peer_t = tensor_from_ptr(peer_ptr_val, shape, dtype_str, self._device, False)
+                    peer_ts.append(peer_t)
+
         if take_ownership:
             self._ptr = ctypes.c_void_p(0)
         else:
             new_ptr_val = cur_ptr_val + bytes_alloc
             self._ptr.value = new_ptr_val
 
-        return t
+        return peer_ts if return_peers else t
 
     @property
     def ptr(self) -> int:
diff --git a/tilelang/utils/tensor.py b/tilelang/utils/tensor.py
@@ -45,14 +45,15 @@ def map_torch_type(intype: str) -> torch.dtype:
 def tensor(shape: Tuple[int, ...],
            dtype: torch.dtype,
            device: Optional[Union[str, torch.device, int]] = None,
-           allocator: Optional[BaseAllocator] = None) -> torch.Tensor:
+           allocator: Optional[BaseAllocator] = None,
+           return_peers: Optional[bool] = None) -> Union[torch.Tensor, list[torch.Tensor]]:
     if allocator is not None:
         assert allocator.initialized(), "Allocator is not initialized"
         if device is not None:
             device = parse_device(device)
             assert allocator.device == device, f"Allocator device must be the " \
                 f"same as the device of the tensor, but got {allocator.device} != {device}"
-        return allocator._allocate_tensor(shape, dtype)
+        return allocator._allocate_tensor(shape, dtype, return_peers)
     else:
         return torch.empty(shape, dtype=dtype, device=device)
 
