From dd2158c2694a63594f23513e0763801f86a2dba7 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Fri, 22 May 2026 23:05:06 +0000 Subject: [PATCH 01/55] Expert Parallelism: common C API + NCCL EP v0.1 backend Signed-off-by: Phuong Nguyen --- .gitmodules | 4 + 3rdparty/nccl | 1 + qa/L1_cpp_distributed/test.sh | 3 + setup.py | 127 +++ tests/cpp_distributed/CMakeLists.txt | 89 +- tests/cpp_distributed/run_test_ep.sh | 137 +++ tests/cpp_distributed/test_ep_common.h | 308 ++++++ tests/cpp_distributed/test_ep_coverage.cu | 379 ++++++++ tests/cpp_distributed/test_ep_init.cu | 64 ++ tests/cpp_distributed/test_ep_pipeline.cu | 890 ++++++++++++++++++ transformer_engine/common/CMakeLists.txt | 90 ++ transformer_engine/common/ep/ep_api.cpp | 76 ++ transformer_engine/common/ep/ep_api_stub.cpp | 61 ++ transformer_engine/common/ep/ep_backend.cpp | 514 ++++++++++ transformer_engine/common/ep/ep_backend.h | 114 +++ .../include/transformer_engine/comm_window.h | 32 + .../common/include/transformer_engine/ep.h | 161 ++++ 17 files changed, 3049 insertions(+), 1 deletion(-) create mode 160000 3rdparty/nccl create mode 100755 tests/cpp_distributed/run_test_ep.sh create mode 100644 tests/cpp_distributed/test_ep_common.h create mode 100644 tests/cpp_distributed/test_ep_coverage.cu create mode 100644 tests/cpp_distributed/test_ep_init.cu create mode 100644 tests/cpp_distributed/test_ep_pipeline.cu create mode 100644 transformer_engine/common/ep/ep_api.cpp create mode 100644 transformer_engine/common/ep/ep_api_stub.cpp create mode 100644 transformer_engine/common/ep/ep_backend.cpp create mode 100644 transformer_engine/common/ep/ep_backend.h create mode 100644 transformer_engine/common/include/transformer_engine/comm_window.h create mode 100644 transformer_engine/common/include/transformer_engine/ep.h diff --git a/.gitmodules b/.gitmodules index 4b188d6bb1..e531c95507 100644 --- a/.gitmodules +++ b/.gitmodules @@ -7,3 +7,7 @@ [submodule "3rdparty/cutlass"] path = 3rdparty/cutlass url = https://github.com/NVIDIA/cutlass.git +[submodule "3rdparty/nccl"] + path = 3rdparty/nccl + url = https://github.com/NVIDIA/nccl.git + branch = v2.30u1 diff --git a/3rdparty/nccl b/3rdparty/nccl new file mode 160000 index 0000000000..6a9bc953ac --- /dev/null +++ b/3rdparty/nccl @@ -0,0 +1 @@ +Subproject commit 6a9bc953ac1c4eef92d5adbe3092d4c2cb0a4c98 diff --git a/qa/L1_cpp_distributed/test.sh b/qa/L1_cpp_distributed/test.sh index 8d767a4efb..7e5ce2cf0d 100755 --- a/qa/L1_cpp_distributed/test.sh +++ b/qa/L1_cpp_distributed/test.sh @@ -14,4 +14,7 @@ if [[ $(nvidia-smi --list-gpus | wc -l) -ge 4 ]]; then cmake -GNinja -S. -Bbuild cmake --build build mpirun --allow-run-as-root --np 4 --oversubscribe ./build/test_comm_gemm + + # EP suites; runner self-skips on pre-Hopper GPUs. + bash ./run_test_ep.sh 4 ./build fi diff --git a/setup.py b/setup.py index 7f6b51c148..9599193ced 100644 --- a/setup.py +++ b/setup.py @@ -84,6 +84,34 @@ def setup_common_extension() -> CMakeExtension: cusolvermp_dir = os.getenv("CUSOLVERMP_HOME", "/usr") cmake_flags.append(f"-DCUSOLVERMP_DIR={cusolvermp_dir}") + # NCCL EP: on by default; auto-disabled if no arch >= 90. + # Set NVTE_BUILD_WITH_NCCL_EP=0/1 to force off/on. + nccl_ep_env = os.getenv("NVTE_BUILD_WITH_NCCL_EP") + explicit_nccl_ep = nccl_ep_env is not None + build_with_nccl_ep = bool(int(nccl_ep_env)) if explicit_nccl_ep else True + + if build_with_nccl_ep: + arch_tokens = [a.strip() for a in str(archs or "").split(";") if a.strip()] + has_hopper_or_newer = any(t.lower() == "native" for t in arch_tokens) or any( + int(t.rstrip("af")) >= 90 for t in arch_tokens if t.rstrip("af").isdigit() + ) + if not has_hopper_or_newer: + if explicit_nccl_ep: + raise RuntimeError( + "NVTE_BUILD_WITH_NCCL_EP=1 requires at least one CUDA arch >= 90 in " + f"NVTE_CUDA_ARCHS (got '{archs}'). Add '90' or unset NVTE_BUILD_WITH_NCCL_EP." + ) + print( + "[NCCL EP] No CUDA arch >= 90 in NVTE_CUDA_ARCHS" + f" ('{archs}'); auto-disabling NCCL EP (nvte_ep_* will throw at runtime)." + ) + build_with_nccl_ep = False + + if build_with_nccl_ep: + build_nccl_ep_submodule() + else: + cmake_flags.append("-DNVTE_WITH_NCCL_EP=OFF") + # Add custom CMake arguments from environment variable nvte_cmake_extra_args = os.getenv("NVTE_CMAKE_EXTRA_ARGS") if nvte_cmake_extra_args: @@ -130,6 +158,105 @@ def setup_requirements() -> Tuple[List[str], List[str]]: return [remove_dups(reqs) for reqs in [install_reqs, test_reqs]] +def _discover_nccl_home() -> str: + """Resolve NCCL_HOME: honor env var, else probe well-known prefixes, else ldconfig.""" + env_home = os.environ.get("NCCL_HOME") + if env_home: + if (Path(env_home) / "include" / "nccl.h").exists(): + return env_home + print( + f"[NCCL EP] WARNING: NCCL_HOME='{env_home}' is set but " + f"'{env_home}/include/nccl.h' was not found; falling back to system probes." + ) + + for cand in ("/opt/nvidia/nccl", "/usr/local/nccl", "/usr"): + p = Path(cand) + if (p / "include" / "nccl.h").exists() and any( + (p / "lib" / name).exists() or (p / "lib64" / name).exists() + for name in ("libnccl.so", "libnccl.so.2") + ): + return str(p) + + try: + out = subprocess.check_output(["ldconfig", "-p"], stderr=subprocess.DEVNULL).decode() + for line in out.splitlines(): + if "libnccl.so" in line and "=>" in line: + lib_path = Path(line.split("=>")[-1].strip()) + root = lib_path.parent.parent + if (root / "include" / "nccl.h").exists(): + return str(root) + except (subprocess.CalledProcessError, FileNotFoundError): + pass + + raise RuntimeError( + "Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix." + ) + + +def build_nccl_ep_submodule() -> str: + """Build libnccl_ep.so from the 3rdparty/nccl submodule. + + NCCL EP is on by default; the system NCCL core (libnccl.so) supplies the + headers and runtime symbols. Returns the submodule build directory. + """ + nccl_root = current_file_path / "3rdparty" / "nccl" + if not (nccl_root / "Makefile").exists(): + raise RuntimeError( + f"NCCL submodule not found at {nccl_root}. " + "Run `git submodule update --init --recursive`." + ) + + build_dir = nccl_root / "build" + nccl_ep_lib = build_dir / "lib" / "libnccl_ep.so" + + archs = cuda_archs() or "90" + arch_list = [] + for a in str(archs).split(";"): + a = a.strip().rstrip("af") + if a and a.isdigit() and int(a) >= 90: + arch_list.append(a) + if not arch_list: + arch_list = ["90"] + gencode = " ".join(f"-gencode=arch=compute_{a},code=sm_{a}" for a in arch_list) + + nproc = os.cpu_count() or 8 + env = os.environ.copy() + env["NVCC_GENCODE"] = gencode + # NCCL EP needs the core NCCL headers + libnccl.so; write NCCL EP build + # outputs to the submodule's local build/ tree. + nccl_home = _discover_nccl_home() + env["NCCL_HOME"] = nccl_home + env["NCCL_EP_BUILDDIR"] = str(build_dir) + + if not nccl_ep_lib.exists(): + print(f"[NCCL EP] Building libnccl_ep.so (gencode='{gencode}')") + subprocess.check_call( + ["make", "-j", str(nproc), "-C", "contrib/nccl_ep", "lib"], + cwd=str(nccl_root), + env=env, + ) + + # TE's CMake expects nccl.h under 3rdparty/nccl/build/include/ for its + # version check. Mirror the top-level host headers from the system NCCL + # install — DON'T mirror nccl_device/ because the submodule ships its own + # newer copy at src/include/nccl_device/ with device-side templates that + # conflict with older system versions, and the JIT include path picks the + # submodule's. + nccl_include = build_dir / "include" + nccl_include.mkdir(parents=True, exist_ok=True) + for cand in (Path(nccl_home) / "include", Path("/usr/include")): + p = Path(cand) + if (p / "nccl.h").exists(): + for name in ("nccl.h", "nccl_net.h", "nccl_tuner.h"): + src = p / name + dst = nccl_include / name + if src.exists() and not dst.exists(): + dst.symlink_to(src) + break + + return str(build_dir) + + def git_check_submodules() -> None: """ Attempt to checkout git submodules automatically during setup. diff --git a/tests/cpp_distributed/CMakeLists.txt b/tests/cpp_distributed/CMakeLists.txt index 44ad7c7384..463ae011a5 100644 --- a/tests/cpp_distributed/CMakeLists.txt +++ b/tests/cpp_distributed/CMakeLists.txt @@ -55,9 +55,14 @@ target_include_directories(test_comm_gemm PRIVATE ${test_comm_gemm_INCLUDES}) find_package(CUDAToolkit REQUIRED) find_package(OpenMP REQUIRED) find_package(MPI REQUIRED) + +# ── NCCL library ────────────────────────────────────────────────────────────── +# Search order: NCCL_HOME env → 3rdparty/nccl submodule build → system paths. +set(NCCL_SUBMODULE_BUILD "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl/build") find_library(NCCL_LIB NAMES nccl libnccl - PATH_SUFFIXES lib + HINTS $ENV{NCCL_HOME}/lib ${NCCL_SUBMODULE_BUILD}/lib + PATH_SUFFIXES lib lib64 REQUIRED) list(APPEND test_comm_gemm_LINKER_LIBS CUDA::cuda_driver @@ -72,5 +77,87 @@ target_link_libraries(test_comm_gemm PUBLIC ${test_comm_gemm_LINKER_LIBS}) target_compile_options(test_comm_gemm PRIVATE -O2 -fopenmp) +# NCCL headers: prefer submodule build output (has the handle_init API), +# then submodule src, then system (CUDA toolkit). +set(NCCL_SUBMODULE_INCLUDE "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl/build/include") +set(NCCL_SUBMODULE_SRC_INCLUDE "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl/src/include") +if(EXISTS "${NCCL_SUBMODULE_INCLUDE}/nccl.h") + set(NCCL_INCLUDE_DIR "${NCCL_SUBMODULE_INCLUDE}") +elseif(EXISTS "${NCCL_SUBMODULE_SRC_INCLUDE}/nccl.h") + set(NCCL_INCLUDE_DIR "${NCCL_SUBMODULE_SRC_INCLUDE}") +elseif(DEFINED ENV{NCCL_HOME}) + set(NCCL_INCLUDE_DIR "$ENV{NCCL_HOME}/include") +endif() + include(GoogleTest) gtest_discover_tests(test_comm_gemm DISCOVERY_TIMEOUT 600) + +# ── EP distributed tests (HT mode) ───────────────────────────────────────── +# No MPI dependency — processes are spawned by run_test_ep.sh with +# --rank / --nranks flags. ncclUniqueId exchange uses a +# shared temp file (see test_ep_common.h for details). +# Headers + libs come from the in-tree 3rdparty/nccl submodule build. +set(NCCL_EP_SUBMODULE_ROOT + "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl") +find_library(NCCL_EP_LIB + NAMES nccl_ep libnccl_ep + HINTS ${NCCL_EP_SUBMODULE_ROOT}/build/lib + NO_DEFAULT_PATH + REQUIRED) + +set(NCCL_EP_INCLUDE_DIR "${NCCL_EP_SUBMODULE_ROOT}/contrib/nccl_ep/include") +if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h") + message(FATAL_ERROR + "NCCL EP header not found at ${NCCL_EP_INCLUDE_DIR}/nccl_ep.h. " + "Run `git submodule update --init --recursive` to checkout 3rdparty/nccl.") +endif() +message(STATUS "EP test: NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}") + +# Collect NCCL include dirs shared by all EP test targets (nccl_ep.h + nccl.h). +set(EP_TEST_NCCL_INCLUDES ${NCCL_EP_INCLUDE_DIR}) +if(DEFINED NCCL_INCLUDE_DIR) + list(APPEND EP_TEST_NCCL_INCLUDES ${NCCL_INCLUDE_DIR}) + message(STATUS "EP test: NCCL headers: ${NCCL_INCLUDE_DIR}") +endif() + +set(EP_TEST_COMMON_INCLUDES + ${EP_TEST_NCCL_INCLUDES} + ../../transformer_engine/common/include + ../../transformer_engine/common + ${CMAKE_CURRENT_SOURCE_DIR}) + +set(EP_TEST_COMMON_LIBS + CUDA::cuda_driver + CUDA::cudart + CUDA::nvrtc + GTest::gtest + ${TE_LIB} + ${NCCL_LIB} + ${NCCL_EP_LIB}) + +# nvrtc symbols are referenced from libtransformer_engine.so but not in its +# DT_NEEDED list (loaded via dlopen in Python). For cpp tests we link nvrtc +# explicitly with --no-as-needed so the linker keeps the dependency. +set(EP_TEST_LINK_OPTS "LINKER:--no-as-needed") + +# ── EP init tests (InitPath, HandleMemSizeQuery) ───────────────────────────── +add_executable(test_ep_init test_ep_init.cu) +target_include_directories(test_ep_init PRIVATE ${EP_TEST_COMMON_INCLUDES}) +target_link_libraries(test_ep_init PUBLIC ${EP_TEST_COMMON_LIBS}) +target_link_options(test_ep_init PUBLIC ${EP_TEST_LINK_OPTS}) + +# ── EP pipeline tests (dispatch, combine, bwd, integrated) ─────────────────── +add_executable(test_ep_pipeline test_ep_pipeline.cu) +target_include_directories(test_ep_pipeline PRIVATE ${EP_TEST_COMMON_INCLUDES}) +target_link_libraries(test_ep_pipeline PUBLIC ${EP_TEST_COMMON_LIBS}) +target_link_options(test_ep_pipeline PUBLIC ${EP_TEST_LINK_OPTS}) + +# ── EP coverage tests (multi-handle, top_k=1, empty experts, negatives, threading) ── +add_executable(test_ep_coverage test_ep_coverage.cu) +target_include_directories(test_ep_coverage PRIVATE ${EP_TEST_COMMON_INCLUDES}) +target_link_libraries(test_ep_coverage PUBLIC ${EP_TEST_COMMON_LIBS}) +target_link_options(test_ep_coverage PUBLIC ${EP_TEST_LINK_OPTS}) + +# Do NOT use gtest_discover_tests — these binaries require multi-process +# launch via run_test_ep.sh, not direct single-process execution. +message(STATUS "EP distributed tests enabled: ${NCCL_EP_LIB}") diff --git a/tests/cpp_distributed/run_test_ep.sh b/tests/cpp_distributed/run_test_ep.sh new file mode 100755 index 0000000000..017d3f807b --- /dev/null +++ b/tests/cpp_distributed/run_test_ep.sh @@ -0,0 +1,137 @@ +#!/usr/bin/env bash +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +# +# Run TE EP distributed unit tests across multiple GPUs. +# +# Spawns one background bash process per GPU (no MPI dependency), matching the +# JAX multi-process launcher style. ncclUniqueId is exchanged via a shared +# temp file (see test_ep_common.h). Each rank builds its own ncclComm_t and +# passes it to nvte_ep_initialize. +# +# Usage: +# bash run_test_ep.sh [num_gpus] [build_dir] +# +# Defaults: +# num_gpus = number of GPUs visible to nvidia-smi +# build_dir = /build +# +# Environment variables: +# GTEST_FILTER — forwarded to all processes (e.g., "EPDispatchTest.*") +# TEST_TIMEOUT_S — per-process timeout in seconds (default: 180) + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +BUILD_DIR="${2:-${SCRIPT_DIR}/build}" +NUM_GPUS="${1:-$(nvidia-smi -L 2>/dev/null | wc -l)}" +TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-180}" + +# Skip cleanly on pre-Hopper: NCCL EP requires SM>=90. +MIN_SM=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null \ + | awk -F. 'NR==1 || ($1*10+$2) 0 && MIN_SM < 90 )); then + echo "NCCL EP requires SM>=90 (lowest visible GPU is SM${MIN_SM}); SKIPPING." + exit 0 +fi + +GTEST_ARGS="${GTEST_FILTER:+--gtest_filter=${GTEST_FILTER}}" +OVERALL_FAIL=0 + +# --------------------------------------------------------------------------- +# run_suite BINARY SUITE_NAME MIN_GPUS +# --------------------------------------------------------------------------- +run_suite() { + local BINARY="$1" + local SUITE_NAME="$2" + local MIN_GPUS="${3:-2}" + + local TEST_BIN="${BUILD_DIR}/${BINARY}" + + if [[ ! -x "${TEST_BIN}" ]]; then + echo "ERROR: binary not found: ${TEST_BIN}" + echo "Build: cd ${SCRIPT_DIR} && mkdir -p build && cd build && cmake .. && make" + OVERALL_FAIL=1 + return + fi + + if (( NUM_GPUS < MIN_GPUS )); then + echo "${SUITE_NAME}: requires ${MIN_GPUS} GPUs, found ${NUM_GPUS}. Skipping." + return + fi + + local TMPDIR_L="${TMPDIR:-/tmp}" + local UID_FILE="${TMPDIR_L}/te_ep_uid_${BINARY}_$$" + rm -f "${UID_FILE}" + + local LOG_DIR + LOG_DIR=$(mktemp -d) + local FAIL=0 + + echo "=== ${SUITE_NAME} ===" + echo " GPUs: ${NUM_GPUS} Binary: ${TEST_BIN}" + echo + + # Spawn one background process per GPU. ncclUniqueId is exchanged via the + # shared UID_FILE. Each process is wrapped in `timeout` to detect hangs early. + local PIDS=() + for i in $(seq 0 $((NUM_GPUS - 1))); do + timeout --foreground --signal=KILL "${TEST_TIMEOUT_S}" \ + "${TEST_BIN}" \ + --rank="${i}" \ + --nranks="${NUM_GPUS}" \ + --uid-file="${UID_FILE}" \ + ${GTEST_ARGS} \ + > "${LOG_DIR}/rank_${i}.log" 2>&1 & + PIDS+=($!) + done + for i in $(seq 0 $((NUM_GPUS - 1))); do + if ! wait "${PIDS[$i]}"; then + local rc=$? + FAIL=1 + if [[ $rc -eq 137 || $rc -eq 124 ]]; then + echo " rank ${i}: TIMEOUT after ${TEST_TIMEOUT_S}s (rc=${rc})" + fi + fi + done + + echo "--- Rank 0 output ---" + cat "${LOG_DIR}/rank_0.log" + + if (( FAIL )); then + for i in $(seq 1 $((NUM_GPUS - 1))); do + echo "--- Rank ${i} output ---" + cat "${LOG_DIR}/rank_${i}.log" + done + echo "=== ${SUITE_NAME}: FAILED ===" + OVERALL_FAIL=1 + else + echo "=== ${SUITE_NAME}: ALL PASSED ===" + fi + + rm -rf "${LOG_DIR}" + rm -f "${UID_FILE}" +} + +# --------------------------------------------------------------------------- +# Cleanup on abort +# --------------------------------------------------------------------------- +cleanup() { rm -f "${TMPDIR:-/tmp}"/te_ep_uid_*_"$$" 2>/dev/null || true; } +trap cleanup EXIT INT TERM + +# --------------------------------------------------------------------------- +# Run all suites +# --------------------------------------------------------------------------- +run_suite "test_ep_init" "EP Init Tests" 2 +run_suite "test_ep_pipeline" "EP Pipeline Tests" 2 +run_suite "test_ep_coverage" "EP Coverage Tests" 2 + +echo +if (( OVERALL_FAIL )); then + echo "=== SOME SUITES FAILED ===" +else + echo "=== ALL SUITES PASSED ===" +fi + +exit "${OVERALL_FAIL}" diff --git a/tests/cpp_distributed/test_ep_common.h b/tests/cpp_distributed/test_ep_common.h new file mode 100644 index 0000000000..77baa92b0c --- /dev/null +++ b/tests/cpp_distributed/test_ep_common.h @@ -0,0 +1,308 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/* + * Shared TE EP test infrastructure. Include once per TU; ep_bootstrap() in + * each test binary's main() populates process-level globals. + * Defaults: 4 experts/rank, hidden_dim=256, max_tokens_per_rank=64. + */ +#pragma once + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +// ── Error-checking macros ───────────────────────────────────────────────────── + +#define CHECK_NCCL(expr) \ + do { \ + ncclResult_t _err = (expr); \ + if (_err != ncclSuccess) \ + FAIL() << "NCCL error " << _err << ": " << ncclGetErrorString(_err); \ + } while (false) + +#define CHECK_CUDA(expr) \ + do { \ + cudaError_t _err = (expr); \ + if (_err != cudaSuccess) \ + FAIL() << "CUDA error " << _err << ": " << cudaGetErrorString(_err); \ + } while (false) + +#define ASSERT_CUDA_OK(expr) \ + do { \ + cudaError_t _err = (expr); \ + if (_err != cudaSuccess) { \ + fprintf(stderr, "CUDA error %d: %s\n", _err, cudaGetErrorString(_err)); \ + exit(EXIT_FAILURE); \ + } \ + } while (false) + +#define ASSERT_NCCL_OK(expr) \ + do { \ + ncclResult_t _err = (expr); \ + if (_err != ncclSuccess) { \ + fprintf(stderr, "NCCL error %d: %s\n", _err, ncclGetErrorString(_err)); \ + exit(EXIT_FAILURE); \ + } \ + } while (false) + +// ── Process-level state ─────────────────────────────────────────────────────── + +static int g_process_id = -1; +static int g_num_processes = -1; +static std::string g_uid_file; + +static int g_sm_major = -1; // set by ep_bootstrap; -1 until then +static int g_ep_size = -1; +static int g_num_experts = -1; +static int g_hidden_dim = 256; +static int g_max_tokens_per_rank = 64; +static bool g_ep_initialized = false; +static ncclComm_t g_ep_comm = nullptr; // owned by harness, destroyed in ep_teardown + +// ── TensorHandle RAII wrapper ───────────────────────────────────────────────── + +// View over a caller-owned device buffer; owns NVTETensor metadata only. Move-only. +struct TensorHandle { + NVTETensor tensor = nullptr; + void* dev_ptr = nullptr; + + ~TensorHandle() { + if (tensor) nvte_destroy_tensor(tensor); + } + + TensorHandle() = default; + TensorHandle(const TensorHandle&) = delete; + TensorHandle& operator=(const TensorHandle&) = delete; + + TensorHandle(TensorHandle&& o) noexcept : tensor(o.tensor), dev_ptr(o.dev_ptr) { + o.tensor = nullptr; o.dev_ptr = nullptr; + } + TensorHandle& operator=(TensorHandle&& o) noexcept { + if (this != &o) { + if (tensor) nvte_destroy_tensor(tensor); + tensor = o.tensor; dev_ptr = o.dev_ptr; + o.tensor = nullptr; o.dev_ptr = nullptr; + } + return *this; + } +}; + +static TensorHandle make_nvte_tensor(void* dev_ptr, + const std::vector& shape, + NVTEDType dtype) { + TensorHandle h; + h.dev_ptr = dev_ptr; + h.tensor = nvte_create_tensor(NVTE_DELAYED_TENSOR_SCALING); + + NVTEShape s; + s.ndim = shape.size(); + for (size_t i = 0; i < shape.size(); ++i) s.data[i] = shape[i]; + + NVTEBasicTensor bt; + bt.data_ptr = dev_ptr; + bt.dtype = dtype; + bt.shape = s; + nvte_set_tensor_param_v2(h.tensor, kNVTERowwiseData, &bt, sizeof(bt)); + + return h; +} + +// RAII owner for a cudaMalloc'd device buffer; frees on destruction. +template +struct DevBuf { + T* ptr = nullptr; + size_t count = 0; + + DevBuf() = default; + explicit DevBuf(size_t n) { alloc(n); } + ~DevBuf() { reset(); } + + DevBuf(const DevBuf&) = delete; + DevBuf& operator=(const DevBuf&) = delete; + DevBuf(DevBuf&& o) noexcept : ptr(o.ptr), count(o.count) { o.ptr = nullptr; o.count = 0; } + DevBuf& operator=(DevBuf&& o) noexcept { + if (this != &o) { reset(); ptr = o.ptr; count = o.count; o.ptr = nullptr; o.count = 0; } + return *this; + } + + void alloc(size_t n) { + reset(); + count = n; + if (n > 0) { + cudaError_t e = cudaMalloc(&ptr, n * sizeof(T)); + if (e != cudaSuccess) { + fprintf(stderr, "DevBuf cudaMalloc(%zu) failed: %s\n", n * sizeof(T), + cudaGetErrorString(e)); + ptr = nullptr; + count = 0; + } + } + } + + void reset() { + if (ptr) { cudaFree(ptr); ptr = nullptr; } + count = 0; + } + + T* get() const { return ptr; } + size_t bytes() const { return count * sizeof(T); } +}; + +// ── Shared routing helper ───────────────────────────────────────────────────── + +// Balanced round-robin routing: token t on rank r maps top_k experts to +// (r * num_local_experts + t * top_k + k) % num_experts +static inline std::vector routing_balanced( + int rank, int num_tokens, int top_k, int num_experts, int num_local_experts) { + std::vector idx(num_tokens * top_k); + for (int t = 0; t < num_tokens; ++t) + for (int k = 0; k < top_k; ++k) + idx[t * top_k + k] = (rank * num_local_experts + t * top_k + k) % num_experts; + return idx; +} + +// ── File-based ncclUniqueId exchange ───────────────────────────────────────── + +static void exchange_unique_id(ncclUniqueId* uid) { + const size_t sz = sizeof(ncclUniqueId); + + if (g_process_id == 0) { + ASSERT_NCCL_OK(ncclGetUniqueId(uid)); + FILE* f = fopen(g_uid_file.c_str(), "wb"); + if (!f) { fprintf(stderr, "Cannot open uid file: %s\n", g_uid_file.c_str()); exit(EXIT_FAILURE); } + fwrite(uid, 1, sz, f); + fclose(f); + } else { + auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(60); + while (true) { + FILE* f = fopen(g_uid_file.c_str(), "rb"); + if (f) { + fseek(f, 0, SEEK_END); + if (static_cast(ftell(f)) >= sz) { + fseek(f, 0, SEEK_SET); + size_t n = fread(uid, 1, sz, f); + fclose(f); + if (n == sz) break; + } else { + fclose(f); + } + } + if (std::chrono::steady_clock::now() > deadline) { + fprintf(stderr, "Process %d: timed out waiting for uid file\n", g_process_id); + exit(EXIT_FAILURE); + } + std::this_thread::sleep_for(std::chrono::milliseconds(50)); + } + } +} + +// ── CLI parsing ─────────────────────────────────────────────────────────────── + +static void ep_parse_args(int argc, char* argv[]) { + for (int i = 1; i < argc; ++i) { + std::string a(argv[i]); + if (a.rfind("--process-id=", 0) == 0) g_process_id = std::stoi(a.substr(13)); + else if (a.rfind("--rank=", 0) == 0) g_process_id = std::stoi(a.substr(7)); + else if (a.rfind("--num-processes=",0)==0) g_num_processes = std::stoi(a.substr(16)); + else if (a.rfind("--nranks=", 0) == 0) g_num_processes = std::stoi(a.substr(9)); + else if (a.rfind("--uid-file=", 0) == 0) g_uid_file = a.substr(11); + } + + if (g_process_id < 0 || g_num_processes <= 0) { + fprintf(stderr, + "Usage: %s --rank=N --nranks=N [--uid-file=path] [gtest flags]\n" + " Aliases: --process-id=N, --num-processes=N\n", + argc > 0 ? argv[0] : "test_ep"); + exit(EXIT_FAILURE); + } + + if (g_uid_file.empty()) { + const char* t = getenv("TMPDIR"); if (!t) t = "/tmp"; + g_uid_file = std::string(t) + "/te_ep_uid_" + std::to_string(g_process_id); + } +} + +// ── Bootstrap / teardown ────────────────────────────────────────────────────── + +// Returns false if the binary should exit without running tests (wrong SM, etc.). +static bool ep_bootstrap(int argc, char* argv[]) { + ep_parse_args(argc, argv); + ::testing::InitGoogleTest(&argc, argv); + + int device_count; + cudaGetDeviceCount(&device_count); + cudaSetDevice(g_process_id % device_count); + + int device, major; + cudaGetDevice(&device); + cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device); + g_sm_major = major; + if (major < 9) { + if (g_process_id == 0) + printf("SKIP: EP requires SM_90+ (device is SM_%d0)\n", major); + return false; + } + if (g_num_processes < 2) { + if (g_process_id == 0) + printf("SKIP: at least 2 processes required\n"); + return false; + } + + g_ep_size = g_num_processes; + g_num_experts = g_ep_size * 4; // 4 experts per rank + + ncclUniqueId uid{}; + exchange_unique_id(&uid); + + NVTEEpGroupConfig group_config{}; + group_config.ep_size = g_ep_size; + group_config.num_experts = g_num_experts; + group_config.max_tokens_per_rank = g_max_tokens_per_rank; + // Worst-case for top_k fan-out: ep_size * max_tokens_per_rank * 2. + group_config.max_recv_tokens_per_rank = g_ep_size * g_max_tokens_per_rank * 2; + group_config.hidden_dim = g_hidden_dim; + + ASSERT_NCCL_OK(ncclCommInitRank(&g_ep_comm, g_num_processes, uid, g_process_id)); + nvte_ep_initialize(static_cast(g_ep_comm), group_config); + + if (g_process_id == 0) { + printf("EP initialized: ep_size=%d num_experts=%d " + "hidden_dim=%d max_tokens_per_rank=%d\n", + g_ep_size, g_num_experts, g_hidden_dim, g_max_tokens_per_rank); + } + + g_ep_initialized = true; + return true; +} + +// Tear down in dependency order: backend's ep_group reads from ep_comm, +// so destroy the group first, then the comm. +static void ep_teardown() { + if (g_ep_initialized) { + nvte_ep_shutdown(); + if (g_ep_comm != nullptr) { + ncclCommDestroy(g_ep_comm); + g_ep_comm = nullptr; + } + g_ep_initialized = false; + } + if (g_process_id == 0) remove(g_uid_file.c_str()); +} diff --git a/tests/cpp_distributed/test_ep_coverage.cu b/tests/cpp_distributed/test_ep_coverage.cu new file mode 100644 index 0000000000..ef7941905d --- /dev/null +++ b/tests/cpp_distributed/test_ep_coverage.cu @@ -0,0 +1,379 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/* + * EP C-API coverage tests (paths not exercised by the pipeline suite). + * + * MultiHandleAllocTest — distinct handle ids; each works end-to-end. + * TopK1Test — top_k=1 dispatch/combine/bwd round-trip. + * EmptyExpertsTest — alignment ∈ {0, 2, 8, 16} with experts receiving 0 tokens. + * NegativeTests — alignment mismatch and null handle_mem must throw. + */ + +#include "test_ep_common.h" + +#include +#include + +// top1 -> expert 0, top2 -> expert 2; leaves local-expert 1 empty between two +// full experts. Requires top_k >= 2 and num_experts >= 3. +static std::vector routing_skip_middle(int num_tokens, int top_k) { + std::vector idx(num_tokens * top_k); + for (int t = 0; t < num_tokens; ++t) { + idx[t * top_k + 0] = 0; + if (top_k >= 2) idx[t * top_k + 1] = 2; + for (int k = 2; k < top_k; ++k) idx[t * top_k + k] = 2 + k; // distinct stragglers + } + return idx; +} + +static std::vector tokens_constant(int num_tokens, int hidden_dim, float val) { + std::vector v(num_tokens * hidden_dim); + nv_bfloat16 b = __float2bfloat16(val); + std::fill(v.begin(), v.end(), b); + return v; +} + +namespace { + +class EpCoverageBase : public ::testing::Test { + protected: + int ep_size_, num_experts_, num_local_experts_, hidden_dim_; + int max_tokens_per_rank_; + + void SetUp() override { + if (g_sm_major < 9) + GTEST_SKIP() << "EP requires SM_90+ (device is SM_" << g_sm_major << "0)"; + ASSERT_GE(g_num_processes, 2); + ASSERT_TRUE(g_ep_initialized); + ep_size_ = g_ep_size; + num_experts_ = g_num_experts; + num_local_experts_ = num_experts_ / ep_size_; + hidden_dim_ = g_hidden_dim; + max_tokens_per_rank_ = g_max_tokens_per_rank; + } + + // Helper: allocate buffers + tensor views for a single dispatch+combine. + struct Bundle { + DevBuf topk_idx; + DevBuf topk_weights; + DevBuf tokens; + DevBuf token_counts; + DevBuf handle_mem; + DevBuf recv_tokens; + DevBuf recv_topk_weights; + DevBuf result; + uint64_t handle_id = 0; + size_t handle_mem_size = 0; + size_t recv_capacity = 0; + }; + + Bundle make_bundle(int num_tokens, int top_k, int num_local_experts, + size_t alignment) { + Bundle b; + b.recv_capacity = static_cast(ep_size_) * max_tokens_per_rank_ * 2; + b.topk_idx.alloc(num_tokens * top_k); + b.topk_weights.alloc(num_tokens * top_k); + b.tokens.alloc(num_tokens * hidden_dim_); + b.token_counts.alloc(num_local_experts); + b.recv_tokens.alloc(b.recv_capacity * hidden_dim_); + b.recv_topk_weights.alloc(b.recv_capacity); + b.result.alloc(num_tokens * hidden_dim_); + NVTEEpLayerConfig cfg{num_local_experts, top_k, alignment}; + b.handle_id = nvte_ep_register_layer(cfg, &b.handle_mem_size); + b.handle_mem.alloc(b.handle_mem_size); + return b; + } +}; + +} // namespace + +// ============================================================================= +// MultiHandleAllocTest: ids are distinct and each is independently usable. +// ============================================================================= + +class MultiHandleAllocTest : public EpCoverageBase {}; + +TEST_F(MultiHandleAllocTest, IdsAreDistinct) { + NVTEEpLayerConfig cfg{num_local_experts_, /*top_k=*/2, /*alignment=*/0}; + const int kN = 8; + std::vector ids(kN); + for (int i = 0; i < kN; ++i) { + size_t sz = 0; + ids[i] = nvte_ep_register_layer(cfg, &sz); + } + for (int i = 0; i < kN; ++i) { + EXPECT_NE(ids[i], 0u) << "handle_id 0 is reserved as \"no id\""; + for (int j = i + 1; j < kN; ++j) + EXPECT_NE(ids[i], ids[j]) << "duplicate id " << ids[i] << " at indices " << i << ", " << j; + } +} + +TEST_F(MultiHandleAllocTest, TwoHandlesCoexist) { + const int num_tokens = 16, top_k = 2; + Bundle a = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); + Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); + + auto h_idx = routing_balanced(g_process_id, num_tokens, top_k, + num_experts_, num_local_experts_); + std::vector h_w(num_tokens * top_k, 1.0f / top_k); + auto h_tok = tokens_constant(num_tokens, hidden_dim_, 0.5f); + for (Bundle* x : {&a, &b}) { + CHECK_CUDA(cudaMemcpy(x->topk_idx.get(), h_idx.data(), + h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(x->topk_weights.get(), h_w.data(), + h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(x->tokens.get(), h_tok.data(), + h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); + } + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + ASSERT_NE(a.handle_id, b.handle_id); + + auto run_one = [&](Bundle& x) { + auto topk_idx = make_nvte_tensor(x.topk_idx.get(), {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); + auto topk_weights = make_nvte_tensor(x.topk_weights.get(), {(size_t)num_tokens, (size_t)top_k}, kNVTEFloat32); + auto token_counts = make_nvte_tensor(x.token_counts.get(), {(size_t)num_local_experts_}, kNVTEInt32); + auto handle_mem = make_nvte_tensor(x.handle_mem.get(), {x.handle_mem_size}, kNVTEByte); + auto tokens = make_nvte_tensor(x.tokens.get(), {(size_t)num_tokens, (size_t)hidden_dim_}, kNVTEBFloat16); + auto recv_tokens = make_nvte_tensor(x.recv_tokens.get(), {x.recv_capacity, (size_t)hidden_dim_}, kNVTEBFloat16); + auto recv_w = make_nvte_tensor(x.recv_topk_weights.get(), {x.recv_capacity}, kNVTEFloat32); + auto result = make_nvte_tensor(x.result.get(), {(size_t)num_tokens, (size_t)hidden_dim_}, kNVTEBFloat16); + NVTEEpHandle h{x.handle_id, handle_mem.tensor}; + ASSERT_NO_THROW(nvte_ep_prepare(h, topk_idx.tensor, token_counts.tensor, + /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(h, topk_idx.tensor, tokens.tensor, + NVTECommWindow{}, topk_weights.tensor, NVTECommWindow{}, + recv_tokens.tensor, NVTECommWindow{}, recv_w.tensor, + NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(h, recv_tokens.tensor, NVTECommWindow{}, + result.tensor, stream)); + }; + run_one(a); + run_one(b); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + // Both round-trips must produce result == top_k * 0.5 = 1.0. + for (Bundle* x : {&a, &b}) { + std::vector h_res(num_tokens * hidden_dim_); + CHECK_CUDA(cudaMemcpy(h_res.data(), x->result.get(), + h_res.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + const int probes[3] = {0, hidden_dim_ / 2, hidden_dim_ - 1}; + for (int t = 0; t < num_tokens; ++t) + for (int p : probes) + EXPECT_NEAR(__bfloat162float(h_res[t * hidden_dim_ + p]), + static_cast(top_k) * 0.5f, 1e-2f); + } + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// ============================================================================= +// TopK1Test: top_k=1 dispatch/combine round-trip, including dispatch_bwd. +// ============================================================================= + +class TopK1Test : public EpCoverageBase {}; + +TEST_F(TopK1Test, RoundTrip) { + const int num_tokens = 16, top_k = 1; + Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); + + auto h_idx = routing_balanced(g_process_id, num_tokens, top_k, + num_experts_, num_local_experts_); + std::vector h_w(num_tokens * top_k, 1.0f); // top_k=1: weight is unity + auto h_tok = tokens_constant(num_tokens, hidden_dim_, 0.25f); + CHECK_CUDA(cudaMemcpy(b.topk_idx.get(), h_idx.data(), + h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(b.topk_weights.get(), h_w.data(), + h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(b.tokens.get(), h_tok.data(), + h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); + + auto topk_idx_t = make_nvte_tensor(b.topk_idx.get(), + {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); + auto topk_weights_t = make_nvte_tensor(b.topk_weights.get(), + {(size_t)num_tokens, (size_t)top_k}, kNVTEFloat32); + auto token_counts_t = make_nvte_tensor(b.token_counts.get(), + {(size_t)num_local_experts_}, kNVTEInt32); + auto handle_mem_t = make_nvte_tensor(b.handle_mem.get(), + {b.handle_mem_size}, kNVTEByte); + auto tokens_t = make_nvte_tensor(b.tokens.get(), + {(size_t)num_tokens, (size_t)hidden_dim_}, kNVTEBFloat16); + auto recv_tokens_t = make_nvte_tensor(b.recv_tokens.get(), + {b.recv_capacity, (size_t)hidden_dim_}, kNVTEBFloat16); + auto recv_w_t = make_nvte_tensor(b.recv_topk_weights.get(), + {b.recv_capacity}, kNVTEFloat32); + auto result_t = make_nvte_tensor(b.result.get(), + {(size_t)num_tokens, (size_t)hidden_dim_}, kNVTEBFloat16); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + NVTEEpHandle h{b.handle_id, handle_mem_t.tensor}; + ASSERT_NO_THROW(nvte_ep_prepare(h, topk_idx_t.tensor, token_counts_t.tensor, + /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(h, topk_idx_t.tensor, + tokens_t.tensor, NVTECommWindow{}, topk_weights_t.tensor, + NVTECommWindow{}, recv_tokens_t.tensor, NVTECommWindow{}, + recv_w_t.tensor, NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(h, recv_tokens_t.tensor, + NVTECommWindow{}, result_t.tensor, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + // top_k=1: combine is unweighted gather, so result[t] == tokens[t]. + std::vector h_res(num_tokens * hidden_dim_); + CHECK_CUDA(cudaMemcpy(h_res.data(), b.result.get(), + h_res.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + const int probes[3] = {0, hidden_dim_ / 2, hidden_dim_ - 1}; + for (int t = 0; t < num_tokens; ++t) + for (int p : probes) + EXPECT_NEAR(__bfloat162float(h_res[t * hidden_dim_ + p]), 0.25f, 1e-2f) + << "tok " << t << " hidden " << p; + + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// ============================================================================= +// EmptyExpertsTest: alignment ∈ {0, 2, 8, 16}, only local-expert 0 receives +// tokens. Round-trip must produce result == top_k * tokens regardless of the +// per-expert padding choice. +// ============================================================================= + +class EmptyExpertsTest : public EpCoverageBase, + public ::testing::WithParamInterface {}; + +TEST_P(EmptyExpertsTest, RoundTripCorrect) { + // routing_skip_middle needs experts {0, 2, ...}; smallest viable num_experts is 3. + ASSERT_GE(num_experts_, 3); + const size_t alignment = GetParam(); + const int num_tokens = 16, top_k = 2; + Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, alignment); + + // top1 -> expert 0, top2 -> expert 2; rank 0's local-expert 1 receives 0 + // tokens between two non-empty experts. + std::vector h_idx = routing_skip_middle(num_tokens, top_k); + std::vector h_w(num_tokens * top_k, 1.0f / top_k); + auto h_tok = tokens_constant(num_tokens, hidden_dim_, 0.3f); + + CHECK_CUDA(cudaMemcpy(b.topk_idx.get(), h_idx.data(), + h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(b.topk_weights.get(), h_w.data(), + h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(b.tokens.get(), h_tok.data(), + h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); + + auto topk_idx_t = make_nvte_tensor(b.topk_idx.get(), + {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); + auto topk_weights_t = make_nvte_tensor(b.topk_weights.get(), + {(size_t)num_tokens, (size_t)top_k}, kNVTEFloat32); + auto token_counts_t = make_nvte_tensor(b.token_counts.get(), + {(size_t)num_local_experts_}, kNVTEInt32); + auto handle_mem_t = make_nvte_tensor(b.handle_mem.get(), + {b.handle_mem_size}, kNVTEByte); + auto tokens_t = make_nvte_tensor(b.tokens.get(), + {(size_t)num_tokens, (size_t)hidden_dim_}, kNVTEBFloat16); + auto recv_tokens_t = make_nvte_tensor(b.recv_tokens.get(), + {b.recv_capacity, (size_t)hidden_dim_}, kNVTEBFloat16); + auto recv_w_t = make_nvte_tensor(b.recv_topk_weights.get(), + {b.recv_capacity}, kNVTEFloat32); + auto result_t = make_nvte_tensor(b.result.get(), + {(size_t)num_tokens, (size_t)hidden_dim_}, kNVTEBFloat16); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + NVTEEpHandle h{b.handle_id, handle_mem_t.tensor}; + ASSERT_NO_THROW(nvte_ep_prepare(h, topk_idx_t.tensor, token_counts_t.tensor, + alignment, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(h, topk_idx_t.tensor, + tokens_t.tensor, NVTECommWindow{}, topk_weights_t.tensor, + NVTECommWindow{}, recv_tokens_t.tensor, NVTECommWindow{}, + recv_w_t.tensor, NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(h, recv_tokens_t.tensor, + NVTECommWindow{}, result_t.tensor, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + // Identity expert + uniform weights: result[t] == top_k * tokens[t]. + std::vector h_res(num_tokens * hidden_dim_); + CHECK_CUDA(cudaMemcpy(h_res.data(), b.result.get(), + h_res.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + const float expected = static_cast(top_k) * 0.3f; + const int probes[3] = {0, hidden_dim_ / 2, hidden_dim_ - 1}; + for (int t = 0; t < num_tokens; ++t) + for (int p : probes) + EXPECT_NEAR(__bfloat162float(h_res[t * hidden_dim_ + p]), expected, 1e-2f) + << "alignment=" << alignment << " tok=" << t << " hidden=" << p; + + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +INSTANTIATE_TEST_SUITE_P(Alignments, EmptyExpertsTest, + ::testing::Values(0, 2, 8, 16)); + +// ============================================================================= +// NegativeTests: prepare/dispatch must surface bad inputs as exceptions. +// ============================================================================= + +class NegativeTests : public EpCoverageBase {}; + +TEST_F(NegativeTests, AlignmentMismatchThrows) { + const int num_tokens = 8, top_k = 2; + // Allocate handle for alignment=0, then call prepare with alignment=16. + Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); + auto h_idx = routing_balanced(g_process_id, num_tokens, top_k, + num_experts_, num_local_experts_); + CHECK_CUDA(cudaMemcpy(b.topk_idx.get(), h_idx.data(), + h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); + + auto topk_idx_t = make_nvte_tensor(b.topk_idx.get(), + {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); + auto token_counts_t = make_nvte_tensor(b.token_counts.get(), + {(size_t)num_local_experts_}, kNVTEInt32); + auto handle_mem_t = make_nvte_tensor(b.handle_mem.get(), + {b.handle_mem_size}, kNVTEByte); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + NVTEEpHandle h{b.handle_id, handle_mem_t.tensor}; + EXPECT_THROW(nvte_ep_prepare(h, topk_idx_t.tensor, token_counts_t.tensor, + /*alignment=*/16, stream), + std::exception); + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +TEST_F(NegativeTests, NullHandleMemThrows) { + const int num_tokens = 8, top_k = 2; + Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); + auto h_idx = routing_balanced(g_process_id, num_tokens, top_k, + num_experts_, num_local_experts_); + CHECK_CUDA(cudaMemcpy(b.topk_idx.get(), h_idx.data(), + h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); + + auto topk_idx_t = make_nvte_tensor(b.topk_idx.get(), + {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); + auto token_counts_t = make_nvte_tensor(b.token_counts.get(), + {(size_t)num_local_experts_}, kNVTEInt32); + // Construct a tensor view backed by a null device pointer. + auto null_hm_t = make_nvte_tensor(nullptr, {b.handle_mem_size}, kNVTEByte); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + NVTEEpHandle h{b.handle_id, null_hm_t.tensor}; + EXPECT_THROW(nvte_ep_prepare(h, topk_idx_t.tensor, token_counts_t.tensor, + /*alignment=*/0, stream), + std::exception); + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// ── main ────────────────────────────────────────────────────────────────────── + +int main(int argc, char* argv[]) { + if (!ep_bootstrap(argc, argv)) return 0; + int ret = RUN_ALL_TESTS(); + ep_teardown(); + return ret; +} diff --git a/tests/cpp_distributed/test_ep_init.cu b/tests/cpp_distributed/test_ep_init.cu new file mode 100644 index 0000000000..08744dfee5 --- /dev/null +++ b/tests/cpp_distributed/test_ep_init.cu @@ -0,0 +1,64 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/* + * Unit tests for EP initialization paths. + * + * Tests: + * EPInitTest/InitPath — backend is live after init, handle_mem_size > 0 + * EPInitTest/NumLocalExperts — handle_mem_size is consistent across num_local_experts values + * + * Run via run_test_ep.sh (both uid and comm init paths are tested by the script). + */ + +#include "test_ep_common.h" + +// ── Fixture ─────────────────────────────────────────────────────────────────── + +class EPInitTest : public ::testing::Test { + protected: + void SetUp() override { + if (g_sm_major < 9) + GTEST_SKIP() << "EP requires SM_90+ (device is SM_" << g_sm_major << "0)"; + ASSERT_GE(g_num_processes, 2) << "EP tests require at least 2 processes"; + ASSERT_TRUE(g_ep_initialized) << "EP not initialized"; + } +}; + +// ── Tests ───────────────────────────────────────────────────────────────────── + +TEST_F(EPInitTest, InitPath) { + int nle = g_num_experts / g_ep_size; + NVTEEpLayerConfig cfg{nle, /*top_k=*/2}; + size_t sz = 0; + (void)nvte_ep_register_layer(cfg, &sz); + ASSERT_GT(sz, 0u) << "handle_mem_size must be > 0 after init"; + + if (g_process_id == 0) { + printf(" handle_mem : %zu bytes\n", sz); + } +} + +TEST_F(EPInitTest, NumLocalExperts) { + // handle_mem_size should be > 0 for any valid num_local_experts value. + for (int nle : {1, g_num_experts / g_ep_size}) { + NVTEEpLayerConfig cfg{nle, /*top_k=*/2}; + size_t sz = 0; + (void)nvte_ep_register_layer(cfg, &sz); + ASSERT_GT(sz, 0u) << "num_local_experts=" << nle; + if (g_process_id == 0) + printf(" nle=%-3d handle_mem_size=%zu bytes\n", nle, sz); + } +} + +// ── main ────────────────────────────────────────────────────────────────────── + +int main(int argc, char* argv[]) { + if (!ep_bootstrap(argc, argv)) return 0; + int ret = RUN_ALL_TESTS(); + ep_teardown(); + return ret; +} diff --git a/tests/cpp_distributed/test_ep_pipeline.cu b/tests/cpp_distributed/test_ep_pipeline.cu new file mode 100644 index 0000000000..41f83a6d11 --- /dev/null +++ b/tests/cpp_distributed/test_ep_pipeline.cu @@ -0,0 +1,890 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/* + * EP pipeline tests: smallest-scope first. + * + * EPDispatchTest/PrepareAndDispatch — exact recv values + per-expert counts + * EPCombineTest/Combine — round-trip: out == top_k * tokens + * EPCombineBwdTest/CombineBwdCheck — exact grad_expert values + * EPDispatchBwdTest/DispatchBwdCheck — exact grad_tokens + * EPDispatchBwdGradWeightsTest/RoundTrip — exact per-(t, k) grad_topk_weights + * EPPipelineTest/FullForwardBackward — fwd + bwd NaN/Inf check + * + * Routing: token t on rank r → expert (r * num_local_experts + t * top_k + k) % num_experts + * Token values: rank r, token t → all hidden dims = (r+1)*0.01 + t*0.001 + * + * Closed-form expected values: + * dispatch recv: multiset of source-token values routed to this rank's experts + * combine: result[t] == top_k * tokens[t] + * combine_bwd: grad_expert[slot] == d_result[t] (no weighting) + * dispatch_bwd: grad_tokens[t] == top_k * d_result[t] + */ + +#include "test_ep_common.h" + +#include +#include +#include +#include + +// ── Deterministic routing helpers ───────────────────────────────────────────── + +// Token value for (rank, t): (rank * num_tokens + t + 1) / 256. Step 1/256 is +// bf16-exact and unique across (rank, t) when rank * num_tokens + t < 256. +static inline float token_value(int rank, int t, int num_tokens) { + return static_cast(rank * num_tokens + t + 1) * (1.0f / 256.0f); +} + +static std::vector generate_tokens(int rank, int num_tokens, int hidden_dim) { + std::vector v(num_tokens * hidden_dim); + for (int t = 0; t < num_tokens; ++t) { + nv_bfloat16 val = __float2bfloat16(token_value(rank, t, num_tokens)); + for (int h = 0; h < hidden_dim; ++h) + v[t * hidden_dim + h] = val; + } + return v; +} + +static std::vector expected_token_counts( + int recv_rank, int num_processes, int num_tokens, int top_k, + int num_experts, int num_local_experts) { + int base = recv_rank * num_local_experts; + std::vector cnt(num_local_experts, 0); + for (int src = 0; src < num_processes; ++src) { + auto idx = routing_balanced(src, num_tokens, top_k, num_experts, num_local_experts); + for (int t = 0; t < num_tokens; ++t) + for (int k = 0; k < top_k; ++k) { + int64_t e = idx[t * top_k + k]; + if (e >= base && e < base + num_local_experts) ++cnt[e - base]; + } + } + return cnt; +} + +static std::vector expected_recv_values_sorted( + int recv_rank, int num_processes, int num_tokens, int top_k, + int num_experts, int num_local_experts) { + int base = recv_rank * num_local_experts; + std::vector vals; + for (int src = 0; src < num_processes; ++src) { + auto idx = routing_balanced(src, num_tokens, top_k, num_experts, num_local_experts); + for (int t = 0; t < num_tokens; ++t) + for (int k = 0; k < top_k; ++k) { + int64_t e = idx[t * top_k + k]; + if (e >= base && e < base + num_local_experts) { + float raw = token_value(src, t, num_tokens); + vals.push_back(__bfloat162float(__float2bfloat16(raw))); + } + } + } + std::sort(vals.begin(), vals.end()); + return vals; +} + +// BF16 has 7 mantissa bits; relative ULP ≈ 2^-7. Use 4× headroom for +// accumulation noise inside dispatch/combine. +static float bf16_tol(float magnitude) { + return 4.f * std::ldexp(std::fabs(magnitude) + 1e-3f, -7); +} + +static bool check_no_nan_inf(const nv_bfloat16* dev, int count, const char* name) { + std::vector h(count); + cudaMemcpy(h.data(), dev, count * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost); + for (int i = 0; i < count; ++i) { + float v = __bfloat162float(h[i]); + if (std::isnan(v) || std::isinf(v)) { + fprintf(stderr, "Rank %d: %s in %s[%d]\n", + g_process_id, std::isnan(v) ? "NaN" : "Inf", name, i); + return false; + } + } + return true; +} + +// ── Forward buffer set with RAII ────────────────────────────────────────────── + +struct EPBuffers { + // Forward + DevBuf topk_idx; + DevBuf topk_weights; + DevBuf tokens; + DevBuf token_counts; + DevBuf handle_mem; + DevBuf recv_tokens; + DevBuf recv_topk_weights; + DevBuf result; + // Backward + DevBuf grad_result; + DevBuf grad_expert; + DevBuf grad_tokens; + DevBuf g_recv_topk_weights; + DevBuf grad_topk_weights; + + uint64_t handle_id = 0; + size_t handle_mem_size = 0; + size_t recv_capacity = 0; + int top_k_ = 0; + + void alloc(int num_tokens, int top_k, int hidden_dim, int num_local_experts, + int ep_size, int max_tokens_per_rank, size_t alignment = 0) { + top_k_ = top_k; + recv_capacity = static_cast(ep_size) * max_tokens_per_rank * 2; + + topk_idx.alloc(num_tokens * top_k); + topk_weights.alloc(num_tokens * top_k); + tokens.alloc(num_tokens * hidden_dim); + token_counts.alloc(num_local_experts); + recv_tokens.alloc(recv_capacity * hidden_dim); + recv_topk_weights.alloc(recv_capacity); + result.alloc(num_tokens * hidden_dim); + + NVTEEpLayerConfig cfg{num_local_experts, top_k, alignment}; + handle_id = nvte_ep_register_layer(cfg, &handle_mem_size); + handle_mem.alloc(handle_mem_size); + + grad_result.alloc(num_tokens * hidden_dim); + grad_expert.alloc(recv_capacity * hidden_dim); + grad_tokens.alloc(num_tokens * hidden_dim); + g_recv_topk_weights.alloc(recv_capacity); + grad_topk_weights.alloc(num_tokens * top_k); + } +}; + +// Bundled NVTETensor views over an EPBuffers — one place to update the shape +// conventions when the C-API evolves. +struct EPTensors { + TensorHandle topk_idx, topk_weights, token_counts, handle_mem, tokens; + TensorHandle recv_tokens, recv_topk_weights, result; + TensorHandle grad_result, grad_expert, grad_tokens; + TensorHandle g_recv_topk_weights, grad_topk_weights; + + EPTensors(EPBuffers& b, int num_tokens, int top_k, int hidden_dim, + int num_local_experts) { + topk_idx = make_nvte_tensor(b.topk_idx.get(), + {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); + topk_weights = make_nvte_tensor(b.topk_weights.get(), + {(size_t)num_tokens, (size_t)top_k}, kNVTEFloat32); + token_counts = make_nvte_tensor(b.token_counts.get(), + {(size_t)num_local_experts}, kNVTEInt32); + handle_mem = make_nvte_tensor(b.handle_mem.get(), + {b.handle_mem_size}, kNVTEByte); + tokens = make_nvte_tensor(b.tokens.get(), + {(size_t)num_tokens, (size_t)hidden_dim}, kNVTEBFloat16); + recv_tokens = make_nvte_tensor(b.recv_tokens.get(), + {b.recv_capacity, (size_t)hidden_dim}, kNVTEBFloat16); + recv_topk_weights = make_nvte_tensor(b.recv_topk_weights.get(), + {b.recv_capacity}, kNVTEFloat32); + result = make_nvte_tensor(b.result.get(), + {(size_t)num_tokens, (size_t)hidden_dim}, kNVTEBFloat16); + grad_result = make_nvte_tensor(b.grad_result.get(), + {(size_t)num_tokens, (size_t)hidden_dim}, kNVTEBFloat16); + grad_expert = make_nvte_tensor(b.grad_expert.get(), + {b.recv_capacity, (size_t)hidden_dim}, kNVTEBFloat16); + grad_tokens = make_nvte_tensor(b.grad_tokens.get(), + {(size_t)num_tokens, (size_t)hidden_dim}, kNVTEBFloat16); + g_recv_topk_weights = make_nvte_tensor(b.g_recv_topk_weights.get(), + {b.recv_capacity}, kNVTEFloat32); + grad_topk_weights = make_nvte_tensor(b.grad_topk_weights.get(), + {(size_t)num_tokens, (size_t)top_k}, kNVTEFloat32); + } +}; + +// ── Shared fixture base ─────────────────────────────────────────────────────── + +class EpOpTestBase : public ::testing::Test { + protected: + int ep_size_, num_experts_, num_local_experts_, hidden_dim_; + int max_tokens_per_rank_, top_k_, num_tokens_; + + void SetUp() override { + if (g_sm_major < 9) + GTEST_SKIP() << "EP requires SM_90+ (device is SM_" << g_sm_major << "0)"; + ASSERT_GE(g_num_processes, 2); + ASSERT_TRUE(g_ep_initialized); + + ep_size_ = g_ep_size; + num_experts_ = g_num_experts; + num_local_experts_ = num_experts_ / ep_size_; + hidden_dim_ = g_hidden_dim; + max_tokens_per_rank_ = g_max_tokens_per_rank; + top_k_ = 2; + num_tokens_ = 32; + } + + void upload_inputs(EPBuffers& buf, int rank = -1) { + if (rank < 0) rank = g_process_id; + auto h_idx = routing_balanced(rank, num_tokens_, top_k_, + num_experts_, num_local_experts_); + std::vector h_w(num_tokens_ * top_k_, 1.0f / top_k_); + auto h_tok = generate_tokens(rank, num_tokens_, hidden_dim_); + + CHECK_CUDA(cudaMemcpy(buf.topk_idx.get(), h_idx.data(), + h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(buf.topk_weights.get(), h_w.data(), + h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(buf.tokens.get(), h_tok.data(), + h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); + } + + NVTEEpLayerConfig layer_config(size_t alignment = 0) const { + return NVTEEpLayerConfig{num_local_experts_, top_k_, alignment}; + } + + // ASSERT_CUDA_OK (fprintf+exit) so this non-void helper stays legal. + int read_total_recv(const EPBuffers& buf) const { + std::vector cnt(num_local_experts_); + ASSERT_CUDA_OK(cudaMemcpy(cnt.data(), buf.token_counts.get(), + num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost)); + int total = 0; + for (int c : cnt) total += c; + return total; + } +}; + +// ============================================================================= +// EPDispatchTest: exact recv values and per-expert counts. +// ============================================================================= + +class EPDispatchTest : public EpOpTestBase {}; + +TEST_F(EPDispatchTest, PrepareAndDispatch) { + EPBuffers buf; + buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, + ep_size_, max_tokens_per_rank_); + upload_inputs(buf); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + + CHECK_CUDA(cudaMemset(buf.recv_tokens.get(), 0, buf.recv_tokens.bytes())); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + uint64_t handle_id = buf.handle_id; + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, + t.tokens.tensor, NVTECommWindow{}, t.topk_weights.tensor, + NVTECommWindow{}, t.recv_tokens.tensor, NVTECommWindow{}, + t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + // 1. Per-expert counts. + std::vector got_counts(num_local_experts_); + CHECK_CUDA(cudaMemcpy(got_counts.data(), buf.token_counts.get(), + num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost)); + auto exp_counts = expected_token_counts(g_process_id, g_num_processes, num_tokens_, top_k_, + num_experts_, num_local_experts_); + int total_recv = 0; + for (int i = 0; i < num_local_experts_; ++i) { + EXPECT_EQ(got_counts[i], exp_counts[i]) << "local expert " << i; + total_recv += exp_counts[i]; + } + ASSERT_LE(total_recv, static_cast(buf.recv_capacity)) + << "total_recv exceeded recv_capacity — overflow would corrupt downstream memory"; + + // 2. Recv values: read only the filled prefix per local-expert zone, not the + // whole recv buffer — avoids false positives from legitimate-zero token values. + std::vector h_recv(buf.recv_capacity * hidden_dim_); + CHECK_CUDA(cudaMemcpy(h_recv.data(), buf.recv_tokens.get(), + h_recv.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + + std::vector got_vals; + got_vals.reserve(total_recv); + size_t slot = 0; + for (int e = 0; e < num_local_experts_; ++e) { + for (int i = 0; i < got_counts[e]; ++i) { + got_vals.push_back(__bfloat162float(h_recv[slot * hidden_dim_])); + ++slot; + } + } + std::sort(got_vals.begin(), got_vals.end()); + + auto exp_vals = expected_recv_values_sorted(g_process_id, g_num_processes, num_tokens_, + top_k_, num_experts_, num_local_experts_); + + ASSERT_EQ(got_vals.size(), exp_vals.size()); + for (size_t i = 0; i < exp_vals.size(); ++i) + EXPECT_NEAR(got_vals[i], exp_vals[i], bf16_tol(exp_vals[i])) + << "recv value mismatch at sorted index " << i; + + // 3. recv_topk_weights: every filled slot must equal the per-token weight (1/top_k). + std::vector h_w(buf.recv_capacity); + CHECK_CUDA(cudaMemcpy(h_w.data(), buf.recv_topk_weights.get(), + h_w.size() * sizeof(float), cudaMemcpyDeviceToHost)); + const float exp_w = 1.0f / static_cast(top_k_); + for (int i = 0; i < total_recv; ++i) + EXPECT_NEAR(h_w[i], exp_w, 1e-6f) << "recv_topk_weights[" << i << "]"; + + if (g_process_id == 0) + printf(" PrepareAndDispatch: passed (recv=%d, values + weights exact)\n", total_recv); + + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// ============================================================================= +// EPCombineTest: round-trip identity expert → result == top_k * tokens. +// ============================================================================= + +class EPCombineTest : public EpOpTestBase {}; + +TEST_F(EPCombineTest, Combine) { + EPBuffers buf; + buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, + ep_size_, max_tokens_per_rank_); + upload_inputs(buf); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + uint64_t handle_id = buf.handle_id; + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, + t.tokens.tensor, NVTECommWindow{}, t.topk_weights.tensor, + NVTECommWindow{}, t.recv_tokens.tensor, NVTECommWindow{}, + t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.recv_tokens.tensor, NVTECommWindow{}, + t.result.tensor, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + std::vector h_result(num_tokens_ * hidden_dim_); + CHECK_CUDA(cudaMemcpy(h_result.data(), buf.result.get(), + h_result.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + auto h_tok = generate_tokens(g_process_id, num_tokens_, hidden_dim_); + // Spot-check 3 hidden-dim positions per token to catch partial-row writes. + const int probes[3] = {0, hidden_dim_ / 2, hidden_dim_ - 1}; + for (int tok = 0; tok < num_tokens_; ++tok) { + float exp = __bfloat162float(h_tok[tok * hidden_dim_]) * static_cast(top_k_); + for (int p : probes) { + float got = __bfloat162float(h_result[tok * hidden_dim_ + p]); + EXPECT_NEAR(got, exp, bf16_tol(exp)) + << "token " << tok << " rank " << g_process_id << " hidden " << p; + } + } + + if (g_process_id == 0) + printf(" Combine: passed (result == top_k * tokens)\n"); + + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// ============================================================================= +// EPCombineBwdTest: filled slots in grad_expert == d_result (unweighted). +// ============================================================================= + +class EPCombineBwdTest : public EpOpTestBase {}; + +TEST_F(EPCombineBwdTest, CombineBwdCheck) { + EPBuffers buf; + buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, + ep_size_, max_tokens_per_rank_); + upload_inputs(buf); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + uint64_t handle_id = buf.handle_id; + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, + t.tokens.tensor, NVTECommWindow{}, t.topk_weights.tensor, + NVTECommWindow{}, t.recv_tokens.tensor, NVTECommWindow{}, + t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.recv_tokens.tensor, NVTECommWindow{}, + t.result.tensor, stream)); + + std::vector h_grad_r(num_tokens_ * hidden_dim_, __float2bfloat16(0.1f)); + CHECK_CUDA(cudaMemcpyAsync(buf.grad_result.get(), h_grad_r.data(), + h_grad_r.size() * sizeof(nv_bfloat16), + cudaMemcpyHostToDevice, stream)); + CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); + + ASSERT_NO_THROW(nvte_ep_combine_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_result.tensor, NVTECommWindow{}, + t.grad_expert.tensor, NVTECommWindow{}, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + int total_recv = read_total_recv(buf); + + std::vector cnt(num_local_experts_); + CHECK_CUDA(cudaMemcpy(cnt.data(), buf.token_counts.get(), + num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost)); + std::vector h_ge(buf.recv_capacity * hidden_dim_); + CHECK_CUDA(cudaMemcpy(h_ge.data(), buf.grad_expert.get(), + h_ge.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + + // Walk filled slots by per-expert zone (no v != 0 heuristic). + const float kExpGrad = 0.1f; + size_t slot = 0; + int filled = 0; + for (int e = 0; e < num_local_experts_; ++e) { + for (int i = 0; i < cnt[e]; ++i) { + float v = __bfloat162float(h_ge[slot * hidden_dim_]); + EXPECT_NEAR(v, kExpGrad, bf16_tol(kExpGrad)) + << "grad_expert expert " << e << " slot " << i << " (linear " << slot << ")"; + ++filled; ++slot; + } + } + EXPECT_EQ(filled, total_recv); + + if (g_process_id == 0) + printf(" CombineBwdCheck: passed (filled=%d)\n", filled); + + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// ============================================================================= +// EPDispatchBwdTest: grad_tokens == top_k * d_result. +// ============================================================================= + +class EPDispatchBwdTest : public EpOpTestBase {}; + +TEST_F(EPDispatchBwdTest, DispatchBwdCheck) { + EPBuffers buf; + buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, + ep_size_, max_tokens_per_rank_); + upload_inputs(buf); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + uint64_t handle_id = buf.handle_id; + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, + t.tokens.tensor, NVTECommWindow{}, t.topk_weights.tensor, + NVTECommWindow{}, t.recv_tokens.tensor, NVTECommWindow{}, + t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.recv_tokens.tensor, NVTECommWindow{}, + t.result.tensor, stream)); + + std::vector h_grad(num_tokens_ * hidden_dim_, __float2bfloat16(0.1f)); + CHECK_CUDA(cudaMemcpyAsync(buf.grad_result.get(), h_grad.data(), + h_grad.size() * sizeof(nv_bfloat16), + cudaMemcpyHostToDevice, stream)); + CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); + CHECK_CUDA(cudaMemsetAsync(buf.g_recv_topk_weights.get(), 0, buf.g_recv_topk_weights.bytes(), stream)); + CHECK_CUDA(cudaMemsetAsync(buf.grad_topk_weights.get(), 0, buf.grad_topk_weights.bytes(), stream)); + + ASSERT_NO_THROW(nvte_ep_combine_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_result.tensor, NVTECommWindow{}, + t.grad_expert.tensor, NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_expert.tensor, NVTECommWindow{}, + t.g_recv_topk_weights.tensor, NVTECommWindow{}, + t.grad_tokens.tensor, t.grad_topk_weights.tensor, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + std::vector h_gt(num_tokens_ * hidden_dim_); + CHECK_CUDA(cudaMemcpy(h_gt.data(), buf.grad_tokens.get(), + h_gt.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + const float kExpGrad = static_cast(top_k_) * 0.1f; + for (int tok = 0; tok < num_tokens_; ++tok) + EXPECT_NEAR(__bfloat162float(h_gt[tok * hidden_dim_]), kExpGrad, bf16_tol(kExpGrad)) + << "grad_tokens token " << tok; + + if (g_process_id == 0) + printf(" DispatchBwdCheck: passed (grad_tokens == %.2f)\n", kExpGrad); + + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// ============================================================================= +// EPDispatchBwdGradWeightsTest: round-trip per-(t, k) weights. +// ============================================================================= + +class EPDispatchBwdGradWeightsTest : public EpOpTestBase {}; + +TEST_F(EPDispatchBwdGradWeightsTest, RoundTrip) { + EPBuffers buf; + buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, + ep_size_, max_tokens_per_rank_); + upload_inputs(buf); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + + // Distinct per-(rank, t, k) weights so each slot carries a unique value. + std::vector h_w(num_tokens_ * top_k_); + for (int tok = 0; tok < num_tokens_; ++tok) + for (int k = 0; k < top_k_; ++k) + h_w[tok * top_k_ + k] = 0.1f + 0.01f * tok + 0.001f * k + + 0.0001f * (g_process_id + 1); + CHECK_CUDA(cudaMemcpy(buf.topk_weights.get(), h_w.data(), + h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + uint64_t handle_id = buf.handle_id; + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); + CHECK_CUDA(cudaMemsetAsync(buf.recv_topk_weights.get(), 0, + buf.recv_topk_weights.bytes(), stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, + t.tokens.tensor, NVTECommWindow{}, t.topk_weights.tensor, + NVTECommWindow{}, t.recv_tokens.tensor, NVTECommWindow{}, + t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); + + // Sentinel: NaN so any (t, k) the bwd kernel fails to write is immediately visible. + std::vector h_nan(num_tokens_ * top_k_, + std::numeric_limits::quiet_NaN()); + CHECK_CUDA(cudaMemcpyAsync(buf.grad_topk_weights.get(), h_nan.data(), + h_nan.size() * sizeof(float), + cudaMemcpyHostToDevice, stream)); + CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); + + // g_recv_topk_weights := recv_topk_weights (the round-trip input). + auto g_recv_t = make_nvte_tensor(buf.recv_topk_weights.get(), + {buf.recv_capacity}, kNVTEFloat32); + ASSERT_NO_THROW(nvte_ep_dispatch_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_expert.tensor, + NVTECommWindow{}, g_recv_t.tensor, NVTECommWindow{}, + t.grad_tokens.tensor, t.grad_topk_weights.tensor, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + std::vector h_grad_w(num_tokens_ * top_k_); + CHECK_CUDA(cudaMemcpy(h_grad_w.data(), buf.grad_topk_weights.get(), + h_grad_w.size() * sizeof(float), cudaMemcpyDeviceToHost)); + + const float kTol = 1e-5f; + int errs = 0, k0_eq_k1 = 0; + for (int tok = 0; tok < num_tokens_; ++tok) { + for (int k = 0; k < top_k_; ++k) { + float got = h_grad_w[tok * top_k_ + k]; + float exp = h_w[tok * top_k_ + k]; + if (std::isnan(got) || std::fabs(got - exp) > kTol) { + if (errs < 8) + fprintf(stderr, "Rank %d: grad_topk_weights[%d, %d]: got %.6f, expected %.6f\n", + g_process_id, tok, k, got, exp); + ++errs; + } + } + if (top_k_ >= 2 && + std::fabs(h_grad_w[tok * top_k_ + 0] - h_grad_w[tok * top_k_ + 1]) < 1e-7f) + ++k0_eq_k1; + } + EXPECT_EQ(errs, 0); + EXPECT_EQ(k0_eq_k1, 0) << "per-token-average regression: grad[t, 0] == grad[t, 1]"; + + if (g_process_id == 0 && errs == 0 && k0_eq_k1 == 0) + printf(" RoundTrip: passed (%d (t, k) gradients)\n", num_tokens_ * top_k_); + + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// ============================================================================= +// Integrated FwdBwd: NaN/Inf check end-to-end. +// ============================================================================= + +class EPPipelineTest : public EpOpTestBase {}; + +TEST_F(EPPipelineTest, FullForwardBackward) { + EPBuffers buf; + buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, + ep_size_, max_tokens_per_rank_); + upload_inputs(buf); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + uint64_t handle_id = buf.handle_id; + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, + t.tokens.tensor, NVTECommWindow{}, t.topk_weights.tensor, + NVTECommWindow{}, t.recv_tokens.tensor, NVTECommWindow{}, + t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.recv_tokens.tensor, NVTECommWindow{}, + t.result.tensor, stream)); + + std::vector h_grad(num_tokens_ * hidden_dim_, __float2bfloat16(0.1f)); + CHECK_CUDA(cudaMemcpyAsync(buf.grad_result.get(), h_grad.data(), + h_grad.size() * sizeof(nv_bfloat16), + cudaMemcpyHostToDevice, stream)); + CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); + CHECK_CUDA(cudaMemsetAsync(buf.g_recv_topk_weights.get(), 0, buf.g_recv_topk_weights.bytes(), stream)); + CHECK_CUDA(cudaMemsetAsync(buf.grad_topk_weights.get(), 0, buf.grad_topk_weights.bytes(), stream)); + + ASSERT_NO_THROW(nvte_ep_combine_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_result.tensor, NVTECommWindow{}, + t.grad_expert.tensor, NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_expert.tensor, NVTECommWindow{}, + t.g_recv_topk_weights.tensor, NVTECommWindow{}, + t.grad_tokens.tensor, t.grad_topk_weights.tensor, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + ASSERT_TRUE(check_no_nan_inf(buf.result.get(), num_tokens_ * hidden_dim_, "result")); + ASSERT_TRUE(check_no_nan_inf(buf.grad_tokens.get(), num_tokens_ * hidden_dim_, "grad_tokens")); + + if (g_process_id == 0) printf(" FullForwardBackward: passed\n"); + + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// ============================================================================= +// EPZeroCopyTest: dispatch/combine with NCCL symmetric-memory windows attached +// to payload tensors (zero-copy fast path via ncclEpTensorCreateFromWindow). +// Symm-mem requirements per spec: input&output of Dispatch, input of Combine, +// input&output of Combine bwd, input of Dispatch bwd. +// ============================================================================= + +namespace { + +// Caller-owned ncclMemAlloc'd buffer with a registered symmetric window. +// Frees in destructor (deregister + ncclMemFree). Non-copyable, move-only. +struct SymmBuf { + void* ptr = nullptr; + size_t bytes = 0; + ncclWindow_t win = nullptr; + + SymmBuf() = default; + SymmBuf(const SymmBuf&) = delete; + SymmBuf& operator=(const SymmBuf&) = delete; + SymmBuf(SymmBuf&& o) noexcept : ptr(o.ptr), bytes(o.bytes), win(o.win) { + o.ptr = nullptr; o.win = nullptr; o.bytes = 0; + } + ~SymmBuf() { + if (win) ncclCommWindowDeregister(g_ep_comm, win); + if (ptr) ncclMemFree(ptr); + } + + void alloc(size_t n_bytes) { + bytes = n_bytes; + ASSERT_NCCL_OK(ncclMemAlloc(&ptr, bytes)); + CHECK_CUDA(cudaMemset(ptr, 0, bytes)); + ASSERT_NCCL_OK(ncclCommWindowRegister(g_ep_comm, ptr, bytes, &win, + NCCL_WIN_COLL_SYMMETRIC)); + } +}; + +// Build an NVTECommWindow descriptor pointing at a SymmBuf's window (offset 0). +static inline NVTECommWindow symm_window(const SymmBuf& b) { + return NVTECommWindow{b.win, /*offset=*/0}; +} + +} // namespace + +class EPZeroCopyTest : public EpOpTestBase {}; + +// Identity round-trip with symm-mem on dispatch i/o + combine input. Bit-exact +// vs HBM reference (same routing, same input). +TEST_F(EPZeroCopyTest, IdentityAllSymm) { + // HBM reference run. + EPBuffers ref_buf; + ref_buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, + ep_size_, max_tokens_per_rank_); + upload_inputs(ref_buf); + EPTensors ref_t(ref_buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + uint64_t ref_hid = ref_buf.handle_id; + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{ref_hid, ref_t.handle_mem.tensor}, ref_t.topk_idx.tensor, ref_t.token_counts.tensor, /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{ref_hid, ref_t.handle_mem.tensor}, ref_t.topk_idx.tensor, + ref_t.tokens.tensor, NVTECommWindow{}, ref_t.topk_weights.tensor, + NVTECommWindow{}, ref_t.recv_tokens.tensor, NVTECommWindow{}, + ref_t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{ref_hid, ref_t.handle_mem.tensor}, ref_t.recv_tokens.tensor, NVTECommWindow{}, + ref_t.result.tensor, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + std::vector ref_recv(ref_buf.recv_capacity * hidden_dim_); + std::vector ref_result(num_tokens_ * hidden_dim_); + CHECK_CUDA(cudaMemcpy(ref_recv.data(), ref_buf.recv_tokens.get(), + ref_recv.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + CHECK_CUDA(cudaMemcpy(ref_result.data(), ref_buf.result.get(), + ref_result.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + + // Symm-mem run: tokens, recv_tokens, combine_input (== recv_tokens) all symm. + EPBuffers sym_buf; // alloc all buffers except the symm ones. + sym_buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, + ep_size_, max_tokens_per_rank_); + upload_inputs(sym_buf); + + SymmBuf sym_tokens, sym_recv; + sym_tokens.alloc(num_tokens_ * hidden_dim_ * sizeof(nv_bfloat16)); + sym_recv .alloc(sym_buf.recv_capacity * hidden_dim_ * sizeof(nv_bfloat16)); + + // Stage same tokens into the symm-mem input. + auto h_tok = generate_tokens(g_process_id, num_tokens_, hidden_dim_); + CHECK_CUDA(cudaMemcpy(sym_tokens.ptr, h_tok.data(), + h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); + + EPTensors sym_t(sym_buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + // Replace the tokens/recv_tokens views with ones pointing at the symm buffers. + sym_t.tokens = make_nvte_tensor(sym_tokens.ptr, + {(size_t)num_tokens_, (size_t)hidden_dim_}, kNVTEBFloat16); + sym_t.recv_tokens = make_nvte_tensor(sym_recv.ptr, + {sym_buf.recv_capacity, (size_t)hidden_dim_}, kNVTEBFloat16); + + uint64_t sym_hid = sym_buf.handle_id; + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{sym_hid, sym_t.handle_mem.tensor}, sym_t.topk_idx.tensor, sym_t.token_counts.tensor, /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{sym_hid, sym_t.handle_mem.tensor}, sym_t.topk_idx.tensor, + sym_t.tokens.tensor, symm_window(sym_tokens), + sym_t.topk_weights.tensor, NVTECommWindow{}, + sym_t.recv_tokens.tensor, symm_window(sym_recv), + sym_t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{sym_hid, sym_t.handle_mem.tensor}, sym_t.recv_tokens.tensor, + symm_window(sym_recv), sym_t.result.tensor, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + std::vector sym_recv_host(sym_buf.recv_capacity * hidden_dim_); + std::vector sym_result(num_tokens_ * hidden_dim_); + CHECK_CUDA(cudaMemcpy(sym_recv_host.data(), sym_recv.ptr, + sym_recv_host.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + CHECK_CUDA(cudaMemcpy(sym_result.data(), sym_buf.result.get(), + sym_result.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + + // Compare per filled recv slot (HBM ref vs symm) and full result. + int total_recv = read_total_recv(sym_buf); + for (int i = 0; i < total_recv * hidden_dim_; ++i) + ASSERT_EQ(__bfloat162float(sym_recv_host[i]), __bfloat162float(ref_recv[i])) + << "recv mismatch at " << i; + for (size_t i = 0; i < sym_result.size(); ++i) + ASSERT_EQ(__bfloat162float(sym_result[i]), __bfloat162float(ref_result[i])) + << "result mismatch at " << i; + + if (g_process_id == 0) + printf(" IdentityAllSymm: passed (recv_slots=%d, bit-exact vs HBM)\n", total_recv); + + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// Same buffers, 2 iterations — catches window-lifecycle regressions where the +// symm-mem registration goes stale between calls. +TEST_F(EPZeroCopyTest, IdentityAllSymmRepeated) { + EPBuffers buf; + buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, + ep_size_, max_tokens_per_rank_); + upload_inputs(buf); + + SymmBuf sym_tokens, sym_recv; + sym_tokens.alloc(num_tokens_ * hidden_dim_ * sizeof(nv_bfloat16)); + sym_recv .alloc(buf.recv_capacity * hidden_dim_ * sizeof(nv_bfloat16)); + auto h_tok = generate_tokens(g_process_id, num_tokens_, hidden_dim_); + CHECK_CUDA(cudaMemcpy(sym_tokens.ptr, h_tok.data(), + h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); + + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + t.tokens = make_nvte_tensor(sym_tokens.ptr, + {(size_t)num_tokens_, (size_t)hidden_dim_}, kNVTEBFloat16); + t.recv_tokens = make_nvte_tensor(sym_recv.ptr, + {buf.recv_capacity, (size_t)hidden_dim_}, kNVTEBFloat16); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + uint64_t handle_id = buf.handle_id; + for (int iter = 0; iter < 2; ++iter) { + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, + t.tokens.tensor, symm_window(sym_tokens), + t.topk_weights.tensor, NVTECommWindow{}, + t.recv_tokens.tensor, symm_window(sym_recv), + t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.recv_tokens.tensor, + symm_window(sym_recv), t.result.tensor, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + std::vector h_res(num_tokens_ * hidden_dim_); + CHECK_CUDA(cudaMemcpy(h_res.data(), buf.result.get(), + h_res.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + for (int tok = 0; tok < num_tokens_; ++tok) { + float exp = __bfloat162float(h_tok[tok * hidden_dim_]) * static_cast(top_k_); + float got = __bfloat162float(h_res[tok * hidden_dim_]); + ASSERT_NEAR(got, exp, bf16_tol(exp)) << "iter " << iter << " tok " << tok; + } + } + + if (g_process_id == 0) + printf(" IdentityAllSymmRepeated: passed (2 iters)\n"); + + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// Full forward+backward with symm-mem on every spec-mandated buffer: +// dispatch i/o, combine input, combine_bwd i/o, dispatch_bwd input. +// TODO: flaky on rank 0 (grad_tokens partial-zero) when run after the prior +// EPZeroCopyTest cases in the same binary; passes in isolation. Re-enable once +// the root cause (likely NCCL EP NVLS write→read coherence on grad_expert) is +// understood. Tracked separately. +TEST_F(EPZeroCopyTest, DISABLED_FullPipelineSymm) { + EPBuffers buf; + buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, + ep_size_, max_tokens_per_rank_); + upload_inputs(buf); + + // Symm-mem: tokens (dispatch input), recv_tokens (dispatch output AND + // combine input), grad_result (combine_bwd input), grad_expert + // (combine_bwd output AND dispatch_bwd input). + SymmBuf sym_tokens, sym_recv, sym_grad_result, sym_grad_expert; + sym_tokens .alloc(num_tokens_ * hidden_dim_ * sizeof(nv_bfloat16)); + sym_recv .alloc(buf.recv_capacity * hidden_dim_ * sizeof(nv_bfloat16)); + sym_grad_result.alloc(num_tokens_ * hidden_dim_ * sizeof(nv_bfloat16)); + sym_grad_expert.alloc(buf.recv_capacity * hidden_dim_ * sizeof(nv_bfloat16)); + + auto h_tok = generate_tokens(g_process_id, num_tokens_, hidden_dim_); + CHECK_CUDA(cudaMemcpy(sym_tokens.ptr, h_tok.data(), + h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); + + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + t.tokens = make_nvte_tensor(sym_tokens.ptr, + {(size_t)num_tokens_, (size_t)hidden_dim_}, kNVTEBFloat16); + t.recv_tokens = make_nvte_tensor(sym_recv.ptr, + {buf.recv_capacity, (size_t)hidden_dim_}, kNVTEBFloat16); + t.grad_result = make_nvte_tensor(sym_grad_result.ptr, + {(size_t)num_tokens_, (size_t)hidden_dim_}, kNVTEBFloat16); + t.grad_expert = make_nvte_tensor(sym_grad_expert.ptr, + {buf.recv_capacity, (size_t)hidden_dim_}, kNVTEBFloat16); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + uint64_t handle_id = buf.handle_id; + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, + t.tokens.tensor, symm_window(sym_tokens), + t.topk_weights.tensor, NVTECommWindow{}, + t.recv_tokens.tensor, symm_window(sym_recv), + t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.recv_tokens.tensor, + symm_window(sym_recv), t.result.tensor, stream)); + + std::vector h_grad(num_tokens_ * hidden_dim_, __float2bfloat16(0.1f)); + CHECK_CUDA(cudaMemcpyAsync(sym_grad_result.ptr, h_grad.data(), + h_grad.size() * sizeof(nv_bfloat16), + cudaMemcpyHostToDevice, stream)); + CHECK_CUDA(cudaMemsetAsync(sym_grad_expert.ptr, 0, sym_grad_expert.bytes, stream)); + CHECK_CUDA(cudaMemsetAsync(buf.g_recv_topk_weights.get(), 0, buf.g_recv_topk_weights.bytes(), stream)); + CHECK_CUDA(cudaMemsetAsync(buf.grad_topk_weights.get(), 0, buf.grad_topk_weights.bytes(), stream)); + + ASSERT_NO_THROW(nvte_ep_combine_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_result.tensor, + symm_window(sym_grad_result), t.grad_expert.tensor, + symm_window(sym_grad_expert), stream)); + ASSERT_NO_THROW(nvte_ep_dispatch_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_expert.tensor, + symm_window(sym_grad_expert), + t.g_recv_topk_weights.tensor, NVTECommWindow{}, + t.grad_tokens.tensor, t.grad_topk_weights.tensor, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + ASSERT_TRUE(check_no_nan_inf(buf.result.get(), num_tokens_ * hidden_dim_, "result")); + ASSERT_TRUE(check_no_nan_inf(buf.grad_tokens.get(), num_tokens_ * hidden_dim_, "grad_tokens")); + + std::vector h_gt(num_tokens_ * hidden_dim_); + CHECK_CUDA(cudaMemcpy(h_gt.data(), buf.grad_tokens.get(), + h_gt.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + const float kExpGrad = static_cast(top_k_) * 0.1f; + for (int tok = 0; tok < num_tokens_; ++tok) + EXPECT_NEAR(__bfloat162float(h_gt[tok * hidden_dim_]), kExpGrad, bf16_tol(kExpGrad)) + << "grad_tokens token " << tok; + + if (g_process_id == 0) printf(" FullPipelineSymm: passed\n"); + + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +// ── main ────────────────────────────────────────────────────────────────────── + +int main(int argc, char* argv[]) { + if (!ep_bootstrap(argc, argv)) return 0; + int ret = RUN_ALL_TESTS(); + ep_teardown(); + return ret; +} diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 8f96432ed8..18c4af7b09 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -437,6 +437,96 @@ if (NVTE_WITH_CUSOLVERMP) message(STATUS "Using cuSolverMp at: ${CUSOLVERMP_DIR}") endif() +# ── NCCL EP (on by default, HT mode only) ───────────────────────────────── +# Set -DNVTE_WITH_NCCL_EP=OFF (or NVTE_BUILD_WITH_NCCL_EP=0 in setup.py) to +# skip NCCL EP entirely — useful on older images whose system NCCL is below +# the 2.30.4 EP minimum. +option(NVTE_WITH_NCCL_EP "Build NCCL EP into libtransformer_engine.so" ON) +if(NVTE_WITH_NCCL_EP) +# SM>=90 and NCCL>=2.30.4 are gated at runtime in EPBackend::initialize. +# ── NCCL EP headers ──────────────────────────────────────────────────────── +# Headers + libs are produced by the in-tree 3rdparty/nccl submodule build +# (auto-built by setup.py via build_nccl_ep_submodule). +set(NCCL_EP_SUBMODULE_ROOT + "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl") +set(NCCL_EP_INCLUDE_DIR "${NCCL_EP_SUBMODULE_ROOT}/contrib/nccl_ep/include") +if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h") + message(FATAL_ERROR + "NCCL EP header not found at ${NCCL_EP_INCLUDE_DIR}/nccl_ep.h. " + "Run `git submodule update --init --recursive` to checkout 3rdparty/nccl.") +endif() +message(STATUS "NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}") + +# ── libnccl_ep.so ────────────────────────────────────────────────────────── +set(NCCL_EP_LIB_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/lib") +find_library(NCCL_EP_LIB + NAMES nccl_ep libnccl_ep + HINTS ${NCCL_EP_LIB_DIR} + NO_DEFAULT_PATH + REQUIRED) + +# ── NCCL + GIN headers ───────────────────────────────────────────────────── +# libnccl.so and all GIN headers (ncclGin.h, ncclWindow_t, ncclDevComm_t) +# ship with the base CUDA Toolkit OR the 3rdparty/nccl submodule build +# (preferred when present; auto-built by setup.py via build_nccl_ep_submodule). +if(NOT NCCL_LIB) + find_library(NCCL_LIB + NAMES nccl libnccl + HINTS ${NCCL_EP_LIB_DIR} ${CUDAToolkit_LIBRARY_DIR} + PATH_SUFFIXES lib lib64 + REQUIRED) +endif() + +set(NCCL_SUBMODULE_INCLUDE + "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl/build/include") +if(EXISTS "${NCCL_SUBMODULE_INCLUDE}/nccl.h") + set(NCCL_INCLUDE_DIRS_FOR_TE ${NCCL_SUBMODULE_INCLUDE}) +else() + set(NCCL_INCLUDE_DIRS_FOR_TE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) +endif() + +# Diagnostic: log detected NCCL header version (minimum enforced at runtime). +find_file(_nvte_nccl_header_path nccl.h + PATHS ${NCCL_INCLUDE_DIRS_FOR_TE} ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} + NO_DEFAULT_PATH) +if(_nvte_nccl_header_path) + file(READ "${_nvte_nccl_header_path}" _nvte_nccl_h) + string(REGEX MATCH "#define[ \t]+NCCL_MAJOR[ \t]+([0-9]+)" _ "${_nvte_nccl_h}") + set(_nvte_nccl_major "${CMAKE_MATCH_1}") + string(REGEX MATCH "#define[ \t]+NCCL_MINOR[ \t]+([0-9]+)" _ "${_nvte_nccl_h}") + set(_nvte_nccl_minor "${CMAKE_MATCH_1}") + string(REGEX MATCH "#define[ \t]+NCCL_PATCH[ \t]+([0-9]+)" _ "${_nvte_nccl_h}") + set(_nvte_nccl_patch "${CMAKE_MATCH_1}") + if(_nvte_nccl_major AND _nvte_nccl_minor AND _nvte_nccl_patch) + message(STATUS "NCCL header: ${_nvte_nccl_header_path} (version ${_nvte_nccl_major}.${_nvte_nccl_minor}.${_nvte_nccl_patch})") + endif() +endif() + +target_include_directories(transformer_engine PRIVATE + ${NCCL_EP_INCLUDE_DIR} + ${NCCL_INCLUDE_DIRS_FOR_TE}) # covers nccl.h + nccl_device/ + +target_link_libraries(transformer_engine PUBLIC + ${NCCL_EP_LIB} + ${NCCL_LIB}) + +# Embed rpath so the installed wheel finds libnccl_ep.so at runtime. +# libnccl.so is already on the system via the Toolkit — no rpath needed for it. +set_target_properties(transformer_engine PROPERTIES + INSTALL_RPATH "$ORIGIN;${NCCL_EP_LIB_DIR}") + +target_sources(transformer_engine PRIVATE + ep/ep_backend.cpp + ep/ep_api.cpp) + +message(STATUS "NCCL EP enabled: ${NCCL_EP_LIB}") +message(STATUS "NCCL EP include: ${NCCL_EP_INCLUDE_DIR}") +else() + # NCCL EP off: export throwing nvte_ep_* stubs so framework bindings link. + target_sources(transformer_engine PRIVATE ep/ep_api_stub.cpp) + message(STATUS "NCCL EP disabled (NVTE_WITH_NCCL_EP=OFF) — using nvte_ep_* stubs") +endif() + # Number of philox4x32 rounds for stochastic rounding (build-time constant). set(NVTE_BUILD_NUM_PHILOX_ROUNDS_STR $ENV{NVTE_BUILD_NUM_PHILOX_ROUNDS}) if (NOT NVTE_BUILD_NUM_PHILOX_ROUNDS_STR) diff --git a/transformer_engine/common/ep/ep_api.cpp b/transformer_engine/common/ep/ep_api.cpp new file mode 100644 index 0000000000..89d8b38607 --- /dev/null +++ b/transformer_engine/common/ep/ep_api.cpp @@ -0,0 +1,76 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file ep_api.cpp + * \brief nvte_ep_* C API: thin delegations to the EPBackend singleton. + */ + +#include +#include + +#include "../common.h" +#include "../util/logging.h" +#include "ep_backend.h" + +using transformer_engine::ep::EPBackend; + +void nvte_ep_initialize(void* ep_comm, NVTEEpGroupConfig group_config) { + NVTE_CHECK(ep_comm != nullptr, "ep_comm must not be null"); + EPBackend::initialize(static_cast(ep_comm), group_config); +} + +void nvte_ep_shutdown(void) { EPBackend::shutdown(); } + +uint64_t nvte_ep_register_layer(NVTEEpLayerConfig layer_config, size_t* handle_mem_size) { + NVTE_CHECK(handle_mem_size != nullptr, "handle_mem_size must not be null"); + return EPBackend::get().register_layer(layer_config, handle_mem_size); +} + +void nvte_ep_prepare(NVTEEpHandle handle, NVTETensor topk_idx, NVTETensor token_counts, + size_t dispatch_output_per_expert_alignment, cudaStream_t stream) { + void* mem_ptr = nvte_tensor_data(handle.mem); + NVTE_CHECK(mem_ptr != nullptr, "handle_mem tensor data must not be null"); + EPBackend::get().prepare(handle.id, topk_idx, token_counts, mem_ptr, + dispatch_output_per_expert_alignment, stream); +} + +void nvte_ep_dispatch(NVTEEpHandle handle, NVTETensor topk_idx, NVTETensor tokens, + NVTECommWindow tokens_win, NVTETensor topk_weights, + NVTECommWindow topk_weights_win, NVTETensor recv_tokens, + NVTECommWindow recv_tokens_win, NVTETensor recv_topk_weights, + NVTECommWindow recv_topk_weights_win, cudaStream_t stream) { + void* mem_ptr = nvte_tensor_data(handle.mem); + NVTE_CHECK(mem_ptr != nullptr, "handle_mem tensor data must not be null"); + EPBackend::get().dispatch(handle.id, mem_ptr, topk_idx, tokens, tokens_win, topk_weights, + topk_weights_win, recv_tokens, recv_tokens_win, recv_topk_weights, + recv_topk_weights_win, stream); +} + +void nvte_ep_combine(NVTEEpHandle handle, NVTETensor expert_out, NVTECommWindow expert_out_win, + NVTETensor result, cudaStream_t stream) { + void* mem_ptr = nvte_tensor_data(handle.mem); + NVTE_CHECK(mem_ptr != nullptr, "handle_mem tensor data must not be null"); + EPBackend::get().combine(handle.id, mem_ptr, expert_out, expert_out_win, result, stream); +} + +void nvte_ep_dispatch_bwd(NVTEEpHandle handle, NVTETensor grad, NVTECommWindow grad_win, + NVTETensor g_recv_topk_weights, NVTECommWindow g_recv_topk_weights_win, + NVTETensor grad_tokens, NVTETensor grad_topk_weights, + cudaStream_t stream) { + void* mem_ptr = nvte_tensor_data(handle.mem); + NVTE_CHECK(mem_ptr != nullptr, "handle_mem tensor data must not be null"); + EPBackend::get().dispatch_bwd(handle.id, mem_ptr, grad, grad_win, g_recv_topk_weights, + g_recv_topk_weights_win, grad_tokens, grad_topk_weights, stream); +} + +void nvte_ep_combine_bwd(NVTEEpHandle handle, NVTETensor grad, NVTECommWindow grad_win, + NVTETensor grad_expert_out, NVTECommWindow grad_expert_out_win, + cudaStream_t stream) { + void* mem_ptr = nvte_tensor_data(handle.mem); + NVTE_CHECK(mem_ptr != nullptr, "handle_mem tensor data must not be null"); + EPBackend::get().combine_bwd(handle.id, mem_ptr, grad, grad_win, grad_expert_out, + grad_expert_out_win, stream); +} diff --git a/transformer_engine/common/ep/ep_api_stub.cpp b/transformer_engine/common/ep/ep_api_stub.cpp new file mode 100644 index 0000000000..fe4127d87d --- /dev/null +++ b/transformer_engine/common/ep/ep_api_stub.cpp @@ -0,0 +1,61 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file ep_api_stub.cpp + * \brief Throwing nvte_ep_* stubs compiled when NVTE_WITH_NCCL_EP=OFF. + */ + +#include + +#include "../util/logging.h" + +namespace { +[[noreturn]] void ep_not_built() { + NVTE_ERROR( + "NCCL EP is not built into this TransformerEngine. Rebuild TE with " + "NVTE_BUILD_WITH_NCCL_EP=1 and CUDA arch >= 90 (e.g. NVTE_CUDA_ARCHS=\"90\")."); +} +} // namespace + +void nvte_ep_initialize(void* /*ep_comm*/, NVTEEpGroupConfig /*group_config*/) { ep_not_built(); } + +void nvte_ep_shutdown(void) {} + +uint64_t nvte_ep_register_layer(NVTEEpLayerConfig /*layer_config*/, size_t* /*handle_mem_size*/) { + ep_not_built(); +} + +void nvte_ep_prepare(NVTEEpHandle /*handle*/, NVTETensor /*topk_idx*/, NVTETensor /*token_counts*/, + size_t /*dispatch_output_per_expert_alignment*/, cudaStream_t /*stream*/) { + ep_not_built(); +} + +void nvte_ep_dispatch(NVTEEpHandle /*handle*/, NVTETensor /*topk_idx*/, NVTETensor /*tokens*/, + NVTECommWindow /*tokens_win*/, NVTETensor /*topk_weights*/, + NVTECommWindow /*topk_weights_win*/, NVTETensor /*recv_tokens*/, + NVTECommWindow /*recv_tokens_win*/, NVTETensor /*recv_topk_weights*/, + NVTECommWindow /*recv_topk_weights_win*/, cudaStream_t /*stream*/) { + ep_not_built(); +} + +void nvte_ep_combine(NVTEEpHandle /*handle*/, NVTETensor /*expert_out*/, + NVTECommWindow /*expert_out_win*/, NVTETensor /*result*/, + cudaStream_t /*stream*/) { + ep_not_built(); +} + +void nvte_ep_dispatch_bwd(NVTEEpHandle /*handle*/, NVTETensor /*grad*/, NVTECommWindow /*grad_win*/, + NVTETensor /*g_recv_topk_weights*/, + NVTECommWindow /*g_recv_topk_weights_win*/, NVTETensor /*grad_tokens*/, + NVTETensor /*grad_topk_weights*/, cudaStream_t /*stream*/) { + ep_not_built(); +} + +void nvte_ep_combine_bwd(NVTEEpHandle /*handle*/, NVTETensor /*grad*/, NVTECommWindow /*grad_win*/, + NVTETensor /*grad_expert_out*/, NVTECommWindow /*grad_expert_out_win*/, + cudaStream_t /*stream*/) { + ep_not_built(); +} diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp new file mode 100644 index 0000000000..ae0f3ab888 --- /dev/null +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -0,0 +1,514 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file ep_backend.cpp + * \brief EPBackend implementation. See ep_backend.h for the op flow. + */ + +#include "ep_backend.h" + +#include +#include +#include +#include +#include +#include + +#include "../common.h" +#include "../util/cuda_runtime.h" +#include "../util/logging.h" + +namespace transformer_engine { +namespace ep { + +namespace { + +// Build a by-value ncclEpTensor_t descriptor. `sizes` is caller-owned and must +// outlive any NCCL EP call that consumes the descriptor. +inline ncclEpTensor_t make_tensor(void* data, unsigned int ndim, ncclDataType_t datatype, + size_t* sizes) { + ncclEpTensor_t t = NCCL_EP_TENSOR_INIT; + t.ndim = ndim; + t.datatype = datatype; + t.data = data; + t.sizes = sizes; + return t; +} + +// Payload descriptor: prefer the symmem window when set, else fall back to the +// NVTETensor's raw device pointer. +inline ncclEpTensor_t make_payload_tensor(const NVTETensor t, const NVTECommWindow& win, + unsigned int ndim, ncclDataType_t datatype, + size_t* sizes) { + ncclEpTensor_t desc = NCCL_EP_TENSOR_INIT; + desc.ndim = ndim; + desc.datatype = datatype; + desc.sizes = sizes; + if (win.window != nullptr) { + desc.win_hdl = win.window; + desc.win_offset = win.offset; + } else { + desc.data = nvte_tensor_data(t); + NVTE_CHECK(desc.data != nullptr, "payload tensor data must not be null"); + } + return desc; +} + +// RAII guard for ncclEpHandle_t — destroys on scope exit, leak-free on throw. +class ScopedEpHandle { + public: + ScopedEpHandle() = default; + explicit ScopedEpHandle(ncclEpHandle_t h) : h_(h) {} + ~ScopedEpHandle() { + if (h_ != nullptr) ncclEpHandleDestroy(h_); + } + ScopedEpHandle(const ScopedEpHandle&) = delete; + ScopedEpHandle& operator=(const ScopedEpHandle&) = delete; + ScopedEpHandle(ScopedEpHandle&& other) noexcept : h_(other.h_) { other.h_ = nullptr; } + ScopedEpHandle& operator=(ScopedEpHandle&& other) noexcept { + if (this != &other) { + if (h_ != nullptr) ncclEpHandleDestroy(h_); + h_ = other.h_; + other.h_ = nullptr; + } + return *this; + } + operator ncclEpHandle_t() const { return h_; } + ncclEpHandle_t get() const { return h_; } + + private: + ncclEpHandle_t h_ = nullptr; +}; + +} // namespace + +// --------------------------------------------------------------------------- +// Singleton + bootstrap +// --------------------------------------------------------------------------- + +EPBackend& EPBackend::instance() { + static EPBackend inst; + return inst; +} + +EPBackend& EPBackend::get() { + EPBackend& inst = instance(); + NVTE_CHECK(inst.initialized_, "EPBackend not initialized. Call nvte_ep_initialize() first."); + return inst; +} + +void EPBackend::validate_config(const NVTEEpGroupConfig& config) { + NVTE_CHECK(config.ep_size > 0, "ep_size must be positive, got ", config.ep_size); + NVTE_CHECK(config.num_experts > 0, "num_experts must be positive, got ", config.num_experts); + NVTE_CHECK(config.max_tokens_per_rank > 0, "max_tokens_per_rank must be positive, got ", + config.max_tokens_per_rank); + NVTE_CHECK(config.max_recv_tokens_per_rank > 0, "max_recv_tokens_per_rank must be positive, got ", + config.max_recv_tokens_per_rank); + NVTE_CHECK(config.hidden_dim > 0, "hidden_dim must be positive, got ", config.hidden_dim); + NVTE_CHECK(config.hidden_dim * sizeof(nv_bfloat16) >= 16, + "hidden_dim * 2 must be >= 16 (NCCL EP 16B row alignment); got hidden_dim=", + config.hidden_dim); + NVTE_CHECK(config.num_experts % config.ep_size == 0, "num_experts (", config.num_experts, + ") must be divisible by ep_size (", config.ep_size, ")"); + NVTE_CHECK(config.max_num_sms >= 0, "max_num_sms must be >= 0 (0 = auto), got ", + config.max_num_sms); + + int device, major; + NVTE_CHECK_CUDA(cudaGetDevice(&device)); + NVTE_CHECK_CUDA(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device)); + NVTE_CHECK(major >= 9, + "NCCL EP requires SM_90+ (Hopper or later), " + "but current device has compute capability ", + major, ".x"); + + // NCCL EP needs CUDA multicast (NVLS); init hangs without it. + NVTE_CHECK(cuda::supports_multicast(device), + "NCCL EP requires CUDA multicast (NVLS) support on device ", device, + " but CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED reports 0."); +} + +void EPBackend::initialize(ncclComm_t ep_comm, NVTEEpGroupConfig config) { + EPBackend& inst = instance(); + std::lock_guard lock(inst.mutex_); + NVTE_CHECK(!inst.initialized_, "EP already initialized. Call initialize only once per process."); + NVTE_CHECK(ep_comm != nullptr, "ep_comm must not be null"); + + // Runtime gate: NCCL >= 2.30.4 (matches the submodule pin). + constexpr int kMinNcclVersion = 23004; + int nccl_version = 0; + NVTE_CHECK_NCCL(ncclGetVersion(&nccl_version)); + NVTE_CHECK(nccl_version >= kMinNcclVersion, "NCCL EP requires NCCL >= 2.30.4, found ", + nccl_version / 10000, ".", (nccl_version / 100) % 100, ".", nccl_version % 100, + " at runtime."); + + validate_config(config); + + int comm_size = 0; + NVTE_CHECK_NCCL(ncclCommCount(ep_comm, &comm_size)); + NVTE_CHECK(comm_size == config.ep_size, "ep_comm size (", comm_size, ") must equal ep_size (", + config.ep_size, "). Pass the EP sub-communicator, not the world comm."); + + inst.init(ep_comm, config); +} + +void EPBackend::shutdown() { + EPBackend& inst = instance(); + std::lock_guard lock(inst.mutex_); + if (!inst.initialized_) return; + inst.handles_.clear(); + // ncclEpGroupDestroy reads from ep_comm_; destroy group while comm is still alive. + if (inst.ep_group_ != nullptr) { + ncclEpGroupDestroy(inst.ep_group_); + inst.ep_group_ = nullptr; + } + inst.ep_comm_ = nullptr; // borrowed — caller destroys + inst.initialized_ = false; +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +ncclDataType_t EPBackend::nvte_dtype_to_nccl(NVTEDType dtype) { + switch (dtype) { + case kNVTEFloat32: + return ncclFloat32; + case kNVTEFloat16: + return ncclFloat16; + case kNVTEBFloat16: + return ncclBfloat16; + case kNVTEInt32: + return ncclInt32; + case kNVTEInt64: + return ncclInt64; + case kNVTEByte: + return ncclUint8; + case kNVTEFloat8E4M3: + return ncclFloat8e4m3; + case kNVTEFloat8E5M2: + return ncclFloat8e5m2; + default: + NVTE_ERROR("Unsupported NVTEDType for NCCL EP conversion: ", static_cast(dtype)); + } + return ncclFloat32; // unreachable +} + +// Open a transient ncclEpHandle over handle_mem. Caller owns the result. +ncclEpHandle_t EPBackend::open_handle(void* handle_mem, size_t handle_mem_size, int num_topk, + size_t dispatch_output_per_expert_alignment) { + size_t hm_sizes[1] = {handle_mem_size}; + ncclEpTensor_t routing_desc = make_tensor(handle_mem, 1, ncclUint8, hm_sizes); + ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; + hcfg.dispatch_output_per_expert_alignment = dispatch_output_per_expert_alignment; + ncclEpHandle_t handle; + NVTE_CHECK_NCCL(ncclEpInitHandle(&handle, ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, num_topk, + &routing_desc)); + return handle; +} + +// --------------------------------------------------------------------------- +// Lifecycle +// --------------------------------------------------------------------------- + +// Static-dtor teardown: skip NCCL calls (CUDA context / borrowed ep_comm_ may +// already be gone) and release in-memory state only. +EPBackend::~EPBackend() { + std::lock_guard lock(mutex_); + if (!initialized_) return; + handles_.clear(); + ep_group_ = nullptr; + ep_comm_ = nullptr; + initialized_ = false; +} + +void EPBackend::init(ncclComm_t ep_comm, NVTEEpGroupConfig group_config) { + NVTE_CHECK(!initialized_, "EPBackend already initialized"); + + group_config_ = group_config; + + ncclEpGroupConfig_t cfg = NCCL_EP_GROUP_CONFIG_INIT; + cfg.algorithm = NCCL_EP_ALGO_HIGH_THROUGHPUT; + cfg.num_experts = static_cast(group_config.num_experts); + cfg.max_dispatch_tokens_per_rank = static_cast(group_config.max_tokens_per_rank); + cfg.max_token_bytes = static_cast(group_config.hidden_dim * sizeof(nv_bfloat16)); + cfg.rdma_buffer_size = NCCL_EP_AUTO; + cfg.num_qp_per_rank = NCCL_EP_AUTO; + cfg.num_channels = NCCL_EP_AUTO; + cfg.max_num_sms = group_config.max_num_sms > 0 + ? static_cast(group_config.max_num_sms) + : NCCL_EP_AUTO; + // Must be > 0; NCCL EP errors out on 0. + cfg.max_recv_tokens_per_rank = static_cast(group_config.max_recv_tokens_per_rank); + + NVTE_CHECK_NCCL(ncclEpCreateGroup(&ep_group_, ep_comm, &cfg)); + + ep_comm_ = ep_comm; + + initialized_ = true; +} + +// --------------------------------------------------------------------------- +// Per-handle_id config cache +// --------------------------------------------------------------------------- + +uint64_t EPBackend::insert_new_entry(size_t handle_mem_size, int top_k, size_t alignment) { + if (handle_cache_cap_ == 0) { + const char* cap_env = std::getenv("NVTE_EP_HANDLE_CACHE_SIZE"); + handle_cache_cap_ = (cap_env != nullptr) ? std::max(1, std::atoi(cap_env)) : 8192; + } + NVTE_CHECK(handles_.size() < handle_cache_cap_, "EP handle cache full (", handle_cache_cap_, + " entries). Raise via NVTE_EP_HANDLE_CACHE_SIZE."); + uint64_t id = next_handle_id_.fetch_add(1, std::memory_order_relaxed); + handles_.emplace(id, HandleEntry{handle_mem_size, alignment, top_k}); + return id; +} + +EPBackend::HandleEntry& EPBackend::lookup_config(uint64_t handle_id) { + auto it = handles_.find(handle_id); + NVTE_CHECK(it != handles_.end(), "ep op on handle_id=", handle_id, + " with no cached config — call ep_prepare first."); + return it->second; +} + +// --------------------------------------------------------------------------- +// Per-step operations +// --------------------------------------------------------------------------- + +uint64_t EPBackend::register_layer(NVTEEpLayerConfig layer_config, size_t* handle_mem_size) { + NVTE_CHECK(initialized_, "EPBackend not initialized"); + NVTE_CHECK(layer_config.top_k > 0, "NVTEEpLayerConfig.top_k must be > 0"); + NVTE_CHECK(handle_mem_size != nullptr, "handle_mem_size must not be null"); + ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; + hcfg.dispatch_output_per_expert_alignment = layer_config.dispatch_output_per_expert_alignment; + size_t hm_size = 0; + NVTE_CHECK_NCCL(ncclEpHandleMemSize(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, &hm_size, + layer_config.top_k)); + *handle_mem_size = hm_size; + std::lock_guard lock(mutex_); + return insert_new_entry(hm_size, layer_config.top_k, + layer_config.dispatch_output_per_expert_alignment); +} + +void EPBackend::prepare(uint64_t handle_id, const NVTETensor topk_idx, NVTETensor token_counts, + void* handle_mem, size_t dispatch_output_per_expert_alignment, + cudaStream_t stream) { + NVTE_CHECK(initialized_, "EPBackend not initialized"); + NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); + + NVTEShape idx_shape = nvte_tensor_shape(topk_idx); + void* idx_data = nvte_tensor_data(topk_idx); + NVTE_CHECK(idx_data != nullptr, "topk_idx data must not be null"); + + const size_t num_tokens = idx_shape.data[0]; + const size_t top_k = idx_shape.ndim > 1 ? idx_shape.data[1] : 1; + const size_t num_local_experts = + static_cast(group_config_.num_experts / group_config_.ep_size); + + size_t idx_sizes[2] = {num_tokens, top_k}; + ncclEpTensor_t nccl_topk_idx = make_tensor(idx_data, 2, ncclInt64, idx_sizes); + + // ncclEpUpdateHandle writes per-expert counts via expert_counters. + size_t cnt_sizes[1] = {num_local_experts}; + ncclEpTensor_t token_counts_desc; + void* token_counts_data = (token_counts != nullptr) ? nvte_tensor_data(token_counts) : nullptr; + if (token_counts_data != nullptr) { + token_counts_desc = make_tensor(token_counts_data, 1, ncclInt32, cnt_sizes); + } + ncclEpLayoutInfo_t layout_info = NCCL_EP_LAYOUT_INFO_INIT; + layout_info.expert_counters = (token_counts_data != nullptr) ? &token_counts_desc : nullptr; + + ScopedEpHandle transient; + { + std::lock_guard lock(mutex_); + HandleEntry& cfg = lookup_config(handle_id); + NVTE_CHECK(cfg.alignment == dispatch_output_per_expert_alignment, + "ep_prepare: alignment mismatch for handle_id=", handle_id, + " (cached=", cfg.alignment, ", got=", dispatch_output_per_expert_alignment, ")"); + transient = + ScopedEpHandle(open_handle(handle_mem, cfg.handle_mem_size, cfg.top_k, cfg.alignment)); + } + NVTE_CHECK_NCCL(ncclEpUpdateHandle(transient, &nccl_topk_idx, &layout_info, stream)); +} + +void EPBackend::dispatch(uint64_t handle_id, void* handle_mem, const NVTETensor topk_idx, + const NVTETensor tokens, const NVTECommWindow& tokens_win, + const NVTETensor topk_weights, const NVTECommWindow& topk_weights_win, + NVTETensor recv_tokens, const NVTECommWindow& recv_tokens_win, + NVTETensor recv_topk_weights, const NVTECommWindow& recv_topk_weights_win, + cudaStream_t stream) { + NVTE_CHECK(initialized_, "EPBackend not initialized"); + NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); + + NVTEShape tok_shape = nvte_tensor_shape(tokens); + NVTEDType tok_dtype = nvte_tensor_type(tokens); + + const size_t num_tokens = tok_shape.data[0]; + const size_t hidden_dim = tok_shape.data[1]; + + size_t tok_sizes[2] = {num_tokens, hidden_dim}; + ncclEpTensor_t nccl_tokens_in = + make_payload_tensor(tokens, tokens_win, 2, nvte_dtype_to_nccl(tok_dtype), tok_sizes); + + const bool is_forward = (topk_weights != nullptr); + + // Routing is cached in handle_mem by ep_prepare; dispatch only needs + // topk_weights to reconstruct the sparse-to-dense prob map. + size_t weights_in_sizes[2] = {0, 0}; + ncclEpTensor_t nccl_topk_weights_in; + if (is_forward) { + NVTE_CHECK(topk_idx != nullptr, "topk_idx required in forward dispatch"); + NVTEShape idx_shape = nvte_tensor_shape(topk_idx); + const size_t top_k = idx_shape.ndim > 1 ? idx_shape.data[1] : 1; + weights_in_sizes[0] = num_tokens; + weights_in_sizes[1] = top_k; + nccl_topk_weights_in = + make_payload_tensor(topk_weights, topk_weights_win, 2, ncclFloat32, weights_in_sizes); + } + + NVTEShape recv_shape = nvte_tensor_shape(recv_tokens); + NVTEDType recv_dtype = nvte_tensor_type(recv_tokens); + + size_t recv_sizes[2] = {recv_shape.data[0], recv_shape.data[1]}; + ncclEpTensor_t nccl_tokens_out = make_payload_tensor(recv_tokens, recv_tokens_win, 2, + nvte_dtype_to_nccl(recv_dtype), recv_sizes); + + size_t weights_out_sizes[1] = {recv_shape.data[0]}; + ncclEpTensor_t nccl_topk_weights_out; + if (is_forward) { + NVTE_CHECK(recv_topk_weights != nullptr, + "recv_topk_weights must not be null in forward dispatch"); + NVTEShape recv_w_shape = nvte_tensor_shape(recv_topk_weights); + NVTE_CHECK(recv_w_shape.ndim == 1, "recv_topk_weights must be 1D [recv_capacity]"); + nccl_topk_weights_out = make_payload_tensor(recv_topk_weights, recv_topk_weights_win, 1, + ncclFloat32, weights_out_sizes); + } + + ncclEpDispatchInputs_t in_struct = NCCL_EP_DISPATCH_INPUTS_INIT; + in_struct.tokens = &nccl_tokens_in; + in_struct.topk_weights = is_forward ? &nccl_topk_weights_in : nullptr; + + ncclEpDispatchOutputs_t out_struct = NCCL_EP_DISPATCH_OUTPUTS_INIT; + out_struct.tokens = &nccl_tokens_out; + out_struct.topk_weights = is_forward ? &nccl_topk_weights_out : nullptr; + + ncclEpDispatchConfig_t dispatch_cfg = NCCL_EP_DISPATCH_CONFIG_INIT; + dispatch_cfg.pass_direction = is_forward ? NCCL_EP_FWD_PASS : NCCL_EP_BWD_PASS; + + ScopedEpHandle transient; + { + std::lock_guard lock(mutex_); + HandleEntry& cfg = lookup_config(handle_id); + transient = + ScopedEpHandle(open_handle(handle_mem, cfg.handle_mem_size, cfg.top_k, cfg.alignment)); + } + NVTE_CHECK_NCCL(ncclEpDispatch(transient, &in_struct, &out_struct, + /*layout_info=*/nullptr, &dispatch_cfg, stream)); +} + +void EPBackend::combine(uint64_t handle_id, void* handle_mem, const NVTETensor expert_out, + const NVTECommWindow& expert_out_win, NVTETensor result, + cudaStream_t stream) { + NVTE_CHECK(initialized_, "EPBackend not initialized"); + NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); + + NVTEShape exp_shape = nvte_tensor_shape(expert_out); + NVTEDType exp_dtype = nvte_tensor_type(expert_out); + + size_t exp_sizes[2] = {exp_shape.data[0], exp_shape.data[1]}; + ncclEpTensor_t nccl_expert_in = + make_payload_tensor(expert_out, expert_out_win, 2, nvte_dtype_to_nccl(exp_dtype), exp_sizes); + + NVTEShape res_shape = nvte_tensor_shape(result); + void* res_data = nvte_tensor_data(result); + NVTEDType res_dtype = nvte_tensor_type(result); + NVTE_CHECK(res_data != nullptr, "result data must not be null"); + + size_t res_sizes[2] = {res_shape.data[0], res_shape.data[1]}; + ncclEpTensor_t nccl_result_out = + make_tensor(res_data, 2, nvte_dtype_to_nccl(res_dtype), res_sizes); + + ncclEpCombineInputs_t in_struct = NCCL_EP_COMBINE_INPUTS_INIT; + in_struct.tokens = &nccl_expert_in; + + ncclEpCombineOutputs_t out_struct = NCCL_EP_COMBINE_OUTPUTS_INIT; + out_struct.tokens = &nccl_result_out; + + ScopedEpHandle transient; + { + std::lock_guard lock(mutex_); + HandleEntry& cfg = lookup_config(handle_id); + transient = + ScopedEpHandle(open_handle(handle_mem, cfg.handle_mem_size, cfg.top_k, cfg.alignment)); + } + NVTE_CHECK_NCCL(ncclEpCombine(transient, &in_struct, &out_struct, /*config=*/nullptr, stream)); +} + +void EPBackend::dispatch_bwd(uint64_t handle_id, void* handle_mem, const NVTETensor grad, + const NVTECommWindow& grad_win, const NVTETensor g_recv_topk_weights, + const NVTECommWindow& g_recv_topk_weights_win, NVTETensor grad_tokens, + NVTETensor grad_topk_weights, cudaStream_t stream) { + NVTE_CHECK(initialized_, "EPBackend not initialized"); + NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); + + NVTEShape g_shape = nvte_tensor_shape(grad); + NVTEDType g_dtype = nvte_tensor_type(grad); + size_t g_sizes[2] = {g_shape.data[0], g_shape.data[1]}; + ncclEpTensor_t nccl_tok_in = + make_payload_tensor(grad, grad_win, 2, nvte_dtype_to_nccl(g_dtype), g_sizes); + + // g_recv_topk_weights must be 1D [recv_capacity] — caller flattens. + NVTEShape gw_shape = nvte_tensor_shape(g_recv_topk_weights); + NVTE_CHECK(gw_shape.ndim == 1, + "g_recv_topk_weights must be 1D [recv_capacity]; caller must flatten leading dims"); + size_t gw_sizes[1] = {gw_shape.data[0]}; + ncclEpTensor_t nccl_w_in = + make_payload_tensor(g_recv_topk_weights, g_recv_topk_weights_win, 1, ncclFloat32, gw_sizes); + + NVTEShape gt_shape = nvte_tensor_shape(grad_tokens); + void* gt_data = nvte_tensor_data(grad_tokens); + NVTE_CHECK(gt_data != nullptr, "grad_tokens data must not be null"); + size_t gt_sizes[2] = {gt_shape.data[0], gt_shape.data[1]}; + ncclEpTensor_t nccl_tok_out = make_tensor(gt_data, 2, nvte_dtype_to_nccl(g_dtype), gt_sizes); + + NVTEShape gtw_shape = nvte_tensor_shape(grad_topk_weights); + void* gtw_data = nvte_tensor_data(grad_topk_weights); + NVTE_CHECK(gtw_data != nullptr, "grad_topk_weights data must not be null"); + NVTE_CHECK(gtw_shape.ndim == 2, "grad_topk_weights must be 2D [T, top_k]"); + size_t gtw_sizes[2] = {gtw_shape.data[0], gtw_shape.data[1]}; + ncclEpTensor_t nccl_w_out = make_tensor(gtw_data, 2, ncclFloat32, gtw_sizes); + + ncclEpCombineInputs_t in_struct = NCCL_EP_COMBINE_INPUTS_INIT; + in_struct.tokens = &nccl_tok_in; + in_struct.topk_weights = &nccl_w_in; + + ncclEpCombineOutputs_t out_struct = NCCL_EP_COMBINE_OUTPUTS_INIT; + out_struct.tokens = &nccl_tok_out; + out_struct.topk_weights = &nccl_w_out; + + ncclEpCombineConfig_t cfg = NCCL_EP_COMBINE_CONFIG_INIT; + cfg.pass_direction = NCCL_EP_BWD_PASS; + + ScopedEpHandle transient; + { + std::lock_guard lock(mutex_); + HandleEntry& entry = lookup_config(handle_id); + transient = ScopedEpHandle( + open_handle(handle_mem, entry.handle_mem_size, entry.top_k, entry.alignment)); + } + NVTE_CHECK_NCCL(ncclEpCombine(transient, &in_struct, &out_struct, &cfg, stream)); +} + +void EPBackend::combine_bwd(uint64_t handle_id, void* handle_mem, const NVTETensor grad, + const NVTECommWindow& grad_win, NVTETensor grad_expert_out, + const NVTECommWindow& grad_expert_out_win, cudaStream_t stream) { + // Backward of combine = reverse-direction dispatch. + dispatch(handle_id, handle_mem, /*topk_idx=*/nullptr, grad, grad_win, /*topk_weights=*/nullptr, + /*topk_weights_win=*/NVTECommWindow{}, grad_expert_out, grad_expert_out_win, + /*recv_topk_weights=*/nullptr, /*recv_topk_weights_win=*/NVTECommWindow{}, stream); +} + +} // namespace ep +} // namespace transformer_engine diff --git a/transformer_engine/common/ep/ep_backend.h b/transformer_engine/common/ep/ep_backend.h new file mode 100644 index 0000000000..18307ebb4f --- /dev/null +++ b/transformer_engine/common/ep/ep_backend.h @@ -0,0 +1,114 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file ep_backend.h + * \brief Internal NCCL EP singleton; not part of the public API. + * + * Per handle_id the cache stores config only (no device pointers), so + * handle_mem may be relocated between ops. Cap: NVTE_EP_HANDLE_CACHE_SIZE + * (default 8192); overflow throws. + */ + +#ifndef TRANSFORMER_ENGINE_COMMON_EP_EP_BACKEND_H_ +#define TRANSFORMER_ENGINE_COMMON_EP_EP_BACKEND_H_ + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +namespace transformer_engine { +namespace ep { + +/*! \brief EP backend singleton — owns the NCCL EP group; borrows the comm. */ +class EPBackend { + public: + /*! \brief Access the singleton. Aborts if not initialized. */ + static EPBackend& get(); + + /*! \brief Bootstrap from an existing EP sub-communicator. + * ep_comm is borrowed; the caller keeps it alive until shutdown() returns + * and must span exactly config.ep_size ranks. + */ + static void initialize(ncclComm_t ep_comm, NVTEEpGroupConfig config); + + /*! \brief Tear down the backend. Idempotent. Does not destroy ep_comm_. */ + static void shutdown(); + + // Host-only: reserve a fresh handle_id, cache the layer config, and report + // the handle_mem buffer size the caller must allocate. + uint64_t register_layer(NVTEEpLayerConfig layer_config, size_t* handle_mem_size); + + void prepare(uint64_t handle_id, const NVTETensor topk_idx, NVTETensor token_counts, + void* handle_mem, size_t dispatch_output_per_expert_alignment, cudaStream_t stream); + + void dispatch(uint64_t handle_id, void* handle_mem, const NVTETensor topk_idx, + const NVTETensor tokens, const NVTECommWindow& tokens_win, + const NVTETensor topk_weights, const NVTECommWindow& topk_weights_win, + NVTETensor recv_tokens, const NVTECommWindow& recv_tokens_win, + NVTETensor recv_topk_weights, const NVTECommWindow& recv_topk_weights_win, + cudaStream_t stream); + + void combine(uint64_t handle_id, void* handle_mem, const NVTETensor expert_out, + const NVTECommWindow& expert_out_win, NVTETensor result, cudaStream_t stream); + + // g_recv_topk_weights: 1D [recv_capacity] f32; grad_topk_weights: 2D [T, top_k] f32. + void dispatch_bwd(uint64_t handle_id, void* handle_mem, const NVTETensor grad, + const NVTECommWindow& grad_win, const NVTETensor g_recv_topk_weights, + const NVTECommWindow& g_recv_topk_weights_win, NVTETensor grad_tokens, + NVTETensor grad_topk_weights, cudaStream_t stream); + + void combine_bwd(uint64_t handle_id, void* handle_mem, const NVTETensor grad, + const NVTECommWindow& grad_win, NVTETensor grad_expert_out, + const NVTECommWindow& grad_expert_out_win, cudaStream_t stream); + + private: + EPBackend() = default; + ~EPBackend(); + EPBackend(const EPBackend&) = delete; + EPBackend& operator=(const EPBackend&) = delete; + + // ep_comm is borrowed — caller retains ownership across the backend lifetime. + void init(ncclComm_t ep_comm, NVTEEpGroupConfig config); + + static EPBackend& instance(); // Meyers singleton accessor + static void validate_config(const NVTEEpGroupConfig& config); + + static ncclDataType_t nvte_dtype_to_nccl(NVTEDType dtype); + // Open a transient ncclEpHandle over handle_mem. num_topk=-1 for paths + // that don't carry per-token weights. + ncclEpHandle_t open_handle(void* handle_mem, size_t handle_mem_size, int num_topk, + size_t dispatch_output_per_expert_alignment); + + ncclEpGroup_t ep_group_{nullptr}; + ncclComm_t ep_comm_{nullptr}; + NVTEEpGroupConfig group_config_{}; + bool initialized_{false}; + std::mutex mutex_; + struct HandleEntry { + size_t handle_mem_size; + size_t alignment; + int top_k; + }; + std::unordered_map handles_; + std::atomic next_handle_id_{1}; // 0 reserved as "no id" + size_t handle_cache_cap_{0}; // set lazily from NVTE_EP_HANDLE_CACHE_SIZE + + // Caller must hold mutex_. Throws on cap overflow. + uint64_t insert_new_entry(size_t handle_mem_size, int top_k, size_t alignment); + HandleEntry& lookup_config(uint64_t handle_id); +}; + +} // namespace ep +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_EP_EP_BACKEND_H_ diff --git a/transformer_engine/common/include/transformer_engine/comm_window.h b/transformer_engine/common/include/transformer_engine/comm_window.h new file mode 100644 index 0000000000..088ea7f0c3 --- /dev/null +++ b/transformer_engine/common/include/transformer_engine/comm_window.h @@ -0,0 +1,32 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file comm_window.h + * \brief Borrowed symmetric-memory window + offset for zero-copy one-sided ops. + * Pass ``{NULL, 0}`` to use the raw-pointer path. + */ + +#ifndef TRANSFORMER_ENGINE_COMM_WINDOW_H_ +#define TRANSFORMER_ENGINE_COMM_WINDOW_H_ + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/*! \brief NCCL window + byte offset for a zero-copy payload tensor. */ +typedef struct { + ncclWindow_t window; /*!< NCCL window, or NULL to use the raw data pointer. */ + uint64_t offset; /*!< Byte offset of the payload within ``window``. */ +} NVTECommWindow; + +#ifdef __cplusplus +} +#endif + +#endif // TRANSFORMER_ENGINE_COMM_WINDOW_H_ diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h new file mode 100644 index 0000000000..8c3a06b5f0 --- /dev/null +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -0,0 +1,161 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file ep.h + * \brief Public C API for Expert Parallelism. Per-step ops are allocation-free + * and CUDA graph-capturable. + */ + +#ifndef TRANSFORMER_ENGINE_EP_H_ +#define TRANSFORMER_ENGINE_EP_H_ + +#include +#include +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* ── Config structs ─────────────────────────────────────────────────────── */ + +/*! \brief Group-level EP configuration (fixed for the EP group lifetime). */ +typedef struct { + int ep_size; /*!< EP world size. */ + int num_experts; /*!< Total experts across all ranks. */ + int max_tokens_per_rank; /*!< Upper bound on tokens this rank sends per dispatch. */ + /*! Upper bound on tokens received per dispatch (worst-case top_k fan-out; must be > 0). */ + int max_recv_tokens_per_rank; + int hidden_dim; /*!< Token hidden dimension. */ + int max_num_sms; /*!< Max SMs for EP kernels. 0 = auto. */ + /*! 0 (default): throw on relocated handle_mem for a cached handle_id. 1: silently rebuild. */ + int allow_handle_mem_reloc; +} NVTEEpGroupConfig; + +/*! \brief Per-layer EP configuration. */ +typedef struct { + int num_local_experts; /*!< Reserved for ABI stability (derived from group config). */ + int top_k; /*!< Per-token expert fan-out. Required. */ + size_t dispatch_output_per_expert_alignment; + /*!< Per-expert zone alignment in tokens (pow2; 0/1 = no padding). Must match + * between nvte_ep_register_layer and nvte_ep_prepare. */ +} NVTEEpLayerConfig; + +/* ── Bootstrap ──────────────────────────────────────────────────────────── */ + +/*! \brief Bootstrap from an existing NCCL EP sub-communicator. Requires SM>=90. + * + * ep_comm is borrowed and must span exactly group_config.ep_size ranks. + * Re-init after shutdown is allowed; double-init throws. + * + * \param[in] ep_comm Opaque ncclComm_t for the EP sub-group. + * \param[in] group_config Group-level EP configuration. + */ +void nvte_ep_initialize(void* ep_comm, NVTEEpGroupConfig group_config); + +/*! \brief Tear down the EP backend. Idempotent. Does not destroy ep_comm. */ +void nvte_ep_shutdown(void); + +/* ── Layer registration (host-only, eager) ───────────────────────────────── */ + +/*! \brief Reserve a handle_id for a layer config and report the handle_mem buffer + * size the caller must allocate. Host-only. + * + * \param[in] layer_config Per-layer EP configuration. + * \param[out] handle_mem_size Bytes the caller must allocate for handle_mem. + * \return uint64_t handle_id (non-zero). + */ +uint64_t nvte_ep_register_layer(NVTEEpLayerConfig layer_config, size_t* handle_mem_size); + +/*! \brief Per-step handle: the registered handle_id paired with its handle_mem buffer. */ +typedef struct { + uint64_t id; /*!< Handle id from nvte_ep_register_layer. */ + NVTETensor mem; /*!< Caller-allocated handle_mem buffer (size from nvte_ep_register_layer). */ +} NVTEEpHandle; + +/* ── Per-step ops (all allocation-free, CUDA graph-capturable) ──────────── */ + +/*! \brief AllGather the routing map; write per-expert counts and cache routing + * metadata in handle.mem for the subsequent dispatch/combine. + * + * \param[in] handle EP handle (id + mem buffer). + * \param[in] topk_idx [T, top_k] int64 routing indices. + * \param[out] token_counts [num_local_experts] int32 counts. + * \param[in] dispatch_output_per_expert_alignment Must match the handle_mem sizing. + * \param[in] stream CUDA stream. + */ +void nvte_ep_prepare(NVTEEpHandle handle, NVTETensor topk_idx, NVTETensor token_counts, + size_t dispatch_output_per_expert_alignment, cudaStream_t stream); + +/*! \brief Dispatch tokens (and routing weights) to expert ranks. + * + * \param[in] handle EP handle (id + mem buffer). + * \param[in] topk_idx [T, top_k] int64 sparse routing indices. + * \param[in] tokens [T, hidden_dim] input tokens. + * \param[in] tokens_win Optional symmem window for ``tokens``. + * \param[in] topk_weights [T, top_k] float32 weights, or null in backward. + * \param[in] topk_weights_win Optional symmem window for ``topk_weights``. + * \param[out] recv_tokens [recv_T, hidden_dim] received tokens. + * \param[in] recv_tokens_win Optional symmem window for ``recv_tokens``. + * \param[out] recv_topk_weights [recv_T] float32 per-slot weights, or null in backward. + * \param[in] recv_topk_weights_win Optional symmem window for ``recv_topk_weights``. + * \param[in] stream CUDA stream. + */ +void nvte_ep_dispatch(NVTEEpHandle handle, NVTETensor topk_idx, NVTETensor tokens, + NVTECommWindow tokens_win, NVTETensor topk_weights, + NVTECommWindow topk_weights_win, NVTETensor recv_tokens, + NVTECommWindow recv_tokens_win, NVTETensor recv_topk_weights, + NVTECommWindow recv_topk_weights_win, cudaStream_t stream); + +/*! \brief Scatter-sum expert outputs back to originating ranks. Unweighted — + * caller must pre-multiply expert_out by recv_topk_weights (and the + * valid-slot mask) before calling. + * + * \param[in] handle EP handle (id + mem buffer). + * \param[in] expert_out [recv_T, hidden_dim] pre-weighted expert outputs. + * \param[in] expert_out_win Optional symmem window for ``expert_out``. + * \param[out] result [T, hidden_dim] combined output. + * \param[in] stream CUDA stream. + */ +void nvte_ep_combine(NVTEEpHandle handle, NVTETensor expert_out, NVTECommWindow expert_out_win, + NVTETensor result, cudaStream_t stream); + +/*! \brief Backward of dispatch — routes token and weight grads back to source. + * + * \param[in] handle EP handle (id + mem buffer). + * \param[in] grad [recv_capacity, hidden_dim] grad w.r.t. recv_tokens. + * \param[in] grad_win Optional symmem window for ``grad``. + * \param[in] g_recv_topk_weights [recv_capacity] f32 grad w.r.t. recv_topk_weights. + * \param[in] g_recv_topk_weights_win Optional symmem window for ``g_recv_topk_weights``. + * \param[out] grad_tokens [T, hidden_dim] grad w.r.t. tokens. + * \param[out] grad_topk_weights [T, top_k] f32 grad w.r.t. topk_weights. + * \param[in] stream CUDA stream. + */ +void nvte_ep_dispatch_bwd(NVTEEpHandle handle, NVTETensor grad, NVTECommWindow grad_win, + NVTETensor g_recv_topk_weights, NVTECommWindow g_recv_topk_weights_win, + NVTETensor grad_tokens, NVTETensor grad_topk_weights, + cudaStream_t stream); + +/*! \brief Backward of combine. Padded slots in grad_expert_out are zeroed. + * + * \param[in] handle EP handle (id + mem buffer). + * \param[in] grad [T, hidden_dim] grad w.r.t. result. + * \param[in] grad_win Optional symmem window for ``grad``. + * \param[out] grad_expert_out [recv_capacity, hidden_dim] grad w.r.t. expert_out. + * \param[in] grad_expert_out_win Optional symmem window for ``grad_expert_out``. + * \param[in] stream CUDA stream. + */ +void nvte_ep_combine_bwd(NVTEEpHandle handle, NVTETensor grad, NVTECommWindow grad_win, + NVTETensor grad_expert_out, NVTECommWindow grad_expert_out_win, + cudaStream_t stream); + +#ifdef __cplusplus +} +#endif + +#endif // TRANSFORMER_ENGINE_EP_H_ From 3a1f9134cb48f91192e79ac0994ec667026a824f Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Sat, 23 May 2026 19:36:55 +0000 Subject: [PATCH 02/55] Expert Parallelism: persistent ncclEpHandle cache with allow_handle_mem_reloc gating Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/test_ep_coverage.cu | 183 ++++++++++++++++++++ transformer_engine/common/ep/ep_backend.cpp | 109 +++++------- transformer_engine/common/ep/ep_backend.h | 8 + 3 files changed, 238 insertions(+), 62 deletions(-) diff --git a/tests/cpp_distributed/test_ep_coverage.cu b/tests/cpp_distributed/test_ep_coverage.cu index ef7941905d..e9e532386c 100644 --- a/tests/cpp_distributed/test_ep_coverage.cu +++ b/tests/cpp_distributed/test_ep_coverage.cu @@ -369,6 +369,189 @@ TEST_F(NegativeTests, NullHandleMemThrows) { CHECK_CUDA(cudaStreamDestroy(stream)); } +// ============================================================================= +// HandleCacheTest: persistent ncclEpHandle is reused across ops on the same +// handle_mem ptr; relocation triggers throw by default and rebuild when +// NVTEEpGroupConfig.allow_handle_mem_reloc=1. +// ============================================================================= + +class HandleCacheTest : public EpCoverageBase {}; + +// Run prepare → dispatch → combine on bundle b. handle_mem_data overrides the +// device ptr used for handle_mem (must be the buffer owned by b unless +// reloc-allowed mode is active). Templated on Bundle because EpCoverageBase:: +// Bundle is declared in a protected section. +template +static void run_round_trip(B& b, void* handle_mem_data, + int num_tokens, int top_k, int num_local_experts, + int hidden_dim, size_t alignment, + cudaStream_t stream) { + auto topk_idx_t = make_nvte_tensor(b.topk_idx.get(), + {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); + auto topk_weights_t = make_nvte_tensor(b.topk_weights.get(), + {(size_t)num_tokens, (size_t)top_k}, kNVTEFloat32); + auto token_counts_t = make_nvte_tensor(b.token_counts.get(), + {(size_t)num_local_experts}, kNVTEInt32); + auto handle_mem_t = make_nvte_tensor(handle_mem_data, + {b.handle_mem_size}, kNVTEByte); + auto tokens_t = make_nvte_tensor(b.tokens.get(), + {(size_t)num_tokens, (size_t)hidden_dim}, kNVTEBFloat16); + auto recv_tokens_t = make_nvte_tensor(b.recv_tokens.get(), + {b.recv_capacity, (size_t)hidden_dim}, kNVTEBFloat16); + auto recv_w_t = make_nvte_tensor(b.recv_topk_weights.get(), + {b.recv_capacity}, kNVTEFloat32); + auto result_t = make_nvte_tensor(b.result.get(), + {(size_t)num_tokens, (size_t)hidden_dim}, kNVTEBFloat16); + + NVTEEpHandle h{b.handle_id, handle_mem_t.tensor}; + nvte_ep_prepare(h, topk_idx_t.tensor, token_counts_t.tensor, alignment, stream); + nvte_ep_dispatch(h, topk_idx_t.tensor, tokens_t.tensor, NVTECommWindow{}, + topk_weights_t.tensor, NVTECommWindow{}, + recv_tokens_t.tensor, NVTECommWindow{}, + recv_w_t.tensor, NVTECommWindow{}, stream); + nvte_ep_combine(h, recv_tokens_t.tensor, NVTECommWindow{}, result_t.tensor, stream); +} + +// Re-bootstrap EP backend with a different allow_handle_mem_reloc setting. +// Reuses the existing g_ep_comm; caller is responsible for restoring defaults. +static void reinit_ep_with_reloc(int allow_reloc) { + nvte_ep_shutdown(); + NVTEEpGroupConfig cfg{}; + cfg.ep_size = g_ep_size; + cfg.num_experts = g_num_experts; + cfg.max_tokens_per_rank = g_max_tokens_per_rank; + cfg.max_recv_tokens_per_rank = g_ep_size * g_max_tokens_per_rank * 2; + cfg.hidden_dim = g_hidden_dim; + cfg.allow_handle_mem_reloc = allow_reloc; + nvte_ep_initialize(static_cast(g_ep_comm), cfg); +} + +TEST_F(HandleCacheTest, ReuseSameMemSucceeds) { + const int num_tokens = 16, top_k = 2; + Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); + + auto h_idx = routing_balanced(g_process_id, num_tokens, top_k, + num_experts_, num_local_experts_); + std::vector h_w(num_tokens * top_k, 1.0f / top_k); + auto h_tok = tokens_constant(num_tokens, hidden_dim_, 0.5f); + CHECK_CUDA(cudaMemcpy(b.topk_idx.get(), h_idx.data(), + h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(b.topk_weights.get(), h_w.data(), + h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(b.tokens.get(), h_tok.data(), + h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + // Two consecutive round-trips on the same handle_mem ptr: first opens the + // cached handle, second hits the cache. Both must succeed and be correct. + for (int iter = 0; iter < 2; ++iter) { + ASSERT_NO_THROW(run_round_trip(b, b.handle_mem.get(), num_tokens, top_k, + num_local_experts_, hidden_dim_, + /*alignment=*/0, stream)); + } + CHECK_CUDA(cudaStreamSynchronize(stream)); + + std::vector h_res(num_tokens * hidden_dim_); + CHECK_CUDA(cudaMemcpy(h_res.data(), b.result.get(), + h_res.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + const int probes[3] = {0, hidden_dim_ / 2, hidden_dim_ - 1}; + for (int t = 0; t < num_tokens; ++t) + for (int p : probes) + EXPECT_NEAR(__bfloat162float(h_res[t * hidden_dim_ + p]), + static_cast(top_k) * 0.5f, 1e-2f); + + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +TEST_F(HandleCacheTest, RelocDefaultThrows) { + // Default bootstrap has allow_handle_mem_reloc=0: a second prepare call on + // the same handle_id with a different handle_mem ptr must throw. + const int num_tokens = 8, top_k = 2; + Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); + DevBuf second_hm(b.handle_mem_size); // distinct device buffer + ASSERT_NE(b.handle_mem.get(), second_hm.get()); + + auto h_idx = routing_balanced(g_process_id, num_tokens, top_k, + num_experts_, num_local_experts_); + CHECK_CUDA(cudaMemcpy(b.topk_idx.get(), h_idx.data(), + h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); + + auto topk_idx_t = make_nvte_tensor(b.topk_idx.get(), + {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); + auto token_counts_t = make_nvte_tensor(b.token_counts.get(), + {(size_t)num_local_experts_}, kNVTEInt32); + auto hm1_t = make_nvte_tensor(b.handle_mem.get(), + {b.handle_mem_size}, kNVTEByte); + auto hm2_t = make_nvte_tensor(second_hm.get(), + {b.handle_mem_size}, kNVTEByte); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + // First prepare seeds the cache. + NVTEEpHandle h1{b.handle_id, hm1_t.tensor}; + ASSERT_NO_THROW(nvte_ep_prepare(h1, topk_idx_t.tensor, token_counts_t.tensor, + /*alignment=*/0, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + // Same handle_id with a different handle_mem ptr must throw. + NVTEEpHandle h2{b.handle_id, hm2_t.tensor}; + EXPECT_THROW(nvte_ep_prepare(h2, topk_idx_t.tensor, token_counts_t.tensor, + /*alignment=*/0, stream), + std::exception); + CHECK_CUDA(cudaStreamDestroy(stream)); +} + +TEST_F(HandleCacheTest, RelocAllowedRebuilds) { + // Re-init EP backend with allow_handle_mem_reloc=1, run two round-trips with + // distinct handle_mem buffers, verify both succeed numerically, restore. + reinit_ep_with_reloc(/*allow_reloc=*/1); + + struct Restore { ~Restore() { reinit_ep_with_reloc(/*allow_reloc=*/0); } } restore; + + const int num_tokens = 16, top_k = 2; + Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); + DevBuf alt_hm(b.handle_mem_size); + ASSERT_NE(b.handle_mem.get(), alt_hm.get()); + + auto h_idx = routing_balanced(g_process_id, num_tokens, top_k, + num_experts_, num_local_experts_); + std::vector h_w(num_tokens * top_k, 1.0f / top_k); + auto h_tok = tokens_constant(num_tokens, hidden_dim_, 0.5f); + CHECK_CUDA(cudaMemcpy(b.topk_idx.get(), h_idx.data(), + h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(b.topk_weights.get(), h_w.data(), + h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); + CHECK_CUDA(cudaMemcpy(b.tokens.get(), h_tok.data(), + h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); + + cudaStream_t stream; + CHECK_CUDA(cudaStreamCreate(&stream)); + + // First on the original handle_mem. + ASSERT_NO_THROW(run_round_trip(b, b.handle_mem.get(), num_tokens, top_k, + num_local_experts_, hidden_dim_, + /*alignment=*/0, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + // Then on the relocated handle_mem — must trigger silent rebuild, not throw. + ASSERT_NO_THROW(run_round_trip(b, alt_hm.get(), num_tokens, top_k, + num_local_experts_, hidden_dim_, + /*alignment=*/0, stream)); + CHECK_CUDA(cudaStreamSynchronize(stream)); + + std::vector h_res(num_tokens * hidden_dim_); + CHECK_CUDA(cudaMemcpy(h_res.data(), b.result.get(), + h_res.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + const int probes[3] = {0, hidden_dim_ / 2, hidden_dim_ - 1}; + for (int t = 0; t < num_tokens; ++t) + for (int p : probes) + EXPECT_NEAR(__bfloat162float(h_res[t * hidden_dim_ + p]), + static_cast(top_k) * 0.5f, 1e-2f); + + CHECK_CUDA(cudaStreamDestroy(stream)); +} + // ── main ────────────────────────────────────────────────────────────────────── int main(int argc, char* argv[]) { diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index ae0f3ab888..6494a86817 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -57,32 +57,6 @@ inline ncclEpTensor_t make_payload_tensor(const NVTETensor t, const NVTECommWind return desc; } -// RAII guard for ncclEpHandle_t — destroys on scope exit, leak-free on throw. -class ScopedEpHandle { - public: - ScopedEpHandle() = default; - explicit ScopedEpHandle(ncclEpHandle_t h) : h_(h) {} - ~ScopedEpHandle() { - if (h_ != nullptr) ncclEpHandleDestroy(h_); - } - ScopedEpHandle(const ScopedEpHandle&) = delete; - ScopedEpHandle& operator=(const ScopedEpHandle&) = delete; - ScopedEpHandle(ScopedEpHandle&& other) noexcept : h_(other.h_) { other.h_ = nullptr; } - ScopedEpHandle& operator=(ScopedEpHandle&& other) noexcept { - if (this != &other) { - if (h_ != nullptr) ncclEpHandleDestroy(h_); - h_ = other.h_; - other.h_ = nullptr; - } - return *this; - } - operator ncclEpHandle_t() const { return h_; } - ncclEpHandle_t get() const { return h_; } - - private: - ncclEpHandle_t h_ = nullptr; -}; - } // namespace // --------------------------------------------------------------------------- @@ -158,6 +132,13 @@ void EPBackend::shutdown() { EPBackend& inst = instance(); std::lock_guard lock(inst.mutex_); if (!inst.initialized_) return; + for (auto& kv : inst.handles_) { + if (kv.second.cached_handle != nullptr) { + ncclEpHandleDestroy(kv.second.cached_handle); + kv.second.cached_handle = nullptr; + kv.second.cached_handle_mem = nullptr; + } + } inst.handles_.clear(); // ncclEpGroupDestroy reads from ep_comm_; destroy group while comm is still alive. if (inst.ep_group_ != nullptr) { @@ -196,7 +177,7 @@ ncclDataType_t EPBackend::nvte_dtype_to_nccl(NVTEDType dtype) { return ncclFloat32; // unreachable } -// Open a transient ncclEpHandle over handle_mem. Caller owns the result. +// Open a fresh ncclEpHandle over handle_mem. Caller (or cache) owns the result. ncclEpHandle_t EPBackend::open_handle(void* handle_mem, size_t handle_mem_size, int num_topk, size_t dispatch_output_per_expert_alignment) { size_t hm_sizes[1] = {handle_mem_size}; @@ -273,6 +254,26 @@ EPBackend::HandleEntry& EPBackend::lookup_config(uint64_t handle_id) { return it->second; } +ncclEpHandle_t EPBackend::get_or_open_handle(HandleEntry& cfg, void* handle_mem) { + if (cfg.cached_handle != nullptr && cfg.cached_handle_mem == handle_mem) { + return cfg.cached_handle; + } + if (cfg.cached_handle != nullptr) { + NVTE_CHECK(group_config_.allow_handle_mem_reloc != 0, + "EP handle_mem relocated for cached handle (old=", + reinterpret_cast(cfg.cached_handle_mem), + ", new=", reinterpret_cast(handle_mem), + "). Set NVTEEpGroupConfig.allow_handle_mem_reloc=1 to allow rebuild."); + ncclEpHandleDestroy(cfg.cached_handle); + cfg.cached_handle = nullptr; + cfg.cached_handle_mem = nullptr; + } + ncclEpHandle_t h = open_handle(handle_mem, cfg.handle_mem_size, cfg.top_k, cfg.alignment); + cfg.cached_handle = h; + cfg.cached_handle_mem = handle_mem; + return h; +} + // --------------------------------------------------------------------------- // Per-step operations // --------------------------------------------------------------------------- @@ -320,17 +321,13 @@ void EPBackend::prepare(uint64_t handle_id, const NVTETensor topk_idx, NVTETenso ncclEpLayoutInfo_t layout_info = NCCL_EP_LAYOUT_INFO_INIT; layout_info.expert_counters = (token_counts_data != nullptr) ? &token_counts_desc : nullptr; - ScopedEpHandle transient; - { - std::lock_guard lock(mutex_); - HandleEntry& cfg = lookup_config(handle_id); - NVTE_CHECK(cfg.alignment == dispatch_output_per_expert_alignment, - "ep_prepare: alignment mismatch for handle_id=", handle_id, - " (cached=", cfg.alignment, ", got=", dispatch_output_per_expert_alignment, ")"); - transient = - ScopedEpHandle(open_handle(handle_mem, cfg.handle_mem_size, cfg.top_k, cfg.alignment)); - } - NVTE_CHECK_NCCL(ncclEpUpdateHandle(transient, &nccl_topk_idx, &layout_info, stream)); + std::lock_guard lock(mutex_); + HandleEntry& cfg = lookup_config(handle_id); + NVTE_CHECK(cfg.alignment == dispatch_output_per_expert_alignment, + "ep_prepare: alignment mismatch for handle_id=", handle_id, + " (cached=", cfg.alignment, ", got=", dispatch_output_per_expert_alignment, ")"); + ncclEpHandle_t h = get_or_open_handle(cfg, handle_mem); + NVTE_CHECK_NCCL(ncclEpUpdateHandle(h, &nccl_topk_idx, &layout_info, stream)); } void EPBackend::dispatch(uint64_t handle_id, void* handle_mem, const NVTETensor topk_idx, @@ -397,14 +394,10 @@ void EPBackend::dispatch(uint64_t handle_id, void* handle_mem, const NVTETensor ncclEpDispatchConfig_t dispatch_cfg = NCCL_EP_DISPATCH_CONFIG_INIT; dispatch_cfg.pass_direction = is_forward ? NCCL_EP_FWD_PASS : NCCL_EP_BWD_PASS; - ScopedEpHandle transient; - { - std::lock_guard lock(mutex_); - HandleEntry& cfg = lookup_config(handle_id); - transient = - ScopedEpHandle(open_handle(handle_mem, cfg.handle_mem_size, cfg.top_k, cfg.alignment)); - } - NVTE_CHECK_NCCL(ncclEpDispatch(transient, &in_struct, &out_struct, + std::lock_guard lock(mutex_); + HandleEntry& cfg = lookup_config(handle_id); + ncclEpHandle_t h = get_or_open_handle(cfg, handle_mem); + NVTE_CHECK_NCCL(ncclEpDispatch(h, &in_struct, &out_struct, /*layout_info=*/nullptr, &dispatch_cfg, stream)); } @@ -436,14 +429,10 @@ void EPBackend::combine(uint64_t handle_id, void* handle_mem, const NVTETensor e ncclEpCombineOutputs_t out_struct = NCCL_EP_COMBINE_OUTPUTS_INIT; out_struct.tokens = &nccl_result_out; - ScopedEpHandle transient; - { - std::lock_guard lock(mutex_); - HandleEntry& cfg = lookup_config(handle_id); - transient = - ScopedEpHandle(open_handle(handle_mem, cfg.handle_mem_size, cfg.top_k, cfg.alignment)); - } - NVTE_CHECK_NCCL(ncclEpCombine(transient, &in_struct, &out_struct, /*config=*/nullptr, stream)); + std::lock_guard lock(mutex_); + HandleEntry& cfg = lookup_config(handle_id); + ncclEpHandle_t h = get_or_open_handle(cfg, handle_mem); + NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, /*config=*/nullptr, stream)); } void EPBackend::dispatch_bwd(uint64_t handle_id, void* handle_mem, const NVTETensor grad, @@ -491,14 +480,10 @@ void EPBackend::dispatch_bwd(uint64_t handle_id, void* handle_mem, const NVTETen ncclEpCombineConfig_t cfg = NCCL_EP_COMBINE_CONFIG_INIT; cfg.pass_direction = NCCL_EP_BWD_PASS; - ScopedEpHandle transient; - { - std::lock_guard lock(mutex_); - HandleEntry& entry = lookup_config(handle_id); - transient = ScopedEpHandle( - open_handle(handle_mem, entry.handle_mem_size, entry.top_k, entry.alignment)); - } - NVTE_CHECK_NCCL(ncclEpCombine(transient, &in_struct, &out_struct, &cfg, stream)); + std::lock_guard lock(mutex_); + HandleEntry& entry = lookup_config(handle_id); + ncclEpHandle_t h = get_or_open_handle(entry, handle_mem); + NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, &cfg, stream)); } void EPBackend::combine_bwd(uint64_t handle_id, void* handle_mem, const NVTETensor grad, diff --git a/transformer_engine/common/ep/ep_backend.h b/transformer_engine/common/ep/ep_backend.h index 18307ebb4f..e82c974c3f 100644 --- a/transformer_engine/common/ep/ep_backend.h +++ b/transformer_engine/common/ep/ep_backend.h @@ -98,6 +98,10 @@ class EPBackend { size_t handle_mem_size; size_t alignment; int top_k; + // Persistent ncclEpHandle bound to cached_handle_mem. Lazily opened on first + // op; reused while handle_mem ptr is unchanged. Destroyed in shutdown(). + ncclEpHandle_t cached_handle{nullptr}; + void* cached_handle_mem{nullptr}; }; std::unordered_map handles_; std::atomic next_handle_id_{1}; // 0 reserved as "no id" @@ -106,6 +110,10 @@ class EPBackend { // Caller must hold mutex_. Throws on cap overflow. uint64_t insert_new_entry(size_t handle_mem_size, int top_k, size_t alignment); HandleEntry& lookup_config(uint64_t handle_id); + // Caller must hold mutex_. Returns the cached handle if handle_mem matches. + // On mismatch: if group_config_.allow_handle_mem_reloc != 0, destroys the + // stale handle and opens a fresh one; otherwise throws. + ncclEpHandle_t get_or_open_handle(HandleEntry& cfg, void* handle_mem); }; } // namespace ep From 0d3a565ca2de98000d1bc38f1cbac9eeffc568b0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 23 May 2026 23:09:15 +0000 Subject: [PATCH 03/55] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/ep/ep_backend.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index 6494a86817..83657943a4 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -324,8 +324,8 @@ void EPBackend::prepare(uint64_t handle_id, const NVTETensor topk_idx, NVTETenso std::lock_guard lock(mutex_); HandleEntry& cfg = lookup_config(handle_id); NVTE_CHECK(cfg.alignment == dispatch_output_per_expert_alignment, - "ep_prepare: alignment mismatch for handle_id=", handle_id, - " (cached=", cfg.alignment, ", got=", dispatch_output_per_expert_alignment, ")"); + "ep_prepare: alignment mismatch for handle_id=", handle_id, " (cached=", cfg.alignment, + ", got=", dispatch_output_per_expert_alignment, ")"); ncclEpHandle_t h = get_or_open_handle(cfg, handle_mem); NVTE_CHECK_NCCL(ncclEpUpdateHandle(h, &nccl_topk_idx, &layout_info, stream)); } From 44255194c5e3258f41f3db8c023ea0c4314e0a60 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 27 May 2026 14:12:53 -0700 Subject: [PATCH 04/55] Build: NCCL_HOME discovery supports Debian/Ubuntu multiarch lib paths Signed-off-by: Phuong Nguyen --- setup.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/setup.py b/setup.py index 9599193ced..d46690e626 100644 --- a/setup.py +++ b/setup.py @@ -169,11 +169,13 @@ def _discover_nccl_home() -> str: f"'{env_home}/include/nccl.h' was not found; falling back to system probes." ) + lib_names = ("libnccl.so", "libnccl.so.2") + # Include Debian/Ubuntu multiarch subdirs (e.g. lib/aarch64-linux-gnu). + lib_subdirs = ("lib", "lib64", "lib/aarch64-linux-gnu", "lib/x86_64-linux-gnu") for cand in ("/opt/nvidia/nccl", "/usr/local/nccl", "/usr"): p = Path(cand) if (p / "include" / "nccl.h").exists() and any( - (p / "lib" / name).exists() or (p / "lib64" / name).exists() - for name in ("libnccl.so", "libnccl.so.2") + (p / sub / name).exists() for sub in lib_subdirs for name in lib_names ): return str(p) @@ -182,9 +184,11 @@ def _discover_nccl_home() -> str: for line in out.splitlines(): if "libnccl.so" in line and "=>" in line: lib_path = Path(line.split("=>")[-1].strip()) - root = lib_path.parent.parent - if (root / "include" / "nccl.h").exists(): - return str(root) + # Walk upward so multiarch layouts (.../lib//libnccl.so) + # resolve to the prefix that contains include/nccl.h. + for root in (lib_path.parent.parent, lib_path.parent.parent.parent): + if (root / "include" / "nccl.h").exists(): + return str(root) except (subprocess.CalledProcessError, FileNotFoundError): pass From e7bb61bc4d697c9c3774e307225129c4d24b49e0 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 27 May 2026 14:26:39 -0700 Subject: [PATCH 05/55] bump NCCL Signed-off-by: Phuong Nguyen --- 3rdparty/nccl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/nccl b/3rdparty/nccl index 6a9bc953ac..146496ac88 160000 --- a/3rdparty/nccl +++ b/3rdparty/nccl @@ -1 +1 @@ -Subproject commit 6a9bc953ac1c4eef92d5adbe3092d4c2cb0a4c98 +Subproject commit 146496ac881bc504ed1a52be0ae7b707ce41e706 From d81e848af3152e624992056f8d1a959cc0a372d7 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 28 May 2026 15:25:16 -0700 Subject: [PATCH 06/55] Expert Parallelism: require token_dtype in NVTEEpGroupConfig and enforce at dispatch Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/test_ep_common.h | 4 ++++ transformer_engine/common/ep/ep_backend.cpp | 21 +++++++++++++++---- .../common/include/transformer_engine/ep.h | 3 +++ 3 files changed, 24 insertions(+), 4 deletions(-) diff --git a/tests/cpp_distributed/test_ep_common.h b/tests/cpp_distributed/test_ep_common.h index 77baa92b0c..ccb20ee3a0 100644 --- a/tests/cpp_distributed/test_ep_common.h +++ b/tests/cpp_distributed/test_ep_common.h @@ -74,6 +74,7 @@ static int g_ep_size = -1; static int g_num_experts = -1; static int g_hidden_dim = 256; static int g_max_tokens_per_rank = 64; +static NVTEDType g_token_dtype = kNVTEBFloat16; static bool g_ep_initialized = false; static ncclComm_t g_ep_comm = nullptr; // owned by harness, destroyed in ep_teardown @@ -224,6 +225,8 @@ static void ep_parse_args(int argc, char* argv[]) { else if (a.rfind("--num-processes=",0)==0) g_num_processes = std::stoi(a.substr(16)); else if (a.rfind("--nranks=", 0) == 0) g_num_processes = std::stoi(a.substr(9)); else if (a.rfind("--uid-file=", 0) == 0) g_uid_file = a.substr(11); + else if (a.rfind("--token-dtype=", 0) == 0) + g_token_dtype = static_cast(std::stoi(a.substr(14))); } if (g_process_id < 0 || g_num_processes <= 0) { @@ -279,6 +282,7 @@ static bool ep_bootstrap(int argc, char* argv[]) { // Worst-case for top_k fan-out: ep_size * max_tokens_per_rank * 2. group_config.max_recv_tokens_per_rank = g_ep_size * g_max_tokens_per_rank * 2; group_config.hidden_dim = g_hidden_dim; + group_config.token_dtype = g_token_dtype; ASSERT_NCCL_OK(ncclCommInitRank(&g_ep_comm, g_num_processes, uid, g_process_id)); nvte_ep_initialize(static_cast(g_ep_comm), group_config); diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index 83657943a4..1e08cb55df 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -82,9 +82,13 @@ void EPBackend::validate_config(const NVTEEpGroupConfig& config) { NVTE_CHECK(config.max_recv_tokens_per_rank > 0, "max_recv_tokens_per_rank must be positive, got ", config.max_recv_tokens_per_rank); NVTE_CHECK(config.hidden_dim > 0, "hidden_dim must be positive, got ", config.hidden_dim); - NVTE_CHECK(config.hidden_dim * sizeof(nv_bfloat16) >= 16, - "hidden_dim * 2 must be >= 16 (NCCL EP 16B row alignment); got hidden_dim=", - config.hidden_dim); + NVTE_CHECK(config.token_dtype >= 0 && config.token_dtype < kNVTENumTypes, + "token_dtype out of range, got ", static_cast(config.token_dtype)); + const size_t elem_bytes = typeToSize(static_cast(config.token_dtype)); + NVTE_CHECK(config.hidden_dim * elem_bytes >= 16, + "hidden_dim * sizeof(token_dtype) must be >= 16 (NCCL EP 16B row alignment); " + "got hidden_dim=", + config.hidden_dim, ", element_bytes=", elem_bytes); NVTE_CHECK(config.num_experts % config.ep_size == 0, "num_experts (", config.num_experts, ") must be divisible by ep_size (", config.ep_size, ")"); NVTE_CHECK(config.max_num_sms >= 0, "max_num_sms must be >= 0 (0 = auto), got ", @@ -214,7 +218,8 @@ void EPBackend::init(ncclComm_t ep_comm, NVTEEpGroupConfig group_config) { cfg.algorithm = NCCL_EP_ALGO_HIGH_THROUGHPUT; cfg.num_experts = static_cast(group_config.num_experts); cfg.max_dispatch_tokens_per_rank = static_cast(group_config.max_tokens_per_rank); - cfg.max_token_bytes = static_cast(group_config.hidden_dim * sizeof(nv_bfloat16)); + const size_t elem_bytes = typeToSize(static_cast(group_config.token_dtype)); + cfg.max_token_bytes = static_cast(group_config.hidden_dim * elem_bytes); cfg.rdma_buffer_size = NCCL_EP_AUTO; cfg.num_qp_per_rank = NCCL_EP_AUTO; cfg.num_channels = NCCL_EP_AUTO; @@ -341,6 +346,10 @@ void EPBackend::dispatch(uint64_t handle_id, void* handle_mem, const NVTETensor NVTEShape tok_shape = nvte_tensor_shape(tokens); NVTEDType tok_dtype = nvte_tensor_type(tokens); + NVTE_CHECK(tok_dtype == group_config_.token_dtype, + "tokens dtype (", static_cast(tok_dtype), + ") does not match group token_dtype (", + static_cast(group_config_.token_dtype), ")"); const size_t num_tokens = tok_shape.data[0]; const size_t hidden_dim = tok_shape.data[1]; @@ -367,6 +376,10 @@ void EPBackend::dispatch(uint64_t handle_id, void* handle_mem, const NVTETensor NVTEShape recv_shape = nvte_tensor_shape(recv_tokens); NVTEDType recv_dtype = nvte_tensor_type(recv_tokens); + NVTE_CHECK(recv_dtype == group_config_.token_dtype, + "recv_tokens dtype (", static_cast(recv_dtype), + ") does not match group token_dtype (", + static_cast(group_config_.token_dtype), ")"); size_t recv_sizes[2] = {recv_shape.data[0], recv_shape.data[1]}; ncclEpTensor_t nccl_tokens_out = make_payload_tensor(recv_tokens, recv_tokens_win, 2, diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index 8c3a06b5f0..ac7f1dbf07 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -35,6 +35,9 @@ typedef struct { int max_num_sms; /*!< Max SMs for EP kernels. 0 = auto. */ /*! 0 (default): throw on relocated handle_mem for a cached handle_id. 1: silently rebuild. */ int allow_handle_mem_reloc; + /*! Token dtype for this EP group. Sizes NCCL EP staging buffers at group + * create and is enforced against tensors passed to nvte_ep_dispatch. */ + NVTEDType token_dtype; } NVTEEpGroupConfig; /*! \brief Per-layer EP configuration. */ From e5eebdcfc09123d3c01a619dfdb71fbe2d392ff3 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 28 May 2026 15:31:47 -0700 Subject: [PATCH 07/55] Expert Parallelism: document ep_comm lifetime, v0.1 single-GPU scope, static layer registration Signed-off-by: Phuong Nguyen --- .../common/include/transformer_engine/ep.h | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index ac7f1dbf07..a1c9305e9b 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -54,8 +54,13 @@ typedef struct { /*! \brief Bootstrap from an existing NCCL EP sub-communicator. Requires SM>=90. * * ep_comm is borrowed and must span exactly group_config.ep_size ranks. + * The caller retains ownership and must keep ep_comm alive until + * nvte_ep_shutdown() returns; destroying it earlier is undefined behavior. * Re-init after shutdown is allowed; double-init throws. * + * v0.1 scope: one EP group per process, bound to the current CUDA device at + * initialize time. Multiple GPUs per process are not supported. + * * \param[in] ep_comm Opaque ncclComm_t for the EP sub-group. * \param[in] group_config Group-level EP configuration. */ @@ -69,6 +74,11 @@ void nvte_ep_shutdown(void); /*! \brief Reserve a handle_id for a layer config and report the handle_mem buffer * size the caller must allocate. Host-only. * + * Registration is intended to be static (once per layer at model init). There is + * no per-layer unregister API; all registrations are released by nvte_ep_shutdown. + * Re-registering the same layer config each step is not supported and will + * eventually exhaust the handle cache (NVTE_EP_HANDLE_CACHE_SIZE, default 8192). + * * \param[in] layer_config Per-layer EP configuration. * \param[out] handle_mem_size Bytes the caller must allocate for handle_mem. * \return uint64_t handle_id (non-zero). From 52c38793f10c041e7873ce4aafb544c3c42ab55a Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 28 May 2026 15:32:48 -0700 Subject: [PATCH 08/55] Expert Parallelism: drop version label from initialize scope note Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/CMakeLists.txt | 43 +- tests/cpp_distributed/run_test_ep.sh | 123 +--- .../{test_ep_pipeline.cu => test_ep.cu} | 643 ++++++++---------- tests/cpp_distributed/test_ep_common.h | 194 +----- tests/cpp_distributed/test_ep_coverage.cu | 562 --------------- tests/cpp_distributed/test_ep_init.cu | 64 -- transformer_engine/common/ep/ep_backend.cpp | 25 +- .../common/include/transformer_engine/ep.h | 13 +- transformer_engine/common/util/logging.h | 8 + 9 files changed, 375 insertions(+), 1300 deletions(-) rename tests/cpp_distributed/{test_ep_pipeline.cu => test_ep.cu} (51%) delete mode 100644 tests/cpp_distributed/test_ep_coverage.cu delete mode 100644 tests/cpp_distributed/test_ep_init.cu diff --git a/tests/cpp_distributed/CMakeLists.txt b/tests/cpp_distributed/CMakeLists.txt index 463ae011a5..191dde5d2d 100644 --- a/tests/cpp_distributed/CMakeLists.txt +++ b/tests/cpp_distributed/CMakeLists.txt @@ -92,10 +92,8 @@ endif() include(GoogleTest) gtest_discover_tests(test_comm_gemm DISCOVERY_TIMEOUT 600) -# ── EP distributed tests (HT mode) ───────────────────────────────────────── -# No MPI dependency — processes are spawned by run_test_ep.sh with -# --rank / --nranks flags. ncclUniqueId exchange uses a -# shared temp file (see test_ep_common.h for details). +# ── EP distributed tests ────────────────────────────────────────────────────── +# Launched via mpirun; ncclUniqueId exchange uses MPI_Bcast (see test_ep_common.h). # Headers + libs come from the in-tree 3rdparty/nccl submodule build. set(NCCL_EP_SUBMODULE_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl") @@ -122,41 +120,28 @@ endif() set(EP_TEST_COMMON_INCLUDES ${EP_TEST_NCCL_INCLUDES} + ${MPI_CXX_INCLUDE_PATH} ../../transformer_engine/common/include ../../transformer_engine/common ${CMAKE_CURRENT_SOURCE_DIR}) +# nvrtc must follow TE_LIB so symbols referenced from libtransformer_engine.so +# (loaded via dlopen in Python; not in its DT_NEEDED) resolve through nvrtc. set(EP_TEST_COMMON_LIBS CUDA::cuda_driver CUDA::cudart - CUDA::nvrtc GTest::gtest ${TE_LIB} + CUDA::nvrtc ${NCCL_LIB} - ${NCCL_EP_LIB}) - -# nvrtc symbols are referenced from libtransformer_engine.so but not in its -# DT_NEEDED list (loaded via dlopen in Python). For cpp tests we link nvrtc -# explicitly with --no-as-needed so the linker keeps the dependency. -set(EP_TEST_LINK_OPTS "LINKER:--no-as-needed") - -# ── EP init tests (InitPath, HandleMemSizeQuery) ───────────────────────────── -add_executable(test_ep_init test_ep_init.cu) -target_include_directories(test_ep_init PRIVATE ${EP_TEST_COMMON_INCLUDES}) -target_link_libraries(test_ep_init PUBLIC ${EP_TEST_COMMON_LIBS}) -target_link_options(test_ep_init PUBLIC ${EP_TEST_LINK_OPTS}) - -# ── EP pipeline tests (dispatch, combine, bwd, integrated) ─────────────────── -add_executable(test_ep_pipeline test_ep_pipeline.cu) -target_include_directories(test_ep_pipeline PRIVATE ${EP_TEST_COMMON_INCLUDES}) -target_link_libraries(test_ep_pipeline PUBLIC ${EP_TEST_COMMON_LIBS}) -target_link_options(test_ep_pipeline PUBLIC ${EP_TEST_LINK_OPTS}) - -# ── EP coverage tests (multi-handle, top_k=1, empty experts, negatives, threading) ── -add_executable(test_ep_coverage test_ep_coverage.cu) -target_include_directories(test_ep_coverage PRIVATE ${EP_TEST_COMMON_INCLUDES}) -target_link_libraries(test_ep_coverage PUBLIC ${EP_TEST_COMMON_LIBS}) -target_link_options(test_ep_coverage PUBLIC ${EP_TEST_LINK_OPTS}) + ${NCCL_EP_LIB} + MPI::MPI_CXX + OpenMP::OpenMP_CXX) + +# ── EP distributed tests (per-op + full pipeline + zero-copy symm) ─────────── +add_executable(test_ep test_ep.cu ../cpp/test_common.cu) +target_include_directories(test_ep PRIVATE ${EP_TEST_COMMON_INCLUDES}) +target_link_libraries(test_ep PUBLIC ${EP_TEST_COMMON_LIBS}) # Do NOT use gtest_discover_tests — these binaries require multi-process # launch via run_test_ep.sh, not direct single-process execution. diff --git a/tests/cpp_distributed/run_test_ep.sh b/tests/cpp_distributed/run_test_ep.sh index 017d3f807b..13e86fa02d 100755 --- a/tests/cpp_distributed/run_test_ep.sh +++ b/tests/cpp_distributed/run_test_ep.sh @@ -3,12 +3,8 @@ # # See LICENSE for license information. # -# Run TE EP distributed unit tests across multiple GPUs. -# -# Spawns one background bash process per GPU (no MPI dependency), matching the -# JAX multi-process launcher style. ncclUniqueId is exchanged via a shared -# temp file (see test_ep_common.h). Each rank builds its own ncclComm_t and -# passes it to nvte_ep_initialize. +# Run TE EP distributed unit tests via mpirun. Each MPI rank pins to one GPU +# (rank % device_count) and exchanges ncclUniqueId through MPI_Bcast. # # Usage: # bash run_test_ep.sh [num_gpus] [build_dir] @@ -18,15 +14,16 @@ # build_dir = /build # # Environment variables: -# GTEST_FILTER — forwarded to all processes (e.g., "EPDispatchTest.*") -# TEST_TIMEOUT_S — per-process timeout in seconds (default: 180) +# GTEST_FILTER — forwarded to all processes (e.g., "EPPipelineTest.*") +# MPIRUN — override the mpirun binary (default: mpirun) +# MPIRUN_EXTRA — extra flags forwarded to mpirun set -euo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" BUILD_DIR="${2:-${SCRIPT_DIR}/build}" NUM_GPUS="${1:-$(nvidia-smi -L 2>/dev/null | wc -l)}" -TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-180}" +MPIRUN="${MPIRUN:-mpirun}" # Skip cleanly on pre-Hopper: NCCL EP requires SM>=90. MIN_SM=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null \ @@ -36,102 +33,22 @@ if (( MIN_SM > 0 && MIN_SM < 90 )); then exit 0 fi -GTEST_ARGS="${GTEST_FILTER:+--gtest_filter=${GTEST_FILTER}}" -OVERALL_FAIL=0 - -# --------------------------------------------------------------------------- -# run_suite BINARY SUITE_NAME MIN_GPUS -# --------------------------------------------------------------------------- -run_suite() { - local BINARY="$1" - local SUITE_NAME="$2" - local MIN_GPUS="${3:-2}" - - local TEST_BIN="${BUILD_DIR}/${BINARY}" - - if [[ ! -x "${TEST_BIN}" ]]; then - echo "ERROR: binary not found: ${TEST_BIN}" - echo "Build: cd ${SCRIPT_DIR} && mkdir -p build && cd build && cmake .. && make" - OVERALL_FAIL=1 - return - fi - - if (( NUM_GPUS < MIN_GPUS )); then - echo "${SUITE_NAME}: requires ${MIN_GPUS} GPUs, found ${NUM_GPUS}. Skipping." - return - fi - - local TMPDIR_L="${TMPDIR:-/tmp}" - local UID_FILE="${TMPDIR_L}/te_ep_uid_${BINARY}_$$" - rm -f "${UID_FILE}" - - local LOG_DIR - LOG_DIR=$(mktemp -d) - local FAIL=0 - - echo "=== ${SUITE_NAME} ===" - echo " GPUs: ${NUM_GPUS} Binary: ${TEST_BIN}" - echo - - # Spawn one background process per GPU. ncclUniqueId is exchanged via the - # shared UID_FILE. Each process is wrapped in `timeout` to detect hangs early. - local PIDS=() - for i in $(seq 0 $((NUM_GPUS - 1))); do - timeout --foreground --signal=KILL "${TEST_TIMEOUT_S}" \ - "${TEST_BIN}" \ - --rank="${i}" \ - --nranks="${NUM_GPUS}" \ - --uid-file="${UID_FILE}" \ - ${GTEST_ARGS} \ - > "${LOG_DIR}/rank_${i}.log" 2>&1 & - PIDS+=($!) - done - for i in $(seq 0 $((NUM_GPUS - 1))); do - if ! wait "${PIDS[$i]}"; then - local rc=$? - FAIL=1 - if [[ $rc -eq 137 || $rc -eq 124 ]]; then - echo " rank ${i}: TIMEOUT after ${TEST_TIMEOUT_S}s (rc=${rc})" - fi - fi - done - - echo "--- Rank 0 output ---" - cat "${LOG_DIR}/rank_0.log" - - if (( FAIL )); then - for i in $(seq 1 $((NUM_GPUS - 1))); do - echo "--- Rank ${i} output ---" - cat "${LOG_DIR}/rank_${i}.log" - done - echo "=== ${SUITE_NAME}: FAILED ===" - OVERALL_FAIL=1 - else - echo "=== ${SUITE_NAME}: ALL PASSED ===" - fi - - rm -rf "${LOG_DIR}" - rm -f "${UID_FILE}" -} +TEST_BIN="${BUILD_DIR}/test_ep" +if [[ ! -x "${TEST_BIN}" ]]; then + echo "ERROR: binary not found: ${TEST_BIN}" + echo "Build: cd ${SCRIPT_DIR} && mkdir -p build && cd build && cmake .. && make" + exit 1 +fi -# --------------------------------------------------------------------------- -# Cleanup on abort -# --------------------------------------------------------------------------- -cleanup() { rm -f "${TMPDIR:-/tmp}"/te_ep_uid_*_"$$" 2>/dev/null || true; } -trap cleanup EXIT INT TERM +if (( NUM_GPUS < 2 )); then + echo "EP Tests: requires at least 2 GPUs, found ${NUM_GPUS}. Skipping." + exit 0 +fi -# --------------------------------------------------------------------------- -# Run all suites -# --------------------------------------------------------------------------- -run_suite "test_ep_init" "EP Init Tests" 2 -run_suite "test_ep_pipeline" "EP Pipeline Tests" 2 -run_suite "test_ep_coverage" "EP Coverage Tests" 2 +GTEST_ARGS="${GTEST_FILTER:+--gtest_filter=${GTEST_FILTER}}" +echo "=== EP Tests ===" +echo " GPUs: ${NUM_GPUS} Binary: ${TEST_BIN}" echo -if (( OVERALL_FAIL )); then - echo "=== SOME SUITES FAILED ===" -else - echo "=== ALL SUITES PASSED ===" -fi -exit "${OVERALL_FAIL}" +"${MPIRUN}" -n "${NUM_GPUS}" ${MPIRUN_EXTRA:-} "${TEST_BIN}" ${GTEST_ARGS} diff --git a/tests/cpp_distributed/test_ep_pipeline.cu b/tests/cpp_distributed/test_ep.cu similarity index 51% rename from tests/cpp_distributed/test_ep_pipeline.cu rename to tests/cpp_distributed/test_ep.cu index 41f83a6d11..bcf4ca3c98 100644 --- a/tests/cpp_distributed/test_ep_pipeline.cu +++ b/tests/cpp_distributed/test_ep.cu @@ -39,10 +39,21 @@ static inline float token_value(int rank, int t, int num_tokens) { return static_cast(rank * num_tokens + t + 1) * (1.0f / 256.0f); } -static std::vector generate_tokens(int rank, int num_tokens, int hidden_dim) { - std::vector v(num_tokens * hidden_dim); +// Per-element host-side conversion helpers used by templated test code. +inline float tok_to_float(nv_bfloat16 v) { return __bfloat162float(v); } +inline float tok_to_float(__half v) { return __half2float(v); } +inline float tok_to_float(float v) { return v; } + +template T tok_from_float(float v); +template <> inline nv_bfloat16 tok_from_float(float v) { return __float2bfloat16(v); } +template <> inline __half tok_from_float<__half> (float v) { return __float2half(v); } +template <> inline float tok_from_float (float v) { return v; } + +template +static std::vector generate_tokens(int rank, int num_tokens, int hidden_dim) { + std::vector v(num_tokens * hidden_dim); for (int t = 0; t < num_tokens; ++t) { - nv_bfloat16 val = __float2bfloat16(token_value(rank, t, num_tokens)); + T val = tok_from_float(token_value(rank, t, num_tokens)); for (int h = 0; h < hidden_dim; ++h) v[t * hidden_dim + h] = val; } @@ -85,17 +96,20 @@ static std::vector expected_recv_values_sorted( return vals; } -// BF16 has 7 mantissa bits; relative ULP ≈ 2^-7. Use 4× headroom for -// accumulation noise inside dispatch/combine. +// 2^-5 relative tolerance for BF16 (matches mantissa precision with margin), +// plus a small atol floor for near-zero expected values. +static constexpr float kBf16Rtol = 1.0f / 32.0f; +static constexpr float kBf16Atol = 1e-3f; static float bf16_tol(float magnitude) { - return 4.f * std::ldexp(std::fabs(magnitude) + 1e-3f, -7); + return kBf16Atol + kBf16Rtol * std::fabs(magnitude); } -static bool check_no_nan_inf(const nv_bfloat16* dev, int count, const char* name) { - std::vector h(count); - cudaMemcpy(h.data(), dev, count * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost); +template +static bool check_no_nan_inf(const T* dev, int count, const char* name) { + std::vector h(count); + cudaMemcpy(h.data(), dev, count * sizeof(T), cudaMemcpyDeviceToHost); for (int i = 0; i < count; ++i) { - float v = __bfloat162float(h[i]); + float v = tok_to_float(h[i]); if (std::isnan(v) || std::isinf(v)) { fprintf(stderr, "Rank %d: %s in %s[%d]\n", g_process_id, std::isnan(v) ? "NaN" : "Inf", name, i); @@ -107,20 +121,21 @@ static bool check_no_nan_inf(const nv_bfloat16* dev, int count, const char* name // ── Forward buffer set with RAII ────────────────────────────────────────────── +template struct EPBuffers { // Forward DevBuf topk_idx; DevBuf topk_weights; - DevBuf tokens; + DevBuf tokens; DevBuf token_counts; DevBuf handle_mem; - DevBuf recv_tokens; + DevBuf recv_tokens; DevBuf recv_topk_weights; - DevBuf result; + DevBuf result; // Backward - DevBuf grad_result; - DevBuf grad_expert; - DevBuf grad_tokens; + DevBuf grad_result; + DevBuf grad_expert; + DevBuf grad_tokens; DevBuf g_recv_topk_weights; DevBuf grad_topk_weights; @@ -154,42 +169,45 @@ struct EPBuffers { } }; -// Bundled NVTETensor views over an EPBuffers — one place to update the shape -// conventions when the C-API evolves. +// Bundled NVTETensor views over an EPBuffers, with the shapes the EP C API +// expects. +template struct EPTensors { - TensorHandle topk_idx, topk_weights, token_counts, handle_mem, tokens; - TensorHandle recv_tokens, recv_topk_weights, result; - TensorHandle grad_result, grad_expert, grad_tokens; - TensorHandle g_recv_topk_weights, grad_topk_weights; + TensorWrapper topk_idx, topk_weights, token_counts, handle_mem, tokens; + TensorWrapper recv_tokens, recv_topk_weights, result; + TensorWrapper grad_result, grad_expert, grad_tokens; + TensorWrapper g_recv_topk_weights, grad_topk_weights; - EPTensors(EPBuffers& b, int num_tokens, int top_k, int hidden_dim, + EPTensors(EPBuffers& b, int num_tokens, int top_k, int hidden_dim, int num_local_experts) { - topk_idx = make_nvte_tensor(b.topk_idx.get(), - {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); - topk_weights = make_nvte_tensor(b.topk_weights.get(), - {(size_t)num_tokens, (size_t)top_k}, kNVTEFloat32); - token_counts = make_nvte_tensor(b.token_counts.get(), - {(size_t)num_local_experts}, kNVTEInt32); - handle_mem = make_nvte_tensor(b.handle_mem.get(), - {b.handle_mem_size}, kNVTEByte); - tokens = make_nvte_tensor(b.tokens.get(), - {(size_t)num_tokens, (size_t)hidden_dim}, kNVTEBFloat16); - recv_tokens = make_nvte_tensor(b.recv_tokens.get(), - {b.recv_capacity, (size_t)hidden_dim}, kNVTEBFloat16); - recv_topk_weights = make_nvte_tensor(b.recv_topk_weights.get(), - {b.recv_capacity}, kNVTEFloat32); - result = make_nvte_tensor(b.result.get(), - {(size_t)num_tokens, (size_t)hidden_dim}, kNVTEBFloat16); - grad_result = make_nvte_tensor(b.grad_result.get(), - {(size_t)num_tokens, (size_t)hidden_dim}, kNVTEBFloat16); - grad_expert = make_nvte_tensor(b.grad_expert.get(), - {b.recv_capacity, (size_t)hidden_dim}, kNVTEBFloat16); - grad_tokens = make_nvte_tensor(b.grad_tokens.get(), - {(size_t)num_tokens, (size_t)hidden_dim}, kNVTEBFloat16); - g_recv_topk_weights = make_nvte_tensor(b.g_recv_topk_weights.get(), - {b.recv_capacity}, kNVTEFloat32); - grad_topk_weights = make_nvte_tensor(b.grad_topk_weights.get(), - {(size_t)num_tokens, (size_t)top_k}, kNVTEFloat32); + constexpr DType kTokDType = test::TypeInfo::dtype; + using Shape = std::vector; + topk_idx = TensorWrapper(b.topk_idx.get(), + Shape{(size_t)num_tokens, (size_t)top_k}, DType::kInt64); + topk_weights = TensorWrapper(b.topk_weights.get(), + Shape{(size_t)num_tokens, (size_t)top_k}, DType::kFloat32); + token_counts = TensorWrapper(b.token_counts.get(), + Shape{(size_t)num_local_experts}, DType::kInt32); + handle_mem = TensorWrapper(b.handle_mem.get(), + Shape{b.handle_mem_size}, DType::kByte); + tokens = TensorWrapper(b.tokens.get(), + Shape{(size_t)num_tokens, (size_t)hidden_dim}, kTokDType); + recv_tokens = TensorWrapper(b.recv_tokens.get(), + Shape{b.recv_capacity, (size_t)hidden_dim}, kTokDType); + recv_topk_weights = TensorWrapper(b.recv_topk_weights.get(), + Shape{b.recv_capacity}, DType::kFloat32); + result = TensorWrapper(b.result.get(), + Shape{(size_t)num_tokens, (size_t)hidden_dim}, kTokDType); + grad_result = TensorWrapper(b.grad_result.get(), + Shape{(size_t)num_tokens, (size_t)hidden_dim}, kTokDType); + grad_expert = TensorWrapper(b.grad_expert.get(), + Shape{b.recv_capacity, (size_t)hidden_dim}, kTokDType); + grad_tokens = TensorWrapper(b.grad_tokens.get(), + Shape{(size_t)num_tokens, (size_t)hidden_dim}, kTokDType); + g_recv_topk_weights = TensorWrapper(b.g_recv_topk_weights.get(), + Shape{b.recv_capacity}, DType::kFloat32); + grad_topk_weights = TensorWrapper(b.grad_topk_weights.get(), + Shape{(size_t)num_tokens, (size_t)top_k}, DType::kFloat32); } }; @@ -215,29 +233,31 @@ class EpOpTestBase : public ::testing::Test { num_tokens_ = 32; } - void upload_inputs(EPBuffers& buf, int rank = -1) { + template + void upload_inputs(EPBuffers& buf, int rank = -1) { if (rank < 0) rank = g_process_id; auto h_idx = routing_balanced(rank, num_tokens_, top_k_, num_experts_, num_local_experts_); std::vector h_w(num_tokens_ * top_k_, 1.0f / top_k_); - auto h_tok = generate_tokens(rank, num_tokens_, hidden_dim_); - - CHECK_CUDA(cudaMemcpy(buf.topk_idx.get(), h_idx.data(), - h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(buf.topk_weights.get(), h_w.data(), - h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(buf.tokens.get(), h_tok.data(), - h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); + auto h_tok = generate_tokens(rank, num_tokens_, hidden_dim_); + + NVTE_CHECK_CUDA(cudaMemcpy(buf.topk_idx.get(), h_idx.data(), + h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(buf.topk_weights.get(), h_w.data(), + h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(buf.tokens.get(), h_tok.data(), + h_tok.size() * sizeof(T), cudaMemcpyHostToDevice)); } NVTEEpLayerConfig layer_config(size_t alignment = 0) const { return NVTEEpLayerConfig{num_local_experts_, top_k_, alignment}; } - // ASSERT_CUDA_OK (fprintf+exit) so this non-void helper stays legal. - int read_total_recv(const EPBuffers& buf) const { + // NVTE_CHECK_CUDA (fprintf+exit) so this non-void helper stays legal. + template + int read_total_recv(const EPBuffers& buf) const { std::vector cnt(num_local_experts_); - ASSERT_CUDA_OK(cudaMemcpy(cnt.data(), buf.token_counts.get(), + NVTE_CHECK_CUDA(cudaMemcpy(cnt.data(), buf.token_counts.get(), num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost)); int total = 0; for (int c : cnt) total += c; @@ -252,28 +272,28 @@ class EpOpTestBase : public ::testing::Test { class EPDispatchTest : public EpOpTestBase {}; TEST_F(EPDispatchTest, PrepareAndDispatch) { - EPBuffers buf; + EPBuffers<> buf; buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, ep_size_, max_tokens_per_rank_); upload_inputs(buf); - EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + EPTensors<> t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); - CHECK_CUDA(cudaMemset(buf.recv_tokens.get(), 0, buf.recv_tokens.bytes())); + NVTE_CHECK_CUDA(cudaMemset(buf.recv_tokens.get(), 0, buf.recv_tokens.bytes())); cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); + NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); uint64_t handle_id = buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, - t.tokens.tensor, NVTECommWindow{}, t.topk_weights.tensor, - NVTECommWindow{}, t.recv_tokens.tensor, NVTECommWindow{}, - t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), t.token_counts.data(), /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), + t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), + NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, + t.recv_topk_weights.data(), NVTECommWindow{}, stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); // 1. Per-expert counts. std::vector got_counts(num_local_experts_); - CHECK_CUDA(cudaMemcpy(got_counts.data(), buf.token_counts.get(), + NVTE_CHECK_CUDA(cudaMemcpy(got_counts.data(), buf.token_counts.get(), num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost)); auto exp_counts = expected_token_counts(g_process_id, g_num_processes, num_tokens_, top_k_, num_experts_, num_local_experts_); @@ -288,7 +308,7 @@ TEST_F(EPDispatchTest, PrepareAndDispatch) { // 2. Recv values: read only the filled prefix per local-expert zone, not the // whole recv buffer — avoids false positives from legitimate-zero token values. std::vector h_recv(buf.recv_capacity * hidden_dim_); - CHECK_CUDA(cudaMemcpy(h_recv.data(), buf.recv_tokens.get(), + NVTE_CHECK_CUDA(cudaMemcpy(h_recv.data(), buf.recv_tokens.get(), h_recv.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); std::vector got_vals; @@ -312,7 +332,7 @@ TEST_F(EPDispatchTest, PrepareAndDispatch) { // 3. recv_topk_weights: every filled slot must equal the per-token weight (1/top_k). std::vector h_w(buf.recv_capacity); - CHECK_CUDA(cudaMemcpy(h_w.data(), buf.recv_topk_weights.get(), + NVTE_CHECK_CUDA(cudaMemcpy(h_w.data(), buf.recv_topk_weights.get(), h_w.size() * sizeof(float), cudaMemcpyDeviceToHost)); const float exp_w = 1.0f / static_cast(top_k_); for (int i = 0; i < total_recv; ++i) @@ -321,7 +341,7 @@ TEST_F(EPDispatchTest, PrepareAndDispatch) { if (g_process_id == 0) printf(" PrepareAndDispatch: passed (recv=%d, values + weights exact)\n", total_recv); - CHECK_CUDA(cudaStreamDestroy(stream)); + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } // ============================================================================= @@ -331,34 +351,32 @@ TEST_F(EPDispatchTest, PrepareAndDispatch) { class EPCombineTest : public EpOpTestBase {}; TEST_F(EPCombineTest, Combine) { - EPBuffers buf; + EPBuffers<> buf; buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, ep_size_, max_tokens_per_rank_); upload_inputs(buf); - EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + EPTensors<> t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); + NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); uint64_t handle_id = buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, - t.tokens.tensor, NVTECommWindow{}, t.topk_weights.tensor, - NVTECommWindow{}, t.recv_tokens.tensor, NVTECommWindow{}, - t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.recv_tokens.tensor, NVTECommWindow{}, - t.result.tensor, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), t.token_counts.data(), /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), + t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), + NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, + t.recv_topk_weights.data(), NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.recv_tokens.data(), NVTECommWindow{}, + t.result.data(), stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); std::vector h_result(num_tokens_ * hidden_dim_); - CHECK_CUDA(cudaMemcpy(h_result.data(), buf.result.get(), + NVTE_CHECK_CUDA(cudaMemcpy(h_result.data(), buf.result.get(), h_result.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); auto h_tok = generate_tokens(g_process_id, num_tokens_, hidden_dim_); - // Spot-check 3 hidden-dim positions per token to catch partial-row writes. - const int probes[3] = {0, hidden_dim_ / 2, hidden_dim_ - 1}; for (int tok = 0; tok < num_tokens_; ++tok) { float exp = __bfloat162float(h_tok[tok * hidden_dim_]) * static_cast(top_k_); - for (int p : probes) { + for (int p = 0; p < hidden_dim_; ++p) { float got = __bfloat162float(h_result[tok * hidden_dim_ + p]); EXPECT_NEAR(got, exp, bf16_tol(exp)) << "token " << tok << " rank " << g_process_id << " hidden " << p; @@ -368,7 +386,7 @@ TEST_F(EPCombineTest, Combine) { if (g_process_id == 0) printf(" Combine: passed (result == top_k * tokens)\n"); - CHECK_CUDA(cudaStreamDestroy(stream)); + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } // ============================================================================= @@ -378,41 +396,41 @@ TEST_F(EPCombineTest, Combine) { class EPCombineBwdTest : public EpOpTestBase {}; TEST_F(EPCombineBwdTest, CombineBwdCheck) { - EPBuffers buf; + EPBuffers<> buf; buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, ep_size_, max_tokens_per_rank_); upload_inputs(buf); - EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + EPTensors<> t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); + NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); uint64_t handle_id = buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, - t.tokens.tensor, NVTECommWindow{}, t.topk_weights.tensor, - NVTECommWindow{}, t.recv_tokens.tensor, NVTECommWindow{}, - t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.recv_tokens.tensor, NVTECommWindow{}, - t.result.tensor, stream)); + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), t.token_counts.data(), /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), + t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), + NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, + t.recv_topk_weights.data(), NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.recv_tokens.data(), NVTECommWindow{}, + t.result.data(), stream)); std::vector h_grad_r(num_tokens_ * hidden_dim_, __float2bfloat16(0.1f)); - CHECK_CUDA(cudaMemcpyAsync(buf.grad_result.get(), h_grad_r.data(), + NVTE_CHECK_CUDA(cudaMemcpyAsync(buf.grad_result.get(), h_grad_r.data(), h_grad_r.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice, stream)); - CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); + NVTE_CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); - ASSERT_NO_THROW(nvte_ep_combine_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_result.tensor, NVTECommWindow{}, - t.grad_expert.tensor, NVTECommWindow{}, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); + ASSERT_NO_THROW(nvte_ep_combine_bwd(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.grad_result.data(), NVTECommWindow{}, + t.grad_expert.data(), NVTECommWindow{}, stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); int total_recv = read_total_recv(buf); std::vector cnt(num_local_experts_); - CHECK_CUDA(cudaMemcpy(cnt.data(), buf.token_counts.get(), + NVTE_CHECK_CUDA(cudaMemcpy(cnt.data(), buf.token_counts.get(), num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost)); std::vector h_ge(buf.recv_capacity * hidden_dim_); - CHECK_CUDA(cudaMemcpy(h_ge.data(), buf.grad_expert.get(), + NVTE_CHECK_CUDA(cudaMemcpy(h_ge.data(), buf.grad_expert.get(), h_ge.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); // Walk filled slots by per-expert zone (no v != 0 heuristic). @@ -421,9 +439,12 @@ TEST_F(EPCombineBwdTest, CombineBwdCheck) { int filled = 0; for (int e = 0; e < num_local_experts_; ++e) { for (int i = 0; i < cnt[e]; ++i) { - float v = __bfloat162float(h_ge[slot * hidden_dim_]); - EXPECT_NEAR(v, kExpGrad, bf16_tol(kExpGrad)) - << "grad_expert expert " << e << " slot " << i << " (linear " << slot << ")"; + for (int p = 0; p < hidden_dim_; ++p) { + float v = __bfloat162float(h_ge[slot * hidden_dim_ + p]); + EXPECT_NEAR(v, kExpGrad, bf16_tol(kExpGrad)) + << "grad_expert expert " << e << " slot " << i + << " (linear " << slot << ") hidden " << p; + } ++filled; ++slot; } } @@ -432,7 +453,7 @@ TEST_F(EPCombineBwdTest, CombineBwdCheck) { if (g_process_id == 0) printf(" CombineBwdCheck: passed (filled=%d)\n", filled); - CHECK_CUDA(cudaStreamDestroy(stream)); + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } // ============================================================================= @@ -442,51 +463,53 @@ TEST_F(EPCombineBwdTest, CombineBwdCheck) { class EPDispatchBwdTest : public EpOpTestBase {}; TEST_F(EPDispatchBwdTest, DispatchBwdCheck) { - EPBuffers buf; + EPBuffers<> buf; buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, ep_size_, max_tokens_per_rank_); upload_inputs(buf); - EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + EPTensors<> t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); + NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); uint64_t handle_id = buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, - t.tokens.tensor, NVTECommWindow{}, t.topk_weights.tensor, - NVTECommWindow{}, t.recv_tokens.tensor, NVTECommWindow{}, - t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.recv_tokens.tensor, NVTECommWindow{}, - t.result.tensor, stream)); + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), t.token_counts.data(), /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), + t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), + NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, + t.recv_topk_weights.data(), NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.recv_tokens.data(), NVTECommWindow{}, + t.result.data(), stream)); std::vector h_grad(num_tokens_ * hidden_dim_, __float2bfloat16(0.1f)); - CHECK_CUDA(cudaMemcpyAsync(buf.grad_result.get(), h_grad.data(), + NVTE_CHECK_CUDA(cudaMemcpyAsync(buf.grad_result.get(), h_grad.data(), h_grad.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice, stream)); - CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); - CHECK_CUDA(cudaMemsetAsync(buf.g_recv_topk_weights.get(), 0, buf.g_recv_topk_weights.bytes(), stream)); - CHECK_CUDA(cudaMemsetAsync(buf.grad_topk_weights.get(), 0, buf.grad_topk_weights.bytes(), stream)); + NVTE_CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); + NVTE_CHECK_CUDA(cudaMemsetAsync(buf.g_recv_topk_weights.get(), 0, buf.g_recv_topk_weights.bytes(), stream)); + NVTE_CHECK_CUDA(cudaMemsetAsync(buf.grad_topk_weights.get(), 0, buf.grad_topk_weights.bytes(), stream)); - ASSERT_NO_THROW(nvte_ep_combine_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_result.tensor, NVTECommWindow{}, - t.grad_expert.tensor, NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_expert.tensor, NVTECommWindow{}, - t.g_recv_topk_weights.tensor, NVTECommWindow{}, - t.grad_tokens.tensor, t.grad_topk_weights.tensor, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); + ASSERT_NO_THROW(nvte_ep_combine_bwd(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.grad_result.data(), NVTECommWindow{}, + t.grad_expert.data(), NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch_bwd(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.grad_expert.data(), NVTECommWindow{}, + t.g_recv_topk_weights.data(), NVTECommWindow{}, + t.grad_tokens.data(), t.grad_topk_weights.data(), stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); std::vector h_gt(num_tokens_ * hidden_dim_); - CHECK_CUDA(cudaMemcpy(h_gt.data(), buf.grad_tokens.get(), + NVTE_CHECK_CUDA(cudaMemcpy(h_gt.data(), buf.grad_tokens.get(), h_gt.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); const float kExpGrad = static_cast(top_k_) * 0.1f; for (int tok = 0; tok < num_tokens_; ++tok) - EXPECT_NEAR(__bfloat162float(h_gt[tok * hidden_dim_]), kExpGrad, bf16_tol(kExpGrad)) - << "grad_tokens token " << tok; + for (int p = 0; p < hidden_dim_; ++p) + EXPECT_NEAR(__bfloat162float(h_gt[tok * hidden_dim_ + p]), kExpGrad, + bf16_tol(kExpGrad)) + << "grad_tokens token " << tok << " hidden " << p; if (g_process_id == 0) printf(" DispatchBwdCheck: passed (grad_tokens == %.2f)\n", kExpGrad); - CHECK_CUDA(cudaStreamDestroy(stream)); + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } // ============================================================================= @@ -496,11 +519,11 @@ TEST_F(EPDispatchBwdTest, DispatchBwdCheck) { class EPDispatchBwdGradWeightsTest : public EpOpTestBase {}; TEST_F(EPDispatchBwdGradWeightsTest, RoundTrip) { - EPBuffers buf; + EPBuffers<> buf; buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, ep_size_, max_tokens_per_rank_); upload_inputs(buf); - EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + EPTensors<> t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); // Distinct per-(rank, t, k) weights so each slot carries a unique value. std::vector h_w(num_tokens_ * top_k_); @@ -508,39 +531,39 @@ TEST_F(EPDispatchBwdGradWeightsTest, RoundTrip) { for (int k = 0; k < top_k_; ++k) h_w[tok * top_k_ + k] = 0.1f + 0.01f * tok + 0.001f * k + 0.0001f * (g_process_id + 1); - CHECK_CUDA(cudaMemcpy(buf.topk_weights.get(), h_w.data(), + NVTE_CHECK_CUDA(cudaMemcpy(buf.topk_weights.get(), h_w.data(), h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); + NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); uint64_t handle_id = buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); - CHECK_CUDA(cudaMemsetAsync(buf.recv_topk_weights.get(), 0, + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), t.token_counts.data(), /*alignment=*/0, stream)); + NVTE_CHECK_CUDA(cudaMemsetAsync(buf.recv_topk_weights.get(), 0, buf.recv_topk_weights.bytes(), stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, - t.tokens.tensor, NVTECommWindow{}, t.topk_weights.tensor, - NVTECommWindow{}, t.recv_tokens.tensor, NVTECommWindow{}, - t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), + t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), + NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, + t.recv_topk_weights.data(), NVTECommWindow{}, stream)); // Sentinel: NaN so any (t, k) the bwd kernel fails to write is immediately visible. std::vector h_nan(num_tokens_ * top_k_, std::numeric_limits::quiet_NaN()); - CHECK_CUDA(cudaMemcpyAsync(buf.grad_topk_weights.get(), h_nan.data(), + NVTE_CHECK_CUDA(cudaMemcpyAsync(buf.grad_topk_weights.get(), h_nan.data(), h_nan.size() * sizeof(float), cudaMemcpyHostToDevice, stream)); - CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); + NVTE_CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); // g_recv_topk_weights := recv_topk_weights (the round-trip input). - auto g_recv_t = make_nvte_tensor(buf.recv_topk_weights.get(), - {buf.recv_capacity}, kNVTEFloat32); - ASSERT_NO_THROW(nvte_ep_dispatch_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_expert.tensor, - NVTECommWindow{}, g_recv_t.tensor, NVTECommWindow{}, - t.grad_tokens.tensor, t.grad_topk_weights.tensor, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); + auto g_recv_t = TensorWrapper(buf.recv_topk_weights.get(), + std::vector{buf.recv_capacity}, DType::kFloat32); + ASSERT_NO_THROW(nvte_ep_dispatch_bwd(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.grad_expert.data(), + NVTECommWindow{}, g_recv_t.data(), NVTECommWindow{}, + t.grad_tokens.data(), t.grad_topk_weights.data(), stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); std::vector h_grad_w(num_tokens_ * top_k_); - CHECK_CUDA(cudaMemcpy(h_grad_w.data(), buf.grad_topk_weights.get(), + NVTE_CHECK_CUDA(cudaMemcpy(h_grad_w.data(), buf.grad_topk_weights.get(), h_grad_w.size() * sizeof(float), cudaMemcpyDeviceToHost)); const float kTol = 1e-5f; @@ -566,57 +589,81 @@ TEST_F(EPDispatchBwdGradWeightsTest, RoundTrip) { if (g_process_id == 0 && errs == 0 && k0_eq_k1 == 0) printf(" RoundTrip: passed (%d (t, k) gradients)\n", num_tokens_ * top_k_); - CHECK_CUDA(cudaStreamDestroy(stream)); + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } // ============================================================================= // Integrated FwdBwd: NaN/Inf check end-to-end. // ============================================================================= -class EPPipelineTest : public EpOpTestBase {}; - -TEST_F(EPPipelineTest, FullForwardBackward) { - EPBuffers buf; - buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, - ep_size_, max_tokens_per_rank_); - upload_inputs(buf); - EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); - - cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); - - uint64_t handle_id = buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, - t.tokens.tensor, NVTECommWindow{}, t.topk_weights.tensor, - NVTECommWindow{}, t.recv_tokens.tensor, NVTECommWindow{}, - t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.recv_tokens.tensor, NVTECommWindow{}, - t.result.tensor, stream)); - - std::vector h_grad(num_tokens_ * hidden_dim_, __float2bfloat16(0.1f)); - CHECK_CUDA(cudaMemcpyAsync(buf.grad_result.get(), h_grad.data(), - h_grad.size() * sizeof(nv_bfloat16), - cudaMemcpyHostToDevice, stream)); - CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); - CHECK_CUDA(cudaMemsetAsync(buf.g_recv_topk_weights.get(), 0, buf.g_recv_topk_weights.bytes(), stream)); - CHECK_CUDA(cudaMemsetAsync(buf.grad_topk_weights.get(), 0, buf.grad_topk_weights.bytes(), stream)); - - ASSERT_NO_THROW(nvte_ep_combine_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_result.tensor, NVTECommWindow{}, - t.grad_expert.tensor, NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_expert.tensor, NVTECommWindow{}, - t.g_recv_topk_weights.tensor, NVTECommWindow{}, - t.grad_tokens.tensor, t.grad_topk_weights.tensor, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); - - ASSERT_TRUE(check_no_nan_inf(buf.result.get(), num_tokens_ * hidden_dim_, "result")); - ASSERT_TRUE(check_no_nan_inf(buf.grad_tokens.get(), num_tokens_ * hidden_dim_, "grad_tokens")); - - if (g_process_id == 0) printf(" FullForwardBackward: passed\n"); +class EPPipelineTest : public EpOpTestBase, public ::testing::WithParamInterface { + protected: + template + void run_full_forward_backward() { + EPBuffers buf; + buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, + ep_size_, max_tokens_per_rank_); + upload_inputs(buf); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + + cudaStream_t stream; + NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); + + uint64_t handle_id = buf.handle_id; + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), t.token_counts.data(), /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), + t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), + NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, + t.recv_topk_weights.data(), NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.recv_tokens.data(), NVTECommWindow{}, + t.result.data(), stream)); + + std::vector h_grad(num_tokens_ * hidden_dim_, tok_from_float(0.1f)); + NVTE_CHECK_CUDA(cudaMemcpyAsync(buf.grad_result.get(), h_grad.data(), + h_grad.size() * sizeof(Tok), + cudaMemcpyHostToDevice, stream)); + NVTE_CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); + NVTE_CHECK_CUDA(cudaMemsetAsync(buf.g_recv_topk_weights.get(), 0, buf.g_recv_topk_weights.bytes(), stream)); + NVTE_CHECK_CUDA(cudaMemsetAsync(buf.grad_topk_weights.get(), 0, buf.grad_topk_weights.bytes(), stream)); + + ASSERT_NO_THROW(nvte_ep_combine_bwd(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.grad_result.data(), NVTECommWindow{}, + t.grad_expert.data(), NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch_bwd(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.grad_expert.data(), NVTECommWindow{}, + t.g_recv_topk_weights.data(), NVTECommWindow{}, + t.grad_tokens.data(), t.grad_topk_weights.data(), stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); + + ASSERT_TRUE(check_no_nan_inf(buf.result.get(), num_tokens_ * hidden_dim_, "result")); + ASSERT_TRUE(check_no_nan_inf(buf.grad_tokens.get(), num_tokens_ * hidden_dim_, "grad_tokens")); + + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); + } +}; - CHECK_CUDA(cudaStreamDestroy(stream)); +TEST_P(EPPipelineTest, FullForwardBackward) { + const DType dtype = GetParam(); + // NCCL EP backend currently asserts ncclBfloat16 in ncclEpDispatch + // (contrib/nccl_ep/nccl_ep.cc); skip FP16/FP32 until the backend supports them. + if (dtype != DType::kBFloat16) { + GTEST_SKIP() << test::typeName(dtype) << " not yet supported by NCCL EP backend"; + } + switch (dtype) { + case DType::kBFloat16: run_full_forward_backward(); break; + case DType::kFloat16: run_full_forward_backward<__half> (); break; + case DType::kFloat32: run_full_forward_backward (); break; + default: FAIL() << "unsupported token dtype " << static_cast(dtype); + } + if (g_process_id == 0) + printf(" FullForwardBackward[%s]: passed\n", test::typeName(dtype).c_str()); } +INSTANTIATE_TEST_SUITE_P( + Dtypes, EPPipelineTest, + ::testing::Values(DType::kBFloat16, DType::kFloat16, DType::kFloat32), + [](const ::testing::TestParamInfo& info) { + return test::typeName(info.param); + }); + // ============================================================================= // EPZeroCopyTest: dispatch/combine with NCCL symmetric-memory windows attached // to payload tensors (zero-copy fast path via ncclEpTensorCreateFromWindow). @@ -646,9 +693,9 @@ struct SymmBuf { void alloc(size_t n_bytes) { bytes = n_bytes; - ASSERT_NCCL_OK(ncclMemAlloc(&ptr, bytes)); - CHECK_CUDA(cudaMemset(ptr, 0, bytes)); - ASSERT_NCCL_OK(ncclCommWindowRegister(g_ep_comm, ptr, bytes, &win, + NVTE_CHECK_NCCL(ncclMemAlloc(&ptr, bytes)); + NVTE_CHECK_CUDA(cudaMemset(ptr, 0, bytes)); + NVTE_CHECK_NCCL(ncclCommWindowRegister(g_ep_comm, ptr, bytes, &win, NCCL_WIN_COLL_SYMMETRIC)); } }; @@ -666,34 +713,34 @@ class EPZeroCopyTest : public EpOpTestBase {}; // vs HBM reference (same routing, same input). TEST_F(EPZeroCopyTest, IdentityAllSymm) { // HBM reference run. - EPBuffers ref_buf; + EPBuffers<> ref_buf; ref_buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, ep_size_, max_tokens_per_rank_); upload_inputs(ref_buf); - EPTensors ref_t(ref_buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + EPTensors<> ref_t(ref_buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); + NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); uint64_t ref_hid = ref_buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{ref_hid, ref_t.handle_mem.tensor}, ref_t.topk_idx.tensor, ref_t.token_counts.tensor, /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{ref_hid, ref_t.handle_mem.tensor}, ref_t.topk_idx.tensor, - ref_t.tokens.tensor, NVTECommWindow{}, ref_t.topk_weights.tensor, - NVTECommWindow{}, ref_t.recv_tokens.tensor, NVTECommWindow{}, - ref_t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{ref_hid, ref_t.handle_mem.tensor}, ref_t.recv_tokens.tensor, NVTECommWindow{}, - ref_t.result.tensor, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{ref_hid, ref_t.handle_mem.data()}, ref_t.topk_idx.data(), ref_t.token_counts.data(), /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{ref_hid, ref_t.handle_mem.data()}, ref_t.topk_idx.data(), + ref_t.tokens.data(), NVTECommWindow{}, ref_t.topk_weights.data(), + NVTECommWindow{}, ref_t.recv_tokens.data(), NVTECommWindow{}, + ref_t.recv_topk_weights.data(), NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{ref_hid, ref_t.handle_mem.data()}, ref_t.recv_tokens.data(), NVTECommWindow{}, + ref_t.result.data(), stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); std::vector ref_recv(ref_buf.recv_capacity * hidden_dim_); std::vector ref_result(num_tokens_ * hidden_dim_); - CHECK_CUDA(cudaMemcpy(ref_recv.data(), ref_buf.recv_tokens.get(), + NVTE_CHECK_CUDA(cudaMemcpy(ref_recv.data(), ref_buf.recv_tokens.get(), ref_recv.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); - CHECK_CUDA(cudaMemcpy(ref_result.data(), ref_buf.result.get(), + NVTE_CHECK_CUDA(cudaMemcpy(ref_result.data(), ref_buf.result.get(), ref_result.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); // Symm-mem run: tokens, recv_tokens, combine_input (== recv_tokens) all symm. - EPBuffers sym_buf; // alloc all buffers except the symm ones. + EPBuffers<> sym_buf; // alloc all buffers except the symm ones. sym_buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, ep_size_, max_tokens_per_rank_); upload_inputs(sym_buf); @@ -704,32 +751,32 @@ TEST_F(EPZeroCopyTest, IdentityAllSymm) { // Stage same tokens into the symm-mem input. auto h_tok = generate_tokens(g_process_id, num_tokens_, hidden_dim_); - CHECK_CUDA(cudaMemcpy(sym_tokens.ptr, h_tok.data(), + NVTE_CHECK_CUDA(cudaMemcpy(sym_tokens.ptr, h_tok.data(), h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); - EPTensors sym_t(sym_buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + EPTensors<> sym_t(sym_buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); // Replace the tokens/recv_tokens views with ones pointing at the symm buffers. - sym_t.tokens = make_nvte_tensor(sym_tokens.ptr, - {(size_t)num_tokens_, (size_t)hidden_dim_}, kNVTEBFloat16); - sym_t.recv_tokens = make_nvte_tensor(sym_recv.ptr, - {sym_buf.recv_capacity, (size_t)hidden_dim_}, kNVTEBFloat16); + sym_t.tokens = TensorWrapper(sym_tokens.ptr, + std::vector{(size_t)num_tokens_, (size_t)hidden_dim_}, DType::kBFloat16); + sym_t.recv_tokens = TensorWrapper(sym_recv.ptr, + std::vector{sym_buf.recv_capacity, (size_t)hidden_dim_}, DType::kBFloat16); uint64_t sym_hid = sym_buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{sym_hid, sym_t.handle_mem.tensor}, sym_t.topk_idx.tensor, sym_t.token_counts.tensor, /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{sym_hid, sym_t.handle_mem.tensor}, sym_t.topk_idx.tensor, - sym_t.tokens.tensor, symm_window(sym_tokens), - sym_t.topk_weights.tensor, NVTECommWindow{}, - sym_t.recv_tokens.tensor, symm_window(sym_recv), - sym_t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{sym_hid, sym_t.handle_mem.tensor}, sym_t.recv_tokens.tensor, - symm_window(sym_recv), sym_t.result.tensor, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); + ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{sym_hid, sym_t.handle_mem.data()}, sym_t.topk_idx.data(), sym_t.token_counts.data(), /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{sym_hid, sym_t.handle_mem.data()}, sym_t.topk_idx.data(), + sym_t.tokens.data(), symm_window(sym_tokens), + sym_t.topk_weights.data(), NVTECommWindow{}, + sym_t.recv_tokens.data(), symm_window(sym_recv), + sym_t.recv_topk_weights.data(), NVTECommWindow{}, stream)); + ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{sym_hid, sym_t.handle_mem.data()}, sym_t.recv_tokens.data(), + symm_window(sym_recv), sym_t.result.data(), stream)); + NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); std::vector sym_recv_host(sym_buf.recv_capacity * hidden_dim_); std::vector sym_result(num_tokens_ * hidden_dim_); - CHECK_CUDA(cudaMemcpy(sym_recv_host.data(), sym_recv.ptr, + NVTE_CHECK_CUDA(cudaMemcpy(sym_recv_host.data(), sym_recv.ptr, sym_recv_host.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); - CHECK_CUDA(cudaMemcpy(sym_result.data(), sym_buf.result.get(), + NVTE_CHECK_CUDA(cudaMemcpy(sym_result.data(), sym_buf.result.get(), sym_result.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); // Compare per filled recv slot (HBM ref vs symm) and full result. @@ -744,141 +791,9 @@ TEST_F(EPZeroCopyTest, IdentityAllSymm) { if (g_process_id == 0) printf(" IdentityAllSymm: passed (recv_slots=%d, bit-exact vs HBM)\n", total_recv); - CHECK_CUDA(cudaStreamDestroy(stream)); -} - -// Same buffers, 2 iterations — catches window-lifecycle regressions where the -// symm-mem registration goes stale between calls. -TEST_F(EPZeroCopyTest, IdentityAllSymmRepeated) { - EPBuffers buf; - buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, - ep_size_, max_tokens_per_rank_); - upload_inputs(buf); - - SymmBuf sym_tokens, sym_recv; - sym_tokens.alloc(num_tokens_ * hidden_dim_ * sizeof(nv_bfloat16)); - sym_recv .alloc(buf.recv_capacity * hidden_dim_ * sizeof(nv_bfloat16)); - auto h_tok = generate_tokens(g_process_id, num_tokens_, hidden_dim_); - CHECK_CUDA(cudaMemcpy(sym_tokens.ptr, h_tok.data(), - h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); - - EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); - t.tokens = make_nvte_tensor(sym_tokens.ptr, - {(size_t)num_tokens_, (size_t)hidden_dim_}, kNVTEBFloat16); - t.recv_tokens = make_nvte_tensor(sym_recv.ptr, - {buf.recv_capacity, (size_t)hidden_dim_}, kNVTEBFloat16); - - cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); - - uint64_t handle_id = buf.handle_id; - for (int iter = 0; iter < 2; ++iter) { - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, - t.tokens.tensor, symm_window(sym_tokens), - t.topk_weights.tensor, NVTECommWindow{}, - t.recv_tokens.tensor, symm_window(sym_recv), - t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.recv_tokens.tensor, - symm_window(sym_recv), t.result.tensor, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); - - std::vector h_res(num_tokens_ * hidden_dim_); - CHECK_CUDA(cudaMemcpy(h_res.data(), buf.result.get(), - h_res.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); - for (int tok = 0; tok < num_tokens_; ++tok) { - float exp = __bfloat162float(h_tok[tok * hidden_dim_]) * static_cast(top_k_); - float got = __bfloat162float(h_res[tok * hidden_dim_]); - ASSERT_NEAR(got, exp, bf16_tol(exp)) << "iter " << iter << " tok " << tok; - } - } - - if (g_process_id == 0) - printf(" IdentityAllSymmRepeated: passed (2 iters)\n"); - - CHECK_CUDA(cudaStreamDestroy(stream)); + NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } -// Full forward+backward with symm-mem on every spec-mandated buffer: -// dispatch i/o, combine input, combine_bwd i/o, dispatch_bwd input. -// TODO: flaky on rank 0 (grad_tokens partial-zero) when run after the prior -// EPZeroCopyTest cases in the same binary; passes in isolation. Re-enable once -// the root cause (likely NCCL EP NVLS write→read coherence on grad_expert) is -// understood. Tracked separately. -TEST_F(EPZeroCopyTest, DISABLED_FullPipelineSymm) { - EPBuffers buf; - buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, - ep_size_, max_tokens_per_rank_); - upload_inputs(buf); - - // Symm-mem: tokens (dispatch input), recv_tokens (dispatch output AND - // combine input), grad_result (combine_bwd input), grad_expert - // (combine_bwd output AND dispatch_bwd input). - SymmBuf sym_tokens, sym_recv, sym_grad_result, sym_grad_expert; - sym_tokens .alloc(num_tokens_ * hidden_dim_ * sizeof(nv_bfloat16)); - sym_recv .alloc(buf.recv_capacity * hidden_dim_ * sizeof(nv_bfloat16)); - sym_grad_result.alloc(num_tokens_ * hidden_dim_ * sizeof(nv_bfloat16)); - sym_grad_expert.alloc(buf.recv_capacity * hidden_dim_ * sizeof(nv_bfloat16)); - - auto h_tok = generate_tokens(g_process_id, num_tokens_, hidden_dim_); - CHECK_CUDA(cudaMemcpy(sym_tokens.ptr, h_tok.data(), - h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); - - EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); - t.tokens = make_nvte_tensor(sym_tokens.ptr, - {(size_t)num_tokens_, (size_t)hidden_dim_}, kNVTEBFloat16); - t.recv_tokens = make_nvte_tensor(sym_recv.ptr, - {buf.recv_capacity, (size_t)hidden_dim_}, kNVTEBFloat16); - t.grad_result = make_nvte_tensor(sym_grad_result.ptr, - {(size_t)num_tokens_, (size_t)hidden_dim_}, kNVTEBFloat16); - t.grad_expert = make_nvte_tensor(sym_grad_expert.ptr, - {buf.recv_capacity, (size_t)hidden_dim_}, kNVTEBFloat16); - - cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); - - uint64_t handle_id = buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, t.token_counts.tensor, /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.topk_idx.tensor, - t.tokens.tensor, symm_window(sym_tokens), - t.topk_weights.tensor, NVTECommWindow{}, - t.recv_tokens.tensor, symm_window(sym_recv), - t.recv_topk_weights.tensor, NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.recv_tokens.tensor, - symm_window(sym_recv), t.result.tensor, stream)); - - std::vector h_grad(num_tokens_ * hidden_dim_, __float2bfloat16(0.1f)); - CHECK_CUDA(cudaMemcpyAsync(sym_grad_result.ptr, h_grad.data(), - h_grad.size() * sizeof(nv_bfloat16), - cudaMemcpyHostToDevice, stream)); - CHECK_CUDA(cudaMemsetAsync(sym_grad_expert.ptr, 0, sym_grad_expert.bytes, stream)); - CHECK_CUDA(cudaMemsetAsync(buf.g_recv_topk_weights.get(), 0, buf.g_recv_topk_weights.bytes(), stream)); - CHECK_CUDA(cudaMemsetAsync(buf.grad_topk_weights.get(), 0, buf.grad_topk_weights.bytes(), stream)); - - ASSERT_NO_THROW(nvte_ep_combine_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_result.tensor, - symm_window(sym_grad_result), t.grad_expert.tensor, - symm_window(sym_grad_expert), stream)); - ASSERT_NO_THROW(nvte_ep_dispatch_bwd(NVTEEpHandle{handle_id, t.handle_mem.tensor}, t.grad_expert.tensor, - symm_window(sym_grad_expert), - t.g_recv_topk_weights.tensor, NVTECommWindow{}, - t.grad_tokens.tensor, t.grad_topk_weights.tensor, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); - - ASSERT_TRUE(check_no_nan_inf(buf.result.get(), num_tokens_ * hidden_dim_, "result")); - ASSERT_TRUE(check_no_nan_inf(buf.grad_tokens.get(), num_tokens_ * hidden_dim_, "grad_tokens")); - - std::vector h_gt(num_tokens_ * hidden_dim_); - CHECK_CUDA(cudaMemcpy(h_gt.data(), buf.grad_tokens.get(), - h_gt.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); - const float kExpGrad = static_cast(top_k_) * 0.1f; - for (int tok = 0; tok < num_tokens_; ++tok) - EXPECT_NEAR(__bfloat162float(h_gt[tok * hidden_dim_]), kExpGrad, bf16_tol(kExpGrad)) - << "grad_tokens token " << tok; - - if (g_process_id == 0) printf(" FullPipelineSymm: passed\n"); - - CHECK_CUDA(cudaStreamDestroy(stream)); -} // ── main ────────────────────────────────────────────────────────────────────── diff --git a/tests/cpp_distributed/test_ep_common.h b/tests/cpp_distributed/test_ep_common.h index ccb20ee3a0..135a39416e 100644 --- a/tests/cpp_distributed/test_ep_common.h +++ b/tests/cpp_distributed/test_ep_common.h @@ -13,157 +13,67 @@ #include #include +#include #include #include +#include #include -#include #include #include #include #include -#include #include #include #include #include +#include "../cpp/test_common.h" +#include "util/logging.h" -// ── Error-checking macros ───────────────────────────────────────────────────── +using transformer_engine::DType; +using transformer_engine::TensorWrapper; -#define CHECK_NCCL(expr) \ - do { \ - ncclResult_t _err = (expr); \ - if (_err != ncclSuccess) \ - FAIL() << "NCCL error " << _err << ": " << ncclGetErrorString(_err); \ - } while (false) - -#define CHECK_CUDA(expr) \ - do { \ - cudaError_t _err = (expr); \ - if (_err != cudaSuccess) \ - FAIL() << "CUDA error " << _err << ": " << cudaGetErrorString(_err); \ - } while (false) - -#define ASSERT_CUDA_OK(expr) \ - do { \ - cudaError_t _err = (expr); \ - if (_err != cudaSuccess) { \ - fprintf(stderr, "CUDA error %d: %s\n", _err, cudaGetErrorString(_err)); \ - exit(EXIT_FAILURE); \ - } \ - } while (false) - -#define ASSERT_NCCL_OK(expr) \ - do { \ - ncclResult_t _err = (expr); \ - if (_err != ncclSuccess) { \ - fprintf(stderr, "NCCL error %d: %s\n", _err, ncclGetErrorString(_err)); \ - exit(EXIT_FAILURE); \ - } \ +#define CHECK_MPI(expr) \ + do { \ + int _err_mpi = (expr); \ + NVTE_CHECK(_err_mpi == MPI_SUCCESS, "MPI error: ", _err_mpi); \ } while (false) // ── Process-level state ─────────────────────────────────────────────────────── static int g_process_id = -1; static int g_num_processes = -1; -static std::string g_uid_file; static int g_sm_major = -1; // set by ep_bootstrap; -1 until then static int g_ep_size = -1; static int g_num_experts = -1; static int g_hidden_dim = 256; static int g_max_tokens_per_rank = 64; -static NVTEDType g_token_dtype = kNVTEBFloat16; +static NVTEDType g_max_token_dtype = kNVTEFloat32; // staging-buffer sizing static bool g_ep_initialized = false; static ncclComm_t g_ep_comm = nullptr; // owned by harness, destroyed in ep_teardown -// ── TensorHandle RAII wrapper ───────────────────────────────────────────────── - -// View over a caller-owned device buffer; owns NVTETensor metadata only. Move-only. -struct TensorHandle { - NVTETensor tensor = nullptr; - void* dev_ptr = nullptr; - - ~TensorHandle() { - if (tensor) nvte_destroy_tensor(tensor); - } - - TensorHandle() = default; - TensorHandle(const TensorHandle&) = delete; - TensorHandle& operator=(const TensorHandle&) = delete; - - TensorHandle(TensorHandle&& o) noexcept : tensor(o.tensor), dev_ptr(o.dev_ptr) { - o.tensor = nullptr; o.dev_ptr = nullptr; - } - TensorHandle& operator=(TensorHandle&& o) noexcept { - if (this != &o) { - if (tensor) nvte_destroy_tensor(tensor); - tensor = o.tensor; dev_ptr = o.dev_ptr; - o.tensor = nullptr; o.dev_ptr = nullptr; - } - return *this; - } -}; - -static TensorHandle make_nvte_tensor(void* dev_ptr, - const std::vector& shape, - NVTEDType dtype) { - TensorHandle h; - h.dev_ptr = dev_ptr; - h.tensor = nvte_create_tensor(NVTE_DELAYED_TENSOR_SCALING); - - NVTEShape s; - s.ndim = shape.size(); - for (size_t i = 0; i < shape.size(); ++i) s.data[i] = shape[i]; - - NVTEBasicTensor bt; - bt.data_ptr = dev_ptr; - bt.dtype = dtype; - bt.shape = s; - nvte_set_tensor_param_v2(h.tensor, kNVTERowwiseData, &bt, sizeof(bt)); - - return h; -} - -// RAII owner for a cudaMalloc'd device buffer; frees on destruction. +// RAII owner for a cudaMalloc'd device buffer; element-count API on top of +// test::CudaPtr. template struct DevBuf { - T* ptr = nullptr; + test::CudaPtr ptr; size_t count = 0; DevBuf() = default; explicit DevBuf(size_t n) { alloc(n); } - ~DevBuf() { reset(); } - - DevBuf(const DevBuf&) = delete; - DevBuf& operator=(const DevBuf&) = delete; - DevBuf(DevBuf&& o) noexcept : ptr(o.ptr), count(o.count) { o.ptr = nullptr; o.count = 0; } - DevBuf& operator=(DevBuf&& o) noexcept { - if (this != &o) { reset(); ptr = o.ptr; count = o.count; o.ptr = nullptr; o.count = 0; } - return *this; - } void alloc(size_t n) { - reset(); count = n; - if (n > 0) { - cudaError_t e = cudaMalloc(&ptr, n * sizeof(T)); - if (e != cudaSuccess) { - fprintf(stderr, "DevBuf cudaMalloc(%zu) failed: %s\n", n * sizeof(T), - cudaGetErrorString(e)); - ptr = nullptr; - count = 0; - } - } + ptr = (n > 0) ? test::cuda_alloc(n * sizeof(T)) : test::CudaPtr{}; } - void reset() { - if (ptr) { cudaFree(ptr); ptr = nullptr; } + ptr.reset(); count = 0; } - T* get() const { return ptr; } + T* get() const { return ptr.get(); } size_t bytes() const { return count * sizeof(T); } }; @@ -180,39 +90,11 @@ static inline std::vector routing_balanced( return idx; } -// ── File-based ncclUniqueId exchange ───────────────────────────────────────── +// ── ncclUniqueId exchange via MPI ───────────────────────────────────────────── static void exchange_unique_id(ncclUniqueId* uid) { - const size_t sz = sizeof(ncclUniqueId); - - if (g_process_id == 0) { - ASSERT_NCCL_OK(ncclGetUniqueId(uid)); - FILE* f = fopen(g_uid_file.c_str(), "wb"); - if (!f) { fprintf(stderr, "Cannot open uid file: %s\n", g_uid_file.c_str()); exit(EXIT_FAILURE); } - fwrite(uid, 1, sz, f); - fclose(f); - } else { - auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(60); - while (true) { - FILE* f = fopen(g_uid_file.c_str(), "rb"); - if (f) { - fseek(f, 0, SEEK_END); - if (static_cast(ftell(f)) >= sz) { - fseek(f, 0, SEEK_SET); - size_t n = fread(uid, 1, sz, f); - fclose(f); - if (n == sz) break; - } else { - fclose(f); - } - } - if (std::chrono::steady_clock::now() > deadline) { - fprintf(stderr, "Process %d: timed out waiting for uid file\n", g_process_id); - exit(EXIT_FAILURE); - } - std::this_thread::sleep_for(std::chrono::milliseconds(50)); - } - } + if (g_process_id == 0) NVTE_CHECK_NCCL(ncclGetUniqueId(uid)); + CHECK_MPI(MPI_Bcast(uid, sizeof(*uid), MPI_BYTE, 0, MPI_COMM_WORLD)); } // ── CLI parsing ─────────────────────────────────────────────────────────────── @@ -220,26 +102,8 @@ static void exchange_unique_id(ncclUniqueId* uid) { static void ep_parse_args(int argc, char* argv[]) { for (int i = 1; i < argc; ++i) { std::string a(argv[i]); - if (a.rfind("--process-id=", 0) == 0) g_process_id = std::stoi(a.substr(13)); - else if (a.rfind("--rank=", 0) == 0) g_process_id = std::stoi(a.substr(7)); - else if (a.rfind("--num-processes=",0)==0) g_num_processes = std::stoi(a.substr(16)); - else if (a.rfind("--nranks=", 0) == 0) g_num_processes = std::stoi(a.substr(9)); - else if (a.rfind("--uid-file=", 0) == 0) g_uid_file = a.substr(11); - else if (a.rfind("--token-dtype=", 0) == 0) - g_token_dtype = static_cast(std::stoi(a.substr(14))); - } - - if (g_process_id < 0 || g_num_processes <= 0) { - fprintf(stderr, - "Usage: %s --rank=N --nranks=N [--uid-file=path] [gtest flags]\n" - " Aliases: --process-id=N, --num-processes=N\n", - argc > 0 ? argv[0] : "test_ep"); - exit(EXIT_FAILURE); - } - - if (g_uid_file.empty()) { - const char* t = getenv("TMPDIR"); if (!t) t = "/tmp"; - g_uid_file = std::string(t) + "/te_ep_uid_" + std::to_string(g_process_id); + if (a.rfind("--max-token-dtype=", 0) == 0) + g_max_token_dtype = static_cast(std::stoi(a.substr(18))); } } @@ -247,6 +111,12 @@ static void ep_parse_args(int argc, char* argv[]) { // Returns false if the binary should exit without running tests (wrong SM, etc.). static bool ep_bootstrap(int argc, char* argv[]) { + int mpi_initialized = 0; + MPI_Initialized(&mpi_initialized); + if (!mpi_initialized) CHECK_MPI(MPI_Init(&argc, &argv)); + CHECK_MPI(MPI_Comm_rank(MPI_COMM_WORLD, &g_process_id)); + CHECK_MPI(MPI_Comm_size(MPI_COMM_WORLD, &g_num_processes)); + ep_parse_args(argc, argv); ::testing::InitGoogleTest(&argc, argv); @@ -282,9 +152,9 @@ static bool ep_bootstrap(int argc, char* argv[]) { // Worst-case for top_k fan-out: ep_size * max_tokens_per_rank * 2. group_config.max_recv_tokens_per_rank = g_ep_size * g_max_tokens_per_rank * 2; group_config.hidden_dim = g_hidden_dim; - group_config.token_dtype = g_token_dtype; + group_config.max_token_dtype = g_max_token_dtype; - ASSERT_NCCL_OK(ncclCommInitRank(&g_ep_comm, g_num_processes, uid, g_process_id)); + NVTE_CHECK_NCCL(ncclCommInitRank(&g_ep_comm, g_num_processes, uid, g_process_id)); nvte_ep_initialize(static_cast(g_ep_comm), group_config); if (g_process_id == 0) { @@ -308,5 +178,7 @@ static void ep_teardown() { } g_ep_initialized = false; } - if (g_process_id == 0) remove(g_uid_file.c_str()); + int finalized = 0; + MPI_Finalized(&finalized); + if (!finalized) MPI_Finalize(); } diff --git a/tests/cpp_distributed/test_ep_coverage.cu b/tests/cpp_distributed/test_ep_coverage.cu deleted file mode 100644 index e9e532386c..0000000000 --- a/tests/cpp_distributed/test_ep_coverage.cu +++ /dev/null @@ -1,562 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -/* - * EP C-API coverage tests (paths not exercised by the pipeline suite). - * - * MultiHandleAllocTest — distinct handle ids; each works end-to-end. - * TopK1Test — top_k=1 dispatch/combine/bwd round-trip. - * EmptyExpertsTest — alignment ∈ {0, 2, 8, 16} with experts receiving 0 tokens. - * NegativeTests — alignment mismatch and null handle_mem must throw. - */ - -#include "test_ep_common.h" - -#include -#include - -// top1 -> expert 0, top2 -> expert 2; leaves local-expert 1 empty between two -// full experts. Requires top_k >= 2 and num_experts >= 3. -static std::vector routing_skip_middle(int num_tokens, int top_k) { - std::vector idx(num_tokens * top_k); - for (int t = 0; t < num_tokens; ++t) { - idx[t * top_k + 0] = 0; - if (top_k >= 2) idx[t * top_k + 1] = 2; - for (int k = 2; k < top_k; ++k) idx[t * top_k + k] = 2 + k; // distinct stragglers - } - return idx; -} - -static std::vector tokens_constant(int num_tokens, int hidden_dim, float val) { - std::vector v(num_tokens * hidden_dim); - nv_bfloat16 b = __float2bfloat16(val); - std::fill(v.begin(), v.end(), b); - return v; -} - -namespace { - -class EpCoverageBase : public ::testing::Test { - protected: - int ep_size_, num_experts_, num_local_experts_, hidden_dim_; - int max_tokens_per_rank_; - - void SetUp() override { - if (g_sm_major < 9) - GTEST_SKIP() << "EP requires SM_90+ (device is SM_" << g_sm_major << "0)"; - ASSERT_GE(g_num_processes, 2); - ASSERT_TRUE(g_ep_initialized); - ep_size_ = g_ep_size; - num_experts_ = g_num_experts; - num_local_experts_ = num_experts_ / ep_size_; - hidden_dim_ = g_hidden_dim; - max_tokens_per_rank_ = g_max_tokens_per_rank; - } - - // Helper: allocate buffers + tensor views for a single dispatch+combine. - struct Bundle { - DevBuf topk_idx; - DevBuf topk_weights; - DevBuf tokens; - DevBuf token_counts; - DevBuf handle_mem; - DevBuf recv_tokens; - DevBuf recv_topk_weights; - DevBuf result; - uint64_t handle_id = 0; - size_t handle_mem_size = 0; - size_t recv_capacity = 0; - }; - - Bundle make_bundle(int num_tokens, int top_k, int num_local_experts, - size_t alignment) { - Bundle b; - b.recv_capacity = static_cast(ep_size_) * max_tokens_per_rank_ * 2; - b.topk_idx.alloc(num_tokens * top_k); - b.topk_weights.alloc(num_tokens * top_k); - b.tokens.alloc(num_tokens * hidden_dim_); - b.token_counts.alloc(num_local_experts); - b.recv_tokens.alloc(b.recv_capacity * hidden_dim_); - b.recv_topk_weights.alloc(b.recv_capacity); - b.result.alloc(num_tokens * hidden_dim_); - NVTEEpLayerConfig cfg{num_local_experts, top_k, alignment}; - b.handle_id = nvte_ep_register_layer(cfg, &b.handle_mem_size); - b.handle_mem.alloc(b.handle_mem_size); - return b; - } -}; - -} // namespace - -// ============================================================================= -// MultiHandleAllocTest: ids are distinct and each is independently usable. -// ============================================================================= - -class MultiHandleAllocTest : public EpCoverageBase {}; - -TEST_F(MultiHandleAllocTest, IdsAreDistinct) { - NVTEEpLayerConfig cfg{num_local_experts_, /*top_k=*/2, /*alignment=*/0}; - const int kN = 8; - std::vector ids(kN); - for (int i = 0; i < kN; ++i) { - size_t sz = 0; - ids[i] = nvte_ep_register_layer(cfg, &sz); - } - for (int i = 0; i < kN; ++i) { - EXPECT_NE(ids[i], 0u) << "handle_id 0 is reserved as \"no id\""; - for (int j = i + 1; j < kN; ++j) - EXPECT_NE(ids[i], ids[j]) << "duplicate id " << ids[i] << " at indices " << i << ", " << j; - } -} - -TEST_F(MultiHandleAllocTest, TwoHandlesCoexist) { - const int num_tokens = 16, top_k = 2; - Bundle a = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); - Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); - - auto h_idx = routing_balanced(g_process_id, num_tokens, top_k, - num_experts_, num_local_experts_); - std::vector h_w(num_tokens * top_k, 1.0f / top_k); - auto h_tok = tokens_constant(num_tokens, hidden_dim_, 0.5f); - for (Bundle* x : {&a, &b}) { - CHECK_CUDA(cudaMemcpy(x->topk_idx.get(), h_idx.data(), - h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(x->topk_weights.get(), h_w.data(), - h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(x->tokens.get(), h_tok.data(), - h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); - } - - cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); - - ASSERT_NE(a.handle_id, b.handle_id); - - auto run_one = [&](Bundle& x) { - auto topk_idx = make_nvte_tensor(x.topk_idx.get(), {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); - auto topk_weights = make_nvte_tensor(x.topk_weights.get(), {(size_t)num_tokens, (size_t)top_k}, kNVTEFloat32); - auto token_counts = make_nvte_tensor(x.token_counts.get(), {(size_t)num_local_experts_}, kNVTEInt32); - auto handle_mem = make_nvte_tensor(x.handle_mem.get(), {x.handle_mem_size}, kNVTEByte); - auto tokens = make_nvte_tensor(x.tokens.get(), {(size_t)num_tokens, (size_t)hidden_dim_}, kNVTEBFloat16); - auto recv_tokens = make_nvte_tensor(x.recv_tokens.get(), {x.recv_capacity, (size_t)hidden_dim_}, kNVTEBFloat16); - auto recv_w = make_nvte_tensor(x.recv_topk_weights.get(), {x.recv_capacity}, kNVTEFloat32); - auto result = make_nvte_tensor(x.result.get(), {(size_t)num_tokens, (size_t)hidden_dim_}, kNVTEBFloat16); - NVTEEpHandle h{x.handle_id, handle_mem.tensor}; - ASSERT_NO_THROW(nvte_ep_prepare(h, topk_idx.tensor, token_counts.tensor, - /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(h, topk_idx.tensor, tokens.tensor, - NVTECommWindow{}, topk_weights.tensor, NVTECommWindow{}, - recv_tokens.tensor, NVTECommWindow{}, recv_w.tensor, - NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(h, recv_tokens.tensor, NVTECommWindow{}, - result.tensor, stream)); - }; - run_one(a); - run_one(b); - CHECK_CUDA(cudaStreamSynchronize(stream)); - - // Both round-trips must produce result == top_k * 0.5 = 1.0. - for (Bundle* x : {&a, &b}) { - std::vector h_res(num_tokens * hidden_dim_); - CHECK_CUDA(cudaMemcpy(h_res.data(), x->result.get(), - h_res.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); - const int probes[3] = {0, hidden_dim_ / 2, hidden_dim_ - 1}; - for (int t = 0; t < num_tokens; ++t) - for (int p : probes) - EXPECT_NEAR(__bfloat162float(h_res[t * hidden_dim_ + p]), - static_cast(top_k) * 0.5f, 1e-2f); - } - CHECK_CUDA(cudaStreamDestroy(stream)); -} - -// ============================================================================= -// TopK1Test: top_k=1 dispatch/combine round-trip, including dispatch_bwd. -// ============================================================================= - -class TopK1Test : public EpCoverageBase {}; - -TEST_F(TopK1Test, RoundTrip) { - const int num_tokens = 16, top_k = 1; - Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); - - auto h_idx = routing_balanced(g_process_id, num_tokens, top_k, - num_experts_, num_local_experts_); - std::vector h_w(num_tokens * top_k, 1.0f); // top_k=1: weight is unity - auto h_tok = tokens_constant(num_tokens, hidden_dim_, 0.25f); - CHECK_CUDA(cudaMemcpy(b.topk_idx.get(), h_idx.data(), - h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(b.topk_weights.get(), h_w.data(), - h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(b.tokens.get(), h_tok.data(), - h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); - - auto topk_idx_t = make_nvte_tensor(b.topk_idx.get(), - {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); - auto topk_weights_t = make_nvte_tensor(b.topk_weights.get(), - {(size_t)num_tokens, (size_t)top_k}, kNVTEFloat32); - auto token_counts_t = make_nvte_tensor(b.token_counts.get(), - {(size_t)num_local_experts_}, kNVTEInt32); - auto handle_mem_t = make_nvte_tensor(b.handle_mem.get(), - {b.handle_mem_size}, kNVTEByte); - auto tokens_t = make_nvte_tensor(b.tokens.get(), - {(size_t)num_tokens, (size_t)hidden_dim_}, kNVTEBFloat16); - auto recv_tokens_t = make_nvte_tensor(b.recv_tokens.get(), - {b.recv_capacity, (size_t)hidden_dim_}, kNVTEBFloat16); - auto recv_w_t = make_nvte_tensor(b.recv_topk_weights.get(), - {b.recv_capacity}, kNVTEFloat32); - auto result_t = make_nvte_tensor(b.result.get(), - {(size_t)num_tokens, (size_t)hidden_dim_}, kNVTEBFloat16); - - cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); - - NVTEEpHandle h{b.handle_id, handle_mem_t.tensor}; - ASSERT_NO_THROW(nvte_ep_prepare(h, topk_idx_t.tensor, token_counts_t.tensor, - /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(h, topk_idx_t.tensor, - tokens_t.tensor, NVTECommWindow{}, topk_weights_t.tensor, - NVTECommWindow{}, recv_tokens_t.tensor, NVTECommWindow{}, - recv_w_t.tensor, NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(h, recv_tokens_t.tensor, - NVTECommWindow{}, result_t.tensor, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); - - // top_k=1: combine is unweighted gather, so result[t] == tokens[t]. - std::vector h_res(num_tokens * hidden_dim_); - CHECK_CUDA(cudaMemcpy(h_res.data(), b.result.get(), - h_res.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); - const int probes[3] = {0, hidden_dim_ / 2, hidden_dim_ - 1}; - for (int t = 0; t < num_tokens; ++t) - for (int p : probes) - EXPECT_NEAR(__bfloat162float(h_res[t * hidden_dim_ + p]), 0.25f, 1e-2f) - << "tok " << t << " hidden " << p; - - CHECK_CUDA(cudaStreamDestroy(stream)); -} - -// ============================================================================= -// EmptyExpertsTest: alignment ∈ {0, 2, 8, 16}, only local-expert 0 receives -// tokens. Round-trip must produce result == top_k * tokens regardless of the -// per-expert padding choice. -// ============================================================================= - -class EmptyExpertsTest : public EpCoverageBase, - public ::testing::WithParamInterface {}; - -TEST_P(EmptyExpertsTest, RoundTripCorrect) { - // routing_skip_middle needs experts {0, 2, ...}; smallest viable num_experts is 3. - ASSERT_GE(num_experts_, 3); - const size_t alignment = GetParam(); - const int num_tokens = 16, top_k = 2; - Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, alignment); - - // top1 -> expert 0, top2 -> expert 2; rank 0's local-expert 1 receives 0 - // tokens between two non-empty experts. - std::vector h_idx = routing_skip_middle(num_tokens, top_k); - std::vector h_w(num_tokens * top_k, 1.0f / top_k); - auto h_tok = tokens_constant(num_tokens, hidden_dim_, 0.3f); - - CHECK_CUDA(cudaMemcpy(b.topk_idx.get(), h_idx.data(), - h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(b.topk_weights.get(), h_w.data(), - h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(b.tokens.get(), h_tok.data(), - h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); - - auto topk_idx_t = make_nvte_tensor(b.topk_idx.get(), - {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); - auto topk_weights_t = make_nvte_tensor(b.topk_weights.get(), - {(size_t)num_tokens, (size_t)top_k}, kNVTEFloat32); - auto token_counts_t = make_nvte_tensor(b.token_counts.get(), - {(size_t)num_local_experts_}, kNVTEInt32); - auto handle_mem_t = make_nvte_tensor(b.handle_mem.get(), - {b.handle_mem_size}, kNVTEByte); - auto tokens_t = make_nvte_tensor(b.tokens.get(), - {(size_t)num_tokens, (size_t)hidden_dim_}, kNVTEBFloat16); - auto recv_tokens_t = make_nvte_tensor(b.recv_tokens.get(), - {b.recv_capacity, (size_t)hidden_dim_}, kNVTEBFloat16); - auto recv_w_t = make_nvte_tensor(b.recv_topk_weights.get(), - {b.recv_capacity}, kNVTEFloat32); - auto result_t = make_nvte_tensor(b.result.get(), - {(size_t)num_tokens, (size_t)hidden_dim_}, kNVTEBFloat16); - - cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); - - NVTEEpHandle h{b.handle_id, handle_mem_t.tensor}; - ASSERT_NO_THROW(nvte_ep_prepare(h, topk_idx_t.tensor, token_counts_t.tensor, - alignment, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(h, topk_idx_t.tensor, - tokens_t.tensor, NVTECommWindow{}, topk_weights_t.tensor, - NVTECommWindow{}, recv_tokens_t.tensor, NVTECommWindow{}, - recv_w_t.tensor, NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(h, recv_tokens_t.tensor, - NVTECommWindow{}, result_t.tensor, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); - - // Identity expert + uniform weights: result[t] == top_k * tokens[t]. - std::vector h_res(num_tokens * hidden_dim_); - CHECK_CUDA(cudaMemcpy(h_res.data(), b.result.get(), - h_res.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); - const float expected = static_cast(top_k) * 0.3f; - const int probes[3] = {0, hidden_dim_ / 2, hidden_dim_ - 1}; - for (int t = 0; t < num_tokens; ++t) - for (int p : probes) - EXPECT_NEAR(__bfloat162float(h_res[t * hidden_dim_ + p]), expected, 1e-2f) - << "alignment=" << alignment << " tok=" << t << " hidden=" << p; - - CHECK_CUDA(cudaStreamDestroy(stream)); -} - -INSTANTIATE_TEST_SUITE_P(Alignments, EmptyExpertsTest, - ::testing::Values(0, 2, 8, 16)); - -// ============================================================================= -// NegativeTests: prepare/dispatch must surface bad inputs as exceptions. -// ============================================================================= - -class NegativeTests : public EpCoverageBase {}; - -TEST_F(NegativeTests, AlignmentMismatchThrows) { - const int num_tokens = 8, top_k = 2; - // Allocate handle for alignment=0, then call prepare with alignment=16. - Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); - auto h_idx = routing_balanced(g_process_id, num_tokens, top_k, - num_experts_, num_local_experts_); - CHECK_CUDA(cudaMemcpy(b.topk_idx.get(), h_idx.data(), - h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); - - auto topk_idx_t = make_nvte_tensor(b.topk_idx.get(), - {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); - auto token_counts_t = make_nvte_tensor(b.token_counts.get(), - {(size_t)num_local_experts_}, kNVTEInt32); - auto handle_mem_t = make_nvte_tensor(b.handle_mem.get(), - {b.handle_mem_size}, kNVTEByte); - - cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); - NVTEEpHandle h{b.handle_id, handle_mem_t.tensor}; - EXPECT_THROW(nvte_ep_prepare(h, topk_idx_t.tensor, token_counts_t.tensor, - /*alignment=*/16, stream), - std::exception); - CHECK_CUDA(cudaStreamDestroy(stream)); -} - -TEST_F(NegativeTests, NullHandleMemThrows) { - const int num_tokens = 8, top_k = 2; - Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); - auto h_idx = routing_balanced(g_process_id, num_tokens, top_k, - num_experts_, num_local_experts_); - CHECK_CUDA(cudaMemcpy(b.topk_idx.get(), h_idx.data(), - h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); - - auto topk_idx_t = make_nvte_tensor(b.topk_idx.get(), - {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); - auto token_counts_t = make_nvte_tensor(b.token_counts.get(), - {(size_t)num_local_experts_}, kNVTEInt32); - // Construct a tensor view backed by a null device pointer. - auto null_hm_t = make_nvte_tensor(nullptr, {b.handle_mem_size}, kNVTEByte); - - cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); - NVTEEpHandle h{b.handle_id, null_hm_t.tensor}; - EXPECT_THROW(nvte_ep_prepare(h, topk_idx_t.tensor, token_counts_t.tensor, - /*alignment=*/0, stream), - std::exception); - CHECK_CUDA(cudaStreamDestroy(stream)); -} - -// ============================================================================= -// HandleCacheTest: persistent ncclEpHandle is reused across ops on the same -// handle_mem ptr; relocation triggers throw by default and rebuild when -// NVTEEpGroupConfig.allow_handle_mem_reloc=1. -// ============================================================================= - -class HandleCacheTest : public EpCoverageBase {}; - -// Run prepare → dispatch → combine on bundle b. handle_mem_data overrides the -// device ptr used for handle_mem (must be the buffer owned by b unless -// reloc-allowed mode is active). Templated on Bundle because EpCoverageBase:: -// Bundle is declared in a protected section. -template -static void run_round_trip(B& b, void* handle_mem_data, - int num_tokens, int top_k, int num_local_experts, - int hidden_dim, size_t alignment, - cudaStream_t stream) { - auto topk_idx_t = make_nvte_tensor(b.topk_idx.get(), - {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); - auto topk_weights_t = make_nvte_tensor(b.topk_weights.get(), - {(size_t)num_tokens, (size_t)top_k}, kNVTEFloat32); - auto token_counts_t = make_nvte_tensor(b.token_counts.get(), - {(size_t)num_local_experts}, kNVTEInt32); - auto handle_mem_t = make_nvte_tensor(handle_mem_data, - {b.handle_mem_size}, kNVTEByte); - auto tokens_t = make_nvte_tensor(b.tokens.get(), - {(size_t)num_tokens, (size_t)hidden_dim}, kNVTEBFloat16); - auto recv_tokens_t = make_nvte_tensor(b.recv_tokens.get(), - {b.recv_capacity, (size_t)hidden_dim}, kNVTEBFloat16); - auto recv_w_t = make_nvte_tensor(b.recv_topk_weights.get(), - {b.recv_capacity}, kNVTEFloat32); - auto result_t = make_nvte_tensor(b.result.get(), - {(size_t)num_tokens, (size_t)hidden_dim}, kNVTEBFloat16); - - NVTEEpHandle h{b.handle_id, handle_mem_t.tensor}; - nvte_ep_prepare(h, topk_idx_t.tensor, token_counts_t.tensor, alignment, stream); - nvte_ep_dispatch(h, topk_idx_t.tensor, tokens_t.tensor, NVTECommWindow{}, - topk_weights_t.tensor, NVTECommWindow{}, - recv_tokens_t.tensor, NVTECommWindow{}, - recv_w_t.tensor, NVTECommWindow{}, stream); - nvte_ep_combine(h, recv_tokens_t.tensor, NVTECommWindow{}, result_t.tensor, stream); -} - -// Re-bootstrap EP backend with a different allow_handle_mem_reloc setting. -// Reuses the existing g_ep_comm; caller is responsible for restoring defaults. -static void reinit_ep_with_reloc(int allow_reloc) { - nvte_ep_shutdown(); - NVTEEpGroupConfig cfg{}; - cfg.ep_size = g_ep_size; - cfg.num_experts = g_num_experts; - cfg.max_tokens_per_rank = g_max_tokens_per_rank; - cfg.max_recv_tokens_per_rank = g_ep_size * g_max_tokens_per_rank * 2; - cfg.hidden_dim = g_hidden_dim; - cfg.allow_handle_mem_reloc = allow_reloc; - nvte_ep_initialize(static_cast(g_ep_comm), cfg); -} - -TEST_F(HandleCacheTest, ReuseSameMemSucceeds) { - const int num_tokens = 16, top_k = 2; - Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); - - auto h_idx = routing_balanced(g_process_id, num_tokens, top_k, - num_experts_, num_local_experts_); - std::vector h_w(num_tokens * top_k, 1.0f / top_k); - auto h_tok = tokens_constant(num_tokens, hidden_dim_, 0.5f); - CHECK_CUDA(cudaMemcpy(b.topk_idx.get(), h_idx.data(), - h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(b.topk_weights.get(), h_w.data(), - h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(b.tokens.get(), h_tok.data(), - h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); - - cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); - - // Two consecutive round-trips on the same handle_mem ptr: first opens the - // cached handle, second hits the cache. Both must succeed and be correct. - for (int iter = 0; iter < 2; ++iter) { - ASSERT_NO_THROW(run_round_trip(b, b.handle_mem.get(), num_tokens, top_k, - num_local_experts_, hidden_dim_, - /*alignment=*/0, stream)); - } - CHECK_CUDA(cudaStreamSynchronize(stream)); - - std::vector h_res(num_tokens * hidden_dim_); - CHECK_CUDA(cudaMemcpy(h_res.data(), b.result.get(), - h_res.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); - const int probes[3] = {0, hidden_dim_ / 2, hidden_dim_ - 1}; - for (int t = 0; t < num_tokens; ++t) - for (int p : probes) - EXPECT_NEAR(__bfloat162float(h_res[t * hidden_dim_ + p]), - static_cast(top_k) * 0.5f, 1e-2f); - - CHECK_CUDA(cudaStreamDestroy(stream)); -} - -TEST_F(HandleCacheTest, RelocDefaultThrows) { - // Default bootstrap has allow_handle_mem_reloc=0: a second prepare call on - // the same handle_id with a different handle_mem ptr must throw. - const int num_tokens = 8, top_k = 2; - Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); - DevBuf second_hm(b.handle_mem_size); // distinct device buffer - ASSERT_NE(b.handle_mem.get(), second_hm.get()); - - auto h_idx = routing_balanced(g_process_id, num_tokens, top_k, - num_experts_, num_local_experts_); - CHECK_CUDA(cudaMemcpy(b.topk_idx.get(), h_idx.data(), - h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); - - auto topk_idx_t = make_nvte_tensor(b.topk_idx.get(), - {(size_t)num_tokens, (size_t)top_k}, kNVTEInt64); - auto token_counts_t = make_nvte_tensor(b.token_counts.get(), - {(size_t)num_local_experts_}, kNVTEInt32); - auto hm1_t = make_nvte_tensor(b.handle_mem.get(), - {b.handle_mem_size}, kNVTEByte); - auto hm2_t = make_nvte_tensor(second_hm.get(), - {b.handle_mem_size}, kNVTEByte); - - cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); - - // First prepare seeds the cache. - NVTEEpHandle h1{b.handle_id, hm1_t.tensor}; - ASSERT_NO_THROW(nvte_ep_prepare(h1, topk_idx_t.tensor, token_counts_t.tensor, - /*alignment=*/0, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); - // Same handle_id with a different handle_mem ptr must throw. - NVTEEpHandle h2{b.handle_id, hm2_t.tensor}; - EXPECT_THROW(nvte_ep_prepare(h2, topk_idx_t.tensor, token_counts_t.tensor, - /*alignment=*/0, stream), - std::exception); - CHECK_CUDA(cudaStreamDestroy(stream)); -} - -TEST_F(HandleCacheTest, RelocAllowedRebuilds) { - // Re-init EP backend with allow_handle_mem_reloc=1, run two round-trips with - // distinct handle_mem buffers, verify both succeed numerically, restore. - reinit_ep_with_reloc(/*allow_reloc=*/1); - - struct Restore { ~Restore() { reinit_ep_with_reloc(/*allow_reloc=*/0); } } restore; - - const int num_tokens = 16, top_k = 2; - Bundle b = make_bundle(num_tokens, top_k, num_local_experts_, /*alignment=*/0); - DevBuf alt_hm(b.handle_mem_size); - ASSERT_NE(b.handle_mem.get(), alt_hm.get()); - - auto h_idx = routing_balanced(g_process_id, num_tokens, top_k, - num_experts_, num_local_experts_); - std::vector h_w(num_tokens * top_k, 1.0f / top_k); - auto h_tok = tokens_constant(num_tokens, hidden_dim_, 0.5f); - CHECK_CUDA(cudaMemcpy(b.topk_idx.get(), h_idx.data(), - h_idx.size() * sizeof(int64_t), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(b.topk_weights.get(), h_w.data(), - h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); - CHECK_CUDA(cudaMemcpy(b.tokens.get(), h_tok.data(), - h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); - - cudaStream_t stream; - CHECK_CUDA(cudaStreamCreate(&stream)); - - // First on the original handle_mem. - ASSERT_NO_THROW(run_round_trip(b, b.handle_mem.get(), num_tokens, top_k, - num_local_experts_, hidden_dim_, - /*alignment=*/0, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); - // Then on the relocated handle_mem — must trigger silent rebuild, not throw. - ASSERT_NO_THROW(run_round_trip(b, alt_hm.get(), num_tokens, top_k, - num_local_experts_, hidden_dim_, - /*alignment=*/0, stream)); - CHECK_CUDA(cudaStreamSynchronize(stream)); - - std::vector h_res(num_tokens * hidden_dim_); - CHECK_CUDA(cudaMemcpy(h_res.data(), b.result.get(), - h_res.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); - const int probes[3] = {0, hidden_dim_ / 2, hidden_dim_ - 1}; - for (int t = 0; t < num_tokens; ++t) - for (int p : probes) - EXPECT_NEAR(__bfloat162float(h_res[t * hidden_dim_ + p]), - static_cast(top_k) * 0.5f, 1e-2f); - - CHECK_CUDA(cudaStreamDestroy(stream)); -} - -// ── main ────────────────────────────────────────────────────────────────────── - -int main(int argc, char* argv[]) { - if (!ep_bootstrap(argc, argv)) return 0; - int ret = RUN_ALL_TESTS(); - ep_teardown(); - return ret; -} diff --git a/tests/cpp_distributed/test_ep_init.cu b/tests/cpp_distributed/test_ep_init.cu deleted file mode 100644 index 08744dfee5..0000000000 --- a/tests/cpp_distributed/test_ep_init.cu +++ /dev/null @@ -1,64 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -/* - * Unit tests for EP initialization paths. - * - * Tests: - * EPInitTest/InitPath — backend is live after init, handle_mem_size > 0 - * EPInitTest/NumLocalExperts — handle_mem_size is consistent across num_local_experts values - * - * Run via run_test_ep.sh (both uid and comm init paths are tested by the script). - */ - -#include "test_ep_common.h" - -// ── Fixture ─────────────────────────────────────────────────────────────────── - -class EPInitTest : public ::testing::Test { - protected: - void SetUp() override { - if (g_sm_major < 9) - GTEST_SKIP() << "EP requires SM_90+ (device is SM_" << g_sm_major << "0)"; - ASSERT_GE(g_num_processes, 2) << "EP tests require at least 2 processes"; - ASSERT_TRUE(g_ep_initialized) << "EP not initialized"; - } -}; - -// ── Tests ───────────────────────────────────────────────────────────────────── - -TEST_F(EPInitTest, InitPath) { - int nle = g_num_experts / g_ep_size; - NVTEEpLayerConfig cfg{nle, /*top_k=*/2}; - size_t sz = 0; - (void)nvte_ep_register_layer(cfg, &sz); - ASSERT_GT(sz, 0u) << "handle_mem_size must be > 0 after init"; - - if (g_process_id == 0) { - printf(" handle_mem : %zu bytes\n", sz); - } -} - -TEST_F(EPInitTest, NumLocalExperts) { - // handle_mem_size should be > 0 for any valid num_local_experts value. - for (int nle : {1, g_num_experts / g_ep_size}) { - NVTEEpLayerConfig cfg{nle, /*top_k=*/2}; - size_t sz = 0; - (void)nvte_ep_register_layer(cfg, &sz); - ASSERT_GT(sz, 0u) << "num_local_experts=" << nle; - if (g_process_id == 0) - printf(" nle=%-3d handle_mem_size=%zu bytes\n", nle, sz); - } -} - -// ── main ────────────────────────────────────────────────────────────────────── - -int main(int argc, char* argv[]) { - if (!ep_bootstrap(argc, argv)) return 0; - int ret = RUN_ALL_TESTS(); - ep_teardown(); - return ret; -} diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index 1e08cb55df..a5ae99b089 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -82,11 +82,11 @@ void EPBackend::validate_config(const NVTEEpGroupConfig& config) { NVTE_CHECK(config.max_recv_tokens_per_rank > 0, "max_recv_tokens_per_rank must be positive, got ", config.max_recv_tokens_per_rank); NVTE_CHECK(config.hidden_dim > 0, "hidden_dim must be positive, got ", config.hidden_dim); - NVTE_CHECK(config.token_dtype >= 0 && config.token_dtype < kNVTENumTypes, - "token_dtype out of range, got ", static_cast(config.token_dtype)); - const size_t elem_bytes = typeToSize(static_cast(config.token_dtype)); + NVTE_CHECK(config.max_token_dtype >= 0 && config.max_token_dtype < kNVTENumTypes, + "max_token_dtype out of range, got ", static_cast(config.max_token_dtype)); + const size_t elem_bytes = typeToSize(static_cast(config.max_token_dtype)); NVTE_CHECK(config.hidden_dim * elem_bytes >= 16, - "hidden_dim * sizeof(token_dtype) must be >= 16 (NCCL EP 16B row alignment); " + "hidden_dim * sizeof(max_token_dtype) must be >= 16 (NCCL EP 16B row alignment); " "got hidden_dim=", config.hidden_dim, ", element_bytes=", elem_bytes); NVTE_CHECK(config.num_experts % config.ep_size == 0, "num_experts (", config.num_experts, @@ -218,7 +218,7 @@ void EPBackend::init(ncclComm_t ep_comm, NVTEEpGroupConfig group_config) { cfg.algorithm = NCCL_EP_ALGO_HIGH_THROUGHPUT; cfg.num_experts = static_cast(group_config.num_experts); cfg.max_dispatch_tokens_per_rank = static_cast(group_config.max_tokens_per_rank); - const size_t elem_bytes = typeToSize(static_cast(group_config.token_dtype)); + const size_t elem_bytes = typeToSize(static_cast(group_config.max_token_dtype)); cfg.max_token_bytes = static_cast(group_config.hidden_dim * elem_bytes); cfg.rdma_buffer_size = NCCL_EP_AUTO; cfg.num_qp_per_rank = NCCL_EP_AUTO; @@ -346,10 +346,10 @@ void EPBackend::dispatch(uint64_t handle_id, void* handle_mem, const NVTETensor NVTEShape tok_shape = nvte_tensor_shape(tokens); NVTEDType tok_dtype = nvte_tensor_type(tokens); - NVTE_CHECK(tok_dtype == group_config_.token_dtype, - "tokens dtype (", static_cast(tok_dtype), - ") does not match group token_dtype (", - static_cast(group_config_.token_dtype), ")"); + NVTE_CHECK(typeToSize(static_cast(tok_dtype)) <= + typeToSize(static_cast(group_config_.max_token_dtype)), + "tokens dtype (", static_cast(tok_dtype), ") wider than group max_token_dtype (", + static_cast(group_config_.max_token_dtype), ")"); const size_t num_tokens = tok_shape.data[0]; const size_t hidden_dim = tok_shape.data[1]; @@ -376,10 +376,11 @@ void EPBackend::dispatch(uint64_t handle_id, void* handle_mem, const NVTETensor NVTEShape recv_shape = nvte_tensor_shape(recv_tokens); NVTEDType recv_dtype = nvte_tensor_type(recv_tokens); - NVTE_CHECK(recv_dtype == group_config_.token_dtype, + NVTE_CHECK(typeToSize(static_cast(recv_dtype)) <= + typeToSize(static_cast(group_config_.max_token_dtype)), "recv_tokens dtype (", static_cast(recv_dtype), - ") does not match group token_dtype (", - static_cast(group_config_.token_dtype), ")"); + ") wider than group max_token_dtype (", + static_cast(group_config_.max_token_dtype), ")"); size_t recv_sizes[2] = {recv_shape.data[0], recv_shape.data[1]}; ncclEpTensor_t nccl_tokens_out = make_payload_tensor(recv_tokens, recv_tokens_win, 2, diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index a1c9305e9b..22e7ec48ac 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -23,6 +23,8 @@ extern "C" { #endif /* ── Config structs ─────────────────────────────────────────────────────── */ +/* TODO: add a struct_size/version field to these configs (and align with other + * TE public structs) once a TE-wide convention for ABI versioning lands. */ /*! \brief Group-level EP configuration (fixed for the EP group lifetime). */ typedef struct { @@ -35,9 +37,10 @@ typedef struct { int max_num_sms; /*!< Max SMs for EP kernels. 0 = auto. */ /*! 0 (default): throw on relocated handle_mem for a cached handle_id. 1: silently rebuild. */ int allow_handle_mem_reloc; - /*! Token dtype for this EP group. Sizes NCCL EP staging buffers at group - * create and is enforced against tensors passed to nvte_ep_dispatch. */ - NVTEDType token_dtype; + /*! Widest token dtype the group will dispatch. Sizes NCCL EP staging buffers + * at group create. Tensors passed to nvte_ep_dispatch may use any dtype whose + * element size is <= sizeof(max_token_dtype). */ + NVTEDType max_token_dtype; } NVTEEpGroupConfig; /*! \brief Per-layer EP configuration. */ @@ -58,8 +61,8 @@ typedef struct { * nvte_ep_shutdown() returns; destroying it earlier is undefined behavior. * Re-init after shutdown is allowed; double-init throws. * - * v0.1 scope: one EP group per process, bound to the current CUDA device at - * initialize time. Multiple GPUs per process are not supported. + * One EP group per process, bound to the current CUDA device at initialize + * time. Multiple GPUs per process are not supported. * * \param[in] ep_comm Opaque ncclComm_t for the EP sub-group. * \param[in] group_config Group-level EP configuration. diff --git a/transformer_engine/common/util/logging.h b/transformer_engine/common/util/logging.h index da8b9b377d..3308bd22e4 100644 --- a/transformer_engine/common/util/logging.h +++ b/transformer_engine/common/util/logging.h @@ -98,6 +98,14 @@ } \ } while (false) +#define NVTE_CHECK_NCCL(expr) \ + do { \ + const ncclResult_t status_NVTE_CHECK_NCCL = (expr); \ + if (status_NVTE_CHECK_NCCL != ncclSuccess) { \ + NVTE_ERROR("NCCL Error: ", ncclGetErrorString(status_NVTE_CHECK_NCCL)); \ + } \ + } while (false) + #ifdef NVTE_WITH_CUBLASMP #define NVTE_CHECK_CUBLASMP(expr) \ From 16265ae2324876f02e8806becd2434edbbf26b21 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 09:42:18 -0700 Subject: [PATCH 09/55] Expert Parallelism: pointer-keyed LRU handle cache; drop register_layer + NVTEEpHandle struct (NVTE_EP_HANDLE_CACHE_SIZE=-1 disables eviction) Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/test_ep.cu | 109 +++++------ transformer_engine/common/ep/ep_api.cpp | 46 ++--- transformer_engine/common/ep/ep_api_stub.cpp | 30 +-- transformer_engine/common/ep/ep_backend.cpp | 185 ++++++++++-------- transformer_engine/common/ep/ep_backend.h | 87 ++++---- .../common/include/transformer_engine/ep.h | 97 +++++---- 6 files changed, 284 insertions(+), 270 deletions(-) diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu index bcf4ca3c98..7f40b36530 100644 --- a/tests/cpp_distributed/test_ep.cu +++ b/tests/cpp_distributed/test_ep.cu @@ -7,15 +7,15 @@ /* * EP pipeline tests: smallest-scope first. * - * EPDispatchTest/PrepareAndDispatch — exact recv values + per-expert counts - * EPCombineTest/Combine — round-trip: out == top_k * tokens - * EPCombineBwdTest/CombineBwdCheck — exact grad_expert values - * EPDispatchBwdTest/DispatchBwdCheck — exact grad_tokens - * EPDispatchBwdGradWeightsTest/RoundTrip — exact per-(t, k) grad_topk_weights - * EPPipelineTest/FullForwardBackward — fwd + bwd NaN/Inf check + * EPDispatchTest/PrepareAndDispatch : exact recv values + per-expert counts + * EPCombineTest/Combine : round-trip: out == top_k * tokens + * EPCombineBwdTest/CombineBwdCheck : exact grad_expert values + * EPDispatchBwdTest/DispatchBwdCheck : exact grad_tokens + * EPDispatchBwdGradWeightsTest/RoundTrip : exact per-(t, k) grad_topk_weights + * EPPipelineTest/FullForwardBackward : fwd + bwd NaN/Inf check * - * Routing: token t on rank r → expert (r * num_local_experts + t * top_k + k) % num_experts - * Token values: rank r, token t → all hidden dims = (r+1)*0.01 + t*0.001 + * Routing: token t on rank r -> expert (r * num_local_experts + t * top_k + k) % num_experts + * Token values: rank r, token t -> all hidden dims = (r+1)*0.01 + t*0.001 * * Closed-form expected values: * dispatch recv: multiset of source-token values routed to this rank's experts @@ -31,7 +31,7 @@ #include #include -// ── Deterministic routing helpers ───────────────────────────────────────────── +// -- Deterministic routing helpers --------------------------------------------- // Token value for (rank, t): (rank * num_tokens + t + 1) / 256. Step 1/256 is // bf16-exact and unique across (rank, t) when rank * num_tokens + t < 256. @@ -119,7 +119,7 @@ static bool check_no_nan_inf(const T* dev, int count, const char* name) { return true; } -// ── Forward buffer set with RAII ────────────────────────────────────────────── +// -- Forward buffer set with RAII ---------------------------------------------- template struct EPBuffers { @@ -139,14 +139,15 @@ struct EPBuffers { DevBuf g_recv_topk_weights; DevBuf grad_topk_weights; - uint64_t handle_id = 0; size_t handle_mem_size = 0; size_t recv_capacity = 0; int top_k_ = 0; + size_t alignment_ = 0; void alloc(int num_tokens, int top_k, int hidden_dim, int num_local_experts, int ep_size, int max_tokens_per_rank, size_t alignment = 0) { top_k_ = top_k; + alignment_ = alignment; recv_capacity = static_cast(ep_size) * max_tokens_per_rank * 2; topk_idx.alloc(num_tokens * top_k); @@ -157,8 +158,7 @@ struct EPBuffers { recv_topk_weights.alloc(recv_capacity); result.alloc(num_tokens * hidden_dim); - NVTEEpLayerConfig cfg{num_local_experts, top_k, alignment}; - handle_id = nvte_ep_register_layer(cfg, &handle_mem_size); + handle_mem_size = nvte_ep_handle_mem_size(NVTEEpLayerConfig{top_k, alignment}); handle_mem.alloc(handle_mem_size); grad_result.alloc(num_tokens * hidden_dim); @@ -178,8 +178,13 @@ struct EPTensors { TensorWrapper grad_result, grad_expert, grad_tokens; TensorWrapper g_recv_topk_weights, grad_topk_weights; + int top_k_ = 0; + size_t alignment_ = 0; + EPTensors(EPBuffers& b, int num_tokens, int top_k, int hidden_dim, int num_local_experts) { + top_k_ = top_k; + alignment_ = b.alignment_; constexpr DType kTokDType = test::TypeInfo::dtype; using Shape = std::vector; topk_idx = TensorWrapper(b.topk_idx.get(), @@ -211,7 +216,7 @@ struct EPTensors { } }; -// ── Shared fixture base ─────────────────────────────────────────────────────── +// -- Shared fixture base ------------------------------------------------------- class EpOpTestBase : public ::testing::Test { protected: @@ -249,10 +254,6 @@ class EpOpTestBase : public ::testing::Test { h_tok.size() * sizeof(T), cudaMemcpyHostToDevice)); } - NVTEEpLayerConfig layer_config(size_t alignment = 0) const { - return NVTEEpLayerConfig{num_local_experts_, top_k_, alignment}; - } - // NVTE_CHECK_CUDA (fprintf+exit) so this non-void helper stays legal. template int read_total_recv(const EPBuffers& buf) const { @@ -283,9 +284,8 @@ TEST_F(EPDispatchTest, PrepareAndDispatch) { cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - uint64_t handle_id = buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), t.token_counts.data(), /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, t.recv_topk_weights.data(), NVTECommWindow{}, stream)); @@ -303,10 +303,10 @@ TEST_F(EPDispatchTest, PrepareAndDispatch) { total_recv += exp_counts[i]; } ASSERT_LE(total_recv, static_cast(buf.recv_capacity)) - << "total_recv exceeded recv_capacity — overflow would corrupt downstream memory"; + << "total_recv exceeded recv_capacity; overflow would corrupt downstream memory"; // 2. Recv values: read only the filled prefix per local-expert zone, not the - // whole recv buffer — avoids false positives from legitimate-zero token values. + // whole recv buffer; avoids false positives from legitimate-zero token values. std::vector h_recv(buf.recv_capacity * hidden_dim_); NVTE_CHECK_CUDA(cudaMemcpy(h_recv.data(), buf.recv_tokens.get(), h_recv.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); @@ -345,7 +345,7 @@ TEST_F(EPDispatchTest, PrepareAndDispatch) { } // ============================================================================= -// EPCombineTest: round-trip identity expert → result == top_k * tokens. +// EPCombineTest: round-trip identity expert -> result == top_k * tokens. // ============================================================================= class EPCombineTest : public EpOpTestBase {}; @@ -360,13 +360,12 @@ TEST_F(EPCombineTest, Combine) { cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - uint64_t handle_id = buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), t.token_counts.data(), /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, t.recv_topk_weights.data(), NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.recv_tokens.data(), NVTECommWindow{}, + ASSERT_NO_THROW(nvte_ep_combine(t.handle_mem.data(), t.recv_tokens.data(), NVTECommWindow{}, t.result.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); @@ -405,13 +404,12 @@ TEST_F(EPCombineBwdTest, CombineBwdCheck) { cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - uint64_t handle_id = buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), t.token_counts.data(), /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, t.recv_topk_weights.data(), NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.recv_tokens.data(), NVTECommWindow{}, + ASSERT_NO_THROW(nvte_ep_combine(t.handle_mem.data(), t.recv_tokens.data(), NVTECommWindow{}, t.result.data(), stream)); std::vector h_grad_r(num_tokens_ * hidden_dim_, __float2bfloat16(0.1f)); @@ -420,7 +418,7 @@ TEST_F(EPCombineBwdTest, CombineBwdCheck) { cudaMemcpyHostToDevice, stream)); NVTE_CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); - ASSERT_NO_THROW(nvte_ep_combine_bwd(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.grad_result.data(), NVTECommWindow{}, + ASSERT_NO_THROW(nvte_ep_combine_bwd(t.handle_mem.data(), t.grad_result.data(), NVTECommWindow{}, t.grad_expert.data(), NVTECommWindow{}, stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); @@ -472,13 +470,12 @@ TEST_F(EPDispatchBwdTest, DispatchBwdCheck) { cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - uint64_t handle_id = buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), t.token_counts.data(), /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, t.recv_topk_weights.data(), NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.recv_tokens.data(), NVTECommWindow{}, + ASSERT_NO_THROW(nvte_ep_combine(t.handle_mem.data(), t.recv_tokens.data(), NVTECommWindow{}, t.result.data(), stream)); std::vector h_grad(num_tokens_ * hidden_dim_, __float2bfloat16(0.1f)); @@ -489,9 +486,9 @@ TEST_F(EPDispatchBwdTest, DispatchBwdCheck) { NVTE_CHECK_CUDA(cudaMemsetAsync(buf.g_recv_topk_weights.get(), 0, buf.g_recv_topk_weights.bytes(), stream)); NVTE_CHECK_CUDA(cudaMemsetAsync(buf.grad_topk_weights.get(), 0, buf.grad_topk_weights.bytes(), stream)); - ASSERT_NO_THROW(nvte_ep_combine_bwd(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.grad_result.data(), NVTECommWindow{}, + ASSERT_NO_THROW(nvte_ep_combine_bwd(t.handle_mem.data(), t.grad_result.data(), NVTECommWindow{}, t.grad_expert.data(), NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch_bwd(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.grad_expert.data(), NVTECommWindow{}, + ASSERT_NO_THROW(nvte_ep_dispatch_bwd(t.handle_mem.data(), t.grad_expert.data(), NVTECommWindow{}, t.g_recv_topk_weights.data(), NVTECommWindow{}, t.grad_tokens.data(), t.grad_topk_weights.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); @@ -537,11 +534,10 @@ TEST_F(EPDispatchBwdGradWeightsTest, RoundTrip) { cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - uint64_t handle_id = buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), t.token_counts.data(), /*alignment=*/0, stream)); + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream)); NVTE_CHECK_CUDA(cudaMemsetAsync(buf.recv_topk_weights.get(), 0, buf.recv_topk_weights.bytes(), stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), + ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, t.recv_topk_weights.data(), NVTECommWindow{}, stream)); @@ -557,7 +553,7 @@ TEST_F(EPDispatchBwdGradWeightsTest, RoundTrip) { // g_recv_topk_weights := recv_topk_weights (the round-trip input). auto g_recv_t = TensorWrapper(buf.recv_topk_weights.get(), std::vector{buf.recv_capacity}, DType::kFloat32); - ASSERT_NO_THROW(nvte_ep_dispatch_bwd(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.grad_expert.data(), + ASSERT_NO_THROW(nvte_ep_dispatch_bwd(t.handle_mem.data(), t.grad_expert.data(), NVTECommWindow{}, g_recv_t.data(), NVTECommWindow{}, t.grad_tokens.data(), t.grad_topk_weights.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); @@ -609,13 +605,12 @@ class EPPipelineTest : public EpOpTestBase, public ::testing::WithParamInterface cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - uint64_t handle_id = buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), t.token_counts.data(), /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.topk_idx.data(), + ASSERT_NO_THROW(nvte_ep_prepare(t.handle_mem.data(), t.topk_idx.data(), t.token_counts.data(), NVTEEpLayerConfig{t.top_k_, t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(t.handle_mem.data(), t.topk_idx.data(), t.tokens.data(), NVTECommWindow{}, t.topk_weights.data(), NVTECommWindow{}, t.recv_tokens.data(), NVTECommWindow{}, t.recv_topk_weights.data(), NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.recv_tokens.data(), NVTECommWindow{}, + ASSERT_NO_THROW(nvte_ep_combine(t.handle_mem.data(), t.recv_tokens.data(), NVTECommWindow{}, t.result.data(), stream)); std::vector h_grad(num_tokens_ * hidden_dim_, tok_from_float(0.1f)); @@ -626,9 +621,9 @@ class EPPipelineTest : public EpOpTestBase, public ::testing::WithParamInterface NVTE_CHECK_CUDA(cudaMemsetAsync(buf.g_recv_topk_weights.get(), 0, buf.g_recv_topk_weights.bytes(), stream)); NVTE_CHECK_CUDA(cudaMemsetAsync(buf.grad_topk_weights.get(), 0, buf.grad_topk_weights.bytes(), stream)); - ASSERT_NO_THROW(nvte_ep_combine_bwd(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.grad_result.data(), NVTECommWindow{}, + ASSERT_NO_THROW(nvte_ep_combine_bwd(t.handle_mem.data(), t.grad_result.data(), NVTECommWindow{}, t.grad_expert.data(), NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch_bwd(NVTEEpHandle{handle_id, t.handle_mem.data()}, t.grad_expert.data(), NVTECommWindow{}, + ASSERT_NO_THROW(nvte_ep_dispatch_bwd(t.handle_mem.data(), t.grad_expert.data(), NVTECommWindow{}, t.g_recv_topk_weights.data(), NVTECommWindow{}, t.grad_tokens.data(), t.grad_topk_weights.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); @@ -722,13 +717,12 @@ TEST_F(EPZeroCopyTest, IdentityAllSymm) { cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); - uint64_t ref_hid = ref_buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{ref_hid, ref_t.handle_mem.data()}, ref_t.topk_idx.data(), ref_t.token_counts.data(), /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{ref_hid, ref_t.handle_mem.data()}, ref_t.topk_idx.data(), + ASSERT_NO_THROW(nvte_ep_prepare(ref_t.handle_mem.data(), ref_t.topk_idx.data(), ref_t.token_counts.data(), NVTEEpLayerConfig{ref_t.top_k_, ref_t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(ref_t.handle_mem.data(), ref_t.topk_idx.data(), ref_t.tokens.data(), NVTECommWindow{}, ref_t.topk_weights.data(), NVTECommWindow{}, ref_t.recv_tokens.data(), NVTECommWindow{}, ref_t.recv_topk_weights.data(), NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{ref_hid, ref_t.handle_mem.data()}, ref_t.recv_tokens.data(), NVTECommWindow{}, + ASSERT_NO_THROW(nvte_ep_combine(ref_t.handle_mem.data(), ref_t.recv_tokens.data(), NVTECommWindow{}, ref_t.result.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); @@ -761,14 +755,13 @@ TEST_F(EPZeroCopyTest, IdentityAllSymm) { sym_t.recv_tokens = TensorWrapper(sym_recv.ptr, std::vector{sym_buf.recv_capacity, (size_t)hidden_dim_}, DType::kBFloat16); - uint64_t sym_hid = sym_buf.handle_id; - ASSERT_NO_THROW(nvte_ep_prepare(NVTEEpHandle{sym_hid, sym_t.handle_mem.data()}, sym_t.topk_idx.data(), sym_t.token_counts.data(), /*alignment=*/0, stream)); - ASSERT_NO_THROW(nvte_ep_dispatch(NVTEEpHandle{sym_hid, sym_t.handle_mem.data()}, sym_t.topk_idx.data(), + ASSERT_NO_THROW(nvte_ep_prepare(sym_t.handle_mem.data(), sym_t.topk_idx.data(), sym_t.token_counts.data(), NVTEEpLayerConfig{sym_t.top_k_, sym_t.alignment_}, stream)); + ASSERT_NO_THROW(nvte_ep_dispatch(sym_t.handle_mem.data(), sym_t.topk_idx.data(), sym_t.tokens.data(), symm_window(sym_tokens), sym_t.topk_weights.data(), NVTECommWindow{}, sym_t.recv_tokens.data(), symm_window(sym_recv), sym_t.recv_topk_weights.data(), NVTECommWindow{}, stream)); - ASSERT_NO_THROW(nvte_ep_combine(NVTEEpHandle{sym_hid, sym_t.handle_mem.data()}, sym_t.recv_tokens.data(), + ASSERT_NO_THROW(nvte_ep_combine(sym_t.handle_mem.data(), sym_t.recv_tokens.data(), symm_window(sym_recv), sym_t.result.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); @@ -795,7 +788,7 @@ TEST_F(EPZeroCopyTest, IdentityAllSymm) { } -// ── main ────────────────────────────────────────────────────────────────────── +// -- main ---------------------------------------------------------------------- int main(int argc, char* argv[]) { if (!ep_bootstrap(argc, argv)) return 0; diff --git a/transformer_engine/common/ep/ep_api.cpp b/transformer_engine/common/ep/ep_api.cpp index 89d8b38607..51d5af77d0 100644 --- a/transformer_engine/common/ep/ep_api.cpp +++ b/transformer_engine/common/ep/ep_api.cpp @@ -24,53 +24,49 @@ void nvte_ep_initialize(void* ep_comm, NVTEEpGroupConfig group_config) { void nvte_ep_shutdown(void) { EPBackend::shutdown(); } -uint64_t nvte_ep_register_layer(NVTEEpLayerConfig layer_config, size_t* handle_mem_size) { - NVTE_CHECK(handle_mem_size != nullptr, "handle_mem_size must not be null"); - return EPBackend::get().register_layer(layer_config, handle_mem_size); +size_t nvte_ep_handle_mem_size(NVTEEpLayerConfig layer_cfg) { + return EPBackend::get().handle_mem_size(layer_cfg); } -void nvte_ep_prepare(NVTEEpHandle handle, NVTETensor topk_idx, NVTETensor token_counts, - size_t dispatch_output_per_expert_alignment, cudaStream_t stream) { - void* mem_ptr = nvte_tensor_data(handle.mem); - NVTE_CHECK(mem_ptr != nullptr, "handle_mem tensor data must not be null"); - EPBackend::get().prepare(handle.id, topk_idx, token_counts, mem_ptr, - dispatch_output_per_expert_alignment, stream); +namespace { +inline void* handle_mem_ptr(NVTETensor mem) { + void* p = nvte_tensor_data(mem); + NVTE_CHECK(p != nullptr, "handle_mem tensor data must not be null"); + return p; } +} // namespace -void nvte_ep_dispatch(NVTEEpHandle handle, NVTETensor topk_idx, NVTETensor tokens, +void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor token_counts, + NVTEEpLayerConfig layer_cfg, cudaStream_t stream) { + EPBackend::get().prepare(handle_mem_ptr(handle_mem), topk_idx, token_counts, layer_cfg, stream); +} + +void nvte_ep_dispatch(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor tokens, NVTECommWindow tokens_win, NVTETensor topk_weights, NVTECommWindow topk_weights_win, NVTETensor recv_tokens, NVTECommWindow recv_tokens_win, NVTETensor recv_topk_weights, NVTECommWindow recv_topk_weights_win, cudaStream_t stream) { - void* mem_ptr = nvte_tensor_data(handle.mem); - NVTE_CHECK(mem_ptr != nullptr, "handle_mem tensor data must not be null"); - EPBackend::get().dispatch(handle.id, mem_ptr, topk_idx, tokens, tokens_win, topk_weights, + EPBackend::get().dispatch(handle_mem_ptr(handle_mem), topk_idx, tokens, tokens_win, topk_weights, topk_weights_win, recv_tokens, recv_tokens_win, recv_topk_weights, recv_topk_weights_win, stream); } -void nvte_ep_combine(NVTEEpHandle handle, NVTETensor expert_out, NVTECommWindow expert_out_win, +void nvte_ep_combine(NVTETensor handle_mem, NVTETensor expert_out, NVTECommWindow expert_out_win, NVTETensor result, cudaStream_t stream) { - void* mem_ptr = nvte_tensor_data(handle.mem); - NVTE_CHECK(mem_ptr != nullptr, "handle_mem tensor data must not be null"); - EPBackend::get().combine(handle.id, mem_ptr, expert_out, expert_out_win, result, stream); + EPBackend::get().combine(handle_mem_ptr(handle_mem), expert_out, expert_out_win, result, stream); } -void nvte_ep_dispatch_bwd(NVTEEpHandle handle, NVTETensor grad, NVTECommWindow grad_win, +void nvte_ep_dispatch_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow grad_win, NVTETensor g_recv_topk_weights, NVTECommWindow g_recv_topk_weights_win, NVTETensor grad_tokens, NVTETensor grad_topk_weights, cudaStream_t stream) { - void* mem_ptr = nvte_tensor_data(handle.mem); - NVTE_CHECK(mem_ptr != nullptr, "handle_mem tensor data must not be null"); - EPBackend::get().dispatch_bwd(handle.id, mem_ptr, grad, grad_win, g_recv_topk_weights, + EPBackend::get().dispatch_bwd(handle_mem_ptr(handle_mem), grad, grad_win, g_recv_topk_weights, g_recv_topk_weights_win, grad_tokens, grad_topk_weights, stream); } -void nvte_ep_combine_bwd(NVTEEpHandle handle, NVTETensor grad, NVTECommWindow grad_win, +void nvte_ep_combine_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow grad_win, NVTETensor grad_expert_out, NVTECommWindow grad_expert_out_win, cudaStream_t stream) { - void* mem_ptr = nvte_tensor_data(handle.mem); - NVTE_CHECK(mem_ptr != nullptr, "handle_mem tensor data must not be null"); - EPBackend::get().combine_bwd(handle.id, mem_ptr, grad, grad_win, grad_expert_out, + EPBackend::get().combine_bwd(handle_mem_ptr(handle_mem), grad, grad_win, grad_expert_out, grad_expert_out_win, stream); } diff --git a/transformer_engine/common/ep/ep_api_stub.cpp b/transformer_engine/common/ep/ep_api_stub.cpp index fe4127d87d..a62416cc7f 100644 --- a/transformer_engine/common/ep/ep_api_stub.cpp +++ b/transformer_engine/common/ep/ep_api_stub.cpp @@ -24,38 +24,40 @@ void nvte_ep_initialize(void* /*ep_comm*/, NVTEEpGroupConfig /*group_config*/) { void nvte_ep_shutdown(void) {} -uint64_t nvte_ep_register_layer(NVTEEpLayerConfig /*layer_config*/, size_t* /*handle_mem_size*/) { - ep_not_built(); -} +size_t nvte_ep_handle_mem_size(NVTEEpLayerConfig /*layer_cfg*/) { ep_not_built(); } -void nvte_ep_prepare(NVTEEpHandle /*handle*/, NVTETensor /*topk_idx*/, NVTETensor /*token_counts*/, - size_t /*dispatch_output_per_expert_alignment*/, cudaStream_t /*stream*/) { +void nvte_ep_prepare(NVTETensor /*handle_mem*/, NVTETensor /*topk_idx*/, + NVTETensor /*token_counts*/, NVTEEpLayerConfig /*layer_cfg*/, + cudaStream_t /*stream*/) { ep_not_built(); } -void nvte_ep_dispatch(NVTEEpHandle /*handle*/, NVTETensor /*topk_idx*/, NVTETensor /*tokens*/, +void nvte_ep_dispatch(NVTETensor /*handle_mem*/, NVTETensor /*topk_idx*/, NVTETensor /*tokens*/, NVTECommWindow /*tokens_win*/, NVTETensor /*topk_weights*/, NVTECommWindow /*topk_weights_win*/, NVTETensor /*recv_tokens*/, NVTECommWindow /*recv_tokens_win*/, NVTETensor /*recv_topk_weights*/, - NVTECommWindow /*recv_topk_weights_win*/, cudaStream_t /*stream*/) { + NVTECommWindow /*recv_topk_weights_win*/, NVTEEpLayerConfig /*layer_cfg*/, + cudaStream_t /*stream*/) { ep_not_built(); } -void nvte_ep_combine(NVTEEpHandle /*handle*/, NVTETensor /*expert_out*/, +void nvte_ep_combine(NVTETensor /*handle_mem*/, NVTETensor /*expert_out*/, NVTECommWindow /*expert_out_win*/, NVTETensor /*result*/, - cudaStream_t /*stream*/) { + NVTEEpLayerConfig /*layer_cfg*/, cudaStream_t /*stream*/) { ep_not_built(); } -void nvte_ep_dispatch_bwd(NVTEEpHandle /*handle*/, NVTETensor /*grad*/, NVTECommWindow /*grad_win*/, - NVTETensor /*g_recv_topk_weights*/, +void nvte_ep_dispatch_bwd(NVTETensor /*handle_mem*/, NVTETensor /*grad*/, + NVTECommWindow /*grad_win*/, NVTETensor /*g_recv_topk_weights*/, NVTECommWindow /*g_recv_topk_weights_win*/, NVTETensor /*grad_tokens*/, - NVTETensor /*grad_topk_weights*/, cudaStream_t /*stream*/) { + NVTETensor /*grad_topk_weights*/, NVTEEpLayerConfig /*layer_cfg*/, + cudaStream_t /*stream*/) { ep_not_built(); } -void nvte_ep_combine_bwd(NVTEEpHandle /*handle*/, NVTETensor /*grad*/, NVTECommWindow /*grad_win*/, - NVTETensor /*grad_expert_out*/, NVTECommWindow /*grad_expert_out_win*/, +void nvte_ep_combine_bwd(NVTETensor /*handle_mem*/, NVTETensor /*grad*/, + NVTECommWindow /*grad_win*/, NVTETensor /*grad_expert_out*/, + NVTECommWindow /*grad_expert_out_win*/, NVTEEpLayerConfig /*layer_cfg*/, cudaStream_t /*stream*/) { ep_not_built(); } diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index a5ae99b089..ae7c0900d6 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -11,7 +11,6 @@ #include "ep_backend.h" #include -#include #include #include #include @@ -136,20 +135,18 @@ void EPBackend::shutdown() { EPBackend& inst = instance(); std::lock_guard lock(inst.mutex_); if (!inst.initialized_) return; - for (auto& kv : inst.handles_) { - if (kv.second.cached_handle != nullptr) { - ncclEpHandleDestroy(kv.second.cached_handle); - kv.second.cached_handle = nullptr; - kv.second.cached_handle_mem = nullptr; - } + for (auto& e : inst.lru_) { + if (e.handle != nullptr) ncclEpHandleDestroy(e.handle); } - inst.handles_.clear(); + inst.lru_.clear(); + inst.index_.clear(); + inst.fallback_layer_cfg_.reset(); // ncclEpGroupDestroy reads from ep_comm_; destroy group while comm is still alive. if (inst.ep_group_ != nullptr) { ncclEpGroupDestroy(inst.ep_group_); inst.ep_group_ = nullptr; } - inst.ep_comm_ = nullptr; // borrowed — caller destroys + inst.ep_comm_ = nullptr; // borrowed; caller destroys inst.initialized_ = false; } @@ -181,7 +178,6 @@ ncclDataType_t EPBackend::nvte_dtype_to_nccl(NVTEDType dtype) { return ncclFloat32; // unreachable } -// Open a fresh ncclEpHandle over handle_mem. Caller (or cache) owns the result. ncclEpHandle_t EPBackend::open_handle(void* handle_mem, size_t handle_mem_size, int num_topk, size_t dispatch_output_per_expert_alignment) { size_t hm_sizes[1] = {handle_mem_size}; @@ -203,7 +199,9 @@ ncclEpHandle_t EPBackend::open_handle(void* handle_mem, size_t handle_mem_size, EPBackend::~EPBackend() { std::lock_guard lock(mutex_); if (!initialized_) return; - handles_.clear(); + lru_.clear(); + index_.clear(); + fallback_layer_cfg_.reset(); ep_group_ = nullptr; ep_comm_ = nullptr; initialized_ = false; @@ -237,83 +235,119 @@ void EPBackend::init(ncclComm_t ep_comm, NVTEEpGroupConfig group_config) { } // --------------------------------------------------------------------------- -// Per-handle_id config cache +// Pointer-keyed LRU cache // --------------------------------------------------------------------------- -uint64_t EPBackend::insert_new_entry(size_t handle_mem_size, int top_k, size_t alignment) { +size_t EPBackend::cache_cap_locked() { if (handle_cache_cap_ == 0) { const char* cap_env = std::getenv("NVTE_EP_HANDLE_CACHE_SIZE"); - handle_cache_cap_ = (cap_env != nullptr) ? std::max(1, std::atoi(cap_env)) : 8192; + if (cap_env != nullptr) { + const long v = std::atol(cap_env); + if (v < 0) { + // Unlimited cache. WAR for JAX until XLA fixes handle_mem + // reloc between runs. + handle_cache_cap_ = SIZE_MAX; + } else { + NVTE_CHECK(v > 0, + "NVTE_EP_HANDLE_CACHE_SIZE=0 is invalid; use -1 for unlimited or a positive " + "cap."); + handle_cache_cap_ = static_cast(v); + } + } else { + handle_cache_cap_ = 4096; + } } - NVTE_CHECK(handles_.size() < handle_cache_cap_, "EP handle cache full (", handle_cache_cap_, - " entries). Raise via NVTE_EP_HANDLE_CACHE_SIZE."); - uint64_t id = next_handle_id_.fetch_add(1, std::memory_order_relaxed); - handles_.emplace(id, HandleEntry{handle_mem_size, alignment, top_k}); - return id; + return handle_cache_cap_; } -EPBackend::HandleEntry& EPBackend::lookup_config(uint64_t handle_id) { - auto it = handles_.find(handle_id); - NVTE_CHECK(it != handles_.end(), "ep op on handle_id=", handle_id, - " with no cached config — call ep_prepare first."); - return it->second; -} +ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerConfig layer_cfg) { + // Update the program-wide fallback cfg so dispatch/combine/_bwd can + // reconstruct the handle on a pointer-cache miss (WAR for XLA buffer reloc + // between runs; one cfg per process). Remove this once XLA preserves the + // handle_mem device pointer across runs. + if (fallback_layer_cfg_.has_value()) { + NVTE_CHECK(fallback_layer_cfg_->top_k == layer_cfg.top_k, + "EP prepare top_k=", layer_cfg.top_k, + " disagrees with process-wide cached top_k=", fallback_layer_cfg_->top_k); + NVTE_CHECK(fallback_layer_cfg_->dispatch_output_per_expert_alignment == + layer_cfg.dispatch_output_per_expert_alignment, + "EP prepare alignment=", layer_cfg.dispatch_output_per_expert_alignment, + " disagrees with process-wide cached alignment=", + fallback_layer_cfg_->dispatch_output_per_expert_alignment); + } else { + fallback_layer_cfg_ = layer_cfg; + } -ncclEpHandle_t EPBackend::get_or_open_handle(HandleEntry& cfg, void* handle_mem) { - if (cfg.cached_handle != nullptr && cfg.cached_handle_mem == handle_mem) { - return cfg.cached_handle; + auto it = index_.find(handle_mem); + if (it != index_.end()) { + lru_.splice(lru_.begin(), lru_, it->second); + return it->second->handle; } - if (cfg.cached_handle != nullptr) { - NVTE_CHECK(group_config_.allow_handle_mem_reloc != 0, - "EP handle_mem relocated for cached handle (old=", - reinterpret_cast(cfg.cached_handle_mem), - ", new=", reinterpret_cast(handle_mem), - "). Set NVTEEpGroupConfig.allow_handle_mem_reloc=1 to allow rebuild."); - ncclEpHandleDestroy(cfg.cached_handle); - cfg.cached_handle = nullptr; - cfg.cached_handle_mem = nullptr; + ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; + hcfg.dispatch_output_per_expert_alignment = layer_cfg.dispatch_output_per_expert_alignment; + size_t hm_size = 0; + NVTE_CHECK_NCCL(ncclEpHandleMemSize(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, &hm_size, + layer_cfg.top_k)); + ncclEpHandle_t h = open_handle(handle_mem, hm_size, layer_cfg.top_k, + layer_cfg.dispatch_output_per_expert_alignment); + lru_.push_front(HandleEntry{handle_mem, h, layer_cfg, hm_size}); + index_.emplace(handle_mem, lru_.begin()); + while (lru_.size() > cache_cap_locked()) { + HandleEntry& victim = lru_.back(); + if (victim.handle != nullptr) ncclEpHandleDestroy(victim.handle); + index_.erase(victim.handle_mem); + lru_.pop_back(); } - ncclEpHandle_t h = open_handle(handle_mem, cfg.handle_mem_size, cfg.top_k, cfg.alignment); - cfg.cached_handle = h; - cfg.cached_handle_mem = handle_mem; return h; } +ncclEpHandle_t EPBackend::lookup_handle_locked(void* handle_mem) { + auto it = index_.find(handle_mem); + if (it != index_.end()) { + lru_.splice(lru_.begin(), lru_, it->second); + return it->second->handle; + } + // Miss: reconstruct from the process-wide cached cfg. XLA may relocate + // handle_mem between runs, breaking the pointer key; the fallback cfg lets + // us open a fresh handle on the new buffer. Drop this branch once XLA + // preserves buffer pointers. + const uintptr_t hm_addr = reinterpret_cast(handle_mem); + NVTE_CHECK(fallback_layer_cfg_.has_value(), "ep op on handle_mem=0x", hm_addr, + " with no cached entry and no prior nvte_ep_prepare; call prepare first."); + return prepare_handle_locked(handle_mem, *fallback_layer_cfg_); +} + // --------------------------------------------------------------------------- // Per-step operations // --------------------------------------------------------------------------- -uint64_t EPBackend::register_layer(NVTEEpLayerConfig layer_config, size_t* handle_mem_size) { +size_t EPBackend::handle_mem_size(NVTEEpLayerConfig layer_cfg) { NVTE_CHECK(initialized_, "EPBackend not initialized"); - NVTE_CHECK(layer_config.top_k > 0, "NVTEEpLayerConfig.top_k must be > 0"); - NVTE_CHECK(handle_mem_size != nullptr, "handle_mem_size must not be null"); + NVTE_CHECK(layer_cfg.top_k > 0, "top_k must be > 0, got ", layer_cfg.top_k); ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; - hcfg.dispatch_output_per_expert_alignment = layer_config.dispatch_output_per_expert_alignment; + hcfg.dispatch_output_per_expert_alignment = layer_cfg.dispatch_output_per_expert_alignment; size_t hm_size = 0; NVTE_CHECK_NCCL(ncclEpHandleMemSize(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, &hm_size, - layer_config.top_k)); - *handle_mem_size = hm_size; - std::lock_guard lock(mutex_); - return insert_new_entry(hm_size, layer_config.top_k, - layer_config.dispatch_output_per_expert_alignment); + layer_cfg.top_k)); + return hm_size; } -void EPBackend::prepare(uint64_t handle_id, const NVTETensor topk_idx, NVTETensor token_counts, - void* handle_mem, size_t dispatch_output_per_expert_alignment, - cudaStream_t stream) { +void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, NVTETensor token_counts, + NVTEEpLayerConfig layer_cfg, cudaStream_t stream) { NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); + NVTE_CHECK(layer_cfg.top_k > 0, "top_k must be > 0, got ", layer_cfg.top_k); NVTEShape idx_shape = nvte_tensor_shape(topk_idx); void* idx_data = nvte_tensor_data(topk_idx); NVTE_CHECK(idx_data != nullptr, "topk_idx data must not be null"); const size_t num_tokens = idx_shape.data[0]; - const size_t top_k = idx_shape.ndim > 1 ? idx_shape.data[1] : 1; + const size_t topk_in = idx_shape.ndim > 1 ? idx_shape.data[1] : 1; const size_t num_local_experts = static_cast(group_config_.num_experts / group_config_.ep_size); - size_t idx_sizes[2] = {num_tokens, top_k}; + size_t idx_sizes[2] = {num_tokens, topk_in}; ncclEpTensor_t nccl_topk_idx = make_tensor(idx_data, 2, ncclInt64, idx_sizes); // ncclEpUpdateHandle writes per-expert counts via expert_counters. @@ -327,20 +361,15 @@ void EPBackend::prepare(uint64_t handle_id, const NVTETensor topk_idx, NVTETenso layout_info.expert_counters = (token_counts_data != nullptr) ? &token_counts_desc : nullptr; std::lock_guard lock(mutex_); - HandleEntry& cfg = lookup_config(handle_id); - NVTE_CHECK(cfg.alignment == dispatch_output_per_expert_alignment, - "ep_prepare: alignment mismatch for handle_id=", handle_id, " (cached=", cfg.alignment, - ", got=", dispatch_output_per_expert_alignment, ")"); - ncclEpHandle_t h = get_or_open_handle(cfg, handle_mem); + ncclEpHandle_t h = prepare_handle_locked(handle_mem, layer_cfg); NVTE_CHECK_NCCL(ncclEpUpdateHandle(h, &nccl_topk_idx, &layout_info, stream)); } -void EPBackend::dispatch(uint64_t handle_id, void* handle_mem, const NVTETensor topk_idx, - const NVTETensor tokens, const NVTECommWindow& tokens_win, - const NVTETensor topk_weights, const NVTECommWindow& topk_weights_win, - NVTETensor recv_tokens, const NVTECommWindow& recv_tokens_win, - NVTETensor recv_topk_weights, const NVTECommWindow& recv_topk_weights_win, - cudaStream_t stream) { +void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTETensor tokens, + const NVTECommWindow& tokens_win, const NVTETensor topk_weights, + const NVTECommWindow& topk_weights_win, NVTETensor recv_tokens, + const NVTECommWindow& recv_tokens_win, NVTETensor recv_topk_weights, + const NVTECommWindow& recv_topk_weights_win, cudaStream_t stream) { NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); @@ -367,9 +396,9 @@ void EPBackend::dispatch(uint64_t handle_id, void* handle_mem, const NVTETensor if (is_forward) { NVTE_CHECK(topk_idx != nullptr, "topk_idx required in forward dispatch"); NVTEShape idx_shape = nvte_tensor_shape(topk_idx); - const size_t top_k = idx_shape.ndim > 1 ? idx_shape.data[1] : 1; + const size_t topk_in = idx_shape.ndim > 1 ? idx_shape.data[1] : 1; weights_in_sizes[0] = num_tokens; - weights_in_sizes[1] = top_k; + weights_in_sizes[1] = topk_in; nccl_topk_weights_in = make_payload_tensor(topk_weights, topk_weights_win, 2, ncclFloat32, weights_in_sizes); } @@ -409,13 +438,12 @@ void EPBackend::dispatch(uint64_t handle_id, void* handle_mem, const NVTETensor dispatch_cfg.pass_direction = is_forward ? NCCL_EP_FWD_PASS : NCCL_EP_BWD_PASS; std::lock_guard lock(mutex_); - HandleEntry& cfg = lookup_config(handle_id); - ncclEpHandle_t h = get_or_open_handle(cfg, handle_mem); + ncclEpHandle_t h = lookup_handle_locked(handle_mem); NVTE_CHECK_NCCL(ncclEpDispatch(h, &in_struct, &out_struct, /*layout_info=*/nullptr, &dispatch_cfg, stream)); } -void EPBackend::combine(uint64_t handle_id, void* handle_mem, const NVTETensor expert_out, +void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, const NVTECommWindow& expert_out_win, NVTETensor result, cudaStream_t stream) { NVTE_CHECK(initialized_, "EPBackend not initialized"); @@ -444,12 +472,11 @@ void EPBackend::combine(uint64_t handle_id, void* handle_mem, const NVTETensor e out_struct.tokens = &nccl_result_out; std::lock_guard lock(mutex_); - HandleEntry& cfg = lookup_config(handle_id); - ncclEpHandle_t h = get_or_open_handle(cfg, handle_mem); + ncclEpHandle_t h = lookup_handle_locked(handle_mem); NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, /*config=*/nullptr, stream)); } -void EPBackend::dispatch_bwd(uint64_t handle_id, void* handle_mem, const NVTETensor grad, +void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, const NVTECommWindow& grad_win, const NVTETensor g_recv_topk_weights, const NVTECommWindow& g_recv_topk_weights_win, NVTETensor grad_tokens, NVTETensor grad_topk_weights, cudaStream_t stream) { @@ -462,7 +489,7 @@ void EPBackend::dispatch_bwd(uint64_t handle_id, void* handle_mem, const NVTETen ncclEpTensor_t nccl_tok_in = make_payload_tensor(grad, grad_win, 2, nvte_dtype_to_nccl(g_dtype), g_sizes); - // g_recv_topk_weights must be 1D [recv_capacity] — caller flattens. + // g_recv_topk_weights must be 1D [recv_capacity]; caller flattens. NVTEShape gw_shape = nvte_tensor_shape(g_recv_topk_weights); NVTE_CHECK(gw_shape.ndim == 1, "g_recv_topk_weights must be 1D [recv_capacity]; caller must flatten leading dims"); @@ -495,17 +522,17 @@ void EPBackend::dispatch_bwd(uint64_t handle_id, void* handle_mem, const NVTETen cfg.pass_direction = NCCL_EP_BWD_PASS; std::lock_guard lock(mutex_); - HandleEntry& entry = lookup_config(handle_id); - ncclEpHandle_t h = get_or_open_handle(entry, handle_mem); + ncclEpHandle_t h = lookup_handle_locked(handle_mem); NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, &cfg, stream)); } -void EPBackend::combine_bwd(uint64_t handle_id, void* handle_mem, const NVTETensor grad, +void EPBackend::combine_bwd(void* handle_mem, const NVTETensor grad, const NVTECommWindow& grad_win, NVTETensor grad_expert_out, const NVTECommWindow& grad_expert_out_win, cudaStream_t stream) { // Backward of combine = reverse-direction dispatch. - dispatch(handle_id, handle_mem, /*topk_idx=*/nullptr, grad, grad_win, /*topk_weights=*/nullptr, - /*topk_weights_win=*/NVTECommWindow{}, grad_expert_out, grad_expert_out_win, + dispatch(handle_mem, /*topk_idx=*/nullptr, grad, grad_win, + /*topk_weights=*/nullptr, /*topk_weights_win=*/NVTECommWindow{}, grad_expert_out, + grad_expert_out_win, /*recv_topk_weights=*/nullptr, /*recv_topk_weights_win=*/NVTECommWindow{}, stream); } diff --git a/transformer_engine/common/ep/ep_backend.h b/transformer_engine/common/ep/ep_backend.h index e82c974c3f..405226646b 100644 --- a/transformer_engine/common/ep/ep_backend.h +++ b/transformer_engine/common/ep/ep_backend.h @@ -7,9 +7,10 @@ /*! \file ep_backend.h * \brief Internal NCCL EP singleton; not part of the public API. * - * Per handle_id the cache stores config only (no device pointers), so - * handle_mem may be relocated between ops. Cap: NVTE_EP_HANDLE_CACHE_SIZE - * (default 8192); overflow throws. + * ncclEpHandles are cached by handle_mem device pointer. nvte_ep_prepare + * seeds the entry with the layer_cfg; dispatch/combine/_bwd look up by + * pointer. Cache cap: NVTE_EP_HANDLE_CACHE_SIZE (default 4096; -1 disables + * LRU eviction). */ #ifndef TRANSFORMER_ENGINE_COMMON_EP_EP_BACKEND_H_ @@ -20,16 +21,17 @@ #include #include -#include #include #include +#include #include +#include #include namespace transformer_engine { namespace ep { -/*! \brief EP backend singleton — owns the NCCL EP group; borrows the comm. */ +/*! \brief EP backend singleton; owns the NCCL EP group, borrows the comm. */ class EPBackend { public: /*! \brief Access the singleton. Aborts if not initialized. */ @@ -44,32 +46,32 @@ class EPBackend { /*! \brief Tear down the backend. Idempotent. Does not destroy ep_comm_. */ static void shutdown(); - // Host-only: reserve a fresh handle_id, cache the layer config, and report - // the handle_mem buffer size the caller must allocate. - uint64_t register_layer(NVTEEpLayerConfig layer_config, size_t* handle_mem_size); + // Host-only: report handle_mem byte size for layer_cfg. + size_t handle_mem_size(NVTEEpLayerConfig layer_cfg); - void prepare(uint64_t handle_id, const NVTETensor topk_idx, NVTETensor token_counts, - void* handle_mem, size_t dispatch_output_per_expert_alignment, cudaStream_t stream); + // Seeds the cache for handle_mem with layer_cfg and runs the routing AllGather. + void prepare(void* handle_mem, const NVTETensor topk_idx, NVTETensor token_counts, + NVTEEpLayerConfig layer_cfg, cudaStream_t stream); - void dispatch(uint64_t handle_id, void* handle_mem, const NVTETensor topk_idx, - const NVTETensor tokens, const NVTECommWindow& tokens_win, - const NVTETensor topk_weights, const NVTECommWindow& topk_weights_win, - NVTETensor recv_tokens, const NVTECommWindow& recv_tokens_win, - NVTETensor recv_topk_weights, const NVTECommWindow& recv_topk_weights_win, - cudaStream_t stream); + // Per-step ops below require a prior prepare(). + void dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTETensor tokens, + const NVTECommWindow& tokens_win, const NVTETensor topk_weights, + const NVTECommWindow& topk_weights_win, NVTETensor recv_tokens, + const NVTECommWindow& recv_tokens_win, NVTETensor recv_topk_weights, + const NVTECommWindow& recv_topk_weights_win, cudaStream_t stream); - void combine(uint64_t handle_id, void* handle_mem, const NVTETensor expert_out, + void combine(void* handle_mem, const NVTETensor expert_out, const NVTECommWindow& expert_out_win, NVTETensor result, cudaStream_t stream); // g_recv_topk_weights: 1D [recv_capacity] f32; grad_topk_weights: 2D [T, top_k] f32. - void dispatch_bwd(uint64_t handle_id, void* handle_mem, const NVTETensor grad, - const NVTECommWindow& grad_win, const NVTETensor g_recv_topk_weights, + void dispatch_bwd(void* handle_mem, const NVTETensor grad, const NVTECommWindow& grad_win, + const NVTETensor g_recv_topk_weights, const NVTECommWindow& g_recv_topk_weights_win, NVTETensor grad_tokens, NVTETensor grad_topk_weights, cudaStream_t stream); - void combine_bwd(uint64_t handle_id, void* handle_mem, const NVTETensor grad, - const NVTECommWindow& grad_win, NVTETensor grad_expert_out, - const NVTECommWindow& grad_expert_out_win, cudaStream_t stream); + void combine_bwd(void* handle_mem, const NVTETensor grad, const NVTECommWindow& grad_win, + NVTETensor grad_expert_out, const NVTECommWindow& grad_expert_out_win, + cudaStream_t stream); private: EPBackend() = default; @@ -77,43 +79,40 @@ class EPBackend { EPBackend(const EPBackend&) = delete; EPBackend& operator=(const EPBackend&) = delete; - // ep_comm is borrowed — caller retains ownership across the backend lifetime. + // ep_comm is borrowed; caller retains ownership across the backend lifetime. void init(ncclComm_t ep_comm, NVTEEpGroupConfig config); static EPBackend& instance(); // Meyers singleton accessor static void validate_config(const NVTEEpGroupConfig& config); static ncclDataType_t nvte_dtype_to_nccl(NVTEDType dtype); - // Open a transient ncclEpHandle over handle_mem. num_topk=-1 for paths + // Open a fresh ncclEpHandle over handle_mem. num_topk=-1 for paths // that don't carry per-token weights. ncclEpHandle_t open_handle(void* handle_mem, size_t handle_mem_size, int num_topk, size_t dispatch_output_per_expert_alignment); + // LRU cache: most-recently-used at the front of lru_, evict from the back. + struct HandleEntry { + void* handle_mem; + ncclEpHandle_t handle; + NVTEEpLayerConfig layer_cfg; + size_t handle_mem_size; + }; + ncclEpGroup_t ep_group_{nullptr}; ncclComm_t ep_comm_{nullptr}; NVTEEpGroupConfig group_config_{}; bool initialized_{false}; std::mutex mutex_; - struct HandleEntry { - size_t handle_mem_size; - size_t alignment; - int top_k; - // Persistent ncclEpHandle bound to cached_handle_mem. Lazily opened on first - // op; reused while handle_mem ptr is unchanged. Destroyed in shutdown(). - ncclEpHandle_t cached_handle{nullptr}; - void* cached_handle_mem{nullptr}; - }; - std::unordered_map handles_; - std::atomic next_handle_id_{1}; // 0 reserved as "no id" - size_t handle_cache_cap_{0}; // set lazily from NVTE_EP_HANDLE_CACHE_SIZE - - // Caller must hold mutex_. Throws on cap overflow. - uint64_t insert_new_entry(size_t handle_mem_size, int top_k, size_t alignment); - HandleEntry& lookup_config(uint64_t handle_id); - // Caller must hold mutex_. Returns the cached handle if handle_mem matches. - // On mismatch: if group_config_.allow_handle_mem_reloc != 0, destroys the - // stale handle and opens a fresh one; otherwise throws. - ncclEpHandle_t get_or_open_handle(HandleEntry& cfg, void* handle_mem); + std::list lru_; + std::unordered_map::iterator> index_; + size_t handle_cache_cap_{0}; // set lazily from NVTE_EP_HANDLE_CACHE_SIZE + std::optional fallback_layer_cfg_; + + // Caller must hold mutex_. + ncclEpHandle_t prepare_handle_locked(void* handle_mem, NVTEEpLayerConfig layer_cfg); + ncclEpHandle_t lookup_handle_locked(void* handle_mem); + size_t cache_cap_locked(); }; } // namespace ep diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index 22e7ec48ac..b18862bb44 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -5,8 +5,13 @@ ************************************************************************/ /*! \file ep.h - * \brief Public C API for Expert Parallelism. Per-step ops are allocation-free - * and CUDA graph-capturable. + * \brief Public C API for Expert Parallelism. Per-step ops are + * allocation-free and CUDA graph-capturable. + * + * Per layer: call nvte_ep_handle_mem_size(layer_cfg) for the buffer size; + * allocate handle_mem as a kByte NVTETensor. Per step: nvte_ep_prepare seeds + * routing, then nvte_ep_dispatch / nvte_ep_combine / _bwd consume it. + * Cache cap: NVTE_EP_HANDLE_CACHE_SIZE (default 4096; -1 disables eviction). */ #ifndef TRANSFORMER_ENGINE_EP_H_ @@ -22,7 +27,7 @@ extern "C" { #endif -/* ── Config structs ─────────────────────────────────────────────────────── */ +/* -- Config structs ------------------------------------------------------- */ /* TODO: add a struct_size/version field to these configs (and align with other * TE public structs) once a TE-wide convention for ABI versioning lands. */ @@ -35,24 +40,23 @@ typedef struct { int max_recv_tokens_per_rank; int hidden_dim; /*!< Token hidden dimension. */ int max_num_sms; /*!< Max SMs for EP kernels. 0 = auto. */ - /*! 0 (default): throw on relocated handle_mem for a cached handle_id. 1: silently rebuild. */ - int allow_handle_mem_reloc; /*! Widest token dtype the group will dispatch. Sizes NCCL EP staging buffers * at group create. Tensors passed to nvte_ep_dispatch may use any dtype whose * element size is <= sizeof(max_token_dtype). */ NVTEDType max_token_dtype; } NVTEEpGroupConfig; -/*! \brief Per-layer EP configuration. */ +/*! \brief Per-layer configuration consumed by nvte_ep_handle_mem_size and + * nvte_ep_prepare. Reserved for future per-call options (fp8 scale, + * overflow policy, ...). + */ typedef struct { - int num_local_experts; /*!< Reserved for ABI stability (derived from group config). */ - int top_k; /*!< Per-token expert fan-out. Required. */ + int top_k; /*!< Per-token expert fan-out (> 0). */ + /*! Per-expert zone alignment in tokens (pow2; 0/1 = none). */ size_t dispatch_output_per_expert_alignment; - /*!< Per-expert zone alignment in tokens (pow2; 0/1 = no padding). Must match - * between nvte_ep_register_layer and nvte_ep_prepare. */ } NVTEEpLayerConfig; -/* ── Bootstrap ──────────────────────────────────────────────────────────── */ +/* -- Bootstrap ------------------------------------------------------------ */ /*! \brief Bootstrap from an existing NCCL EP sub-communicator. Requires SM>=90. * @@ -72,45 +76,36 @@ void nvte_ep_initialize(void* ep_comm, NVTEEpGroupConfig group_config); /*! \brief Tear down the EP backend. Idempotent. Does not destroy ep_comm. */ void nvte_ep_shutdown(void); -/* ── Layer registration (host-only, eager) ───────────────────────────────── */ +/* -- Layer sizing (host-only) --------------------------------------------- */ -/*! \brief Reserve a handle_id for a layer config and report the handle_mem buffer - * size the caller must allocate. Host-only. - * - * Registration is intended to be static (once per layer at model init). There is - * no per-layer unregister API; all registrations are released by nvte_ep_shutdown. - * Re-registering the same layer config each step is not supported and will - * eventually exhaust the handle cache (NVTE_EP_HANDLE_CACHE_SIZE, default 8192). +/*! \brief Report the handle_mem byte size required for the given layer config. + * Host-only; cheap to call. The caller allocates the buffer and passes + * it back to ep ops as the handle_mem argument. top_k comes from the + * active NVTEEpGroupConfig. * - * \param[in] layer_config Per-layer EP configuration. - * \param[out] handle_mem_size Bytes the caller must allocate for handle_mem. - * \return uint64_t handle_id (non-zero). + * \param[in] layer_cfg Per-call layer configuration. + * \return size in bytes for the handle_mem buffer. */ -uint64_t nvte_ep_register_layer(NVTEEpLayerConfig layer_config, size_t* handle_mem_size); - -/*! \brief Per-step handle: the registered handle_id paired with its handle_mem buffer. */ -typedef struct { - uint64_t id; /*!< Handle id from nvte_ep_register_layer. */ - NVTETensor mem; /*!< Caller-allocated handle_mem buffer (size from nvte_ep_register_layer). */ -} NVTEEpHandle; +size_t nvte_ep_handle_mem_size(NVTEEpLayerConfig layer_cfg); -/* ── Per-step ops (all allocation-free, CUDA graph-capturable) ──────────── */ +/* -- Per-step ops (all allocation-free, CUDA graph-capturable) ------------ */ /*! \brief AllGather the routing map; write per-expert counts and cache routing - * metadata in handle.mem for the subsequent dispatch/combine. + * metadata in handle_mem for the subsequent dispatch/combine. * - * \param[in] handle EP handle (id + mem buffer). - * \param[in] topk_idx [T, top_k] int64 routing indices. - * \param[out] token_counts [num_local_experts] int32 counts. - * \param[in] dispatch_output_per_expert_alignment Must match the handle_mem sizing. - * \param[in] stream CUDA stream. + * \param[in] handle_mem uint8 routing-state buffer. + * \param[in] topk_idx [T, top_k] int64 routing indices. + * \param[out] token_counts [num_local_experts] int32 counts. + * \param[in] layer_cfg Per-call layer configuration. + * \param[in] stream CUDA stream. */ -void nvte_ep_prepare(NVTEEpHandle handle, NVTETensor topk_idx, NVTETensor token_counts, - size_t dispatch_output_per_expert_alignment, cudaStream_t stream); +void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor token_counts, + NVTEEpLayerConfig layer_cfg, cudaStream_t stream); -/*! \brief Dispatch tokens (and routing weights) to expert ranks. +/*! \brief Dispatch tokens (and routing weights) to expert ranks. Requires a + * prior nvte_ep_prepare. * - * \param[in] handle EP handle (id + mem buffer). + * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] topk_idx [T, top_k] int64 sparse routing indices. * \param[in] tokens [T, hidden_dim] input tokens. * \param[in] tokens_win Optional symmem window for ``tokens``. @@ -122,28 +117,29 @@ void nvte_ep_prepare(NVTEEpHandle handle, NVTETensor topk_idx, NVTETensor token_ * \param[in] recv_topk_weights_win Optional symmem window for ``recv_topk_weights``. * \param[in] stream CUDA stream. */ -void nvte_ep_dispatch(NVTEEpHandle handle, NVTETensor topk_idx, NVTETensor tokens, +void nvte_ep_dispatch(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor tokens, NVTECommWindow tokens_win, NVTETensor topk_weights, NVTECommWindow topk_weights_win, NVTETensor recv_tokens, NVTECommWindow recv_tokens_win, NVTETensor recv_topk_weights, NVTECommWindow recv_topk_weights_win, cudaStream_t stream); -/*! \brief Scatter-sum expert outputs back to originating ranks. Unweighted — +/*! \brief Scatter-sum expert outputs back to originating ranks. Unweighted; * caller must pre-multiply expert_out by recv_topk_weights (and the - * valid-slot mask) before calling. + * valid-slot mask) before calling. Requires a prior nvte_ep_prepare. * - * \param[in] handle EP handle (id + mem buffer). + * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] expert_out [recv_T, hidden_dim] pre-weighted expert outputs. * \param[in] expert_out_win Optional symmem window for ``expert_out``. * \param[out] result [T, hidden_dim] combined output. * \param[in] stream CUDA stream. */ -void nvte_ep_combine(NVTEEpHandle handle, NVTETensor expert_out, NVTECommWindow expert_out_win, +void nvte_ep_combine(NVTETensor handle_mem, NVTETensor expert_out, NVTECommWindow expert_out_win, NVTETensor result, cudaStream_t stream); -/*! \brief Backward of dispatch — routes token and weight grads back to source. +/*! \brief Backward of dispatch; routes token and weight grads back to source. + * Requires a prior nvte_ep_prepare. * - * \param[in] handle EP handle (id + mem buffer). + * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] grad [recv_capacity, hidden_dim] grad w.r.t. recv_tokens. * \param[in] grad_win Optional symmem window for ``grad``. * \param[in] g_recv_topk_weights [recv_capacity] f32 grad w.r.t. recv_topk_weights. @@ -152,21 +148,22 @@ void nvte_ep_combine(NVTEEpHandle handle, NVTETensor expert_out, NVTECommWindow * \param[out] grad_topk_weights [T, top_k] f32 grad w.r.t. topk_weights. * \param[in] stream CUDA stream. */ -void nvte_ep_dispatch_bwd(NVTEEpHandle handle, NVTETensor grad, NVTECommWindow grad_win, +void nvte_ep_dispatch_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow grad_win, NVTETensor g_recv_topk_weights, NVTECommWindow g_recv_topk_weights_win, NVTETensor grad_tokens, NVTETensor grad_topk_weights, cudaStream_t stream); /*! \brief Backward of combine. Padded slots in grad_expert_out are zeroed. + * Requires a prior nvte_ep_prepare. * - * \param[in] handle EP handle (id + mem buffer). + * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] grad [T, hidden_dim] grad w.r.t. result. * \param[in] grad_win Optional symmem window for ``grad``. * \param[out] grad_expert_out [recv_capacity, hidden_dim] grad w.r.t. expert_out. * \param[in] grad_expert_out_win Optional symmem window for ``grad_expert_out``. * \param[in] stream CUDA stream. */ -void nvte_ep_combine_bwd(NVTEEpHandle handle, NVTETensor grad, NVTECommWindow grad_win, +void nvte_ep_combine_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow grad_win, NVTETensor grad_expert_out, NVTECommWindow grad_expert_out_win, cudaStream_t stream); From d04bf5f7246cf262f5578971c8a3f92254d66e58 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 11:26:28 -0700 Subject: [PATCH 10/55] bump nccl to latest v0.1 Signed-off-by: Phuong Nguyen --- 3rdparty/nccl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/nccl b/3rdparty/nccl index 146496ac88..b245138bf6 160000 --- a/3rdparty/nccl +++ b/3rdparty/nccl @@ -1 +1 @@ -Subproject commit 146496ac881bc504ed1a52be0ae7b707ce41e706 +Subproject commit b245138bf6ccb6c2b1f41a723e7b17c5e3b7c28b From e3995845a9550d65f8323acf398fe14d21d18c0d Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 16:46:27 -0700 Subject: [PATCH 11/55] tests/cpp_distributed: drop unused NCCL EP header include path Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/CMakeLists.txt | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/tests/cpp_distributed/CMakeLists.txt b/tests/cpp_distributed/CMakeLists.txt index 191dde5d2d..79e4a8337a 100644 --- a/tests/cpp_distributed/CMakeLists.txt +++ b/tests/cpp_distributed/CMakeLists.txt @@ -103,16 +103,8 @@ find_library(NCCL_EP_LIB NO_DEFAULT_PATH REQUIRED) -set(NCCL_EP_INCLUDE_DIR "${NCCL_EP_SUBMODULE_ROOT}/contrib/nccl_ep/include") -if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h") - message(FATAL_ERROR - "NCCL EP header not found at ${NCCL_EP_INCLUDE_DIR}/nccl_ep.h. " - "Run `git submodule update --init --recursive` to checkout 3rdparty/nccl.") -endif() -message(STATUS "EP test: NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}") - -# Collect NCCL include dirs shared by all EP test targets (nccl_ep.h + nccl.h). -set(EP_TEST_NCCL_INCLUDES ${NCCL_EP_INCLUDE_DIR}) +# Tests use TE's public wrapper, not nccl_ep.h. +set(EP_TEST_NCCL_INCLUDES "") if(DEFINED NCCL_INCLUDE_DIR) list(APPEND EP_TEST_NCCL_INCLUDES ${NCCL_INCLUDE_DIR}) message(STATUS "EP test: NCCL headers: ${NCCL_INCLUDE_DIR}") From 14e51090fb24df31769a7787171ea3ae324bc3ed Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 16:58:00 -0700 Subject: [PATCH 12/55] common/ep: fold nvte_ep_* stubs into ep_api.cpp under #if NVTE_WITH_NCCL_EP Signed-off-by: Phuong Nguyen --- transformer_engine/common/CMakeLists.txt | 5 +- transformer_engine/common/ep/ep_api.cpp | 62 ++++++++++++++++++- transformer_engine/common/ep/ep_api_stub.cpp | 63 -------------------- 3 files changed, 63 insertions(+), 67 deletions(-) delete mode 100644 transformer_engine/common/ep/ep_api_stub.cpp diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 18c4af7b09..863fbe5118 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -518,12 +518,13 @@ set_target_properties(transformer_engine PROPERTIES target_sources(transformer_engine PRIVATE ep/ep_backend.cpp ep/ep_api.cpp) +target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_NCCL_EP) message(STATUS "NCCL EP enabled: ${NCCL_EP_LIB}") message(STATUS "NCCL EP include: ${NCCL_EP_INCLUDE_DIR}") else() - # NCCL EP off: export throwing nvte_ep_* stubs so framework bindings link. - target_sources(transformer_engine PRIVATE ep/ep_api_stub.cpp) + # NCCL EP off: ep_api.cpp's #else branch exports throwing nvte_ep_* stubs. + target_sources(transformer_engine PRIVATE ep/ep_api.cpp) message(STATUS "NCCL EP disabled (NVTE_WITH_NCCL_EP=OFF) — using nvte_ep_* stubs") endif() diff --git a/transformer_engine/common/ep/ep_api.cpp b/transformer_engine/common/ep/ep_api.cpp index 51d5af77d0..1f29af743d 100644 --- a/transformer_engine/common/ep/ep_api.cpp +++ b/transformer_engine/common/ep/ep_api.cpp @@ -6,13 +6,20 @@ /*! \file ep_api.cpp * \brief nvte_ep_* C API: thin delegations to the EPBackend singleton. + * + * When NVTE_WITH_NCCL_EP is undefined, the entry points become throwing + * stubs so framework bindings still link without NCCL EP support. */ -#include #include -#include "../common.h" #include "../util/logging.h" + +#if defined(NVTE_WITH_NCCL_EP) + +#include + +#include "../common.h" #include "ep_backend.h" using transformer_engine::ep::EPBackend; @@ -70,3 +77,54 @@ void nvte_ep_combine_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow EPBackend::get().combine_bwd(handle_mem_ptr(handle_mem), grad, grad_win, grad_expert_out, grad_expert_out_win, stream); } + +#else // !NVTE_WITH_NCCL_EP — throwing stubs. + +namespace { +[[noreturn]] void ep_not_built() { + NVTE_ERROR( + "NCCL EP is not built into this TransformerEngine. Rebuild TE with " + "NVTE_BUILD_WITH_NCCL_EP=1 and CUDA arch >= 90 (e.g. NVTE_CUDA_ARCHS=\"90\")."); +} +} // namespace + +void nvte_ep_initialize(void* /*ep_comm*/, NVTEEpGroupConfig /*group_config*/) { ep_not_built(); } + +void nvte_ep_shutdown(void) {} + +size_t nvte_ep_handle_mem_size(NVTEEpLayerConfig /*layer_cfg*/) { ep_not_built(); } + +void nvte_ep_prepare(NVTETensor /*handle_mem*/, NVTETensor /*topk_idx*/, + NVTETensor /*token_counts*/, NVTEEpLayerConfig /*layer_cfg*/, + cudaStream_t /*stream*/) { + ep_not_built(); +} + +void nvte_ep_dispatch(NVTETensor /*handle_mem*/, NVTETensor /*topk_idx*/, NVTETensor /*tokens*/, + NVTECommWindow /*tokens_win*/, NVTETensor /*topk_weights*/, + NVTECommWindow /*topk_weights_win*/, NVTETensor /*recv_tokens*/, + NVTECommWindow /*recv_tokens_win*/, NVTETensor /*recv_topk_weights*/, + NVTECommWindow /*recv_topk_weights_win*/, cudaStream_t /*stream*/) { + ep_not_built(); +} + +void nvte_ep_combine(NVTETensor /*handle_mem*/, NVTETensor /*expert_out*/, + NVTECommWindow /*expert_out_win*/, NVTETensor /*result*/, + cudaStream_t /*stream*/) { + ep_not_built(); +} + +void nvte_ep_dispatch_bwd(NVTETensor /*handle_mem*/, NVTETensor /*grad*/, + NVTECommWindow /*grad_win*/, NVTETensor /*g_recv_topk_weights*/, + NVTECommWindow /*g_recv_topk_weights_win*/, NVTETensor /*grad_tokens*/, + NVTETensor /*grad_topk_weights*/, cudaStream_t /*stream*/) { + ep_not_built(); +} + +void nvte_ep_combine_bwd(NVTETensor /*handle_mem*/, NVTETensor /*grad*/, + NVTECommWindow /*grad_win*/, NVTETensor /*grad_expert_out*/, + NVTECommWindow /*grad_expert_out_win*/, cudaStream_t /*stream*/) { + ep_not_built(); +} + +#endif // NVTE_WITH_NCCL_EP diff --git a/transformer_engine/common/ep/ep_api_stub.cpp b/transformer_engine/common/ep/ep_api_stub.cpp deleted file mode 100644 index a62416cc7f..0000000000 --- a/transformer_engine/common/ep/ep_api_stub.cpp +++ /dev/null @@ -1,63 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -/*! \file ep_api_stub.cpp - * \brief Throwing nvte_ep_* stubs compiled when NVTE_WITH_NCCL_EP=OFF. - */ - -#include - -#include "../util/logging.h" - -namespace { -[[noreturn]] void ep_not_built() { - NVTE_ERROR( - "NCCL EP is not built into this TransformerEngine. Rebuild TE with " - "NVTE_BUILD_WITH_NCCL_EP=1 and CUDA arch >= 90 (e.g. NVTE_CUDA_ARCHS=\"90\")."); -} -} // namespace - -void nvte_ep_initialize(void* /*ep_comm*/, NVTEEpGroupConfig /*group_config*/) { ep_not_built(); } - -void nvte_ep_shutdown(void) {} - -size_t nvte_ep_handle_mem_size(NVTEEpLayerConfig /*layer_cfg*/) { ep_not_built(); } - -void nvte_ep_prepare(NVTETensor /*handle_mem*/, NVTETensor /*topk_idx*/, - NVTETensor /*token_counts*/, NVTEEpLayerConfig /*layer_cfg*/, - cudaStream_t /*stream*/) { - ep_not_built(); -} - -void nvte_ep_dispatch(NVTETensor /*handle_mem*/, NVTETensor /*topk_idx*/, NVTETensor /*tokens*/, - NVTECommWindow /*tokens_win*/, NVTETensor /*topk_weights*/, - NVTECommWindow /*topk_weights_win*/, NVTETensor /*recv_tokens*/, - NVTECommWindow /*recv_tokens_win*/, NVTETensor /*recv_topk_weights*/, - NVTECommWindow /*recv_topk_weights_win*/, NVTEEpLayerConfig /*layer_cfg*/, - cudaStream_t /*stream*/) { - ep_not_built(); -} - -void nvte_ep_combine(NVTETensor /*handle_mem*/, NVTETensor /*expert_out*/, - NVTECommWindow /*expert_out_win*/, NVTETensor /*result*/, - NVTEEpLayerConfig /*layer_cfg*/, cudaStream_t /*stream*/) { - ep_not_built(); -} - -void nvte_ep_dispatch_bwd(NVTETensor /*handle_mem*/, NVTETensor /*grad*/, - NVTECommWindow /*grad_win*/, NVTETensor /*g_recv_topk_weights*/, - NVTECommWindow /*g_recv_topk_weights_win*/, NVTETensor /*grad_tokens*/, - NVTETensor /*grad_topk_weights*/, NVTEEpLayerConfig /*layer_cfg*/, - cudaStream_t /*stream*/) { - ep_not_built(); -} - -void nvte_ep_combine_bwd(NVTETensor /*handle_mem*/, NVTETensor /*grad*/, - NVTECommWindow /*grad_win*/, NVTETensor /*grad_expert_out*/, - NVTECommWindow /*grad_expert_out_win*/, NVTEEpLayerConfig /*layer_cfg*/, - cudaStream_t /*stream*/) { - ep_not_built(); -} From cc34089e1f8f3fe9d246f448676df65b8ada4786 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 19:32:05 -0700 Subject: [PATCH 13/55] common/ep: dlopen libnccl_ep.so so libtransformer_engine.so loads without it Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/CMakeLists.txt | 16 +--- transformer_engine/common/CMakeLists.txt | 16 +++- transformer_engine/common/ep/ep_backend.cpp | 32 ++++---- .../common/ep/ep_nccl_loader.cpp | 76 +++++++++++++++++++ transformer_engine/common/ep/ep_nccl_loader.h | 48 ++++++++++++ 5 files changed, 157 insertions(+), 31 deletions(-) create mode 100644 transformer_engine/common/ep/ep_nccl_loader.cpp create mode 100644 transformer_engine/common/ep/ep_nccl_loader.h diff --git a/tests/cpp_distributed/CMakeLists.txt b/tests/cpp_distributed/CMakeLists.txt index 79e4a8337a..a2ffa82bde 100644 --- a/tests/cpp_distributed/CMakeLists.txt +++ b/tests/cpp_distributed/CMakeLists.txt @@ -94,16 +94,9 @@ gtest_discover_tests(test_comm_gemm DISCOVERY_TIMEOUT 600) # ── EP distributed tests ────────────────────────────────────────────────────── # Launched via mpirun; ncclUniqueId exchange uses MPI_Bcast (see test_ep_common.h). -# Headers + libs come from the in-tree 3rdparty/nccl submodule build. -set(NCCL_EP_SUBMODULE_ROOT - "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl") -find_library(NCCL_EP_LIB - NAMES nccl_ep libnccl_ep - HINTS ${NCCL_EP_SUBMODULE_ROOT}/build/lib - NO_DEFAULT_PATH - REQUIRED) - -# Tests use TE's public wrapper, not nccl_ep.h. +# The test binary only uses NCCL core symbols (ncclMemAlloc, ncclCommWindow*); +# all ncclEp* calls live behind TE's public , which +# resolves libnccl_ep.so via dlopen in libtransformer_engine.so itself. set(EP_TEST_NCCL_INCLUDES "") if(DEFINED NCCL_INCLUDE_DIR) list(APPEND EP_TEST_NCCL_INCLUDES ${NCCL_INCLUDE_DIR}) @@ -126,7 +119,6 @@ set(EP_TEST_COMMON_LIBS ${TE_LIB} CUDA::nvrtc ${NCCL_LIB} - ${NCCL_EP_LIB} MPI::MPI_CXX OpenMP::OpenMP_CXX) @@ -137,4 +129,4 @@ target_link_libraries(test_ep PUBLIC ${EP_TEST_COMMON_LIBS}) # Do NOT use gtest_discover_tests — these binaries require multi-process # launch via run_test_ep.sh, not direct single-process execution. -message(STATUS "EP distributed tests enabled: ${NCCL_EP_LIB}") +message(STATUS "EP distributed tests enabled (TE backend dlopens libnccl_ep.so)") diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 863fbe5118..9aa099ebed 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -458,6 +458,10 @@ endif() message(STATUS "NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}") # ── libnccl_ep.so ────────────────────────────────────────────────────────── +# Resolved at runtime via dlopen in ep/ep_nccl_loader.cpp — NOT link-time bound, +# so libtransformer_engine.so still loads on systems missing libnccl_ep.so or +# with too-old NCCL. Locate the build artifact only to keep it on the rpath +# (so dlopen by SONAME finds the bundled copy via DT_RUNPATH). set(NCCL_EP_LIB_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/lib") find_library(NCCL_EP_LIB NAMES nccl_ep libnccl_ep @@ -506,9 +510,12 @@ target_include_directories(transformer_engine PRIVATE ${NCCL_EP_INCLUDE_DIR} ${NCCL_INCLUDE_DIRS_FOR_TE}) # covers nccl.h + nccl_device/ +# libnccl_ep.so is dlopen'd from ep_nccl_loader.cpp, so do NOT link it here. +# libnccl.so stays direct-linked: only ancient symbols (ncclGetVersion, +# ncclCommCount, ncclGetErrorString) are referenced from this TU. target_link_libraries(transformer_engine PUBLIC - ${NCCL_EP_LIB} - ${NCCL_LIB}) + ${NCCL_LIB} + ${CMAKE_DL_LIBS}) # Embed rpath so the installed wheel finds libnccl_ep.so at runtime. # libnccl.so is already on the system via the Toolkit — no rpath needed for it. @@ -517,10 +524,11 @@ set_target_properties(transformer_engine PROPERTIES target_sources(transformer_engine PRIVATE ep/ep_backend.cpp - ep/ep_api.cpp) + ep/ep_api.cpp + ep/ep_nccl_loader.cpp) target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_NCCL_EP) -message(STATUS "NCCL EP enabled: ${NCCL_EP_LIB}") +message(STATUS "NCCL EP enabled (dlopen at runtime): ${NCCL_EP_LIB}") message(STATUS "NCCL EP include: ${NCCL_EP_INCLUDE_DIR}") else() # NCCL EP off: ep_api.cpp's #else branch exports throwing nvte_ep_* stubs. diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index ae7c0900d6..4e2555389e 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -19,6 +19,7 @@ #include "../common.h" #include "../util/cuda_runtime.h" #include "../util/logging.h" +#include "ep_nccl_loader.h" namespace transformer_engine { namespace ep { @@ -135,15 +136,16 @@ void EPBackend::shutdown() { EPBackend& inst = instance(); std::lock_guard lock(inst.mutex_); if (!inst.initialized_) return; + const auto& nccl = loader::fns(); for (auto& e : inst.lru_) { - if (e.handle != nullptr) ncclEpHandleDestroy(e.handle); + if (e.handle != nullptr) nccl.HandleDestroy(e.handle); } inst.lru_.clear(); inst.index_.clear(); inst.fallback_layer_cfg_.reset(); // ncclEpGroupDestroy reads from ep_comm_; destroy group while comm is still alive. if (inst.ep_group_ != nullptr) { - ncclEpGroupDestroy(inst.ep_group_); + nccl.GroupDestroy(inst.ep_group_); inst.ep_group_ = nullptr; } inst.ep_comm_ = nullptr; // borrowed; caller destroys @@ -185,8 +187,8 @@ ncclEpHandle_t EPBackend::open_handle(void* handle_mem, size_t handle_mem_size, ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; hcfg.dispatch_output_per_expert_alignment = dispatch_output_per_expert_alignment; ncclEpHandle_t handle; - NVTE_CHECK_NCCL(ncclEpInitHandle(&handle, ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, num_topk, - &routing_desc)); + NVTE_CHECK_NCCL(loader::fns().InitHandle(&handle, ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, + num_topk, &routing_desc)); return handle; } @@ -227,7 +229,7 @@ void EPBackend::init(ncclComm_t ep_comm, NVTEEpGroupConfig group_config) { // Must be > 0; NCCL EP errors out on 0. cfg.max_recv_tokens_per_rank = static_cast(group_config.max_recv_tokens_per_rank); - NVTE_CHECK_NCCL(ncclEpCreateGroup(&ep_group_, ep_comm, &cfg)); + NVTE_CHECK_NCCL(loader::fns().CreateGroup(&ep_group_, ep_comm, &cfg)); ep_comm_ = ep_comm; @@ -286,15 +288,15 @@ ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerCon ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; hcfg.dispatch_output_per_expert_alignment = layer_cfg.dispatch_output_per_expert_alignment; size_t hm_size = 0; - NVTE_CHECK_NCCL(ncclEpHandleMemSize(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, &hm_size, - layer_cfg.top_k)); + NVTE_CHECK_NCCL(loader::fns().HandleMemSize(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, + &hm_size, layer_cfg.top_k)); ncclEpHandle_t h = open_handle(handle_mem, hm_size, layer_cfg.top_k, layer_cfg.dispatch_output_per_expert_alignment); lru_.push_front(HandleEntry{handle_mem, h, layer_cfg, hm_size}); index_.emplace(handle_mem, lru_.begin()); while (lru_.size() > cache_cap_locked()) { HandleEntry& victim = lru_.back(); - if (victim.handle != nullptr) ncclEpHandleDestroy(victim.handle); + if (victim.handle != nullptr) loader::fns().HandleDestroy(victim.handle); index_.erase(victim.handle_mem); lru_.pop_back(); } @@ -327,8 +329,8 @@ size_t EPBackend::handle_mem_size(NVTEEpLayerConfig layer_cfg) { ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; hcfg.dispatch_output_per_expert_alignment = layer_cfg.dispatch_output_per_expert_alignment; size_t hm_size = 0; - NVTE_CHECK_NCCL(ncclEpHandleMemSize(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, &hm_size, - layer_cfg.top_k)); + NVTE_CHECK_NCCL(loader::fns().HandleMemSize(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, + &hm_size, layer_cfg.top_k)); return hm_size; } @@ -362,7 +364,7 @@ void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, NVTETensor std::lock_guard lock(mutex_); ncclEpHandle_t h = prepare_handle_locked(handle_mem, layer_cfg); - NVTE_CHECK_NCCL(ncclEpUpdateHandle(h, &nccl_topk_idx, &layout_info, stream)); + NVTE_CHECK_NCCL(loader::fns().UpdateHandle(h, &nccl_topk_idx, &layout_info, stream)); } void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTETensor tokens, @@ -439,8 +441,8 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE std::lock_guard lock(mutex_); ncclEpHandle_t h = lookup_handle_locked(handle_mem); - NVTE_CHECK_NCCL(ncclEpDispatch(h, &in_struct, &out_struct, - /*layout_info=*/nullptr, &dispatch_cfg, stream)); + NVTE_CHECK_NCCL(loader::fns().Dispatch(h, &in_struct, &out_struct, + /*layout_info=*/nullptr, &dispatch_cfg, stream)); } void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, @@ -473,7 +475,7 @@ void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, std::lock_guard lock(mutex_); ncclEpHandle_t h = lookup_handle_locked(handle_mem); - NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, /*config=*/nullptr, stream)); + NVTE_CHECK_NCCL(loader::fns().Combine(h, &in_struct, &out_struct, /*config=*/nullptr, stream)); } void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, @@ -523,7 +525,7 @@ void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, std::lock_guard lock(mutex_); ncclEpHandle_t h = lookup_handle_locked(handle_mem); - NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, &cfg, stream)); + NVTE_CHECK_NCCL(loader::fns().Combine(h, &in_struct, &out_struct, &cfg, stream)); } void EPBackend::combine_bwd(void* handle_mem, const NVTETensor grad, diff --git a/transformer_engine/common/ep/ep_nccl_loader.cpp b/transformer_engine/common/ep/ep_nccl_loader.cpp new file mode 100644 index 0000000000..20c9e6f8bf --- /dev/null +++ b/transformer_engine/common/ep/ep_nccl_loader.cpp @@ -0,0 +1,76 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "ep_nccl_loader.h" + +#include + +#include "../util/logging.h" + +namespace transformer_engine { +namespace ep { +namespace loader { + +namespace { + +constexpr const char* kSonames[] = {"libnccl_ep.so.0", "libnccl_ep.so"}; + +void* try_dlopen(std::string& last_err) { + for (const char* name : kSonames) { + dlerror(); + void* h = dlopen(name, RTLD_LAZY | RTLD_LOCAL); + if (h != nullptr) return h; + if (const char* e = dlerror()) last_err = e; + } + return nullptr; +} + +template +Fn resolve(void* lib, const char* sym) { + dlerror(); + void* p = dlsym(lib, sym); + const char* err = dlerror(); + NVTE_CHECK(err == nullptr && p != nullptr, + "libnccl_ep.so is loaded but symbol '", sym, "' could not be resolved", + (err != nullptr ? std::string(": ") + err : std::string{}), + ". The runtime libnccl_ep.so is older than the version TransformerEngine " + "was built against; upgrade NCCL EP or rebuild TE with -DNVTE_WITH_NCCL_EP=OFF."); + return reinterpret_cast(p); +} + +NcclEpFns load_or_throw() { + std::string last_err; + void* lib = try_dlopen(last_err); + NVTE_CHECK(lib != nullptr, + "Failed to load libnccl_ep.so (", + (last_err.empty() ? "no error message" : last_err), + "). NCCL EP requires libnccl_ep.so (>= 0.0.1) and NCCL >= 2.30.4 at runtime. " + "Install the NCCL EP shared library, or rebuild TransformerEngine with " + "-DNVTE_WITH_NCCL_EP=OFF to disable EP support."); + NcclEpFns fns{}; + fns.InitHandle = resolve(lib, "ncclEpInitHandle"); + fns.CreateGroup = resolve(lib, "ncclEpCreateGroup"); + fns.GroupDestroy = resolve(lib, "ncclEpGroupDestroy"); + fns.HandleDestroy = resolve(lib, "ncclEpHandleDestroy"); + fns.HandleMemSize = resolve(lib, "ncclEpHandleMemSize"); + fns.UpdateHandle = resolve(lib, "ncclEpUpdateHandle"); + fns.Dispatch = resolve(lib, "ncclEpDispatch"); + fns.Combine = resolve(lib, "ncclEpCombine"); + return fns; +} + +} // namespace + +const NcclEpFns& fns() { + // Function-local static: thread-safe one-shot init; re-throws on every call + // if initialization fails, so a missing library is surfaced consistently. + static const NcclEpFns table = load_or_throw(); + return table; +} + +} // namespace loader +} // namespace ep +} // namespace transformer_engine diff --git a/transformer_engine/common/ep/ep_nccl_loader.h b/transformer_engine/common/ep/ep_nccl_loader.h new file mode 100644 index 0000000000..8ffb437ed8 --- /dev/null +++ b/transformer_engine/common/ep/ep_nccl_loader.h @@ -0,0 +1,48 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file ep_nccl_loader.h + * \brief Lazy dlopen-based resolver for libnccl_ep.so. + * + * libtransformer_engine.so is not link-time bound to libnccl_ep.so. The first + * call to ep::loader::fns() opens it via dlopen and dlsyms the ncclEp* + * entry points the EP backend uses. If the library or any symbol cannot be + * resolved (e.g. libnccl_ep.so is missing, or system NCCL is older than the + * EP minimum so libnccl_ep.so's own DT_NEEDED chain fails), the call throws + * NVTE_ERROR with remediation instead of preventing libtransformer_engine.so + * from loading. + */ + +#ifndef TRANSFORMER_ENGINE_COMMON_EP_EP_NCCL_LOADER_H_ +#define TRANSFORMER_ENGINE_COMMON_EP_EP_NCCL_LOADER_H_ + +#include + +namespace transformer_engine { +namespace ep { +namespace loader { + +struct NcclEpFns { + decltype(&::ncclEpInitHandle) InitHandle; + decltype(&::ncclEpCreateGroup) CreateGroup; + decltype(&::ncclEpGroupDestroy) GroupDestroy; + decltype(&::ncclEpHandleDestroy) HandleDestroy; + decltype(&::ncclEpHandleMemSize) HandleMemSize; + decltype(&::ncclEpUpdateHandle) UpdateHandle; + decltype(&::ncclEpDispatch) Dispatch; + decltype(&::ncclEpCombine) Combine; +}; + +/*! \brief Resolve libnccl_ep.so on first call; cache the table thereafter. + * Thread-safe; throws NVTE_ERROR if the library or any symbol is missing. + */ +const NcclEpFns& fns(); + +} // namespace loader +} // namespace ep +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_COMMON_EP_EP_NCCL_LOADER_H_ From 09ffadce7e01eda939bc5be6334e333cc5d4080d Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 19:50:55 -0700 Subject: [PATCH 14/55] common/ep: add BUILD_RPATH=NCCL_EP_LIB_DIR for in-tree dev builds Signed-off-by: Phuong Nguyen --- transformer_engine/common/CMakeLists.txt | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 9aa099ebed..a1acdb2d60 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -517,9 +517,14 @@ target_link_libraries(transformer_engine PUBLIC ${NCCL_LIB} ${CMAKE_DL_LIBS}) -# Embed rpath so the installed wheel finds libnccl_ep.so at runtime. +# Embed rpath so dlopen("libnccl_ep.so.0") finds the bundled lib via DT_RUNPATH: +# - BUILD_RPATH: covers the in-tree build artifact (CMake no longer auto-adds +# NCCL_EP_LIB_DIR since libnccl_ep is not in target_link_libraries anymore). +# - INSTALL_RPATH: $ORIGIN covers the wheel layout (libnccl_ep.so sits next to +# libtransformer_engine.so); NCCL_EP_LIB_DIR is a dev fallback. # libnccl.so is already on the system via the Toolkit — no rpath needed for it. set_target_properties(transformer_engine PROPERTIES + BUILD_RPATH "${NCCL_EP_LIB_DIR}" INSTALL_RPATH "$ORIGIN;${NCCL_EP_LIB_DIR}") target_sources(transformer_engine PRIVATE From d1340d7865c9bd47f5ec18a52e9d632b9f1eae3d Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 19:55:31 -0700 Subject: [PATCH 15/55] common/ep: polish ep.h docstrings; drop unused NVTE_CHECK_NCCL from logging.h Signed-off-by: Phuong Nguyen --- .../common/include/transformer_engine/ep.h | 91 ++++++++++++------- transformer_engine/common/util/logging.h | 8 -- 2 files changed, 58 insertions(+), 41 deletions(-) diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index b18862bb44..7ddeea6761 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -33,16 +33,20 @@ extern "C" { /*! \brief Group-level EP configuration (fixed for the EP group lifetime). */ typedef struct { - int ep_size; /*!< EP world size. */ - int num_experts; /*!< Total experts across all ranks. */ - int max_tokens_per_rank; /*!< Upper bound on tokens this rank sends per dispatch. */ - /*! Upper bound on tokens received per dispatch (worst-case top_k fan-out; must be > 0). */ + /*! EP world size. */ + int ep_size; + /*! Total experts across all ranks. */ + int num_experts; + /*! Upper bound on tokens this rank sends per dispatch. */ + int max_tokens_per_rank; + /*! Upper bound on tokens this rank receives per dispatch (must be > 0). */ int max_recv_tokens_per_rank; - int hidden_dim; /*!< Token hidden dimension. */ - int max_num_sms; /*!< Max SMs for EP kernels. 0 = auto. */ - /*! Widest token dtype the group will dispatch. Sizes NCCL EP staging buffers - * at group create. Tensors passed to nvte_ep_dispatch may use any dtype whose - * element size is <= sizeof(max_token_dtype). */ + /*! Token hidden dimension. */ + int hidden_dim; + /*! Max SMs for EP kernels. 0 = auto. */ + int max_num_sms; + /*! Widest token dtype the group will dispatch; sizes staging buffers. + * Per-dispatch tensors may use any dtype with element size <= this. */ NVTEDType max_token_dtype; } NVTEEpGroupConfig; @@ -51,22 +55,23 @@ typedef struct { * overflow policy, ...). */ typedef struct { - int top_k; /*!< Per-token expert fan-out (> 0). */ - /*! Per-expert zone alignment in tokens (pow2; 0/1 = none). */ + /*! Per-token expert fan-out (> 0). */ + int top_k; + /*! Per-expert recv-slab alignment in tokens (power of two; 0/1 disables). + * When > 1, each expert's slab in recv_tokens is zero-padded up to a + * multiple of this for downstream per-expert GEMM alignment. */ size_t dispatch_output_per_expert_alignment; } NVTEEpLayerConfig; /* -- Bootstrap ------------------------------------------------------------ */ -/*! \brief Bootstrap from an existing NCCL EP sub-communicator. Requires SM>=90. +/*! \brief Bootstrap the EP backend from an existing NCCL EP sub-communicator. + * Requires SM>=90. * - * ep_comm is borrowed and must span exactly group_config.ep_size ranks. - * The caller retains ownership and must keep ep_comm alive until - * nvte_ep_shutdown() returns; destroying it earlier is undefined behavior. - * Re-init after shutdown is allowed; double-init throws. - * - * One EP group per process, bound to the current CUDA device at initialize - * time. Multiple GPUs per process are not supported. + * ep_comm is borrowed and must span exactly group_config.ep_size ranks. The + * caller retains ownership and must keep it alive until nvte_ep_shutdown() + * returns. Re-init after shutdown is allowed; double-init throws. One EP + * group per process, bound to the current CUDA device. * * \param[in] ep_comm Opaque ncclComm_t for the EP sub-group. * \param[in] group_config Group-level EP configuration. @@ -79,9 +84,11 @@ void nvte_ep_shutdown(void); /* -- Layer sizing (host-only) --------------------------------------------- */ /*! \brief Report the handle_mem byte size required for the given layer config. - * Host-only; cheap to call. The caller allocates the buffer and passes - * it back to ep ops as the handle_mem argument. top_k comes from the - * active NVTEEpGroupConfig. + * + * handle_mem is a per-layer kByte routing-state buffer; allocate once and + * thread the same pointer through every prepare/dispatch/combine/_bwd call + * for that layer (the backend keys its cache on the pointer). Host-only; + * size is stable for a given (group, layer) pair. * * \param[in] layer_cfg Per-call layer configuration. * \return size in bytes for the handle_mem buffer. @@ -90,8 +97,13 @@ size_t nvte_ep_handle_mem_size(NVTEEpLayerConfig layer_cfg); /* -- Per-step ops (all allocation-free, CUDA graph-capturable) ------------ */ -/*! \brief AllGather the routing map; write per-expert counts and cache routing - * metadata in handle_mem for the subsequent dispatch/combine. +/*! \brief Seed handle_mem with this step's routing plan. + * + * AllGathers topk_idx across the EP group and stages per-expert offsets and + * counts into handle_mem so the matching dispatch/combine/_bwd can run with + * no further routing computation. Must precede every dispatch/combine/_bwd + * that uses this handle_mem. token_counts becomes host-valid after a stream + * sync. * * \param[in] handle_mem uint8 routing-state buffer. * \param[in] topk_idx [T, top_k] int64 routing indices. @@ -102,8 +114,12 @@ size_t nvte_ep_handle_mem_size(NVTEEpLayerConfig layer_cfg); void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor token_counts, NVTEEpLayerConfig layer_cfg, cudaStream_t stream); -/*! \brief Dispatch tokens (and routing weights) to expert ranks. Requires a - * prior nvte_ep_prepare. +/*! \brief Dispatch tokens (and routing weights) to expert ranks. + * + * Each local token is sent to its top_k destinations; recv_tokens is laid out + * expert-major (contiguous per-expert slabs, padded per layer_cfg). The + * *_win arguments enable zero-copy via symmem windows; pass NVTECommWindow{} + * when unused. Requires a prior nvte_ep_prepare on this handle_mem. * * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] topk_idx [T, top_k] int64 sparse routing indices. @@ -123,9 +139,12 @@ void nvte_ep_dispatch(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor tok NVTECommWindow recv_tokens_win, NVTETensor recv_topk_weights, NVTECommWindow recv_topk_weights_win, cudaStream_t stream); -/*! \brief Scatter-sum expert outputs back to originating ranks. Unweighted; - * caller must pre-multiply expert_out by recv_topk_weights (and the - * valid-slot mask) before calling. Requires a prior nvte_ep_prepare. +/*! \brief Scatter-sum expert outputs back to originating ranks. + * + * Inverse of dispatch: the top_k destination slots for token t are summed + * into result[t]. Sums are unweighted: pre-scale expert_out by + * recv_topk_weights (and the valid-slot mask) before calling. Requires a + * prior nvte_ep_prepare on this handle_mem. * * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] expert_out [recv_T, hidden_dim] pre-weighted expert outputs. @@ -136,8 +155,11 @@ void nvte_ep_dispatch(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor tok void nvte_ep_combine(NVTETensor handle_mem, NVTETensor expert_out, NVTECommWindow expert_out_win, NVTETensor result, cudaStream_t stream); -/*! \brief Backward of dispatch; routes token and weight grads back to source. - * Requires a prior nvte_ep_prepare. +/*! \brief Backward of dispatch: route per-recv-slot grads back to source. + * + * Sums the top_k recv-slot grads into grad_tokens[t]; scatters per-slot + * recv-weight grads into grad_topk_weights[t, k]. Padded recv slots + * contribute nothing. Requires a prior nvte_ep_prepare on this handle_mem. * * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] grad [recv_capacity, hidden_dim] grad w.r.t. recv_tokens. @@ -153,8 +175,11 @@ void nvte_ep_dispatch_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow NVTETensor grad_tokens, NVTETensor grad_topk_weights, cudaStream_t stream); -/*! \brief Backward of combine. Padded slots in grad_expert_out are zeroed. - * Requires a prior nvte_ep_prepare. +/*! \brief Backward of combine: replicate each source-token grad to its recv + * slots from the forward. + * + * Padded recv slots in grad_expert_out are zeroed. Requires a prior + * nvte_ep_prepare on this handle_mem. * * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] grad [T, hidden_dim] grad w.r.t. result. diff --git a/transformer_engine/common/util/logging.h b/transformer_engine/common/util/logging.h index 3308bd22e4..da8b9b377d 100644 --- a/transformer_engine/common/util/logging.h +++ b/transformer_engine/common/util/logging.h @@ -98,14 +98,6 @@ } \ } while (false) -#define NVTE_CHECK_NCCL(expr) \ - do { \ - const ncclResult_t status_NVTE_CHECK_NCCL = (expr); \ - if (status_NVTE_CHECK_NCCL != ncclSuccess) { \ - NVTE_ERROR("NCCL Error: ", ncclGetErrorString(status_NVTE_CHECK_NCCL)); \ - } \ - } while (false) - #ifdef NVTE_WITH_CUBLASMP #define NVTE_CHECK_CUBLASMP(expr) \ From ac75685ab0e104434b54f781ad6532f94bbbbd25 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 19:57:42 -0700 Subject: [PATCH 16/55] common/ep: expose zero_copy in NVTEEpGroupConfig; map to NCCL_EP_ZERO_COPY_{ON,OFF} Signed-off-by: Phuong Nguyen --- transformer_engine/common/ep/ep_backend.cpp | 1 + transformer_engine/common/include/transformer_engine/ep.h | 6 +++++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index 4e2555389e..5167600754 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -228,6 +228,7 @@ void EPBackend::init(ncclComm_t ep_comm, NVTEEpGroupConfig group_config) { : NCCL_EP_AUTO; // Must be > 0; NCCL EP errors out on 0. cfg.max_recv_tokens_per_rank = static_cast(group_config.max_recv_tokens_per_rank); + cfg.zero_copy = group_config.zero_copy ? NCCL_EP_ZERO_COPY_ON : NCCL_EP_ZERO_COPY_OFF; NVTE_CHECK_NCCL(loader::fns().CreateGroup(&ep_group_, ep_comm, &cfg)); diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index 7ddeea6761..5682e9fdb6 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -48,6 +48,10 @@ typedef struct { /*! Widest token dtype the group will dispatch; sizes staging buffers. * Per-dispatch tensors may use any dtype with element size <= this. */ NVTEDType max_token_dtype; + /*! Zero-copy dispatch/combine. When nonzero, payload tensors must be backed + * by NVTECommWindow handles and transfer in place (no staging copies); + * 0 (default) = staged. */ + int zero_copy; } NVTEEpGroupConfig; /*! \brief Per-layer configuration consumed by nvte_ep_handle_mem_size and @@ -142,7 +146,7 @@ void nvte_ep_dispatch(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor tok /*! \brief Scatter-sum expert outputs back to originating ranks. * * Inverse of dispatch: the top_k destination slots for token t are summed - * into result[t]. Sums are unweighted: pre-scale expert_out by + * into result[t]. Sums are unweighted; pre-scale expert_out by * recv_topk_weights (and the valid-slot mask) before calling. Requires a * prior nvte_ep_prepare on this handle_mem. * From 871db30d4d17894005aed8399c9d0a6d939af931 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 20:06:15 -0700 Subject: [PATCH 17/55] tests/cpp_distributed: exercise zero_copy=ON in EPZeroCopyTest.IdentityAllSymm Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/test_ep.cu | 17 ++++++++++++++++- tests/cpp_distributed/test_ep_common.h | 16 ++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu index 7f40b36530..5c2b65a0e9 100644 --- a/tests/cpp_distributed/test_ep.cu +++ b/tests/cpp_distributed/test_ep.cu @@ -702,7 +702,17 @@ static inline NVTECommWindow symm_window(const SymmBuf& b) { } // namespace -class EPZeroCopyTest : public EpOpTestBase {}; +// The symm-window path needs the EP backend bootstrapped with zero_copy=ON +// (so dispatch-output / combine-input must be window-backed). Tests do the +// HBM reference phase under the suite default (OFF) and rebootstrap to ON for +// the symm phase via ep_reinitialize(); TearDown restores OFF for the rest +// of the suite. +class EPZeroCopyTest : public EpOpTestBase { + protected: + void TearDown() override { + if (g_ep_initialized) ep_reinitialize(/*zero_copy=*/0); + } +}; // Identity round-trip with symm-mem on dispatch i/o + combine input. Bit-exact // vs HBM reference (same routing, same input). @@ -733,6 +743,11 @@ TEST_F(EPZeroCopyTest, IdentityAllSymm) { NVTE_CHECK_CUDA(cudaMemcpy(ref_result.data(), ref_buf.result.get(), ref_result.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + // Switch backend to zero_copy=ON for the symm phase. The HBM ref outputs + // are already host-side; cached handles tied to ref_t are evicted by the + // shutdown inside ep_reinitialize. + ep_reinitialize(/*zero_copy=*/1); + // Symm-mem run: tokens, recv_tokens, combine_input (== recv_tokens) all symm. EPBuffers<> sym_buf; // alloc all buffers except the symm ones. sym_buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, diff --git a/tests/cpp_distributed/test_ep_common.h b/tests/cpp_distributed/test_ep_common.h index 135a39416e..cfae30012e 100644 --- a/tests/cpp_distributed/test_ep_common.h +++ b/tests/cpp_distributed/test_ep_common.h @@ -167,6 +167,22 @@ static bool ep_bootstrap(int argc, char* argv[]) { return true; } +// Re-bootstrap the EP backend on the existing g_ep_comm with a new zero_copy +// setting. Used by tests that need the symmem zero-copy fast path. +static void ep_reinitialize(int zero_copy) { + if (!g_ep_initialized) return; + nvte_ep_shutdown(); + NVTEEpGroupConfig group_config{}; + group_config.ep_size = g_ep_size; + group_config.num_experts = g_num_experts; + group_config.max_tokens_per_rank = g_max_tokens_per_rank; + group_config.max_recv_tokens_per_rank = g_ep_size * g_max_tokens_per_rank * 2; + group_config.hidden_dim = g_hidden_dim; + group_config.max_token_dtype = g_max_token_dtype; + group_config.zero_copy = zero_copy; + nvte_ep_initialize(static_cast(g_ep_comm), group_config); +} + // Tear down in dependency order: backend's ep_group reads from ep_comm, // so destroy the group first, then the comm. static void ep_teardown() { From 8be7397d41c9a17f1b23e497eae8dd5dc2c91e8a Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 20:07:22 -0700 Subject: [PATCH 18/55] tests/cpp_distributed: tighten EPZeroCopyTest comments Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/test_ep.cu | 11 +++-------- tests/cpp_distributed/test_ep_common.h | 2 +- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu index 5c2b65a0e9..1a67644d06 100644 --- a/tests/cpp_distributed/test_ep.cu +++ b/tests/cpp_distributed/test_ep.cu @@ -702,11 +702,8 @@ static inline NVTECommWindow symm_window(const SymmBuf& b) { } // namespace -// The symm-window path needs the EP backend bootstrapped with zero_copy=ON -// (so dispatch-output / combine-input must be window-backed). Tests do the -// HBM reference phase under the suite default (OFF) and rebootstrap to ON for -// the symm phase via ep_reinitialize(); TearDown restores OFF for the rest -// of the suite. +// Tests rebootstrap the backend to zero_copy=ON for the symm phase via +// ep_reinitialize(); TearDown restores OFF for the rest of the suite. class EPZeroCopyTest : public EpOpTestBase { protected: void TearDown() override { @@ -743,9 +740,7 @@ TEST_F(EPZeroCopyTest, IdentityAllSymm) { NVTE_CHECK_CUDA(cudaMemcpy(ref_result.data(), ref_buf.result.get(), ref_result.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); - // Switch backend to zero_copy=ON for the symm phase. The HBM ref outputs - // are already host-side; cached handles tied to ref_t are evicted by the - // shutdown inside ep_reinitialize. + // Switch backend to zero_copy=ON for the symm phase. ep_reinitialize(/*zero_copy=*/1); // Symm-mem run: tokens, recv_tokens, combine_input (== recv_tokens) all symm. diff --git a/tests/cpp_distributed/test_ep_common.h b/tests/cpp_distributed/test_ep_common.h index cfae30012e..a2c2821528 100644 --- a/tests/cpp_distributed/test_ep_common.h +++ b/tests/cpp_distributed/test_ep_common.h @@ -168,7 +168,7 @@ static bool ep_bootstrap(int argc, char* argv[]) { } // Re-bootstrap the EP backend on the existing g_ep_comm with a new zero_copy -// setting. Used by tests that need the symmem zero-copy fast path. +// setting. static void ep_reinitialize(int zero_copy) { if (!g_ep_initialized) return; nvte_ep_shutdown(); From 1e1329d527a79bffdc98a4485ec6b84693fef406 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 20:09:28 -0700 Subject: [PATCH 19/55] common/CMakeLists: correct NCCL resolution comment (not bundled with CUDA Toolkit) Signed-off-by: Phuong Nguyen --- transformer_engine/common/CMakeLists.txt | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index a1acdb2d60..a9c618fe18 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -470,9 +470,11 @@ find_library(NCCL_EP_LIB REQUIRED) # ── NCCL + GIN headers ───────────────────────────────────────────────────── -# libnccl.so and all GIN headers (ncclGin.h, ncclWindow_t, ncclDevComm_t) -# ship with the base CUDA Toolkit OR the 3rdparty/nccl submodule build -# (preferred when present; auto-built by setup.py via build_nccl_ep_submodule). +# libnccl.so and the GIN headers (ncclGin.h, ncclWindow_t, ncclDevComm_t) are +# resolved from, in order of preference: the 3rdparty/nccl submodule build +# (auto-built by setup.py via build_nccl_ep_submodule), a system NCCL install +# (e.g. apt libnccl2 / libnccl-dev), or the NVIDIA NCCL pip wheel +# (nvidia-nccl-cu1*). NCCL is NOT part of the base CUDA Toolkit. if(NOT NCCL_LIB) find_library(NCCL_LIB NAMES nccl libnccl @@ -522,7 +524,8 @@ target_link_libraries(transformer_engine PUBLIC # NCCL_EP_LIB_DIR since libnccl_ep is not in target_link_libraries anymore). # - INSTALL_RPATH: $ORIGIN covers the wheel layout (libnccl_ep.so sits next to # libtransformer_engine.so); NCCL_EP_LIB_DIR is a dev fallback. -# libnccl.so is already on the system via the Toolkit — no rpath needed for it. +# libnccl.so is resolved from the system NCCL install or NVIDIA NCCL pip wheel +# (the dynamic linker finds it via its own configured paths); no rpath needed. set_target_properties(transformer_engine PROPERTIES BUILD_RPATH "${NCCL_EP_LIB_DIR}" INSTALL_RPATH "$ORIGIN;${NCCL_EP_LIB_DIR}") From 4ae83636d5f48300166409c9a0955fb421661290 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 20:10:35 -0700 Subject: [PATCH 20/55] common/CMakeLists: shorten NCCL/GIN headers comments Signed-off-by: Phuong Nguyen --- transformer_engine/common/CMakeLists.txt | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index a9c618fe18..ebebbc190e 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -470,11 +470,8 @@ find_library(NCCL_EP_LIB REQUIRED) # ── NCCL + GIN headers ───────────────────────────────────────────────────── -# libnccl.so and the GIN headers (ncclGin.h, ncclWindow_t, ncclDevComm_t) are -# resolved from, in order of preference: the 3rdparty/nccl submodule build -# (auto-built by setup.py via build_nccl_ep_submodule), a system NCCL install -# (e.g. apt libnccl2 / libnccl-dev), or the NVIDIA NCCL pip wheel -# (nvidia-nccl-cu1*). NCCL is NOT part of the base CUDA Toolkit. +# libnccl.so + nccl.h: 3rdparty/nccl/build/ if pre-built, else a system NCCL +# install or NVIDIA NCCL pip wheel. setup.py does not build core NCCL. if(NOT NCCL_LIB) find_library(NCCL_LIB NAMES nccl libnccl @@ -524,8 +521,7 @@ target_link_libraries(transformer_engine PUBLIC # NCCL_EP_LIB_DIR since libnccl_ep is not in target_link_libraries anymore). # - INSTALL_RPATH: $ORIGIN covers the wheel layout (libnccl_ep.so sits next to # libtransformer_engine.so); NCCL_EP_LIB_DIR is a dev fallback. -# libnccl.so is resolved from the system NCCL install or NVIDIA NCCL pip wheel -# (the dynamic linker finds it via its own configured paths); no rpath needed. +# libnccl.so: resolved by the dynamic linker via its configured paths; no rpath needed. set_target_properties(transformer_engine PROPERTIES BUILD_RPATH "${NCCL_EP_LIB_DIR}" INSTALL_RPATH "$ORIGIN;${NCCL_EP_LIB_DIR}") From 28d9470b5603e8330984bbb7307de8cb51a7cb35 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 20:22:53 -0700 Subject: [PATCH 21/55] setup,common: bundle libnccl_ep.so.0 next to libtransformer_engine.so for wheel install Signed-off-by: Phuong Nguyen --- .gitignore | 3 +++ setup.py | 13 ++++++++++++- transformer_engine/common/CMakeLists.txt | 11 ++++------- 3 files changed, 19 insertions(+), 8 deletions(-) diff --git a/.gitignore b/.gitignore index 8a627a7e76..7b03d79cc7 100644 --- a/.gitignore +++ b/.gitignore @@ -43,3 +43,6 @@ tensor_dumps/ artifacts/ .DS_Store .claude/ + +# NCCL EP shared library staged by setup.py for wheel packaging. +/transformer_engine/libnccl_ep.so* diff --git a/setup.py b/setup.py index d46690e626..1cef3a1e13 100644 --- a/setup.py +++ b/setup.py @@ -240,6 +240,16 @@ def build_nccl_ep_submodule() -> str: env=env, ) + # Stage libnccl_ep.so.0 alongside libtransformer_engine.so so $ORIGIN-rpath + # finds it in the installed wheel. + soname = "libnccl_ep.so.0" + src = (build_dir / "lib" / soname).resolve() + dst = current_file_path / "transformer_engine" / soname + if dst.is_symlink() or dst.exists(): + dst.unlink() + shutil.copy2(src, dst) + print(f"[NCCL EP] Bundled {dst} ({src.stat().st_size // (1 << 20)} MB)") + # TE's CMake expects nccl.h under 3rdparty/nccl/build/include/ for its # version check. Mirror the top-level host headers from the system NCCL # install — DON'T mirror nccl_device/ because the submodule ships its own @@ -341,7 +351,8 @@ def git_check_submodules() -> None: else: install_requires, test_requires = setup_requirements() ext_modules = [setup_common_extension()] - package_data = {"": ["VERSION.txt"]} + # libnccl_ep.so.0 is staged by build_nccl_ep_submodule(); ship it. + package_data = {"": ["VERSION.txt"], "transformer_engine": ["libnccl_ep.so*"]} include_package_data = True extras_require = {"test": test_requires} diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index ebebbc190e..c1cddd3158 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -516,15 +516,12 @@ target_link_libraries(transformer_engine PUBLIC ${NCCL_LIB} ${CMAKE_DL_LIBS}) -# Embed rpath so dlopen("libnccl_ep.so.0") finds the bundled lib via DT_RUNPATH: -# - BUILD_RPATH: covers the in-tree build artifact (CMake no longer auto-adds -# NCCL_EP_LIB_DIR since libnccl_ep is not in target_link_libraries anymore). -# - INSTALL_RPATH: $ORIGIN covers the wheel layout (libnccl_ep.so sits next to -# libtransformer_engine.so); NCCL_EP_LIB_DIR is a dev fallback. -# libnccl.so: resolved by the dynamic linker via its configured paths; no rpath needed. +# rpath for dlopen("libnccl_ep.so.0"): in-tree build dir for dev, $ORIGIN for +# the wheel (libnccl_ep.so.0 ships beside libtransformer_engine.so). +# libnccl.so: resolved by the dynamic linker via its configured paths. set_target_properties(transformer_engine PROPERTIES BUILD_RPATH "${NCCL_EP_LIB_DIR}" - INSTALL_RPATH "$ORIGIN;${NCCL_EP_LIB_DIR}") + INSTALL_RPATH "$ORIGIN") target_sources(transformer_engine PRIVATE ep/ep_backend.cpp From c307209c87969fb2ad5f4c173e73e9092f484975 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 20:24:17 -0700 Subject: [PATCH 22/55] .gitmodules: drop nccl branch pin and align indentation with other submodules Signed-off-by: Phuong Nguyen --- .gitmodules | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.gitmodules b/.gitmodules index e531c95507..495d8e3fe7 100644 --- a/.gitmodules +++ b/.gitmodules @@ -9,5 +9,4 @@ url = https://github.com/NVIDIA/cutlass.git [submodule "3rdparty/nccl"] path = 3rdparty/nccl - url = https://github.com/NVIDIA/nccl.git - branch = v2.30u1 + url = https://github.com/NVIDIA/nccl.git From 3a0f230b11560cf7d02ec15508d9075b402ba475 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 20:37:53 -0700 Subject: [PATCH 23/55] setup: gate NCCL EP build on arch >= 90 or native; drop sm_90 fallback Signed-off-by: Phuong Nguyen --- setup.py | 44 +++++++++++++++++--------------------------- 1 file changed, 17 insertions(+), 27 deletions(-) diff --git a/setup.py b/setup.py index 1cef3a1e13..9c90ca7afd 100644 --- a/setup.py +++ b/setup.py @@ -84,29 +84,18 @@ def setup_common_extension() -> CMakeExtension: cusolvermp_dir = os.getenv("CUSOLVERMP_HOME", "/usr") cmake_flags.append(f"-DCUSOLVERMP_DIR={cusolvermp_dir}") - # NCCL EP: on by default; auto-disabled if no arch >= 90. - # Set NVTE_BUILD_WITH_NCCL_EP=0/1 to force off/on. - nccl_ep_env = os.getenv("NVTE_BUILD_WITH_NCCL_EP") - explicit_nccl_ep = nccl_ep_env is not None - build_with_nccl_ep = bool(int(nccl_ep_env)) if explicit_nccl_ep else True - + # NCCL EP (Hopper+): on by default; auto-skipped when no arch >= 90 is + # targeted. Set NVTE_BUILD_WITH_NCCL_EP=0 to force off. + build_with_nccl_ep = bool(int(os.getenv("NVTE_BUILD_WITH_NCCL_EP", "1"))) if build_with_nccl_ep: arch_tokens = [a.strip() for a in str(archs or "").split(";") if a.strip()] - has_hopper_or_newer = any(t.lower() == "native" for t in arch_tokens) or any( - int(t.rstrip("af")) >= 90 for t in arch_tokens if t.rstrip("af").isdigit() + has_hopper_or_newer = any( + t.lower() == "native" or (t.rstrip("af").isdigit() and int(t.rstrip("af")) >= 90) + for t in arch_tokens ) if not has_hopper_or_newer: - if explicit_nccl_ep: - raise RuntimeError( - "NVTE_BUILD_WITH_NCCL_EP=1 requires at least one CUDA arch >= 90 in " - f"NVTE_CUDA_ARCHS (got '{archs}'). Add '90' or unset NVTE_BUILD_WITH_NCCL_EP." - ) - print( - "[NCCL EP] No CUDA arch >= 90 in NVTE_CUDA_ARCHS" - f" ('{archs}'); auto-disabling NCCL EP (nvte_ep_* will throw at runtime)." - ) + print(f"[NCCL EP] No arch >= 90 in NVTE_CUDA_ARCHS ('{archs}'); skipping build.") build_with_nccl_ep = False - if build_with_nccl_ep: build_nccl_ep_submodule() else: @@ -213,15 +202,16 @@ def build_nccl_ep_submodule() -> str: build_dir = nccl_root / "build" nccl_ep_lib = build_dir / "lib" / "libnccl_ep.so" - archs = cuda_archs() or "90" - arch_list = [] - for a in str(archs).split(";"): - a = a.strip().rstrip("af") - if a and a.isdigit() and int(a) >= 90: - arch_list.append(a) - if not arch_list: - arch_list = ["90"] - gencode = " ".join(f"-gencode=arch=compute_{a},code=sm_{a}" for a in arch_list) + # Caller gates on arch >= 90 or "native"; let nvcc resolve "native". + arch_tokens = [a.strip() for a in str(cuda_archs() or "").split(";") if a.strip()] + if any(t.lower() == "native" for t in arch_tokens): + gencode = "-arch=native" + else: + arch_list = [ + t.rstrip("af") for t in arch_tokens + if t.rstrip("af").isdigit() and int(t.rstrip("af")) >= 90 + ] + gencode = " ".join(f"-gencode=arch=compute_{a},code=sm_{a}" for a in arch_list) nproc = os.cpu_count() or 8 env = os.environ.copy() From 6ec55bb5b4ac25ed8d242dbf5b3f08dc8f1adeae Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 20:54:04 -0700 Subject: [PATCH 24/55] common,setup,tests: discover nccl.h via find_path/NCCL_INCLUDE_DIR; drop submodule header mirror Signed-off-by: Phuong Nguyen --- setup.py | 28 ++++----------- tests/cpp_distributed/CMakeLists.txt | 33 ++++++----------- transformer_engine/common/CMakeLists.txt | 46 ++++++++++-------------- 3 files changed, 35 insertions(+), 72 deletions(-) diff --git a/setup.py b/setup.py index 9c90ca7afd..92c6c70672 100644 --- a/setup.py +++ b/setup.py @@ -97,7 +97,8 @@ def setup_common_extension() -> CMakeExtension: print(f"[NCCL EP] No arch >= 90 in NVTE_CUDA_ARCHS ('{archs}'); skipping build.") build_with_nccl_ep = False if build_with_nccl_ep: - build_nccl_ep_submodule() + nccl_home = build_nccl_ep_submodule() + cmake_flags.append(f"-DNCCL_INCLUDE_DIR={nccl_home}/include") else: cmake_flags.append("-DNVTE_WITH_NCCL_EP=OFF") @@ -189,8 +190,9 @@ def _discover_nccl_home() -> str: def build_nccl_ep_submodule() -> str: """Build libnccl_ep.so from the 3rdparty/nccl submodule. - NCCL EP is on by default; the system NCCL core (libnccl.so) supplies the - headers and runtime symbols. Returns the submodule build directory. + Returns the discovered NCCL core install prefix (the path that contains + include/nccl.h and lib/libnccl.so), which the caller passes to CMake as + NCCL_INCLUDE_DIR for TE's own NCCL link. """ nccl_root = current_file_path / "3rdparty" / "nccl" if not (nccl_root / "Makefile").exists(): @@ -240,25 +242,7 @@ def build_nccl_ep_submodule() -> str: shutil.copy2(src, dst) print(f"[NCCL EP] Bundled {dst} ({src.stat().st_size // (1 << 20)} MB)") - # TE's CMake expects nccl.h under 3rdparty/nccl/build/include/ for its - # version check. Mirror the top-level host headers from the system NCCL - # install — DON'T mirror nccl_device/ because the submodule ships its own - # newer copy at src/include/nccl_device/ with device-side templates that - # conflict with older system versions, and the JIT include path picks the - # submodule's. - nccl_include = build_dir / "include" - nccl_include.mkdir(parents=True, exist_ok=True) - for cand in (Path(nccl_home) / "include", Path("/usr/include")): - p = Path(cand) - if (p / "nccl.h").exists(): - for name in ("nccl.h", "nccl_net.h", "nccl_tuner.h"): - src = p / name - dst = nccl_include / name - if src.exists() and not dst.exists(): - dst.symlink_to(src) - break - - return str(build_dir) + return nccl_home def git_check_submodules() -> None: diff --git a/tests/cpp_distributed/CMakeLists.txt b/tests/cpp_distributed/CMakeLists.txt index a2ffa82bde..8660a2baff 100644 --- a/tests/cpp_distributed/CMakeLists.txt +++ b/tests/cpp_distributed/CMakeLists.txt @@ -56,12 +56,16 @@ find_package(CUDAToolkit REQUIRED) find_package(OpenMP REQUIRED) find_package(MPI REQUIRED) -# ── NCCL library ────────────────────────────────────────────────────────────── -# Search order: NCCL_HOME env → 3rdparty/nccl submodule build → system paths. -set(NCCL_SUBMODULE_BUILD "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl/build") +# ── NCCL core ──────────────────────────────────────────────────────────────── +# Pass -DNCCL_INCLUDE_DIR=/include; falls back to well-known prefixes. +find_path(NCCL_INCLUDE_DIR nccl.h + HINTS /opt/nvidia/nccl/include /usr/local/nccl/include) +if(NOT NCCL_INCLUDE_DIR) + message(FATAL_ERROR + "nccl.h not found. Pass -DNCCL_INCLUDE_DIR=/include.") +endif() find_library(NCCL_LIB NAMES nccl libnccl - HINTS $ENV{NCCL_HOME}/lib ${NCCL_SUBMODULE_BUILD}/lib PATH_SUFFIXES lib lib64 REQUIRED) list(APPEND test_comm_gemm_LINKER_LIBS @@ -77,18 +81,6 @@ target_link_libraries(test_comm_gemm PUBLIC ${test_comm_gemm_LINKER_LIBS}) target_compile_options(test_comm_gemm PRIVATE -O2 -fopenmp) -# NCCL headers: prefer submodule build output (has the handle_init API), -# then submodule src, then system (CUDA toolkit). -set(NCCL_SUBMODULE_INCLUDE "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl/build/include") -set(NCCL_SUBMODULE_SRC_INCLUDE "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl/src/include") -if(EXISTS "${NCCL_SUBMODULE_INCLUDE}/nccl.h") - set(NCCL_INCLUDE_DIR "${NCCL_SUBMODULE_INCLUDE}") -elseif(EXISTS "${NCCL_SUBMODULE_SRC_INCLUDE}/nccl.h") - set(NCCL_INCLUDE_DIR "${NCCL_SUBMODULE_SRC_INCLUDE}") -elseif(DEFINED ENV{NCCL_HOME}) - set(NCCL_INCLUDE_DIR "$ENV{NCCL_HOME}/include") -endif() - include(GoogleTest) gtest_discover_tests(test_comm_gemm DISCOVERY_TIMEOUT 600) @@ -97,14 +89,9 @@ gtest_discover_tests(test_comm_gemm DISCOVERY_TIMEOUT 600) # The test binary only uses NCCL core symbols (ncclMemAlloc, ncclCommWindow*); # all ncclEp* calls live behind TE's public , which # resolves libnccl_ep.so via dlopen in libtransformer_engine.so itself. -set(EP_TEST_NCCL_INCLUDES "") -if(DEFINED NCCL_INCLUDE_DIR) - list(APPEND EP_TEST_NCCL_INCLUDES ${NCCL_INCLUDE_DIR}) - message(STATUS "EP test: NCCL headers: ${NCCL_INCLUDE_DIR}") -endif() - +message(STATUS "EP test: NCCL headers: ${NCCL_INCLUDE_DIR}") set(EP_TEST_COMMON_INCLUDES - ${EP_TEST_NCCL_INCLUDES} + ${NCCL_INCLUDE_DIR} ${MPI_CXX_INCLUDE_PATH} ../../transformer_engine/common/include ../../transformer_engine/common diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index c1cddd3158..b4862f9b67 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -469,45 +469,37 @@ find_library(NCCL_EP_LIB NO_DEFAULT_PATH REQUIRED) -# ── NCCL + GIN headers ───────────────────────────────────────────────────── -# libnccl.so + nccl.h: 3rdparty/nccl/build/ if pre-built, else a system NCCL -# install or NVIDIA NCCL pip wheel. setup.py does not build core NCCL. +# ── NCCL core: nccl.h + libnccl.so ───────────────────────────────────────── +# setup.py passes -DNCCL_INCLUDE_DIR; standalone CMake falls back to probing +# well-known NCCL install prefixes. +find_path(NCCL_INCLUDE_DIR nccl.h + HINTS /opt/nvidia/nccl/include /usr/local/nccl/include) +if(NOT NCCL_INCLUDE_DIR) + message(FATAL_ERROR + "nccl.h not found. Pass -DNCCL_INCLUDE_DIR=/include.") +endif() if(NOT NCCL_LIB) find_library(NCCL_LIB NAMES nccl libnccl - HINTS ${NCCL_EP_LIB_DIR} ${CUDAToolkit_LIBRARY_DIR} PATH_SUFFIXES lib lib64 REQUIRED) endif() -set(NCCL_SUBMODULE_INCLUDE - "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl/build/include") -if(EXISTS "${NCCL_SUBMODULE_INCLUDE}/nccl.h") - set(NCCL_INCLUDE_DIRS_FOR_TE ${NCCL_SUBMODULE_INCLUDE}) -else() - set(NCCL_INCLUDE_DIRS_FOR_TE ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}) -endif() - # Diagnostic: log detected NCCL header version (minimum enforced at runtime). -find_file(_nvte_nccl_header_path nccl.h - PATHS ${NCCL_INCLUDE_DIRS_FOR_TE} ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES} - NO_DEFAULT_PATH) -if(_nvte_nccl_header_path) - file(READ "${_nvte_nccl_header_path}" _nvte_nccl_h) - string(REGEX MATCH "#define[ \t]+NCCL_MAJOR[ \t]+([0-9]+)" _ "${_nvte_nccl_h}") - set(_nvte_nccl_major "${CMAKE_MATCH_1}") - string(REGEX MATCH "#define[ \t]+NCCL_MINOR[ \t]+([0-9]+)" _ "${_nvte_nccl_h}") - set(_nvte_nccl_minor "${CMAKE_MATCH_1}") - string(REGEX MATCH "#define[ \t]+NCCL_PATCH[ \t]+([0-9]+)" _ "${_nvte_nccl_h}") - set(_nvte_nccl_patch "${CMAKE_MATCH_1}") - if(_nvte_nccl_major AND _nvte_nccl_minor AND _nvte_nccl_patch) - message(STATUS "NCCL header: ${_nvte_nccl_header_path} (version ${_nvte_nccl_major}.${_nvte_nccl_minor}.${_nvte_nccl_patch})") - endif() +file(READ "${NCCL_INCLUDE_DIR}/nccl.h" _nvte_nccl_h) +string(REGEX MATCH "#define[ \t]+NCCL_MAJOR[ \t]+([0-9]+)" _ "${_nvte_nccl_h}") +set(_nvte_nccl_major "${CMAKE_MATCH_1}") +string(REGEX MATCH "#define[ \t]+NCCL_MINOR[ \t]+([0-9]+)" _ "${_nvte_nccl_h}") +set(_nvte_nccl_minor "${CMAKE_MATCH_1}") +string(REGEX MATCH "#define[ \t]+NCCL_PATCH[ \t]+([0-9]+)" _ "${_nvte_nccl_h}") +set(_nvte_nccl_patch "${CMAKE_MATCH_1}") +if(_nvte_nccl_major AND _nvte_nccl_minor AND _nvte_nccl_patch) + message(STATUS "NCCL header: ${NCCL_INCLUDE_DIR}/nccl.h (version ${_nvte_nccl_major}.${_nvte_nccl_minor}.${_nvte_nccl_patch})") endif() target_include_directories(transformer_engine PRIVATE ${NCCL_EP_INCLUDE_DIR} - ${NCCL_INCLUDE_DIRS_FOR_TE}) # covers nccl.h + nccl_device/ + ${NCCL_INCLUDE_DIR}) # libnccl_ep.so is dlopen'd from ep_nccl_loader.cpp, so do NOT link it here. # libnccl.so stays direct-linked: only ancient symbols (ncclGetVersion, From 754cae94f3c477133a5817164ebc502fd89c1419 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 20:59:33 -0700 Subject: [PATCH 25/55] common/ep: simplify make_nccl_ep_tensor to take NVTETensor and optional CommWindow Signed-off-by: Phuong Nguyen --- transformer_engine/common/ep/ep_backend.cpp | 147 +++++--------------- transformer_engine/common/ep/ep_backend.h | 3 + 2 files changed, 39 insertions(+), 111 deletions(-) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index 5167600754..e275334eb6 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -26,33 +26,18 @@ namespace ep { namespace { -// Build a by-value ncclEpTensor_t descriptor. `sizes` is caller-owned and must -// outlive any NCCL EP call that consumes the descriptor. -inline ncclEpTensor_t make_tensor(void* data, unsigned int ndim, ncclDataType_t datatype, - size_t* sizes) { - ncclEpTensor_t t = NCCL_EP_TENSOR_INIT; - t.ndim = ndim; - t.datatype = datatype; - t.data = data; - t.sizes = sizes; - return t; -} - -// Payload descriptor: prefer the symmem window when set, else fall back to the -// NVTETensor's raw device pointer. -inline ncclEpTensor_t make_payload_tensor(const NVTETensor t, const NVTECommWindow& win, - unsigned int ndim, ncclDataType_t datatype, - size_t* sizes) { +inline ncclEpTensor_t make_nccl_ep_tensor(const NVTETensor t, const NVTECommWindow& win = {}) { + NVTEShape shape = nvte_tensor_shape(t); ncclEpTensor_t desc = NCCL_EP_TENSOR_INIT; - desc.ndim = ndim; - desc.datatype = datatype; - desc.sizes = sizes; + desc.ndim = shape.ndim; + desc.sizes = shape.data; + desc.datatype = EPBackend::nvte_dtype_to_nccl(nvte_tensor_type(t)); if (win.window != nullptr) { desc.win_hdl = win.window; desc.win_offset = win.offset; } else { desc.data = nvte_tensor_data(t); - NVTE_CHECK(desc.data != nullptr, "payload tensor data must not be null"); + NVTE_CHECK(desc.data != nullptr, "tensor data must not be null"); } return desc; } @@ -183,7 +168,11 @@ ncclDataType_t EPBackend::nvte_dtype_to_nccl(NVTEDType dtype) { ncclEpHandle_t EPBackend::open_handle(void* handle_mem, size_t handle_mem_size, int num_topk, size_t dispatch_output_per_expert_alignment) { size_t hm_sizes[1] = {handle_mem_size}; - ncclEpTensor_t routing_desc = make_tensor(handle_mem, 1, ncclUint8, hm_sizes); + ncclEpTensor_t routing_desc = NCCL_EP_TENSOR_INIT; + routing_desc.ndim = 1; + routing_desc.datatype = ncclUint8; + routing_desc.data = handle_mem; + routing_desc.sizes = hm_sizes; ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; hcfg.dispatch_output_per_expert_alignment = dispatch_output_per_expert_alignment; ncclEpHandle_t handle; @@ -341,27 +330,15 @@ void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, NVTETensor NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); NVTE_CHECK(layer_cfg.top_k > 0, "top_k must be > 0, got ", layer_cfg.top_k); - NVTEShape idx_shape = nvte_tensor_shape(topk_idx); - void* idx_data = nvte_tensor_data(topk_idx); - NVTE_CHECK(idx_data != nullptr, "topk_idx data must not be null"); - - const size_t num_tokens = idx_shape.data[0]; - const size_t topk_in = idx_shape.ndim > 1 ? idx_shape.data[1] : 1; - const size_t num_local_experts = - static_cast(group_config_.num_experts / group_config_.ep_size); - - size_t idx_sizes[2] = {num_tokens, topk_in}; - ncclEpTensor_t nccl_topk_idx = make_tensor(idx_data, 2, ncclInt64, idx_sizes); + ncclEpTensor_t nccl_topk_idx = make_nccl_ep_tensor(topk_idx); // ncclEpUpdateHandle writes per-expert counts via expert_counters. - size_t cnt_sizes[1] = {num_local_experts}; ncclEpTensor_t token_counts_desc; - void* token_counts_data = (token_counts != nullptr) ? nvte_tensor_data(token_counts) : nullptr; - if (token_counts_data != nullptr) { - token_counts_desc = make_tensor(token_counts_data, 1, ncclInt32, cnt_sizes); + if (token_counts != nullptr) { + token_counts_desc = make_nccl_ep_tensor(token_counts); } ncclEpLayoutInfo_t layout_info = NCCL_EP_LAYOUT_INFO_INIT; - layout_info.expert_counters = (token_counts_data != nullptr) ? &token_counts_desc : nullptr; + layout_info.expert_counters = (token_counts != nullptr) ? &token_counts_desc : nullptr; std::lock_guard lock(mutex_); ncclEpHandle_t h = prepare_handle_locked(handle_mem, layer_cfg); @@ -376,37 +353,11 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); - NVTEShape tok_shape = nvte_tensor_shape(tokens); NVTEDType tok_dtype = nvte_tensor_type(tokens); NVTE_CHECK(typeToSize(static_cast(tok_dtype)) <= typeToSize(static_cast(group_config_.max_token_dtype)), "tokens dtype (", static_cast(tok_dtype), ") wider than group max_token_dtype (", static_cast(group_config_.max_token_dtype), ")"); - - const size_t num_tokens = tok_shape.data[0]; - const size_t hidden_dim = tok_shape.data[1]; - - size_t tok_sizes[2] = {num_tokens, hidden_dim}; - ncclEpTensor_t nccl_tokens_in = - make_payload_tensor(tokens, tokens_win, 2, nvte_dtype_to_nccl(tok_dtype), tok_sizes); - - const bool is_forward = (topk_weights != nullptr); - - // Routing is cached in handle_mem by ep_prepare; dispatch only needs - // topk_weights to reconstruct the sparse-to-dense prob map. - size_t weights_in_sizes[2] = {0, 0}; - ncclEpTensor_t nccl_topk_weights_in; - if (is_forward) { - NVTE_CHECK(topk_idx != nullptr, "topk_idx required in forward dispatch"); - NVTEShape idx_shape = nvte_tensor_shape(topk_idx); - const size_t topk_in = idx_shape.ndim > 1 ? idx_shape.data[1] : 1; - weights_in_sizes[0] = num_tokens; - weights_in_sizes[1] = topk_in; - nccl_topk_weights_in = - make_payload_tensor(topk_weights, topk_weights_win, 2, ncclFloat32, weights_in_sizes); - } - - NVTEShape recv_shape = nvte_tensor_shape(recv_tokens); NVTEDType recv_dtype = nvte_tensor_type(recv_tokens); NVTE_CHECK(typeToSize(static_cast(recv_dtype)) <= typeToSize(static_cast(group_config_.max_token_dtype)), @@ -414,19 +365,22 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE ") wider than group max_token_dtype (", static_cast(group_config_.max_token_dtype), ")"); - size_t recv_sizes[2] = {recv_shape.data[0], recv_shape.data[1]}; - ncclEpTensor_t nccl_tokens_out = make_payload_tensor(recv_tokens, recv_tokens_win, 2, - nvte_dtype_to_nccl(recv_dtype), recv_sizes); + ncclEpTensor_t nccl_tokens_in = make_nccl_ep_tensor(tokens, tokens_win); + ncclEpTensor_t nccl_tokens_out = make_nccl_ep_tensor(recv_tokens, recv_tokens_win); - size_t weights_out_sizes[1] = {recv_shape.data[0]}; + // Routing is cached in handle_mem by ep_prepare; dispatch only needs + // topk_weights to reconstruct the sparse-to-dense prob map. + const bool is_forward = (topk_weights != nullptr); + ncclEpTensor_t nccl_topk_weights_in; ncclEpTensor_t nccl_topk_weights_out; if (is_forward) { + NVTE_CHECK(topk_idx != nullptr, "topk_idx required in forward dispatch"); NVTE_CHECK(recv_topk_weights != nullptr, "recv_topk_weights must not be null in forward dispatch"); - NVTEShape recv_w_shape = nvte_tensor_shape(recv_topk_weights); - NVTE_CHECK(recv_w_shape.ndim == 1, "recv_topk_weights must be 1D [recv_capacity]"); - nccl_topk_weights_out = make_payload_tensor(recv_topk_weights, recv_topk_weights_win, 1, - ncclFloat32, weights_out_sizes); + NVTE_CHECK(nvte_tensor_shape(recv_topk_weights).ndim == 1, + "recv_topk_weights must be 1D [recv_capacity]"); + nccl_topk_weights_in = make_nccl_ep_tensor(topk_weights, topk_weights_win); + nccl_topk_weights_out = make_nccl_ep_tensor(recv_topk_weights, recv_topk_weights_win); } ncclEpDispatchInputs_t in_struct = NCCL_EP_DISPATCH_INPUTS_INIT; @@ -452,21 +406,8 @@ void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); - NVTEShape exp_shape = nvte_tensor_shape(expert_out); - NVTEDType exp_dtype = nvte_tensor_type(expert_out); - - size_t exp_sizes[2] = {exp_shape.data[0], exp_shape.data[1]}; - ncclEpTensor_t nccl_expert_in = - make_payload_tensor(expert_out, expert_out_win, 2, nvte_dtype_to_nccl(exp_dtype), exp_sizes); - - NVTEShape res_shape = nvte_tensor_shape(result); - void* res_data = nvte_tensor_data(result); - NVTEDType res_dtype = nvte_tensor_type(result); - NVTE_CHECK(res_data != nullptr, "result data must not be null"); - - size_t res_sizes[2] = {res_shape.data[0], res_shape.data[1]}; - ncclEpTensor_t nccl_result_out = - make_tensor(res_data, 2, nvte_dtype_to_nccl(res_dtype), res_sizes); + ncclEpTensor_t nccl_expert_in = make_nccl_ep_tensor(expert_out, expert_out_win); + ncclEpTensor_t nccl_result_out = make_nccl_ep_tensor(result); ncclEpCombineInputs_t in_struct = NCCL_EP_COMBINE_INPUTS_INIT; in_struct.tokens = &nccl_expert_in; @@ -486,32 +427,16 @@ void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); - NVTEShape g_shape = nvte_tensor_shape(grad); - NVTEDType g_dtype = nvte_tensor_type(grad); - size_t g_sizes[2] = {g_shape.data[0], g_shape.data[1]}; - ncclEpTensor_t nccl_tok_in = - make_payload_tensor(grad, grad_win, 2, nvte_dtype_to_nccl(g_dtype), g_sizes); - // g_recv_topk_weights must be 1D [recv_capacity]; caller flattens. - NVTEShape gw_shape = nvte_tensor_shape(g_recv_topk_weights); - NVTE_CHECK(gw_shape.ndim == 1, + NVTE_CHECK(nvte_tensor_shape(g_recv_topk_weights).ndim == 1, "g_recv_topk_weights must be 1D [recv_capacity]; caller must flatten leading dims"); - size_t gw_sizes[1] = {gw_shape.data[0]}; - ncclEpTensor_t nccl_w_in = - make_payload_tensor(g_recv_topk_weights, g_recv_topk_weights_win, 1, ncclFloat32, gw_sizes); - - NVTEShape gt_shape = nvte_tensor_shape(grad_tokens); - void* gt_data = nvte_tensor_data(grad_tokens); - NVTE_CHECK(gt_data != nullptr, "grad_tokens data must not be null"); - size_t gt_sizes[2] = {gt_shape.data[0], gt_shape.data[1]}; - ncclEpTensor_t nccl_tok_out = make_tensor(gt_data, 2, nvte_dtype_to_nccl(g_dtype), gt_sizes); - - NVTEShape gtw_shape = nvte_tensor_shape(grad_topk_weights); - void* gtw_data = nvte_tensor_data(grad_topk_weights); - NVTE_CHECK(gtw_data != nullptr, "grad_topk_weights data must not be null"); - NVTE_CHECK(gtw_shape.ndim == 2, "grad_topk_weights must be 2D [T, top_k]"); - size_t gtw_sizes[2] = {gtw_shape.data[0], gtw_shape.data[1]}; - ncclEpTensor_t nccl_w_out = make_tensor(gtw_data, 2, ncclFloat32, gtw_sizes); + NVTE_CHECK(nvte_tensor_shape(grad_topk_weights).ndim == 2, + "grad_topk_weights must be 2D [T, top_k]"); + + ncclEpTensor_t nccl_tok_in = make_nccl_ep_tensor(grad, grad_win); + ncclEpTensor_t nccl_w_in = make_nccl_ep_tensor(g_recv_topk_weights, g_recv_topk_weights_win); + ncclEpTensor_t nccl_tok_out = make_nccl_ep_tensor(grad_tokens); + ncclEpTensor_t nccl_w_out = make_nccl_ep_tensor(grad_topk_weights); ncclEpCombineInputs_t in_struct = NCCL_EP_COMBINE_INPUTS_INIT; in_struct.tokens = &nccl_tok_in; diff --git a/transformer_engine/common/ep/ep_backend.h b/transformer_engine/common/ep/ep_backend.h index 405226646b..616e105c71 100644 --- a/transformer_engine/common/ep/ep_backend.h +++ b/transformer_engine/common/ep/ep_backend.h @@ -85,7 +85,10 @@ class EPBackend { static EPBackend& instance(); // Meyers singleton accessor static void validate_config(const NVTEEpGroupConfig& config); + public: static ncclDataType_t nvte_dtype_to_nccl(NVTEDType dtype); + + private: // Open a fresh ncclEpHandle over handle_mem. num_topk=-1 for paths // that don't carry per-token weights. ncclEpHandle_t open_handle(void* handle_mem, size_t handle_mem_size, int num_topk, From 7b9cadb07db1432aee614e707e67f1f886489977 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 21:09:31 -0700 Subject: [PATCH 26/55] common/ep: move te_dtype_to_nccl_dtype out of EPBackend into anon namespace Signed-off-by: Phuong Nguyen --- transformer_engine/common/ep/ep_backend.cpp | 54 ++++++++++----------- transformer_engine/common/ep/ep_backend.h | 4 -- 2 files changed, 25 insertions(+), 33 deletions(-) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index e275334eb6..423a546c03 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -26,12 +26,36 @@ namespace ep { namespace { +ncclDataType_t te_dtype_to_nccl_dtype(NVTEDType dtype) { + switch (dtype) { + case kNVTEFloat32: + return ncclFloat32; + case kNVTEFloat16: + return ncclFloat16; + case kNVTEBFloat16: + return ncclBfloat16; + case kNVTEInt32: + return ncclInt32; + case kNVTEInt64: + return ncclInt64; + case kNVTEByte: + return ncclUint8; + case kNVTEFloat8E4M3: + return ncclFloat8e4m3; + case kNVTEFloat8E5M2: + return ncclFloat8e5m2; + default: + NVTE_ERROR("Unsupported NVTEDType for NCCL dtype conversion: ", static_cast(dtype)); + } + return ncclFloat32; // unreachable +} + inline ncclEpTensor_t make_nccl_ep_tensor(const NVTETensor t, const NVTECommWindow& win = {}) { NVTEShape shape = nvte_tensor_shape(t); ncclEpTensor_t desc = NCCL_EP_TENSOR_INIT; desc.ndim = shape.ndim; desc.sizes = shape.data; - desc.datatype = EPBackend::nvte_dtype_to_nccl(nvte_tensor_type(t)); + desc.datatype = te_dtype_to_nccl_dtype(nvte_tensor_type(t)); if (win.window != nullptr) { desc.win_hdl = win.window; desc.win_offset = win.offset; @@ -137,34 +161,6 @@ void EPBackend::shutdown() { inst.initialized_ = false; } -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -ncclDataType_t EPBackend::nvte_dtype_to_nccl(NVTEDType dtype) { - switch (dtype) { - case kNVTEFloat32: - return ncclFloat32; - case kNVTEFloat16: - return ncclFloat16; - case kNVTEBFloat16: - return ncclBfloat16; - case kNVTEInt32: - return ncclInt32; - case kNVTEInt64: - return ncclInt64; - case kNVTEByte: - return ncclUint8; - case kNVTEFloat8E4M3: - return ncclFloat8e4m3; - case kNVTEFloat8E5M2: - return ncclFloat8e5m2; - default: - NVTE_ERROR("Unsupported NVTEDType for NCCL EP conversion: ", static_cast(dtype)); - } - return ncclFloat32; // unreachable -} - ncclEpHandle_t EPBackend::open_handle(void* handle_mem, size_t handle_mem_size, int num_topk, size_t dispatch_output_per_expert_alignment) { size_t hm_sizes[1] = {handle_mem_size}; diff --git a/transformer_engine/common/ep/ep_backend.h b/transformer_engine/common/ep/ep_backend.h index 616e105c71..ffb95ab845 100644 --- a/transformer_engine/common/ep/ep_backend.h +++ b/transformer_engine/common/ep/ep_backend.h @@ -85,10 +85,6 @@ class EPBackend { static EPBackend& instance(); // Meyers singleton accessor static void validate_config(const NVTEEpGroupConfig& config); - public: - static ncclDataType_t nvte_dtype_to_nccl(NVTEDType dtype); - - private: // Open a fresh ncclEpHandle over handle_mem. num_topk=-1 for paths // that don't carry per-token weights. ncclEpHandle_t open_handle(void* handle_mem, size_t handle_mem_size, int num_topk, From 8c187e5ddf55efb3123078131926680634897b91 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 21:10:55 -0700 Subject: [PATCH 27/55] common/ep: reword multicast check; drop NVLS framing Signed-off-by: Phuong Nguyen --- transformer_engine/common/ep/ep_backend.cpp | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index 423a546c03..0c7fc95c83 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -111,10 +111,8 @@ void EPBackend::validate_config(const NVTEEpGroupConfig& config) { "but current device has compute capability ", major, ".x"); - // NCCL EP needs CUDA multicast (NVLS); init hangs without it. NVTE_CHECK(cuda::supports_multicast(device), - "NCCL EP requires CUDA multicast (NVLS) support on device ", device, - " but CU_DEVICE_ATTRIBUTE_MULTICAST_SUPPORTED reports 0."); + "NCCL EP requires CUDA multicast support on device ", device); } void EPBackend::initialize(ncclComm_t ep_comm, NVTEEpGroupConfig config) { From 268194366972950395688efe8d6511e469bf8f03 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 21:21:06 -0700 Subject: [PATCH 28/55] common,tests: replace unicode em-dash and box-drawing chars with ASCII in EP files Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/CMakeLists.txt | 8 ++++---- tests/cpp_distributed/run_test_ep.sh | 6 +++--- tests/cpp_distributed/test_ep_common.h | 10 +++++----- transformer_engine/common/CMakeLists.txt | 14 +++++++------- transformer_engine/common/ep/ep_api.cpp | 2 +- 5 files changed, 20 insertions(+), 20 deletions(-) diff --git a/tests/cpp_distributed/CMakeLists.txt b/tests/cpp_distributed/CMakeLists.txt index 8660a2baff..e65c298e15 100644 --- a/tests/cpp_distributed/CMakeLists.txt +++ b/tests/cpp_distributed/CMakeLists.txt @@ -56,7 +56,7 @@ find_package(CUDAToolkit REQUIRED) find_package(OpenMP REQUIRED) find_package(MPI REQUIRED) -# ── NCCL core ──────────────────────────────────────────────────────────────── +# -- NCCL core ---------------------------------------------------------------- # Pass -DNCCL_INCLUDE_DIR=/include; falls back to well-known prefixes. find_path(NCCL_INCLUDE_DIR nccl.h HINTS /opt/nvidia/nccl/include /usr/local/nccl/include) @@ -84,7 +84,7 @@ target_compile_options(test_comm_gemm PRIVATE -O2 -fopenmp) include(GoogleTest) gtest_discover_tests(test_comm_gemm DISCOVERY_TIMEOUT 600) -# ── EP distributed tests ────────────────────────────────────────────────────── +# -- EP distributed tests ------------------------------------------------------ # Launched via mpirun; ncclUniqueId exchange uses MPI_Bcast (see test_ep_common.h). # The test binary only uses NCCL core symbols (ncclMemAlloc, ncclCommWindow*); # all ncclEp* calls live behind TE's public , which @@ -109,11 +109,11 @@ set(EP_TEST_COMMON_LIBS MPI::MPI_CXX OpenMP::OpenMP_CXX) -# ── EP distributed tests (per-op + full pipeline + zero-copy symm) ─────────── +# -- EP distributed tests (per-op + full pipeline + zero-copy symm) ----------- add_executable(test_ep test_ep.cu ../cpp/test_common.cu) target_include_directories(test_ep PRIVATE ${EP_TEST_COMMON_INCLUDES}) target_link_libraries(test_ep PUBLIC ${EP_TEST_COMMON_LIBS}) -# Do NOT use gtest_discover_tests — these binaries require multi-process +# Do NOT use gtest_discover_tests - these binaries require multi-process # launch via run_test_ep.sh, not direct single-process execution. message(STATUS "EP distributed tests enabled (TE backend dlopens libnccl_ep.so)") diff --git a/tests/cpp_distributed/run_test_ep.sh b/tests/cpp_distributed/run_test_ep.sh index 13e86fa02d..1c4432531c 100755 --- a/tests/cpp_distributed/run_test_ep.sh +++ b/tests/cpp_distributed/run_test_ep.sh @@ -14,9 +14,9 @@ # build_dir = /build # # Environment variables: -# GTEST_FILTER — forwarded to all processes (e.g., "EPPipelineTest.*") -# MPIRUN — override the mpirun binary (default: mpirun) -# MPIRUN_EXTRA — extra flags forwarded to mpirun +# GTEST_FILTER - forwarded to all processes (e.g., "EPPipelineTest.*") +# MPIRUN - override the mpirun binary (default: mpirun) +# MPIRUN_EXTRA - extra flags forwarded to mpirun set -euo pipefail diff --git a/tests/cpp_distributed/test_ep_common.h b/tests/cpp_distributed/test_ep_common.h index a2c2821528..b2421ffd10 100644 --- a/tests/cpp_distributed/test_ep_common.h +++ b/tests/cpp_distributed/test_ep_common.h @@ -40,7 +40,7 @@ using transformer_engine::TensorWrapper; NVTE_CHECK(_err_mpi == MPI_SUCCESS, "MPI error: ", _err_mpi); \ } while (false) -// ── Process-level state ─────────────────────────────────────────────────────── +// -- Process-level state ------------------------------------------------------- static int g_process_id = -1; static int g_num_processes = -1; @@ -77,7 +77,7 @@ struct DevBuf { size_t bytes() const { return count * sizeof(T); } }; -// ── Shared routing helper ───────────────────────────────────────────────────── +// -- Shared routing helper ----------------------------------------------------- // Balanced round-robin routing: token t on rank r maps top_k experts to // (r * num_local_experts + t * top_k + k) % num_experts @@ -90,14 +90,14 @@ static inline std::vector routing_balanced( return idx; } -// ── ncclUniqueId exchange via MPI ───────────────────────────────────────────── +// -- ncclUniqueId exchange via MPI --------------------------------------------- static void exchange_unique_id(ncclUniqueId* uid) { if (g_process_id == 0) NVTE_CHECK_NCCL(ncclGetUniqueId(uid)); CHECK_MPI(MPI_Bcast(uid, sizeof(*uid), MPI_BYTE, 0, MPI_COMM_WORLD)); } -// ── CLI parsing ─────────────────────────────────────────────────────────────── +// -- CLI parsing --------------------------------------------------------------- static void ep_parse_args(int argc, char* argv[]) { for (int i = 1; i < argc; ++i) { @@ -107,7 +107,7 @@ static void ep_parse_args(int argc, char* argv[]) { } } -// ── Bootstrap / teardown ────────────────────────────────────────────────────── +// -- Bootstrap / teardown ------------------------------------------------------ // Returns false if the binary should exit without running tests (wrong SM, etc.). static bool ep_bootstrap(int argc, char* argv[]) { diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index b4862f9b67..f7af26a2bf 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -437,14 +437,14 @@ if (NVTE_WITH_CUSOLVERMP) message(STATUS "Using cuSolverMp at: ${CUSOLVERMP_DIR}") endif() -# ── NCCL EP (on by default, HT mode only) ───────────────────────────────── +# -- NCCL EP (on by default, HT mode only) --------------------------------- # Set -DNVTE_WITH_NCCL_EP=OFF (or NVTE_BUILD_WITH_NCCL_EP=0 in setup.py) to -# skip NCCL EP entirely — useful on older images whose system NCCL is below +# skip NCCL EP entirely - useful on older images whose system NCCL is below # the 2.30.4 EP minimum. option(NVTE_WITH_NCCL_EP "Build NCCL EP into libtransformer_engine.so" ON) if(NVTE_WITH_NCCL_EP) # SM>=90 and NCCL>=2.30.4 are gated at runtime in EPBackend::initialize. -# ── NCCL EP headers ──────────────────────────────────────────────────────── +# -- NCCL EP headers -------------------------------------------------------- # Headers + libs are produced by the in-tree 3rdparty/nccl submodule build # (auto-built by setup.py via build_nccl_ep_submodule). set(NCCL_EP_SUBMODULE_ROOT @@ -457,8 +457,8 @@ if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h") endif() message(STATUS "NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}") -# ── libnccl_ep.so ────────────────────────────────────────────────────────── -# Resolved at runtime via dlopen in ep/ep_nccl_loader.cpp — NOT link-time bound, +# -- libnccl_ep.so ---------------------------------------------------------- +# Resolved at runtime via dlopen in ep/ep_nccl_loader.cpp - NOT link-time bound, # so libtransformer_engine.so still loads on systems missing libnccl_ep.so or # with too-old NCCL. Locate the build artifact only to keep it on the rpath # (so dlopen by SONAME finds the bundled copy via DT_RUNPATH). @@ -469,7 +469,7 @@ find_library(NCCL_EP_LIB NO_DEFAULT_PATH REQUIRED) -# ── NCCL core: nccl.h + libnccl.so ───────────────────────────────────────── +# -- NCCL core: nccl.h + libnccl.so ----------------------------------------- # setup.py passes -DNCCL_INCLUDE_DIR; standalone CMake falls back to probing # well-known NCCL install prefixes. find_path(NCCL_INCLUDE_DIR nccl.h @@ -526,7 +526,7 @@ message(STATUS "NCCL EP include: ${NCCL_EP_INCLUDE_DIR}") else() # NCCL EP off: ep_api.cpp's #else branch exports throwing nvte_ep_* stubs. target_sources(transformer_engine PRIVATE ep/ep_api.cpp) - message(STATUS "NCCL EP disabled (NVTE_WITH_NCCL_EP=OFF) — using nvte_ep_* stubs") + message(STATUS "NCCL EP disabled (NVTE_WITH_NCCL_EP=OFF) - using nvte_ep_* stubs") endif() # Number of philox4x32 rounds for stochastic rounding (build-time constant). diff --git a/transformer_engine/common/ep/ep_api.cpp b/transformer_engine/common/ep/ep_api.cpp index 1f29af743d..b8cf04aa4a 100644 --- a/transformer_engine/common/ep/ep_api.cpp +++ b/transformer_engine/common/ep/ep_api.cpp @@ -78,7 +78,7 @@ void nvte_ep_combine_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow grad_expert_out_win, stream); } -#else // !NVTE_WITH_NCCL_EP — throwing stubs. +#else // !NVTE_WITH_NCCL_EP - throwing stubs. namespace { [[noreturn]] void ep_not_built() { From 6cf7cef6decb0500de9c4700e81de738b7607ea0 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 8 Jun 2026 21:26:14 -0700 Subject: [PATCH 29/55] bump nccl to latest v0.1 Signed-off-by: Phuong Nguyen --- 3rdparty/nccl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/nccl b/3rdparty/nccl index b245138bf6..9d22d5dfec 160000 --- a/3rdparty/nccl +++ b/3rdparty/nccl @@ -1 +1 @@ -Subproject commit b245138bf6ccb6c2b1f41a723e7b17c5e3b7c28b +Subproject commit 9d22d5dfec8391ee65b56df139d471f8e08e921e From 71859eb3340bcd37562b0593ba02ff482100bfee Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 9 Jun 2026 04:37:11 +0000 Subject: [PATCH 30/55] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- setup.py | 3 ++- transformer_engine/common/ep/ep_backend.cpp | 13 ++++++------- transformer_engine/common/ep/ep_backend.h | 4 ++-- transformer_engine/common/ep/ep_nccl_loader.cpp | 8 +++----- 4 files changed, 13 insertions(+), 15 deletions(-) diff --git a/setup.py b/setup.py index 92c6c70672..1e3bae2307 100644 --- a/setup.py +++ b/setup.py @@ -210,7 +210,8 @@ def build_nccl_ep_submodule() -> str: gencode = "-arch=native" else: arch_list = [ - t.rstrip("af") for t in arch_tokens + t.rstrip("af") + for t in arch_tokens if t.rstrip("af").isdigit() and int(t.rstrip("af")) >= 90 ] gencode = " ".join(f"-gencode=arch=compute_{a},code=sm_{a}" for a in arch_list) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index 0c7fc95c83..6c73a0d74a 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -111,8 +111,8 @@ void EPBackend::validate_config(const NVTEEpGroupConfig& config) { "but current device has compute capability ", major, ".x"); - NVTE_CHECK(cuda::supports_multicast(device), - "NCCL EP requires CUDA multicast support on device ", device); + NVTE_CHECK(cuda::supports_multicast(device), "NCCL EP requires CUDA multicast support on device ", + device); } void EPBackend::initialize(ncclComm_t ep_comm, NVTEEpGroupConfig config) { @@ -252,8 +252,7 @@ ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerCon // between runs; one cfg per process). Remove this once XLA preserves the // handle_mem device pointer across runs. if (fallback_layer_cfg_.has_value()) { - NVTE_CHECK(fallback_layer_cfg_->top_k == layer_cfg.top_k, - "EP prepare top_k=", layer_cfg.top_k, + NVTE_CHECK(fallback_layer_cfg_->top_k == layer_cfg.top_k, "EP prepare top_k=", layer_cfg.top_k, " disagrees with process-wide cached top_k=", fallback_layer_cfg_->top_k); NVTE_CHECK(fallback_layer_cfg_->dispatch_output_per_expert_alignment == layer_cfg.dispatch_output_per_expert_alignment, @@ -448,9 +447,9 @@ void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, NVTE_CHECK_NCCL(loader::fns().Combine(h, &in_struct, &out_struct, &cfg, stream)); } -void EPBackend::combine_bwd(void* handle_mem, const NVTETensor grad, - const NVTECommWindow& grad_win, NVTETensor grad_expert_out, - const NVTECommWindow& grad_expert_out_win, cudaStream_t stream) { +void EPBackend::combine_bwd(void* handle_mem, const NVTETensor grad, const NVTECommWindow& grad_win, + NVTETensor grad_expert_out, const NVTECommWindow& grad_expert_out_win, + cudaStream_t stream) { // Backward of combine = reverse-direction dispatch. dispatch(handle_mem, /*topk_idx=*/nullptr, grad, grad_win, /*topk_weights=*/nullptr, /*topk_weights_win=*/NVTECommWindow{}, grad_expert_out, diff --git a/transformer_engine/common/ep/ep_backend.h b/transformer_engine/common/ep/ep_backend.h index ffb95ab845..ea9aa019fa 100644 --- a/transformer_engine/common/ep/ep_backend.h +++ b/transformer_engine/common/ep/ep_backend.h @@ -60,8 +60,8 @@ class EPBackend { const NVTECommWindow& recv_tokens_win, NVTETensor recv_topk_weights, const NVTECommWindow& recv_topk_weights_win, cudaStream_t stream); - void combine(void* handle_mem, const NVTETensor expert_out, - const NVTECommWindow& expert_out_win, NVTETensor result, cudaStream_t stream); + void combine(void* handle_mem, const NVTETensor expert_out, const NVTECommWindow& expert_out_win, + NVTETensor result, cudaStream_t stream); // g_recv_topk_weights: 1D [recv_capacity] f32; grad_topk_weights: 2D [T, top_k] f32. void dispatch_bwd(void* handle_mem, const NVTETensor grad, const NVTECommWindow& grad_win, diff --git a/transformer_engine/common/ep/ep_nccl_loader.cpp b/transformer_engine/common/ep/ep_nccl_loader.cpp index 20c9e6f8bf..8374acd7b3 100644 --- a/transformer_engine/common/ep/ep_nccl_loader.cpp +++ b/transformer_engine/common/ep/ep_nccl_loader.cpp @@ -33,9 +33,8 @@ Fn resolve(void* lib, const char* sym) { dlerror(); void* p = dlsym(lib, sym); const char* err = dlerror(); - NVTE_CHECK(err == nullptr && p != nullptr, - "libnccl_ep.so is loaded but symbol '", sym, "' could not be resolved", - (err != nullptr ? std::string(": ") + err : std::string{}), + NVTE_CHECK(err == nullptr && p != nullptr, "libnccl_ep.so is loaded but symbol '", sym, + "' could not be resolved", (err != nullptr ? std::string(": ") + err : std::string{}), ". The runtime libnccl_ep.so is older than the version TransformerEngine " "was built against; upgrade NCCL EP or rebuild TE with -DNVTE_WITH_NCCL_EP=OFF."); return reinterpret_cast(p); @@ -44,8 +43,7 @@ Fn resolve(void* lib, const char* sym) { NcclEpFns load_or_throw() { std::string last_err; void* lib = try_dlopen(last_err); - NVTE_CHECK(lib != nullptr, - "Failed to load libnccl_ep.so (", + NVTE_CHECK(lib != nullptr, "Failed to load libnccl_ep.so (", (last_err.empty() ? "no error message" : last_err), "). NCCL EP requires libnccl_ep.so (>= 0.0.1) and NCCL >= 2.30.4 at runtime. " "Install the NCCL EP shared library, or rebuild TransformerEngine with " From c9ca69e25a1104a436b5ee97e0b14e3684b4a6cb Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 9 Jun 2026 15:43:25 -0700 Subject: [PATCH 31/55] nccl commit to 2.31.0a4-1 Signed-off-by: Phuong Nguyen --- 3rdparty/nccl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/nccl b/3rdparty/nccl index 9d22d5dfec..808d2433dd 160000 --- a/3rdparty/nccl +++ b/3rdparty/nccl @@ -1 +1 @@ -Subproject commit 9d22d5dfec8391ee65b56df139d471f8e08e921e +Subproject commit 808d2433dda3cccc80f8172a94a6b117359e7102 From f9b3029d3d5104c10410dea4a4536c487b5b62a9 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 9 Jun 2026 15:55:04 -0700 Subject: [PATCH 32/55] common/CMakeLists: point NCCL_EP_INCLUDE_DIR at build/include staged headers Signed-off-by: Phuong Nguyen --- transformer_engine/common/CMakeLists.txt | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index f7af26a2bf..d22a572968 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -449,11 +449,12 @@ if(NVTE_WITH_NCCL_EP) # (auto-built by setup.py via build_nccl_ep_submodule). set(NCCL_EP_SUBMODULE_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl") -set(NCCL_EP_INCLUDE_DIR "${NCCL_EP_SUBMODULE_ROOT}/contrib/nccl_ep/include") +set(NCCL_EP_INCLUDE_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/include") if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h") message(FATAL_ERROR "NCCL EP header not found at ${NCCL_EP_INCLUDE_DIR}/nccl_ep.h. " - "Run `git submodule update --init --recursive` to checkout 3rdparty/nccl.") + "setup.py builds 3rdparty/nccl/contrib/nccl_ep/ via make, which stages " + "nccl_ep.h + nccl_ep/ into build/include/.") endif() message(STATUS "NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}") From 041598a9a8a7a79d284ca2a4c4009392b998cb62 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 9 Jun 2026 18:16:51 -0700 Subject: [PATCH 33/55] common/CMakeLists: clarify NCCL EP missing-header instructions Signed-off-by: Phuong Nguyen --- transformer_engine/common/CMakeLists.txt | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index d22a572968..6f5117ef08 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -453,8 +453,7 @@ set(NCCL_EP_INCLUDE_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/include") if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h") message(FATAL_ERROR "NCCL EP header not found at ${NCCL_EP_INCLUDE_DIR}/nccl_ep.h. " - "setup.py builds 3rdparty/nccl/contrib/nccl_ep/ via make, which stages " - "nccl_ep.h + nccl_ep/ into build/include/.") + "Run `git submodule update --init --recursive` and rebuild TE.") endif() message(STATUS "NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}") From f8a2ad99e8aff71c53a738448f2e109d156e9d11 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 9 Jun 2026 18:38:17 -0700 Subject: [PATCH 34/55] common/ep: use int64_t instead of long for handle-cache size env (cpplint runtime/int) Signed-off-by: Phuong Nguyen --- transformer_engine/common/ep/ep_backend.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index 6c73a0d74a..7ebac7d3e5 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -11,6 +11,7 @@ #include "ep_backend.h" #include +#include #include #include #include @@ -228,7 +229,7 @@ size_t EPBackend::cache_cap_locked() { if (handle_cache_cap_ == 0) { const char* cap_env = std::getenv("NVTE_EP_HANDLE_CACHE_SIZE"); if (cap_env != nullptr) { - const long v = std::atol(cap_env); + const int64_t v = static_cast(std::atol(cap_env)); if (v < 0) { // Unlimited cache. WAR for JAX until XLA fixes handle_mem // reloc between runs. From 9125623034321b040ae3bc48706165c8352c94c9 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 9 Jun 2026 19:11:22 -0700 Subject: [PATCH 35/55] common/ep: fix dangling sizes pointer in make_nccl_ep_tensor (NVTEShape lifetime) Signed-off-by: Phuong Nguyen --- transformer_engine/common/ep/ep_backend.cpp | 46 ++++++++++++++------- 1 file changed, 30 insertions(+), 16 deletions(-) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index 7ebac7d3e5..f4d46eac8e 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -51,11 +51,14 @@ ncclDataType_t te_dtype_to_nccl_dtype(NVTEDType dtype) { return ncclFloat32; // unreachable } -inline ncclEpTensor_t make_nccl_ep_tensor(const NVTETensor t, const NVTECommWindow& win = {}) { - NVTEShape shape = nvte_tensor_shape(t); +// shape_out is caller-owned; desc.sizes aliases shape_out.data and must +// outlive the NCCL EP call. +inline ncclEpTensor_t make_nccl_ep_tensor(const NVTETensor t, NVTEShape& shape_out, + const NVTECommWindow& win = {}) { + shape_out = nvte_tensor_shape(t); ncclEpTensor_t desc = NCCL_EP_TENSOR_INIT; - desc.ndim = shape.ndim; - desc.sizes = shape.data; + desc.ndim = shape_out.ndim; + desc.sizes = shape_out.data; desc.datatype = te_dtype_to_nccl_dtype(nvte_tensor_type(t)); if (win.window != nullptr) { desc.win_hdl = win.window; @@ -324,12 +327,14 @@ void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, NVTETensor NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); NVTE_CHECK(layer_cfg.top_k > 0, "top_k must be > 0, got ", layer_cfg.top_k); - ncclEpTensor_t nccl_topk_idx = make_nccl_ep_tensor(topk_idx); + NVTEShape topk_idx_shape; + ncclEpTensor_t nccl_topk_idx = make_nccl_ep_tensor(topk_idx, topk_idx_shape); // ncclEpUpdateHandle writes per-expert counts via expert_counters. + NVTEShape token_counts_shape; ncclEpTensor_t token_counts_desc; if (token_counts != nullptr) { - token_counts_desc = make_nccl_ep_tensor(token_counts); + token_counts_desc = make_nccl_ep_tensor(token_counts, token_counts_shape); } ncclEpLayoutInfo_t layout_info = NCCL_EP_LAYOUT_INFO_INIT; layout_info.expert_counters = (token_counts != nullptr) ? &token_counts_desc : nullptr; @@ -359,12 +364,15 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE ") wider than group max_token_dtype (", static_cast(group_config_.max_token_dtype), ")"); - ncclEpTensor_t nccl_tokens_in = make_nccl_ep_tensor(tokens, tokens_win); - ncclEpTensor_t nccl_tokens_out = make_nccl_ep_tensor(recv_tokens, recv_tokens_win); + NVTEShape tokens_shape, recv_tokens_shape; + ncclEpTensor_t nccl_tokens_in = make_nccl_ep_tensor(tokens, tokens_shape, tokens_win); + ncclEpTensor_t nccl_tokens_out = make_nccl_ep_tensor(recv_tokens, recv_tokens_shape, + recv_tokens_win); // Routing is cached in handle_mem by ep_prepare; dispatch only needs // topk_weights to reconstruct the sparse-to-dense prob map. const bool is_forward = (topk_weights != nullptr); + NVTEShape topk_weights_shape, recv_topk_weights_shape; ncclEpTensor_t nccl_topk_weights_in; ncclEpTensor_t nccl_topk_weights_out; if (is_forward) { @@ -373,8 +381,10 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE "recv_topk_weights must not be null in forward dispatch"); NVTE_CHECK(nvte_tensor_shape(recv_topk_weights).ndim == 1, "recv_topk_weights must be 1D [recv_capacity]"); - nccl_topk_weights_in = make_nccl_ep_tensor(topk_weights, topk_weights_win); - nccl_topk_weights_out = make_nccl_ep_tensor(recv_topk_weights, recv_topk_weights_win); + nccl_topk_weights_in = make_nccl_ep_tensor(topk_weights, topk_weights_shape, + topk_weights_win); + nccl_topk_weights_out = make_nccl_ep_tensor(recv_topk_weights, recv_topk_weights_shape, + recv_topk_weights_win); } ncclEpDispatchInputs_t in_struct = NCCL_EP_DISPATCH_INPUTS_INIT; @@ -400,8 +410,10 @@ void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); - ncclEpTensor_t nccl_expert_in = make_nccl_ep_tensor(expert_out, expert_out_win); - ncclEpTensor_t nccl_result_out = make_nccl_ep_tensor(result); + NVTEShape expert_out_shape, result_shape; + ncclEpTensor_t nccl_expert_in = make_nccl_ep_tensor(expert_out, expert_out_shape, + expert_out_win); + ncclEpTensor_t nccl_result_out = make_nccl_ep_tensor(result, result_shape); ncclEpCombineInputs_t in_struct = NCCL_EP_COMBINE_INPUTS_INIT; in_struct.tokens = &nccl_expert_in; @@ -427,10 +439,12 @@ void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, NVTE_CHECK(nvte_tensor_shape(grad_topk_weights).ndim == 2, "grad_topk_weights must be 2D [T, top_k]"); - ncclEpTensor_t nccl_tok_in = make_nccl_ep_tensor(grad, grad_win); - ncclEpTensor_t nccl_w_in = make_nccl_ep_tensor(g_recv_topk_weights, g_recv_topk_weights_win); - ncclEpTensor_t nccl_tok_out = make_nccl_ep_tensor(grad_tokens); - ncclEpTensor_t nccl_w_out = make_nccl_ep_tensor(grad_topk_weights); + NVTEShape grad_shape, g_recv_w_shape, grad_tokens_shape, grad_w_shape; + ncclEpTensor_t nccl_tok_in = make_nccl_ep_tensor(grad, grad_shape, grad_win); + ncclEpTensor_t nccl_w_in = make_nccl_ep_tensor(g_recv_topk_weights, g_recv_w_shape, + g_recv_topk_weights_win); + ncclEpTensor_t nccl_tok_out = make_nccl_ep_tensor(grad_tokens, grad_tokens_shape); + ncclEpTensor_t nccl_w_out = make_nccl_ep_tensor(grad_topk_weights, grad_w_shape); ncclEpCombineInputs_t in_struct = NCCL_EP_COMBINE_INPUTS_INIT; in_struct.tokens = &nccl_tok_in; From 5064d10a67d358731471815d228df1388d9afe7d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 10 Jun 2026 02:14:55 +0000 Subject: [PATCH 36/55] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/ep/ep_backend.cpp | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index f4d46eac8e..b43b01fa73 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -366,8 +366,8 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE NVTEShape tokens_shape, recv_tokens_shape; ncclEpTensor_t nccl_tokens_in = make_nccl_ep_tensor(tokens, tokens_shape, tokens_win); - ncclEpTensor_t nccl_tokens_out = make_nccl_ep_tensor(recv_tokens, recv_tokens_shape, - recv_tokens_win); + ncclEpTensor_t nccl_tokens_out = + make_nccl_ep_tensor(recv_tokens, recv_tokens_shape, recv_tokens_win); // Routing is cached in handle_mem by ep_prepare; dispatch only needs // topk_weights to reconstruct the sparse-to-dense prob map. @@ -381,10 +381,9 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE "recv_topk_weights must not be null in forward dispatch"); NVTE_CHECK(nvte_tensor_shape(recv_topk_weights).ndim == 1, "recv_topk_weights must be 1D [recv_capacity]"); - nccl_topk_weights_in = make_nccl_ep_tensor(topk_weights, topk_weights_shape, - topk_weights_win); - nccl_topk_weights_out = make_nccl_ep_tensor(recv_topk_weights, recv_topk_weights_shape, - recv_topk_weights_win); + nccl_topk_weights_in = make_nccl_ep_tensor(topk_weights, topk_weights_shape, topk_weights_win); + nccl_topk_weights_out = + make_nccl_ep_tensor(recv_topk_weights, recv_topk_weights_shape, recv_topk_weights_win); } ncclEpDispatchInputs_t in_struct = NCCL_EP_DISPATCH_INPUTS_INIT; @@ -411,8 +410,7 @@ void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); NVTEShape expert_out_shape, result_shape; - ncclEpTensor_t nccl_expert_in = make_nccl_ep_tensor(expert_out, expert_out_shape, - expert_out_win); + ncclEpTensor_t nccl_expert_in = make_nccl_ep_tensor(expert_out, expert_out_shape, expert_out_win); ncclEpTensor_t nccl_result_out = make_nccl_ep_tensor(result, result_shape); ncclEpCombineInputs_t in_struct = NCCL_EP_COMBINE_INPUTS_INIT; @@ -441,8 +439,8 @@ void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, NVTEShape grad_shape, g_recv_w_shape, grad_tokens_shape, grad_w_shape; ncclEpTensor_t nccl_tok_in = make_nccl_ep_tensor(grad, grad_shape, grad_win); - ncclEpTensor_t nccl_w_in = make_nccl_ep_tensor(g_recv_topk_weights, g_recv_w_shape, - g_recv_topk_weights_win); + ncclEpTensor_t nccl_w_in = + make_nccl_ep_tensor(g_recv_topk_weights, g_recv_w_shape, g_recv_topk_weights_win); ncclEpTensor_t nccl_tok_out = make_nccl_ep_tensor(grad_tokens, grad_tokens_shape); ncclEpTensor_t nccl_w_out = make_nccl_ep_tensor(grad_topk_weights, grad_w_shape); From dbfad0b81b085e22ed298e7782bf2d65afdf0b31 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 10 Jun 2026 11:22:12 -0700 Subject: [PATCH 37/55] common/ep: static-link libnccl_ep.a into libtransformer_engine.so; drop dlopen loader and runtime bundling Signed-off-by: Phuong Nguyen --- setup.py | 24 +----- tests/cpp_distributed/CMakeLists.txt | 6 +- transformer_engine/__init__.py | 29 ++++++++ transformer_engine/common/CMakeLists.txt | 39 ++++------ transformer_engine/common/ep/ep_backend.cpp | 34 ++++----- .../common/ep/ep_nccl_loader.cpp | 74 ------------------- transformer_engine/common/ep/ep_nccl_loader.h | 48 ------------ 7 files changed, 69 insertions(+), 185 deletions(-) delete mode 100644 transformer_engine/common/ep/ep_nccl_loader.cpp delete mode 100644 transformer_engine/common/ep/ep_nccl_loader.h diff --git a/setup.py b/setup.py index 1e3bae2307..c9370073b2 100644 --- a/setup.py +++ b/setup.py @@ -188,12 +188,7 @@ def _discover_nccl_home() -> str: def build_nccl_ep_submodule() -> str: - """Build libnccl_ep.so from the 3rdparty/nccl submodule. - - Returns the discovered NCCL core install prefix (the path that contains - include/nccl.h and lib/libnccl.so), which the caller passes to CMake as - NCCL_INCLUDE_DIR for TE's own NCCL link. - """ + """Build libnccl_ep.a from the 3rdparty/nccl submodule and return NCCL_HOME.""" nccl_root = current_file_path / "3rdparty" / "nccl" if not (nccl_root / "Makefile").exists(): raise RuntimeError( @@ -202,7 +197,7 @@ def build_nccl_ep_submodule() -> str: ) build_dir = nccl_root / "build" - nccl_ep_lib = build_dir / "lib" / "libnccl_ep.so" + nccl_ep_lib = build_dir / "lib" / "libnccl_ep.a" # Caller gates on arch >= 90 or "native"; let nvcc resolve "native". arch_tokens = [a.strip() for a in str(cuda_archs() or "").split(";") if a.strip()] @@ -226,23 +221,13 @@ def build_nccl_ep_submodule() -> str: env["NCCL_EP_BUILDDIR"] = str(build_dir) if not nccl_ep_lib.exists(): - print(f"[NCCL EP] Building libnccl_ep.so (gencode='{gencode}')") + print(f"[NCCL EP] Building libnccl_ep.a (gencode='{gencode}')") subprocess.check_call( ["make", "-j", str(nproc), "-C", "contrib/nccl_ep", "lib"], cwd=str(nccl_root), env=env, ) - # Stage libnccl_ep.so.0 alongside libtransformer_engine.so so $ORIGIN-rpath - # finds it in the installed wheel. - soname = "libnccl_ep.so.0" - src = (build_dir / "lib" / soname).resolve() - dst = current_file_path / "transformer_engine" / soname - if dst.is_symlink() or dst.exists(): - dst.unlink() - shutil.copy2(src, dst) - print(f"[NCCL EP] Bundled {dst} ({src.stat().st_size // (1 << 20)} MB)") - return nccl_home @@ -326,8 +311,7 @@ def git_check_submodules() -> None: else: install_requires, test_requires = setup_requirements() ext_modules = [setup_common_extension()] - # libnccl_ep.so.0 is staged by build_nccl_ep_submodule(); ship it. - package_data = {"": ["VERSION.txt"], "transformer_engine": ["libnccl_ep.so*"]} + package_data = {"": ["VERSION.txt"]} include_package_data = True extras_require = {"test": test_requires} diff --git a/tests/cpp_distributed/CMakeLists.txt b/tests/cpp_distributed/CMakeLists.txt index e65c298e15..0ea97eb68d 100644 --- a/tests/cpp_distributed/CMakeLists.txt +++ b/tests/cpp_distributed/CMakeLists.txt @@ -87,8 +87,8 @@ gtest_discover_tests(test_comm_gemm DISCOVERY_TIMEOUT 600) # -- EP distributed tests ------------------------------------------------------ # Launched via mpirun; ncclUniqueId exchange uses MPI_Bcast (see test_ep_common.h). # The test binary only uses NCCL core symbols (ncclMemAlloc, ncclCommWindow*); -# all ncclEp* calls live behind TE's public , which -# resolves libnccl_ep.so via dlopen in libtransformer_engine.so itself. +# all ncclEp* calls live behind TE's public , which is +# statically linked into libtransformer_engine.so. message(STATUS "EP test: NCCL headers: ${NCCL_INCLUDE_DIR}") set(EP_TEST_COMMON_INCLUDES ${NCCL_INCLUDE_DIR} @@ -116,4 +116,4 @@ target_link_libraries(test_ep PUBLIC ${EP_TEST_COMMON_LIBS}) # Do NOT use gtest_discover_tests - these binaries require multi-process # launch via run_test_ep.sh, not direct single-process execution. -message(STATUS "EP distributed tests enabled (TE backend dlopens libnccl_ep.so)") +message(STATUS "EP distributed tests enabled (NCCL EP statically linked into libtransformer_engine.so)") diff --git a/transformer_engine/__init__.py b/transformer_engine/__init__.py index 0175f04e2e..d84fd81360 100644 --- a/transformer_engine/__init__.py +++ b/transformer_engine/__init__.py @@ -6,10 +6,39 @@ # pylint: disable=unused-import +import ctypes +import functools import os from importlib import metadata import transformer_engine.common +# Minimum NCCL version for the statically-linked NCCL EP backend. +_NCCL_EP_MIN_VERSION = (2, 30, 4) + + +@functools.lru_cache(maxsize=1) +def is_nccl_ep_available() -> bool: + """Return True if the runtime libnccl.so is new enough for NCCL EP.""" + try: + libnccl = ctypes.CDLL("libnccl.so.2", mode=ctypes.RTLD_GLOBAL) + ver = ctypes.c_int(0) + libnccl.ncclGetVersion(ctypes.byref(ver)) + except (OSError, AttributeError): + return False + v = ver.value + cur = (v // 10000, (v // 100) % 100, v % 100) + return cur >= _NCCL_EP_MIN_VERSION + + +def require_nccl_ep() -> None: + """Raise RuntimeError if NCCL EP cannot run on the current libnccl.""" + if not is_nccl_ep_available(): + mn = ".".join(str(x) for x in _NCCL_EP_MIN_VERSION) + raise RuntimeError( + f"NCCL EP requires NCCL >= {mn} at runtime; upgrade libnccl.so or " + "rebuild Transformer Engine with NVTE_BUILD_WITH_NCCL_EP=0." + ) + try: from . import pytorch except ImportError: diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 6f5117ef08..dba040efe6 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -457,14 +457,14 @@ if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h") endif() message(STATUS "NCCL EP headers: ${NCCL_EP_INCLUDE_DIR}") -# -- libnccl_ep.so ---------------------------------------------------------- -# Resolved at runtime via dlopen in ep/ep_nccl_loader.cpp - NOT link-time bound, -# so libtransformer_engine.so still loads on systems missing libnccl_ep.so or -# with too-old NCCL. Locate the build artifact only to keep it on the rpath -# (so dlopen by SONAME finds the bundled copy via DT_RUNPATH). +# -- libnccl_ep.a ----------------------------------------------------------- +# Statically linked into libtransformer_engine.so. EPBackend::initialize checks +# NCCL >= 2.30.4 before any nccl_ep call, so the newer NCCL symbols nccl_ep +# imports stay unresolved (and harmless) under default ELF lazy binding when +# the gate trips. LD_BIND_NOW environments lose this property. set(NCCL_EP_LIB_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/lib") -find_library(NCCL_EP_LIB - NAMES nccl_ep libnccl_ep +find_file(NCCL_EP_LIB + NAMES libnccl_ep.a HINTS ${NCCL_EP_LIB_DIR} NO_DEFAULT_PATH REQUIRED) @@ -501,27 +501,20 @@ target_include_directories(transformer_engine PRIVATE ${NCCL_EP_INCLUDE_DIR} ${NCCL_INCLUDE_DIR}) -# libnccl_ep.so is dlopen'd from ep_nccl_loader.cpp, so do NOT link it here. -# libnccl.so stays direct-linked: only ancient symbols (ncclGetVersion, -# ncclCommCount, ncclGetErrorString) are referenced from this TU. -target_link_libraries(transformer_engine PUBLIC - ${NCCL_LIB} - ${CMAKE_DL_LIBS}) - -# rpath for dlopen("libnccl_ep.so.0"): in-tree build dir for dev, $ORIGIN for -# the wheel (libnccl_ep.so.0 ships beside libtransformer_engine.so). -# libnccl.so: resolved by the dynamic linker via its configured paths. -set_target_properties(transformer_engine PROPERTIES - BUILD_RPATH "${NCCL_EP_LIB_DIR}" - INSTALL_RPATH "$ORIGIN") +# libnccl.so direct symbols (ncclGetVersion etc.) come from libnccl_ep.a's +# DT_NEEDED chain plus this TU's own references. CUDA::cuda_driver must follow +# the static archive on the link line so --as-needed records libcuda.so.1. +target_link_libraries(transformer_engine PUBLIC ${NCCL_LIB}) +target_link_libraries(transformer_engine PRIVATE + -Wl,--whole-archive ${NCCL_EP_LIB} -Wl,--no-whole-archive + CUDA::cuda_driver) target_sources(transformer_engine PRIVATE ep/ep_backend.cpp - ep/ep_api.cpp - ep/ep_nccl_loader.cpp) + ep/ep_api.cpp) target_compile_definitions(transformer_engine PRIVATE NVTE_WITH_NCCL_EP) -message(STATUS "NCCL EP enabled (dlopen at runtime): ${NCCL_EP_LIB}") +message(STATUS "NCCL EP enabled (static link): ${NCCL_EP_LIB}") message(STATUS "NCCL EP include: ${NCCL_EP_INCLUDE_DIR}") else() # NCCL EP off: ep_api.cpp's #else branch exports throwing nvte_ep_* stubs. diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index b43b01fa73..b52a1f6a37 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -17,10 +17,11 @@ #include #include +#include + #include "../common.h" #include "../util/cuda_runtime.h" #include "../util/logging.h" -#include "ep_nccl_loader.h" namespace transformer_engine { namespace ep { @@ -147,16 +148,15 @@ void EPBackend::shutdown() { EPBackend& inst = instance(); std::lock_guard lock(inst.mutex_); if (!inst.initialized_) return; - const auto& nccl = loader::fns(); for (auto& e : inst.lru_) { - if (e.handle != nullptr) nccl.HandleDestroy(e.handle); + if (e.handle != nullptr) ncclEpHandleDestroy(e.handle); } inst.lru_.clear(); inst.index_.clear(); inst.fallback_layer_cfg_.reset(); // ncclEpGroupDestroy reads from ep_comm_; destroy group while comm is still alive. if (inst.ep_group_ != nullptr) { - nccl.GroupDestroy(inst.ep_group_); + ncclEpGroupDestroy(inst.ep_group_); inst.ep_group_ = nullptr; } inst.ep_comm_ = nullptr; // borrowed; caller destroys @@ -174,8 +174,8 @@ ncclEpHandle_t EPBackend::open_handle(void* handle_mem, size_t handle_mem_size, ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; hcfg.dispatch_output_per_expert_alignment = dispatch_output_per_expert_alignment; ncclEpHandle_t handle; - NVTE_CHECK_NCCL(loader::fns().InitHandle(&handle, ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, - num_topk, &routing_desc)); + NVTE_CHECK_NCCL(ncclEpInitHandle(&handle, ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, num_topk, + &routing_desc)); return handle; } @@ -217,7 +217,7 @@ void EPBackend::init(ncclComm_t ep_comm, NVTEEpGroupConfig group_config) { cfg.max_recv_tokens_per_rank = static_cast(group_config.max_recv_tokens_per_rank); cfg.zero_copy = group_config.zero_copy ? NCCL_EP_ZERO_COPY_ON : NCCL_EP_ZERO_COPY_OFF; - NVTE_CHECK_NCCL(loader::fns().CreateGroup(&ep_group_, ep_comm, &cfg)); + NVTE_CHECK_NCCL(ncclEpCreateGroup(&ep_group_, ep_comm, &cfg)); ep_comm_ = ep_comm; @@ -275,15 +275,15 @@ ncclEpHandle_t EPBackend::prepare_handle_locked(void* handle_mem, NVTEEpLayerCon ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; hcfg.dispatch_output_per_expert_alignment = layer_cfg.dispatch_output_per_expert_alignment; size_t hm_size = 0; - NVTE_CHECK_NCCL(loader::fns().HandleMemSize(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, - &hm_size, layer_cfg.top_k)); + NVTE_CHECK_NCCL(ncclEpHandleMemSize(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, &hm_size, + layer_cfg.top_k)); ncclEpHandle_t h = open_handle(handle_mem, hm_size, layer_cfg.top_k, layer_cfg.dispatch_output_per_expert_alignment); lru_.push_front(HandleEntry{handle_mem, h, layer_cfg, hm_size}); index_.emplace(handle_mem, lru_.begin()); while (lru_.size() > cache_cap_locked()) { HandleEntry& victim = lru_.back(); - if (victim.handle != nullptr) loader::fns().HandleDestroy(victim.handle); + if (victim.handle != nullptr) ncclEpHandleDestroy(victim.handle); index_.erase(victim.handle_mem); lru_.pop_back(); } @@ -316,8 +316,8 @@ size_t EPBackend::handle_mem_size(NVTEEpLayerConfig layer_cfg) { ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; hcfg.dispatch_output_per_expert_alignment = layer_cfg.dispatch_output_per_expert_alignment; size_t hm_size = 0; - NVTE_CHECK_NCCL(loader::fns().HandleMemSize(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, - &hm_size, layer_cfg.top_k)); + NVTE_CHECK_NCCL(ncclEpHandleMemSize(ep_group_, NCCL_EP_LAYOUT_EXPERT_MAJOR, &hcfg, &hm_size, + layer_cfg.top_k)); return hm_size; } @@ -341,7 +341,7 @@ void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, NVTETensor std::lock_guard lock(mutex_); ncclEpHandle_t h = prepare_handle_locked(handle_mem, layer_cfg); - NVTE_CHECK_NCCL(loader::fns().UpdateHandle(h, &nccl_topk_idx, &layout_info, stream)); + NVTE_CHECK_NCCL(ncclEpUpdateHandle(h, &nccl_topk_idx, &layout_info, stream)); } void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTETensor tokens, @@ -399,8 +399,8 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE std::lock_guard lock(mutex_); ncclEpHandle_t h = lookup_handle_locked(handle_mem); - NVTE_CHECK_NCCL(loader::fns().Dispatch(h, &in_struct, &out_struct, - /*layout_info=*/nullptr, &dispatch_cfg, stream)); + NVTE_CHECK_NCCL(ncclEpDispatch(h, &in_struct, &out_struct, + /*layout_info=*/nullptr, &dispatch_cfg, stream)); } void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, @@ -421,7 +421,7 @@ void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, std::lock_guard lock(mutex_); ncclEpHandle_t h = lookup_handle_locked(handle_mem); - NVTE_CHECK_NCCL(loader::fns().Combine(h, &in_struct, &out_struct, /*config=*/nullptr, stream)); + NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, /*config=*/nullptr, stream)); } void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, @@ -457,7 +457,7 @@ void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, std::lock_guard lock(mutex_); ncclEpHandle_t h = lookup_handle_locked(handle_mem); - NVTE_CHECK_NCCL(loader::fns().Combine(h, &in_struct, &out_struct, &cfg, stream)); + NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, &cfg, stream)); } void EPBackend::combine_bwd(void* handle_mem, const NVTETensor grad, const NVTECommWindow& grad_win, diff --git a/transformer_engine/common/ep/ep_nccl_loader.cpp b/transformer_engine/common/ep/ep_nccl_loader.cpp deleted file mode 100644 index 8374acd7b3..0000000000 --- a/transformer_engine/common/ep/ep_nccl_loader.cpp +++ /dev/null @@ -1,74 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include "ep_nccl_loader.h" - -#include - -#include "../util/logging.h" - -namespace transformer_engine { -namespace ep { -namespace loader { - -namespace { - -constexpr const char* kSonames[] = {"libnccl_ep.so.0", "libnccl_ep.so"}; - -void* try_dlopen(std::string& last_err) { - for (const char* name : kSonames) { - dlerror(); - void* h = dlopen(name, RTLD_LAZY | RTLD_LOCAL); - if (h != nullptr) return h; - if (const char* e = dlerror()) last_err = e; - } - return nullptr; -} - -template -Fn resolve(void* lib, const char* sym) { - dlerror(); - void* p = dlsym(lib, sym); - const char* err = dlerror(); - NVTE_CHECK(err == nullptr && p != nullptr, "libnccl_ep.so is loaded but symbol '", sym, - "' could not be resolved", (err != nullptr ? std::string(": ") + err : std::string{}), - ". The runtime libnccl_ep.so is older than the version TransformerEngine " - "was built against; upgrade NCCL EP or rebuild TE with -DNVTE_WITH_NCCL_EP=OFF."); - return reinterpret_cast(p); -} - -NcclEpFns load_or_throw() { - std::string last_err; - void* lib = try_dlopen(last_err); - NVTE_CHECK(lib != nullptr, "Failed to load libnccl_ep.so (", - (last_err.empty() ? "no error message" : last_err), - "). NCCL EP requires libnccl_ep.so (>= 0.0.1) and NCCL >= 2.30.4 at runtime. " - "Install the NCCL EP shared library, or rebuild TransformerEngine with " - "-DNVTE_WITH_NCCL_EP=OFF to disable EP support."); - NcclEpFns fns{}; - fns.InitHandle = resolve(lib, "ncclEpInitHandle"); - fns.CreateGroup = resolve(lib, "ncclEpCreateGroup"); - fns.GroupDestroy = resolve(lib, "ncclEpGroupDestroy"); - fns.HandleDestroy = resolve(lib, "ncclEpHandleDestroy"); - fns.HandleMemSize = resolve(lib, "ncclEpHandleMemSize"); - fns.UpdateHandle = resolve(lib, "ncclEpUpdateHandle"); - fns.Dispatch = resolve(lib, "ncclEpDispatch"); - fns.Combine = resolve(lib, "ncclEpCombine"); - return fns; -} - -} // namespace - -const NcclEpFns& fns() { - // Function-local static: thread-safe one-shot init; re-throws on every call - // if initialization fails, so a missing library is surfaced consistently. - static const NcclEpFns table = load_or_throw(); - return table; -} - -} // namespace loader -} // namespace ep -} // namespace transformer_engine diff --git a/transformer_engine/common/ep/ep_nccl_loader.h b/transformer_engine/common/ep/ep_nccl_loader.h deleted file mode 100644 index 8ffb437ed8..0000000000 --- a/transformer_engine/common/ep/ep_nccl_loader.h +++ /dev/null @@ -1,48 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -/*! \file ep_nccl_loader.h - * \brief Lazy dlopen-based resolver for libnccl_ep.so. - * - * libtransformer_engine.so is not link-time bound to libnccl_ep.so. The first - * call to ep::loader::fns() opens it via dlopen and dlsyms the ncclEp* - * entry points the EP backend uses. If the library or any symbol cannot be - * resolved (e.g. libnccl_ep.so is missing, or system NCCL is older than the - * EP minimum so libnccl_ep.so's own DT_NEEDED chain fails), the call throws - * NVTE_ERROR with remediation instead of preventing libtransformer_engine.so - * from loading. - */ - -#ifndef TRANSFORMER_ENGINE_COMMON_EP_EP_NCCL_LOADER_H_ -#define TRANSFORMER_ENGINE_COMMON_EP_EP_NCCL_LOADER_H_ - -#include - -namespace transformer_engine { -namespace ep { -namespace loader { - -struct NcclEpFns { - decltype(&::ncclEpInitHandle) InitHandle; - decltype(&::ncclEpCreateGroup) CreateGroup; - decltype(&::ncclEpGroupDestroy) GroupDestroy; - decltype(&::ncclEpHandleDestroy) HandleDestroy; - decltype(&::ncclEpHandleMemSize) HandleMemSize; - decltype(&::ncclEpUpdateHandle) UpdateHandle; - decltype(&::ncclEpDispatch) Dispatch; - decltype(&::ncclEpCombine) Combine; -}; - -/*! \brief Resolve libnccl_ep.so on first call; cache the table thereafter. - * Thread-safe; throws NVTE_ERROR if the library or any symbol is missing. - */ -const NcclEpFns& fns(); - -} // namespace loader -} // namespace ep -} // namespace transformer_engine - -#endif // TRANSFORMER_ENGINE_COMMON_EP_EP_NCCL_LOADER_H_ From 811714a59d191315fceea9163d91ed3bcfc1e2ec Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 10 Jun 2026 11:43:18 -0700 Subject: [PATCH 38/55] common/ep: lock handle_mem_size, guard tensor ndim, forward-declare ncclWindow in public header Signed-off-by: Phuong Nguyen --- transformer_engine/common/ep/ep_backend.cpp | 14 +++++++++++++- .../include/transformer_engine/comm_window.h | 10 +++++++--- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index b52a1f6a37..3094350db1 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -311,8 +311,9 @@ ncclEpHandle_t EPBackend::lookup_handle_locked(void* handle_mem) { // --------------------------------------------------------------------------- size_t EPBackend::handle_mem_size(NVTEEpLayerConfig layer_cfg) { - NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(layer_cfg.top_k > 0, "top_k must be > 0, got ", layer_cfg.top_k); + std::lock_guard lock(mutex_); + NVTE_CHECK(initialized_, "EPBackend not initialized"); ncclEpHandleConfig_t hcfg = NCCL_EP_HANDLE_CONFIG_INIT; hcfg.dispatch_output_per_expert_alignment = layer_cfg.dispatch_output_per_expert_alignment; size_t hm_size = 0; @@ -326,6 +327,7 @@ void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, NVTETensor NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); NVTE_CHECK(layer_cfg.top_k > 0, "top_k must be > 0, got ", layer_cfg.top_k); + NVTE_CHECK(nvte_tensor_shape(topk_idx).ndim == 2, "topk_idx must be 2D [T, top_k]"); NVTEShape topk_idx_shape; ncclEpTensor_t nccl_topk_idx = make_nccl_ep_tensor(topk_idx, topk_idx_shape); @@ -351,6 +353,9 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE const NVTECommWindow& recv_topk_weights_win, cudaStream_t stream) { NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); + NVTE_CHECK(nvte_tensor_shape(tokens).ndim == 2, "tokens must be 2D [T, hidden_dim]"); + NVTE_CHECK(nvte_tensor_shape(recv_tokens).ndim == 2, + "recv_tokens must be 2D [recv_T, hidden_dim]"); NVTEDType tok_dtype = nvte_tensor_type(tokens); NVTE_CHECK(typeToSize(static_cast(tok_dtype)) <= @@ -377,6 +382,8 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE ncclEpTensor_t nccl_topk_weights_out; if (is_forward) { NVTE_CHECK(topk_idx != nullptr, "topk_idx required in forward dispatch"); + NVTE_CHECK(nvte_tensor_shape(topk_idx).ndim == 2, "topk_idx must be 2D [T, top_k]"); + NVTE_CHECK(nvte_tensor_shape(topk_weights).ndim == 2, "topk_weights must be 2D [T, top_k]"); NVTE_CHECK(recv_topk_weights != nullptr, "recv_topk_weights must not be null in forward dispatch"); NVTE_CHECK(nvte_tensor_shape(recv_topk_weights).ndim == 1, @@ -408,6 +415,9 @@ void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, cudaStream_t stream) { NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); + NVTE_CHECK(nvte_tensor_shape(expert_out).ndim == 2, + "expert_out must be 2D [recv_T, hidden_dim]"); + NVTE_CHECK(nvte_tensor_shape(result).ndim == 2, "result must be 2D [T, hidden_dim]"); NVTEShape expert_out_shape, result_shape; ncclEpTensor_t nccl_expert_in = make_nccl_ep_tensor(expert_out, expert_out_shape, expert_out_win); @@ -430,6 +440,8 @@ void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, NVTETensor grad_topk_weights, cudaStream_t stream) { NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); + NVTE_CHECK(nvte_tensor_shape(grad).ndim == 2, "grad must be 2D [recv_capacity, hidden_dim]"); + NVTE_CHECK(nvte_tensor_shape(grad_tokens).ndim == 2, "grad_tokens must be 2D [T, hidden_dim]"); // g_recv_topk_weights must be 1D [recv_capacity]; caller flattens. NVTE_CHECK(nvte_tensor_shape(g_recv_topk_weights).ndim == 1, diff --git a/transformer_engine/common/include/transformer_engine/comm_window.h b/transformer_engine/common/include/transformer_engine/comm_window.h index 088ea7f0c3..ef750344e7 100644 --- a/transformer_engine/common/include/transformer_engine/comm_window.h +++ b/transformer_engine/common/include/transformer_engine/comm_window.h @@ -12,17 +12,21 @@ #ifndef TRANSFORMER_ENGINE_COMM_WINDOW_H_ #define TRANSFORMER_ENGINE_COMM_WINDOW_H_ -#include #include #ifdef __cplusplus extern "C" { #endif +/* Forward declaration of NCCL's opaque window struct so this header does not + * pull in ; matches NCCL's own typedef + * (struct ncclWindow_vidmem* ncclWindow_t). */ +struct ncclWindow_vidmem; + /*! \brief NCCL window + byte offset for a zero-copy payload tensor. */ typedef struct { - ncclWindow_t window; /*!< NCCL window, or NULL to use the raw data pointer. */ - uint64_t offset; /*!< Byte offset of the payload within ``window``. */ + struct ncclWindow_vidmem* window; /*!< NCCL window, or NULL to use the raw data pointer. */ + uint64_t offset; /*!< Byte offset of the payload within ``window``. */ } NVTECommWindow; #ifdef __cplusplus From a3ae3dfa342672004fe7c1d3c059e3ac2db11662 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 10 Jun 2026 11:48:12 -0700 Subject: [PATCH 39/55] qa,tests/cpp_distributed: emit per-rank JUnit XML and isolate sub-test failures in L1 CI Signed-off-by: Phuong Nguyen --- qa/L1_cpp_distributed/test.sh | 34 ++++++++++++++++++++++------ tests/cpp_distributed/run_test_ep.sh | 11 ++++++++- 2 files changed, 37 insertions(+), 8 deletions(-) diff --git a/qa/L1_cpp_distributed/test.sh b/qa/L1_cpp_distributed/test.sh index 7e5ce2cf0d..1b0501d1be 100755 --- a/qa/L1_cpp_distributed/test.sh +++ b/qa/L1_cpp_distributed/test.sh @@ -2,19 +2,39 @@ # # See LICENSE for license information. -set -e +function test_fail() { + RET=1 + FAILED_CASES="$FAILED_CASES $1" + echo "Error: sub-test failed: $1" +} + +RET=0 +FAILED_CASES="" # Find TE : ${TE_PATH:=/opt/transformerengine} +: ${XML_LOG_DIR:=/logs} +mkdir -p "$XML_LOG_DIR" + TE_LIB_PATH=$(pip3 show transformer-engine | grep -E "Location:|Editable project location:" | tail -n 1 | awk '{print $NF}') export LD_LIBRARY_PATH=$TE_LIB_PATH:$LD_LIBRARY_PATH if [[ $(nvidia-smi --list-gpus | wc -l) -ge 4 ]]; then - cd $TE_PATH/tests/cpp_distributed - cmake -GNinja -S. -Bbuild - cmake --build build - mpirun --allow-run-as-root --np 4 --oversubscribe ./build/test_comm_gemm + cd $TE_PATH/tests/cpp_distributed + cmake -GNinja -S. -Bbuild || test_fail "configure" + cmake --build build || test_fail "build" + + # test_comm_gemm: per-rank XML to avoid a write race on a shared path. + mpirun --allow-run-as-root --np 4 --oversubscribe bash -c \ + "exec ./build/test_comm_gemm --gtest_output=xml:$XML_LOG_DIR/cpp_distributed_test_comm_gemm.rank\${OMPI_COMM_WORLD_RANK}.xml" \ + || test_fail "test_comm_gemm" + + # EP suites; runner self-skips on pre-Hopper GPUs. + GTEST_XML_PREFIX="$XML_LOG_DIR/cpp_distributed_test_ep" \ + bash ./run_test_ep.sh 4 ./build || test_fail "test_ep" +fi - # EP suites; runner self-skips on pre-Hopper GPUs. - bash ./run_test_ep.sh 4 ./build +if [ "$RET" -ne 0 ]; then + echo "FAILED sub-tests:$FAILED_CASES" fi +exit $RET diff --git a/tests/cpp_distributed/run_test_ep.sh b/tests/cpp_distributed/run_test_ep.sh index 1c4432531c..d737df5ad8 100755 --- a/tests/cpp_distributed/run_test_ep.sh +++ b/tests/cpp_distributed/run_test_ep.sh @@ -15,6 +15,8 @@ # # Environment variables: # GTEST_FILTER - forwarded to all processes (e.g., "EPPipelineTest.*") +# GTEST_XML_PREFIX - if set, each rank writes JUnit XML to +# ${GTEST_XML_PREFIX}.rank.xml # MPIRUN - override the mpirun binary (default: mpirun) # MPIRUN_EXTRA - extra flags forwarded to mpirun @@ -51,4 +53,11 @@ echo "=== EP Tests ===" echo " GPUs: ${NUM_GPUS} Binary: ${TEST_BIN}" echo -"${MPIRUN}" -n "${NUM_GPUS}" ${MPIRUN_EXTRA:-} "${TEST_BIN}" ${GTEST_ARGS} +if [[ -n "${GTEST_XML_PREFIX:-}" ]]; then + # bash -c so OMPI_COMM_WORLD_RANK expands per-rank, avoiding a write race + # on a single shared output path. + "${MPIRUN}" -n "${NUM_GPUS}" ${MPIRUN_EXTRA:-} bash -c \ + "exec '${TEST_BIN}' ${GTEST_ARGS} --gtest_output=xml:${GTEST_XML_PREFIX}.rank\${OMPI_COMM_WORLD_RANK}.xml" +else + "${MPIRUN}" -n "${NUM_GPUS}" ${MPIRUN_EXTRA:-} "${TEST_BIN}" ${GTEST_ARGS} +fi From 12100dd3fc6626929e700484ddba230971df927b Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 10 Jun 2026 18:50:40 +0000 Subject: [PATCH 40/55] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/__init__.py | 1 + transformer_engine/common/ep/ep_backend.cpp | 7 +++---- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/transformer_engine/__init__.py b/transformer_engine/__init__.py index d84fd81360..b1bf648aba 100644 --- a/transformer_engine/__init__.py +++ b/transformer_engine/__init__.py @@ -39,6 +39,7 @@ def require_nccl_ep() -> None: "rebuild Transformer Engine with NVTE_BUILD_WITH_NCCL_EP=0." ) + try: from . import pytorch except ImportError: diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index 3094350db1..c2f02d4734 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -10,6 +10,8 @@ #include "ep_backend.h" +#include + #include #include #include @@ -17,8 +19,6 @@ #include #include -#include - #include "../common.h" #include "../util/cuda_runtime.h" #include "../util/logging.h" @@ -415,8 +415,7 @@ void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, cudaStream_t stream) { NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); - NVTE_CHECK(nvte_tensor_shape(expert_out).ndim == 2, - "expert_out must be 2D [recv_T, hidden_dim]"); + NVTE_CHECK(nvte_tensor_shape(expert_out).ndim == 2, "expert_out must be 2D [recv_T, hidden_dim]"); NVTE_CHECK(nvte_tensor_shape(result).ndim == 2, "result must be 2D [T, hidden_dim]"); NVTEShape expert_out_shape, result_shape; From ad07937e9923dfb04eb28cfe374c1327e83c1745 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 10 Jun 2026 12:37:55 -0700 Subject: [PATCH 41/55] common/ep: harden NCCL EP init/probe paths and CI build short-circuit Signed-off-by: Phuong Nguyen --- build_tools/utils.py | 41 ++++++++++++++++++++ qa/L1_cpp_distributed/test.sh | 29 ++++++++------ setup.py | 42 +-------------------- transformer_engine/__init__.py | 40 ++++++++++++++------ transformer_engine/common/ep/ep_backend.cpp | 22 +++++------ transformer_engine/common/ep/ep_backend.h | 3 +- 6 files changed, 102 insertions(+), 75 deletions(-) diff --git a/build_tools/utils.py b/build_tools/utils.py index f2548b4de6..9b94f25040 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -224,6 +224,47 @@ def nvcc_path() -> Tuple[str, str]: return nvcc_bin +@functools.lru_cache(maxsize=None) +def nccl_home_path() -> str: + """Returns the NCCL install prefix containing include/nccl.h and libnccl.so. + + Throws RuntimeError if NCCL is not found. + """ + # Honor env var when its include is valid. + env_home = os.environ.get("NCCL_HOME") + if env_home: + if (Path(env_home) / "include" / "nccl.h").is_file(): + return env_home + print( + f"[NCCL EP] WARNING: NCCL_HOME='{env_home}' is set but " + f"'{env_home}/include/nccl.h' was not found; falling back to system probes." + ) + + lib_names = ("libnccl.so", "libnccl.so.2") + lib_subdirs = ("lib", "lib64", "lib/aarch64-linux-gnu", "lib/x86_64-linux-gnu") + for cand in ("/opt/nvidia/nccl", "/usr/local/nccl", "/usr"): + p = Path(cand) + if (p / "include" / "nccl.h").is_file() and any( + (p / sub / name).exists() for sub in lib_subdirs for name in lib_names + ): + return str(p) + + try: + out = subprocess.check_output(["ldconfig", "-p"], stderr=subprocess.DEVNULL).decode() + for line in out.splitlines(): + if "libnccl.so" in line and "=>" in line: + lib_path = Path(line.split("=>")[-1].strip()) + for root in (lib_path.parent.parent, lib_path.parent.parent.parent): + if (root / "include" / "nccl.h").is_file(): + return str(root) + except (subprocess.CalledProcessError, FileNotFoundError): + pass + + raise RuntimeError( + "Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix." + ) + + @functools.lru_cache(maxsize=None) def get_cuda_include_dirs() -> Tuple[str, str]: """Returns the CUDA header directory.""" diff --git a/qa/L1_cpp_distributed/test.sh b/qa/L1_cpp_distributed/test.sh index 1b0501d1be..8c34b6941a 100755 --- a/qa/L1_cpp_distributed/test.sh +++ b/qa/L1_cpp_distributed/test.sh @@ -21,17 +21,24 @@ export LD_LIBRARY_PATH=$TE_LIB_PATH:$LD_LIBRARY_PATH if [[ $(nvidia-smi --list-gpus | wc -l) -ge 4 ]]; then cd $TE_PATH/tests/cpp_distributed - cmake -GNinja -S. -Bbuild || test_fail "configure" - cmake --build build || test_fail "build" - - # test_comm_gemm: per-rank XML to avoid a write race on a shared path. - mpirun --allow-run-as-root --np 4 --oversubscribe bash -c \ - "exec ./build/test_comm_gemm --gtest_output=xml:$XML_LOG_DIR/cpp_distributed_test_comm_gemm.rank\${OMPI_COMM_WORLD_RANK}.xml" \ - || test_fail "test_comm_gemm" - - # EP suites; runner self-skips on pre-Hopper GPUs. - GTEST_XML_PREFIX="$XML_LOG_DIR/cpp_distributed_test_ep" \ - bash ./run_test_ep.sh 4 ./build || test_fail "test_ep" + build_ok=1 + cmake -GNinja -S. -Bbuild || { test_fail "configure"; build_ok=0; } + if [[ $build_ok -eq 1 ]]; then + cmake --build build || { test_fail "build"; build_ok=0; } + fi + + # Run tests only when the build succeeded; otherwise mpirun on a missing or + # stale binary can hang across all ranks and mask the real failure. + if [[ $build_ok -eq 1 ]]; then + # test_comm_gemm: per-rank XML to avoid a write race on a shared path. + mpirun --allow-run-as-root --np 4 --oversubscribe bash -c \ + "exec ./build/test_comm_gemm --gtest_output=xml:$XML_LOG_DIR/cpp_distributed_test_comm_gemm.rank\${OMPI_COMM_WORLD_RANK}.xml" \ + || test_fail "test_comm_gemm" + + # EP suites; runner self-skips on pre-Hopper GPUs. + GTEST_XML_PREFIX="$XML_LOG_DIR/cpp_distributed_test_ep" \ + bash ./run_test_ep.sh 4 ./build || test_fail "test_ep" + fi fi if [ "$RET" -ne 0 ]; then diff --git a/setup.py b/setup.py index c9370073b2..c212fda1e7 100644 --- a/setup.py +++ b/setup.py @@ -22,6 +22,7 @@ cuda_version, cusolvermp_pypi_package_name, get_frameworks, + nccl_home_path, remove_dups, min_python_version_str, ) @@ -148,45 +149,6 @@ def setup_requirements() -> Tuple[List[str], List[str]]: return [remove_dups(reqs) for reqs in [install_reqs, test_reqs]] -def _discover_nccl_home() -> str: - """Resolve NCCL_HOME: honor env var, else probe well-known prefixes, else ldconfig.""" - env_home = os.environ.get("NCCL_HOME") - if env_home: - if (Path(env_home) / "include" / "nccl.h").exists(): - return env_home - print( - f"[NCCL EP] WARNING: NCCL_HOME='{env_home}' is set but " - f"'{env_home}/include/nccl.h' was not found; falling back to system probes." - ) - - lib_names = ("libnccl.so", "libnccl.so.2") - # Include Debian/Ubuntu multiarch subdirs (e.g. lib/aarch64-linux-gnu). - lib_subdirs = ("lib", "lib64", "lib/aarch64-linux-gnu", "lib/x86_64-linux-gnu") - for cand in ("/opt/nvidia/nccl", "/usr/local/nccl", "/usr"): - p = Path(cand) - if (p / "include" / "nccl.h").exists() and any( - (p / sub / name).exists() for sub in lib_subdirs for name in lib_names - ): - return str(p) - - try: - out = subprocess.check_output(["ldconfig", "-p"], stderr=subprocess.DEVNULL).decode() - for line in out.splitlines(): - if "libnccl.so" in line and "=>" in line: - lib_path = Path(line.split("=>")[-1].strip()) - # Walk upward so multiarch layouts (.../lib//libnccl.so) - # resolve to the prefix that contains include/nccl.h. - for root in (lib_path.parent.parent, lib_path.parent.parent.parent): - if (root / "include" / "nccl.h").exists(): - return str(root) - except (subprocess.CalledProcessError, FileNotFoundError): - pass - - raise RuntimeError( - "Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix." - ) - - def build_nccl_ep_submodule() -> str: """Build libnccl_ep.a from the 3rdparty/nccl submodule and return NCCL_HOME.""" nccl_root = current_file_path / "3rdparty" / "nccl" @@ -216,7 +178,7 @@ def build_nccl_ep_submodule() -> str: env["NVCC_GENCODE"] = gencode # NCCL EP needs the core NCCL headers + libnccl.so; write NCCL EP build # outputs to the submodule's local build/ tree. - nccl_home = _discover_nccl_home() + nccl_home = nccl_home_path() env["NCCL_HOME"] = nccl_home env["NCCL_EP_BUILDDIR"] = str(build_dir) diff --git a/transformer_engine/__init__.py b/transformer_engine/__init__.py index b1bf648aba..dcf7f5e73c 100644 --- a/transformer_engine/__init__.py +++ b/transformer_engine/__init__.py @@ -10,6 +10,7 @@ import functools import os from importlib import metadata +from typing import Optional, Tuple import transformer_engine.common # Minimum NCCL version for the statically-linked NCCL EP backend. @@ -17,26 +18,41 @@ @functools.lru_cache(maxsize=1) -def is_nccl_ep_available() -> bool: - """Return True if the runtime libnccl.so is new enough for NCCL EP.""" +def _nccl_runtime_version() -> Optional[Tuple[int, int, int]]: + """Return runtime (major, minor, patch) from libnccl.so.2, or None if unavailable.""" try: - libnccl = ctypes.CDLL("libnccl.so.2", mode=ctypes.RTLD_GLOBAL) - ver = ctypes.c_int(0) - libnccl.ncclGetVersion(ctypes.byref(ver)) + libnccl = ctypes.CDLL("libnccl.so.2", mode=ctypes.RTLD_LOCAL) + ncclGetVersion = libnccl.ncclGetVersion except (OSError, AttributeError): - return False + return None + ver = ctypes.c_int(0) + if ncclGetVersion(ctypes.byref(ver)) != 0: + return None v = ver.value - cur = (v // 10000, (v // 100) % 100, v % 100) - return cur >= _NCCL_EP_MIN_VERSION + return (v // 10000, (v // 100) % 100, v % 100) + + +def is_nccl_ep_available() -> bool: + """Return True if the runtime libnccl.so meets the NCCL EP minimum.""" + cur = _nccl_runtime_version() + return cur is not None and cur >= _NCCL_EP_MIN_VERSION def require_nccl_ep() -> None: """Raise RuntimeError if NCCL EP cannot run on the current libnccl.""" - if not is_nccl_ep_available(): - mn = ".".join(str(x) for x in _NCCL_EP_MIN_VERSION) + mn = ".".join(str(x) for x in _NCCL_EP_MIN_VERSION) + cur = _nccl_runtime_version() + if cur is None: + raise RuntimeError( + f"NCCL EP requires libnccl.so.2 (>= {mn}); could not load libnccl.so.2 " + "or query its version. Install NCCL or rebuild Transformer Engine with " + "NVTE_BUILD_WITH_NCCL_EP=0." + ) + if cur < _NCCL_EP_MIN_VERSION: raise RuntimeError( - f"NCCL EP requires NCCL >= {mn} at runtime; upgrade libnccl.so or " - "rebuild Transformer Engine with NVTE_BUILD_WITH_NCCL_EP=0." + f"NCCL EP requires NCCL >= {mn} at runtime; found " + f"{'.'.join(str(x) for x in cur)}. Upgrade libnccl.so or rebuild " + "Transformer Engine with NVTE_BUILD_WITH_NCCL_EP=0." ) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index c2f02d4734..c759d221e1 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -13,6 +13,7 @@ #include #include +#include #include #include #include @@ -99,25 +100,24 @@ void EPBackend::validate_config(const NVTEEpGroupConfig& config) { NVTE_CHECK(config.max_token_dtype >= 0 && config.max_token_dtype < kNVTENumTypes, "max_token_dtype out of range, got ", static_cast(config.max_token_dtype)); const size_t elem_bytes = typeToSize(static_cast(config.max_token_dtype)); - NVTE_CHECK(config.hidden_dim * elem_bytes >= 16, + const size_t row_bytes = static_cast(config.hidden_dim) * elem_bytes; + NVTE_CHECK(row_bytes >= 16, "hidden_dim * sizeof(max_token_dtype) must be >= 16 (NCCL EP 16B row alignment); " "got hidden_dim=", config.hidden_dim, ", element_bytes=", elem_bytes); + // NCCL EP packs row size into ncclEpGroupConfig::max_token_bytes (unsigned int). + NVTE_CHECK(row_bytes <= static_cast(UINT_MAX), + "hidden_dim * sizeof(max_token_dtype) exceeds 4 GiB; got ", row_bytes, " bytes"); NVTE_CHECK(config.num_experts % config.ep_size == 0, "num_experts (", config.num_experts, ") must be divisible by ep_size (", config.ep_size, ")"); NVTE_CHECK(config.max_num_sms >= 0, "max_num_sms must be >= 0 (0 = auto), got ", config.max_num_sms); - int device, major; - NVTE_CHECK_CUDA(cudaGetDevice(&device)); - NVTE_CHECK_CUDA(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, device)); - NVTE_CHECK(major >= 9, - "NCCL EP requires SM_90+ (Hopper or later), " - "but current device has compute capability ", - major, ".x"); - - NVTE_CHECK(cuda::supports_multicast(device), "NCCL EP requires CUDA multicast support on device ", - device); + const int sm = cuda::sm_arch(); + NVTE_CHECK(sm >= 90, + "NCCL EP requires SM_90+ (Hopper or later), but current device is SM_", sm); + NVTE_CHECK(cuda::supports_multicast(), + "NCCL EP requires CUDA multicast support on the current device"); } void EPBackend::initialize(ncclComm_t ep_comm, NVTEEpGroupConfig config) { diff --git a/transformer_engine/common/ep/ep_backend.h b/transformer_engine/common/ep/ep_backend.h index ea9aa019fa..828e724822 100644 --- a/transformer_engine/common/ep/ep_backend.h +++ b/transformer_engine/common/ep/ep_backend.h @@ -21,6 +21,7 @@ #include #include +#include #include #include #include @@ -101,7 +102,7 @@ class EPBackend { ncclEpGroup_t ep_group_{nullptr}; ncclComm_t ep_comm_{nullptr}; NVTEEpGroupConfig group_config_{}; - bool initialized_{false}; + std::atomic initialized_{false}; std::mutex mutex_; std::list lru_; std::unordered_map::iterator> index_; From d835573be7645fcbdefbed49ec7042854991eb12 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 10 Jun 2026 13:21:34 -0700 Subject: [PATCH 42/55] common/ep: drop Sphinx backticks and consolidate EP docstring Signed-off-by: Phuong Nguyen --- transformer_engine/common/ep/ep_backend.h | 7 +------ .../include/transformer_engine/comm_window.h | 13 ++++++------- .../common/include/transformer_engine/ep.h | 18 +++++++++--------- 3 files changed, 16 insertions(+), 22 deletions(-) diff --git a/transformer_engine/common/ep/ep_backend.h b/transformer_engine/common/ep/ep_backend.h index 828e724822..2325baafca 100644 --- a/transformer_engine/common/ep/ep_backend.h +++ b/transformer_engine/common/ep/ep_backend.h @@ -5,12 +5,7 @@ ************************************************************************/ /*! \file ep_backend.h - * \brief Internal NCCL EP singleton; not part of the public API. - * - * ncclEpHandles are cached by handle_mem device pointer. nvte_ep_prepare - * seeds the entry with the layer_cfg; dispatch/combine/_bwd look up by - * pointer. Cache cap: NVTE_EP_HANDLE_CACHE_SIZE (default 4096; -1 disables - * LRU eviction). + * \brief Internal NCCL EP singleton; not part of the public API. See ep.h. */ #ifndef TRANSFORMER_ENGINE_COMMON_EP_EP_BACKEND_H_ diff --git a/transformer_engine/common/include/transformer_engine/comm_window.h b/transformer_engine/common/include/transformer_engine/comm_window.h index ef750344e7..424c350bbd 100644 --- a/transformer_engine/common/include/transformer_engine/comm_window.h +++ b/transformer_engine/common/include/transformer_engine/comm_window.h @@ -5,8 +5,8 @@ ************************************************************************/ /*! \file comm_window.h - * \brief Borrowed symmetric-memory window + offset for zero-copy one-sided ops. - * Pass ``{NULL, 0}`` to use the raw-pointer path. + * \brief NCCL symmetric-memory window handle for zero-copy ops. Pass + * {NULL, 0} to use the raw-pointer path. */ #ifndef TRANSFORMER_ENGINE_COMM_WINDOW_H_ @@ -18,15 +18,14 @@ extern "C" { #endif -/* Forward declaration of NCCL's opaque window struct so this header does not - * pull in ; matches NCCL's own typedef - * (struct ncclWindow_vidmem* ncclWindow_t). */ +/* Forward-declare NCCL's opaque window struct so this header does not pull in + * ; matches NCCL's typedef (struct ncclWindow_vidmem* ncclWindow_t). */ struct ncclWindow_vidmem; -/*! \brief NCCL window + byte offset for a zero-copy payload tensor. */ +/*! \brief NCCL window plus byte offset for a zero-copy payload tensor. */ typedef struct { struct ncclWindow_vidmem* window; /*!< NCCL window, or NULL to use the raw data pointer. */ - uint64_t offset; /*!< Byte offset of the payload within ``window``. */ + uint64_t offset; /*!< Byte offset of the payload within window. */ } NVTECommWindow; #ifdef __cplusplus diff --git a/transformer_engine/common/include/transformer_engine/ep.h b/transformer_engine/common/include/transformer_engine/ep.h index 5682e9fdb6..8928b92825 100644 --- a/transformer_engine/common/include/transformer_engine/ep.h +++ b/transformer_engine/common/include/transformer_engine/ep.h @@ -128,13 +128,13 @@ void nvte_ep_prepare(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor toke * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] topk_idx [T, top_k] int64 sparse routing indices. * \param[in] tokens [T, hidden_dim] input tokens. - * \param[in] tokens_win Optional symmem window for ``tokens``. + * \param[in] tokens_win Optional symmem window for tokens. * \param[in] topk_weights [T, top_k] float32 weights, or null in backward. - * \param[in] topk_weights_win Optional symmem window for ``topk_weights``. + * \param[in] topk_weights_win Optional symmem window for topk_weights. * \param[out] recv_tokens [recv_T, hidden_dim] received tokens. - * \param[in] recv_tokens_win Optional symmem window for ``recv_tokens``. + * \param[in] recv_tokens_win Optional symmem window for recv_tokens. * \param[out] recv_topk_weights [recv_T] float32 per-slot weights, or null in backward. - * \param[in] recv_topk_weights_win Optional symmem window for ``recv_topk_weights``. + * \param[in] recv_topk_weights_win Optional symmem window for recv_topk_weights. * \param[in] stream CUDA stream. */ void nvte_ep_dispatch(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor tokens, @@ -152,7 +152,7 @@ void nvte_ep_dispatch(NVTETensor handle_mem, NVTETensor topk_idx, NVTETensor tok * * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] expert_out [recv_T, hidden_dim] pre-weighted expert outputs. - * \param[in] expert_out_win Optional symmem window for ``expert_out``. + * \param[in] expert_out_win Optional symmem window for expert_out. * \param[out] result [T, hidden_dim] combined output. * \param[in] stream CUDA stream. */ @@ -167,9 +167,9 @@ void nvte_ep_combine(NVTETensor handle_mem, NVTETensor expert_out, NVTECommWindo * * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] grad [recv_capacity, hidden_dim] grad w.r.t. recv_tokens. - * \param[in] grad_win Optional symmem window for ``grad``. + * \param[in] grad_win Optional symmem window for grad. * \param[in] g_recv_topk_weights [recv_capacity] f32 grad w.r.t. recv_topk_weights. - * \param[in] g_recv_topk_weights_win Optional symmem window for ``g_recv_topk_weights``. + * \param[in] g_recv_topk_weights_win Optional symmem window for g_recv_topk_weights. * \param[out] grad_tokens [T, hidden_dim] grad w.r.t. tokens. * \param[out] grad_topk_weights [T, top_k] f32 grad w.r.t. topk_weights. * \param[in] stream CUDA stream. @@ -187,9 +187,9 @@ void nvte_ep_dispatch_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow * * \param[in] handle_mem uint8 routing-state buffer (from prepare). * \param[in] grad [T, hidden_dim] grad w.r.t. result. - * \param[in] grad_win Optional symmem window for ``grad``. + * \param[in] grad_win Optional symmem window for grad. * \param[out] grad_expert_out [recv_capacity, hidden_dim] grad w.r.t. expert_out. - * \param[in] grad_expert_out_win Optional symmem window for ``grad_expert_out``. + * \param[in] grad_expert_out_win Optional symmem window for grad_expert_out. * \param[in] stream CUDA stream. */ void nvte_ep_combine_bwd(NVTETensor handle_mem, NVTETensor grad, NVTECommWindow grad_win, From 3a8758a2707d66c5d7a1da52c28fa1e7055f0c67 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 10 Jun 2026 13:26:03 -0700 Subject: [PATCH 43/55] setup: keep _discover_nccl_home local; revert utils.py move Signed-off-by: Phuong Nguyen --- build_tools/utils.py | 41 ----------------------------------------- setup.py | 42 ++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 40 insertions(+), 43 deletions(-) diff --git a/build_tools/utils.py b/build_tools/utils.py index 9b94f25040..f2548b4de6 100644 --- a/build_tools/utils.py +++ b/build_tools/utils.py @@ -224,47 +224,6 @@ def nvcc_path() -> Tuple[str, str]: return nvcc_bin -@functools.lru_cache(maxsize=None) -def nccl_home_path() -> str: - """Returns the NCCL install prefix containing include/nccl.h and libnccl.so. - - Throws RuntimeError if NCCL is not found. - """ - # Honor env var when its include is valid. - env_home = os.environ.get("NCCL_HOME") - if env_home: - if (Path(env_home) / "include" / "nccl.h").is_file(): - return env_home - print( - f"[NCCL EP] WARNING: NCCL_HOME='{env_home}' is set but " - f"'{env_home}/include/nccl.h' was not found; falling back to system probes." - ) - - lib_names = ("libnccl.so", "libnccl.so.2") - lib_subdirs = ("lib", "lib64", "lib/aarch64-linux-gnu", "lib/x86_64-linux-gnu") - for cand in ("/opt/nvidia/nccl", "/usr/local/nccl", "/usr"): - p = Path(cand) - if (p / "include" / "nccl.h").is_file() and any( - (p / sub / name).exists() for sub in lib_subdirs for name in lib_names - ): - return str(p) - - try: - out = subprocess.check_output(["ldconfig", "-p"], stderr=subprocess.DEVNULL).decode() - for line in out.splitlines(): - if "libnccl.so" in line and "=>" in line: - lib_path = Path(line.split("=>")[-1].strip()) - for root in (lib_path.parent.parent, lib_path.parent.parent.parent): - if (root / "include" / "nccl.h").is_file(): - return str(root) - except (subprocess.CalledProcessError, FileNotFoundError): - pass - - raise RuntimeError( - "Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix." - ) - - @functools.lru_cache(maxsize=None) def get_cuda_include_dirs() -> Tuple[str, str]: """Returns the CUDA header directory.""" diff --git a/setup.py b/setup.py index c212fda1e7..c9370073b2 100644 --- a/setup.py +++ b/setup.py @@ -22,7 +22,6 @@ cuda_version, cusolvermp_pypi_package_name, get_frameworks, - nccl_home_path, remove_dups, min_python_version_str, ) @@ -149,6 +148,45 @@ def setup_requirements() -> Tuple[List[str], List[str]]: return [remove_dups(reqs) for reqs in [install_reqs, test_reqs]] +def _discover_nccl_home() -> str: + """Resolve NCCL_HOME: honor env var, else probe well-known prefixes, else ldconfig.""" + env_home = os.environ.get("NCCL_HOME") + if env_home: + if (Path(env_home) / "include" / "nccl.h").exists(): + return env_home + print( + f"[NCCL EP] WARNING: NCCL_HOME='{env_home}' is set but " + f"'{env_home}/include/nccl.h' was not found; falling back to system probes." + ) + + lib_names = ("libnccl.so", "libnccl.so.2") + # Include Debian/Ubuntu multiarch subdirs (e.g. lib/aarch64-linux-gnu). + lib_subdirs = ("lib", "lib64", "lib/aarch64-linux-gnu", "lib/x86_64-linux-gnu") + for cand in ("/opt/nvidia/nccl", "/usr/local/nccl", "/usr"): + p = Path(cand) + if (p / "include" / "nccl.h").exists() and any( + (p / sub / name).exists() for sub in lib_subdirs for name in lib_names + ): + return str(p) + + try: + out = subprocess.check_output(["ldconfig", "-p"], stderr=subprocess.DEVNULL).decode() + for line in out.splitlines(): + if "libnccl.so" in line and "=>" in line: + lib_path = Path(line.split("=>")[-1].strip()) + # Walk upward so multiarch layouts (.../lib//libnccl.so) + # resolve to the prefix that contains include/nccl.h. + for root in (lib_path.parent.parent, lib_path.parent.parent.parent): + if (root / "include" / "nccl.h").exists(): + return str(root) + except (subprocess.CalledProcessError, FileNotFoundError): + pass + + raise RuntimeError( + "Could not locate NCCL core (nccl.h + libnccl.so). Set NCCL_HOME to the install prefix." + ) + + def build_nccl_ep_submodule() -> str: """Build libnccl_ep.a from the 3rdparty/nccl submodule and return NCCL_HOME.""" nccl_root = current_file_path / "3rdparty" / "nccl" @@ -178,7 +216,7 @@ def build_nccl_ep_submodule() -> str: env["NVCC_GENCODE"] = gencode # NCCL EP needs the core NCCL headers + libnccl.so; write NCCL EP build # outputs to the submodule's local build/ tree. - nccl_home = nccl_home_path() + nccl_home = _discover_nccl_home() env["NCCL_HOME"] = nccl_home env["NCCL_EP_BUILDDIR"] = str(build_dir) From a1ddadbbf98824730aaf080890b1ab317ada19f9 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 10 Jun 2026 14:04:00 -0700 Subject: [PATCH 44/55] qa/L1_cpp_distributed: build and run each suite independently Signed-off-by: Phuong Nguyen --- qa/L1_cpp_distributed/test.sh | 40 ++++++++++++++++++++--------------- 1 file changed, 23 insertions(+), 17 deletions(-) diff --git a/qa/L1_cpp_distributed/test.sh b/qa/L1_cpp_distributed/test.sh index 8c34b6941a..4122163f69 100755 --- a/qa/L1_cpp_distributed/test.sh +++ b/qa/L1_cpp_distributed/test.sh @@ -21,23 +21,29 @@ export LD_LIBRARY_PATH=$TE_LIB_PATH:$LD_LIBRARY_PATH if [[ $(nvidia-smi --list-gpus | wc -l) -ge 4 ]]; then cd $TE_PATH/tests/cpp_distributed - build_ok=1 - cmake -GNinja -S. -Bbuild || { test_fail "configure"; build_ok=0; } - if [[ $build_ok -eq 1 ]]; then - cmake --build build || { test_fail "build"; build_ok=0; } - fi - - # Run tests only when the build succeeded; otherwise mpirun on a missing or - # stale binary can hang across all ranks and mask the real failure. - if [[ $build_ok -eq 1 ]]; then - # test_comm_gemm: per-rank XML to avoid a write race on a shared path. - mpirun --allow-run-as-root --np 4 --oversubscribe bash -c \ - "exec ./build/test_comm_gemm --gtest_output=xml:$XML_LOG_DIR/cpp_distributed_test_comm_gemm.rank\${OMPI_COMM_WORLD_RANK}.xml" \ - || test_fail "test_comm_gemm" - - # EP suites; runner self-skips on pre-Hopper GPUs. - GTEST_XML_PREFIX="$XML_LOG_DIR/cpp_distributed_test_ep" \ - bash ./run_test_ep.sh 4 ./build || test_fail "test_ep" + configure_ok=1 + cmake -GNinja -S. -Bbuild || { test_fail "configure"; configure_ok=0; } + + # Build each suite independently so one suite's build failure does not + # mask the other's results. Skip mpirun when the binary is missing. + if [[ $configure_ok -eq 1 ]]; then + comm_gemm_ok=1 + ep_ok=1 + cmake --build build --target test_comm_gemm || { test_fail "test_comm_gemm_build"; comm_gemm_ok=0; } + cmake --build build --target test_ep || { test_fail "test_ep_build"; ep_ok=0; } + + if [[ $comm_gemm_ok -eq 1 ]]; then + # Per-rank XML to avoid a write race on a shared path. + mpirun --allow-run-as-root --np 4 --oversubscribe bash -c \ + "exec ./build/test_comm_gemm --gtest_output=xml:$XML_LOG_DIR/cpp_distributed_test_comm_gemm.rank\${OMPI_COMM_WORLD_RANK}.xml" \ + || test_fail "test_comm_gemm" + fi + + if [[ $ep_ok -eq 1 ]]; then + # EP suites; runner self-skips on pre-Hopper GPUs. + GTEST_XML_PREFIX="$XML_LOG_DIR/cpp_distributed_test_ep" \ + bash ./run_test_ep.sh 4 ./build || test_fail "test_ep" + fi fi fi From 343b4edfa6c12e2a889fafee3445c8201c965573 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 10 Jun 2026 16:40:55 -0700 Subject: [PATCH 45/55] qa/L1_cpp_distributed: pass --allow-run-as-root --oversubscribe to mpirun in run_test_ep.sh Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/run_test_ep.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/cpp_distributed/run_test_ep.sh b/tests/cpp_distributed/run_test_ep.sh index d737df5ad8..d486d45f8a 100755 --- a/tests/cpp_distributed/run_test_ep.sh +++ b/tests/cpp_distributed/run_test_ep.sh @@ -56,8 +56,8 @@ echo if [[ -n "${GTEST_XML_PREFIX:-}" ]]; then # bash -c so OMPI_COMM_WORLD_RANK expands per-rank, avoiding a write race # on a single shared output path. - "${MPIRUN}" -n "${NUM_GPUS}" ${MPIRUN_EXTRA:-} bash -c \ + "${MPIRUN}" --allow-run-as-root --oversubscribe -n "${NUM_GPUS}" ${MPIRUN_EXTRA:-} bash -c \ "exec '${TEST_BIN}' ${GTEST_ARGS} --gtest_output=xml:${GTEST_XML_PREFIX}.rank\${OMPI_COMM_WORLD_RANK}.xml" else - "${MPIRUN}" -n "${NUM_GPUS}" ${MPIRUN_EXTRA:-} "${TEST_BIN}" ${GTEST_ARGS} + "${MPIRUN}" --allow-run-as-root --oversubscribe -n "${NUM_GPUS}" ${MPIRUN_EXTRA:-} "${TEST_BIN}" ${GTEST_ARGS} fi From aab70705bd6e4b9b6774e2bc85fad7fffc9cb22f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 10 Jun 2026 23:44:41 +0000 Subject: [PATCH 46/55] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- transformer_engine/common/ep/ep_backend.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index c759d221e1..3445cb6672 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -114,8 +114,7 @@ void EPBackend::validate_config(const NVTEEpGroupConfig& config) { config.max_num_sms); const int sm = cuda::sm_arch(); - NVTE_CHECK(sm >= 90, - "NCCL EP requires SM_90+ (Hopper or later), but current device is SM_", sm); + NVTE_CHECK(sm >= 90, "NCCL EP requires SM_90+ (Hopper or later), but current device is SM_", sm); NVTE_CHECK(cuda::supports_multicast(), "NCCL EP requires CUDA multicast support on the current device"); } From c80bcbc6a4c3290b8b9184e95ab89468189a2a4e Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Wed, 10 Jun 2026 16:49:08 -0700 Subject: [PATCH 47/55] common/ep: check initialized_ under mutex_ in per-step EPBackend ops Signed-off-by: Phuong Nguyen --- transformer_engine/common/ep/ep_backend.cpp | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index 3445cb6672..b1fe8baf83 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -323,7 +323,6 @@ size_t EPBackend::handle_mem_size(NVTEEpLayerConfig layer_cfg) { void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, NVTETensor token_counts, NVTEEpLayerConfig layer_cfg, cudaStream_t stream) { - NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); NVTE_CHECK(layer_cfg.top_k > 0, "top_k must be > 0, got ", layer_cfg.top_k); NVTE_CHECK(nvte_tensor_shape(topk_idx).ndim == 2, "topk_idx must be 2D [T, top_k]"); @@ -341,6 +340,7 @@ void EPBackend::prepare(void* handle_mem, const NVTETensor topk_idx, NVTETensor layout_info.expert_counters = (token_counts != nullptr) ? &token_counts_desc : nullptr; std::lock_guard lock(mutex_); + NVTE_CHECK(initialized_, "EPBackend not initialized"); ncclEpHandle_t h = prepare_handle_locked(handle_mem, layer_cfg); NVTE_CHECK_NCCL(ncclEpUpdateHandle(h, &nccl_topk_idx, &layout_info, stream)); } @@ -350,7 +350,6 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE const NVTECommWindow& topk_weights_win, NVTETensor recv_tokens, const NVTECommWindow& recv_tokens_win, NVTETensor recv_topk_weights, const NVTECommWindow& recv_topk_weights_win, cudaStream_t stream) { - NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); NVTE_CHECK(nvte_tensor_shape(tokens).ndim == 2, "tokens must be 2D [T, hidden_dim]"); NVTE_CHECK(nvte_tensor_shape(recv_tokens).ndim == 2, @@ -404,6 +403,7 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE dispatch_cfg.pass_direction = is_forward ? NCCL_EP_FWD_PASS : NCCL_EP_BWD_PASS; std::lock_guard lock(mutex_); + NVTE_CHECK(initialized_, "EPBackend not initialized"); ncclEpHandle_t h = lookup_handle_locked(handle_mem); NVTE_CHECK_NCCL(ncclEpDispatch(h, &in_struct, &out_struct, /*layout_info=*/nullptr, &dispatch_cfg, stream)); @@ -412,7 +412,6 @@ void EPBackend::dispatch(void* handle_mem, const NVTETensor topk_idx, const NVTE void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, const NVTECommWindow& expert_out_win, NVTETensor result, cudaStream_t stream) { - NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); NVTE_CHECK(nvte_tensor_shape(expert_out).ndim == 2, "expert_out must be 2D [recv_T, hidden_dim]"); NVTE_CHECK(nvte_tensor_shape(result).ndim == 2, "result must be 2D [T, hidden_dim]"); @@ -428,6 +427,7 @@ void EPBackend::combine(void* handle_mem, const NVTETensor expert_out, out_struct.tokens = &nccl_result_out; std::lock_guard lock(mutex_); + NVTE_CHECK(initialized_, "EPBackend not initialized"); ncclEpHandle_t h = lookup_handle_locked(handle_mem); NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, /*config=*/nullptr, stream)); } @@ -436,7 +436,6 @@ void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, const NVTECommWindow& grad_win, const NVTETensor g_recv_topk_weights, const NVTECommWindow& g_recv_topk_weights_win, NVTETensor grad_tokens, NVTETensor grad_topk_weights, cudaStream_t stream) { - NVTE_CHECK(initialized_, "EPBackend not initialized"); NVTE_CHECK(handle_mem != nullptr, "handle_mem must not be null"); NVTE_CHECK(nvte_tensor_shape(grad).ndim == 2, "grad must be 2D [recv_capacity, hidden_dim]"); NVTE_CHECK(nvte_tensor_shape(grad_tokens).ndim == 2, "grad_tokens must be 2D [T, hidden_dim]"); @@ -466,6 +465,7 @@ void EPBackend::dispatch_bwd(void* handle_mem, const NVTETensor grad, cfg.pass_direction = NCCL_EP_BWD_PASS; std::lock_guard lock(mutex_); + NVTE_CHECK(initialized_, "EPBackend not initialized"); ncclEpHandle_t h = lookup_handle_locked(handle_mem); NVTE_CHECK_NCCL(ncclEpCombine(h, &in_struct, &out_struct, &cfg, stream)); } From f4c8646e5d2ef4e1eea255f53f8d26f747122291 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 11 Jun 2026 08:01:23 -0700 Subject: [PATCH 48/55] setup/ep: rename NVTE_BUILD_WITH_NCCL_EP to NVTE_WITH_NCCL_EP, harden NCCL EP build Signed-off-by: Phuong Nguyen --- .gitignore | 3 - setup.py | 76 ++++++++++++++++++++---- transformer_engine/__init__.py | 4 +- transformer_engine/common/CMakeLists.txt | 2 +- transformer_engine/common/ep/ep_api.cpp | 2 +- 5 files changed, 67 insertions(+), 20 deletions(-) diff --git a/.gitignore b/.gitignore index 7b03d79cc7..8a627a7e76 100644 --- a/.gitignore +++ b/.gitignore @@ -43,6 +43,3 @@ tensor_dumps/ artifacts/ .DS_Store .claude/ - -# NCCL EP shared library staged by setup.py for wheel packaging. -/transformer_engine/libnccl_ep.so* diff --git a/setup.py b/setup.py index c9370073b2..9fb962bbc7 100644 --- a/setup.py +++ b/setup.py @@ -85,8 +85,10 @@ def setup_common_extension() -> CMakeExtension: cmake_flags.append(f"-DCUSOLVERMP_DIR={cusolvermp_dir}") # NCCL EP (Hopper+): on by default; auto-skipped when no arch >= 90 is - # targeted. Set NVTE_BUILD_WITH_NCCL_EP=0 to force off. - build_with_nccl_ep = bool(int(os.getenv("NVTE_BUILD_WITH_NCCL_EP", "1"))) + # targeted. Set NVTE_WITH_NCCL_EP=0 to force off. + nccl_ep_env = os.getenv("NVTE_WITH_NCCL_EP") + nccl_ep_explicit = nccl_ep_env is not None + build_with_nccl_ep = bool(int(nccl_ep_env if nccl_ep_explicit else "1")) if build_with_nccl_ep: arch_tokens = [a.strip() for a in str(archs or "").split(";") if a.strip()] has_hopper_or_newer = any( @@ -94,6 +96,11 @@ def setup_common_extension() -> CMakeExtension: for t in arch_tokens ) if not has_hopper_or_newer: + if nccl_ep_explicit: + raise RuntimeError( + f"NVTE_WITH_NCCL_EP=1 was set but NVTE_CUDA_ARCHS ('{archs}') " + "contains no arch >= 90. NCCL EP requires Hopper or newer." + ) print(f"[NCCL EP] No arch >= 90 in NVTE_CUDA_ARCHS ('{archs}'); skipping build.") build_with_nccl_ep = False if build_with_nccl_ep: @@ -162,6 +169,22 @@ def _discover_nccl_home() -> str: lib_names = ("libnccl.so", "libnccl.so.2") # Include Debian/Ubuntu multiarch subdirs (e.g. lib/aarch64-linux-gnu). lib_subdirs = ("lib", "lib64", "lib/aarch64-linux-gnu", "lib/x86_64-linux-gnu") + + # pip-installed NCCL (nvidia-nccl-cu* wheel) lives under nvidia/nccl in + # site-packages and has no top-level include/lib layout. + try: + import importlib.util + + spec = importlib.util.find_spec("nvidia.nccl") + if spec is not None and spec.submodule_search_locations: + pip_root = Path(next(iter(spec.submodule_search_locations))) + if (pip_root / "include" / "nccl.h").exists() and any( + (pip_root / sub / name).exists() for sub in lib_subdirs for name in lib_names + ): + return str(pip_root) + except (ImportError, ValueError): + pass + for cand in ("/opt/nvidia/nccl", "/usr/local/nccl", "/usr"): p = Path(cand) if (p / "include" / "nccl.h").exists() and any( @@ -198,18 +221,32 @@ def build_nccl_ep_submodule() -> str: build_dir = nccl_root / "build" nccl_ep_lib = build_dir / "lib" / "libnccl_ep.a" + gencode_stamp = build_dir / "lib" / "libnccl_ep.gencode" - # Caller gates on arch >= 90 or "native"; let nvcc resolve "native". + # Caller gates on arch >= 90 or "native"; expand "native" to the host's + # actual sm_XX so the build stamp distinguishes machines. arch_tokens = [a.strip() for a in str(cuda_archs() or "").split(";") if a.strip()] - if any(t.lower() == "native" for t in arch_tokens): - gencode = "-arch=native" - else: - arch_list = [ - t.rstrip("af") - for t in arch_tokens - if t.rstrip("af").isdigit() and int(t.rstrip("af")) >= 90 - ] - gencode = " ".join(f"-gencode=arch=compute_{a},code=sm_{a}" for a in arch_list) + arch_list: list[str] = [] + for t in arch_tokens: + if t.lower() == "native": + try: + out = subprocess.check_output( + ["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], + stderr=subprocess.DEVNULL, + ).decode() + except (subprocess.CalledProcessError, FileNotFoundError) as e: + raise RuntimeError( + "NVTE_CUDA_ARCHS=native requires nvidia-smi to resolve the host arch." + ) from e + for line in out.splitlines(): + cap = line.strip().replace(".", "") + if cap.isdigit() and int(cap) >= 90 and cap not in arch_list: + arch_list.append(cap) + else: + bare = t.rstrip("af") + if bare.isdigit() and int(bare) >= 90 and bare not in arch_list: + arch_list.append(bare) + gencode = " ".join(f"-gencode=arch=compute_{a},code=sm_{a}" for a in arch_list) nproc = os.cpu_count() or 8 env = os.environ.copy() @@ -220,13 +257,26 @@ def build_nccl_ep_submodule() -> str: env["NCCL_HOME"] = nccl_home env["NCCL_EP_BUILDDIR"] = str(build_dir) - if not nccl_ep_lib.exists(): + prev_gencode = gencode_stamp.read_text().strip() if gencode_stamp.exists() else None + if not nccl_ep_lib.exists() or prev_gencode != gencode: + if nccl_ep_lib.exists() and prev_gencode != gencode: + print( + f"[NCCL EP] gencode changed ('{prev_gencode}' -> '{gencode}'); " + "rebuilding libnccl_ep.a" + ) + subprocess.check_call( + ["make", "-C", "contrib/nccl_ep", "clean"], + cwd=str(nccl_root), + env=env, + ) print(f"[NCCL EP] Building libnccl_ep.a (gencode='{gencode}')") subprocess.check_call( ["make", "-j", str(nproc), "-C", "contrib/nccl_ep", "lib"], cwd=str(nccl_root), env=env, ) + gencode_stamp.parent.mkdir(parents=True, exist_ok=True) + gencode_stamp.write_text(gencode) return nccl_home diff --git a/transformer_engine/__init__.py b/transformer_engine/__init__.py index dcf7f5e73c..91fd7f6f1a 100644 --- a/transformer_engine/__init__.py +++ b/transformer_engine/__init__.py @@ -46,13 +46,13 @@ def require_nccl_ep() -> None: raise RuntimeError( f"NCCL EP requires libnccl.so.2 (>= {mn}); could not load libnccl.so.2 " "or query its version. Install NCCL or rebuild Transformer Engine with " - "NVTE_BUILD_WITH_NCCL_EP=0." + "NVTE_WITH_NCCL_EP=0." ) if cur < _NCCL_EP_MIN_VERSION: raise RuntimeError( f"NCCL EP requires NCCL >= {mn} at runtime; found " f"{'.'.join(str(x) for x in cur)}. Upgrade libnccl.so or rebuild " - "Transformer Engine with NVTE_BUILD_WITH_NCCL_EP=0." + "Transformer Engine with NVTE_WITH_NCCL_EP=0." ) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index dba040efe6..28cb71dcf9 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -438,7 +438,7 @@ if (NVTE_WITH_CUSOLVERMP) endif() # -- NCCL EP (on by default, HT mode only) --------------------------------- -# Set -DNVTE_WITH_NCCL_EP=OFF (or NVTE_BUILD_WITH_NCCL_EP=0 in setup.py) to +# Set -DNVTE_WITH_NCCL_EP=OFF (or NVTE_WITH_NCCL_EP=0 in setup.py) to # skip NCCL EP entirely - useful on older images whose system NCCL is below # the 2.30.4 EP minimum. option(NVTE_WITH_NCCL_EP "Build NCCL EP into libtransformer_engine.so" ON) diff --git a/transformer_engine/common/ep/ep_api.cpp b/transformer_engine/common/ep/ep_api.cpp index b8cf04aa4a..66ee3dc8d9 100644 --- a/transformer_engine/common/ep/ep_api.cpp +++ b/transformer_engine/common/ep/ep_api.cpp @@ -84,7 +84,7 @@ namespace { [[noreturn]] void ep_not_built() { NVTE_ERROR( "NCCL EP is not built into this TransformerEngine. Rebuild TE with " - "NVTE_BUILD_WITH_NCCL_EP=1 and CUDA arch >= 90 (e.g. NVTE_CUDA_ARCHS=\"90\")."); + "NVTE_WITH_NCCL_EP=1 and CUDA arch >= 90 (e.g. NVTE_CUDA_ARCHS=\"90\")."); } } // namespace From 0716d318b5731fbaa2ab5c223d02832a9cea3f74 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 11 Jun 2026 08:09:10 -0700 Subject: [PATCH 49/55] tests/cpp_distributed: anchor NCCL discovery on libnccl, derive include dir from its prefix Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/CMakeLists.txt | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/tests/cpp_distributed/CMakeLists.txt b/tests/cpp_distributed/CMakeLists.txt index 0ea97eb68d..13b6242816 100644 --- a/tests/cpp_distributed/CMakeLists.txt +++ b/tests/cpp_distributed/CMakeLists.txt @@ -57,17 +57,30 @@ find_package(OpenMP REQUIRED) find_package(MPI REQUIRED) # -- NCCL core ---------------------------------------------------------------- -# Pass -DNCCL_INCLUDE_DIR=/include; falls back to well-known prefixes. -find_path(NCCL_INCLUDE_DIR nccl.h - HINTS /opt/nvidia/nccl/include /usr/local/nccl/include) -if(NOT NCCL_INCLUDE_DIR) - message(FATAL_ERROR - "nccl.h not found. Pass -DNCCL_INCLUDE_DIR=/include.") -endif() +# Anchor on libnccl and derive nccl.h from the same install prefix so the +# header and library can't drift across installs. find_library(NCCL_LIB NAMES nccl libnccl + HINTS /opt/nvidia/nccl/lib /opt/nvidia/nccl/lib64 + /usr/local/nccl/lib /usr/local/nccl/lib64 PATH_SUFFIXES lib lib64 REQUIRED) +get_filename_component(_nccl_lib_dir "${NCCL_LIB}" DIRECTORY) +set(NCCL_PREFIX "${_nccl_lib_dir}") +while(NCCL_PREFIX AND NOT EXISTS "${NCCL_PREFIX}/include/nccl.h") + get_filename_component(_nccl_parent "${NCCL_PREFIX}" DIRECTORY) + if(_nccl_parent STREQUAL NCCL_PREFIX) + break() + endif() + set(NCCL_PREFIX "${_nccl_parent}") +endwhile() +find_path(NCCL_INCLUDE_DIR nccl.h + HINTS "${NCCL_PREFIX}/include" + NO_DEFAULT_PATH) +if(NOT NCCL_INCLUDE_DIR) + message(FATAL_ERROR + "nccl.h not found under the prefix of ${NCCL_LIB}.") +endif() list(APPEND test_comm_gemm_LINKER_LIBS CUDA::cuda_driver CUDA::cudart From 4230f08c31cda4fbc6ca236511be302832b9d9be Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 11 Jun 2026 08:23:02 -0700 Subject: [PATCH 50/55] tests/cpp_distributed: template EP unit tests on token dtype; switch routing to global counter Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/test_ep.cu | 189 +++++++++++++++---------- tests/cpp_distributed/test_ep_common.h | 7 +- 2 files changed, 116 insertions(+), 80 deletions(-) diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu index 1a67644d06..c7fee7720c 100644 --- a/tests/cpp_distributed/test_ep.cu +++ b/tests/cpp_distributed/test_ep.cu @@ -14,7 +14,7 @@ * EPDispatchBwdGradWeightsTest/RoundTrip : exact per-(t, k) grad_topk_weights * EPPipelineTest/FullForwardBackward : fwd + bwd NaN/Inf check * - * Routing: token t on rank r -> expert (r * num_local_experts + t * top_k + k) % num_experts + * Routing: token t on rank r -> expert (r * num_tokens * top_k + t * top_k + k) % num_experts * Token values: rank r, token t -> all hidden dims = (r+1)*0.01 + t*0.001 * * Closed-form expected values: @@ -76,6 +76,7 @@ static std::vector expected_token_counts( return cnt; } +template static std::vector expected_recv_values_sorted( int recv_rank, int num_processes, int num_tokens, int top_k, int num_experts, int num_local_experts) { @@ -88,7 +89,7 @@ static std::vector expected_recv_values_sorted( int64_t e = idx[t * top_k + k]; if (e >= base && e < base + num_local_experts) { float raw = token_value(src, t, num_tokens); - vals.push_back(__bfloat162float(__float2bfloat16(raw))); + vals.push_back(tok_to_float(tok_from_float(raw))); } } } @@ -266,18 +267,33 @@ class EpOpTestBase : public ::testing::Test { } }; +// Pull non-dependent base members into the typed-test scope as local consts so +// the bodies can reference them unqualified. +#define EP_PULL_FIXTURE() \ + const int ep_size_ = this->ep_size_; \ + const int num_experts_ = this->num_experts_; \ + const int num_local_experts_ = this->num_local_experts_; \ + const int hidden_dim_ = this->hidden_dim_; \ + const int max_tokens_per_rank_ = this->max_tokens_per_rank_; \ + const int top_k_ = this->top_k_; \ + const int num_tokens_ = this->num_tokens_ + // ============================================================================= // EPDispatchTest: exact recv values and per-expert counts. // ============================================================================= -class EPDispatchTest : public EpOpTestBase {}; +template class EPDispatchTest : public EpOpTestBase {}; +using EPBf16Only = ::testing::Types; +TYPED_TEST_SUITE(EPDispatchTest, EPBf16Only); -TEST_F(EPDispatchTest, PrepareAndDispatch) { - EPBuffers<> buf; +TYPED_TEST(EPDispatchTest, PrepareAndDispatch) { + using Tok = TypeParam; + EP_PULL_FIXTURE(); + EPBuffers buf; buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, ep_size_, max_tokens_per_rank_); - upload_inputs(buf); - EPTensors<> t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + this->template upload_inputs(buf); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); NVTE_CHECK_CUDA(cudaMemset(buf.recv_tokens.get(), 0, buf.recv_tokens.bytes())); @@ -307,27 +323,27 @@ TEST_F(EPDispatchTest, PrepareAndDispatch) { // 2. Recv values: read only the filled prefix per local-expert zone, not the // whole recv buffer; avoids false positives from legitimate-zero token values. - std::vector h_recv(buf.recv_capacity * hidden_dim_); + std::vector h_recv(buf.recv_capacity * hidden_dim_); NVTE_CHECK_CUDA(cudaMemcpy(h_recv.data(), buf.recv_tokens.get(), - h_recv.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + h_recv.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); std::vector got_vals; got_vals.reserve(total_recv); size_t slot = 0; for (int e = 0; e < num_local_experts_; ++e) { for (int i = 0; i < got_counts[e]; ++i) { - got_vals.push_back(__bfloat162float(h_recv[slot * hidden_dim_])); + got_vals.push_back(tok_to_float(h_recv[slot * hidden_dim_])); ++slot; } } std::sort(got_vals.begin(), got_vals.end()); - auto exp_vals = expected_recv_values_sorted(g_process_id, g_num_processes, num_tokens_, - top_k_, num_experts_, num_local_experts_); + auto exp_vals = expected_recv_values_sorted(g_process_id, g_num_processes, num_tokens_, + top_k_, num_experts_, num_local_experts_); ASSERT_EQ(got_vals.size(), exp_vals.size()); for (size_t i = 0; i < exp_vals.size(); ++i) - EXPECT_NEAR(got_vals[i], exp_vals[i], bf16_tol(exp_vals[i])) + EXPECT_EQ(got_vals[i], exp_vals[i]) << "recv value mismatch at sorted index " << i; // 3. recv_topk_weights: every filled slot must equal the per-token weight (1/top_k). @@ -348,14 +364,17 @@ TEST_F(EPDispatchTest, PrepareAndDispatch) { // EPCombineTest: round-trip identity expert -> result == top_k * tokens. // ============================================================================= -class EPCombineTest : public EpOpTestBase {}; +template class EPCombineTest : public EpOpTestBase {}; +TYPED_TEST_SUITE(EPCombineTest, EPBf16Only); -TEST_F(EPCombineTest, Combine) { - EPBuffers<> buf; +TYPED_TEST(EPCombineTest, Combine) { + using Tok = TypeParam; + EP_PULL_FIXTURE(); + EPBuffers buf; buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, ep_size_, max_tokens_per_rank_); - upload_inputs(buf); - EPTensors<> t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + this->template upload_inputs(buf); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); @@ -369,14 +388,14 @@ TEST_F(EPCombineTest, Combine) { t.result.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - std::vector h_result(num_tokens_ * hidden_dim_); + std::vector h_result(num_tokens_ * hidden_dim_); NVTE_CHECK_CUDA(cudaMemcpy(h_result.data(), buf.result.get(), - h_result.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); - auto h_tok = generate_tokens(g_process_id, num_tokens_, hidden_dim_); + h_result.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); + auto h_tok = generate_tokens(g_process_id, num_tokens_, hidden_dim_); for (int tok = 0; tok < num_tokens_; ++tok) { - float exp = __bfloat162float(h_tok[tok * hidden_dim_]) * static_cast(top_k_); + float exp = tok_to_float(h_tok[tok * hidden_dim_]) * static_cast(top_k_); for (int p = 0; p < hidden_dim_; ++p) { - float got = __bfloat162float(h_result[tok * hidden_dim_ + p]); + float got = tok_to_float(h_result[tok * hidden_dim_ + p]); EXPECT_NEAR(got, exp, bf16_tol(exp)) << "token " << tok << " rank " << g_process_id << " hidden " << p; } @@ -392,14 +411,17 @@ TEST_F(EPCombineTest, Combine) { // EPCombineBwdTest: filled slots in grad_expert == d_result (unweighted). // ============================================================================= -class EPCombineBwdTest : public EpOpTestBase {}; +template class EPCombineBwdTest : public EpOpTestBase {}; +TYPED_TEST_SUITE(EPCombineBwdTest, EPBf16Only); -TEST_F(EPCombineBwdTest, CombineBwdCheck) { - EPBuffers<> buf; +TYPED_TEST(EPCombineBwdTest, CombineBwdCheck) { + using Tok = TypeParam; + EP_PULL_FIXTURE(); + EPBuffers buf; buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, ep_size_, max_tokens_per_rank_); - upload_inputs(buf); - EPTensors<> t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + this->template upload_inputs(buf); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); @@ -412,9 +434,9 @@ TEST_F(EPCombineBwdTest, CombineBwdCheck) { ASSERT_NO_THROW(nvte_ep_combine(t.handle_mem.data(), t.recv_tokens.data(), NVTECommWindow{}, t.result.data(), stream)); - std::vector h_grad_r(num_tokens_ * hidden_dim_, __float2bfloat16(0.1f)); + std::vector h_grad_r(num_tokens_ * hidden_dim_, tok_from_float(0.1f)); NVTE_CHECK_CUDA(cudaMemcpyAsync(buf.grad_result.get(), h_grad_r.data(), - h_grad_r.size() * sizeof(nv_bfloat16), + h_grad_r.size() * sizeof(Tok), cudaMemcpyHostToDevice, stream)); NVTE_CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); @@ -422,23 +444,23 @@ TEST_F(EPCombineBwdTest, CombineBwdCheck) { t.grad_expert.data(), NVTECommWindow{}, stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - int total_recv = read_total_recv(buf); + int total_recv = this->template read_total_recv(buf); std::vector cnt(num_local_experts_); NVTE_CHECK_CUDA(cudaMemcpy(cnt.data(), buf.token_counts.get(), num_local_experts_ * sizeof(int32_t), cudaMemcpyDeviceToHost)); - std::vector h_ge(buf.recv_capacity * hidden_dim_); + std::vector h_ge(buf.recv_capacity * hidden_dim_); NVTE_CHECK_CUDA(cudaMemcpy(h_ge.data(), buf.grad_expert.get(), - h_ge.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + h_ge.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); // Walk filled slots by per-expert zone (no v != 0 heuristic). - const float kExpGrad = 0.1f; + const float kExpGrad = tok_to_float(tok_from_float(0.1f)); size_t slot = 0; int filled = 0; for (int e = 0; e < num_local_experts_; ++e) { for (int i = 0; i < cnt[e]; ++i) { for (int p = 0; p < hidden_dim_; ++p) { - float v = __bfloat162float(h_ge[slot * hidden_dim_ + p]); + float v = tok_to_float(h_ge[slot * hidden_dim_ + p]); EXPECT_NEAR(v, kExpGrad, bf16_tol(kExpGrad)) << "grad_expert expert " << e << " slot " << i << " (linear " << slot << ") hidden " << p; @@ -458,14 +480,17 @@ TEST_F(EPCombineBwdTest, CombineBwdCheck) { // EPDispatchBwdTest: grad_tokens == top_k * d_result. // ============================================================================= -class EPDispatchBwdTest : public EpOpTestBase {}; +template class EPDispatchBwdTest : public EpOpTestBase {}; +TYPED_TEST_SUITE(EPDispatchBwdTest, EPBf16Only); -TEST_F(EPDispatchBwdTest, DispatchBwdCheck) { - EPBuffers<> buf; +TYPED_TEST(EPDispatchBwdTest, DispatchBwdCheck) { + using Tok = TypeParam; + EP_PULL_FIXTURE(); + EPBuffers buf; buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, ep_size_, max_tokens_per_rank_); - upload_inputs(buf); - EPTensors<> t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + this->template upload_inputs(buf); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); @@ -478,9 +503,9 @@ TEST_F(EPDispatchBwdTest, DispatchBwdCheck) { ASSERT_NO_THROW(nvte_ep_combine(t.handle_mem.data(), t.recv_tokens.data(), NVTECommWindow{}, t.result.data(), stream)); - std::vector h_grad(num_tokens_ * hidden_dim_, __float2bfloat16(0.1f)); + std::vector h_grad(num_tokens_ * hidden_dim_, tok_from_float(0.1f)); NVTE_CHECK_CUDA(cudaMemcpyAsync(buf.grad_result.get(), h_grad.data(), - h_grad.size() * sizeof(nv_bfloat16), + h_grad.size() * sizeof(Tok), cudaMemcpyHostToDevice, stream)); NVTE_CHECK_CUDA(cudaMemsetAsync(buf.grad_expert.get(), 0, buf.grad_expert.bytes(), stream)); NVTE_CHECK_CUDA(cudaMemsetAsync(buf.g_recv_topk_weights.get(), 0, buf.g_recv_topk_weights.bytes(), stream)); @@ -493,13 +518,13 @@ TEST_F(EPDispatchBwdTest, DispatchBwdCheck) { t.grad_tokens.data(), t.grad_topk_weights.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - std::vector h_gt(num_tokens_ * hidden_dim_); + std::vector h_gt(num_tokens_ * hidden_dim_); NVTE_CHECK_CUDA(cudaMemcpy(h_gt.data(), buf.grad_tokens.get(), - h_gt.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); - const float kExpGrad = static_cast(top_k_) * 0.1f; + h_gt.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); + const float kExpGrad = static_cast(top_k_) * tok_to_float(tok_from_float(0.1f)); for (int tok = 0; tok < num_tokens_; ++tok) for (int p = 0; p < hidden_dim_; ++p) - EXPECT_NEAR(__bfloat162float(h_gt[tok * hidden_dim_ + p]), kExpGrad, + EXPECT_NEAR(tok_to_float(h_gt[tok * hidden_dim_ + p]), kExpGrad, bf16_tol(kExpGrad)) << "grad_tokens token " << tok << " hidden " << p; @@ -513,21 +538,25 @@ TEST_F(EPDispatchBwdTest, DispatchBwdCheck) { // EPDispatchBwdGradWeightsTest: round-trip per-(t, k) weights. // ============================================================================= -class EPDispatchBwdGradWeightsTest : public EpOpTestBase {}; +template class EPDispatchBwdGradWeightsTest : public EpOpTestBase {}; +TYPED_TEST_SUITE(EPDispatchBwdGradWeightsTest, EPBf16Only); -TEST_F(EPDispatchBwdGradWeightsTest, RoundTrip) { - EPBuffers<> buf; +TYPED_TEST(EPDispatchBwdGradWeightsTest, RoundTrip) { + using Tok = TypeParam; + EP_PULL_FIXTURE(); + EPBuffers buf; buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, ep_size_, max_tokens_per_rank_); - upload_inputs(buf); - EPTensors<> t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + this->template upload_inputs(buf); + EPTensors t(buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); // Distinct per-(rank, t, k) weights so each slot carries a unique value. + // Global integer counter over (rank, tok, k) keeps every slot unique. std::vector h_w(num_tokens_ * top_k_); for (int tok = 0; tok < num_tokens_; ++tok) for (int k = 0; k < top_k_; ++k) - h_w[tok * top_k_ + k] = 0.1f + 0.01f * tok + 0.001f * k + - 0.0001f * (g_process_id + 1); + h_w[tok * top_k_ + k] = static_cast( + (g_process_id * num_tokens_ + tok) * top_k_ + k + 1); NVTE_CHECK_CUDA(cudaMemcpy(buf.topk_weights.get(), h_w.data(), h_w.size() * sizeof(float), cudaMemcpyHostToDevice)); @@ -704,22 +733,28 @@ static inline NVTECommWindow symm_window(const SymmBuf& b) { // Tests rebootstrap the backend to zero_copy=ON for the symm phase via // ep_reinitialize(); TearDown restores OFF for the rest of the suite. +template class EPZeroCopyTest : public EpOpTestBase { protected: void TearDown() override { if (g_ep_initialized) ep_reinitialize(/*zero_copy=*/0); } }; +TYPED_TEST_SUITE(EPZeroCopyTest, EPBf16Only); // Identity round-trip with symm-mem on dispatch i/o + combine input. Bit-exact // vs HBM reference (same routing, same input). -TEST_F(EPZeroCopyTest, IdentityAllSymm) { +TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { + using Tok = TypeParam; + EP_PULL_FIXTURE(); + constexpr DType kTokDType = test::TypeInfo::dtype; + // HBM reference run. - EPBuffers<> ref_buf; + EPBuffers ref_buf; ref_buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, ep_size_, max_tokens_per_rank_); - upload_inputs(ref_buf); - EPTensors<> ref_t(ref_buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + this->template upload_inputs(ref_buf); + EPTensors ref_t(ref_buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); cudaStream_t stream; NVTE_CHECK_CUDA(cudaStreamCreate(&stream)); @@ -733,37 +768,37 @@ TEST_F(EPZeroCopyTest, IdentityAllSymm) { ref_t.result.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - std::vector ref_recv(ref_buf.recv_capacity * hidden_dim_); - std::vector ref_result(num_tokens_ * hidden_dim_); + std::vector ref_recv(ref_buf.recv_capacity * hidden_dim_); + std::vector ref_result(num_tokens_ * hidden_dim_); NVTE_CHECK_CUDA(cudaMemcpy(ref_recv.data(), ref_buf.recv_tokens.get(), - ref_recv.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + ref_recv.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); NVTE_CHECK_CUDA(cudaMemcpy(ref_result.data(), ref_buf.result.get(), - ref_result.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + ref_result.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); // Switch backend to zero_copy=ON for the symm phase. ep_reinitialize(/*zero_copy=*/1); // Symm-mem run: tokens, recv_tokens, combine_input (== recv_tokens) all symm. - EPBuffers<> sym_buf; // alloc all buffers except the symm ones. + EPBuffers sym_buf; // alloc all buffers except the symm ones. sym_buf.alloc(num_tokens_, top_k_, hidden_dim_, num_local_experts_, ep_size_, max_tokens_per_rank_); - upload_inputs(sym_buf); + this->template upload_inputs(sym_buf); SymmBuf sym_tokens, sym_recv; - sym_tokens.alloc(num_tokens_ * hidden_dim_ * sizeof(nv_bfloat16)); - sym_recv .alloc(sym_buf.recv_capacity * hidden_dim_ * sizeof(nv_bfloat16)); + sym_tokens.alloc(num_tokens_ * hidden_dim_ * sizeof(Tok)); + sym_recv .alloc(sym_buf.recv_capacity * hidden_dim_ * sizeof(Tok)); // Stage same tokens into the symm-mem input. - auto h_tok = generate_tokens(g_process_id, num_tokens_, hidden_dim_); + auto h_tok = generate_tokens(g_process_id, num_tokens_, hidden_dim_); NVTE_CHECK_CUDA(cudaMemcpy(sym_tokens.ptr, h_tok.data(), - h_tok.size() * sizeof(nv_bfloat16), cudaMemcpyHostToDevice)); + h_tok.size() * sizeof(Tok), cudaMemcpyHostToDevice)); - EPTensors<> sym_t(sym_buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); + EPTensors sym_t(sym_buf, num_tokens_, top_k_, hidden_dim_, num_local_experts_); // Replace the tokens/recv_tokens views with ones pointing at the symm buffers. sym_t.tokens = TensorWrapper(sym_tokens.ptr, - std::vector{(size_t)num_tokens_, (size_t)hidden_dim_}, DType::kBFloat16); + std::vector{(size_t)num_tokens_, (size_t)hidden_dim_}, kTokDType); sym_t.recv_tokens = TensorWrapper(sym_recv.ptr, - std::vector{sym_buf.recv_capacity, (size_t)hidden_dim_}, DType::kBFloat16); + std::vector{sym_buf.recv_capacity, (size_t)hidden_dim_}, kTokDType); ASSERT_NO_THROW(nvte_ep_prepare(sym_t.handle_mem.data(), sym_t.topk_idx.data(), sym_t.token_counts.data(), NVTEEpLayerConfig{sym_t.top_k_, sym_t.alignment_}, stream)); ASSERT_NO_THROW(nvte_ep_dispatch(sym_t.handle_mem.data(), sym_t.topk_idx.data(), @@ -775,20 +810,20 @@ TEST_F(EPZeroCopyTest, IdentityAllSymm) { symm_window(sym_recv), sym_t.result.data(), stream)); NVTE_CHECK_CUDA(cudaStreamSynchronize(stream)); - std::vector sym_recv_host(sym_buf.recv_capacity * hidden_dim_); - std::vector sym_result(num_tokens_ * hidden_dim_); + std::vector sym_recv_host(sym_buf.recv_capacity * hidden_dim_); + std::vector sym_result(num_tokens_ * hidden_dim_); NVTE_CHECK_CUDA(cudaMemcpy(sym_recv_host.data(), sym_recv.ptr, - sym_recv_host.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + sym_recv_host.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); NVTE_CHECK_CUDA(cudaMemcpy(sym_result.data(), sym_buf.result.get(), - sym_result.size() * sizeof(nv_bfloat16), cudaMemcpyDeviceToHost)); + sym_result.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); // Compare per filled recv slot (HBM ref vs symm) and full result. - int total_recv = read_total_recv(sym_buf); + int total_recv = this->template read_total_recv(sym_buf); for (int i = 0; i < total_recv * hidden_dim_; ++i) - ASSERT_EQ(__bfloat162float(sym_recv_host[i]), __bfloat162float(ref_recv[i])) + ASSERT_EQ(tok_to_float(sym_recv_host[i]), tok_to_float(ref_recv[i])) << "recv mismatch at " << i; for (size_t i = 0; i < sym_result.size(); ++i) - ASSERT_EQ(__bfloat162float(sym_result[i]), __bfloat162float(ref_result[i])) + ASSERT_EQ(tok_to_float(sym_result[i]), tok_to_float(ref_result[i])) << "result mismatch at " << i; if (g_process_id == 0) diff --git a/tests/cpp_distributed/test_ep_common.h b/tests/cpp_distributed/test_ep_common.h index b2421ffd10..d5e006cef6 100644 --- a/tests/cpp_distributed/test_ep_common.h +++ b/tests/cpp_distributed/test_ep_common.h @@ -80,13 +80,14 @@ struct DevBuf { // -- Shared routing helper ----------------------------------------------------- // Balanced round-robin routing: token t on rank r maps top_k experts to -// (r * num_local_experts + t * top_k + k) % num_experts +// (r * num_tokens * top_k + t * top_k + k) % num_experts +// i.e. a single global counter over all (rank, t, k) triples mod num_experts. static inline std::vector routing_balanced( - int rank, int num_tokens, int top_k, int num_experts, int num_local_experts) { + int rank, int num_tokens, int top_k, int num_experts, int /*num_local_experts*/) { std::vector idx(num_tokens * top_k); for (int t = 0; t < num_tokens; ++t) for (int k = 0; k < top_k; ++k) - idx[t * top_k + k] = (rank * num_local_experts + t * top_k + k) % num_experts; + idx[t * top_k + k] = (rank * num_tokens * top_k + t * top_k + k) % num_experts; return idx; } From 053032554cb669f760152c0525a6d3c34ec694b8 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 11 Jun 2026 08:36:20 -0700 Subject: [PATCH 51/55] ep: drop misleading rebuild advice in NCCL EP errors and submodule header version log Signed-off-by: Phuong Nguyen --- transformer_engine/__init__.py | 8 +++----- transformer_engine/common/CMakeLists.txt | 12 ------------ 2 files changed, 3 insertions(+), 17 deletions(-) diff --git a/transformer_engine/__init__.py b/transformer_engine/__init__.py index 91fd7f6f1a..480a2e9a06 100644 --- a/transformer_engine/__init__.py +++ b/transformer_engine/__init__.py @@ -44,15 +44,13 @@ def require_nccl_ep() -> None: cur = _nccl_runtime_version() if cur is None: raise RuntimeError( - f"NCCL EP requires libnccl.so.2 (>= {mn}); could not load libnccl.so.2 " - "or query its version. Install NCCL or rebuild Transformer Engine with " - "NVTE_WITH_NCCL_EP=0." + f"NCCL EP requires NCCL >= {mn}; could not load libnccl.so.2 or query its " + "version. Install NCCL or ensure libnccl.so.2 is on the loader path." ) if cur < _NCCL_EP_MIN_VERSION: raise RuntimeError( f"NCCL EP requires NCCL >= {mn} at runtime; found " - f"{'.'.join(str(x) for x in cur)}. Upgrade libnccl.so or rebuild " - "Transformer Engine with NVTE_WITH_NCCL_EP=0." + f"{'.'.join(str(x) for x in cur)}. Upgrade NCCL to a compatible version." ) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 28cb71dcf9..edb8c5e109 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -485,18 +485,6 @@ if(NOT NCCL_LIB) REQUIRED) endif() -# Diagnostic: log detected NCCL header version (minimum enforced at runtime). -file(READ "${NCCL_INCLUDE_DIR}/nccl.h" _nvte_nccl_h) -string(REGEX MATCH "#define[ \t]+NCCL_MAJOR[ \t]+([0-9]+)" _ "${_nvte_nccl_h}") -set(_nvte_nccl_major "${CMAKE_MATCH_1}") -string(REGEX MATCH "#define[ \t]+NCCL_MINOR[ \t]+([0-9]+)" _ "${_nvte_nccl_h}") -set(_nvte_nccl_minor "${CMAKE_MATCH_1}") -string(REGEX MATCH "#define[ \t]+NCCL_PATCH[ \t]+([0-9]+)" _ "${_nvte_nccl_h}") -set(_nvte_nccl_patch "${CMAKE_MATCH_1}") -if(_nvte_nccl_major AND _nvte_nccl_minor AND _nvte_nccl_patch) - message(STATUS "NCCL header: ${NCCL_INCLUDE_DIR}/nccl.h (version ${_nvte_nccl_major}.${_nvte_nccl_minor}.${_nvte_nccl_patch})") -endif() - target_include_directories(transformer_engine PRIVATE ${NCCL_EP_INCLUDE_DIR} ${NCCL_INCLUDE_DIR}) From 70b216176ead3929d57893154fa64686f5ff1267 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 11 Jun 2026 08:48:01 -0700 Subject: [PATCH 52/55] common/ep: drop multicast support check; NCCL EP does not require multicast Signed-off-by: Phuong Nguyen --- transformer_engine/common/ep/ep_backend.cpp | 2 -- 1 file changed, 2 deletions(-) diff --git a/transformer_engine/common/ep/ep_backend.cpp b/transformer_engine/common/ep/ep_backend.cpp index b1fe8baf83..f1510693bb 100644 --- a/transformer_engine/common/ep/ep_backend.cpp +++ b/transformer_engine/common/ep/ep_backend.cpp @@ -115,8 +115,6 @@ void EPBackend::validate_config(const NVTEEpGroupConfig& config) { const int sm = cuda::sm_arch(); NVTE_CHECK(sm >= 90, "NCCL EP requires SM_90+ (Hopper or later), but current device is SM_", sm); - NVTE_CHECK(cuda::supports_multicast(), - "NCCL EP requires CUDA multicast support on the current device"); } void EPBackend::initialize(ncclComm_t ep_comm, NVTEEpGroupConfig config) { From d1bd816f09f258689eca0092ac0244e413d128c8 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Thu, 11 Jun 2026 22:26:46 -0700 Subject: [PATCH 53/55] setup/ep: error out when NVTE_CUDA_ARCHS yields no SM>=90 arch for NCCL EP Signed-off-by: Phuong Nguyen --- setup.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/setup.py b/setup.py index 9fb962bbc7..c63d76f5ff 100644 --- a/setup.py +++ b/setup.py @@ -246,6 +246,11 @@ def build_nccl_ep_submodule() -> str: bare = t.rstrip("af") if bare.isdigit() and int(bare) >= 90 and bare not in arch_list: arch_list.append(bare) + if not arch_list: + raise RuntimeError( + f"NCCL EP requires Hopper or newer (SM >= 90); none found in NVTE_CUDA_ARCHS={cuda_archs()!r}. " + "Re-run with NVTE_WITH_NCCL_EP=0 to skip the NCCL EP build (the rest of TE still builds)." + ) gencode = " ".join(f"-gencode=arch=compute_{a},code=sm_{a}" for a in arch_list) nproc = os.cpu_count() or 8 From 128c5b68cbd155544e97bf61252143b5c2d1e6bb Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 12 Jun 2026 05:48:53 +0000 Subject: [PATCH 54/55] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- setup.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index c63d76f5ff..64ed120268 100644 --- a/setup.py +++ b/setup.py @@ -248,8 +248,9 @@ def build_nccl_ep_submodule() -> str: arch_list.append(bare) if not arch_list: raise RuntimeError( - f"NCCL EP requires Hopper or newer (SM >= 90); none found in NVTE_CUDA_ARCHS={cuda_archs()!r}. " - "Re-run with NVTE_WITH_NCCL_EP=0 to skip the NCCL EP build (the rest of TE still builds)." + "NCCL EP requires Hopper or newer (SM >= 90); none found in" + f" NVTE_CUDA_ARCHS={cuda_archs()!r}. Re-run with NVTE_WITH_NCCL_EP=0 to skip the NCCL" + " EP build (the rest of TE still builds)." ) gencode = " ".join(f"-gencode=arch=compute_{a},code=sm_{a}" for a in arch_list) From 3459461ed89b178c4932d7b2558f5b312bdaeffc Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Fri, 12 Jun 2026 12:12:03 -0700 Subject: [PATCH 55/55] make core to be RTLD_LAZY Signed-off-by: Phuong Nguyen --- transformer_engine/common/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/common/__init__.py b/transformer_engine/common/__init__.py index fd2d146616..42b458bfc5 100644 --- a/transformer_engine/common/__init__.py +++ b/transformer_engine/common/__init__.py @@ -380,7 +380,7 @@ def _load_cuda_library(lib_name: str): @functools.lru_cache(maxsize=None) def _load_core_library(): """Load shared library with Transformer Engine C extensions""" - return ctypes.CDLL(_get_shared_object_file("core"), mode=ctypes.RTLD_GLOBAL) + return ctypes.CDLL(_get_shared_object_file("core"), mode=ctypes.RTLD_GLOBAL | os.RTLD_LAZY) if "NVTE_PROJECT_BUILDING" not in os.environ or bool(int(os.getenv("NVTE_RELEASE_BUILD", "0"))):