[JAX] Expert Parallelism: JAX primitives + VJPs

[phu0ngng]

Summary

Third PR in the TE Expert Parallelism (EP) series, built on top of #3034. Lands the JAX bindings: an XLA FFI layer over the nvte_ep_* C API, a Python wrapper with custom_vjp for autograd, mesh-aware sharding rules, a multi-process test suite, and an end-to-end MoE example. NCCL ncclEpDispatch/ncclEpCombine are exposed as XLA primitives and work with CUDA-graph capture.

Implementation

Public Python API (transformer_engine/jax/ep.py)

from transformer_engine.jax.ep import (
    EpHandle,        # opaque (id, handle_mem) pair from ep_prepare
    ep_bootstrap,    # one-shot per-process: init NCCL comm + nvte_ep_initialize
    ep_dispatch,     # custom_vjp-wrapped dispatch 
    ep_combine,      # custom_vjp-wrapped combine

ep_dispatch / ep_combine are jax.custom_vjp functions: forward is the FFI primitive, backward calls the matching nvte_ep_*_bwd FFI primitive directly (no ep_prepare in the bwd — routing state is already cached in handle.mem). Note that ep_dispatch also calls ep_prepare in the forward path, which all-gathers and preprocesses routing maps.

XLA FFI bindings (transformer_engine/jax/csrc/extensions/ep.cpp)

Five XLA_FFI_DEFINE_HANDLER_SYMBOL entries — EpPrepareHandler, EpDispatchHandler, EpCombineHandler, EpDispatchBwdHandler, EpCombineBwdHandler — each calling the corresponding nvte_ep_* C entry point. All marked FFI_CudaGraph_Traits so they capture cleanly. handle_id is a static FFI attribute baked at jit trace time.

Primitives + Python layer (transformer_engine/jax/cpp_extensions/ep.py, +951 lines)

Standard TE primitive plumbing: abstract_eval (shape/dtype inference), lowering, impl, outer_primitive registration, and partitioning rules so the EP collective is treated as a single sharded op by XLA (no spurious resharding around it).

Sharding (transformer_engine/jax/sharding.py, +12 lines)

Adds the EP mesh axis to the global mesh resource set so downstream sharding rules can reference it.

Build wiring (build_tools/jax.py, +41 lines)

Threads NCCL EP linkage through the JAX transformer_engine_jax extension. No new top-level build flags; rides on the parent PR's NVTE_BUILD_WITH_NCCL_EP.

Tests & example

• tests/jax/test_multi_process_ep.py (+690 lines): 13 tests covering bootstrap, ep_prepare shape/handle contracts, primitive-level dispatch/combine identity (uniform + skewed routing), custom_vjp fwd+bwd correctness, and HLO inspection (must not insert XLA collectives outside the EP FFI).
• tests/jax/multi_process_launch_ep.sh: 4-rank launcher; sets XLA_FLAGS to keep XLA command-buffer capture off for the EP FFI sequence (NCCL EP graph-destroy interaction).
• examples/jax/ep/ep_moe.py (+394 lines) + run_test_ep.sh: end-to-end MoE with EP, dp=ep=2 mesh, includes a ref-comparison --check that verifies fwd+bwd vs a single-process reference.

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 JAX Expert Parallelism bindings: five XLA FFI handlers over the nvte_ep_* C API, a Python wrapper with jax.custom_vjp for autograd, mesh-aware SPMD sharding rules, a 13-test multi-process suite, and an end-to-end MoE example. The implementation is substantial (~3,800 lines) and well-structured.

• transformer_engine/jax/cpp_extensions/ep.py: ~960-line primitive layer registering EpPrepare, EpDispatch, EpCombine (fwd+bwd) as JAX primitives with abstract_eval, lowering, impl, batcher, partition, and Shardy rules.
• transformer_engine/jax/ep.py: Public API — ep_bootstrap (one-shot NCCL comm init via ctypes UID exchange + XLA collective fallback), ep_dispatch / ep_combine as jax.custom_vjp functions with EP-mesh-aware residual handling.
• transformer_engine/jax/csrc/extensions/ep.cpp: XLA FFI C++ handlers using a shared_ptr<EpResources> weak-ref lifetime model; all five handlers marked FFI_CudaGraph_Traits for CUDA-graph compatibility.

Confidence Score: 5/5

The new JAX EP layer is safe to merge; no new blocking defects were found beyond what has already been flagged in earlier review rounds.

The core forward and backward logic, custom_vjp residual handling, and SPMD partition rules are all correct. The only new findings are non-blocking: a broad exception swallow in _allgather_uid that masks bootstrap failures, EpCombineBwdPrimitive.partition omitting the sharding validation every other EP partition function performs, and _dispatch_bwd capturing out_partition_spec from mesh context at trace time rather than threading it as a nondiff arg.

The backward path in transformer_engine/jax/ep.py (_dispatch_bwd) and EpCombineBwdPrimitive.partition in transformer_engine/jax/cpp_extensions/ep.py deserve a second look for sharding-spec consistency.

Important Files Changed

Filename	Overview
transformer_engine/jax/ep.py	New file: high-level ep_bootstrap, ep_dispatch (custom_vjp), ep_combine (custom_vjp) APIs; backward's out_partition_spec is trace-time-captured rather than threaded as a nondiff arg (inconsistency with ep_combine); broad exception swallow in _allgather_uid.
transformer_engine/jax/cpp_extensions/ep.py	New file: ~960-line primitive layer for EpPrepare/Dispatch/Combine (fwd+bwd); EpCombineBwdPrimitive.partition lacks input-sharding validation unlike every other EP partition function.
transformer_engine/jax/csrc/extensions/ep.cpp	New file: XLA FFI handlers for 5 EP operations; double-lock pattern in SetEpBootstrapParams (race window on re-init already flagged); topk_weights always wrapped as float32 (acknowledged by developer as intentional).
build_tools/jax.py	+3 lines: defaults NVTE_BUILD_WITH_NCCL_EP=1 (EP always on), but arch guard logic and hard-fail on missing submodule are inconsistent with setup.py's auto-disable behaviour (flagged in previous review cycles).
transformer_engine/jax/sharding.py	Adds ep_axis_size() helper and expands MeshResource.ep_resource docstring; straightforward, no issues.
tests/jax/test_multi_process_ep.py	New file: 742-line multi-process test suite with 13 cases covering bootstrap, primitive shapes, VJP correctness (closed-form), and HLO collective guard; good coverage.
transformer_engine/jax/csrc/extensions/pybind.cpp	+31 lines: registers EP FFI handlers and pybind11 bindings under NVTE_WITH_NCCL_EP guard; correct use of instantiate+execute pair for stateful FFI.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant User as User code
    participant Bootstrap as ep_bootstrap
    participant Allgather as _allgather_uid
    participant NCCL as ncclGetUniqueId (ctypes)
    participant TE_JAX as transformer_engine_jax (C++)
    participant Dispatch as ep_dispatch (custom_vjp)
    participant Combine as ep_combine (custom_vjp)
    participant FFI as XLA FFI (ep.cpp)

    User->>Bootstrap: ep_bootstrap(world_size, rank, ...)
    Bootstrap->>NCCL: ncclGetUniqueId (color-root only)
    Bootstrap->>Allgather: _allgather_uid(uid_arr, world_size)
    Allgather-->>Bootstrap: all_uids[world_size, 128]
    Bootstrap->>TE_JAX: set_ep_bootstrap_params(uid_bytes, ep_size, ...)
    TE_JAX->>FFI: AcquireEpResources - ncclCommInitRank + nvte_ep_initialize
    Bootstrap->>Bootstrap: set_ep_config(EpConfig(...))

    User->>Dispatch: ep_dispatch(cfg, topk_idx, tokens, topk_weights, recv_cap)
    Dispatch->>FFI: EpPrepareFFI (nvte_ep_prepare - handle_mem, token_counts)
    Dispatch->>FFI: EpDispatchFFI (nvte_ep_dispatch - recv_tokens, recv_topk_weights)
    Dispatch-->>User: (recv_tokens, recv_topk_weights, handle_mem, token_counts)

    User->>Combine: ep_combine(cfg, handle_mem, token_counts, expert_out, ...)
    Combine->>FFI: EpCombineFFI (nvte_ep_combine - result)
    Combine-->>User: result [T, H]

    Note over Dispatch,FFI: Backward (custom_vjp)
    Dispatch->>FFI: EpDispatchBwdFFI (nvte_ep_dispatch_bwd)
    Combine->>FFI: EpCombineBwdFFI (nvte_ep_combine_bwd)

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant User as User code
    participant Bootstrap as ep_bootstrap
    participant Allgather as _allgather_uid
    participant NCCL as ncclGetUniqueId (ctypes)
    participant TE_JAX as transformer_engine_jax (C++)
    participant Dispatch as ep_dispatch (custom_vjp)
    participant Combine as ep_combine (custom_vjp)
    participant FFI as XLA FFI (ep.cpp)

    User->>Bootstrap: ep_bootstrap(world_size, rank, ...)
    Bootstrap->>NCCL: ncclGetUniqueId (color-root only)
    Bootstrap->>Allgather: _allgather_uid(uid_arr, world_size)
    Allgather-->>Bootstrap: all_uids[world_size, 128]
    Bootstrap->>TE_JAX: set_ep_bootstrap_params(uid_bytes, ep_size, ...)
    TE_JAX->>FFI: AcquireEpResources - ncclCommInitRank + nvte_ep_initialize
    Bootstrap->>Bootstrap: set_ep_config(EpConfig(...))

    User->>Dispatch: ep_dispatch(cfg, topk_idx, tokens, topk_weights, recv_cap)
    Dispatch->>FFI: EpPrepareFFI (nvte_ep_prepare - handle_mem, token_counts)
    Dispatch->>FFI: EpDispatchFFI (nvte_ep_dispatch - recv_tokens, recv_topk_weights)
    Dispatch-->>User: (recv_tokens, recv_topk_weights, handle_mem, token_counts)

    User->>Combine: ep_combine(cfg, handle_mem, token_counts, expert_out, ...)
    Combine->>FFI: EpCombineFFI (nvte_ep_combine - result)
    Combine-->>User: result [T, H]

    Note over Dispatch,FFI: Backward (custom_vjp)
    Dispatch->>FFI: EpDispatchBwdFFI (nvte_ep_dispatch_bwd)
    Combine->>FFI: EpCombineBwdFFI (nvte_ep_combine_bwd)

Loading

_{Reviews (19): Last reviewed commit: "jax/ep: resolve test path from launcher ..." | Re-trigger Greptile}

[mgoldfarb-nvidia]

nit: can we return FFI InvalidArgument instead of a NVTE_CHECK for these inputs?

[phu0ngng]

This is probably a good idea. I suggest we make another follow-up MR to do so for all the FFIs.

[phu0ngng]

I would appreciate your help to review this PR @tdophung @jberchtold-nvidia!
Please focus on the changes in the JAX side, as the TE/Common ones will be discussed in #3034

[jberchtold-nvidia]

LGTM pending CI

[phu0ngng]

/te-ci JAX L1

[phu0ngng]

/te-ci JAX L1

[phu0ngng]

pipeline 54608589

[jberchtold-nvidia]

Since the TE core library and TE/JAX may be built separately, e.g. when pip installing from PyPi the TE core is pre-built but TE/JAX is a source dist and rebuilt. In that case, it's possible TE core is built with NVTE_WITH_NCCL_EP support but when installing TE/JAX the user doesn't have this flag set.

Is this intended? If not, would it be better to have TE core expose a nvte_is_nccl_ep_supported() and we call that here instead of using the NVTE_WITH_NCCL_EP macro?

[jberchtold-nvidia]

Can we add a description of each arg into the docstring? I think I could infer what "world_size" is but would be good to be explicit

[phu0ngng]

pipeline 54620197