[DOC] Added 2 examples of triton-distributrd for reference

python/tutorials/tle/test_tle_intra_node_allgather.py

@@ -0,0 +1,229 @@
+"""
+Intra-node AllGather with FlagTree TLE device remote pointers.
+This tutorial implements a single-node all-gather operator using
+FlagTree TLE.
+Run with a FlagTree environment, for example:
+
+    export FLAGCX_MEM_ENABLE=1
+    export FLAGCX_USE_HETERO_COMM=1
+    export FLAGCX_VMM_ENABLE=0
+    export FLAGCX_P2P_DISABLE=1
+    export CUDA_VISIBLE_DEVICES=0,1
+    # Optional: set FLAGCX_IB_HCA to the HCA list for your machine.
+    torchrun \
+    --nproc_per_node="${NPROC_PER_NODE}" \
+    --nnodes=1 \
+    --node_rank=0 \
+    --master_addr="${MASTER_ADDR}" \
+    --master_port="${MASTER_PORT}" \
+    "${SCRIPT_DIR}/test_tle_intra_node_allgather.py"
+
+If you explicitly disabled distributed support with USE_FLAGCX=0, USE_DIST=0,
+or USE_TLE_DIST=0, it might be necessary to reset these settings before running this tutorial.
+"""
+
+import os
+
+import torch
+import torch.distributed as dist
+import triton
+import triton.language as tl
+import triton.experimental.tle.language as tle
+
+
+@triton.jit
+def _all_gather_push_2d_kernel(
+    local_ptr,
+    ag_ptr,
+    ag_dev_mem,
+    dev_comm_dptr,  # DevComm handle, used to query the current rank within the kernel
+    mesh: tl.constexpr,
+    ELEM_PER_RANK: tl.constexpr,
+    BLOCK: tl.constexpr,
+):
+    peer = tl.program_id(0)
+    block_id = tl.program_id(1)
+    local_rank = tle.shard_id(mesh, "device", comm_ptr=dev_comm_dptr)
+    dst_base = local_rank * ELEM_PER_RANK
+    offsets = block_id * BLOCK + tl.arange(0, BLOCK)
+    mask = offsets < ELEM_PER_RANK
+    vals = tl.load(local_ptr + offsets, mask=mask, other=0.0)
+
+    if peer != local_rank:
+        dst_ptr = tle.remote(
+            ag_dev_mem,
+            shard_id=peer,
+            space="device",
+            dtype=ag_ptr.dtype.element_ty,
+            offset=dst_base,
+        )
+        tl.store(dst_ptr + offsets, vals, mask=mask)
+    else:
+        tl.store(ag_ptr + dst_base + offsets, vals, mask=mask)
+
+
+@triton.jit(do_not_specialize=["signal_target"])
+def _all_gather_signal_kernel(
+    signal_ptr,
+    signal_dev_mem,
+    dev_comm_dptr,
+    mesh: tl.constexpr,
+    signal_target,
+):
+    peer = tl.program_id(0)
+    local_rank = tle.shard_id(mesh, "device", comm_ptr=dev_comm_dptr)
+
+    if peer != local_rank:
+        remote_signal_ptr = tle.remote(
+            signal_dev_mem,
+            shard_id=peer,
+            space="device",
+            dtype=tl.int32,
+            offset=local_rank,
+        )
+        # Publish completion with system-scope release semantics. The receiver
+        # uses an acquire atomic before consuming the corresponding shard.
+        tl.atomic_xchg(
+            remote_signal_ptr,
+            signal_target,
+            sem="release",
+            scope="sys",
+        )
+    else:
+        tl.atomic_xchg(
+            signal_ptr + local_rank,
+            signal_target,
+            sem="release",
+            scope="sys",
+        )
+
+
+@triton.jit(do_not_specialize=["signal_target"])
+def _all_gather_wait_kernel(
+    signal_ptr,
+    local_rank,
+    signal_target,
+    WORLD_SIZE: tl.constexpr,
+):
+    """Wait until every remote shard in this rank's output is ready."""
+    peer = tl.program_id(0)
+    if peer < WORLD_SIZE and peer != local_rank:
+        # atomic_add(0) is an acquire load expressed with Triton's public
+        # atomic API. GE is required because a faster peer may already have
+        # published a later epoch.
+        observed = tl.atomic_add(
+            signal_ptr + peer,
+            0,
+            sem="acquire",
+            scope="sys",
+        )
+        while observed < signal_target:
+            observed = tl.atomic_add(
+                signal_ptr + peer,
+                0,
+                sem="acquire",
+                scope="sys",
+            )
+
+
+def _rank_print(rank: int, *items):
+    dist.barrier()
+    for cur_rank in range(dist.get_world_size()):
+        if cur_rank == rank:
+            print(*items, flush=True)
+        dist.barrier()
+
+
+def main():
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    torch.cuda.set_device(local_rank)
+
+    mem_pool = tle.get_mem_pool()
+    if mem_pool is None:
+        raise RuntimeError("FlagCX memory pool is unavailable; check FlagCX build and environment variables.")
+
+    rank = dist.get_rank()  # Obtain rank and world_size
+    world_size = dist.get_world_size()
+    local_world_size = int(os.getenv("LOCAL_WORLD_SIZE", str(world_size)))
+    assert world_size == local_world_size, "This tutorial is designed for a single node"
+
+    M = 8192
+    N = 12288
+    assert M % world_size == 0
+    m_per_rank = M // world_size
+    dtype = torch.float16
+    device = torch.device("cuda", local_rank)
+
+    local_data = torch.randn((m_per_rank, N), dtype=dtype, device=device)
+
+    with torch.cuda.use_mem_pool(mem_pool):
+        ag_buffer = torch.empty((M, N), dtype=dtype, device=device)
+        signal = torch.empty((world_size, ), dtype=torch.int32, device=device)
+
+    dev_comm_dptr, ag_dev_mem = tle.create_comm_tensor(ag_buffer)
+    _, signal_dev_mem = tle.create_comm_tensor(signal)
+    # ag_dev_mem is the device-side DevMem handle address created by FlagCX/TLE,
+    # signal_dev_mem is used to remotely write to the peer's signal[local_rank].
+    # dev_comm_dptr is used on the device side by tle.shard_id(..., comm_ptr=...) to query the current rank.
+
+    golden = torch.empty((M, N), dtype=dtype, device=device)
+    dist.all_gather_into_tensor(golden, local_data)
+
+    ag_buffer.fill_(-1)
+    signal.zero_()
+
+    torch.cuda.synchronize()
+    dist.barrier()
+
+    elem_per_rank = m_per_rank * N
+    block = 1024
+    num_blocks = triton.cdiv(elem_per_rank, block)
+
+    # 2D copy grid: split each peer transfer into independent chunks.
+    # The signal is written by a second kernel so it is ordered after all copy chunks in this stream.
+    copy_grid = (world_size, num_blocks)
+    signal_grid = (world_size, )
+    mesh = tle.device_mesh(tle.MeshConfig(device=world_size))
+    signal_target = 1
+
+    def launch_tle_all_gather():
+        _all_gather_push_2d_kernel[copy_grid](
+            local_data,
+            ag_buffer,
+            ag_dev_mem,
+            dev_comm_dptr,
+            mesh,
+            ELEM_PER_RANK=elem_per_rank,
+            BLOCK=block,
+            num_warps=4,
+        )
+        _all_gather_signal_kernel[signal_grid](
+            signal,
+            signal_dev_mem,
+            dev_comm_dptr,
+            mesh,
+            signal_target,
+            num_warps=4,
+        )
+        _all_gather_wait_kernel[signal_grid](
+            signal,
+            rank,
+            signal_target,
+            WORLD_SIZE=world_size,
+            num_warps=1,
+        )
+
+    launch_tle_all_gather()
+    torch.cuda.synchronize()
+    dist.barrier()
+
+    _rank_print(rank, f"Rank {rank} FlagTree Result:", ag_buffer)
+    _rank_print(rank, f"Rank {rank} FlagTree Signal:", signal)
+    assert torch.allclose(golden, ag_buffer, atol=1e-5, rtol=1e-5)
+    _rank_print(rank, f"Rank {rank} Pass!")
+
+    tle.cleanup_communicator()
+
+
+if __name__ == "__main__":
+    main()

python/tutorials/tle/test_tle_intra_node_allgather.sh

@@ -0,0 +1,30 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+
+export FLAGCX_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_6,mlx5_7,mlx5_8,mlx5_9
+export FLAGCX_MEM_ENABLE=1
+export FLAGCX_USE_HETERO_COMM=1
+export FLAGCX_VMM_ENABLE=0
+export FLAGCX_P2P_DISABLE=1
+export CUDA_VISIBLE_DEVICES=0,1,2,3
+
+
+if [[ "${CLEAR_TRITON_CACHE:-0}" == "1" ]]; then
+    rm -rf "${TRITON_CACHE_DIR:-${HOME}/.triton/cache}"
+fi
+
+NPROC_PER_NODE=4
+MASTER_ADDR=localhost
+MASTER_PORT=8333
+
+
+torchrun \
+    --nproc_per_node="${NPROC_PER_NODE}" \
+    --nnodes=1 \
+    --node_rank=0 \
+    --master_addr="${MASTER_ADDR}" \
+    --master_port="${MASTER_PORT}" \
+    "${SCRIPT_DIR}/test_tle_intra_node_allgather.py"