[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 layer for Expert Parallelism: five XLA FFI handlers (EpPrepare, EpDispatch, EpCombine, and their backward variants) wrapping the nvte_ep_* C API, jax.custom_vjp-wrapped Python entry points with mesh-aware SPMD sharding rules, a new nccl_ep_enabled() build helper, a 13-test multi-process suite, and an end-to-end MoE example.

• Custom VJP wrappers (ep_dispatch, ep_combine) correctly thread handle_mem through residuals and re-pin cotangent sharding in the backward via with_sharding_constraint, avoiding XLA-transpose shard-dropping.
• XLA FFI handlers in ep.cpp perform on-stream int32→int64 topk_idx upcasting into a pre-allocated scratch buffer and cleanly manage NCCL comm lifetime via EpInstanceState shared-pointer anchors tied to executable lifetimes.
• nccl_ep_enabled() centralizes the Hopper+ arch gate that was previously duplicated between setup.py and build_tools/jax.py, fixing the previously flagged inconsistency between the two build paths.

Confidence Score: 4/5

Safe to merge with the caveat that direct callers of the public ep_combine_fwd primitive outside the ep_combine custom_vjp wrapper will silently receive partially-uninitialized output in SPMD when out_partition_spec is omitted.

Several previously-flagged issues (the assert on NCCL UID, the libnccl.so.2 hardcode, the stale-communicator re-init path) are still present but were already under discussion. A new finding: EpCombinePrimitive.partition's else branch, added to stop the crash the previous review found, substitutes silent data corruption by passing the global token count as the per-shard output size, causing the C++ combine op to fill only the local token slots and leave the rest of the buffer uninitialised before XLA replicates the result.

transformer_engine/jax/cpp_extensions/ep.py (EpCombinePrimitive.partition else-branch) and transformer_engine/jax/ep.py (assert-based NCCL UID check, hardcoded libnccl.so.2)

Important Files Changed

Filename	Overview
transformer_engine/jax/ep.py	New file: public ep_bootstrap, ep_dispatch (custom_vjp), ep_combine (custom_vjp). Validation guards and sharding helpers look correct; assert-based NCCL UID check and hardcoded libnccl.so.2 are still present (flagged in earlier reviews).
transformer_engine/jax/cpp_extensions/ep.py	New file: five JAX primitives with abstract-eval, lowering, impl, batcher, partition, and shardy rules. EpCombinePrimitive.partition's None-path silently produces wrong SPMD output (P1). Remaining primitives' sharding logic is consistent with global-vs-per-shard conventions.
transformer_engine/jax/csrc/extensions/ep.cpp	New file: five XLA FFI handlers wrapped around nvte_ep_* C API. Workspace sizing for int32 to int64 upcast is byte-exact. ncclCommInitRank holds g_ep_mu during construction, but the intended workflow serializes bootstrap before JIT compilation.
build_tools/utils.py	New nccl_ep_enabled() helper centralizes arch-gate logic previously duplicated in setup.py and jax.py. native arch treated as Hopper+ (P2 false-positive); otherwise logic matches the original setup.py behavior.
tests/jax/test_multi_process_ep.py	Comprehensive 13-test multi-process test suite covering bootstrap, prepare, dispatch/combine identity and custom_vjp fwd+bwd, and HLO collective guard. Tests use a shared class-level bootstrap to avoid repeated NCCL init.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant User as Python (ep.py)
    participant Prim as JAX Primitives
    participant FFI as XLA FFI (ep.cpp)
    participant TE as TE Common
    participant NCCL as NCCL EP

    User->>User: ep_bootstrap() allgather UID, set_ep_bootstrap_params
    FFI->>NCCL: ncclCommInitRank
    FFI->>TE: nvte_ep_initialize

    Note over User,NCCL: Forward pass
    User->>Prim: ep_dispatch(cfg, topk_idx, tokens, weights, cap)
    Prim->>FFI: EpPrepareHandler
    FFI->>TE: nvte_ep_prepare
    Prim->>FFI: EpDispatchHandler
    FFI->>NCCL: nvte_ep_dispatch
    Prim-->>User: recv_tokens, recv_weights, handle_mem, token_counts

    User->>User: expert MLP
    User->>Prim: ep_combine(cfg, handle_mem, token_counts, expert_out, T)
    Prim->>FFI: EpCombineHandler
    FFI->>NCCL: nvte_ep_combine
    Prim-->>User: result tokens

    Note over User,NCCL: Backward pass
    User->>Prim: _combine_bwd
    FFI->>NCCL: nvte_ep_combine_bwd
    User->>Prim: _dispatch_bwd
    FFI->>NCCL: nvte_ep_dispatch_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 Python (ep.py)
    participant Prim as JAX Primitives
    participant FFI as XLA FFI (ep.cpp)
    participant TE as TE Common
    participant NCCL as NCCL EP

    User->>User: ep_bootstrap() allgather UID, set_ep_bootstrap_params
    FFI->>NCCL: ncclCommInitRank
    FFI->>TE: nvte_ep_initialize

    Note over User,NCCL: Forward pass
    User->>Prim: ep_dispatch(cfg, topk_idx, tokens, weights, cap)
    Prim->>FFI: EpPrepareHandler
    FFI->>TE: nvte_ep_prepare
    Prim->>FFI: EpDispatchHandler
    FFI->>NCCL: nvte_ep_dispatch
    Prim-->>User: recv_tokens, recv_weights, handle_mem, token_counts

    User->>User: expert MLP
    User->>Prim: ep_combine(cfg, handle_mem, token_counts, expert_out, T)
    Prim->>FFI: EpCombineHandler
    FFI->>NCCL: nvte_ep_combine
    Prim-->>User: result tokens

    Note over User,NCCL: Backward pass
    User->>Prim: _combine_bwd
    FFI->>NCCL: nvte_ep_combine_bwd
    User->>Prim: _dispatch_bwd
    FFI->>NCCL: nvte_ep_dispatch_bwd

Loading

_{Reviews (24): Last reviewed commit: "Merge branch 'main' into phuong/ep-3-jax" | 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]

/te-ci JAX L1

[jberchtold-nvidia]

LGTM pending CI

[phu0ngng]

/te-ci JAX L1

[greptile-apps]

out_partition_spec=None in SPMD produces silent data corruption instead of a crash

The else branch (reached when ep_combine_fwd is called with the default out_partition_spec=None in a multi-device SPMD context) sets per_shard_leading = out_leading_shape (the global token count, e.g. (8,)) instead of the per-shard count (e.g. (2,) with ep=4), and assigns out_sharding = PartitionSpec() (fully replicated). When sharded_impl runs, it allocates an (8, H) output buffer on each device but nvte_ep_combine writes only the local 2-token result, leaving 6 rows uninitialized, and XLA then treats that partial buffer as the authoritative replicated global result, propagating garbage values silently.

The previous review flagged the crash that existed at this code site before the else branch was added and suggested raise ValueError("out_partition_spec must be specified in SPMD mode"). That guard was not applied, so the failure mode changed from an early loud error to a silent correctness failure. The ep_combine custom_vjp path is safe because _combine_fwd always calls _default_out_partition_spec(), but the public ep_combine_fwd primitive (exported in __all__) is still broken for direct SPMD callers.