[JAX] Expert Parallelism: JAX primitives + VJPs by phu0ngng · Pull Request #3036 · NVIDIA/TransformerEngine

phu0ngng · 2026-05-22T03:01:54Z

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 · 2026-05-22T03:10:53Z

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

%%{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

_{Reviews (24): Last reviewed commit: "Merge branch 'main' into phuong/ep-3-jax" | Re-trigger Greptile}

mgoldfarb-nvidia · 2026-05-22T15:17:01Z

+Error_Type EpPrepareFFI(cudaStream_t stream, Buffer_Type topk_idx, Result_Type token_counts,
+                        Result_Type handle_mem, Result_Type workspace, EpPrepareConfig config) {
+  auto topk_dims = topk_idx.dimensions();
+  NVTE_CHECK(topk_dims.size() >= 2,


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

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

phu0ngng · 2026-05-22T15:53:02Z

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

PR NVIDIA#3034 commit 9b225cb added a required NVTEEpGroupConfig.max_token_dtype field. The C++ backend (ep_backend.cpp:349) enforces typeToSize(tok_dtype) <= typeToSize(max_token_dtype) at every dispatch, and the field is also used at group create to size the NCCL EP staging buffers (ep_backend.cpp:221-222). PR NVIDIA#3036's JAX bootstrap (SetEpBootstrapParams / ep_bootstrap) was written before this field existed and never set it, so any JAX EP group landed with the zero-initialized default (kByte = 1 byte). Any bf16/fp16 dispatch from JAX then failed immediately with: tokens dtype (6) wider than group max_token_dtype (0) This commit threads max_token_dtype end-to-end: - transformer_engine/jax/csrc/extensions.h update SetEpBootstrapParams declaration to match the new arity. - transformer_engine/jax/csrc/extensions/ep.cpp add max_token_dtype to EpBootstrapParams and SetEpBootstrapParams; forward it into NVTEEpGroupConfig in the EpResources ctor. - transformer_engine/jax/csrc/extensions/pybind.cpp add the matching pybind11::arg("max_token_dtype") = 0. - transformer_engine/jax/ep.py add max_token_dtype kwarg to ep_bootstrap, convert numpy dtype to NVTEDType int, forward to the C++ setter. Carried on the te-ep-fixes branch until PR NVIDIA#3036 exposes the field upstream. See PR NVIDIA#3034 (commit 9b225cb, ep.h:43) for the field definition.

Reset 33 local commits onto phuong/ep-3-jax @ c34771d (her latest with EpConfig + EpLayerConfig API, NCCL bumped to 808d2433) and re-applied the three deltas uniquely ours: * transformer_engine/jax/moe.py: replaces upstream's multi-backend MoE block with our TE-EP-only single-custom-vjp rewrite. Adapted to her new API surface: tex.EpLayerConfig replaces tex.ep_make_handle (no more EpHandle pool/cache); 5 EP callsites rewired (cfg passed in place of handle, ep_prepare arg order swapped, top_k= dropped from ep_dispatch_bwd since it's now in cfg. * tests/jax/test_te_ep_moe.py: TE-EP MoE test (kept), with ep_bootstrap kwargs ep_size= and allow_handle_mem_reloc= dropped (no longer supported; ep_size is derived from mesh axes and the handle_mem reloc gating is gone). * tests/jax/run_te_ep_moe.sh: multi-process launcher (kept). Pre-sync state preserved at branch teddy/te_ep_integration.backup-pre-phuong-sync. EOF )

phu0ngng · 2026-06-11T00:02:56Z

/te-ci JAX L1

phu0ngng · 2026-06-24T07:23:14Z

/te-ci JAX L1

…ition methods Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

phu0ngng · 2026-06-25T13:53:01Z

/te-ci JAX L1

jberchtold-nvidia

LGTM pending CI

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

phu0ngng · 2026-06-26T12:13:47Z

/te-ci JAX L1

greptile-apps · 2026-06-27T09:13:52Z

+        mesh,
+        arg_infos,
+        result_infos,
+    ):
+        del result_infos
+        eo_spec = arg_infos[1].sharding.spec
+        if not _ep_spec_ok(eo_spec, trailing_count=2):
+            raise NotImplementedError(
+                "EpCombine: expert_out must be sharded as PartitionSpec(ep_resource,"
+                " None, None) (or ((dp, ep), None, None) when dp/fsdp is set)"
+                f" over [num_procs, recv_pr, H]; got spec={eo_spec}."
+            )
+        if out_partition_spec is not None:
+            per_shard_leading = _leading_per_shard(out_leading_shape, out_partition_spec[0], mesh)
+            out_sharding = NamedSharding(mesh, PartitionSpec(*out_partition_spec))
+        else:
+            per_shard_leading = out_leading_shape
+            out_sharding = NamedSharding(mesh, PartitionSpec())
+        arg_shardings = tuple(a.sharding for a in arg_infos)
+
+        def sharded_impl(handle_mem, expert_out):
+            return EpCombinePrimitive.impl(
+                handle_mem,
+                expert_out,
+                top_k,
+                dispatch_output_per_expert_alignment,
+                per_shard_leading,
+                out_partition_spec,
+            )
+
+        return mesh, sharded_impl, out_sharding, arg_shardings
+
+    @staticmethod
+    def shardy_sharding_rule(*args):
+        # Signature: (*static_args, mesh, value_types, result_types). Static args:
+        # (top_k, dispatch_alignment, out_leading_shape, out_partition_spec).
+        result_types = args[-1]


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.

* Expert Parallelism: JAX primitives + VJPs --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

Reset 33 local commits onto phuong/ep-3-jax @ c34771d (her latest with EpConfig + EpLayerConfig API, NCCL bumped to 808d2433) and re-applied the three deltas uniquely ours: * transformer_engine/jax/moe.py: replaces upstream's multi-backend MoE block with our TE-EP-only single-custom-vjp rewrite. Adapted to her new API surface: tex.EpLayerConfig replaces tex.ep_make_handle (no more EpHandle pool/cache); 5 EP callsites rewired (cfg passed in place of handle, ep_prepare arg order swapped, top_k= dropped from ep_dispatch_bwd since it's now in cfg. * tests/jax/test_te_ep_moe.py: TE-EP MoE test (kept), with ep_bootstrap kwargs ep_size= and allow_handle_mem_reloc= dropped (no longer supported; ep_size is derived from mesh axes and the handle_mem reloc gating is gone). * tests/jax/run_te_ep_moe.sh: multi-process launcher (kept). Pre-sync state preserved at branch teddy/te_ep_integration.backup-pre-phuong-sync. EOF ) Signed-off-by: tdophung <tdophung@nvidia.com>

After the upstream PR NVIDIA#3036 resync the moe() API surface lost PermutationBackend (TE-EP is the only backend now), gate_inside_vjp (always True), and the per-call quantizer_sets knob (quantization flows through the standard TE autocast / with_quantizer_set context). It also gained apply_topk_weights_early and renamed the wrapper's private _align_size to the public align_size the test suite already uses. The Flax _MoEBlock wrapper was still passing the old kwargs, which broke every test that touched the wrapper. Wrapper changes: * drop "from ..moe import PermutationBackend" plus the dataclass field, the isinstance(..., PermutationBackend) validation in __post_init__, and the pass-through to moe(). * drop "from ..quantize import noop_quantizer_set" and the quantizer_sets=(noop, noop, noop) pass-through. * drop gate_inside_vjp=True. * rename _align_size: int = 0 -> align_size: int = 0 (matches what tests/jax/test_te_ep_moe.py already passes). * add apply_topk_weights_early: bool = False and pass it through to moe(). * refresh class docstring: drop permutation_backend / _align_size / quantizer_sets descriptions, add apply_topk_weights_early / align_size, note that quantization currently flows only through fp8_autocast. Signed-off-by: tdophung <tdophung@nvidia.com>

…and inline justifications) Responds to jberchtold-nvidia's PR NVIDIA#3116 review threads on ``transformer_engine/jax/moe.py``. All changes are confined to a single file because each review thread targets a localized region and splitting mid-file would risk reordering bugs. Per review thread: 1. "Why do we need _with_sharding_constraint_cast_bwd? I haven't seen something like this required for our other VJPs." -- Expand the helper's docstring to spell out exactly why MoE needs it: unlike LN+MLP, the MoE bwd composes a bf16 cotangent from ep_dispatch_bwd with an fp32 cotangent from fused_topk_with_score_function_bwd (which the fwd's logits_2d -> fp32 promotion forces). Without the cast, ``d_x`` surfaces at fp32 even when ``x`` is bf16, doubling activation grad bandwidth and breaking any downstream LN bwd that pins a bf16 layout. (Review thread "Why do we need this utility function?".) 2. "Why is this dtype casting required? I don't recall us needing it for the non-MoE LNMLP block." -- Expand the comment above the bwd activation fp32 promotion to explain the MoE-specific math: LN+MLP's silu sits behind a downstream LN that absorbs the bf16 rounding error, while MoE's silu sits on the *expert* side of routing -- the bf16 rounding rides directly into expert_outputs and is summed across topk experts by ep_combine. Bf16 silu alone drifts ~1% vs fp32 silu and compounds through wo->combine into the ~1.4% per-element parity gap we measured against the pure-JAX softmax reference. Mirroring the fwd's fp32 promotion in the bwd keeps silu' in lock-step with silu. (Review thread on "# Activation bwd. Mirror the fwd's fp32 promotion of silu+multiply".) 3. "Do we have a use-case for user-specified alignments beyond 128 currently? ... it'd make sense to instead hardcode _ALIGN_SIZE = 128 as a constant at the top of the file for now to simplify this MoEBlock API. We can always expand the API to support a user-specified align size in the future." -- Implement the suggestion. Drop ``align_size`` from ``_moe_fwd_rule`` / ``_moe_bwd_rule`` / ``_moe`` / public ``moe()``; shift the ``custom_vjp`` ``nondiff_argnums`` from ``range(9, 27)`` -> ``range(9, 26)``; replace ``effective_align = max(int(align_size), 128)`` with the new module-level ``_ALIGN_SIZE = 128`` constant. Trim the ``moe()`` docstring accordingly. (Review thread on "natural_spe = num_ep * max_tokens_per_rank".) 4. "Which axis name inputs are physical mesh axes and why can be logical axes? ... No need to make any changes for now, I just want to assess which are which and then we can discuss if it makes sense to support logical on some/all or if some are required to be physical axes." -- Add an "Axis-name parameters" section to ``moe()``'s docstring listing which kwargs are physical mesh axes (``ep_axis``, ``data_parallelism_axes`` -- they index ``Mesh.shape`` directly to compute ``num_ep`` / ``dp_size`` and to construct the ``P((dp..., ep), None, None)`` for ``jax.lax.with_sharding_constraint``) vs logical axes (``input_axes``, ``gate_kernel_axes``, ``wi_kernel_axes``, ``wo_kernel_axes`` -- resolved via the Flax logical-axis rules). Also document why ``ep_axis`` / ``data_parallelism_axes`` are intentionally non-logical: the EP comm-group construction (``dp_color = rank // ep_size``) and the bootstrap signature check both require concrete integer sizes. (Review thread on "batch_pspec_axis = (*data_parallelism_axes, ep_axis)".) 5. "Is this NaN filtering a debugging artifact or something we need in the final version?" -- Strengthen the inline comment above ``sparse_probs = jnp.where(jnp.isnan(sparse_probs), 0, ...)`` to explicitly call this out as a CORRECTNESS REQUIREMENT, not a debugging artifact: it covers the sigmoid+K>1 underflow path where top-K sigmoid scores all round to zero and the ``weights / (weights.sum + 1e-20)`` normalisation emits NaN. Observationally the filter is a no-op on the dense unit-test distributions, but it must stay in for sparse / production routing. (Review thread on "sparse_probs = jnp.where(jnp.isnan(sparse_probs), ...).") Not addressed in this commit (intentional): * Review thread on the ``align_size: int = 0`` placeholder in ``flax/moe.py`` ("Placeholder comment for me to fix this so align_size is inferred automatically based on the recipe and doesn't need to be specified by the user"). That's jberchtold's own follow-up. * Review thread on the explicit ``tree_flatten`` / ``tree_unflatten`` on ``_Ctx`` ("better to use the ``@flax_struct.dataclass``"). Deferred to a separate, testable commit because changing a ``custom_vjp`` residual's pytree registration touches subtle ordering / None-handling semantics that warrant their own bisect surface. * Review thread on ``use_bias`` / ``use_expert_bias`` renames -- handled in the immediately preceding commit ``jax/flax,tests: rename use_bias/use_expert_bias for symmetry``. * Review thread on the ``expert_bias`` fp32 init -- already resolved during the Phuong PR NVIDIA#3036 resync (the redundant ``jnp.float32`` second-dtype argument on ``self.param`` was dropped; ``expert_bias`` now lives at ``self.dtype``). Signed-off-by: tdophung <tdophung@nvidia.com>

…successor test_moe_vjp.py and test_multiprocess_moe_vjp.py both import PermutationBackend from transformer_engine.jax.moe -- an API that was removed during the Phuong PR NVIDIA#3036 resync. Both files have been dead-on-import ever since; the multiprocess launcher run_multiprocess_moe_vjp.sh only points at the dead test. test_te_ep_moe.py (the TE-EP-only custom_vjp suite) already covers everything the legacy files exercised that is still meaningful: fwd, bwd parity vs the pure-JAX reference, aux loss, both score functions, multi-process. The legacy parametrize axis (PermutationBackend.PURE_JAX vs TRITON) no longer exists. * Delete tests/jax/test_moe_vjp.py * Delete tests/jax/test_multiprocess_moe_vjp.py * Delete tests/jax/run_multiprocess_moe_vjp.sh * qa/L0_jax_distributed_unittest/test.sh: switch the MoE VJP distributed suite invocation from run_multiprocess_moe_vjp.sh / test_multiprocess_moe_vjp.py to run_te_ep_moe.sh / test_te_ep_moe.py. * tests/jax/conftest.py: docstring reference updated. * tests/jax/test_te_ep_moe.py: drop stale "successor to ..." aside and the "mirroring run_multiprocess_moe_vjp.sh" parenthetical. Net: -981 / +9. Signed-off-by: tdophung <tdophung@nvidia.com>

…erplate Per reviewer feedback (Jaberchtold on PR NVIDIA#3036): the manual tree_flatten / tree_unflatten on _Ctx duplicate exactly what @flax.struct.dataclass auto-generates, and the permutation dataclasses elsewhere in this module already use flax.struct. Switching to @flax.struct.dataclass: * Removes ~75 lines of mechanical tree_flatten / tree_unflatten that have to be kept in sync with the field list by hand. * Keeps cfg as the single static field via flax.struct.field(pytree_node=False), so the fwd -> bwd boundary behavior under jax.custom_vjp is unchanged. * Drops two now-unused imports (dataclasses.dataclass, jax.tree_util.register_pytree_node_class) and adds flax.struct. Field order and the (children, aux_data) split are byte-equivalent to the previous manual implementation, so the pytree treedef seen by jax.custom_vjp is identical. Signed-off-by: tdophung <tdophung@nvidia.com>

Reset 33 local commits onto phuong/ep-3-jax @ c34771d (her latest with EpConfig + EpLayerConfig API, NCCL bumped to 808d2433) and re-applied the three deltas uniquely ours: * transformer_engine/jax/moe.py: replaces upstream's multi-backend MoE block with our TE-EP-only single-custom-vjp rewrite. Adapted to her new API surface: tex.EpLayerConfig replaces tex.ep_make_handle (no more EpHandle pool/cache); 5 EP callsites rewired (cfg passed in place of handle, ep_prepare arg order swapped, top_k= dropped from ep_dispatch_bwd since it's now in cfg. * tests/jax/test_te_ep_moe.py: TE-EP MoE test (kept), with ep_bootstrap kwargs ep_size= and allow_handle_mem_reloc= dropped (no longer supported; ep_size is derived from mesh axes and the handle_mem reloc gating is gone). * tests/jax/run_te_ep_moe.sh: multi-process launcher (kept). Pre-sync state preserved at branch teddy/te_ep_integration.backup-pre-phuong-sync. EOF ) Signed-off-by: Teddy Do <tdophung@nvidia.com>

After the upstream PR NVIDIA#3036 resync the moe() API surface lost PermutationBackend (TE-EP is the only backend now), gate_inside_vjp (always True), and the per-call quantizer_sets knob (quantization flows through the standard TE autocast / with_quantizer_set context). It also gained apply_topk_weights_early and renamed the wrapper's private _align_size to the public align_size the test suite already uses. The Flax _MoEBlock wrapper was still passing the old kwargs, which broke every test that touched the wrapper. Wrapper changes: * drop "from ..moe import PermutationBackend" plus the dataclass field, the isinstance(..., PermutationBackend) validation in __post_init__, and the pass-through to moe(). * drop "from ..quantize import noop_quantizer_set" and the quantizer_sets=(noop, noop, noop) pass-through. * drop gate_inside_vjp=True. * rename _align_size: int = 0 -> align_size: int = 0 (matches what tests/jax/test_te_ep_moe.py already passes). * add apply_topk_weights_early: bool = False and pass it through to moe(). * refresh class docstring: drop permutation_backend / _align_size / quantizer_sets descriptions, add apply_topk_weights_early / align_size, note that quantization currently flows only through fp8_autocast. Signed-off-by: Teddy Do <tdophung@nvidia.com>

…and inline justifications) Responds to jberchtold-nvidia's PR NVIDIA#3116 review threads on ``transformer_engine/jax/moe.py``. All changes are confined to a single file because each review thread targets a localized region and splitting mid-file would risk reordering bugs. Per review thread: 1. "Why do we need _with_sharding_constraint_cast_bwd? I haven't seen something like this required for our other VJPs." -- Expand the helper's docstring to spell out exactly why MoE needs it: unlike LN+MLP, the MoE bwd composes a bf16 cotangent from ep_dispatch_bwd with an fp32 cotangent from fused_topk_with_score_function_bwd (which the fwd's logits_2d -> fp32 promotion forces). Without the cast, ``d_x`` surfaces at fp32 even when ``x`` is bf16, doubling activation grad bandwidth and breaking any downstream LN bwd that pins a bf16 layout. (Review thread "Why do we need this utility function?".) 2. "Why is this dtype casting required? I don't recall us needing it for the non-MoE LNMLP block." -- Expand the comment above the bwd activation fp32 promotion to explain the MoE-specific math: LN+MLP's silu sits behind a downstream LN that absorbs the bf16 rounding error, while MoE's silu sits on the *expert* side of routing -- the bf16 rounding rides directly into expert_outputs and is summed across topk experts by ep_combine. Bf16 silu alone drifts ~1% vs fp32 silu and compounds through wo->combine into the ~1.4% per-element parity gap we measured against the pure-JAX softmax reference. Mirroring the fwd's fp32 promotion in the bwd keeps silu' in lock-step with silu. (Review thread on "# Activation bwd. Mirror the fwd's fp32 promotion of silu+multiply".) 3. "Do we have a use-case for user-specified alignments beyond 128 currently? ... it'd make sense to instead hardcode _ALIGN_SIZE = 128 as a constant at the top of the file for now to simplify this MoEBlock API. We can always expand the API to support a user-specified align size in the future." -- Implement the suggestion. Drop ``align_size`` from ``_moe_fwd_rule`` / ``_moe_bwd_rule`` / ``_moe`` / public ``moe()``; shift the ``custom_vjp`` ``nondiff_argnums`` from ``range(9, 27)`` -> ``range(9, 26)``; replace ``effective_align = max(int(align_size), 128)`` with the new module-level ``_ALIGN_SIZE = 128`` constant. Trim the ``moe()`` docstring accordingly. (Review thread on "natural_spe = num_ep * max_tokens_per_rank".) 4. "Which axis name inputs are physical mesh axes and why can be logical axes? ... No need to make any changes for now, I just want to assess which are which and then we can discuss if it makes sense to support logical on some/all or if some are required to be physical axes." -- Add an "Axis-name parameters" section to ``moe()``'s docstring listing which kwargs are physical mesh axes (``ep_axis``, ``data_parallelism_axes`` -- they index ``Mesh.shape`` directly to compute ``num_ep`` / ``dp_size`` and to construct the ``P((dp..., ep), None, None)`` for ``jax.lax.with_sharding_constraint``) vs logical axes (``input_axes``, ``gate_kernel_axes``, ``wi_kernel_axes``, ``wo_kernel_axes`` -- resolved via the Flax logical-axis rules). Also document why ``ep_axis`` / ``data_parallelism_axes`` are intentionally non-logical: the EP comm-group construction (``dp_color = rank // ep_size``) and the bootstrap signature check both require concrete integer sizes. (Review thread on "batch_pspec_axis = (*data_parallelism_axes, ep_axis)".) 5. "Is this NaN filtering a debugging artifact or something we need in the final version?" -- Strengthen the inline comment above ``sparse_probs = jnp.where(jnp.isnan(sparse_probs), 0, ...)`` to explicitly call this out as a CORRECTNESS REQUIREMENT, not a debugging artifact: it covers the sigmoid+K>1 underflow path where top-K sigmoid scores all round to zero and the ``weights / (weights.sum + 1e-20)`` normalisation emits NaN. Observationally the filter is a no-op on the dense unit-test distributions, but it must stay in for sparse / production routing. (Review thread on "sparse_probs = jnp.where(jnp.isnan(sparse_probs), ...).") Not addressed in this commit (intentional): * Review thread on the ``align_size: int = 0`` placeholder in ``flax/moe.py`` ("Placeholder comment for me to fix this so align_size is inferred automatically based on the recipe and doesn't need to be specified by the user"). That's jberchtold's own follow-up. * Review thread on the explicit ``tree_flatten`` / ``tree_unflatten`` on ``_Ctx`` ("better to use the ``@flax_struct.dataclass``"). Deferred to a separate, testable commit because changing a ``custom_vjp`` residual's pytree registration touches subtle ordering / None-handling semantics that warrant their own bisect surface. * Review thread on ``use_bias`` / ``use_expert_bias`` renames -- handled in the immediately preceding commit ``jax/flax,tests: rename use_bias/use_expert_bias for symmetry``. * Review thread on the ``expert_bias`` fp32 init -- already resolved during the Phuong PR NVIDIA#3036 resync (the redundant ``jnp.float32`` second-dtype argument on ``self.param`` was dropped; ``expert_bias`` now lives at ``self.dtype``). Signed-off-by: Teddy Do <tdophung@nvidia.com>

…successor test_moe_vjp.py and test_multiprocess_moe_vjp.py both import PermutationBackend from transformer_engine.jax.moe -- an API that was removed during the Phuong PR NVIDIA#3036 resync. Both files have been dead-on-import ever since; the multiprocess launcher run_multiprocess_moe_vjp.sh only points at the dead test. test_te_ep_moe.py (the TE-EP-only custom_vjp suite) already covers everything the legacy files exercised that is still meaningful: fwd, bwd parity vs the pure-JAX reference, aux loss, both score functions, multi-process. The legacy parametrize axis (PermutationBackend.PURE_JAX vs TRITON) no longer exists. * Delete tests/jax/test_moe_vjp.py * Delete tests/jax/test_multiprocess_moe_vjp.py * Delete tests/jax/run_multiprocess_moe_vjp.sh * qa/L0_jax_distributed_unittest/test.sh: switch the MoE VJP distributed suite invocation from run_multiprocess_moe_vjp.sh / test_multiprocess_moe_vjp.py to run_te_ep_moe.sh / test_te_ep_moe.py. * tests/jax/conftest.py: docstring reference updated. * tests/jax/test_te_ep_moe.py: drop stale "successor to ..." aside and the "mirroring run_multiprocess_moe_vjp.sh" parenthetical. Net: -981 / +9. Signed-off-by: Teddy Do <tdophung@nvidia.com>

…erplate Per reviewer feedback (Jaberchtold on PR NVIDIA#3036): the manual tree_flatten / tree_unflatten on _Ctx duplicate exactly what @flax.struct.dataclass auto-generates, and the permutation dataclasses elsewhere in this module already use flax.struct. Switching to @flax.struct.dataclass: * Removes ~75 lines of mechanical tree_flatten / tree_unflatten that have to be kept in sync with the field list by hand. * Keeps cfg as the single static field via flax.struct.field(pytree_node=False), so the fwd -> bwd boundary behavior under jax.custom_vjp is unchanged. * Drops two now-unused imports (dataclasses.dataclass, jax.tree_util.register_pytree_node_class) and adds flax.struct. Field order and the (children, aux_data) split are byte-equivalent to the previous manual implementation, so the pytree treedef seen by jax.custom_vjp is identical. Signed-off-by: Teddy Do <tdophung@nvidia.com>

* [JAX] Resync onto upstream PR #3036, restore TE-EP-only MoE block Reset 33 local commits onto phuong/ep-3-jax @ c34771d (her latest with EpConfig + EpLayerConfig API, NCCL bumped to 808d2433) and re-applied the three deltas uniquely ours: * transformer_engine/jax/moe.py: replaces upstream's multi-backend MoE block with our TE-EP-only single-custom-vjp rewrite. Adapted to her new API surface: tex.EpLayerConfig replaces tex.ep_make_handle (no more EpHandle pool/cache); 5 EP callsites rewired (cfg passed in place of handle, ep_prepare arg order swapped, top_k= dropped from ep_dispatch_bwd since it's now in cfg. * tests/jax/test_te_ep_moe.py: TE-EP MoE test (kept), with ep_bootstrap kwargs ep_size= and allow_handle_mem_reloc= dropped (no longer supported; ep_size is derived from mesh axes and the handle_mem reloc gating is gone). * tests/jax/run_te_ep_moe.sh: multi-process launcher (kept). Pre-sync state preserved at branch teddy/te_ep_integration.backup-pre-phuong-sync. EOF ) Signed-off-by: Teddy Do <tdophung@nvidia.com> * tests/jax: trim TE-EP MoE suite (drop bootstrap, flax-wrapper, bias-zero) * drop ``TestZZZTeEpMoeBootstrap``: the re-bootstrap mismatch is a one-line guard in ``ep_bootstrap`` and not the MoE block's concern; exercising it from this suite also taints the per-process NCCL bootstrap cache for the rest of the file with no real upside. * drop ``TestTeEpMoEBlockFlax::test_init_apply_parity``: every config in ``_CONFIGS`` already runs ``MoEBlock`` (the Flax wrapper) end-to-end via ``test_forward`` / ``test_backward``, so this was a duplicate of ``softmax`` parity in another wrapper -- leave wrapper refactors to devs without paying for an extra CI run each time. * drop ``sigmoid-bias-zero``: with a zero-init bias buffer the routing math collapses to the no-bias case, so ``sigmoid`` already covers that numerical path. The bias-aware codepath is still exercised by ``sigmoid-bias-strong`` (non-zero bias). * refresh the module-level docstring to list intentional non-coverage so future readers don't re-add these tests. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/router: fix two bwd custom_partitioning bugs (aux-loss rank, topk closure) Two unrelated one-line bugs in the bwd custom_partitioning machinery that only surface once the MoE block's aux-loss path is lifted out of shard_map (the custom_partitioning_sharding_rule check is skipped under shard_map, which is why these never tripped before). 1. FusedMoEAuxLossBwdPrimitive.shardy_sharding_rule: ``grad_aux_loss`` is the cotangent of a scalar loss and is rank-0; declaring it with a spurious ``grad_one`` factor gave it rank-1 and tripped JAX's custom_partitioning_sharding_rule rank check at global view. Change the rule's third operand entry to empty: "const_buf_one, num_experts, grad_one -> i num_experts" -> "const_buf_one, num_experts, -> i num_experts" 2. FusedTopkWithScoreFunctionBwdPrimitive.partition: ``del result_infos, routing_map_format`` removed ``routing_map_format`` from the enclosing scope before the nested ``sharded_impl`` closure was invoked. Python closures resolve names at call time, not definition time, so when XLA finally invoked ``sharded_impl`` for the bwd partitioned impl it raised ``NameError: cannot access free variable 'routing_map_format'``. Drop ``routing_map_format`` from the ``del`` and leave a NOTE so future cleanups don't reintroduce the bug. Sibling partition methods (fwd topk, both aux-loss directions) already only ``del result_infos`` and need no change. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/ep: skip size-1 dp/fsdp axis in _ep_outer_axis A dp_resource or fsdp_resource that exists in the active mesh resource config but is sized 1 in the actual mesh would still be returned by ``_ep_outer_axis()``, pinning EP-output PartitionSpecs to a degenerate axis. JAX collapses size-1 mesh axes during lowering, which made the EP-output specs reference an axis that no longer exists at runtime -- breaking shard_map output stitching on configs where DP or FSDP is optional. Treat a size-1 axis as absent: prefer dp -> fsdp, but only when the candidate axis is actually sized > 1 in the current mesh. Falls back to the previous behaviour when no axis is configured at all. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/flax: realign _MoEBlock with post-resync moe() signature After the upstream PR #3036 resync the moe() API surface lost PermutationBackend (TE-EP is the only backend now), gate_inside_vjp (always True), and the per-call quantizer_sets knob (quantization flows through the standard TE autocast / with_quantizer_set context). It also gained apply_topk_weights_early and renamed the wrapper's private _align_size to the public align_size the test suite already uses. The Flax _MoEBlock wrapper was still passing the old kwargs, which broke every test that touched the wrapper. Wrapper changes: * drop "from ..moe import PermutationBackend" plus the dataclass field, the isinstance(..., PermutationBackend) validation in __post_init__, and the pass-through to moe(). * drop "from ..quantize import noop_quantizer_set" and the quantizer_sets=(noop, noop, noop) pass-through. * drop gate_inside_vjp=True. * rename _align_size: int = 0 -> align_size: int = 0 (matches what tests/jax/test_te_ep_moe.py already passes). * add apply_topk_weights_early: bool = False and pass it through to moe(). * refresh class docstring: drop permutation_backend / _align_size / quantizer_sets descriptions, add apply_topk_weights_early / align_size, note that quantization currently flows only through fp8_autocast. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: plumb token_counts to grouped_gemm and zero 0-token wgrad slices Two correctness fixes for the TE-EP MoE custom_vjp that together let the bwd parity tests pass on 0-token-globally experts, and drop a workaround that is no longer needed. (1) Plumb per-expert padded token_counts into grouped_gemm group_sizes. NCCL EP HT dispatch lays out recv_tokens expert-major as [expert_0_padded | expert_1_padded | ... | overalloc_tail] where each per-expert block already includes the dispatch_output_per_expert_alignment zero-padding and only the trailing overalloc tail (slack between sum(token_counts) and the worst-case recv_pr) is unused. Previously _ffn_fwd_per_shard built a static local_group_sizes = jnp.full((num_local_experts,), slots_per_expert), which over-counted by the overalloc tail and forced cuBLAS to run the GEMM for every group including 0-token-routed experts. Pipe the real per-shard token_counts (1, num_local_experts) from ep_prepare through _moe_fwd_rule (added to ffn_in_specs/ffn_in_args with ep2_spec), into _ffn_fwd_per_shard as token_counts_local, and reshape into local_group_sizes for both grouped_quantize and grouped_gemm. cuBLAS now skips both 0-token experts and the trailing overalloc tail. Mirror the residual spec change on the bwd (local_group_sizes residual moves from P() to ep2_spec). (2) Per-group jnp.where zero-fill on wgrad outputs. cuBLAS grouped_gemm skips groups with size_g == 0 without zero-filling the corresponding out[g, :, :] slice (cublaslt_grouped_gemm.cu lines 2086/2096). For a shard hosting an expert that received zero tokens globally, d_wo / d_wi_combined for that expert is left uninit, which propagates NaN straight into the user's optimizer state. Add wgrad_group_active = (local_group_sizes > 0)[:, None, None] in _ffn_bwd_per_shard and apply via jnp.where on d_wo (right after the wo wgrad) and d_wi_combined (right after the fused wi_0+wi_1 wgrad). Mask shape is (num_local_experts, 1, 1) so cost is negligible. (3) Drop the lax.cond zero-init guard on r_tok in _moe_fwd_rule._body. Previously a jax.lax.cond(jnp.any(r_w != 0), identity, zeros_like) wrapper around recv_tokens worked around tex.ep_dispatch_fwd leaving the recv buffer uninit on fully-empty-receiver ranks. With (1) in place, cuBLAS skips experts whose group_sizes == 0 and the per-row trailing tail of dispatched recv_tokens is unread by every downstream consumer (subsequent grouped_gemms read only sum(group_sizes) rows; ep_combine and ep_dispatch_bwd are handle_mem-aware). The only per-row consumer that would propagate the tail is grouped_dbias (per-row segment_sum), which only runs when has_bias=True, and that FFN bias path is currently gated upstream (cuBLAS grouped_gemm has no fused bias on Hopper yet; PR 3083 adds the pure-JAX bias add). With (2) handling the user-visible wgrad-NaN risk on 0-token experts, the lax.cond is now redundant. Replace with a NOTE pointing at the two follow-ups that would force its reintroduction: - a future caller that reads the full recv tile non-group-aware (e.g. an inspect probe), or - the FFN bias path landing, which would resurrect grouped_dbias. Also rewrite the _ffn_fwd_per_shard and _ffn_bwd_per_shard docstrings to spell out the per-row vs per-group uninit semantics so the next person debugging a NaN here has the invariants written down. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/flax,tests: rename use_bias/use_expert_bias for symmetry (PR #3116) Address jberchtold-nvidia's PR #3116 nit "rename use_bias -> use_ffn_bias and use_expert_bias -> use_expert_routing_bias". The two flags are siblings (they enable two different bias buffers) but the old names suggested ``use_bias`` was the general fallback, which wasn't the intent. The new names make the FFN-vs-routing distinction obvious from the call site. * transformer_engine/jax/flax/moe.py use_bias -> use_ffn_bias (dataclass field + branch in __call__ + docstring entry) use_expert_bias -> use_expert_routing_bias (same) * tests/jax/test_te_ep_moe.py _make_block(use_expert_bias=...) -> use_expert_routing_bias sigmoid-bias-strong config key updated _reference_kwargs_from_config now reads use_expert_routing_bias ``_MoEBlock`` is still the experimental underscore-prefixed alias (no public ``MoEBlock`` export yet), so the rename is API-safe. The pre-resync legacy tests (``test_moe_vjp.py``, ``test_multiprocess_moe_vjp.py``) are intentionally not updated -- they already reference removed APIs like ``PermutationBackend`` and need a separate post-resync cleanup pass. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: address PR #3116 review feedback (hardcode align + expand inline justifications) Responds to jberchtold-nvidia's PR #3116 review threads on ``transformer_engine/jax/moe.py``. All changes are confined to a single file because each review thread targets a localized region and splitting mid-file would risk reordering bugs. Per review thread: 1. "Why do we need _with_sharding_constraint_cast_bwd? I haven't seen something like this required for our other VJPs." -- Expand the helper's docstring to spell out exactly why MoE needs it: unlike LN+MLP, the MoE bwd composes a bf16 cotangent from ep_dispatch_bwd with an fp32 cotangent from fused_topk_with_score_function_bwd (which the fwd's logits_2d -> fp32 promotion forces). Without the cast, ``d_x`` surfaces at fp32 even when ``x`` is bf16, doubling activation grad bandwidth and breaking any downstream LN bwd that pins a bf16 layout. (Review thread "Why do we need this utility function?".) 2. "Why is this dtype casting required? I don't recall us needing it for the non-MoE LNMLP block." -- Expand the comment above the bwd activation fp32 promotion to explain the MoE-specific math: LN+MLP's silu sits behind a downstream LN that absorbs the bf16 rounding error, while MoE's silu sits on the *expert* side of routing -- the bf16 rounding rides directly into expert_outputs and is summed across topk experts by ep_combine. Bf16 silu alone drifts ~1% vs fp32 silu and compounds through wo->combine into the ~1.4% per-element parity gap we measured against the pure-JAX softmax reference. Mirroring the fwd's fp32 promotion in the bwd keeps silu' in lock-step with silu. (Review thread on "# Activation bwd. Mirror the fwd's fp32 promotion of silu+multiply".) 3. "Do we have a use-case for user-specified alignments beyond 128 currently? ... it'd make sense to instead hardcode _ALIGN_SIZE = 128 as a constant at the top of the file for now to simplify this MoEBlock API. We can always expand the API to support a user-specified align size in the future." -- Implement the suggestion. Drop ``align_size`` from ``_moe_fwd_rule`` / ``_moe_bwd_rule`` / ``_moe`` / public ``moe()``; shift the ``custom_vjp`` ``nondiff_argnums`` from ``range(9, 27)`` -> ``range(9, 26)``; replace ``effective_align = max(int(align_size), 128)`` with the new module-level ``_ALIGN_SIZE = 128`` constant. Trim the ``moe()`` docstring accordingly. (Review thread on "natural_spe = num_ep * max_tokens_per_rank".) 4. "Which axis name inputs are physical mesh axes and why can be logical axes? ... No need to make any changes for now, I just want to assess which are which and then we can discuss if it makes sense to support logical on some/all or if some are required to be physical axes." -- Add an "Axis-name parameters" section to ``moe()``'s docstring listing which kwargs are physical mesh axes (``ep_axis``, ``data_parallelism_axes`` -- they index ``Mesh.shape`` directly to compute ``num_ep`` / ``dp_size`` and to construct the ``P((dp..., ep), None, None)`` for ``jax.lax.with_sharding_constraint``) vs logical axes (``input_axes``, ``gate_kernel_axes``, ``wi_kernel_axes``, ``wo_kernel_axes`` -- resolved via the Flax logical-axis rules). Also document why ``ep_axis`` / ``data_parallelism_axes`` are intentionally non-logical: the EP comm-group construction (``dp_color = rank // ep_size``) and the bootstrap signature check both require concrete integer sizes. (Review thread on "batch_pspec_axis = (*data_parallelism_axes, ep_axis)".) 5. "Is this NaN filtering a debugging artifact or something we need in the final version?" -- Strengthen the inline comment above ``sparse_probs = jnp.where(jnp.isnan(sparse_probs), 0, ...)`` to explicitly call this out as a CORRECTNESS REQUIREMENT, not a debugging artifact: it covers the sigmoid+K>1 underflow path where top-K sigmoid scores all round to zero and the ``weights / (weights.sum + 1e-20)`` normalisation emits NaN. Observationally the filter is a no-op on the dense unit-test distributions, but it must stay in for sparse / production routing. (Review thread on "sparse_probs = jnp.where(jnp.isnan(sparse_probs), ...).") Not addressed in this commit (intentional): * Review thread on the ``align_size: int = 0`` placeholder in ``flax/moe.py`` ("Placeholder comment for me to fix this so align_size is inferred automatically based on the recipe and doesn't need to be specified by the user"). That's jberchtold's own follow-up. * Review thread on the explicit ``tree_flatten`` / ``tree_unflatten`` on ``_Ctx`` ("better to use the ``@flax_struct.dataclass``"). Deferred to a separate, testable commit because changing a ``custom_vjp`` residual's pytree registration touches subtle ordering / None-handling semantics that warrant their own bisect surface. * Review thread on ``use_bias`` / ``use_expert_bias`` renames -- handled in the immediately preceding commit ``jax/flax,tests: rename use_bias/use_expert_bias for symmetry``. * Review thread on the ``expert_bias`` fp32 init -- already resolved during the Phuong PR #3036 resync (the redundant ``jnp.float32`` second-dtype argument on ``self.param`` was dropped; ``expert_bias`` now lives at ``self.dtype``). Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: strip PR-response framing from comments; drop sparse_probs NaN sanitizer * Rewrite the inline justifications added in 078a7d80 so each one reads as standalone code documentation, not as a reply to a reviewer: drop "per PR #3116 review", "review feedback", "Renamed from ... per PR ..." and similar PR/thread references from moe.py, flax/moe.py, and tests/jax/test_te_ep_moe.py. Technical content (why the fp32 promotion is needed for the MoE silu+multiply, why _with_sharding_constraint_cast_bwd exists, physical-vs-logical axis split in moe() docstring, the 128 alignment rationale) is preserved and reframed to be useful to a reader who has no PR context. * Drop the jnp.where(jnp.isnan(sparse_probs), 0, sparse_probs) guard. Tracing fused_topk_with_score_function.cu shows the kernel divides by sum_scores + 1e-20, so finite non-negative sigmoid scores cannot produce NaN here; the filter was only defense against upstream NaNs, which would mask a real regression if anything ever did start producing them. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: drop fp32 island around silu+multiply (fwd, bwd, reference) The SwiGLU intermediate (activation inputs gate_proj_out/up_proj_out, silu+multiply, and activation output) was previously promoted to fp32 in _ffn_fwd_per_shard and again in _ffn_bwd_per_shard, then cast back to the wi/wo GEMM dtype. The promotion bought nothing: the activation inputs come out of the wi grouped_gemm in bf16, the activation output is consumed by the wo GEMM (or wo's quantizer for FP8/FP4) in the same dtype, and storing higher precision than either consumer is wasted bandwidth. * _ffn_fwd_per_shard: drop the .astype(jnp.float32) on gate_proj_out and up_proj_out and the trailing .astype(sorted_x.dtype). The multiply now stays in the wi GEMM output dtype end-to-end. * _ffn_bwd_per_shard: symmetric simplification. jax.vjp(act_fn, ...) runs at bf16, both d_intermediate * silu' and d_intermediate * up stay at bf16, no casts. silu' is now consistent with silu (both bf16) so the chain rule composes cleanly without the prior fp32 detour. * tests/jax/test_te_ep_moe.py::_pure_jax_moe_reference: drop the matching fp32 silu in the parity reference so the test compares bf16-vs-bf16. Parity tolerance was not loosened; expect the comparison to tighten now that both sides round silu identically. Also fix an inaccurate inline comment at the apply_topk_weights_early fwd branch: the bf16 requirement on expert_outputs is enforced by ep_bootstrap (which rejects max_token_dtype != bf16 and sizes the NCCL EP HT mega-buffer for 2-byte slots accordingly), not by a runtime assert in the combine FFI. Signed-off-by: Teddy Do <tdophung@nvidia.com> * remove useless comments Signed-off-by: Teddy Do <tdophung@nvidia.com> * tests/jax: remove legacy MoE VJP tests + launcher; point CI at TE-EP successor test_moe_vjp.py and test_multiprocess_moe_vjp.py both import PermutationBackend from transformer_engine.jax.moe -- an API that was removed during the Phuong PR #3036 resync. Both files have been dead-on-import ever since; the multiprocess launcher run_multiprocess_moe_vjp.sh only points at the dead test. test_te_ep_moe.py (the TE-EP-only custom_vjp suite) already covers everything the legacy files exercised that is still meaningful: fwd, bwd parity vs the pure-JAX reference, aux loss, both score functions, multi-process. The legacy parametrize axis (PermutationBackend.PURE_JAX vs TRITON) no longer exists. * Delete tests/jax/test_moe_vjp.py * Delete tests/jax/test_multiprocess_moe_vjp.py * Delete tests/jax/run_multiprocess_moe_vjp.sh * qa/L0_jax_distributed_unittest/test.sh: switch the MoE VJP distributed suite invocation from run_multiprocess_moe_vjp.sh / test_multiprocess_moe_vjp.py to run_te_ep_moe.sh / test_te_ep_moe.py. * tests/jax/conftest.py: docstring reference updated. * tests/jax/test_te_ep_moe.py: drop stale "successor to ..." aside and the "mirroring run_multiprocess_moe_vjp.sh" parenthetical. Net: -981 / +9. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: swap _Ctx to @flax.struct.dataclass, drop manual pytree boilerplate Per reviewer feedback (Jaberchtold on PR #3036): the manual tree_flatten / tree_unflatten on _Ctx duplicate exactly what @flax.struct.dataclass auto-generates, and the permutation dataclasses elsewhere in this module already use flax.struct. Switching to @flax.struct.dataclass: * Removes ~75 lines of mechanical tree_flatten / tree_unflatten that have to be kept in sync with the field list by hand. * Keeps cfg as the single static field via flax.struct.field(pytree_node=False), so the fwd -> bwd boundary behavior under jax.custom_vjp is unchanged. * Drops two now-unused imports (dataclasses.dataclass, jax.tree_util.register_pytree_node_class) and adds flax.struct. Field order and the (children, aux_data) split are byte-equivalent to the previous manual implementation, so the pytree treedef seen by jax.custom_vjp is identical. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: drop bwd recv_topk_weights NaN sanitizer; trust the dispatch contract Mirrors the sparse_probs NaN-sanitizer removal in fe44697: we trust ep_dispatch_fwd's contract that recv_topk_weights does not contain NaN, and would rather see NaN propagate (catching a contract violation immediately) than silently sanitize it. The mask_bool dance itself stays: ctx.expert_outputs and grad_pre_combine still carry NaN at padded slots (ep_dispatch_fwd leaves uninit memory in recv_tokens, FFN and combine_bwd propagate it), and IEEE NaN * 0 = NaN means jnp.where is structurally needed to overwrite padded positions with literal zeros before the sum reduction. What changed: * Drop `recv_w_clean = jnp.where(jnp.isnan(...), 0, ...)` and thread ctx.recv_topk_weights directly into w / mask_bool. * Replace the NaN-defensive comment block with a shorter note that explains the structural reason the mask is still needed (NaN in expert_outputs / grad_pre_combine at padded slots), without claiming anything about recv_topk_weights. Addresses Greptile P1 by removing the asymmetry (fwd had no sanitizer, bwd did) -- chosen direction is "remove the bwd sanitizer", matching the project-wide stance of trusting kernel contracts rather than papering over violations. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: assert output dtype; tests cover d_x parity (dtype + values) Two related dtype-contract changes: 1. moe.py: one-line assert at the moe() return path that output.dtype == x.dtype. Cheap structural guard against any future bug that lets the public output drift wider than the user-supplied input dtype. 2. test_te_ep_moe.py: extend test_backward to also check d_x, the gradient propagated back to the previous layer in backprop. _grad_step now uses jax.grad(loss_fn, argnums=(0, 1)) and returns (grads_variables, grad_x); the reference path does the same so we can compare. d_x is checked for: * shape == x.shape * dtype == x.dtype (protects the _with_sharding_constraint_cast_bwd wrapper that casts the fp32-promoted gate path back to the primal dtype on bwd; a regression in that wrapper would silently double activation gradient bandwidth) * finiteness + non-zero * numerical parity vs the pure-JAX reference d_x Addresses jberchtold review comment on test_te_ep_moe.py:650 ("we also need to check the final propagated gradient that will be passed onto the next layer in backprop"). test_combined_loss_grads is adjusted to ``grads, _`` unpacking; it doesn't need d_x for its main+aux finiteness check. Signed-off-by: Teddy Do <tdophung@nvidia.com> * tests/jax/test_te_ep_moe: strip docstring to just "what this suite covers" Drops two paragraphs whose content was agent-flavoured PR-review notes rather than user-facing test docs: * The final "FP8 / MXFP8 deferred" paragraph that referenced an internal review artifact (``.pr3036-review/INTEGRATION_DESIGN.md``) not in the repo. * The "Intentional non-coverage" section that explained which tests deliberately do not exist (no Flax-wrapper smoke, no re-bootstrap-mismatch test) and why -- exactly the kind of defensive / forward-looking justification prose CLAUDE.md says to keep out of the codebase. The remaining docstring covers what readers actually need: how to launch the suite, what each test class exercises, and a short note on the parametrize-vs-class layout. Addresses jberchtold review comment on test_te_ep_moe.py:54. Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: address TE EP alignment review feedback Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: fix early topk weighting padded-slot masking Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: remove unused EP mesh size Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: tighten TE EP recv capacity bound Signed-off-by: Teddy Do <tdophung@nvidia.com> * jax/moe: simplify late TE EP weighting Signed-off-by: Teddy Do <tdophung@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * jax/moe: reduce padded-slot recv weight masking Signed-off-by: Teddy Do <tdophung@nvidia.com> --------- Signed-off-by: Teddy Do <tdophung@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

phu0ngng requested review from jberchtold-nvidia and ptrendx as code owners May 22, 2026 03:01

phu0ngng mentioned this pull request May 22, 2026

[JAX] Dispatch/Combine for BF16 Tensor #3025

Open

greptile-apps Bot reviewed May 22, 2026

View reviewed changes

Comment thread build_tools/jax.py Outdated

Comment thread build_tools/jax.py Outdated

Comment thread transformer_engine/jax/cpp_extensions/ep.py Outdated

mgoldfarb-nvidia reviewed May 22, 2026

View reviewed changes

phu0ngng requested a review from tdophung May 22, 2026 15:51

jberchtold-nvidia reviewed May 22, 2026

View reviewed changes

Comment thread examples/jax/ep/ep_moe.py Outdated

Comment thread tests/jax/multi_process_launch_ep.sh Outdated

Comment thread transformer_engine/jax/cpp_extensions/ep.py Outdated

Comment thread transformer_engine/jax/cpp_extensions/ep.py Outdated

Comment thread examples/jax/ep/ep_moe.py

phu0ngng force-pushed the phuong/ep-3-jax branch from eaaffbf to c502227 Compare May 22, 2026 23:08

greptile-apps Bot reviewed May 22, 2026

View reviewed changes

Comment thread transformer_engine/jax/ep.py Outdated

Comment thread transformer_engine/jax/ep.py Outdated

phu0ngng force-pushed the phuong/ep-3-jax branch from c502227 to 69b0196 Compare May 22, 2026 23:19

greptile-apps Bot reviewed May 22, 2026

View reviewed changes

Comment thread transformer_engine/jax/ep.py Outdated

phu0ngng force-pushed the phuong/ep-3-jax branch from 8cb8de4 to e79efa9 Compare May 23, 2026 23:13

jberchtold-nvidia previously approved these changes May 26, 2026

View reviewed changes

phu0ngng dismissed jberchtold-nvidia’s stale review via 8fa5bdb May 28, 2026 23:42

tdophung reviewed May 29, 2026

View reviewed changes

Comment thread transformer_engine/jax/ep.py

Comment thread transformer_engine/jax/ep.py

Comment thread transformer_engine/jax/ep.py Outdated

Comment thread transformer_engine/jax/cpp_extensions/ep.py Outdated

greptile-apps Bot reviewed May 29, 2026

View reviewed changes

Comment thread transformer_engine/jax/csrc/extensions/ep.cpp

tdophung reviewed May 29, 2026

View reviewed changes

Comment thread transformer_engine/jax/cpp_extensions/ep.py Outdated

tdophung reviewed Jun 4, 2026

View reviewed changes

Comment thread transformer_engine/jax/csrc/extensions.h Outdated

phu0ngng force-pushed the phuong/ep-3-jax branch from e927903 to 5b0e9b7 Compare June 8, 2026 20:58

greptile-apps Bot reviewed Jun 8, 2026

View reviewed changes

Comment thread transformer_engine/jax/cpp_extensions/ep.py Outdated

greptile-apps Bot reviewed Jun 8, 2026

View reviewed changes

Comment thread transformer_engine/jax/csrc/extensions/ep.cpp

phu0ngng force-pushed the phuong/ep-3-jax branch 2 times, most recently from 06f8a13 to c34771d Compare June 10, 2026 15:24

phu0ngng force-pushed the phuong/ep-3-jax branch from c34771d to 351b9df Compare June 10, 2026 23:56

greptile-apps Bot reviewed Jun 24, 2026

View reviewed changes

Comment thread transformer_engine/jax/cpp_extensions/ep.py Outdated

phu0ngng added the 2.7.0 label Jun 24, 2026

Guard against None out_partition_spec in EP combine/dispatch-bwd part…

db8d43b

…ition methods Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

jberchtold-nvidia previously approved these changes Jun 25, 2026

View reviewed changes

Fix test_two_layer_dispatch_no_handle_aliasing: add A->B data edge

0a29355

Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

phu0ngng dismissed jberchtold-nvidia’s stale review via 0a29355 June 26, 2026 10:50

phu0ngng added 2.17 and removed 2.7.0 labels Jun 26, 2026

jberchtold-nvidia approved these changes Jun 26, 2026

View reviewed changes

Merge branch 'main' into phuong/ep-3-jax

9187d81

phu0ngng merged commit 42f37ff into NVIDIA:main Jun 27, 2026
8 of 13 checks passed

phu0ngng deleted the phuong/ep-3-jax branch June 27, 2026 08:57

greptile-apps Bot reviewed Jun 27, 2026

View reviewed changes

KshitijLakhani pushed a commit that referenced this pull request Jun 29, 2026

[JAX] Expert Parallelism: JAX primitives + VJPs (#3036)

563b9bd

* Expert Parallelism: JAX primitives + VJPs --------- Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

This was referenced Jul 15, 2026

Experiment: Single-rack B200 MoE MFU above 20% marin-community/marin#7012

Open

Experiment: TransformerEngine NCCL_EP expert-parallel MoE on B200 marin-community/marin#7331

Open

Uh oh!

Conversation

phu0ngng commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Implementation

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

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

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

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

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

Tests & example

Type of change

Checklist:

Uh oh!

greptile-apps Bot commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 4/5

Important Files Changed

Sequence Diagram

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mgoldfarb-nvidia May 22, 2026

Choose a reason for hiding this comment

Uh oh!

phu0ngng May 22, 2026

Choose a reason for hiding this comment

Uh oh!

phu0ngng commented May 22, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jberchtold-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

phu0ngng commented Jun 11, 2026

Uh oh!

phu0ngng commented Jun 24, 2026

Uh oh!

Uh oh!

phu0ngng commented Jun 25, 2026

Uh oh!

jberchtold-nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

phu0ngng commented Jun 26, 2026

Uh oh!

Uh oh!

greptile-apps Bot Jun 27, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

phu0ngng commented May 22, 2026 •

edited

Loading

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

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

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

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

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

greptile-apps Bot commented May 22, 2026 •

edited

Loading