ep: PyTorch wrapper, autograd ops, symm-mem zero-copy bindings + distributed tests/example

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Author: phu0ngng

build_tools/pytorch.py

@@ -76,6 +76,12 @@ def setup_pytorch_extension(
 
     setup_mpi_flags(include_dirs, cxx_flags)
 
+    # Mirror the NCCL EP gate from setup.py / common CMake. When disabled, the
+    # ep.cpp source no-ops at the #ifdef boundary; without the define it would
+    # produce undefined references to nvte_ep_*.
+    if bool(int(os.getenv("NVTE_BUILD_WITH_NCCL_EP", "1"))):
+        cxx_flags.append("-DNVTE_WITH_NCCL_EP")
+
     library_dirs = []
     libraries = []
     if bool(int(os.getenv("NVTE_ENABLE_NVSHMEM", 0))):

examples/pytorch/ep/ep_moe.py

@@ -0,0 +1,281 @@
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+"""End-to-end MoE example: dispatch -> batched expert linear -> combine, fwd + bwd.
+
+One process per GPU. Launched via run_test_ep.sh (torchrun).
+"""
+
+import argparse
+import os
+import sys
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+from transformer_engine.pytorch.ep import (
+    EpHandle,
+    ep_bootstrap,
+    ep_dispatch,
+    ep_combine,
+    symm_mem_alloc,
+)
+
+
+def _parse_args():
+    p = argparse.ArgumentParser(description="TE-PyTorch EP MoE example (fwd + bwd)")
+    p.add_argument("--num-tokens", type=int, default=8, help="Per-rank token count.")
+    p.add_argument("--top-k", type=int, default=2)
+    p.add_argument("--hidden", type=int, default=32)
+    p.add_argument("--hidden-out", type=int, default=32)
+    p.add_argument("--num-experts", type=int, default=None)
+    p.add_argument("--check", action="store_true", default=True)
+    p.add_argument(
+        "--benchmark",
+        action="store_true",
+        help="Time fwd+bwd over both regular HBM and NCCL symm-mem payload buffers.",
+    )
+    p.add_argument("--benchmark-iters", type=int, default=20)
+    p.add_argument("--benchmark-warmup", type=int, default=5)
+    return p.parse_args()
+
+
+def _make_routing(rank, T, K, E, num_local_experts):
+    """Deterministic: topk_idx[t, k] = (rank*NLE + t*K + k) % E."""
+    topk_idx = np.empty((T, K), dtype=np.int32)
+    for t in range(T):
+        for k in range(K):
+            topk_idx[t, k] = (rank * num_local_experts + t * K + k) % E
+    return topk_idx
+
+
+def _batched_expert_linear(recv_tokens, kernels, num_local_experts):
+    """Per-expert linear over [recv_pr // NLE] tokens per expert."""
+    recv_pr, _H = recv_tokens.shape
+    H_out = kernels.shape[-1]
+    slots_per_expert = recv_pr // num_local_experts
+    grouped = recv_tokens.view(num_local_experts, slots_per_expert, recv_tokens.shape[-1])
+    # kernels: [num_local_experts, H, H_out]
+    out = torch.bmm(grouped, kernels.to(grouped.dtype))
+    return out.view(recv_pr, H_out)
+
+
+def _reference_moe(tokens, topk_idx, topk_w, kernels):
+    T, K = topk_idx.shape
+    H_out = kernels.shape[-1]
+    out = np.zeros((T, H_out), dtype=np.float32)
+    for t in range(T):
+        tok = tokens[t].astype(np.float32)
+        for k in range(K):
+            e = int(topk_idx[t, k])
+            out[t] += float(topk_w[t, k]) * (tok @ kernels[e].astype(np.float32))
+    return out
+
+
+def _reference_grad(tokens, topk_idx, topk_w, kernels):
+    T, K = topk_idx.shape
+    H = tokens.shape[-1]
+    ref_out = _reference_moe(tokens, topk_idx, topk_w, kernels)
+    grad = np.zeros((T, H), dtype=np.float32)
+    for t in range(T):
+        mixed = np.zeros_like(kernels[0])
+        for k in range(K):
+            mixed = mixed + float(topk_w[t, k]) * kernels[int(topk_idx[t, k])]
+        grad[t] = ref_out[t] @ mixed.T
+    return ref_out, grad
+
+
+def main():
+    args = _parse_args()
+
+    dist.init_process_group(backend="nccl")
+    rank = dist.get_rank()
+    world_size = dist.get_world_size()
+    torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", rank)))
+    device = torch.device("cuda", torch.cuda.current_device())
+
+    major, minor = torch.cuda.get_device_capability()
+    if major * 10 + minor < 90:
+        if rank == 0:
+            print(f"[ep_moe] SKIPPED: EP requires SM>=90 (got SM{major}{minor})")
+        dist.destroy_process_group()
+        return
+
+    if world_size < 4:
+        if rank == 0:
+            print(f"[ep_moe] SKIPPED: EP requires >= 4 ranks (got {world_size})")
+        dist.destroy_process_group()
+        return
+
+    ep_size = world_size
+    num_experts = args.num_experts if args.num_experts is not None else world_size
+    assert num_experts % ep_size == 0
+    num_local_experts = num_experts // ep_size
+    T = args.num_tokens
+    recv_pr = ep_size * T * args.top_k
+
+    ep_group = dist.new_group(ranks=list(range(world_size)), backend="nccl")
+    ep_bootstrap(
+        ep_group,
+        num_experts=num_experts,
+        max_tokens_per_rank=T,
+        recv_capacity_per_rank=recv_pr,
+        hidden_dim=args.hidden,
+    )
+
+    rng = np.random.default_rng(seed=42 + rank)
+    tokens_np = (rng.standard_normal((T, args.hidden), dtype=np.float32) * 0.5).astype(np.float32)
+    topk_idx_np = _make_routing(rank, T, args.top_k, num_experts, num_local_experts)
+    w_np = np.full((T, args.top_k), 1.0 / args.top_k, dtype=np.float32)
+    # Same seed on every rank → identical kernel array everywhere.
+    kr = np.random.default_rng(seed=42)
+    kernels_np = (
+        kr.standard_normal((num_experts, args.hidden, args.hidden_out), dtype=np.float32)
+        * (1.0 / np.sqrt(args.hidden))
+    ).astype(np.float32)
+
+    tokens = (
+        torch.from_numpy(tokens_np).to(device=device, dtype=torch.bfloat16).requires_grad_(True)
+    )
+    topk_idx = torch.from_numpy(topk_idx_np).to(device)
+    topk_w = torch.from_numpy(w_np).to(device)
+    kernels_local = torch.from_numpy(
+        kernels_np[rank * num_local_experts : (rank + 1) * num_local_experts]
+    ).to(device=device, dtype=torch.bfloat16)
+
+    handle = EpHandle(
+        top_k=args.top_k,
+        max_tokens_per_rank=T,
+        recv_capacity_per_rank=recv_pr,
+        hidden_dim=args.hidden,
+        num_local_experts=num_local_experts,
+    )
+    recv_tokens = torch.empty(recv_pr, args.hidden, dtype=torch.bfloat16, device=device)
+    recv_w = torch.empty(recv_pr, dtype=torch.float32, device=device)
+    token_counts = torch.empty(num_local_experts, dtype=torch.int32, device=device)
+    result = torch.empty(T, args.hidden_out, dtype=torch.bfloat16, device=device)
+    grad_tokens = torch.empty_like(tokens)
+    grad_topk_w = torch.empty_like(topk_w)
+    grad_eo = torch.empty(recv_pr, args.hidden_out, dtype=torch.bfloat16, device=device)
+
+    recv_t, recv_w_out, _tc = ep_dispatch(
+        handle,
+        topk_idx,
+        tokens,
+        topk_w,
+        recv_tokens,
+        recv_w,
+        token_counts,
+        grad_tokens,
+        grad_topk_w,
+    )
+    expert_out = _batched_expert_linear(recv_t, kernels_local, num_local_experts)
+    out = ep_combine(handle, expert_out, recv_w_out, result, grad_eo)
+
+    loss = 0.5 * (out.float() ** 2).sum()
+    loss.backward()
+    torch.cuda.synchronize()
+
+    if rank == 0:
+        print(
+            f"[ep_moe] loss={float(loss):.4f} grad_tokens.shape={tuple(tokens.grad.shape)} "
+            f"ep={ep_size} num_experts={num_experts} recv_pr={recv_pr}"
+        )
+
+    if args.benchmark:
+        # Time forward dispatch + expert + combine over both regular HBM and
+        # NCCL symm-mem payloads. Recorded only — perf ratio depends heavily on
+        # topology and is not CI-gated.
+        import time
+
+        def _time(label, tokens_buf, recv_tokens_buf, result_buf):
+            torch.cuda.synchronize()
+            dist.barrier()
+            for _ in range(args.benchmark_warmup):
+                rt, rw, _tc = ep_dispatch(
+                    handle,
+                    topk_idx,
+                    tokens_buf,
+                    topk_w,
+                    recv_tokens_buf,
+                    recv_w,
+                    token_counts,
+                    grad_tokens,
+                    grad_topk_w,
+                )
+                eo = _batched_expert_linear(rt, kernels_local, num_local_experts)
+                ep_combine(handle, eo, rw, result_buf, grad_eo)
+            torch.cuda.synchronize()
+            dist.barrier()
+            t0 = time.perf_counter()
+            for _ in range(args.benchmark_iters):
+                rt, rw, _tc = ep_dispatch(
+                    handle,
+                    topk_idx,
+                    tokens_buf,
+                    topk_w,
+                    recv_tokens_buf,
+                    recv_w,
+                    token_counts,
+                    grad_tokens,
+                    grad_topk_w,
+                )
+                eo = _batched_expert_linear(rt, kernels_local, num_local_experts)
+                ep_combine(handle, eo, rw, result_buf, grad_eo)
+            torch.cuda.synchronize()
+            dt_ms = (time.perf_counter() - t0) * 1000.0 / args.benchmark_iters
+            if rank == 0:
+                print(
+                    f"[ep_moe --benchmark] {label}: {dt_ms:.3f} ms/iter "
+                    f"(iters={args.benchmark_iters})"
+                )
+            return dt_ms
+
+        # 1) HBM baseline reuses the existing buffers (already allocated above).
+        hbm_ms = _time("regular HBM", tokens.detach(), recv_tokens, result)
+
+        # 2) Symm-mem variant: reallocate as symm-mem (collective on every rank).
+        try:
+            tokens_sm = symm_mem_alloc((T, args.hidden), torch.bfloat16, ep_group, device=device)
+            recv_tokens_sm = symm_mem_alloc(
+                (recv_pr, args.hidden), torch.bfloat16, ep_group, device=device
+            )
+            result_sm = torch.empty_like(result)
+            tokens_sm.copy_(tokens.detach())
+            symm_ms = _time("symm-mem", tokens_sm, recv_tokens_sm, result_sm)
+            if rank == 0:
+                print(f"[ep_moe --benchmark] speedup: {hbm_ms / symm_ms:.2f}x")
+        except RuntimeError as e:
+            if rank == 0:
+                print(f"[ep_moe --benchmark] symm-mem path skipped: {e}")
+
+    if args.check:
+        # Gather across ranks for a global reference comparison.
+        global_tokens = [torch.empty_like(tokens) for _ in range(world_size)]
+        global_topk_idx = [torch.empty_like(topk_idx) for _ in range(world_size)]
+        global_topk_w = [torch.empty_like(topk_w) for _ in range(world_size)]
+        global_out = [torch.empty_like(out) for _ in range(world_size)]
+        global_grad = [torch.empty_like(tokens.grad) for _ in range(world_size)]
+        dist.all_gather(global_tokens, tokens.detach())
+        dist.all_gather(global_topk_idx, topk_idx)
+        dist.all_gather(global_topk_w, topk_w)
+        dist.all_gather(global_out, out.detach())
+        dist.all_gather(global_grad, tokens.grad)
+        if rank == 0:
+            all_tokens = torch.cat(global_tokens).float().cpu().numpy()
+            all_idx = torch.cat(global_topk_idx).cpu().numpy()
+            all_w = torch.cat(global_topk_w).cpu().numpy()
+            all_out = torch.cat(global_out).float().cpu().numpy()
+            all_grad = torch.cat(global_grad).float().cpu().numpy()
+            ref_out, ref_grad = _reference_grad(all_tokens, all_idx, all_w, kernels_np)
+            np.testing.assert_allclose(all_out, ref_out, rtol=5e-2, atol=5e-2)
+            np.testing.assert_allclose(all_grad, ref_grad, rtol=5e-2, atol=5e-2)
+            print(f"[ep_moe] --check PASSED (ref_out.sum()={float(ref_out.sum()):.4f})")
+
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()
+    sys.exit(0)

examples/pytorch/ep/run_test_ep.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# See LICENSE for license information.
+
+set -uo pipefail
+
+DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l)
+NUM_GPUS="${NUM_GPUS:-${DETECTED_GPUS}}"
+if [ "${NUM_GPUS}" -lt 4 ]; then
+  echo "EP requires >= 4 GPUs (found ${NUM_GPUS}); SKIPPING."
+  exit 0
+fi
+if [ "${NUM_GPUS}" -gt 8 ]; then NUM_GPUS=8; fi
+
+: ${TE_PATH:=/opt/transformerengine}
+: ${TEST_TIMEOUT_S:=120}
+
+SCRIPT="${TE_PATH}/examples/pytorch/ep/ep_moe.py"
+export PYTHONPATH="${TE_PATH}${PYTHONPATH:+:${PYTHONPATH}}"
+
+# Stage JIT cubins on tmpfs for fast iteration.
+: ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"}
+export NCCL_EP_JIT_CACHE_DIR
+mkdir -p "$NCCL_EP_JIT_CACHE_DIR"
+
+echo "*** Executing ep_moe.py across ${NUM_GPUS} GPUs (timeout=${TEST_TIMEOUT_S}s) ***"
+timeout --foreground --signal=KILL "${TEST_TIMEOUT_S}" \
+  torchrun --standalone --nnodes=1 --nproc-per-node="${NUM_GPUS}" \
+  "${SCRIPT}" --check 2>&1 | tee stdout_ep_moe.txt
+RC=${PIPESTATUS[0]}
+
+RET=0
+if [ "${RC}" -ne 0 ]; then RET=1; fi
+if grep -qE "FAILED|Traceback" stdout_ep_moe.txt; then RET=1; fi
+rm -f stdout_ep_moe.txt
+exit $RET