Expert Parallelism: common C API + NCCL EP backend

[phu0ngng]

Summary

First PR in the TE Expert Parallelism (EP) series. Lands the common C API and NCCL EP backend that later framework PRs (PyTorch, JAX) build on. No Python bindings yet — common-lib foundation plus build wiring only. Build/load works on any arch; SM and NCCL version gates fire at runtime.

Every network-bound payload tensor takes an optional NVTECommWindow. When the window is provided, the backend uses NCCL EP's symmetric-memory zero-copy path, which skips the D2D Memcpy from the user buffers to the Symmetric Staging Buffers.

Implementation

Public C API (transformer_engine/common/include/transformer_engine/{ep.h,comm_window.h})

Types: NVTEEpGroupConfig, NVTEEpLayerConfig, NVTEEpHandle, NVTECommWindow (side-band {ncclWindow_t window, size_t offset}; NCCL peer handles are not carried on NVTETensor).

Lifecycle (host-only, eager):

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);

• nvte_ep_initialize — borrow an external ncclComm_t for the EP sub-group and init the singleton backend.

• nvte_ep_shutdown — tear down the backend; idempotent; does not destroy ep_comm.

• nvte_ep_register_layer — reserve a handle_id for a layer config and report the handle_mem buffer size the caller must allocate. The pair {id, mem} becomes the per-step NVTEEpHandle.

Per-step (allocation-free, CUDA-graph capturable)

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_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 nvte_ep_combine(NVTEEpHandle handle, NVTETensor expert_out, NVTECommWindow expert_out_win,
                     NVTETensor result, cudaStream_t 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 nvte_ep_combine_bwd(NVTEEpHandle handle, NVTETensor grad, NVTECommWindow grad_win,
                         NVTETensor grad_expert_out, NVTECommWindow grad_expert_out_win,
                         cudaStream_t stream);

• nvte_ep_prepare — all-gather the routing map and write routing maps to handle.mem.
• nvte_ep_dispatch — scatter tokens and routing weights from source ranks to expert ranks. tokens, topk_weights, recv_tokens, recv_topk_weights each accept an optional symm-mem window for zero-copy.
• nvte_ep_combine — scatter-sum expert outputs back to source ranks (unweighted; caller pre-multiplies by recv_topk_weights). expert_out accepts a window.
• nvte_ep_dispatch_bwd — backward of dispatch; routes token and weight grads back to source. grad and g_recv_topk_weights accept windows; the gathered outputs (grad_tokens, grad_topk_weights).
• nvte_ep_combine_bwd — backward of combine; grad and grad_expert_out accept windows. Padded slots in grad_expert_out are zeroed.

Backend + build

• NCCL EP backend (transformer_engine/common/ep/): EPBackend singleton, HT-mode dispatch/combine over NCCL EP (libnccl_ep.so), group/layer registration. Internal helper make_payload_tensor() builds the per-call ncclEpTensor_t: when the caller's NVTECommWindow.window != nullptr it sets win_hdl + win_offset (zero-copy); otherwise it sets data from nvte_tensor_data(t) (HBM fallback).
• Runtime gates (in EPBackend::initialize): SM>=90 (via cudaDeviceGetAttribute), NCCL>=2.30.4 (via ncclGetVersion), CUDA multicast/NVLS support.
• Stub path: when NVTE_WITH_NCCL_EP=OFF, ep/ep_api_stub.cpp provides throwing nvte_ep_* stubs so framework bindings link unconditionally; failure surfaces at first nvte_ep_initialize.
• Build wiring
  • setup.py builds libnccl_ep.so from 3rdparty/nccl by default; auto-disables NCCL EP when no requested CUDA arch >= 90. Explicit NVTE_BUILD_WITH_NCCL_EP=1 with all archs < 90 is treated as user error NVTE_BUILD_WITH_NCCL_EP=0 to opt out.
  • NCCL_HOME resolved dynamically: explicit env → /opt/nvidia/nccl, /usr/local/nccl, /usr → ldconfig -p fallback.

Testing

• C++ distributed tests under tests/cpp_distributed/.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Summary

This PR lands the Expert Parallelism (EP) common C API and NCCL EP backend for TransformerEngine, providing the foundation for upcoming PyTorch and JAX framework bindings. It adds the nvte_ep_* C API in ep.h/comm_window.h, an EPBackend singleton backed by a pointer-keyed LRU handle cache, and build wiring that auto-gates on SM ≥ 90 / NCCL ≥ 2.30.4 at both configure and runtime.

• Public C API (ep.h, comm_window.h): lifecycle (initialize/shutdown/handle_mem_size), per-step CUDA-graph-capturable ops (prepare/dispatch/combine and their backward passes), and NVTECommWindow for optional zero-copy symmetric-memory payloads. The stubs for the NVTE_WITH_NCCL_EP=OFF build path are co-located in ep_api.cpp under a preprocessor guard, ensuring signature consistency.
• EPBackend singleton (ep_backend.cpp/ep_backend.h): implements an LRU handle cache keyed on handle_mem device pointer (with an XLA buffer-relocation fallback via fallback_layer_cfg_), dtype-width validation against the group's max_token_dtype, and a caller-owned NVTEShape pattern to keep ncclEpTensor_t.sizes pointers valid across NCCL calls.
• Build wiring (setup.py, CMakeLists.txt): builds libnccl_ep.a from the 3rdparty/nccl submodule, expands NVTE_CUDA_ARCHS=native via nvidia-smi, raises a RuntimeError when no SM ≥ 90 arch is found after expansion, and warns when NCCL_HOME points to a path without nccl.h.

Confidence Score: 5/5

New foundational EP backend with no blocking correctness issues in the current code; the previously identified races, dangling pointers, and stub-signature mismatches appear addressed.

All previously raised issues (initialized_ TOCTOU, ep_group_ mutex races, desc.sizes lifetime, stub signature mismatch, comm_window.h NCCL header dependency, max_token_bytes hardcoding, empty arch_list guard) appear resolved in the current revision. The new findings are limited to an LRU-eviction latency concern and a missing first-call guard for mismatched layer_cfg — both non-blocking robustness issues that don't affect correct usage of the API.

transformer_engine/common/ep/ep_backend.cpp — LRU eviction path and prepare_handle_locked deserve a second look before the framework bindings land on top of this.

Important Files Changed

Filename	Overview
transformer_engine/common/ep/ep_backend.cpp	Core NCCL EP backend singleton — implements LRU handle cache, per-step ops (prepare/dispatch/combine/bwd), mutex discipline, and dtype validation. Previous thread issues (dangling sizes pointer, initialized_ TOCTOU, ep_group_ races, register_layer config mismatch) appear addressed; one new concern about ncclEpHandleDestroy called under mutex during LRU eviction.
transformer_engine/common/ep/ep_api.cpp	C API shim: thin delegations to EPBackend when built with NCCL EP; complete throwing stubs in the same file under !NVTE_WITH_NCCL_EP guard. Stub signatures match ep.h; the previously flagged stub-divergence issue is resolved by co-locating stubs in this file.
transformer_engine/common/include/transformer_engine/ep.h	Public C API header — well-structured with clear Doxygen, NVTEEpGroupConfig/NVTEEpLayerConfig, and lifecycle + per-step function declarations. Layer config mismatch between handle_mem_size and prepare is not enforced; no ABI version field (acknowledged TODO).
transformer_engine/common/include/transformer_engine/comm_window.h	Public header for NVTECommWindow. Uses forward declaration of struct ncclWindow_vidmem instead of #include <nccl.h>, which resolves the build-time NCCL dependency but couples the public ABI to an internal NCCL type name.
setup.py	Build wiring for NCCL EP submodule: native arch expansion via nvidia-smi, empty arch_list guard (RuntimeError), NCCL_HOME warning on invalid path. Previously flagged issues with native-arch handling and empty gencode appear addressed.
transformer_engine/common/ep/ep_backend.h	Internal EPBackend singleton declaration. LRU cache (list + unordered_map), per-member mutex discipline, atomic initialized_ flag, and fallback_layer_cfg_ optional for XLA buffer-relocation workaround are all clearly documented.
tests/cpp_distributed/test_ep.cu	Distributed C++ tests covering prepare/dispatch/combine and their backward passes with closed-form expected values. Deterministic routing and BF16/FP16/FP32 type coverage look solid.
qa/L1_cpp_distributed/test.sh	QA harness updated to build test_comm_gemm and test_ep independently, accumulate failures instead of hard-stopping, and write JUnit XML. Error handling is improved over the original set -e approach.
tests/cpp_distributed/run_test_ep.sh	MPI launch wrapper for EP tests. Pre-Hopper SM check via nvidia-smi exits cleanly with code 0; GTEST_XML_PREFIX per-rank output avoids write races. set -euo pipefail is appropriate here.

Sequence Diagram

sequenceDiagram
    participant FW as Framework (PyTorch/JAX)
    participant API as nvte_ep_* (ep_api.cpp)
    participant BE as EPBackend singleton
    participant Cache as LRU Handle Cache
    participant NCCL as NCCL EP (ncclEp*)

    FW->>API: nvte_ep_initialize(ep_comm, group_config)
    API->>BE: EPBackend::initialize()
    BE->>NCCL: ncclGetVersion() [version gate]
    BE->>NCCL: ncclEpCreateGroup(ep_group_)

    FW->>API: nvte_ep_handle_mem_size(layer_cfg)
    API->>BE: handle_mem_size() [acquires mutex]
    BE->>NCCL: ncclEpHandleMemSize()
    BE-->>FW: hm_size bytes

    Note over FW: allocate handle_mem[hm_size]

    FW->>API: nvte_ep_prepare(handle_mem, topk_idx, ...)
    API->>BE: prepare() [acquires mutex]
    BE->>Cache: prepare_handle_locked(handle_mem, layer_cfg)
    Cache->>NCCL: ncclEpInitHandle() [on cache miss]
    BE->>NCCL: ncclEpUpdateHandle() [routing AllGather, async stream]

    FW->>API: nvte_ep_dispatch(handle_mem, tokens, ...)
    API->>BE: dispatch() [acquires mutex]
    BE->>Cache: lookup_handle_locked(handle_mem)
    BE->>NCCL: ncclEpDispatch() [async stream]

    FW->>API: nvte_ep_combine(handle_mem, expert_out, ...)
    API->>BE: combine() [acquires mutex]
    BE->>Cache: lookup_handle_locked(handle_mem)
    BE->>NCCL: ncclEpCombine() [async stream]

    FW->>API: nvte_ep_shutdown()
    API->>BE: EPBackend::shutdown() [acquires mutex]
    BE->>NCCL: ncclEpHandleDestroy() [each LRU entry]
    BE->>NCCL: ncclEpGroupDestroy(ep_group_)

Loading

_{Reviews (21): Last reviewed commit: "make core to be RTLD_LAZY" | Re-trigger Greptile}

[phu0ngng]

/te-ci L1

[phu0ngng]

Pipeline 54600556

[timmoon10]

Nit: expert_parallel.h or expert_parallelism.h would be a less ambiguous file name, especially in a public-facing header.

Asking a few AIs "What does EP mean in deep learning?":

• Claude: "Expectation Propagation", "Energy Propagation", or "Episode"
• Gemini: "Epoch", various non-MoE models (EP-DNN, EP-PINN, EP-RNN)
• ChatGPT: "Expert Parallelism", "Epoch", "Embedded Projection"

[timmoon10]

This could be a standalone utility funciton in build_utils/utils.py. It's actually quite general and not specific to NCCL EP.

[timmoon10]

Flagging that we compile libnccl_ep.a at a different time than when we compile libtransformer_engine.so.

• libnccl_ep.a: Compiled during the Setuptools configuration stage, while processing setup.py.
• libtransformer_engine.so: Compiled during the Setuptools build stage, configured via a custom setuptools.Extension.

It's not obviously wrong, but it does seem messy. I suppose there are also some edge case problems (if NCCL is not installed but is listed as a dependency, compiling NCCL EP in the configuration stage will fail but building during the build stage can succeed).

It seems the "proper" approach would be to treat libnccl_ep.a compilation as basically a part of compiling libtransformer_engine.so. However, CMakeExtension is a very general class for CMake projects and I'm not sure how messy it would be to make add the NCCL EP compilation logic.

[timmoon10]

This seems likely to break users with wimpy nodes. We can reuse an existing API to control build parallelism:

Suggested change

	nproc = os.cpu_count() or 8
	nproc = get_max_jobs_for_parallel_build()

[phu0ngng]

Pipeline 54600556 - all tests passed.

[timmoon10]

This still has some minor issues, but deadlines are what they are and the changes are well-contained. We should incorporate suggestions into future PRs.