diff --git a/.gitignore b/.gitignore index 5a1b0ce57f..44e82485da 100644 --- a/.gitignore +++ b/.gitignore @@ -128,3 +128,6 @@ maint/host_checks/logs/* # uv Lockfile uv.lock + +# Private working notes: cluster hostnames, internal IPs, absolute paths. +CLAUDE.md diff --git a/cmake/FindNCCLGin.cmake b/cmake/FindNCCLGin.cmake new file mode 100644 index 0000000000..efd9be6097 --- /dev/null +++ b/cmake/FindNCCLGin.cmake @@ -0,0 +1,147 @@ +# Detect an NCCL installation that provides the Device API (GIN). +# +# GIN (GPU-Initiated Networking) is what TileScale's inter-node put/signal path +# compiles against. It is *not* present in every NCCL: the device headers first +# ship in 2.28.7, and a pip `nvidia-nccl-cu12` wheel may be several minor +# versions behind that. Rather than gate on the version macro alone, this module +# requires all three things that the device path actually needs: +# +# 1. nccl.h -- host API, and the version macros +# 2. nccl_device/gin.h -- the ncclGin device class +# 3. ncclDevCommCreate -- host-side device-comm bootstrap, in libnccl +# +# A tree can satisfy (1) and report a new enough version while missing (2)/(3), +# which is why the symbol check is not skipped when the header is found. +# +# Sets: NCCLGin_FOUND, NCCL_INCLUDE_DIR, NCCL_LIBRARY, NCCL_VERSION_STRING +# +# Hint with -DNCCL_ROOT=, or let it fall back to the active Python +# environment's nvidia/nccl wheel and then the CUDA toolkit prefix. + +set(_nccl_gin_min_version "2.28.7") + +# Candidate prefixes, most specific first. +set(_nccl_hints "") +if(NCCL_ROOT) + list(APPEND _nccl_hints "${NCCL_ROOT}") +endif() +if(DEFINED ENV{NCCL_ROOT}) + list(APPEND _nccl_hints "$ENV{NCCL_ROOT}") +endif() + +# pip-installed NCCL lives under /nvidia/nccl. Ask the +# interpreter rather than globbing, so we track the env actually in use. +if(Python3_EXECUTABLE OR Python_EXECUTABLE) + if(Python3_EXECUTABLE) + set(_nccl_py "${Python3_EXECUTABLE}") + else() + set(_nccl_py "${Python_EXECUTABLE}") + endif() + execute_process( + COMMAND "${_nccl_py}" -c + "import os,sysconfig;p=os.path.join(sysconfig.get_paths()['purelib'],'nvidia','nccl');print(p if os.path.isdir(p) else '')" + OUTPUT_VARIABLE _nccl_pip_dir + OUTPUT_STRIP_TRAILING_WHITESPACE + ERROR_QUIET) + if(_nccl_pip_dir) + list(APPEND _nccl_hints "${_nccl_pip_dir}") + endif() +endif() + +if(CUDAToolkit_LIBRARY_ROOT) + list(APPEND _nccl_hints "${CUDAToolkit_LIBRARY_ROOT}") +endif() + +find_path(NCCL_INCLUDE_DIR nccl.h + HINTS ${_nccl_hints} + PATH_SUFFIXES include + DOC "Directory containing nccl.h") + +# pip wheels ship only the versioned soname -- there is no libnccl.so +# development symlink -- so the bare "nccl" name that find_library derives is not +# enough. List the versioned file explicitly, and keep the unversioned name first +# so a real system/dev install still wins. +find_library(NCCL_LIBRARY + NAMES nccl libnccl.so.2 libnccl.so.2.dylib + HINTS ${_nccl_hints} + PATH_SUFFIXES lib lib64 + DOC "NCCL shared library") + +set(NCCL_VERSION_STRING "") +set(_nccl_has_gin_header FALSE) +set(_nccl_has_devcomm FALSE) + +if(NCCL_INCLUDE_DIR) + # Version macros. NCCL_VERSION_CODE is not usable here because it is a macro + # expression, so read the three components directly. + foreach(_part MAJOR MINOR PATCH) + file(STRINGS "${NCCL_INCLUDE_DIR}/nccl.h" _line + REGEX "^#define NCCL_${_part} +[0-9]+") + if(_line) + string(REGEX MATCH "[0-9]+" _nccl_${_part} "${_line}") + else() + set(_nccl_${_part} 0) + endif() + endforeach() + set(NCCL_VERSION_STRING "${_nccl_MAJOR}.${_nccl_MINOR}.${_nccl_PATCH}") + + if(EXISTS "${NCCL_INCLUDE_DIR}/nccl_device/gin.h") + set(_nccl_has_gin_header TRUE) + endif() +endif() + +# ncclDevCommCreate is the host entry point the GIN path needs; a tree can carry +# the header and still be linked against a runtime that does not export it. +if(NCCL_LIBRARY AND _nccl_has_gin_header) + if(UNIX AND NOT APPLE) + find_program(_nccl_nm NAMES nm) + if(_nccl_nm) + execute_process( + COMMAND "${_nccl_nm}" -D --defined-only "${NCCL_LIBRARY}" + OUTPUT_VARIABLE _nccl_syms ERROR_QUIET) + if(_nccl_syms MATCHES "ncclDevCommCreate") + set(_nccl_has_devcomm TRUE) + endif() + else() + # No nm available: trust the header plus version gate instead of + # silently disabling GIN on a stripped-down build image. + set(_nccl_has_devcomm TRUE) + endif() + else() + set(_nccl_has_devcomm TRUE) + endif() +endif() + +set(NCCLGin_FOUND FALSE) +if(NCCL_INCLUDE_DIR AND NCCL_LIBRARY AND _nccl_has_gin_header + AND _nccl_has_devcomm + AND NOT NCCL_VERSION_STRING VERSION_LESS _nccl_gin_min_version) + set(NCCLGin_FOUND TRUE) +endif() + +if(NCCLGin_FOUND) + message(STATUS "NCCL GIN: enabled (NCCL ${NCCL_VERSION_STRING} at ${NCCL_INCLUDE_DIR})") +elseif(NCCL_INCLUDE_DIR) + # Found NCCL but cannot use the device path. Say which check failed -- + # "GIN disabled" with no reason is the hard case to debug on a cluster. + # Order matters: the symbol check is skipped when the library is missing, so + # test for the library before blaming its exports. + if(NOT NCCL_LIBRARY) + set(_why "found headers at ${NCCL_INCLUDE_DIR} but no libnccl alongside them") + elseif(NOT _nccl_has_gin_header) + set(_why "no nccl_device/gin.h (needs NCCL >= ${_nccl_gin_min_version})") + elseif(NOT _nccl_has_devcomm) + set(_why "libnccl does not export ncclDevCommCreate") + elseif(NCCL_VERSION_STRING VERSION_LESS _nccl_gin_min_version) + set(_why "NCCL ${NCCL_VERSION_STRING} < ${_nccl_gin_min_version}") + else() + set(_why "incomplete installation") + endif() + message(STATUS "NCCL GIN: disabled -- ${_why}. " + "Inter-node kernels will fall back to intra-node paths.") +else() + message(STATUS "NCCL GIN: disabled -- NCCL not found. " + "Set -DNCCL_ROOT= to enable inter-node support.") +endif() + +mark_as_advanced(NCCL_INCLUDE_DIR NCCL_LIBRARY) diff --git a/docs/distributed_api_reference.md b/docs/distributed_api_reference.md index 8694ff2e6f..833b9acb09 100644 --- a/docs/distributed_api_reference.md +++ b/docs/distributed_api_reference.md @@ -1,10 +1,14 @@ # TileScale Distributed API Reference -TileScale extends TileLang with experimental single-node, multi-GPU CUDA -primitives. The supported process model is one process per local GPU on one -host, using an NCCL process group for host-side coordination. Multi-node -execution, NVSHMEM, and a general-purpose distributed runtime are outside the -current scope. +TileScale extends TileLang with experimental multi-GPU CUDA primitives. The +supported process model is one process per local GPU, using an NCCL process +group for host-side coordination. + +Multi-node execution is supported through the NCCL Device API ("GIN", +GPU-Initiated Networking), which exposes one-sided RDMA callable from inside a +kernel. Intra-node peers keep using IPC or VMM peer pointers; inter-node peers +are reached through a registered NCCL window instead. NVSHMEM and a +general-purpose distributed runtime remain outside the current scope. The Python distribution is `tilescale`, while the import namespace remains `tilelang`. Install the optional host dependencies with @@ -18,12 +22,18 @@ The available memory path depends on the host and visible GPUs: |------|--------------|-----------| | CUDA IPC | Same host, CUDA peer access between participating GPUs | Used when VMM is disabled or unavailable | | VMM fabric | CUDA driver API 12.4 or newer, fabric-handle support on every visible GPU, and an accessible NVIDIA IMEX channel | Auto-selected for distributed allocators when the runtime probe succeeds | -| Multicast | VMM fabric requirements plus multicast-capable GPUs and fabric, normally an NVSwitch system | Enabled only when `mcast_size` is explicitly requested | +| Multicast | A VMM allocator plus multicast-capable GPUs, normally an NVSwitch system. Fabric handles are *not* required: without an IMEX channel the multicast object is shared as a POSIX file descriptor | Enabled only when `mcast_size` is explicitly requested | +| GIN (inter-node) | NCCL 2.28.7 or newer with `nccl_device/gin.h` and `ncclDevCommCreate`, a VMM-backed arena, and a working RDMA fabric | Attempted automatically when more than one node is detected; `TILESCALE_USE_GIN=1` makes an unavailable Device API a hard error | `_supports_vmm_fabric()` and `_supports_multicast()` are runtime probes, not portable feature guarantees. Both inspect all CUDA-visible devices, so set `CUDA_VISIBLE_DEVICES` to the exact local rank set before starting processes. +Both probes reach `cuCtxGetDevice` and therefore report `False` with no current +CUDA context. `cuInit` plus a primary-context retain is not sufficient; force a +context first (for example `torch.zeros(1, device="cuda")`) or a capable machine +will be misreported as incapable. + Fabric handles require an NVIDIA IMEX channel that is accessible inside the host or container. On a compatible Linux driver installation, an administrator can run: @@ -48,10 +58,25 @@ rank, world_size, group = init_dist( num_local_ranks, master_port=None, ) + +rank, world_size, group, node_info = init_dist( + local_rank, + num_local_ranks, + return_node_info=True, +) ``` `init_dist` sets `cuda:local_rank`, creates an NCCL process group, and returns -its rank, size, and `dist.group.WORLD`. +its rank, size, and `dist.group.WORLD`. With `return_node_info=True` it also +returns a `NodeTopology` describing `num_nodes`, `node_rank`, `local_rank` and +`local_world_size`, which the allocator needs to decide whether to set up GIN. +Pass it on as `get_allocator(..., node_info=node_info)`. + +Topology is read from torchrun-style `LOCAL_WORLD_SIZE`/`GROUP_RANK` when those +are present, and otherwise from `NNODES`/`NODE_RANK`. + +`init_dist` sets `NCCL_IB_DISABLE=1` by default. **Clear it for any inter-node +run**, or NCCL will refuse to use the RDMA fabric that GIN depends on. The current implementation requires a contiguous rank-to-device mapping: process rank `i` uses CUDA ordinal `i`, and participating devices must appear @@ -88,6 +113,25 @@ set the variable. Port precedence is: `master_port` argument, `TILESCALE_MASTER_PORT`, `MASTER_PORT`, then `8361`. +Multi-node launches additionally need `WORLD_SIZE`, `LOCAL_WORLD_SIZE`, and +either `NNODES`/`NODE_RANK` or torchrun's `GROUP_RANK`, with `MASTER_ADDR` set to +an address the other nodes can reach. `NCCL_IB_DISABLE` must be cleared. + +Variables read outside `init_dist`: + +| Variable | Effective default | Behavior | +|----------|-------------------|----------| +| `TILESCALE_USE_VMM` | unset | `1` forces VMM, any other value forces CUDA IPC | +| `TILESCALE_USE_GIN` | unset | `1` requests arena window registration even on a single node, and makes an unavailable Device API a hard error rather than a warning | +| `TILESCALE_NCCL_LIB` | unset | Path to a GIN-capable `libnccl.so.2`, for when the ambient NCCL predates the Device API | +| `TILESCALE_GIN_CONTEXTS` | `8` | Requested `ginContextCount`. One context is one QP per peer; the count is a hint and may be granted in part | +| `TILESCALE_GIN_SIGNALS` | `32` | Requested `ginSignalCount`, i.e. how many independent signals a kernel may use | +| `TILESCALE_GIN_COUNTERS` | `32` | Requested `ginCounterCount` | + +The GIN resource counts affect device memory but not start-up latency: +`ncclDevCommCreate` measures the same at 1, 4 and 8 contexts, because its cost is +transport initialisation rather than per-context setup. + ## Distributed Allocator ```python @@ -129,6 +173,14 @@ The selection order is: Forcing VMM does not provide an IPC fallback if fabric allocation fails. +GIN requires a VMM allocator: `ncclCommWindowRegister` rejects `cudaMalloc` +memory, so window registration fails with the IPC backend. When fabric handles +are unavailable the arena is still VMM-backed, created with +`CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR` and shared between local ranks by +duplicating that descriptor with `pidfd_getfd`. That path needs ptrace-level +access to sibling ranks (true for one job under one user) and a kernel with +`pidfd_getfd`, i.e. Linux 5.6 or newer. + ### Multicast Allocation Passing `mcast_size` requires a distributed VMM allocator and a successful @@ -138,6 +190,22 @@ multicast capability probe. The internal instructions; each rank writes its own contribution through the local physical view. +The multicast buffer is a **separate allocation from the arena**, and only the +arena is registered as an NCCL window. Anything a GIN put reads or writes must +therefore come from `tilelang.tensor(..., allocator=...)`, not from +`_allocate_mcast_tensor`. A kernel that reduces through the switch and then +sends over the fabric needs one buffer of each. + +`mcast_size` is the total multicast capacity and is consumed by a bump pointer +with no free, so size it for every multicast tensor the process will allocate. +A caller that allocates two (an allreduce reducing its input and broadcasting +its output, for instance) must request both up front. + +Sharing the multicast object across processes uses fabric handles when an IMEX +channel is available and POSIX file descriptors otherwise; `_multicast_uses_fd()` +reports which route the C++ side selected, so the allocator does not re-probe +and risk disagreeing with it. + `allocator.close()` must be called collectively before destroying the process group, especially for multicast allocations. The allocator is also a context manager. @@ -239,6 +307,46 @@ GPU-local and cross-block barrier helpers include `T.init_barrier_gpu`, `T.barrier_blocks`, and `T.sync_blocks`. Cross-GPU correctness still depends on using peer-visible storage and the appropriate system-scope ordering. +### GIN Inter-Node Operations + +These lower to the NCCL Device API and are the only kernel-side path to a peer +on another node. They live under `T.nccl_gin.*` and require the allocator to +have registered its arena as an NCCL window. + +| API | Purpose | +|-----|---------| +| `T.nccl_gin.put(src, dst, size, peer, scope="block")` | One-sided RDMA write into `peer`'s symmetric buffer | +| `T.nccl_gin.put_signal(src, dst, size, peer, signal_id, scope="block")` | Same, plus increment `signal_id` on the destination once the payload has landed | +| `T.nccl_gin.signal(peer, signal_id, scope="block")` | Increment a remote signal with no payload | +| `T.nccl_gin.wait_signal(least, signal_id, scope="block")` | Block until the cumulative count for `signal_id` reaches `least` | +| `T.nccl_gin.flush(scope="block")` | Wait until this rank's source buffers are reusable | + +`src` and `dst` are ordinary buffer element references; the lowering converts +them to `(window, offset)` pairs using the symmetric arena, so the offset that +names bytes locally names the same bytes on the peer. `peer` is a **global** +rank. + +Three properties of the signal mechanism decide how these are used: + +- **Signals are cumulative and a wait does not consume them.** A compile-time + `least` is therefore satisfied on every launch after the first, which silently + turns the wait into a no-op. Pass the target as a kernel argument that the host + advances per launch. +- **Signal state is per context.** A put issued on sender context *i* increments + the receiver's signal through context *i*, so a CTA sees only `1/contexts` of + the arrivals. `wait_signal` divides `least` by the device-side + `context_span()` for this reason. A wait grid narrower than the sender's rounds + the target down, to zero in the worst case; a wider one parks CTAs on contexts + nothing signals. +- **The requested context count is a hint.** `ncclDevCommRequirements.ginContextCount` + may be granted in part, so the divisor must come from the device rather than + from the host's request. + +Do not wait on a signal from inside a large-grid kernel. A waiting CTA occupies +an SM, so if the CTAs that would issue the matching puts are still queued behind +it, nothing progresses. Use separate streams, or confine the wait to a +small-grid kernel. + ### Multimem Operations Multimem operations require a valid multicast allocation and compatible @@ -258,6 +366,20 @@ for a plain `multimem_tma_store`. The signal type is inferred from `addr`; there is no `dtype_tag` argument. The direct `multimem_ld_reduce` and `multimem_red` paths currently expose ADD; unsupported PTX dtype/operation combinations are rejected during lowering. + +Two constraints shape how these are called in practice: + +- **A float16/bfloat16 region must be exactly one contiguous pair per thread**, + because those dtypes lower to packed x2 instructions. In tile terms the region + width must equal `2 * threads`; any other width lets a neighbouring `T.copy` + infer a wider vectorisation, and layout inference then fails with *"requires + the local fragment layout to preserve canonical pair ownership"*. Work per + thread therefore has to come from looping over tiles, not from widening one. +- **The multicast region must be provably in bounds at compile time.** A runtime + offset into it is rejected with *"multimem packed multicast region must be + provably in bounds or use a tile-aligned all-or-none dynamic partition"*, so a + kernel that publishes a varying slice needs the offset as a compile-time + constant, i.e. one specialisation per slice. Direct multimem operations and signals require SM90+ and CUDA Toolkit 12.1+ (PTX 8.1+). Packed `float16`/`bfloat16` load-reduce additionally requires CUDA Toolkit 12.2+ because its `.acc::f32` form was introduced in PTX 8.2. Packed @@ -339,6 +461,16 @@ releases. Every registered FFI name below is prefixed with | `_sync_vmm_handles(rank, device_ids, buffer_ptrs_gpu_addr, handles) -> None` | `sync_vmm_handles(rank: int64, num_ranks: int64, buffer_ptrs_gpu_addr: int64, packed_handles: Bytes) -> void` | | `_supports_vmm_fabric() -> bool` | `supports_vmm_fabric() -> bool` | | `_supports_multicast() -> bool` | `supports_multicast() -> bool` | +| `_create_vmm_fd_handle(ptr) -> bytes` | `create_vmm_fd_handle(ptr: int64) -> Bytes` | +| `_open_vmm_fd_handle(handle) -> int` | `open_vmm_fd_handle(handle: Bytes) -> int64` | + +`create_vmm_fd_handle` exports the allocation as a POSIX file descriptor and +returns `size | pid | fd`; `open_vmm_fd_handle` duplicates that descriptor into +the caller with `pidfd_open` plus `pidfd_getfd`, then imports and maps it. This +is the route used when fabric handles are unavailable, and it avoids the +unix-socket rendezvous an `SCM_RIGHTS` exchange would need. The exporting +process keeps its descriptor open for its lifetime, because a peer may import at +any point before teardown. The Python sync wrappers pack one handle per rank and derive `num_ranks` from `len(device_ids)`. `buffer_ptrs_gpu_addr` is the integer device address of a @@ -358,6 +490,13 @@ currently ignored. | `_mc_release_handle(handle) -> None` | `(int64) -> void` | | `_mc_unmap(ptr, size, num_devices) -> None` | `(int64, int64, int64) -> void` | | `_mc_get_aligned_size(size, num_devices) -> int` | `(int64, int64) -> int64` | +| `_mc_export_fd_handle(handle) -> bytes` | `(int64) -> Bytes` | +| `_mc_open_fd_handle(handle_bytes) -> int` | `(Bytes) -> int64` | +| `_multicast_uses_fd() -> bool` | `multicast_uses_fd() -> bool` | + +`mc_create` requests whichever handle type `multicast_uses_fd()` reports, and the +matching export/import pair must be used with it. The fd pair carries `pid | fd` +and is duplicated with `pidfd_getfd`, as for the arena. ### Tensor From Pointer @@ -390,3 +529,19 @@ The maintained examples and tests are under `examples/distributed` and `testing/python/distributed`. This reference intentionally does not label every example as universally working because results depend on GPU architecture, topology, CUDA toolkit, driver, and IMEX configuration. + +Two inter-node limitations are known and will be met in normal use rather than +at the edges: + +- **Some `put_signal` transfer sizes fail to lower**, with + `Can't fetch the lanes of a scalable vector at a compile time`. The trigger is + the element count, not the byte count, and the failing set is not contiguous. + Callers that derive a transfer size from a buffer length should be prepared to + adjust it; halving the chunk count doubles the transfer and generally moves off + a bad value. +- **`T.barrier_blocks` is single-node** despite documenting a rendezvous across + "every rank". It is lowered with the global rank and world size, and the device + side takes `get_remote_base_ptr` of each participant, which returns 0 for a peer + it considers inter-node. In a multi-node job the inter-node slots become null + system atomics. Inter-node ordering should come from GIN signals, which need no + barrier; a node-local barrier variant does not exist yet. diff --git a/docs/distributed_autotune_design.md b/docs/distributed_autotune_design.md new file mode 100644 index 0000000000..add48c4066 --- /dev/null +++ b/docs/distributed_autotune_design.md @@ -0,0 +1,305 @@ +# Multi-process autotuning for the 2D inter-node collectives + +Design notes, not an implementation. The subject is the tuning surface of +`examples/distributed/internode/internode_2d.py` — `Allgather2D`, `ReduceScatter2D`, +`Allreduce2D` and the two fused GEMM wrappers — at 16 GPUs over two nodes. The question is +what a tuner for them has to look like, given that every candidate evaluation is a 16-rank +collective and that the measurement is noisy enough to lie. Everything numeric below is +quoted from measurements already recorded in `CLAUDE.md` and the module docstrings, or is +arithmetic over per-unit costs recorded there; nothing here is a new measurement. + +## 1. Why the hand-tuned constants have to go + +The knobs are not independent of each other, of the collective, or of the size. + +**`--mc-tiles` splits the collectives into two families.** Allgather wants 32 tiles per +multicast CTA at every size measured: 0.164 ms against 0.233 at 4 tiles for 48 MiB, 0.501 +against 0.585 for 240 MiB. Allreduce wants the *opposite* at small sizes — 0.282 ms at 4 +tiles against 0.361 at 32 for 48 MiB — then agrees at large ones, 0.944 at 32 against 1.086 +at 4 for 240 MiB. Reduce-scatter behaves like allreduce at 48 MiB (0.138 ms at the +size-scaled value against 0.194 at 32) and like allgather at 120 and 240. The explanation in +`internode_2d` is convincing — allgather's intra half is two publishes and no reduce, so it +is switch-bound and wants few fat CTAs, while the reduce-carrying collectives are +occupancy-bound once the shard is small — but it explains a *shape*, not a formula, and the +crossover is not derivable from it. + +**`--rail-groups` and `--gin-contexts` are one knob, not two.** Allgather is fastest at 2 +groups on 4 contexts (0.500 ms) and slower at 8 groups on 1 context (0.561); allreduce +inverts it, 0.944 ms at 8 groups against 1.015 at 2, because it has two fabric hops to +overlap. They are also structurally coupled: a group's grid must be a whole number of chunks +*and* a multiple of the context count, because `wait_signal` divides the grid-wide target by +the granted context count and a narrower grid rounds it down. So "more groups" is often only +purchasable by spending contexts, and that trade is collective-specific. + +**Every closed form tried regressed something.** `pick_mc_tiles` scales one tile per 480 KB +of shard; refining it to 240 KB regressed allreduce at 120 MiB, 0.747 ms against 0.521. +`MIN_GROUP_BYTES = 1_500_000` is the same kind of artefact — allreduce at 8 groups is 0.95 ms +on a 15.7 MB shard and 0.61 ms on a 3.1 MB shard where the fabric alone needs 0.13, i.e. half +a millisecond of pure launch overhead. Both are fits to three sizes on one cluster. + +**And the noise would let a naive sweep enshrine itself.** Allreduce at 240 MiB, one config, +six runs: 0.944, 0.951, 1.015, 1.028, 1.034, 1.049 ms — about 10% spread. The torch baseline +drifts as much or more over the same window (240 MiB: 0.764, 0.783, 0.993; 120 MiB: 0.559, +0.622, 0.710), because the cluster is shared and other tenants use the same NICs. Several +differences quoted above — 0.944 against 1.015, say — are barely outside that band from a +single pair of runs. + +Three consequences, and they are the constraints that matter most: + +1. **Repeat and take a robust statistic.** Take the **minimum** across repeats of the + per-candidate summary, and within one `do_bench` call the **median** of reps rather than + the mean it currently returns. Minimum across repeats because the noise is one-sided + contention: nothing makes a collective faster than an uncontended fabric allows, so the + fastest observation best estimates the config's own cost while a mean mixes in other + tenants' traffic. Median within a call because one rep can be perturbed by a stray + barrier. The counter-argument is real — if a config is *intermittently* bad the minimum + hides it — so the spread must be reported alongside, never discarded. +2. **Re-time the baseline next to each candidate**, not once per sweep. `bench_vs_torch` + already times torch before *and* after each run for this reason and `report_tuning` + warns above 15% drift. Compare ratios, not milliseconds: over those six allreduce runs + the millisecond column moved 11% while the ratio stayed inside 0.96–1.01x. +3. **Treat any decision under ~5% as undecided** and break ties toward the simpler config + (fewer groups, fewer launches). A tuner that reports a winner without a tie band gives a + different answer every run, which is worse than a constant. + +## 2. Why neither existing tuner transfers + +### 2.1 The `--tune` helpers in `internode_common.py` + +`tune_grid` / `report_tuning` / `per_launch_signals` are a real in-process sweep and they +got the flat collectives tuned, so they are the right ancestor. But they are wired only +into the three *flat* examples; the 2D examples inherit `--tune`, `--tune-chunks` and +`--tune-contexts` from `add_common_args` and never read them, so those flags are silently +inert on the 2D path today. Lifting them is not just a matter of calling them. + +**The grid is the wrong shape.** `tune_grid` knows `chunks`, `gin_contexts`, `threads`. The +2D surface adds `rail_groups`, `mc_tiles`, `mc_threads`, `intra_chunks`, and its single +validity rule (`chunks % contexts == 0`) is one of six. + +**The signal budget does not survive.** `tune_grid` gives each candidate virgin signal slots, +which is correct and is the only reason sweeping contexts in one process is safe: signals are +cumulative, nothing resets them, and the device divides the accumulated total by its *current* +context span, so a slot reused under a different span yields a wrong target from then on. But +a 2D candidate needs one signal per rail group, and a composed allreduce one per group in each +half — at `--rail-groups 8` that is 16 of the 32 provisioned signals, so virgin slots allow +**two candidates per process** against a 19 s bootstrap. Either (i) the driver carries a +running total per slot forward and only reuses a slot for candidates with the same +`gin_contexts`, or (ii) `TILESCALE_GIN_SIGNALS` is raised (the override exists in +`nccl_window.py`, default 32). I could not determine whether a larger request is granted: +contexts are documented as a hint and clamped (8 asked, 4 granted) and the device exposes +`probe.nContexts`, but nothing reports a granted *signal* count. Option (i) works today. + +**Buffers cannot be reallocated per candidate.** Each collective allocates its arena tensors +in `__init__`, and `BaseAllocator._allocate_tensor_locked` is a bump allocator with no free — +only `close()` releases anything. Worse, the multicast buffer is sized *exactly* at +`Context(mcast_bytes=numel*itemsize)`, so a second `ctx.mcast_tensor()` fails outright. A +driver must allocate once and rebuild only kernels between candidates, which means separating +buffer allocation from kernel construction in `_Base` — small, but required. + +**`args` is mutated, and the effective config differs from the requested one.** +`_Base.__init__` does `args.mc_tiles = pick_mc_tiles(...)`, freezing the auto-scaling after +the first candidate, so each candidate needs a copy. And three things rewrite the config +during construction — `workable_chunks()` halves `chunks`, `MIN_GROUP_BYTES` backs +`rail_groups` off, `pick_mc_tiles` fills `mc_tiles` in — so the driver must record what the +object actually built, or the report attributes a time to a config that never ran and two +"different" candidates are silently the same one measured twice. + +### 2.2 TileLang's `AutoTuner` (`tilelang/autotuner/tuner.py`) + +A well-built single-process tuner in which essentially every mechanism is hostile to a +collective evaluation. It compiles on a `ThreadPoolExecutor` and benchmarks from worker +threads, optionally across several local GPUs, so candidate order is nondeterministic; it +enforces a per-candidate timeout via `SIGALRM` or an async exception injected into the +benchmark thread and on timeout or error **skips that candidate**; it picks a winner from +rank-local latencies; and it caches to disk under a key derived from the function source and +the config list. Each of those is a deadlock at 16 ranks: + +- **Every evaluation is collective at three levels.** `ctx.compile` passes + `compile_once=True, compile_group=ctx.group`, and `_maybe_compile_once` does a + `dist.all_gather_object` **per compile** — even on a `KernelCache` memory-cache hit. + `do_bench` barriers before and after and, with `barrier_comm_profiling` on, does a + `torch.cuda._sleep` plus a `dist.all_reduce` **per rep**. `check` does an `all_reduce`. + So ranks must agree on the candidate list, the order, how many kernels each candidate + compiles, and the rep count. A rank that skips, reorders or times out and moves on does + not lose performance — it desynchronises the collective stream and the job hangs. +- **Rank-local winners are incoherent, not merely suboptimal.** A config is a property of + the group: `chunks`, `gin_contexts` and `rail_groups` all appear in the sender's grid + *and* the receiver's wait target. Two ranks on different configs wait for signal totals + nobody will post. +- **Timeouts cannot be per-rank.** An async exception in one rank's benchmark thread leaves + 15 peers in a barrier. The only safe deadline is an agreed one: decided up front, + uniformly, aborting the sweep rather than a candidate. +- **Its cache key is wrong here** — function source and config list, not world size, node + shape or fabric. See §3(c). + +Extending it would also put a distributed policy into core TileLang, which the project's +own layering rule keeps out (distributed work lives in `tilelang/distributed/` and the +examples tree so the tree stays rebaseable on upstream). + +One correction to `workable_chunks`'s comment, which overstates the symmetry of the compile +hazard. In `_maybe_compile_once` the compile root catches its own exception and ships the +traceback through the `all_gather_object`, so every rank raises a `RuntimeError` carrying +it — a lowering failure **on the root** fails the job cleanly rather than hanging. The +dangerous case is asymmetric: a rank raising *outside* that window (a non-root rank's own +`cached()` call after the gather, an `expect=` assertion, divergent control flow before it) +leaves peers blocked. Not academic here — `run_internode.sh` deliberately uses a different +interpreter and NCCL directory on than on the other nodes, so "the same source lowers +the same way on every rank" is an assumption about two toolchains. Hence the rule: **probe +locally with a plain `tilelang.compile`, agree on the verdict collectively, then compile.** + +### 2.3 Configs that are silently wrong, not slow + +Why every candidate needs verifying rather than just timing. All these failure modes are +"fast and plausible": + +- A rail grid narrower than the granted context count rounds the wait target down — to zero + in the worst case — and the wait becomes a no-op. This is the mechanism behind the + historical 332 GB/s reading that exceeded the PCIe egress ceiling. +- Overlapping signal *ranges* (not merely reusing an id) let one collective's wait be + satisfied by another's arrivals; this corrupted exactly the second half of the allreduce + output when the allgather started at `SIGNAL_PHASE2` instead of + `SIGNAL_DATA + rs.signals_used`. +- A reused slot whose device total is stale while the host's running total restarts at zero + satisfies the first wait instantly — exactly the hazard slot reuse (§2.1) introduces if + the totals are not carried. +- Two grid computations truncate rather than raise. `mc_bcast_kernel` uses + `ctas = (span // block_N) // tiles_per_cta` with `span = shard/groups`, while `_Base` + only checks `shard_numel % (2*mc_threads*mc_tiles)` — the *unsliced* shard — so with + `rail_groups > 1` a candidate can publish less than its span, or get `ctas == 0` and + publish nothing. Similarly `rs_sum_kernel`'s chunk width is + `span // (intra_chunks // groups)`, which drops a tail unless `intra_chunks` divides the + shard and `groups` divides `intra_chunks`; neither is checked on the multimem path. Both + become pruning rules below, and arguably should also become assertions. + +A check is cheap next to a timing run — one launch, one torch reference, one `all_reduce` — +so there is no reason to make it optional. + +### 2.4 Bootstrap forces one process for the whole sweep + +A rank costs ~19 s to bring up: 5.7 s `init_dist`, 12 s allocator of which 6.4 s is +`ncclDevCommCreate`. That 6.4 s is flat against the requested resource counts (6.53 / 6.90 / +6.36 s at 8 / 4 / 1 contexts), so it is DOCA/IBGDA setup inside NCCL and not reducible from +here. Fork-per-candidate pays it per candidate; in-process pays it once. That settles the +shape: **one process lifetime, many candidates, all ranks in lockstep.** Two corollaries: +`TL_PG_TIMEOUT_SEC` (180 in `run_internode.sh`) and `RANK_TIMEOUT` (900) both need raising, +since 15 ranks sit inside a collective while one compiles; and the JIT cache must be warm, +since 22 kernels cost 1.5 s as cache hits (~0.08 s each) and "minutes" cold, and one cold +compile past the group timeout aborts the run. + +## 3. Design options + +**(a) Extend `tilelang.autotuner` with a collective evaluation hook** — deterministic +iteration, a group barrier around evaluation, rank 0 decides and broadcasts; in exchange the +sweep inherits progress bars, the result dataclass and the disk cache. But the pieces that +must be disabled (thread pool, multi-GPU workers, per-candidate timeout, skip-on-error, +rank-local winner, cache key) are most of the class, and what survives is a loop and a table. +Worth revisiting *after* (b) exists and the policy has stopped moving, as an upstreamable +"collective evaluation" strategy. + +**(b) A lockstep sweep driver in `internode_common.py`** — a `sweep()` that all ranks call +and walk identically. Rank 0 builds the candidate list and broadcasts it with +`dist.broadcast_object_list`; every rank then iterates it in the same order. Per candidate: +probe validity locally with plain `tilelang.compile`, `all_reduce` the verdict so the skip +is unanimous, build kernels against pre-allocated buffers, verify with `check()`, time with +`do_bench` sandwiched between two baseline timings. Rank 0 accumulates the table and +broadcasts the winner. This is the natural extension of what already works, every +constraint in §2 maps onto an explicit line of it rather than a disabled feature of +something else, and it can be debugged under `run_2d_proxy.sh` — useless for timings, but +it catches every correctness and lockstep bug a real two-node run can. + +**(c) An offline cache consulted at startup** — persist the winner as JSON keyed by problem +and topology, and have `_Base.__init__` consult it instead of +`pick_mc_tiles`/`MIN_GROUP_BYTES` when an entry matches. This is the actual goal; (b) is the +machine that fills it. The key must carry +`(collective, algo, intra mode, world_size, local_world_size, num_nodes, dtype, numel)` plus +a fabric fingerprint — the granted GIN context count (4 here, not the 8 requested), the NCCL +version, and whether multicast came up on fabric or POSIX-FD handles. Node identity matters +too, since a NIC shared with another tenant changes not just the numbers but which knob +appears to matter; I would store hostnames as metadata and warn on them rather than key on +them, since keying makes every entry single-use. Staleness must be detectable, and the repo +documents why to worry: `KernelCache._generate_key` does not hash +`src/tl_templates/cuda/distributed/*.h`, so editing a device template leaves cached binaries +stale under an identical key. A tuning cache inherits that and adds one — a config that won +against one `internode_2d.py` need not win against the next. So an entry should carry the +tilelang version, a hash of `internode_2d.py`, a hash of the device templates, the date, and +the ms *and* baseline ms it won with: hash mismatch is a hard invalidation, and a large +divergence between the recorded ratio and the observed one is a warning that the fit is stale. + +**Recommendation: build (b), then (c) on top of it, and leave (a) for later.** In order: +separate buffer allocation from kernel construction in `_Base`; add driver-owned per-slot +signal running totals; write `sweep()` next to `tune_grid` and wire `--tune` into the three +2D examples; debug under `run_2d_proxy.sh`; then add the JSON cache and the `_Base` lookup. +The first four are the day of work; (c) is a few hours more and is where the value lands, +because it is what removes the constants from the source. + +## 4. Practicalities + +**Search space.** `--chunks` over {4, 8} (16 and above hit the put-size lowering bug on the +2D path, and `workable_chunks` halves down anyway). `--gin-contexts` over {1, 2, 4}, since +4 is what the devcomm grants. `--rail-groups` over {1, 2, 4, 8}. `--mc-threads` × +`--mc-tiles` over {256, 512} × {4, 8, 16, 32} **as a pair, never as two independent axes**: +the climb is non-monotonic at low thread counts (at 256 threads, 8 tiles gives 337 GB/s and +16 gives 317, while 512/32 gives 397). `--intra-chunks` over {512, 1024, 2048} — it looks +pull-path-only but `ReduceScatter2D` and `Allreduce2D` use it for their sum-kernel grids on +the multimem path too. `--no-overlap` is a diagnostic, not a candidate. + +**Pruned by construction, before any compile:** `chunks % groups == 0` and +`(chunks // groups) % gin_contexts == 0`; `shard_numel % (2*mc_threads*mc_tiles*groups) == 0` +with `(span // block_N) // mc_tiles >= 1` (the span-level rules from §2.3, stronger than +what `_Base` checks); `shard_numel % intra_chunks == 0` and `intra_chunks % groups == 0`. +The put size must lower, which can only be probed, not predicted — the bad set is not +contiguous (30720 lowers, 32768–61440 fail, 65536 and 81920 lower, 98304 fails) — and the +verdict depends only on `(shard_numel, chunks)`, so it is shared across the other knobs and +worth caching in-process. `MIN_GROUP_BYTES` should be a *default, not a filter*: it is one +of the fits the tuner exists to replace, so a sweep that enforces it can never discover it +is wrong. Let small slices in and let the measurement reject them. + +**Sweep cost.** At 240 MiB bf16 on 16 ranks (shard 7 864 320 elements) the divisibility +rules leave 15 valid `(chunks, groups, contexts)` triples, all 8 `(mc_threads, mc_tiles)` +pairs and all 3 `intra_chunks` values — 360 combinations, too many for a repeated protocol. +So stage it: the 15 fabric triples at default `mc_*`, then the 8 multicast pairs at the +winning triple, then the 3 `intra_chunks`, then re-confirm the top few of stage 1 at the +winning `mc_*`. About 26 candidates plus a confirmation pass. This is coordinate descent and +can miss a genuine interaction; the confirmation pass is the cheap partial defence, and the +full 360 stays available for an occasional overnight run. Per candidate, warm cache: the +allgather half compiles roughly `4G+1` kernels (9 at 2 groups, 33 at 8) and the +reduce-scatter half `3G+2` (8 at 2, 26 at 8), so a composed allreduce is up to ~59 — at +~0.08 s per cache hit, 0.7–4.7 s of compile, each still paying one `all_gather_object`. +Timing is three `do_bench` calls (baseline, candidate, baseline) at warmup 20 / rep 50, and +each rep carries a `torch.cuda._sleep(2e7)` cycles, a 256 MB L2 flush and a barrier +`all_reduce`, so a call is dominated by its own overhead rather than the ~1 ms collective — +order half a second to a second each. Call it 5–10 s per candidate: the staged sweep is +3–5 min per repeat, 10–15 at three repeats, plus the 19 s bootstrap; three collectives at +three sizes is an hour or two, a held-nodes window rather than a research project. The cost +is dominated by measurement repetition rather than compilation, which is where it belongs. + +**Measurement protocol.** Per candidate: verify, time the baseline, time the candidate, time +the baseline again; report the ratio against `min(pre, post)` and the drift between them. +Repeat the whole candidate loop *R* times (default 3) as an outer pass rather than repeating +each candidate three times in a row, so a bad ten minutes on the fabric penalises all +candidates rather than whichever three landed in it. Summarise by the minimum of the +per-pass ratios and report the spread. Discard the *pass*, not the candidate, if its two +baseline readings disagree by more than ~15% (`report_tuning` already warns there). Never +compare across passes in milliseconds. + +**Reporting.** One row per candidate: the **effective** config (post `workable_chunks` / +`MIN_GROUP_BYTES` / `pick_mc_tiles`), PASS/FAIL, per-pass ratios, summary ratio, spread, +baseline drift. Then the winner *and every candidate within 5% of it*, so a human sees the +plateau rather than a single number, plus the count of pruned combinations with the reason +for each. Print the proposed cache entry verbatim so it can be diffed against the tree — a +sweep that recommends the current defaults is a useful result and should be legible as one. + +## 5. What I could not determine from the code + +- Whether `TILESCALE_GIN_SIGNALS` above 32 is granted. Contexts are a hint and get clamped; + nothing reports a granted signal count. Needs a probe run with `TL_GIN_DEBUG=1`. +- Whether the two per-node toolchains ever disagree about which put sizes lower. Collective + agreement on the verdict makes it safe either way, but it would be worth knowing whether + the `all_reduce` ever actually vetoes anything. +- How the answer moves with node count. Everything above is 2×8. `--algo merged` measures + 1.296 ms against the composed path's 1.086 at two nodes and is expected to win at more, + so the algorithm itself belongs in the search space once a third node exists — and the + key's `num_nodes` field is load-bearing, not decorative. +- Whether the `mc_bcast` / `rs_sum` truncation cases in §2.3 are reachable at the sizes + actually in use, or only at ones the current defaults never produce. Either way a tuner + will reach them, since it explores exactly the corners the defaults avoid. diff --git a/docs/internode_optimization_log.md b/docs/internode_optimization_log.md new file mode 100644 index 0000000000..e96f65dc91 --- /dev/null +++ b/docs/internode_optimization_log.md @@ -0,0 +1,202 @@ +# Inter-node kernel optimisation log + +One section per kernel, worked one at a time. Each records what was tried, what it measured, +and what was concluded — including the attempts that lost, because those are the expensive +knowledge and the reason not to retry them. + +## How to read the numbers + +Absolute times on this cluster drift ~8% with other tenants on the same NICs: one allreduce +configuration measured 0.944 / 0.951 / 1.015 / 1.028 / 1.034 / 1.049 ms unchanged, while the +torch baseline over those same runs moved 0.950 → 1.021. **Ratios are the only stable +quantity**, because torch is timed immediately before and after every candidate. Any +difference under ~5% needs repetition before it means anything, and several early +conclusions in this file's history were noise. + +Runs only happen when two nodes are *fully* idle; a `PreToolUse` guard enforces that and the +`gpu-window` skill queues work for the next window. + +## Sweeping + +Every knob is tunable from one flag, and the sweep runs in a **single process**, because +start-up is ~19 s per rank (6.4 s of it `ncclDevCommCreate`, which does not shrink) so a +process per candidate would spend nearly all its time initialising: + + example_internode_allgather_2d.py --numel 16777216 \ + --sweep "mc_tiles=8,16,32;chunks=4,8;gin_contexts=1,2,4" + +Tunable: `chunks`, `rail_groups`, `gin_contexts`, `mc_threads`, `mc_tiles`, `intra_chunks`, +`threads`. The spec is a **cartesian product**, not one axis at a time, because the knobs do +not compose -- see round 1, where `--chunks 4` measured 209.4 GB/s alone and 185.2 in +combination. Coordinate descent converges on the wrong point here. + +**Candidates are measured round-robin over `--sweep-passes` passes (default 3), each keeping +its best pass.** This is required, not a refinement. Measuring each candidate once in sequence +ranks by *position*, not by configuration: + +| sweep | winner | +|-------|--------| +| `mc_tiles=8,16,32` | 8, at 222.0 GB/s | +| `mc_tiles=32,16,8` | 32, at 218.1 GB/s | +| `mc_threads=512,512,512` (control) | position 1 (1.73x), then 1.55x, then 1.44x | + +The control is the proof: the same configuration three times degraded monotonically by 20%. +That is the GPUs' clocks drooping under sustained load, and it is far larger than the +differences being compared -- a single-pass sweep would have confidently reported whichever +candidate happened to be listed first. Round-robin spreads each candidate across the droop +curve and the per-candidate minimum takes its least-throttled sample; the control then reads +4.4%, and the tool prints a "within the noise floor -- treat these as tied" note when the +whole spread is under 5%. + +Three further properties that make it trustworthy rather than merely convenient: + +- **torch is re-timed either side of every candidate** and the printed `drift` column shows + the disagreement. A candidate whose drift is comparable to its margin proved nothing. +- **Buffers are allocated once** by the first candidate and handed to the rest; only kernels + are rebuilt. Required, not an optimisation: the arena is a bump allocator with no free and + the multicast buffer is sized exactly, so a second allocation exhausts it. +- **GIN signal counters are shared across candidates.** They live on the device and accumulate + for the process lifetime, so a candidate that restarted its count would have its wait + satisfied by an earlier candidate's arrivals and would silently measure nothing. + +Invalid tuples are skipped rather than fatal, and the validity rules are deterministic +functions of the candidate, so every rank rejects the same ones and the collective compiles +stay in lockstep. + +The reported configuration is the **effective** one, after the internal rewrites +(`workable_chunks`, the `MIN_GROUP_BYTES` cap, `pick_mc_tiles`) -- otherwise the table names a +configuration that did not run. + +## Status + +| # | kernel | vs torch | vs triton-dist | state | +|---|--------|----------|----------------|-------| +| 1 | allgather | 1.49–1.78x | 1.30x @64 MB, 0.93x @32 MB | round 1 done | +| 2 | reduce_scatter | 1.34–1.45x | no comparison yet | not started | +| 3 | allreduce | 1.01–1.12x | no comparison yet | not started | +| 4 | ag_gemm | 1.08–1.34x | their test fails on our launcher | not started | +| 5 | gemm_rs | 1.10–1.24x | their test fails on our launcher | not started | + +Baselines are torch NCCL, timed in-run, and triton-dist where its kernel runs inter-node. +Their pull-mode (`dl.symm_at`) kernels are P2P-only by construction and have no inter-node +number; their push kernels do. + +## The one idea not yet tried: SM specialisation + +Every overlap attempt so far has lost, and always the same way — a host `dist.barrier` costs +30–50 µs and a launch 5–10 µs against phases of 0.1–0.2 ms, so adding either exceeds the +fabric time being hidden: + +| attempt | result | +|---|--------| +| merged allreduce, one fabric hop carrying all node partials | 1.296 ms vs composed 1.044 | +| ag_gemm `--mode pipeline`, after the collective got faster | 0.441 vs serial 0.373 | +| gemm_rs `--mode pipeline`, swizzle-adapted remote-first | 0.476 vs serial 0.429 | + +The one overlap that *did* pay — rail-group pipelining, 403 → 471 GB/s — adds no barrier at +all, ordering everything with device-side GIN signals. + +Triton-distributed's answer is **SM specialisation**: a persistent kernel launched with +`grid = min(NUM_SMS, total_tiles)`, with a `gemm_sm` parameter reserving part of the SMs, so +some CTAs communicate while others compute inside one launch. That avoids both costs at once +— no host barrier, and no extra launch — and it is safe from the wait-inside-a-large-grid +deadlock precisely because the grid is capped at one CTA per SM, making every participant +resident. + +This is the missing prerequisite alternative to a device-side barrier (which does not exist +here: `T.barrier_blocks` is lowered with global ranks and takes `get_remote_base_ptr` of each +participant, returning 0 for inter-node peers). Either mechanism would unblock the three +attempts above; SM specialisation needs no new primitive. + +Caveat from an earlier attempt: a persistent GEMM capped at 132 CTAs fixed a hang and was +*slower than serial*, so the cap itself costs something. The open question is whether a +proper split — comm CTAs sized to the fabric, compute CTAs taking the rest — recovers more +than the cap costs. That has never been measured. + +--- + +## 1. allgather + +`example_internode_allgather_2d.py`, rail-aligned GIN + NVSwitch multicast. + +### Where it stands + +``` + 32 MB 64 MB 128 MB 240 MB +ours 203.6 316.3 400.4 471.5 GB/s +torch 121.8 188.5 257.7 309.4 +triton-dist 234.3 243.3 -- -- +``` + +Beats torch everywhere by 1.5–1.68x, and beats triton-dist from 64 MB up (1.30x). **Loses at +32 MB, 0.87x.** Their curve is nearly flat across 32→64 MB (234 → 243) while ours scales +steeply (204 → 316), which is the signature of a latency-dominated regime: their single-kernel +design pays less fixed cost than our multi-launch one. Crossover is near 48 MB. + +### Settled, do not retry + +- **The 32 MB config is already optimal among the existing knobs.** `--rail-groups 2` gives + 0.171 ms and `--mc-tiles 4` gives 0.167 against the default's 0.154. An earlier reading of + 158.2 GB/s that suggested a configuration deficit was contention — torch read 86.6 GB/s in + that same session against 122 on an idle pair. +- `--rail-groups 2` with 4 contexts is the optimum at 240 MB (0.500 ms) and 8 groups on 1 + context is worse (0.561); the reverse holds for allreduce, so the depth is per collective. +- allgather's intra half is two publishes and no reduce, so it is switch-bound rather than + occupancy-bound and wants fatter CTAs than the reduce-carrying collectives (48 MB: 0.164 ms + at 32 tiles against 0.233 at 4). **But "32 at every size" was wrong** -- see round 1: at + 32 MB the curve peaks at 16 and 32 is past it. Only 4 and 32 had been compared at 48 MB. + +### Round 1: knob sweep at 32 MB (2026-08-05) + +Eleven configurations, torch steady at 121.9-123.1 GB/s throughout, so these are comparable. + +| config | GB/s | vs default | +|--------|------|-----------| +| **`--mc-tiles 16`** | **216.9** | **+11.5%** | +| `--mc-tiles 8` | 210.4 | +8.2% | +| `--chunks 4` | 209.4 | +7.7% | +| `--gin-contexts 2` | 208.4 | +7.1% | +| `--mc-threads 256` | 206.6 | +6.2% | +| `--mc-threads 1024` | 202.5 | +4.1% | +| `--gin-contexts 1` | 199.4 | +2.5% | +| default (`--mc-tiles 32`) | 194.5 | -- | +| `--threads 512` | 194.4 | 0% | +| `--mc-tiles 2` | 183.7 | -5.6% | + +**The knobs do not compose.** Every combination of the individual winners came out worse than +the best single change: + +| combination | GB/s | +|-------------|------| +| `--chunks 4 --gin-contexts 2 --mc-tiles 8` | 206.9 | +| `--chunks 4 --gin-contexts 2 --mc-threads 256` | 194.0 | +| `--chunks 4 --mc-tiles 8` | 185.2 | +| `--chunks 4 --gin-contexts 2 --mc-tiles 8 --mc-threads 256` | 183.6 | + +That matters beyond this kernel: it rules out coordinate descent for the autotuner, since each +knob's optimum moves when another changes. The search has to sweep tuples, which is a larger +budget than the design notes assumed. + +Result: 194.5 -> **216.9 GB/s at 32 MB, 1.78x torch**, from one flag. Against triton-dist's +234.3 that closes 0.87x -> 0.93x, though their figure is from an earlier session; a +same-session re-run has twice produced no output, its sweep apparently exceeding the run +timeout at 500 iterations. + +`mc_tiles` is therefore not "32 everywhere" for allgather as previously concluded -- that was +inferred from 48 MB, where only 4 and 32 had been compared. The default is now a threshold on +shard size, fitted to three points and no finer. + +The small-size settings do **not** transfer upward: at 240 MB the combination gives 401.0 GB/s +against the default's 460.3. At 64 MB it gives 327.2 against 316.3, inside noise. + +### Open + +The remaining 32 MB gap is fixed cost, so the candidates reduce launches rather than bandwidth: + +1. Fold `publish_own` and `publish_remote` into one kernel with a slot-indexed grid (-1 launch). +2. Fold the rail put+wait and the publish into one persistent SM-specialised kernel (-2 + launches and the closing barrier). +3. A node-local device barrier, which needs the primitive written first. + +A `--phases` run at 32 MB should apportion the fixed cost first -- without it all three are +guesses about where the 0.145 ms goes. diff --git a/examples/distributed/internode/example_internode_ag_gemm_2d.py b/examples/distributed/internode/example_internode_ag_gemm_2d.py new file mode 100644 index 0000000000..a5109af09c --- /dev/null +++ b/examples/distributed/internode/example_internode_ag_gemm_2d.py @@ -0,0 +1,201 @@ +"""Fused inter-node allgather + GEMM, on the hierarchical (2D) collective. + +``C = allgather(A_shard) @ B``: ``A`` is sharded across ranks along ``M``, ``B`` is +replicated. The flat version of this (``example_internode_ag_gemm.py``) is built on the +flat allgather and so inherits its collapse at 16 GPUs; this one uses ``Allgather2D`` +and the tcgen05 GEMM. + +Three modes: + +* ``serial`` -- allgather, then one GEMM over all ``M`` rows. Forfeits the point of + fusing, and is here as the reference. +* ``overlap`` -- GEMM this rank's own row block during the collective, then the other + fifteen after. Hides only 1/16 of the compute, so it is worth ~3%. +* ``pipeline`` -- GEMM a whole *node's* row blocks as soon as that node's rows are + complete. Since global rank is ``node * lws + local`` and the row block index is the + global rank, a node's rows are **contiguous**, so this is one GEMM launch per node + rather than per rank. Our own node's rows are complete after the intra-node broadcast + alone -- they never touch the fabric -- so half the GEMM (at 2 nodes) runs while the + fabric hop is still in flight. + +What is being hidden +-------------------- +Unfused torch is ``all_gather`` then ``matmul``, strictly serial: 0.204 + 0.170 ms at the +default shape. Serial fusion only inherits the collective's advantage. Overlap puts the +floor at ``max(comm, gemm)``, which also stops cuBLAS being 8% faster than our GEMM from +mattering, since that GEMM is covering the network. + +Measured at 16 GPUs, M=8192 N=4096 K=4096, and the ordering **inverted** once the collective +got faster: with the slower collective pipeline 0.430 ms beat serial 0.481; with rail-group +pipelining, serial 0.373 (736.8 TF) beats pipeline 0.441 (622.9) against torch's 0.500 +(550.0) -- 1.34x against 1.13x. Less exposed fabric leaves less to hide, and the extra +barrier and split launches stop paying. Re-measure whenever the collective changes. + +Done with stream and event ordering, never an in-kernel signal wait: see the deadlock note +in the API reference. Multimem only -- on the pull path ``publish_own`` reads siblings' +shards rather than writing to them, so the ordering differs; falls back to ``overlap``. + +The GEMM wants 256 threads, not the collective's 1024, or warp specialisation overflows the +block limit. +""" + +# NOTE: no `from __future__ import annotations` here -- see internode_2d. +import argparse +import os + +import torch +import torch.distributed as dist + +from internode_2d import Allgather2D, add_2d_args, pick_intra +from internode_common import ( + Context, + TL_DTYPES, + TORCH_DTYPES, + add_common_args, + bench_vs_torch, + check, + prepare_env, +) +from internode_gemm_sm100 import tcgen05_gemm_range_kernel + + +def main() -> int: + parser = add_2d_args(add_common_args(argparse.ArgumentParser(description=__doc__))) + parser.add_argument("--m-per-rank", type=int, default=512) + parser.add_argument("--n", type=int, default=4096) + parser.add_argument("--k", type=int, default=4096) + parser.add_argument("--block-m", type=int, default=128) + parser.add_argument("--block-k", type=int, default=64) + parser.add_argument("--gemm-block-n", type=int, default=256) + # serial is the default again, and the reversal is the point. When the collective was + # slower, pipelining the fabric hop under our own node's GEMM won by 12% (0.430 ms + # against 0.481). Once rail-group pipelining made the collective ~50% faster there was + # far less fabric left to hide, and the pipeline's extra barrier and split GEMM launches + # cost more than they save: serial 0.373 ms / 736.8 TF against pipeline 0.441 / 622.9, + # so 1.34x torch against 1.13x. Re-measure whenever the collective changes -- the answer + # is a function of the comm/compute balance, not a property of this kernel. + parser.add_argument("--mode", choices=("serial", "overlap", "pipeline"), + default="serial") + args = parser.parse_args() + + prepare_env() + world = int(os.environ.get("WORLD_SIZE", torch.cuda.device_count())) + M_per_rank, N, K = args.m_per_rank, args.n, args.k + M = M_per_rank * world + # The collective moves A, so --numel is derived rather than given. + args.numel = M * K + + itemsize = torch.empty((), dtype=TORCH_DTYPES[args.dtype]).element_size() + intra = pick_intra(args.intra) + # A_full + B + C in the arena, plus a multicast copy of A_full. + arena = int(2.5 * (M * K + K * N + M * N) * itemsize) + (1 << 26) + ctx = Context(arena_bytes=arena, + mcast_bytes=M * K * itemsize if intra == "multimem" else 0) + + torch_dtype, tl_dtype = TORCH_DTYPES[args.dtype], TL_DTYPES[args.dtype] + ag = Allgather2D(ctx, M * K, torch_dtype, tl_dtype, args, intra=intra) + lws = ctx.local_world_size + nodes = ctx.world_size // lws + mode = args.mode + if mode == "pipeline" and intra != "multimem": + ctx.log(" note: --mode pipeline needs multimem; falling back to overlap") + mode = "overlap" + + ctx.log( + f"ag_gemm_2d: world={ctx.world_size} M={M} (per-rank {M_per_rank}) N={N} K={K} " + f"mode={mode} intra={intra} chunks={args.chunks} dtype={args.dtype}" + ) + + gemm_full = gemm_block = gemm_node = None + if mode == "serial": + gemm_full = ctx.compile( + tcgen05_gemm_range_kernel(M, N, K, M, block_M=args.block_m, + block_N=args.gemm_block_n, block_K=args.block_k) + ) + elif mode == "overlap": + gemm_block = ctx.compile( + tcgen05_gemm_range_kernel(M, N, K, M_per_rank, block_M=args.block_m, + block_N=args.gemm_block_n, block_K=args.block_k) + ) + else: + # One node's worth of rows, which is contiguous -- see Allgather2D.rows_of_node. + gemm_node = ctx.compile( + tcgen05_gemm_range_kernel(M, N, K, lws * M_per_rank, block_M=args.block_m, + block_N=args.gemm_block_n, block_K=args.block_k) + ) + + B = ctx.tensor((K, N), torch_dtype) + C = ctx.tensor((M, N), torch_dtype) + A_shard = ag.shard.view(M_per_rank, K) + A_full = ag.out.view(M, K) + A_shard.copy_((torch.randn(M_per_rank, K, device=B.device) * 0.02).to(torch_dtype)) + B.copy_((torch.randn(K, N, device=B.device) * 0.02).to(torch_dtype)) + C.zero_() + + gemm_stream = torch.cuda.Stream() + comm_stream = torch.cuda.Stream() + + my_node = ctx.rank // lws + + def launch(): + main_stream = torch.cuda.current_stream() + if mode == "serial": + ag.launch() + gemm_full(A_full, B, C, 0) + return + if mode == "overlap": + # Our own rows are already local, so their GEMM waits on nothing. + gemm_stream.wait_stream(main_stream) + with torch.cuda.stream(gemm_stream): + gemm_block(A_full, B, C, ctx.rank * M_per_rank) + ag.launch() + main_stream.wait_stream(gemm_stream) + for peer in range(ctx.world_size): + if peer != ctx.rank: + gemm_block(A_full, B, C, peer * M_per_rank) + return + # pipeline: the fabric hop runs on its own stream while we finish and then + # consume our own node's rows, which never cross the fabric. + ag.rail_hop(stream=comm_stream) + ag.publish_own() + # Our node's rows are complete once every sibling has published. This barrier is + # ordered after publish_own on the main stream, so it does not wait for the + # fabric hop -- that is still in flight on comm_stream. + dist.barrier(ctx.group) + gemm_node(A_full, B, C, ag.rows_of_node(my_node, M_per_rank)) + main_stream.wait_stream(comm_stream) + ag.consume_groups() + dist.barrier(ctx.group) + for n in range(nodes): + if n != my_node: + gemm_node(A_full, B, C, ag.rows_of_node(n, M_per_rank)) + + torch.cuda.synchronize() + dist.barrier(ctx.group) + launch() + torch.cuda.synchronize() + + ref_a = torch.empty(M, K, dtype=torch_dtype, device=B.device) + dist.all_gather_into_tensor(ref_a.view(-1), ag.shard, group=ctx.group) + ref = (ref_a.float() @ B.float()).to(torch_dtype) + failures = check(ctx, C, ref, "ag_gemm_2d") + + if not args.no_bench and failures == 0: + ref_buf = torch.empty(M, K, dtype=torch_dtype, device=B.device) + out_buf = torch.empty(M, N, dtype=torch_dtype, device=B.device) + + def run_ref(): + dist.all_gather_into_tensor(ref_buf.view(-1), ag.shard, group=ctx.group) + torch.matmul(ref_buf, B, out=out_buf) + + bench_vs_torch(ctx, args, "ag_gemm_2d", launch, run_ref, 0, + tflops=2 * M * N * K / 1e12) + + ctx.close() + if ctx.is_leader: + print("PASS" if failures == 0 else f"FAIL: {failures} rank(s) mismatched", flush=True) + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/distributed/internode/example_internode_allgather_2d.py b/examples/distributed/internode/example_internode_allgather_2d.py new file mode 100644 index 0000000000..47325ada64 --- /dev/null +++ b/examples/distributed/internode/example_internode_allgather_2d.py @@ -0,0 +1,125 @@ +"""Hierarchical (2D) inter-node allgather: rail-aligned GIN plus an NVSwitch broadcast. + +The flat allgather in ``example_internode_allgather.py`` pushes every rank's shard to +all ``world_size - 1`` peers over GIN. That is optimal with one GPU per node and +catastrophic with eight: each rank puts 15 shards onto the NIC, including the 7 for +siblings on the same machine, and it measures 2.7 GB/s against torch's 309. + +The kernels live in ``internode_2d`` -- shared with the other three collectives, which +are the same two halves recombined -- so this file is just the shape, the reference and +the measurement. Read ``internode_2d.Allgather2D`` for what overlaps and why, and the +module docstring there for the three traps worth knowing before editing any of it. + +Measured, 16 GPUs on 2 nodes, 240 MB bf16: + +``` +flat 2.7 GB/s +2D pull 358.6 GB/s torch 309 1.16x +2D multimem 403.6 GB/s torch 309 1.31x (triton-dist push_2d: 290) +``` + +against a roofline of ~700 GB/s -- one shard over this rank's own NIC +(15.7 MB / 47.6 GB/s = 0.33 ms) and fourteen over NVLink (220 MB / 670 GB/s = 0.33 ms), +two legs that happen to balance at 8 GPUs against 8 400-Gbps NICs. ``--phases`` prints +where the time actually goes. +""" + +# NOTE: no `from __future__ import annotations` here -- see internode_2d. +import argparse + +import torch +import torch.distributed as dist + +from internode_2d import ( + Allgather2D, + add_2d_args, + pick_intra, + report_phases, + run_sweep, +) +from internode_common import ( + Context, + TL_DTYPES, + TORCH_DTYPES, + add_common_args, + bench_vs_torch, + check, + prepare_env, +) + + +def main() -> int: + parser = add_2d_args(add_common_args(argparse.ArgumentParser(description=__doc__))) + parser.add_argument("--phases", action="store_true", + help="time each kernel on its own, to locate headroom") + # Allgather's intra-node half is two multicast publishes and no reduce, so it is + # switch-bound rather than occupancy-bound and wants fatter CTAs than the + # reduce-carrying collectives: 32 tiles gives 0.164 ms against 0.233 at 4 for 48 MB, + # and 0.501 against 0.585 for 240 MB. + # + # But not at the smallest sizes. At 32 MB (a 2 MB shard) the curve peaks at 16, and 32 is + # well past it: tiles 2/8/16/32 measure 183.7 / 210.4 / 216.9 / 194.5 GB/s. So this is a + # threshold fitted to three shard sizes -- 2 MB wants 16, 3 MB and 15.7 MB want 32 -- and + # nothing finer. --mc-tiles overrides it. + parser.set_defaults(mc_tiles=0) + args = parser.parse_args() + + prepare_env() + # The multicast buffer is sized before the allocator exists, so the output length + # has to be known here rather than after Context. + itemsize = torch.empty((), dtype=TORCH_DTYPES[args.dtype]).element_size() + intra = pick_intra(args.intra) + ctx = Context(mcast_bytes=args.numel * itemsize if intra == "multimem" else 0) + + torch_dtype, tl_dtype = TORCH_DTYPES[args.dtype], TL_DTYPES[args.dtype] + nodes = ctx.world_size // ctx.local_world_size + ctx.log( + f"allgather_2d: world={ctx.world_size} nodes={nodes} local={ctx.local_world_size} " + f"numel={args.numel} shard={args.numel // ctx.world_size} chunks={args.chunks} " + f"intra={intra} mc_threads={args.mc_threads} mc_tiles={args.mc_tiles} " + f"overlap={not args.no_overlap} contexts={args.gin_contexts} dtype={args.dtype}" + ) + + if not args.mc_tiles: + args.mc_tiles = 16 if (args.numel // ctx.world_size) * itemsize < 3_000_000 else 32 + def make(cand, buffers=None, signal_state=None): + return Allgather2D(ctx, cand.numel, torch_dtype, tl_dtype, cand, intra=intra, + buffers=buffers, signal_state=signal_state) + + ag = make(args) + ag.shard.copy_( + torch.arange(ag.shard_numel, device=ag.shard.device, dtype=torch.float32) + .to(torch_dtype) + ctx.rank * 1000.0 + ) + ref = torch.empty_like(ag.out) + run_ref = lambda: dist.all_gather_into_tensor(ref, ag.shard, group=ctx.group) + moved = ag.shard.numel() * ag.shard.element_size() * (ctx.world_size - 1) + + def verify(coll, label): + dist.all_gather_into_tensor(ref, coll.shard, group=ctx.group) + return check(ctx, coll.out, ref, label) + + if args.sweep: + # Buffers come from the first candidate; every later one rebuilds kernels only. + run_sweep(ctx, args, make, verify, lambda c: c.launch, run_ref, moved, + "allgather_2d", buffers=ag.buffers) + else: + torch.cuda.synchronize() + dist.barrier(ctx.group) + ag.launch() + torch.cuda.synchronize() + failures = check(ctx, ag.out, ref, "allgather_2d") + if failures == 0: + if args.phases: + report_phases(ctx, ag, args) + if not args.no_bench: + bench_vs_torch(ctx, args, "allgather_2d", ag.launch, run_ref, moved) + + ctx.close() + if ctx.is_leader and not args.sweep: + print("PASS" if failures == 0 else f"FAIL: {failures} rank(s) mismatched", flush=True) + return 1 if (not args.sweep and failures) else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/distributed/internode/example_internode_allreduce_2d.py b/examples/distributed/internode/example_internode_allreduce_2d.py new file mode 100644 index 0000000000..04b1780af0 --- /dev/null +++ b/examples/distributed/internode/example_internode_allreduce_2d.py @@ -0,0 +1,153 @@ +"""Hierarchical (2D) inter-node allreduce: 2D reduce-scatter then 2D allgather. + +The flat allreduce in ``example_internode_allreduce.py`` has every rank push to all +``world_size - 1`` peers over GIN, which is fine with one GPU per node and collapses +with eight -- see ``internode_2d`` for the measurement and the fix. + +There is nothing new here: allreduce *is* reduce-scatter followed by allgather, and +both halves are already 2D and already beat torch on their own. So this example is +composition, and the interesting parts are the two seams: + +* **The shard buffer is shared, not copied.** ``Allgather2D`` takes the reduce-scatter's + output as its input, so the two halves hand over in place. Both live in the arena, so + this is just a pointer. +* **The halves need disjoint GIN signal *ranges*.** Signal state is cumulative and a wait + does not consume it, so an overlap lets one half's wait be satisfied by the other's + bytes -- silently, and only under repetition. Each half occupies ``signals_used`` of + them (one per pipelined group), so the allgather starts at + ``SIGNAL_DATA + rs.signals_used`` rather than at a hardcoded second signal. Getting this + wrong by one corrupted exactly the second half of the output. 32 signals are + provisioned. + +Volume-wise this is the right decomposition for large buffers: two-shot moves +``2(W-1)N/W`` bytes against one-shot's ``(W-1)N``. One-shot only wins at small ``W`` +and even then not here, because it reduces over the full buffer rather than a shard -- +see the flat example's ``--algo oneshot`` for that measurement. +""" + +# NOTE: no `from __future__ import annotations` here -- see internode_2d. +import argparse +import os + +import torch +import torch.distributed as dist + +from internode_2d import ( + Allgather2D, + Allreduce2D, + ReduceScatter2D, + add_2d_args, + fused_allreduce_launch, + pick_intra, +) +from internode_common import ( + SIGNAL_DATA, + Context, + TL_DTYPES, + TORCH_DTYPES, + add_common_args, + bench_vs_torch, + check, + prepare_env, +) + + +def main() -> int: + parser = add_2d_args(add_common_args(argparse.ArgumentParser(description=__doc__))) + # composed wins on measurement: 1.086 ms against merged's 1.296. The single hop saves + # a serialisation but the NVLink publish cost is identical, and with only `nodes` slots + # the merged pipeline is too coarse-grained to make up the difference. + parser.add_argument("--algo", choices=("fused", "merged", "composed"), default="fused", + help="merged: one fabric hop carrying every node partial. " + "composed: reduce-scatter then allgather, which is simpler " + "and works without multicast, but the halves cannot overlap") + # Allreduce wants a finer pipeline than the one-hop collectives: it has two fabric + # hops to overlap, so 8 groups beats 2 (0.95 ms vs 1.02). Each group's grid must be a + # multiple of the GIN context count, and 8 groups of chunks=8 leaves one chunk per + # group, hence one context. Allgather goes the other way -- it prefers 2 groups on 4 + # contexts (0.500 ms vs 0.561 at 8/1) -- so these defaults are per example, not global. + parser.set_defaults(rail_groups=8, gin_contexts=1) + args = parser.parse_args() + + prepare_env() + # The multicast buffer is sized before the allocator exists. Both halves want one + # -- the reduce-scatter reduces its input through the switch, the allgather + # broadcasts into its output -- hence two buffers of `numel`. + world = int(os.environ.get("WORLD_SIZE", torch.cuda.device_count())) + itemsize = torch.empty((), dtype=TORCH_DTYPES[args.dtype]).element_size() + intra = pick_intra(args.intra) + if args.algo == "merged" and intra != "multimem": + args.algo = "composed" + if args.algo == "fused" and intra != "multimem": + args.algo = "composed" + # Both algorithms want two multicast buffers of `numel`: one reduced through the + # switch, one broadcast out of it. + ctx = Context(mcast_bytes=2 * args.numel * itemsize if intra == "multimem" else 0) + + torch_dtype, tl_dtype = TORCH_DTYPES[args.dtype], TL_DTYPES[args.dtype] + nodes = ctx.world_size // ctx.local_world_size + ctx.log( + f"allreduce_2d: world={ctx.world_size} nodes={nodes} local={ctx.local_world_size} " + f"numel={args.numel} shard={args.numel // ctx.world_size} chunks={args.chunks} " + f"algo={args.algo} intra={intra} mc_threads={args.mc_threads} mc_tiles={args.mc_tiles} " + f"overlap={not args.no_overlap} contexts={args.gin_contexts} dtype={args.dtype}" + ) + + if args.algo == "merged": + ar = Allreduce2D(ctx, args.numel, torch_dtype, tl_dtype, args, intra=intra) + inp, out, launch = ar.inp, ar.out, ar.launch + else: + rs = ReduceScatter2D(ctx, args.numel, torch_dtype, tl_dtype, args, intra=intra, + signal_id=SIGNAL_DATA) + # Hand over in place, and on a *disjoint signal range* -- not merely a different + # signal. Each half occupies one signal per pipelined group, so with + # --rail-groups 2 the reduce-scatter holds {0,1}; starting the allgather at + # SIGNAL_PHASE2 == 1 overlapped it and its group-0 wait was satisfied by the + # reduce-scatter's group-1 arrivals, corrupting the second half of the output. + ag = Allgather2D(ctx, args.numel, torch_dtype, tl_dtype, args, intra=intra, + signal_id=SIGNAL_DATA + rs.signals_used, shard=rs.out) + inp, out = rs.inp, ag.out + + if args.algo == "fused": + def launch(): + fused_allreduce_launch(rs, ag, ctx) + else: + def launch(): + rs.launch() + ag.launch() + # Small magnitudes: a bf16 sum over 16 ranks of arange values would land outside any + # sensible tolerance. + inp.copy_( + (torch.arange(args.numel, device=inp.device, dtype=torch.float32) % 7 + ctx.rank) + .to(torch_dtype) + ) + out.zero_() + + torch.cuda.synchronize() + dist.barrier(ctx.group) + launch() + torch.cuda.synchronize() + + ref = inp.clone() + dist.all_reduce(ref, op=dist.ReduceOp.SUM, group=ctx.group) + failures = check(ctx, out, ref, "allreduce_2d") + + if not args.no_bench and failures == 0: + ref_buf = torch.empty_like(inp) + + def run_ref(): + ref_buf.copy_(inp) + dist.all_reduce(ref_buf, op=dist.ReduceOp.SUM, group=ctx.group) + + # Allreduce convention: 2(W-1)/W, both directions of the two-shot exchange. + moved = 2 * inp.numel() * inp.element_size() * (ctx.world_size - 1) // ctx.world_size + bench_vs_torch(ctx, args, "allreduce_2d", launch, run_ref, moved) + + ctx.close() + if ctx.is_leader: + print("PASS" if failures == 0 else f"FAIL: {failures} rank(s) mismatched", flush=True) + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/distributed/internode/example_internode_gemm_rs_2d.py b/examples/distributed/internode/example_internode_gemm_rs_2d.py new file mode 100644 index 0000000000..0ecc8f1c42 --- /dev/null +++ b/examples/distributed/internode/example_internode_gemm_rs_2d.py @@ -0,0 +1,176 @@ +"""Fused inter-node GEMM + reduce-scatter, on the hierarchical (2D) collective. + +``C_shard = reduce_scatter(A @ B)``, with ``K`` sharded across ranks: every rank holds +``A (M x K/W)`` and ``B (K/W x N)``, computes a full-size *partial* product, and the +collective sums the partials and scatters the result along ``M``. + +The GEMM writes straight into the collective's input buffer, which on the multimem path +is the multicast allocation -- so there is no copy between compute and communication. +That is the one integration detail worth noting: ``ReduceScatter2D.inp`` is this rank's +own view of that buffer, and it is a perfectly ordinary tensor to write into. + +``--mode pipeline``: overlap by node slot, adapted from triton-dist's tile swizzle +-------------------------------------------------------------------------------- +The dependency runs the wrong way for the trick AG-GEMM uses -- here the GEMM *produces* +what is sent -- but it can still be split. Output row block *g* belongs to global rank *g*, +and one node's ranks occupy a contiguous block of ``M``, so compute the rows the *other* +node needs first, hand them to the fabric, and compute our own while they fly. + +That remote-first ordering is Triton-distributed's ``swizzle_tiled_m_with_padding``, which +renumbers GEMM tiles so a rank computes rank *r+1*'s block first and its own last. **Their +exact swizzle does not transfer**: theirs rotates per rank, which suits per-peer pushes where +every block has one destination, whereas ``multimem.ld_reduce`` needs *every* local rank to +have written a segment, so a per-rank rotation leaves each rank ready on a segment its +siblings are not. Ours needs the same idea at coarser grain in a *common* order, one barrier +per node slot. + +**And it loses, so ``serial`` stays the default.** At ``k-per-rank 2048`` serial 0.334 ms / +411 TF against pipeline 0.365 / 376; at 4096, serial 0.429 / 640 against 0.476 / 578. Two +costs swallow the gain: the two barriers are ~30-50 us each, and splitting one GEMM into two +half-height launches worsens the persistent tcgen05 grid's wave quantisation. Together they +exceed the fabric time hidden -- even though comm (~0.237 ms, near-fixed since the output is +``M*N`` whatever ``K`` is) and compute (0.031 ms at K/rank 512 rising to ~0.248 at 4096) are +comparable at the larger shapes. + +The flag stays because the balance moves with barrier cost: with a device-side barrier this +should invert. Three overlap attempts have now reached that same conclusion. +""" + +# NOTE: no `from __future__ import annotations` here -- see internode_2d. +import argparse +import os + +import torch +import torch.distributed as dist + +from internode_2d import ReduceScatter2D, add_2d_args, pick_intra +from internode_common import ( + Context, + TL_DTYPES, + TORCH_DTYPES, + add_common_args, + bench_vs_torch, + check, + prepare_env, +) +from internode_gemm_sm100 import tcgen05_gemm_range_kernel + + +def main() -> int: + parser = add_2d_args(add_common_args(argparse.ArgumentParser(description=__doc__))) + parser.add_argument("--m", type=int, default=8192) + parser.add_argument("--n", type=int, default=4096) + parser.add_argument("--k-per-rank", type=int, default=512) + parser.add_argument("--block-m", type=int, default=128) + parser.add_argument("--block-k", type=int, default=64) + parser.add_argument("--gemm-block-n", type=int, default=256) + parser.add_argument("--mode", choices=("serial", "pipeline"), default="serial", + help="pipeline computes the other node's rows first and reduces them " + "while our own node's rows are still being computed") + args = parser.parse_args() + + prepare_env() + world = int(os.environ.get("WORLD_SIZE", torch.cuda.device_count())) + M, N, K_per_rank = args.m, args.n, args.k_per_rank + args.numel = M * N + + itemsize = torch.empty((), dtype=TORCH_DTYPES[args.dtype]).element_size() + intra = pick_intra(args.intra) + arena = int(3.5 * (M * K_per_rank + K_per_rank * N + M * N) * itemsize) + (1 << 26) + ctx = Context(arena_bytes=arena, + mcast_bytes=M * N * itemsize if intra == "multimem" else 0) + + torch_dtype, tl_dtype = TORCH_DTYPES[args.dtype], TL_DTYPES[args.dtype] + ctx.log( + f"gemm_rs_2d: world={ctx.world_size} M={M} N={N} K={K_per_rank}/rank " + f"intra={intra} mode={args.mode} chunks={args.chunks} dtype={args.dtype}" + ) + + rs = ReduceScatter2D(ctx, M * N, torch_dtype, tl_dtype, args, intra=intra) + + nodes = ctx.world_size // ctx.local_world_size + rows_per_node = M // nodes + if args.mode == "pipeline" and M % nodes: + raise SystemExit(f"--m {M} must divide by {nodes} nodes for --mode pipeline") + gemm = ctx.compile( + tcgen05_gemm_range_kernel(M, N, K_per_rank, M, block_M=args.block_m, + block_N=args.gemm_block_n, block_K=args.block_k) + ) + gemm_node = ctx.compile( + tcgen05_gemm_range_kernel(M, N, K_per_rank, rows_per_node, block_M=args.block_m, + block_N=args.gemm_block_n, block_K=args.block_k) + ) if args.mode == "pipeline" else None + + A = ctx.tensor((M, K_per_rank), torch_dtype) + B = ctx.tensor((K_per_rank, N), torch_dtype) + A.copy_((torch.randn(M, K_per_rank, device=A.device) * 0.02).to(torch_dtype)) + B.copy_((torch.randn(K_per_rank, N, device=A.device) * 0.02).to(torch_dtype)) + # The GEMM writes the collective's input directly -- no staging copy. + partial = rs.inp.view(M, N) + + my_node = ctx.rank // ctx.local_world_size + + def launch_pipeline(): + # The other node's rows are what the fabric needs, so they go first. + for n in range(nodes): + if n == my_node: + continue + gemm_node(A, B, partial, n * rows_per_node) + dist.barrier(ctx.group) # every rank has finished those rows + rs.reduce_remote() + targets = rs.issue_puts() # fabric starts + gemm_node(A, B, partial, my_node * rows_per_node) # overlaps the transfer + dist.barrier(ctx.group) + rs.reduce_own() + for g, target in enumerate(targets): + rs.finish_group(g, target) + + def launch(): + if args.mode == "pipeline": + launch_pipeline() + return + gemm(A, B, partial, 0) + # The GEMM *produces* the collective's input, and the reduce reads every local + # rank's copy of it through the multicast VA. Stream order only sequences our own + # GEMM before our own reduce -- it says nothing about a sibling's GEMM, so + # without this fence we reduce whatever a slower sibling has written so far. + # + # ReduceScatter2D advertises "no barrier needed", and that is true when the input + # is filled once before the loop, as in the standalone example. A producer inside + # the loop breaks that precondition. It passed on the 8-GPU single-node proxy -- + # where ranks stay tightly synchronised -- and mismatched on 11 of 16 ranks + # across two nodes, which is exactly the shape of a skew-dependent race. + dist.barrier(ctx.group) + rs.launch() + + torch.cuda.synchronize() + dist.barrier(ctx.group) + launch() + torch.cuda.synchronize() + + ref_full = (A.float() @ B.float()).to(torch_dtype) + ref = torch.empty(M // ctx.world_size * N, dtype=torch_dtype, device=A.device) + dist.reduce_scatter_tensor(ref, ref_full.reshape(-1).contiguous(), + op=dist.ReduceOp.SUM, group=ctx.group) + failures = check(ctx, rs.out, ref, "gemm_rs_2d") + + if not args.no_bench and failures == 0: + gemm_buf = torch.empty(M, N, dtype=torch_dtype, device=A.device) + ref_buf = torch.empty_like(ref) + + def run_ref(): + torch.matmul(A, B, out=gemm_buf) + dist.reduce_scatter_tensor(ref_buf, gemm_buf.reshape(-1), + op=dist.ReduceOp.SUM, group=ctx.group) + + bench_vs_torch(ctx, args, "gemm_rs_2d", launch, run_ref, 0, + tflops=2 * M * N * K_per_rank / 1e12) + + ctx.close() + if ctx.is_leader: + print("PASS" if failures == 0 else f"FAIL: {failures} rank(s) mismatched", flush=True) + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/distributed/internode/example_internode_reduce_scatter_2d.py b/examples/distributed/internode/example_internode_reduce_scatter_2d.py new file mode 100644 index 0000000000..5324339cf1 --- /dev/null +++ b/examples/distributed/internode/example_internode_reduce_scatter_2d.py @@ -0,0 +1,103 @@ +"""Hierarchical (2D) inter-node reduce-scatter: NVSwitch reduce, then rail-aligned GIN. + +The transpose of the 2D allgather, and the flat version collapses at 16 GPUs for the +same reason: every rank pushes a slice to all 15 peers over GIN, including the 7 for +siblings on the same machine. + +The kernels live in ``internode_2d``. Read ``internode_2d.ReduceScatter2D`` for why this +direction needs no barrier at all, and the module docstring there for the traps. + +Measured, 16 GPUs on 2 nodes, 240 MB bf16: + +``` +flat 2.5 GB/s +2D pull 293.9 GB/s torch 319 0.92x +2D multimem 397.2 GB/s torch 322 1.24x +``` + +The pull path structurally cannot reach torch: reducing on the consumer means moving +``(lws-1)/lws`` of the input over NVLink into a staging buffer and reading it back, +~210 MB of NVLink plus ~500 MB of avoidable HBM per rank. That was exactly the 8% it +lost, and no tuning removes it. ``multimem.ld_reduce`` reduces in the switch instead -- +the same NVLS mechanism torch is using here, which is what makes the comparison fair +rather than us fighting a switch-reduce with a copy loop. +""" + +# NOTE: no `from __future__ import annotations` here -- see internode_2d. +import argparse + +import torch +import torch.distributed as dist + +from internode_2d import ReduceScatter2D, add_2d_args, pick_intra, report_phases +from internode_common import ( + Context, + TL_DTYPES, + TORCH_DTYPES, + add_common_args, + bench_vs_torch, + check, + prepare_env, +) + + +def main() -> int: + parser = add_2d_args(add_common_args(argparse.ArgumentParser(description=__doc__))) + parser.add_argument("--phases", action="store_true", + help="time each kernel on its own, to locate headroom") + args = parser.parse_args() + + prepare_env() + itemsize = torch.empty((), dtype=TORCH_DTYPES[args.dtype]).element_size() + intra = pick_intra(args.intra) + ctx = Context(mcast_bytes=args.numel * itemsize if intra == "multimem" else 0) + + torch_dtype, tl_dtype = TORCH_DTYPES[args.dtype], TL_DTYPES[args.dtype] + nodes = ctx.world_size // ctx.local_world_size + ctx.log( + f"reduce_scatter_2d: world={ctx.world_size} nodes={nodes} " + f"local={ctx.local_world_size} numel={args.numel} " + f"shard={args.numel // ctx.world_size} chunks={args.chunks} intra={intra} " + f"mc_threads={args.mc_threads} mc_tiles={args.mc_tiles} " + f"intra_chunks={args.intra_chunks} overlap={not args.no_overlap} " + f"contexts={args.gin_contexts} dtype={args.dtype}" + ) + + rs = ReduceScatter2D(ctx, args.numel, torch_dtype, tl_dtype, args, intra=intra) + # Small magnitudes: a bf16 sum over 16 ranks of arange values would land outside any + # sensible tolerance. + rs.inp.copy_( + (torch.arange(args.numel, device=rs.inp.device, dtype=torch.float32) % 7 + ctx.rank) + .to(torch_dtype) + ) + + torch.cuda.synchronize() + dist.barrier(ctx.group) + rs.launch() + torch.cuda.synchronize() + + ref = torch.empty_like(rs.out) + dist.reduce_scatter_tensor(ref, rs.inp, op=dist.ReduceOp.SUM, group=ctx.group) + failures = check(ctx, rs.out, ref, "reduce_scatter_2d") + + if failures == 0: + if args.phases: + report_phases(ctx, rs, args) + if not args.no_bench: + ref_buf = torch.empty_like(rs.out) + moved = rs.out.numel() * rs.out.element_size() * (ctx.world_size - 1) + bench_vs_torch( + ctx, args, "reduce_scatter_2d", rs.launch, + lambda: dist.reduce_scatter_tensor(ref_buf, rs.inp, op=dist.ReduceOp.SUM, + group=ctx.group), + moved, + ) + + ctx.close() + if ctx.is_leader: + print("PASS" if failures == 0 else f"FAIL: {failures} rank(s) mismatched", flush=True) + return 1 if failures else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/distributed/internode/internode_2d.py b/examples/distributed/internode/internode_2d.py new file mode 100644 index 0000000000..8c37eca0a3 --- /dev/null +++ b/examples/distributed/internode/internode_2d.py @@ -0,0 +1,1384 @@ +"""Hierarchical (2D) inter-node collectives, as reusable pieces. + +Everything here exists because a *flat* collective collapses once there is more than +one GPU per node: each rank puts ``world_size - 1`` shards onto the NIC, including the +ones destined for siblings on the same machine. At 16 GPUs that measures ~2.6 GB/s +against torch's ~310. + +The fix is to split every collective along the topology, and all four collectives in +this directory reduce to the same two halves: + +* **rail** -- the fabric hop. Rank ``(node, local)`` exchanges only with the *same* + ``local`` index on other nodes, ``peer = other_node * lws + local``. Every rank + drives its own NIC with no fan-out, and only ``1/lws`` of the data crosses. +* **intra** -- finish inside the node. Either through the NVSwitch (``multimem.st`` to + broadcast, ``multimem.ld_reduce`` to reduce, one instruction reaching every local + rank) or, without multicast, by reading peers' buffers with ``T.get_block``. + +Allgather is rail-then-broadcast, reduce-scatter is reduce-then-rail, allreduce is the +two composed, and the fused GEMM kernels wrap one of them around +``tcgen05_gemm_range_kernel``. So they share these kernels rather than restating them. + +Measured at 16 GPUs on 2 nodes, 240 MB bf16 (bus bandwidth, ``nbytes*(W-1)/W/time``): + +``` + flat 2D pull 2D multimem torch best vs torch +allgather 2.7 358.6 403.6 309.0 1.31x +reduce_scatter 2.5 293.9 397.2 321.6 1.24x +``` + +The roofline is ~700 GB/s: each GPU must take one shard over its own NIC +(15.7 MB / 47.6 GB/s = 0.33 ms) and 14 over NVLink (220 MB / 670 GB/s = 0.33 ms), and +those two legs are almost perfectly balanced at 8 GPUs against 8 400-Gbps NICs. The +remaining gap is that the second intra-node phase cannot start until the fabric hop +has landed; see ``Allgather2D`` for what overlaps and what does not. + +Traps +----- +``docs/distributed_api_reference.md`` documents the ones that belong to the APIs +themselves: peer arguments are global ranks, GIN signals are cumulative and +per-context, a multimem region must be one contiguous pair per thread with a +compile-time offset, and waiting on a signal inside a large-grid kernel +deadlocks. The one specific to this module: + +**Keep the per-put size constant when splitting the fabric hop.** ``--rail-groups`` +divides the hop into groups that each keep the same message size and simply use +fewer concurrent CTAs. Splitting the payload into more, smaller messages instead +measured 4x slower, because RDMA is bandwidth-bound only once messages are large. +""" + +# NOTE: no `from __future__ import annotations` here. T.prim_func resolves parameter +# annotations at runtime via get_type_hints, and PEP 563 would turn `T.Tensor(...)` +# into a string evaluated against module globals, where the closure locals do not +# exist. +import argparse +import functools + +import torch +import torch.distributed as dist + +import tilelang +import tilelang.language as T + +from internode_common import SIGNAL_DATA, fp32_sum + +# Smallest per-group slice worth pipelining; below this the extra launches cost more than +# the overlap saves. See _Base.__init__ for the measurement. +MIN_GROUP_BYTES = 1_500_000 + +# --mc-tiles has to scale with the buffer for the reduce-carrying collectives, and getting +# it wrong is expensive in both directions. Measured at 16 GPUs: allreduce wants 32 tiles +# per CTA at 240 MiB (0.944 ms against 1.086 at 4) and 4 at 48 MiB (0.282 against 0.361 at +# 32) -- a fat-CTA publish starves the GPU once the shard is small, and a thin one wastes +# scheduling once it is large. reduce_scatter behaves the same way (48 MiB: 0.138 ms at the +# scaled value against 0.194 at 32). +# +# Allgather does *not*: it wants 32 at every size measured (48 MiB 0.164 ms at 32 against +# 0.233 at 4; 240 MiB 0.501 against 0.585). Its intra-node half is two publishes and no +# reduce, so it is switch-bound rather than occupancy-bound, and fewer fatter CTAs win. +# Hence the example overrides this rather than a single global law covering both. +# One tile per 480 KB of shard. 240 KB was tried and regressed allreduce at 120 MiB +# (0.747 ms against 0.521), so the coarser divisor stands. This is a fit to three sizes, +# not a law: the true optimum also depends on the group count and on which collective, and +# --mc-tiles overrides it. See the tuning table in CLAUDE.md for what is left on the table. +MC_TILE_BYTES_PER_TILE = 480 << 10 + + +def pick_mc_tiles(shard_numel, itemsize, requested): + """Tiles per multicast CTA: as asked if given, else scaled to the shard. See above.""" + if requested: + return requested + want = max(1, (shard_numel * itemsize) // MC_TILE_BYTES_PER_TILE) + tiles = 1 + while tiles * 2 <= min(want, 32): + tiles *= 2 + return tiles + + +# --------------------------------------------------------------------- fabric hop + + +def rail_put_kernel( + shard_numel: int, chunks: int, threads: int, world_size: int, local_world_size: int, + dtype: str, signal_id: int = SIGNAL_DATA, src_per_node: bool = False, + wait: bool = False, chunk_lo: int = 0, chunk_count: int = 0, +): + """Rail-aligned put of one shard to the same local index on every other node. + + ``chunk_lo``/``chunk_count`` restrict the launch to chunks + ``[chunk_lo, chunk_lo + chunk_count)``, keeping the per-put size and dropping only the + number of concurrent puts -- see the module docstring on why that distinction matters. + + The inbox is indexed by *sender node*, so slots stay disjoint with no rotation + arithmetic: only rank ``(n, l)`` ever writes our slot ``n``. + + ``src_per_node`` selects what is sent. Allgather sends the same shard to every rail + peer (one source slot); reduce-scatter sends a different per-node partial to each + (``nodes`` source slots, indexed by destination node). + + ``wait=True`` folds the arrival wait into this launch. Leave it False when the + consumer is a separate kernel that can wait itself, which lets the put kernel + retire and free its SMs while the RDMA is still in flight. + """ + nodes = world_size // local_world_size + chunk_numel = shard_numel // chunks + src_slots = nodes if src_per_node else 1 + grid = chunk_count or chunks + + @T.prim_func + def main( + src: T.Tensor((src_slots * shard_numel,), dtype), + inbox: T.Tensor((nodes * shard_numel,), dtype), + rank: T.int32, + signal_target: T.int32, + ): + with T.Kernel(grid, threads=threads) as bx: + base = (chunk_lo + bx) * chunk_numel + local_rank = rank % local_world_size + node = rank // local_world_size + for step in range(nodes - 1): + n = (node + step + 1) % nodes + src_off = (n * shard_numel if src_per_node else 0) + base + T.nccl_gin.put_signal( + src=src[src_off], + dst=inbox[node * shard_numel + base], + size=chunk_numel, + peer=n * local_world_size + local_rank, + signal_id=signal_id, + scope="block", + ) + if wait: + T.nccl_gin.wait_signal(least=signal_target, signal_id=signal_id, + scope="block") + + return main + + +def rail_wait_kernel( + shard_numel: int, chunks: int, threads: int, world_size: int, local_world_size: int, + dtype: str, signal_id: int = SIGNAL_DATA, chunk_count: int = 0, +): + """Wait for one group's rail arrivals, nothing else. + + Separate from the put so a group's puts can be issued and its CTAs retire, leaving the + RDMA in flight while the previous group is published over NVLink. + + The grid must match the *sender's* chunk count, for the per-context signal reason in + the API reference; hence ``chunks % gin_contexts == 0`` per group, checked host side. + """ + nodes = world_size // local_world_size + grid = chunk_count or chunks + + @T.prim_func + def main( + inbox: T.Tensor((nodes * shard_numel,), dtype), + signal_target: T.int32, + ): + with T.Kernel(grid, threads=threads) as bx: + T.nccl_gin.wait_signal(least=signal_target, signal_id=signal_id, scope="block") + # Keep `inbox` in the signature: the wait is what makes it readable, so the + # dependency is real even though this kernel does not touch the bytes. + if bx >= grid: + inbox[0] = inbox[0] + + return main + + +def rs_sum_kernel( + shard_numel: int, chunks: int, threads: int, world_size: int, local_world_size: int, + dtype: str, elem_lo: int = 0, span_numel: int = 0, +): + """Reduce-scatter tail for one group: our own partial plus this group's arrivals. + + Like ``allreduce_sum_kernel`` it reads ``partial`` as a term instead of copying it into + the inbox first, which saves writing and re-reading a shard of HBM, and it carries no + ``wait_signal`` so its grid is free of the sender's chunk count. + """ + nodes = world_size // local_world_size + span = span_numel or shard_numel + chunk_numel = span // chunks + + @T.prim_func + def main( + partial: T.Tensor((nodes * shard_numel,), dtype), + inbox: T.Tensor((nodes * shard_numel,), dtype), + out: T.Tensor((shard_numel,), dtype), + rank: T.int32, + ): + with T.Kernel(chunks, threads=threads) as bx: + base = elem_lo + bx * chunk_numel + node = rank // local_world_size + for i in T.Parallel(chunk_numel): + out[base + i] = T.cast( + T.cast(partial[node * shard_numel + base + i], "float32") + fp32_sum( + nodes - 1, + lambda k: T.cast( + inbox[((node + k + 1) % nodes) * shard_numel + base + i], + "float32"), + ), + dtype, + ) + + return main + + +def allreduce_sum_kernel( + vec_numel: int, shard_numel: int, chunks: int, threads: int, world_size: int, + local_world_size: int, dtype: str, slot: int, +): + """Sum one node slot of the merged allreduce: our own partial plus the arrivals. + + Two things this deliberately does not do. + + It does **not** copy our own partial into the inbox first, as the two-hop tail did -- + it just reads ``partial`` as one more term. That saves writing and re-reading a whole + shard of HBM per slot. + + It does **not** wait on the signal; ``rail_wait_kernel`` does that separately. Folding + the wait in would tie this grid to the sender's chunk count, and the first version did + exactly that: 4 CTAs for a whole shard of copy-and-sum, which made the merged + allreduce *slower* than composing the two halves (1.559 ms against 1.086). The wait + needs a narrow grid matched to the sender's contexts; the arithmetic wants a wide one. + """ + nodes = world_size // local_world_size + lo = slot * shard_numel + chunk_numel = shard_numel // chunks + + @T.prim_func + def main( + partial: T.Tensor((vec_numel,), dtype), + inbox: T.Tensor((nodes * vec_numel,), dtype), + reduced: T.Tensor((vec_numel,), dtype), + rank: T.int32, + ): + with T.Kernel(chunks, threads=threads) as bx: + base = lo + bx * chunk_numel + node = rank // local_world_size + for i in T.Parallel(chunk_numel): + reduced[base + i] = T.cast( + T.cast(partial[base + i], "float32") + fp32_sum( + nodes - 1, + lambda k: T.cast( + inbox[((node + k + 1) % nodes) * vec_numel + base + i], + "float32"), + ), + dtype, + ) + + return main + + +# ------------------------------------------------------- intra-node via NVSwitch + + +def mc_bcast_kernel( + shard_numel: int, threads: int, world_size: int, local_world_size: int, dtype: str, + slots: str, tiles_per_cta: int = 32, span_numel: int = 0, elem_lo: int = 0, + node_slot: int = -1, +): + """Publish the shards we own into every local rank's slot with one ``multimem.st``. + + ``span_numel``/``elem_lo`` publish only that slice of each shard, for the pipelined + path. Both are compile-time because the multicast region must be provably in bounds, + so a group gets its own compiled kernel rather than an offset argument. + + ``slots="own"`` publishes our own shard, which exists before the collective starts + and so is ordered against nothing. ``slots="remote"`` publishes what arrived over + the fabric. See trap 2 in the module docstring for why the tile width is fixed. + """ + nodes = world_size // local_world_size + groups = 1 if node_slot >= 0 else {"own": 1, "remote": nodes - 1}.get(slots, nodes) + block_N = 2 * threads + span = span_numel or shard_numel + ctas = (span // block_N) // tiles_per_cta + + @T.prim_func + def main( + src: T.Tensor(((1 if slots == "own" else nodes) * shard_numel,), dtype), + out_mc: T.Tensor((world_size * shard_numel,), dtype), + rank: T.int32, + ): + with T.Kernel(groups * ctas, threads=threads) as bx: + c = bx % ctas + k = bx // ctas + local_rank = rank % local_world_size + node = rank // local_world_size + if node_slot >= 0: + n = node_slot # one absolute slot: the per-slot allreduce pipeline + elif slots == "own": + n = node + elif slots == "remote": + n = (node + k + 1) % nodes + else: + n = k # "all": every node slot, for the merged allreduce + src_base = (0 if slots == "own" else n * shard_numel) + elem_lo + dst_base = (n * local_world_size + local_rank) * shard_numel + elem_lo + buf = T.alloc_fragment((block_N,), dtype) + for j in T.serial(tiles_per_cta): + off = (c * tiles_per_cta + j) * block_N + T.copy(src[src_base + off:src_base + off + block_N], buf) + T.multimem_st(buf, out_mc[dst_base + off:dst_base + off + block_N]) + + return main + + +def mc_reduce_kernel( + shard_numel: int, threads: int, world_size: int, local_world_size: int, dtype: str, + slots: str, tiles_per_cta: int = 32, +): + """Sum every local rank's segment for our rail index, reduced by the NVSwitch. + + One ``multimem.ld_reduce`` against the multicast VA returns the sum over every + bound device, so the rank reads only the bytes it keeps instead of pulling all + ``lws`` contributions into a staging buffer and reading them back. That staging + traffic -- ~210 MB of NVLink plus ~500 MB of HBM per rank -- is exactly what kept + the portable path below under torch. + + ``slots="remote"`` handles the node slots whose partials cross the fabric; + ``slots="own"`` handles the one we keep, which no peer waits for. ``"all"`` does both in + a single launch, which is what a small buffer wants: splitting them exists only to + overlap ``own`` with the fabric, and once the transfer is short that overlap is worth + less than the launch it costs. + """ + nodes = world_size // local_world_size + groups = {"remote": nodes - 1, "own": 1}.get(slots, nodes) + block_N = 2 * threads + ctas = (shard_numel // block_N) // tiles_per_cta + + @T.prim_func + def main( + inp_mc: T.Tensor((world_size * shard_numel,), dtype), + partial: T.Tensor((nodes * shard_numel,), dtype), + rank: T.int32, + ): + with T.Kernel(groups * ctas, threads=threads) as bx: + c = bx % ctas + k = bx // ctas + local_rank = rank % local_world_size + node = rank // local_world_size + if slots == "remote": + n = (node + k + 1) % nodes + elif slots == "own": + n = node + else: + n = k # "all": every node slot in one launch + src_base = (n * local_world_size + local_rank) * shard_numel + dst_base = n * shard_numel + acc = T.alloc_fragment((block_N,), dtype) + for j in T.serial(tiles_per_cta): + off = (c * tiles_per_cta + j) * block_N + T.multimem_ld_reduce( + inp_mc[src_base + off:src_base + off + block_N], acc, + reduce_op=T.MultimemReduceOp.ADD, + ) + T.copy(acc, partial[dst_base + off:dst_base + off + block_N]) + + return main + + +# ------------------------------------------------------ intra-node, portable path + + +def pull_bcast_kernel( + shard_numel: int, chunks: int, threads: int, world_size: int, local_world_size: int, + dtype: str, slots: str, +): + """Portable broadcast: read each sibling's buffer instead of publishing to it. + + ``slots="own"`` reads siblings' inputs to fill our own node's output slots; + ``slots="remote"`` reads their outputs for the other nodes' slots. See trap 1 for + why ``src_pe`` must be a global rank. + """ + nodes = world_size // local_world_size + chunk_numel = shard_numel // chunks + groups = 1 if slots == "own" else nodes - 1 + blocks = (local_world_size - 1) * groups * chunks + + @T.prim_func + def main( + shard: T.Tensor((shard_numel,), dtype), + out: T.Tensor((world_size * shard_numel,), dtype), + rank: T.int32, + ): + with T.Kernel(blocks, threads=threads) as bx: + c = bx % chunks + k = (bx // chunks) % groups + step = (bx // chunks) // groups + local_rank = rank % local_world_size + node = rank // local_world_size + # Rotate so concurrent CTAs do not all read the same sibling first. + lp = (local_rank + step + 1) % local_world_size + peer = rank - local_rank + lp + if slots == "own": + T.get_block( + src=T.address_of(shard[c * chunk_numel]), + dst=T.address_of(out[(node * local_world_size + lp) * shard_numel + + c * chunk_numel]), + size=chunk_numel, src_pe=peer, + ) + else: + n = (node + k + 1) % nodes + off = (n * local_world_size + lp) * shard_numel + c * chunk_numel + T.get_block( + src=T.address_of(out[off]), dst=T.address_of(out[off]), + size=chunk_numel, src_pe=peer, + ) + + return main + + +def pull_reduce_kernel( + shard_numel: int, chunks: int, threads: int, world_size: int, local_world_size: int, + dtype: str, slots: str, +): + """Portable reduce: stage every sibling's slice, then sum the slots. + + Needs a ``scratch`` of ``world_size * shard_numel``, which is the cost that keeps + it behind ``mc_reduce_kernel``. Its grid is one CTA per (node slot, chunk), so + ``chunks`` here should be much larger than the rail kernel's -- 16 CTAs on 148 SMs + measures 99.9 GB/s against 293.9 at 1024. + """ + nodes = world_size // local_world_size + chunk_numel = shard_numel // chunks + groups = nodes - 1 if slots == "remote" else 1 + + @T.prim_func + def main( + inp: T.Tensor((world_size * shard_numel,), dtype), + scratch: T.Tensor((world_size * shard_numel,), dtype), + partial: T.Tensor((nodes * shard_numel,), dtype), + rank: T.int32, + ): + with T.Kernel(groups * chunks, threads=threads) as bx: + c = bx % chunks + k = bx // chunks + local_rank = rank % local_world_size + node = rank // local_world_size + node_base = rank - local_rank + n = (node + k + 1) % nodes if slots == "remote" else node + base = c * chunk_numel + # Every sibling holds the slice for our rail index at the same symmetric + # offset, so only the peer changes across the loop. + src_off = (n * local_world_size + local_rank) * shard_numel + base + for step in range(local_world_size): + lp = (local_rank + step) % local_world_size + T.get_block( + src=T.address_of(inp[src_off]), + dst=T.address_of(scratch[(n * local_world_size + lp) * shard_numel + base]), + size=chunk_numel, src_pe=node_base + lp, + ) + # cp_block spreads the copy over the whole CTA, so the sum must not start + # before every thread's share has landed. + T.sync_threads() + for i in T.Parallel(chunk_numel): + partial[n * shard_numel + base + i] = T.cast( + fp32_sum( + local_world_size, + lambda s: T.cast( + scratch[(n * local_world_size + s) * shard_numel + base + i], + "float32"), + ), + dtype, + ) + + return main + + +# ------------------------------------------------------------------ host drivers + + +def workable_chunks(shard_numel, threads, world_size, local_world_size, dtype, + chunks, gin_contexts, log=None): + """Largest chunk count <= ``chunks`` whose put size actually lowers. + + Some ``put_signal`` sizes fail to lower (see the API reference); which ones depends on + ``numel``, so ordinary buffer lengths are otherwise unusable -- 240 MB works at + ``--chunks 8`` while 120 MB does not, its 491520-element put landing in the bad set. + Halving the count doubles the put size and moves off it, and larger puts suit RDMA + anyway. + + Probed with a plain ``tilelang.compile``, not ``ctx.compile``: each node runs a + different interpreter and NCCL, so "every rank rejects the same size" is an assumption + about two toolchains rather than a guarantee, and probing locally keeps the retry loop + from depending on it. + """ + while True: + try: + tilelang.compile(rail_put_kernel(shard_numel, chunks, threads, world_size, + local_world_size, dtype, wait=True)) + return chunks + except Exception as exc: # noqa: BLE001 - only this lowering failure is retried + if "lanes of a scalable" not in str(exc) or chunks <= gin_contexts: + raise + chunks //= 2 + if log: + log(f" put size did not lower; retrying with --chunks {chunks}") + + +def nodes_of(ctx) -> int: + return ctx.world_size // ctx.local_world_size + + +def pick_intra(mode: str) -> str: + """Resolve ``auto`` against the hardware. See Context.supports_multicast.""" + from internode_common import Context + + if mode != "auto": + return mode + return "multimem" if Context.supports_multicast() else "pull" + + +class _Base: + """Shared plumbing: shapes, signal bookkeeping, side stream.""" + + def __init__(self, ctx, shard_numel, torch_dtype, tl_dtype, args, intra, + signal_id=SIGNAL_DATA, signal_state=None): + self.ctx, self.args, self.intra = ctx, args, intra + self.signal_id = signal_id + # GIN signal counters live on the device and keep accumulating for the process + # lifetime, so a second collective reusing a signal id must continue the count + # rather than restart it. A sweep passes one dict through every candidate; on its + # own each collective just owns a private one. + self.signal_state = {} if signal_state is None else signal_state + self.lws = ctx.local_world_size + self.nodes = ctx.world_size // self.lws + self.shard_numel = shard_numel + self.torch_dtype, self.tl_dtype = torch_dtype, tl_dtype + if self.nodes < 2: + raise SystemExit("the 2D collectives need >= 2 nodes") + itemsize = torch.empty((), dtype=torch_dtype).element_size() + # Resolved onto self, never back onto `args`. Several of these knobs are rewritten + # from the requested value (pick_mc_tiles here, workable_chunks and the + # MIN_GROUP_BYTES cap below), and mutating the shared namespace would make a caller + # that builds more than one collective -- an autotuner sweeping candidates, or + # allreduce building both halves -- silently inherit the previous one's rewrites and + # report a config it did not run. + self.mc_tiles = pick_mc_tiles(shard_numel, itemsize, args.mc_tiles) + self.chunks = workable_chunks(shard_numel, args.threads, ctx.world_size, self.lws, + tl_dtype, args.chunks, args.gin_contexts, ctx.log) + if self.chunks % args.gin_contexts: + raise SystemExit( + f"--chunks {self.chunks} must be a multiple of --gin-contexts " + f"{args.gin_contexts}: wait_signal divides the target by the granted " + f"context count and would round it down") + if intra == "multimem": + # self.mc_tiles, not args.mc_tiles: the latter is 0 in auto mode. + unit = 2 * args.mc_threads * self.mc_tiles + if shard_numel % unit: + raise SystemExit( + f"shard {shard_numel} must be a multiple of " + f"2*--mc-threads*--mc-tiles = {unit}") + elif shard_numel % args.intra_chunks: + raise SystemExit( + f"shard {shard_numel} must be a multiple of --intra-chunks " + f"{args.intra_chunks}") + # Cap the pipeline depth by *slice size*, not just by the divisibility rules. + # More groups hide more of the NVLink phases, but each group costs a put, a wait, a + # sum and a publish launch -- roughly 6 launches -- and once a slice is small those + # launches dominate. Measured: allreduce with 8 groups is 0.95 ms at a 15.7 MB + # shard (1.96 MB per slice) and 0.61 ms at a 3.1 MB shard, where the fabric alone + # needs only 0.13 -- i.e. 0.5 ms of pure overhead, and 0.50x torch. + want = args.rail_groups + shard_bytes = shard_numel * itemsize + reason = "" + while want > 1 and shard_bytes // want < MIN_GROUP_BYTES: + want //= 2 + reason = (f"a {shard_bytes // args.rail_groups / 1e6:.2f} MB slice is too small " + f"to pay for its launches") + # A group's grid must be a whole number of chunks and a multiple of the context + # count, so a chunk count reduced by workable_chunks() caps the depth too. + while want > 1 and (self.chunks % want + or (self.chunks // want) % args.gin_contexts): + want //= 2 + reason = reason or (f"--chunks {self.chunks} cannot be split {args.rail_groups} " + f"ways at {args.gin_contexts} contexts") + if want != args.rail_groups: + ctx.log(f" rail-groups {args.rail_groups} -> {want}: {reason}") + self.rail_groups = want + # Only rail peers signal us. + self.per_launch = (self.nodes - 1) * self.chunks + self.side = torch.cuda.Stream() + + def phases(self): + """[(name, callable, nvlink_bytes_per_gpu)] for --phases timing. + + Timing a kernel alone is not the same as its cost inside the loop: the rail + launch measured on its own exposes the full RDMA round trip every iteration + (0.62 ms), where in steady state consecutive iterations overlap and its real + contribution is ~0.39 ms. Subtracting the intra phases from the total is the + honest way to attribute the fabric time; this is for finding which phase to + attack, not for a bandwidth claim. + """ + raise NotImplementedError + + @property + def signals_used(self) -> int: + """How many consecutive GIN signals this collective occupies, from signal_id. + + One per pipelined group. A composition must space its halves by this much: + arrivals are cumulative and a wait does not consume them, so an overlapping range + lets one half's wait be satisfied by the other half's bytes -- which corrupted + exactly the second half of the allreduce output before this existed. + """ + return getattr(self, "groups", 1) + + def _bump(self): + sid = self.signal_id + self.signal_state[sid] = self.signal_state.get(sid, 0) + self.per_launch + return self.signal_state[sid] + + def _rail_args(self, group=0): # noqa: D401 + # Each pipelined group needs its own signal: arrivals are cumulative and a wait + # does not consume them, so a shared signal would let group 0's wait be satisfied + # by group 1's bytes. 32 signals are provisioned. + return (self.shard_numel, self.chunks, self.args.threads, + self.ctx.world_size, self.lws, self.tl_dtype, self.signal_id + group) + + def _mc_args(self): + return (self.shard_numel, self.args.mc_threads, self.ctx.world_size, self.lws, + self.tl_dtype) + + def _pull_args(self, chunks): + return (self.shard_numel, chunks, self.args.threads, self.ctx.world_size, + self.lws, self.tl_dtype) + + +class Allgather2D(_Base): + """``shard`` on every rank -> the concatenation of all shards, on every rank. + + Buffers are owned here: write input into ``.shard`` and read ``.out``. + + What overlaps, and why that is where the speed is + ------------------------------------------------- + A rank's *own* shard already exists, so publishing it is ordered against nothing + and runs on a side stream **concurrently with the fabric hop**. Only the other + nodes' shards depend on the network. Serially instead: 311 GB/s against 404. + + ``pub_remote`` cannot join that overlap -- it publishes bytes that have not arrived + yet -- and at 0.19 ms it is the whole remaining gap to the ~700 GB/s roofline. + Pipelining it into the rail hop is the next optimisation, but the rail chunking has + to stay coarse: chunking it finely to overlap made reduce-scatter 4x *slower*, + because each launch then moved a few MB from a couple of CTAs. + """ + + def __init__(self, ctx, numel, torch_dtype, tl_dtype, args, intra="auto", + signal_id=SIGNAL_DATA, shard=None, buffers=None, signal_state=None): + intra = pick_intra(intra) + if numel % ctx.world_size: + raise SystemExit(f"numel {numel} must be divisible by {ctx.world_size}") + super().__init__(ctx, numel // ctx.world_size, torch_dtype, tl_dtype, args, intra, + signal_id, signal_state) + self.numel = numel + use_mc = intra == "multimem" + + # --- fabric hop --- + # G == 1: put and wait are one kernel; with nothing to overlap, a separate wait + # launch is pure latency. + # G > 1: one put kernel and one wait kernel per group, each on its own signal, so + # group g's arrival can be published over NVLink while group g+1 is still in + # flight. Both nodes issue groups in order on one stream, and a put from sender + # CTA b signals through context b % contexts, so per-QP ordering makes the groups + # arrive in order -- which is what lets an early group be published early. + self.groups = self.rail_groups if use_mc else 1 + per_group = self.chunks // self.groups + if self.groups > 1: + if self.chunks % self.groups or per_group % args.gin_contexts: + raise SystemExit( + f"--chunks {self.chunks} must divide by --rail-groups {self.groups} " + f"into a multiple of --gin-contexts {args.gin_contexts}; got " + f"{per_group} chunks per group") + self.rail_puts = [ + ctx.compile( + rail_put_kernel(*self._rail_args(g), chunk_lo=g * per_group, + chunk_count=per_group), + expect=("tl::gin::put_signal_addr",), gin_contexts=args.gin_contexts, + ) + for g in range(self.groups) + ] + self.rail_waits = [ + ctx.compile( + rail_wait_kernel(*self._rail_args(g), chunk_count=per_group), + expect=("tl::gin::wait_signal",), gin_contexts=args.gin_contexts, + ) + for g in range(self.groups) + ] + else: + self.rail = ctx.compile( + rail_put_kernel(*self._rail_args(), wait=True), + expect=("tl::gin::put_signal_addr", "tl::gin::wait_signal"), + gin_contexts=args.gin_contexts, + ) + self.per_group_signals = (nodes_of(ctx) - 1) * per_group + + if use_mc: + mc = functools.partial(mc_bcast_kernel, *self._mc_args(), + tiles_per_cta=self.mc_tiles) + span = self.shard_numel // self.groups + self.group_numel = span + # One kernel per group: the slice offset must be compile-time, see + # mc_bcast_kernel. Our own shard is sliced too, so a fused producer can publish + # group g while group g's fabric hop is in flight instead of publishing the + # whole shard between the two hops. + self.pub_own_k = ctx.compile(mc(slots="own")) + self.pub_own_ks = [ + ctx.compile(mc(slots="own", span_numel=span, elem_lo=g * span)) + for g in range(self.groups) + ] if self.groups > 1 else [self.pub_own_k] + self.pub_remote_ks = [ + ctx.compile(mc(slots="remote", span_numel=span, elem_lo=g * span)) + for g in range(self.groups) + ] + else: + build = functools.partial(pull_bcast_kernel, *self._pull_args(self.chunks)) + self.pub_own_k = ctx.compile(build(slots="own")) + self.pub_remote_ks = [ctx.compile(build(slots="remote"))] + + # `buffers` lets a caller hand back a previous instance's allocations. No buffer + # shape depends on a knob -- only grids and tile widths do -- so a sweep can + # allocate once and rebuild kernels per candidate. That is not merely an + # optimisation: the arena is a bump allocator with no free and the multicast + # buffer is sized exactly, so a second allocation would exhaust it. + if buffers is not None: + self.__dict__.update(buffers) + self.buffers = buffers + self._use_mc = use_mc + return + # `shard` lets a caller feed an existing arena tensor straight in -- allreduce + # passes the reduce-scatter's output, which avoids a full-shard copy between + # the two halves. + self.shard = ctx.tensor((self.shard_numel,), torch_dtype) if shard is None else shard + if use_mc: + # A GIN put must target the registered arena window, and the output has to + # live in the multicast buffer -- different allocations, so the fabric hop + # lands in `railbuf` and is published from there. + self.out_mc, self.out = ctx.mcast_tensor((numel,), torch_dtype) + self.inbox = ctx.tensor((self.nodes * self.shard_numel,), torch_dtype) + self.inbox.zero_() + else: + self.out = ctx.tensor((numel,), torch_dtype) + self.out_mc = self.inbox = self.out + self.out.zero_() + self._use_mc = use_mc + self.buffers = {k: getattr(self, k) for k in + ("shard", "out", "out_mc", "inbox")} + + def phases(self): + shard_bytes = self.shard_numel * self.shard.element_size() + if self._use_mc: + return [ + ("rail_nic", lambda: self.rail_hop(), 0), + ("pub_own", self.publish_own, shard_bytes * self.lws), + ("pub_remote", lambda: [self.publish_remote(g) for g in range(self.groups)], + shard_bytes * self.lws * (self.nodes - 1)), + ] + return [ + ("rail_nic", lambda: self.rail(self.shard, self.out, self.ctx.rank, self._bump()), 0), + ("pull_own", lambda: self.pub_own_k(self.shard, self.out, self.ctx.rank), + shard_bytes * (self.lws - 1)), + ("pull_remote", lambda: self.pub_remote_ks[0](self.shard, self.out, self.ctx.rank), + shard_bytes * (self.lws - 1) * (self.nodes - 1)), + ] + + # --- steps, exposed so a fused kernel can interleave compute between them --- + # + # launch() below is the plain path. A consumer that wants to compute on the rows it + # already has, while the fabric hop is still running, needs the steps separately: + # see example_internode_ag_gemm_2d.py --mode pipeline. Multimem only, because on the + # pull path publish_own reads siblings' *shards* and publish_remote reads their + # *out*, so neither is a pure local write and the staging is different. + + def rail_hop(self, stream=None): + """Start the fabric hop. + + With one group this also waits, since there is nothing to overlap. With several, + it only *issues* every group's puts -- in group order on one stream, so per-QP + ordering makes them arrive in order -- and the waits are left to + ``consume_groups``, which interleaves them with the NVLink publishes. + """ + self._targets = [self._bump_group(g) for g in range(self.groups)] + run = self._issue_all if self.groups > 1 else self._rail_one + if stream is None: + run() + return + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + + def _rail_one(self): + self.rail(self.shard, self.inbox, self.ctx.rank, self._targets[0]) + + def _issue_all(self): + for g in range(self.groups): + self.issue_group(g, self._targets[g]) + + def bump_groups(self): + """Reserve this launch's signal targets without issuing anything yet.""" + self._targets = [self._bump_group(g) for g in range(self.groups)] + return self._targets + + def issue_group(self, group, target): + """Start group `group`'s fabric put. Its input slice must already be final.""" + self.rail_puts[group](self.shard, self.inbox, self.ctx.rank, target) + + def wait_group(self, group, target): + self.rail_waits[group](self.inbox, target) + + def consume_groups(self): + """Publish each group over NVLink as soon as that group has landed. + + This is where the pipelining pays: group g's ~0.1 ms of multicast broadcast runs + while group g+1 is still crossing the fabric, instead of the whole broadcast + sitting after the whole hop. + """ + if self.groups == 1: + self.publish_remote() + return + for g in range(self.groups): + self.rail_waits[g](self.inbox, self._targets[g]) + self.publish_remote(g) + + def _bump_group(self, group): + sid = self.signal_id + group + self.signal_state[sid] = self.signal_state.get(sid, 0) + self.per_group_signals + return self.signal_state[sid] + + def publish_own(self, group=None): + """Broadcast our own shard to every local rank. Ordered against nothing. + + ``group=None`` publishes the whole shard in one launch, which is what the plain + path wants. A group index publishes just that slice, for the fused allreduce. + """ + if group is None: + self.pub_own_k(self.shard, self.out_mc, self.ctx.rank) + else: + self.pub_own_ks[group](self.shard, self.out_mc, self.ctx.rank) + + def publish_remote(self, group=0): + """Broadcast what arrived over the fabric. Needs that group's wait to have run.""" + self.pub_remote_ks[group](self.inbox, self.out_mc, self.ctx.rank) + + def rows_of_node(self, node, rows_per_rank): + """First output row belonging to `node`. + + Global rank is ``node * lws + local``, and row block index is global rank, so a + node's row blocks are *contiguous* -- which is why the pipelined GEMM needs one + launch per node rather than one per rank. + """ + return node * self.lws * rows_per_rank + + def launch(self): + ctx, args = self.ctx, self.args + main_stream = torch.cuda.current_stream() + if self._use_mc: + if args.no_overlap: + self.rail_hop() + self.publish_own() + else: + # Our own shard exists already, so publishing it races the fabric rather + # than waiting behind it. + self.side.wait_stream(main_stream) + with torch.cuda.stream(self.side): + self.publish_own() + self.rail_hop() + main_stream.wait_stream(self.side) + self.consume_groups() + # We published into every sibling and they into us, so the output is + # complete only once they have all finished. + dist.barrier(ctx.group) + else: + # The pull path needs its own slot present in `out`, and the rail kernel + # writes the sender's global-rank slot, so route it straight there. + target = self._bump() + if args.no_overlap: + self.rail(self.shard, self.out, ctx.rank, target) + dist.barrier(ctx.group) + self.pub_own_k(self.shard, self.out, ctx.rank) + else: + self.side.wait_stream(main_stream) + with torch.cuda.stream(self.side): + self.pub_own_k(self.shard, self.out, ctx.rank) + self.rail(self.shard, self.out, ctx.rank, target) + main_stream.wait_stream(self.side) + dist.barrier(ctx.group) + self.pub_remote_ks[0](self.shard, self.out, ctx.rank) + + +class ReduceScatter2D(_Base): + """Sum of every rank's ``.inp``, restricted to this rank's shard, into ``.out``. + + Needs **no barrier** *provided the input is not written in the same iteration*: the + intra phase reads siblings' input buffers, which the collective itself never writes, + and the rail phase sends only bytes this rank produced, with the GIN signal proving + arrival. Stream order is then sufficient. + + A caller that **produces** ``inp`` each iteration -- a fused GEMM, say -- breaks that + precondition and must fence between producing and launching, because the reduce reads + every local rank's copy and stream order says nothing about a sibling's producer. See + example_internode_gemm_rs_2d.py, which mismatched 11 of 16 ranks without it. + + Overlap runs along the node axis -- reduce the other nodes' slots, start their + transfer, then reduce our own slot, which nothing on the network waits for. That is + worth only ~4%, because the put kernel returns once the RDMA is *issued* and the + flight time was already absorbed by the wait in the tail kernel. + """ + + def __init__(self, ctx, numel, torch_dtype, tl_dtype, args, intra="auto", + signal_id=SIGNAL_DATA): + intra = pick_intra(intra) + if numel % ctx.world_size: + raise SystemExit(f"numel {numel} must be divisible by {ctx.world_size}") + super().__init__(ctx, numel // ctx.world_size, torch_dtype, tl_dtype, args, intra, + signal_id) + self.numel = numel + use_mc = intra == "multimem" + + if use_mc: + build = functools.partial(mc_reduce_kernel, *self._mc_args(), + tiles_per_cta=self.mc_tiles) + else: + build = functools.partial(pull_reduce_kernel, *self._pull_args(args.intra_chunks)) + # With one group the fabric hop is a single transfer, so there is nothing for the + # own-slot reduce to hide behind and splitting it just costs a launch. + # self.rail_groups is the depth after the size cap, not what was asked for. + self.fused_reduce = self.rail_groups == 1 and use_mc + if self.fused_reduce: + self.red_all = ctx.compile(build(slots="all")) + else: + self.red_remote = ctx.compile(build(slots="remote")) + self.red_own = ctx.compile(build(slots="own")) + # Grouped fabric hop, same idea as Allgather2D: group g's arithmetic runs while + # group g+1 is still crossing the fabric. Per-put size is unchanged. + self.groups = self.rail_groups + per_group = self.chunks // self.groups + if self.chunks % self.groups or per_group % args.gin_contexts: + raise SystemExit( + f"--chunks {self.chunks} must divide by --rail-groups {self.groups} into a " + f"multiple of --gin-contexts {args.gin_contexts}; got {per_group}") + span = self.shard_numel // self.groups + self.put_k, self.wait_k, self.sum_k = [], [], [] + for g in range(self.groups): + self.put_k.append(ctx.compile( + rail_put_kernel(*self._rail_args(g), src_per_node=True, + chunk_lo=g * per_group, chunk_count=per_group), + expect=("tl::gin::put_signal_addr",), gin_contexts=args.gin_contexts)) + self.wait_k.append(ctx.compile( + rail_wait_kernel(*self._rail_args(g), chunk_count=per_group), + expect=("tl::gin::wait_signal",), gin_contexts=args.gin_contexts)) + self.sum_k.append(ctx.compile( + rs_sum_kernel(self.shard_numel, args.intra_chunks // self.groups, + args.threads, ctx.world_size, self.lws, tl_dtype, + elem_lo=g * span, span_numel=span))) + self.per_group_signals = (self.nodes - 1) * per_group + + if use_mc: + self.inp_mc, self.inp = ctx.mcast_tensor((numel,), torch_dtype) + self.scratch = None + else: + self.inp = ctx.tensor((numel,), torch_dtype) + self.inp_mc = self.inp + self.scratch = ctx.tensor((numel,), torch_dtype) + self.partial = ctx.tensor((self.nodes * self.shard_numel,), torch_dtype) + self.inbox = ctx.tensor((self.nodes * self.shard_numel,), torch_dtype) + self.out = ctx.tensor((self.shard_numel,), torch_dtype) + for t in (self.scratch, self.partial, self.inbox, self.out): + if t is not None: + t.zero_() + self._use_mc = use_mc + + def _reduce(self, kernel): + if self._use_mc: + kernel(self.inp_mc, self.partial, self.ctx.rank) + else: + kernel(self.inp, self.scratch, self.partial, self.ctx.rank) + + def phases(self): + shard_bytes = self.shard_numel * self.out.element_size() + return [ + ("rail_put", lambda: self.put_k[0](self.partial, self.inbox, self.ctx.rank, + self._bump()), 0), + ("red_remote", lambda: self._reduce(self.red_remote), + shard_bytes * self.lws * (self.nodes - 1)), + ("red_own", lambda: self._reduce(self.red_own), shard_bytes * self.lws), + ] + + def start(self): + """Reduce over NVLink and get the fabric moving; return the per-group targets. + + Split out of ``launch`` so a consumer can interleave work with the per-group + finishes -- the fused allreduce starts the allgather's hop for group g the moment + this rank's group g is summed, instead of after every group. + """ + ctx, args = self.ctx, self.args + main_stream = torch.cuda.current_stream() + targets = [self._bump_group(g) for g in range(self.groups)] + issue = lambda: [self.put_k[g](self.partial, self.inbox, ctx.rank, targets[g]) + for g in range(self.groups)] + if self.fused_reduce: + self._reduce(self.red_all) + issue() + return targets + self._reduce(self.red_remote) + if args.no_overlap: + self._reduce(self.red_own) + issue() + else: + ready = torch.cuda.Event() + ready.record(main_stream) + self.side.wait_event(ready) + with torch.cuda.stream(self.side): + issue() + # Nothing on the network waits for our own slot, so it reduces in flight. + self._reduce(self.red_own) + main_stream.wait_stream(self.side) + return targets + + # --- staged entry points, for a producer that wants to interleave compute --- + # + # A fused GEMM can compute the rows one node slot needs, hand that slot to the fabric, + # and compute the next slot's rows while the first is in flight. See + # example_internode_gemm_rs_2d.py --mode pipeline. + + def reduce_remote(self): + """NVSwitch-reduce the node slots whose partials cross the fabric.""" + self._reduce(self.red_remote) + + def reduce_own(self): + """NVSwitch-reduce the slot we keep. Nothing on the network waits for it.""" + self._reduce(self.red_own) + + def issue_puts(self): + """Send every group's partial for the remote slots; returns the per-group targets.""" + targets = [self._bump_group(g) for g in range(self.groups)] + for g in range(self.groups): + self.put_k[g](self.partial, self.inbox, self.ctx.rank, targets[g]) + return targets + + def finish_group(self, group, target): + """Wait for group `group`'s arrivals and sum it into ``out``.""" + self.wait_k[group](self.inbox, target) + self.sum_k[group](self.partial, self.inbox, self.out, self.ctx.rank) + + def launch(self): + for g, target in enumerate(self.start()): + self.finish_group(g, target) + + +class Allreduce2D(_Base): + """Sum of every rank's ``.inp``, on every rank, into ``.out`` -- in one fabric hop. + + Composing ReduceScatter2D with Allgather2D is correct and simple, and it measured + 1.088 ms against torch's 0.949 (0.87x). The two halves do not overlap at all -- the + total is exactly their sum -- because the allgather cannot start until the + reduce-scatter has produced the shard it broadcasts. + + So merge them. The insight is that the rail peer can send its partials for **every** + node slot rather than only the one this rank owns, which lets this rank finish all + slots locally and removes the second hop entirely: + + 1. ``mc_reduce`` -- NVSwitch-reduce over local ranks, giving this rank a partial for + each of the ``nodes`` slots at its own rail index. + 2. ``rail`` -- one hop, carrying that whole ``nodes * shard`` vector. + 3. ``tail`` -- add our own vector, wait, sum the ``nodes`` arrivals: now every slot at + our rail index holds the global sum. + 4. ``mc_bcast(slots="all")`` -- publish all of them to every local rank. + + Fabric bytes are identical to the two-hop version, and the floor is the same: a 2-node + allreduce must move the node-sum each way, ``N/lws`` per rank, so 31.4 MB at 47.6 GB/s + = 0.66 ms at this size. What the merge buys is one serialisation instead of two, one + barrier instead of three, and far fewer launches. + """ + + def __init__(self, ctx, numel, torch_dtype, tl_dtype, args, intra="auto", + signal_id=SIGNAL_DATA): + intra = pick_intra(intra) + if intra != "multimem": + raise SystemExit("the merged allreduce needs multimem; use --algo composed") + if numel % ctx.world_size: + raise SystemExit(f"numel {numel} must be divisible by {ctx.world_size}") + super().__init__(ctx, numel // ctx.world_size, torch_dtype, tl_dtype, args, intra, + signal_id) + self.numel = numel + vec = self.nodes * self.shard_numel + shard = self.shard_numel + cps = self.chunks // self.nodes # chunks per node slot + if self.chunks % self.nodes or cps % args.gin_contexts: + raise SystemExit( + f"--chunks {self.chunks} must split into {self.nodes} slots of a multiple " + f"of --gin-contexts {args.gin_contexts}; got {cps} per slot") + + mc = functools.partial(mc_reduce_kernel, *self._mc_args(), + tiles_per_cta=self.mc_tiles) + self.red_remote = ctx.compile(mc(slots="remote")) + self.red_own = ctx.compile(mc(slots="own")) + + # One put / tail / publish per node slot, each on its own signal. Slot m occupies + # a contiguous chunk range of the vector, so the existing chunk_lo/chunk_count + # slicing covers it. The offsets have to be compile-time (multimem needs a + # provably in-bounds region), so this is one kernel per absolute slot and the host + # picks which to call -- our own node index is known there. + rail = (vec, self.chunks, args.threads, ctx.world_size, self.lws, tl_dtype) + self.put_k, self.wait_k, self.sum_k, self.pub_k = [], [], [], [] + for m in range(self.nodes): + self.put_k.append(ctx.compile( + rail_put_kernel(*rail, signal_id + m, chunk_lo=m * cps, chunk_count=cps), + expect=("tl::gin::put_signal_addr",), gin_contexts=args.gin_contexts)) + self.wait_k.append(ctx.compile( + rail_wait_kernel(*rail, signal_id + m, chunk_count=cps), + expect=("tl::gin::wait_signal",), gin_contexts=args.gin_contexts)) + self.sum_k.append(ctx.compile( + allreduce_sum_kernel(vec, shard, args.intra_chunks, args.threads, + ctx.world_size, self.lws, tl_dtype, m))) + self.pub_k.append(ctx.compile( + mc_bcast_kernel(*self._mc_args(), slots="all", + tiles_per_cta=self.mc_tiles, node_slot=m))) + + self.inp_mc, self.inp = ctx.mcast_tensor((numel,), torch_dtype) + self.out_mc, self.out = ctx.mcast_tensor((numel,), torch_dtype) + self.partial = ctx.tensor((vec,), torch_dtype) + self.inbox = ctx.tensor((self.nodes * vec,), torch_dtype) + self.reduced = ctx.tensor((vec,), torch_dtype) + for t in (self.partial, self.inbox, self.reduced, self.out): + t.zero_() + self.per_slot_signals = (self.nodes - 1) * cps + self._stargets = [0] * self.nodes + self.my_node = ctx.rank // self.lws + # Send the slots we can send first, then ours: red_own is still writing our own + # slot while the first put is in flight, and sending it early was a real race -- + # it corrupted exactly the second half of the output. + self.slot_order = [m for m in range(self.nodes) if m != self.my_node] + [self.my_node] + + def launch(self): + ctx = self.ctx + main_stream = torch.cuda.current_stream() + targets = [] + for m in range(self.nodes): + self._stargets[m] += self.per_slot_signals + targets.append(self._stargets[m]) + + # Reduce the slots the fabric needs first, so the hop starts as early as possible. + self.red_remote(self.inp_mc, self.partial, ctx.rank) + for m in self.slot_order[:-1]: + self.put_k[m](self.partial, self.inbox, ctx.rank, targets[m]) + # Our own slot: only now may it be sent. + self.red_own(self.inp_mc, self.partial, ctx.rank) + self.put_k[self.my_node](self.partial, self.inbox, ctx.rank, targets[self.my_node]) + + # Per-slot pipeline: sum and publish each slot as its arrival lands, in the order + # they were sent, so slot A's NVLink publish overlaps slot B's transfer. + for m in self.slot_order: + self.wait_k[m](self.inbox, targets[m]) + self.sum_k[m](self.partial, self.inbox, self.reduced, ctx.rank) + self.pub_k[m](self.reduced, self.out_mc, ctx.rank) + # We published into every sibling and they into us. + dist.barrier(ctx.group) + + +def report_phases(ctx, coll, args): + """Time each kernel of a 2D collective on its own. See ``_Base.phases``.""" + from tilelang.distributed.bench import do_bench + + for name, fn, nvlink_bytes in coll.phases(): + dist.barrier(ctx.group) + ms = do_bench(fn, warmup=args.warmup, rep=args.rep, group=ctx.group) + if nvlink_bytes: + ctx.log(f" {name:12s} {ms:7.3f} ms {nvlink_bytes / 1e6 / ms:6.1f} GB/s " + f"NVLink/GPU") + else: + # One shard per rank crosses the fabric, on that rank's own NIC. + mb = coll.shard_numel * coll.out.element_size() / 1e6 + ctx.log(f" {name:12s} {ms:7.3f} ms {mb / ms:6.1f} GB/s per NIC " + f"(isolated; overstates -- see phases())") + + +def fused_allreduce_launch(rs, ag, ctx): + """Reduce-scatter and allgather with their fabric hops overlapped. + + Composed serially the two halves add up exactly -- 0.550 + 0.500 ms -- because the + allgather cannot start until the reduce-scatter has produced the shard it broadcasts. + But that is only true *per group*: the allgather's hop for group g needs nothing but + the reduce-scatter's sum for group g. So push each group across the fabric as soon as + it is summed, and the second hop overlaps the first's remaining groups instead of + following all of them. + + The two halves must already hold disjoint signal ranges; see the example. + """ + if ag.groups == 1: + # Nothing to interleave, so take the fewest launches: the allgather's put and wait + # are one kernel here, and the publishes are whole-shard. At small sizes this path + # is what wins -- the pipeline's extra launches cost more than its overlap saves. + for g, target in enumerate(rs.start()): + rs.finish_group(g, target) + ag.rail_hop() + ag.publish_own() + ag.publish_remote() + dist.barrier(ctx.group) + return + rs_targets = rs.start() + ag_targets = ag.bump_groups() + for g, target in enumerate(rs_targets): + rs.finish_group(g, target) + # out slice g is final, and ag.shard *is* rs.out, so this group can fly now -- + # and publishing our own copy of that slice needs no network at all, so it runs + # while the slice is crossing the fabric rather than between the two hops. + ag.issue_group(g, ag_targets[g]) + ag.publish_own(g) + for g, target in enumerate(ag_targets): + ag.wait_group(g, target) + ag.publish_remote(g) + dist.barrier(ctx.group) + + +SWEEPABLE = ("chunks", "rail_groups", "gin_contexts", "mc_threads", "mc_tiles", + "intra_chunks", "threads") + + +def parse_sweep(spec): + """``"chunks=4,8;mc_tiles=8,16"`` -> the 4-element cartesian product, as dicts. + + Tuples rather than one knob at a time, because the knobs do not compose: at 32 MB + ``--chunks 4`` alone measured 209.4 GB/s and 185.2 combined with ``--mc-tiles 8``, so a + coordinate-descent sweep converges on the wrong point. See + docs/internode_optimization_log.md. + """ + import itertools + + axes = [] + for part in (p for p in spec.split(";") if p.strip()): + key, _, vals = part.partition("=") + key = key.strip().replace("-", "_") + if key not in SWEEPABLE: + raise SystemExit(f"--sweep: {key!r} is not tunable; pick from {SWEEPABLE}") + axes.append([(key, int(v)) for v in vals.split(",") if v.strip()]) + return [dict(combo) for combo in itertools.product(*axes)] if axes else [{}] + + +def run_sweep(ctx, args, make, verify, launch_of, run_ref, moved, name, buffers=None): + """Time every candidate in ``args.sweep`` in one process, then rank them. + + One process matters: start-up is ~19 s per rank, 6.4 s of it ncclDevCommCreate which does + not shrink, so a process per candidate would spend nearly all its time initialising. + + **Candidates are measured round-robin over several passes, and each keeps its best pass.** + Measuring each candidate once in sequence does not work: position dominates the result. + Sweeping tiles 8/16/32 ranked 8 first; reversing to 32/16/8 ranked 32 first; and the + control -- the same configuration three times -- degraded monotonically, 1.73x then 1.55x + then 1.44x. That is the GPUs' clocks drooping under sustained load, and it is larger than + the differences being compared. Round-robin spreads each candidate across the droop curve + and the min across passes takes each one's least-throttled sample. + + Buffers come from the instance the caller already built; without that, candidate 1 + allocates a second time and the exactly-sized multicast buffer is exhausted. + + Every rank walks the same list in the same order, and validity is a deterministic function + of the candidate, so a rejected candidate is rejected identically everywhere -- which keeps + the collective compiles and barriers in lockstep. An unbuildable candidate is skipped, not + fatal. + """ + from tilelang.distributed.bench import do_bench + + cands = parse_sweep(args.sweep) + passes = max(1, args.sweep_passes) + ctx.log(f"{name}: {len(cands)} candidate(s) x {passes} pass(es), round-robin, " + f"in one process") + signal_state = {} + built, skipped = [], [] + for over in cands: + cand = argparse.Namespace(**vars(args)) + for k, v in over.items(): + setattr(cand, k, v) + try: + coll = make(cand, buffers, signal_state) + except SystemExit as exc: # invalid tuple: same verdict on every rank + skipped.append((over, str(exc))) + continue + if buffers is None: + buffers = coll.buffers + built.append({"over": over, "coll": coll, "launch": launch_of(coll), + "eff": (f"chunks={coll.chunks} groups={coll.groups} " + f"ctx={cand.gin_contexts} mc_threads={cand.mc_threads} " + f"mc_tiles={coll.mc_tiles}"), + "ms": float("inf"), "ref": float("inf"), "bad": 0}) + for over, why in skipped: + ctx.log(f" skipped {over}: {why}") + + # Correctness once per candidate; timing round-robin so no candidate owns a position. + for b in built: + torch.cuda.synchronize() + dist.barrier(ctx.group) + b["launch"]() + torch.cuda.synchronize() + b["bad"] = verify(b["coll"], f"{name}{b['over']}") + built = [b for b in built if not b["bad"]] + + for p in range(passes): + for b in built: + pre = do_bench(run_ref, warmup=args.warmup, rep=args.rep, group=ctx.group) + ms = do_bench(b["launch"], warmup=args.warmup, rep=args.rep, group=ctx.group) + post = do_bench(run_ref, warmup=args.warmup, rep=args.rep, group=ctx.group) + if ms < b["ms"]: + b["ms"], b["ref"] = ms, min(pre, post) + ctx.log(f" pass {p + 1} {b['over']}: {ms:7.3f} ms " + f"{moved / (ms * 1e-3) / 1e9:6.1f} GB/s drift " + f"{abs(pre - post) / min(pre, post) * 100:4.1f}%") + dist.barrier(ctx.group) + + if ctx.is_leader and built: + built.sort(key=lambda b: b["ms"]) + print(f"\n===== {name}: {len(built)}/{len(cands)} candidates, best of " + f"{passes} passes =====", flush=True) + for b in built: + print(f" {b['ms']:7.3f} ms {moved / (b['ms'] * 1e-3) / 1e9:6.1f} GB/s " + f"{b['ref'] / b['ms']:5.2f}x {b['over'] or 'defaults'} [{b['eff']}]", + flush=True) + best = built[0] + print(f" best: {best['over'] or 'defaults'} at " + f"{best['ref'] / best['ms']:.2f}x torch", flush=True) + spread = built[-1]["ms"] / built[0]["ms"] - 1 + if spread < 0.05: + print(f" NOTE: spread is only {spread * 100:.1f}%, within the noise floor -- " + f"treat these as tied", flush=True) + return built + + +def add_2d_args(parser): + """Knobs shared by every 2D example. Defaults are the measured optima at 16 GPUs.""" + parser.add_argument("--intra", choices=("multimem", "pull", "auto"), default="auto", + help="intra-node half: NVSwitch multimem, the portable " + "get_block pull, or multimem when the hardware allows it") + parser.add_argument("--mc-threads", type=int, default=512, + help="threads per CTA on the multimem path; the tile is " + "2*threads, fixed by the packed-x2 fragment layout") + parser.add_argument("--mc-tiles", type=int, default=0, + help="contiguous tiles each multimem CTA loops over; the tile width " + "is pinned, so this is the only work-per-thread knob. 0 scales " + "it to the shard, which matters: the best value is 32 at " + "240 MiB and 4 at 48 MiB") + parser.add_argument("--intra-chunks", type=int, default=1024, + help="chunking of the pull path's intra phase; sets its grid, " + "and is independent of --chunks because it carries no signal") + # 2 measured best on allgather: 472.8 GB/s against 395.2 unsplit (1.53x torch vs + # 1.28x). 4 needs --gin-contexts 2 to keep a group's grid a multiple of the context + # count, and losing contexts costs more than the extra group gains (434.0). Raising + # --chunks to 16 to keep 4 groups at 4 contexts hits the put-size lowering bug. + parser.add_argument("--rail-groups", type=int, default=2, + help="split the fabric hop into this many groups so each group's " + "NVLink publish overlaps the next group's transfer; the " + "per-put size is unchanged, only per-group parallelism drops") + parser.add_argument("--no-overlap", action="store_true", + help="run the phases serially, to show what the overlap buys") + parser.add_argument("--sweep", default=None, metavar="SPEC", + help='tune in one process, e.g. "mc_tiles=8,16,32;chunks=4,8". ' + "Sweeps the cartesian product because the knobs do not " + f"compose. Tunable: {', '.join(SWEEPABLE)}") + parser.add_argument("--sweep-passes", type=int, default=3, + help="round-robin passes over the candidates; each keeps its best. " + "More than one is required, not optional: clock droop makes a " + "single sequential pass rank by position rather than by config") + # add_common_args defaults --chunks to 64, which suits the flat collectives. The + # rail kernel here wants far fewer, larger messages: at 64 it fails to lower with + # "Can't fetch the lanes of a scalable vector", and in the single-node proxy + # gemm_rs_2d returned a *wrong* answer on one rank rather than erroring. Root cause + # not yet found, so default to the tuned value and treat large chunk counts as + # unsupported here. + parser.set_defaults(chunks=8) + return parser diff --git a/examples/distributed/internode/internode_common.py b/examples/distributed/internode/internode_common.py new file mode 100644 index 0000000000..4b5add384e --- /dev/null +++ b/examples/distributed/internode/internode_common.py @@ -0,0 +1,366 @@ +"""Shared setup for the inter-node GIN collectives. + +Every example here follows the same shape: bring up a process group with the +network enabled, allocate symmetric buffers from the TileScale allocator (GIN can +only address the registered arena), run a kernel that moves bytes with +``T.nccl_gin``, check the result against torch, and time both. + +The pieces live here rather than in each example because the environment setup is +easy to get subtly wrong -- ``init_dist`` disables InfiniBand by default, and a +cudaMalloc arena cannot be registered as an NCCL window -- and a silent mistake +in either shows up as a passing test that moved no data over the fabric. +""" + +from __future__ import annotations + +import argparse +import functools +import operator +import os + +import torch +import torch.distributed as dist + +import tilelang +import tilelang.language as T + +# Signal slots. The reduce-scatter half of allreduce and its allgather half must +# not share a slot: signals are cumulative totals, so two phases counting into +# one slot cannot be told apart. +SIGNAL_DATA = 0 +SIGNAL_PHASE2 = 1 + + +def fp32_sum(count: int, term): + """Fold ``term(0) + ... + term(count-1)`` into one add-expression. + + ``term`` maps a source index to a PrimExpr; callers cast to float32 inside it, + because a bf16 running sum over 16 ranks loses enough low bits to fail a + tolerance check. + + Both the loop and the fold live here, outside any traced function, and that + placement is the point. Two shapes that look more natural both fail inside a + ``T.prim_func``: + + * ``acc = T.cast(...)`` then ``acc = acc + ...`` -- in a kernel body the eager + builder treats assignment of a PrimExpr as a TIR variable *declaration*, so + the reassignment emits a second variable (``acc_1``) and the enclosing + ``T.Parallel`` frame rejects it. + * ``fp32_sum([... for src in range(n)])`` -- the builder rewrites ``for`` + statements *and comprehension for-clauses* into TIR loops, so the + comprehension raises ``'ForFrame' object is not iterable``. + + Calling a plain Python helper sidesteps both: the fold runs at trace time and + only its result, a single expression, is emitted. That keeps the copy + vectorisable. + """ + return functools.reduce(operator.add, (term(k) for k in range(count))) + + +def prepare_env() -> None: + """Set the environment GIN needs, before ``init_dist`` is imported or run. + + ``init_dist`` sets ``NCCL_IB_DISABLE=1`` unless it is already set, which would + keep every transfer inside shared memory and make an "inter-node" benchmark + measure nothing. The VMM and GIN flags are hard requirements of window + registration; asserting them here turns a missing Device API into an error at + startup instead of a null devcomm read inside a kernel. + """ + os.environ["NCCL_IB_DISABLE"] = "0" + os.environ.setdefault("TILESCALE_USE_VMM", "1") + os.environ.setdefault("TILESCALE_USE_GIN", "1") + os.environ.setdefault("NCCL_DEBUG", "ERROR") + + +def add_common_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: + parser.add_argument( + "--numel", + type=int, + default=1 << 22, + help="elements in the collective's logical input, summed over ranks", + ) + parser.add_argument("--block", type=int, default=8192, help="elements per CTA") + # Chunks exist to spread a peer's transfer across GIN contexts (one context + # is one QP per peer, so one NIC), NOT to parallelise within a channel. + # Splitting further than the context count only shrinks messages. + parser.add_argument( + "--chunks", + type=int, + default=64, + # Tuned on two idle nodes, 64 MB shard: the reduce variants climb with + # chunk count and peak at 64 (allreduce 41.5 -> 47.2 GB/s from 4 to 64, + # reduce_scatter 44.0 -> 46.2 from 16 to 64; 128 is worse), because their + # CTAs also carry the reduction and want more of them. Allgather is flat + # from 2 to 64, so 64 is a safe shared default. + help="puts per peer; each becomes one CTA issuing one large put", + ) + parser.add_argument( + "--gin-contexts", + type=int, + default=4, + # Contexts are insurance against a busy NIC, not a win on an idle one. + # On two *idle* nodes one context already reaches line rate and extra + # contexts cost ~1% (allgather 47.6 GB/s at 1 vs 47.0 at 4). On a NIC + # shared with another tenant's job, one context collapsed to 23 GB/s while + # 2-4 held ~44. Defaulting to 4 trades that 1% for the contended case. + # The device clamps to what the devcomm granted (4 here, though the + # allocator asks for 8) and scales the wait target to match. + help="-DTL_GIN_CONTEXTS: spread CTAs over n GIN contexts (QPs); 1 pins to context 0", + ) + parser.add_argument( + "--signal-div", + type=int, + default=0, + help="DEBUG ONLY: divide the wait target; under-waits, so any result is invalid", + ) + parser.add_argument( + "--wait-ctx0", + action="store_true", + help="-DTL_GIN_WAIT_CTX0: puts spread over contexts, every wait on context 0", + ) + # 1024 threads matches triton-dist's num_warps=32 for its inter-node send + # blocks: one CTA cooperatively driving one large put. + parser.add_argument("--threads", type=int, default=1024) + parser.add_argument("--dtype", choices=("fp32", "bf16", "fp16"), default="bf16") + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--rep", type=int, default=50) + parser.add_argument("--no-bench", action="store_true", help="check correctness only") + parser.add_argument("--print-source", action="store_true") + parser.add_argument( + "--tune", + action="store_true", + help="sweep the grid below in one process, verify each, report the best vs torch", + ) + parser.add_argument("--tune-chunks", default="2,4,8,16", help="--tune: chunk counts to try") + parser.add_argument("--tune-contexts", default="1,2,4", help="--tune: GIN context counts") + parser.add_argument("--tune-threads", default="1024", help="--tune: threads per CTA") + return parser + + +# GIN_SIGNAL_COUNT in nccl_window.py. Signals are "guaranteed to start at id=0", +# so ids [0, 32) are usable. +MAX_SIGNALS = 32 + + + + + + + + +TORCH_DTYPES = {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16} +TL_DTYPES = {"fp32": "float32", "bf16": "bfloat16", "fp16": "float16"} + + +class Context: + """Process group, allocator and topology for one rank.""" + + def __init__(self, arena_bytes: int = 1 << 30, mcast_bytes: int = 0): + import time + + from tilelang.distributed.host import init_dist + + self._t0 = time.perf_counter() + + self.local_rank = int(os.environ.get("LOCAL_RANK", 0)) + self.local_world_size = int( + os.environ.get("LOCAL_WORLD_SIZE", torch.cuda.device_count()) + ) + # Staged so a hang can be attributed. init_dist, allocator construction + # (which creates the GIN devcomm and registers the arena window) and + # compile are all collective; without markers they are one opaque block. + if os.environ.get("TL_STAGE_TRACE"): + print(f"[rank?] + 0.00s init_dist: enter local_rank={self.local_rank}", + flush=True) + self.rank, self.world_size, self.group, self.node_info = init_dist( + self.local_rank, self.local_world_size, return_node_info=True + ) + self.trace("init_dist: done") + self.num_nodes = self.node_info.num_nodes if self.node_info is not None else 1 + self.trace( + f"allocator: enter bytes={arena_bytes} nodes={self.num_nodes} " + f"world={self.world_size} (devcomm + window register)" + ) + # A non-zero mcast_bytes adds an NVSwitch multicast buffer, which is a + # *separate* allocation from the arena: only the arena is a GIN window, so + # anything a GIN put reads must still come from ctx.tensor(). + self.allocator = tilelang.get_allocator( + size=arena_bytes, + device="cuda", + is_distributed=True, + local_rank=self.local_rank, + num_local_ranks=self.local_world_size, + group=self.group, + node_info=self.node_info, + **({"mcast_size": mcast_bytes} if mcast_bytes else {}), + ) + self.trace("allocator: done (arena window live)") + + @property + def is_leader(self) -> bool: + return self.rank == 0 + + def tensor(self, shape, dtype: torch.dtype): + """Allocate from the arena. Required: only the arena is a GIN window.""" + return tilelang.tensor(shape, dtype, allocator=self.allocator) + + def mcast_tensor(self, shape, dtype: torch.dtype): + """Allocate from the multicast buffer; needs ``mcast_bytes`` at construction. + + Returns ``(mc, local)``. Pass ``mc`` to a kernel using ``T.multimem_*`` -- it + is the multicast VA, so one instruction reaches every local rank and the + NVSwitch does the fan-in or fan-out. Write payload through ``local``, which + is this rank's own physical view, and read it back the same way. + """ + return self.allocator._allocate_mcast_tensor(tuple(shape), dtype) + + @staticmethod + def supports_multicast() -> bool: + """Whether an NVSwitch multicast object can be created on this device. + + The probe calls ``cuCtxGetDevice``, so it reports False with no current + context -- and this is normally called *before* ``Context()`` exists, to size + the multicast buffer. Establishing the context first is therefore part of the + check, not incidental to it: without the touch below the answer is a silent + False on hardware that fully supports multicast. + """ + from tilelang.distributed.shared_memory import _supports_multicast + + torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0))) + torch.zeros(1, device="cuda") + return bool(_supports_multicast()) + + def log(self, msg: str) -> None: + if self.is_leader: + print(msg, flush=True) + + def trace(self, msg: str) -> None: + """Print from every rank, tagged. + + ``log`` is leader-only, so a non-leader rank cannot report progress at + all: it looks identical whether it hung in setup, hung in compile, or + ran fine. That ambiguity is what made the first two-node hang + undiagnosable. Enabled by TL_STAGE_TRACE=1 to keep normal runs quiet. + """ + if os.environ.get("TL_STAGE_TRACE"): + import time + + dt = time.perf_counter() - getattr(self, "_t0", time.perf_counter()) + print(f"[rank{getattr(self, 'rank', '?')}] +{dt:6.2f}s {msg}", flush=True) + + def compile(self, func, *, expect: tuple[str, ...] = (), gin_contexts: int | None = None, + wait_ctx0: bool = False): + # Every rank checks the tokens, not just the leader. compile_once makes + # this a collective, so a leader-only assertion aborts rank 0 while the + # others march on into close()'s barrier and hang until the outer + # timeout -- turning a clear assertion failure into a mystery stall. + # + # gin_contexts becomes -DTL_GIN_CONTEXTS=n. compile_flags is part of the + # cache key, so each setting gets its own cache entry -- which is also + # the only thing that keeps a sweep honest, since the key does not cover + # the device headers this define lives in. + flags = None if gin_contexts is None else [f"-DTL_GIN_CONTEXTS={int(gin_contexts)}"] + if wait_ctx0: + flags = (flags or []) + ["-DTL_GIN_WAIT_CTX0=1"] + if os.environ.get("TL_GIN_DEBUG"): + flags = (flags or []) + ["-DTL_GIN_DEBUG=1"] + self.trace(f"compile: enter (collective in compile_once) flags={flags}") + kernel = tilelang.compile( + func, compile_once=True, compile_group=self.group, compile_flags=flags + ) + self.trace("compile: lowered") + if expect: + source = kernel.get_kernel_source() + # Without this the kernel still compiles and silently moves nothing, + # which would read as a fast and correct-looking result. + for token in expect: + assert token in source, f"lowering did not emit {token!r}" + assert "nccl_gin.h" in source, "generated code is missing the GIN header" + kernel.initialize(allocator=self.allocator) + self.trace("compile: initialized") + return kernel + + def close(self) -> None: + # allocator.close() is collective, so every rank must reach it even if + # this rank's check failed. Guard the barrier: if a peer already died, + # blocking here forever converts its error into a timeout on this rank + # and buries the real message. + self.trace("close: barrier") + try: + dist.barrier(self.group) + except Exception as exc: # noqa: BLE001 - report and keep tearing down + print(f"[rank{self.rank}] close: barrier failed: {exc}", flush=True) + self.trace("close: allocator") + self.allocator.close() + dist.destroy_process_group() + self.trace("close: done") + + +def check(ctx: Context, got: torch.Tensor, want: torch.Tensor, name: str) -> int: + """Compare on every rank and aggregate, so one bad rank fails the run.""" + # bf16 accumulation order differs between a tree reduction and our linear + # one, so compare with a tolerance rather than exactly. + if got.dtype in (torch.bfloat16, torch.float16): + ok = torch.allclose(got.float(), want.float(), rtol=6e-2, atol=6e-2) + else: + ok = torch.allclose(got, want, rtol=1e-5, atol=1e-5) + if not ok: + diff = (got.float() - want.float()).abs() + bad = (diff > 6e-2).nonzero().flatten() + print( + f"[rank {ctx.rank}] {name} MISMATCH: {bad.numel()}/{got.numel()} differ, " + f"max |diff| {diff.max().item():.4g}, first at {bad[0].item() if bad.numel() else -1}", + flush=True, + ) + status = torch.tensor([0 if ok else 1], device=got.device, dtype=torch.int32) + dist.all_reduce(status, group=ctx.group) + failures = int(status.item()) + if failures == 0: + ctx.log(f"{name}: correct on all {ctx.world_size} ranks") + return failures + + +def bench_vs_torch(ctx: Context, args, name: str, launch, run_ref, moved: int, + tflops: float = 0.0) -> None: + """Time ours against torch, with torch measured *both* before and after. + + On a shared cluster torch's NCCL reading is cold on its first use and drifts with + other tenants' traffic on the same NICs -- by more than the effect being measured. + The first configuration of a batched sweep reads 131-186 GB/s where every later one + reads ~309. So both readings are printed: if they disagree, discard the ratio. + """ + from tilelang.distributed.bench import do_bench + + dist.barrier(ctx.group) + pre_ms = do_bench(run_ref, warmup=args.warmup, rep=args.rep, group=ctx.group) + tl_ms = do_bench(launch, warmup=args.warmup, rep=args.rep, group=ctx.group) + post_ms = do_bench(run_ref, warmup=args.warmup, rep=args.rep, group=ctx.group) + ref_ms = min(pre_ms, post_ms) + if tflops: + ctx.log( + f"{name:16s} tilescale {tl_ms:8.3f} ms {tflops / (tl_ms * 1e-3):7.1f} TF" + f" | torch {ref_ms:8.3f} ms {tflops / (ref_ms * 1e-3):7.1f} TF" + f" | speedup {ref_ms / tl_ms:6.2f}x" + ) + else: + report(ctx, name, tl_ms, ref_ms, moved) + ctx.log(f" torch reps: {pre_ms:.3f} / {post_ms:.3f} ms (drift shows contention)") + + +def report(ctx: Context, name: str, tl_ms: float, ref_ms: float, moved_bytes: int) -> None: + """Print both timings plus the bus bandwidth each implies. + + ``moved_bytes`` is the payload that has to cross a rank's network link, not + the buffer size, so the number is comparable between collectives with + different algorithmic volumes. + """ + if not ctx.is_leader: + return + tl_gbps = moved_bytes / (tl_ms * 1e-3) / 1e9 + ref_gbps = moved_bytes / (ref_ms * 1e-3) / 1e9 + speedup = ref_ms / tl_ms if tl_ms > 0 else float("nan") + print( + f"{name:<16} tilescale {tl_ms:8.3f} ms {tl_gbps:7.1f} GB/s | " + f"torch {ref_ms:8.3f} ms {ref_gbps:7.1f} GB/s | speedup {speedup:5.2f}x", + flush=True, + ) diff --git a/examples/distributed/internode/internode_gemm_sm100.py b/examples/distributed/internode/internode_gemm_sm100.py new file mode 100644 index 0000000000..e96c4b9eb7 --- /dev/null +++ b/examples/distributed/internode/internode_gemm_sm100.py @@ -0,0 +1,164 @@ +"""Blackwell (sm_100) warp-specialised persistent GEMM over a row range. + +Ported from ``examples/gemm_sm100/gemm_tcgen5mma_ws_persistent.py``, which reaches +1650 TFLOP/s on a B200 against cuBLAS's 1419 -- 1.16x. The naive +``T.copy``/``T.gemm`` loop this replaces managed only 470 TFLOP/s, so the fused +inter-node kernels were bottlenecked on compute by 3.5x, not on the network. + +What makes it fast, all of which the naive version lacks: + +* **tcgen05 MMA into tensor memory.** Accumulators live in TMEM + (``T.alloc_tmem``) rather than registers, so a 128x256 tile does not blow the + register budget. Two of them, alternated by wave parity, let the next MMA start + while the previous result is still being drained. +* **Warp specialisation.** Warp 0 issues TMA loads, warp 1 issues the MMAs, warps + 4-7 run the epilogue. Each stays resident on its own job instead of the whole + CTA marching through load-compute-store in lockstep. +* **Persistent grid.** One CTA per SM, striding over output tiles, so tile setup + is paid once per SM rather than once per tile. +* **Explicit mbarrier phases.** ``num_stages`` deep pipelining with parity + computed by hand; this is what overlaps the TMA loads with the MMAs. + +The row range is what this file adds. ``m_offset`` is a runtime argument so the +host can launch the same kernel over locally-owned rows and over each peer's rows +separately -- the split-launch overlap the AG-GEMM example relies on. Everything +else follows the upstream example. + +Shape constraints inherited from it: ``K % (2 * block_K) == 0`` and +``n_blocks % (2 * group_size) == 0``. +""" + +# NOTE: no `from __future__ import annotations` here. This file defines a +# T.prim_func, and PEP 563 would turn `T.Tensor((M, K), in_dtype)` into a string +# evaluated against module globals, where the closure locals M/K/in_dtype do not +# exist -- "NameError: name 'M' is not defined" at trace time. +import tilelang.language as T +from tilelang.carver.arch import driver + + +def tcgen05_gemm_range_kernel( + M: int, + N: int, + K: int, + m_rows: int, + block_M: int = 128, + block_N: int = 256, + block_K: int = 64, + store_block_N: int = 64, + num_stages: int = 4, + group_size: int = 8, + in_dtype: str = "bfloat16", + out_dtype: str = "bfloat16", + accum_dtype: str = "float32", +): + """``C[m_offset : m_offset + m_rows] = A[same rows] @ B``. + + ``m_rows`` is compile-time (it fixes the grid and wave count); ``m_offset`` is + a kernel argument so one compiled kernel serves every rank's row block. + """ + sm_num = driver.get_num_sms() + m_blocks = T.ceildiv(m_rows, block_M) + n_blocks = T.ceildiv(N, block_N) + k_blocks = T.ceildiv(K, block_K) + waves = T.ceildiv(m_blocks * n_blocks, sm_num) + if K % (2 * block_K): + raise ValueError(f"K={K} must be a multiple of 2*block_K={2 * block_K}") + if n_blocks % (2 * group_size): + raise ValueError( + f"n_blocks={n_blocks} must be a multiple of 2*group_size={2 * group_size}; " + "adjust --gemm-block-n or group_size" + ) + + @T.prim_func + def main( + A: T.Tensor((M, K), in_dtype), + B: T.Tensor((K, N), in_dtype), + C: T.Tensor((M, N), out_dtype), + m_offset: T.int32, + ): + with T.Kernel(sm_num, threads=256) as block_id: + A_shared = T.alloc_shared((num_stages, block_M, block_K), in_dtype) + B_shared = T.alloc_shared((num_stages, block_K, block_N), in_dtype) + # Double-buffered accumulators so wave w+1's MMAs overlap wave w's drain. + C_tmem_0 = T.alloc_tmem([block_M, block_N], accum_dtype) + C_tmem_1 = T.alloc_tmem([block_M, block_N], accum_dtype) + C_local = T.alloc_fragment((block_M, block_N), accum_dtype) + C_shared = T.alloc_shared((block_M, store_block_N), out_dtype) + loaded = T.alloc_barrier([32] * num_stages) + consumed = T.alloc_barrier([1] * num_stages) + tmem_full = T.alloc_barrier([1] * 2) + tmem_empty = T.alloc_barrier([128] * 2) + + tx = T.get_thread_binding() + + if tx < 32: # warp 0: TMA loads + for w in T.unroll(waves): + tile_id = sm_num * w + block_id + bx = (tile_id // group_size) % m_blocks + by = (tile_id % group_size) + (tile_id // group_size) // m_blocks * group_size + if bx * block_M < m_rows and by * block_N < N: + row = m_offset + bx * block_M + for k in T.serial(k_blocks): + phase = w * k_blocks + k + T.mbarrier_wait_parity(consumed[phase % num_stages], + ((phase // num_stages) & 1) ^ 1) + T.tma_copy( + A[row:row + block_M, k * block_K:(k + 1) * block_K], + A_shared[phase % num_stages, :, :], + barrier=loaded[phase % num_stages], + ) + T.tma_copy( + B[k * block_K:(k + 1) * block_K, + by * block_N:(by + 1) * block_N], + B_shared[phase % num_stages, :, :], + barrier=loaded[phase % num_stages], + ) + T.mbarrier_arrive(loaded[phase % num_stages]) + + elif tx < 64: # warp 1: issue tcgen05 MMAs + for w in T.unroll(waves): + tile_id = sm_num * w + block_id + bx = (tile_id // group_size) % m_blocks + by = (tile_id % group_size) + (tile_id // group_size) // m_blocks * group_size + if bx * block_M < m_rows and by * block_N < N: + T.mbarrier_wait_parity(tmem_empty[w & 1], ((w // 2) & 1) ^ 1) + for k in T.serial(k_blocks): + phase = w * k_blocks + k + T.mbarrier_wait_parity(loaded[phase % num_stages], + (phase // num_stages) & 1) + if w & 1 == 0: + T.tcgen05_gemm( + A_shared[phase % num_stages, :, :], + B_shared[phase % num_stages, :, :], + C_tmem_0, False, False, + mbar=consumed[phase % num_stages], + clear_accum=k == 0, + ) + else: + T.tcgen05_gemm( + A_shared[phase % num_stages, :, :], + B_shared[phase % num_stages, :, :], + C_tmem_1, False, False, + mbar=consumed[phase % num_stages], + clear_accum=k == 0, + ) + T.tcgen05_mma_arrive(tmem_full[w & 1]) + + elif 128 <= tx < 256: # warps 4-7: epilogue + for w in T.unroll(waves): + tile_id = sm_num * w + block_id + bx = (tile_id // group_size) % m_blocks + by = (tile_id % group_size) + (tile_id // group_size) // m_blocks * group_size + if bx * block_M < m_rows and by * block_N < N: + row = m_offset + bx * block_M + T.mbarrier_wait_parity(tmem_full[w & 1], (w // 2) & 1) + if (w & 1) == 0: + T.copy(C_tmem_0, C_local) + else: + T.copy(C_tmem_1, C_local) + T.mbarrier_arrive(tmem_empty[w & 1]) + for i in T.unroll(T.ceildiv(block_N, store_block_N)): + T.copy(C_local[:, i * store_block_N:(i + 1) * store_block_N], C_shared) + T.copy(C_shared, C[row, by * block_N + i * store_block_N]) + + return main diff --git a/examples/distributed/internode/run_internode.sh b/examples/distributed/internode/run_internode.sh new file mode 100755 index 0000000000..5eba853da3 --- /dev/null +++ b/examples/distributed/internode/run_internode.sh @@ -0,0 +1,122 @@ +#!/bin/bash +# Launch an inter-node collective example across two nodes, one GPU each. +# +# Usage: run_internode.sh