Skip to content

[Common] Migrate NCCL EP submodule to NVIDIA/nccl-extensions #3222

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
phu0ngng merged 8 commits into from
Jul 22, 2026
Merged

[Common] Migrate NCCL EP submodule to NVIDIA/nccl-extensions #3222

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
9c51a5d
[PyTorch] Migrate NCCL EP submodule to NVIDIA/nccl-extensions
phu0ngng Jul 20, 2026
d7ea20a
[PyTorch] Drop PYTHONPATH override from EP test, example, and bench l…
phu0ngng Jul 20, 2026
428ccee
[JAX] Fix stale comment on LD_LIBRARY_PATH in EP launchers (needed fo…
phu0ngng Jul 20, 2026
9791468
Merge branch 'main' into ep/pyt_api
phu0ngng Jul 21, 2026
6366ca7
[PyTorch] Update stale nccl_ep.cc path in test_ep.cu comment after nc…
phu0ngng Jul 20, 2026
a64b320
[Common] Drop cross-mode recv comparison in EP zero-copy IdentityAllS…
phu0ngng Jul 21, 2026
75bd865
[Common] Rename 3rdparty/nccl submodule directory to nccl-extensions
phu0ngng Jul 21, 2026
0c949d9
Merge branch 'main' into ep/pyt_api
phu0ngng Jul 22, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .gitmodules
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,5 +5,5 @@
path = 3rdparty/cutlass
url = https://github.com/NVIDIA/cutlass.git
[submodule "3rdparty/nccl"]
path = 3rdparty/nccl
url = https://github.com/NVIDIA/nccl.git
path = 3rdparty/nccl-extensions
url = https://github.com/NVIDIA/nccl-extensions.git
Comment on lines 7 to +9

@greptile-apps greptile-apps Bot Jul 22, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Submodule name/path mismatch after rename

The section header still reads [submodule "3rdparty/nccl"] while path was changed to 3rdparty/nccl-extensions. Git uses the name (the string in square brackets) to locate the module cache under .git/modules/<name> and to key the entry in .git/config. Existing clones will have .git/config with [submodule "3rdparty/nccl"] path = 3rdparty/nccl — after a git pull, git will see a stale .git/config entry pointing at the old path, and git submodule update --init --recursive can silently leave 3rdparty/nccl-extensions uninitialised (or error out) on those clones. A clean rename requires updating the section header to [submodule "3rdparty/nccl-extensions"] so name and path are in sync.

1 change: 0 additions & 1 deletion 3rdparty/nccl
Open in desktop
Submodule nccl deleted from b87848
1 change: 1 addition & 0 deletions 3rdparty/nccl-extensions
Open in desktop
Submodule nccl-extensions added at 2c6135
4 changes: 2 additions & 2 deletions examples/jax/ep/run_test_ep.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}')
Expand Down
2 changes: 0 additions & 2 deletions examples/pytorch/ep/bench/run_ep_bench.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}}"
Expand Down
4 changes: 2 additions & 2 deletions examples/pytorch/ep/bench/run_nccl_ep_bench.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
1 change: 0 additions & 1 deletion examples/pytorch/ep/run_test_ep.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)"}
Expand Down
12 changes: 6 additions & 6 deletions setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -207,11 +207,11 @@ 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():
"""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 submodule not found at {nccl_root}. "
f"NCCL EP submodule not found at {nccl_root}. "
"Run `git submodule update --init --recursive`."
)

Expand Down Expand Up @@ -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,
)
Expand Down
21 changes: 6 additions & 15 deletions tests/cpp_distributed/test_ep.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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";
}
Expand Down Expand Up @@ -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();
Expand All @@ -776,10 +777,7 @@ TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) {
ref_t.result.data(), stream));
NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));

std::vector<Tok> ref_recv(ref_buf.recv_capacity * hidden_dim_);
std::vector<Tok> 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));

Expand Down Expand Up @@ -818,24 +816,17 @@ TYPED_TEST(EPZeroCopyTest, IdentityAllSymm) {
symm_window(sym_recv), sym_t.result.data(), stream));
NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));

std::vector<Tok> sym_recv_host(sym_buf.recv_capacity * hidden_dim_);
std::vector<Tok> 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<Tok>(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));
}
Expand Down
4 changes: 2 additions & 2 deletions tests/jax/multi_process_launch_ep.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}')
Expand Down
2 changes: 0 additions & 2 deletions tests/pytorch/distributed/run_test_ep.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions transformer_engine/common/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Loading