From 9c51a5d74ceb2f8ab26e913366065004a1829282 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 20 Jul 2026 11:46:16 -0700 Subject: [PATCH 1/6] [PyTorch] Migrate NCCL EP submodule to NVIDIA/nccl-extensions Signed-off-by: Phuong Nguyen --- .gitmodules | 2 +- 3rdparty/nccl | 2 +- setup.py | 8 ++++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/.gitmodules b/.gitmodules index 07647e915e..7004ea922b 100644 --- a/.gitmodules +++ b/.gitmodules @@ -6,4 +6,4 @@ url = https://github.com/NVIDIA/cutlass.git [submodule "3rdparty/nccl"] path = 3rdparty/nccl - url = https://github.com/NVIDIA/nccl.git + url = https://github.com/NVIDIA/nccl-extensions.git diff --git a/3rdparty/nccl b/3rdparty/nccl index b87848fbc5..2c6135a721 160000 --- a/3rdparty/nccl +++ b/3rdparty/nccl @@ -1 +1 @@ -Subproject commit b87848fbc52da65b5a898b4ac6633fcf51cec4ed +Subproject commit 2c6135a721824ff792af7b72900b0ab758fa1f98 diff --git a/setup.py b/setup.py index 150d92969c..6beb690415 100644 --- a/setup.py +++ b/setup.py @@ -209,9 +209,9 @@ def _discover_nccl_home() -> str: 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" - if not (nccl_root / "Makefile").exists(): + if not (nccl_root / "nccl_ep" / "Makefile").exists(): raise RuntimeError( - f"NCCL submodule not found at {nccl_root}. " + f"NCCL EP submodule not found at {nccl_root}. " "Run `git submodule update --init --recursive`." ) @@ -267,13 +267,13 @@ def build_nccl_ep_submodule() -> str: "rebuilding libnccl_ep.a" ) subprocess.check_call( - ["make", "-C", "contrib/nccl_ep", "clean"], + ["make", "-C", "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"], + ["make", "-j", str(nproc), "-C", "nccl_ep", "lib"], cwd=str(nccl_root), env=env, ) From d7ea20af564dd43ba2ed463e17e201fc715e90c4 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 20 Jul 2026 11:59:54 -0700 Subject: [PATCH 2/6] [PyTorch] Drop PYTHONPATH override from EP test, example, and bench launchers Signed-off-by: Phuong Nguyen --- examples/pytorch/ep/bench/run_ep_bench.sh | 2 -- examples/pytorch/ep/run_test_ep.sh | 1 - tests/pytorch/distributed/run_test_ep.sh | 2 -- 3 files changed, 5 deletions(-) diff --git a/examples/pytorch/ep/bench/run_ep_bench.sh b/examples/pytorch/ep/bench/run_ep_bench.sh index fefecd7fa9..3b0977e4c3 100755 --- a/examples/pytorch/ep/bench/run_ep_bench.sh +++ b/examples/pytorch/ep/bench/run_ep_bench.sh @@ -26,10 +26,8 @@ if [ "${NSYS}" -eq 1 ] && [ "${KINETO}" -eq 1 ]; then fi SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" RESULTS="${SCRIPT_DIR}/results" mkdir -p "${RESULTS}" -export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) NUM_GPUS="${NUM_GPUS:-${DETECTED_GPUS}}" diff --git a/examples/pytorch/ep/run_test_ep.sh b/examples/pytorch/ep/run_test_ep.sh index 13b41f4cb2..d8e6b50556 100755 --- a/examples/pytorch/ep/run_test_ep.sh +++ b/examples/pytorch/ep/run_test_ep.sh @@ -17,7 +17,6 @@ if [ "${NUM_GPUS}" -gt 8 ]; then NUM_GPUS=8; fi : ${TEST_TIMEOUT_S:=120} SCRIPT="${TE_PATH}/examples/pytorch/ep/ep_moe.py" -export PYTHONPATH="${TE_PATH}${PYTHONPATH:+:${PYTHONPATH}}" # Stage JIT cubins on tmpfs for fast iteration. : ${NCCL_EP_JIT_CACHE_DIR:="${TMPDIR:-/tmp}/nccl_ep_jit_cache_$(id -u)"} diff --git a/tests/pytorch/distributed/run_test_ep.sh b/tests/pytorch/distributed/run_test_ep.sh index 68b691f787..62c9a0207f 100755 --- a/tests/pytorch/distributed/run_test_ep.sh +++ b/tests/pytorch/distributed/run_test_ep.sh @@ -9,8 +9,6 @@ set -uo pipefail SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" -export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" DETECTED_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) if [ "${DETECTED_GPUS}" -lt 4 ]; then From 428cceea6f5ad22ce1e3b973a89a31a6664b7aee Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 20 Jul 2026 12:04:05 -0700 Subject: [PATCH 3/6] [JAX] Fix stale comment on LD_LIBRARY_PATH in EP launchers (needed for libtransformer_engine.so, not libnccl_ep) Signed-off-by: Phuong Nguyen --- examples/jax/ep/run_test_ep.sh | 4 ++-- tests/jax/multi_process_launch_ep.sh | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/jax/ep/run_test_ep.sh b/examples/jax/ep/run_test_ep.sh index 1305ca6fd1..86aa6ca087 100755 --- a/examples/jax/ep/run_test_ep.sh +++ b/examples/jax/ep/run_test_ep.sh @@ -32,8 +32,8 @@ export PYTHONPATH="${TE_PATH}${PYTHONPATH:+:${PYTHONPATH}}" COORD="${COORD:-127.0.0.1:12345}" TEST_TIMEOUT_S="${TEST_TIMEOUT_S:-300}" -# Editable installs don't embed rpath; libtransformer_engine.so needs -# libnccl_ep.so.0 from the TE editable location at dlopen time. +# Editable installs don't embed rpath; the TE JAX extension needs +# libtransformer_engine.so from the TE editable location at dlopen time. TE_LIB_PATH=$(pip3 show transformer-engine 2>/dev/null \ | grep -E "Location:|Editable project location:" \ | tail -n 1 | awk '{print $NF}') diff --git a/tests/jax/multi_process_launch_ep.sh b/tests/jax/multi_process_launch_ep.sh index ff89f712eb..8547d77f2b 100755 --- a/tests/jax/multi_process_launch_ep.sh +++ b/tests/jax/multi_process_launch_ep.sh @@ -17,8 +17,8 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)" export PYTHONPATH="${TE_REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" -# Editable installs don't embed rpath; libtransformer_engine.so needs -# libnccl_ep.so.0 from the TE editable location at dlopen time. +# Editable installs don't embed rpath; the TE JAX extension needs +# libtransformer_engine.so from the TE editable location at dlopen time. TE_LIB_PATH=$(pip3 show transformer-engine 2>/dev/null \ | grep -E "Location:|Editable project location:" \ | tail -n 1 | awk '{print $NF}') From 6366ca759dd367a188ff5c61397edd86d59abb9e Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Mon, 20 Jul 2026 12:05:26 -0700 Subject: [PATCH 4/6] [PyTorch] Update stale nccl_ep.cc path in test_ep.cu comment after nccl-extensions migration Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/test_ep.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu index 7dbbcdce9d..da7c8cc251 100644 --- a/tests/cpp_distributed/test_ep.cu +++ b/tests/cpp_distributed/test_ep.cu @@ -675,7 +675,7 @@ class EPPipelineTest : public EpOpTestBase, public ::testing::WithParamInterface 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. + // (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"; } From a64b32098544882e2a784fa1ab4ae85da298baef Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 21 Jul 2026 14:10:59 -0700 Subject: [PATCH 5/6] [Common] Drop cross-mode recv comparison in EP zero-copy IdentityAllSymm test Signed-off-by: Phuong Nguyen --- tests/cpp_distributed/test_ep.cu | 19 +++++-------------- 1 file changed, 5 insertions(+), 14 deletions(-) diff --git a/tests/cpp_distributed/test_ep.cu b/tests/cpp_distributed/test_ep.cu index da7c8cc251..47732196ed 100644 --- a/tests/cpp_distributed/test_ep.cu +++ b/tests/cpp_distributed/test_ep.cu @@ -750,8 +750,9 @@ class EPZeroCopyTest : public EpOpTestBase { }; 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). +// Identity round-trip with symm-mem on dispatch i/o + combine input. The combined +// result is bit-exact vs the HBM reference; the intermediate recv buffer is not, +// since zero-copy and HBM dispatch use different per-expert layouts. TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) { using Tok = TypeParam; EP_PULL_FIXTURE(); @@ -776,10 +777,7 @@ TYPED_TEST(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_); - NVTE_CHECK_CUDA(cudaMemcpy(ref_recv.data(), ref_buf.recv_tokens.get(), - ref_recv.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); NVTE_CHECK_CUDA(cudaMemcpy(ref_result.data(), ref_buf.result.get(), ref_result.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); @@ -818,24 +816,17 @@ TYPED_TEST(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_); - NVTE_CHECK_CUDA(cudaMemcpy(sym_recv_host.data(), sym_recv.ptr, - sym_recv_host.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); NVTE_CHECK_CUDA(cudaMemcpy(sym_result.data(), sym_buf.result.get(), sym_result.size() * sizeof(Tok), cudaMemcpyDeviceToHost)); - // Compare per filled recv slot (HBM ref vs symm) and full result. - int total_recv = this->template read_total_recv(sym_buf); - for (int i = 0; i < total_recv * hidden_dim_; ++i) - ASSERT_EQ(tok_to_float(sym_recv_host[i]), tok_to_float(ref_recv[i])) - << "recv mismatch at " << i; + // Combined result is the cross-mode invariant (see note above). for (size_t i = 0; i < sym_result.size(); ++i) ASSERT_EQ(tok_to_float(sym_result[i]), tok_to_float(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); + printf(" IdentityAllSymm: passed (result bit-exact vs HBM)\n"); NVTE_CHECK_CUDA(cudaStreamDestroy(stream)); } From 75bd865ba2eda30e0381938420f760d54cccf423 Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 21 Jul 2026 14:11:17 -0700 Subject: [PATCH 6/6] [Common] Rename 3rdparty/nccl submodule directory to nccl-extensions Signed-off-by: Phuong Nguyen --- .gitmodules | 2 +- 3rdparty/{nccl => nccl-extensions} | 0 examples/pytorch/ep/bench/run_nccl_ep_bench.sh | 4 ++-- setup.py | 4 ++-- transformer_engine/common/CMakeLists.txt | 4 ++-- 5 files changed, 7 insertions(+), 7 deletions(-) rename 3rdparty/{nccl => nccl-extensions} (100%) diff --git a/.gitmodules b/.gitmodules index 7004ea922b..88d431cba5 100644 --- a/.gitmodules +++ b/.gitmodules @@ -5,5 +5,5 @@ path = 3rdparty/cutlass url = https://github.com/NVIDIA/cutlass.git [submodule "3rdparty/nccl"] - path = 3rdparty/nccl + path = 3rdparty/nccl-extensions url = https://github.com/NVIDIA/nccl-extensions.git diff --git a/3rdparty/nccl b/3rdparty/nccl-extensions similarity index 100% rename from 3rdparty/nccl rename to 3rdparty/nccl-extensions diff --git a/examples/pytorch/ep/bench/run_nccl_ep_bench.sh b/examples/pytorch/ep/bench/run_nccl_ep_bench.sh index 8f6da04a00..ac5fcacc25 100755 --- a/examples/pytorch/ep/bench/run_nccl_ep_bench.sh +++ b/examples/pytorch/ep/bench/run_nccl_ep_bench.sh @@ -23,8 +23,8 @@ TE_REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../.." && pwd)" RESULTS="${SCRIPT_DIR}/results" mkdir -p "${RESULTS}" -BIN="${TE_REPO_ROOT}/3rdparty/nccl/build/test/nccl_ep/ep_bench" -LIB="${TE_REPO_ROOT}/3rdparty/nccl/build/lib" +BIN="${TE_REPO_ROOT}/3rdparty/nccl-extensions/build/test/nccl_ep/ep_bench" +LIB="${TE_REPO_ROOT}/3rdparty/nccl-extensions/build/lib" [ -x "${BIN}" ] || { echo "ep_bench not built at ${BIN}" >&2; exit 2; } NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) diff --git a/setup.py b/setup.py index 6beb690415..c0c37364a9 100644 --- a/setup.py +++ b/setup.py @@ -207,8 +207,8 @@ def _discover_nccl_home() -> str: 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" + """Build libnccl_ep.a from the 3rdparty/nccl-extensions submodule and return NCCL_HOME.""" + nccl_root = current_file_path / "3rdparty" / "nccl-extensions" if not (nccl_root / "nccl_ep" / "Makefile").exists(): raise RuntimeError( f"NCCL EP submodule not found at {nccl_root}. " diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index d09876c990..8eba515e5e 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -446,10 +446,10 @@ 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 +# Headers + libs are produced by the in-tree 3rdparty/nccl-extensions submodule build # (auto-built by setup.py via build_nccl_ep_submodule). set(NCCL_EP_SUBMODULE_ROOT - "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl") + "${CMAKE_CURRENT_SOURCE_DIR}/../../3rdparty/nccl-extensions") set(NCCL_EP_INCLUDE_DIR "${NCCL_EP_SUBMODULE_ROOT}/build/include") if(NOT EXISTS "${NCCL_EP_INCLUDE_DIR}/nccl_ep.h") message(FATAL_ERROR