[JAX] EP Dispatch with overflow detection option

[phu0ngng]

Description

Adds a way to size EP recv buffers below the worst case and still stay correct. ep_prepare / ep_dispatch now expose total_recv_tokens, the per-rank pre-drop recv-slot demand, so callers can tell exactly when a step's routing exceeds recv_capacity_per_rank; no need to conservatively provision every rank for the theoretical maximum.

Bootstrap also gains an opt-indrop_on_overflow policy so an occasional overflowing step is handled gracefully and keeps running (the default behavior is unchanged).

Together these let users pick a tighter, cheaper recv capacity and monitor total_recv_tokens to confirm it holds or to drive their own capacity policy.

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

Changes

• Bootstrap:ep_bootstrap gains a drop_on_overflow flag, plumbed through set_ep_bootstrap_params (pybind) and NVTEEpGroupConfig.drop_on_overflow.
  • New output:ep_prepare now returns (token_counts, total_recv_tokens, handle_mem); ep_dispatch's
    custom_vjp primal correspondingly returns total_recv_tokens as a non-differentiable output.
  • MoE:moe() and the Flax _MoEBlock now return(output, aux_loss, total_recv_tokens); the extra output is non-differentiable
  • Teardown: adds ep_finalize() (and reset_ep_config) to tear down the EP
    communicator so a process can re-ep_bootstrap with a different config.

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]

Positional return contracts break

Existing callers that unpack the documented two-value ep_prepare, four-value ep_dispatch, or two-value MoE results now receive an additional positional value unconditionally, causing ValueError: too many values to unpack even when overflow dropping is disabled.

Knowledge Base Used:JAX Fused Layers: Functions and Flax Modules

[greptile-apps]

Greptile Summary

Adds optional expert-parallel overflow dropping and exposes each rank’s pre-drop receive demand.

• Plumbs drop_on_overflow through JAX bootstrap, pybind, and the native EP configuration.
• Adds total_recv_tokens outputs to EP prepare/dispatch and fused MoE interfaces.
• Adds EP communicator finalization and re-bootstrap support.
• Extends distributed tests and examples for overflow handling and the new metadata.

Confidence Score: 4/5

The PR is not yet safe to merge because its public EP and MoE functions still unconditionally change positional return arities and break existing unpacking callers.

Existing callers written against the previous ep_prepare, ep_dispatch, moe, or Flax MoE tuple contracts receive additional values even when overflow dropping is disabled, producing runtime unpacking failures; the previously reported compatibility issue remains in the current code.

Files Needing Attention: transformer_engine/jax/ep.py, transformer_engine/jax/moe.py, transformer_engine/jax/flax/moe.py

Important Files Changed

Filename	Overview
transformer_engine/jax/ep.py	Adds overflow policy plumbing, receive-demand reporting, and explicit EP teardown.
transformer_engine/jax/cpp_extensions/ep.py	Extends the EP prepare primitive with a sharded total_recv_tokens result and adds cached-config reset support.
transformer_engine/jax/csrc/extensions/ep.cpp	Passes the overflow policy into native initialization and returns pre-drop receive totals through the FFI handler.
transformer_engine/jax/moe.py	Propagates total_recv_tokens through the fused MoE custom-VJP forward result while discarding its cotangent.
transformer_engine/jax/flax/moe.py	Exposes the receive-demand metadata through the Flax MoE block.
tests/jax/test_multi_process_ep.py	Adds multi-process coverage for overflow dropping, demand reporting, finalization, and re-bootstrap.

Sequence Diagram

sequenceDiagram
participant User
participant JAX as JAX EP/MoE API
participant FFI as JAX FFI
participant EP as Native EP communicator
User->>JAX: ep_bootstrap(drop_on_overflow)
JAX->>FFI: set bootstrap parameters
FFI->>EP: initialize communicator and buffers
User->>JAX: ep_prepare / ep_dispatch
JAX->>FFI: routing metadata and tensors
FFI->>EP: prepare and dispatch
EP-->>FFI: routed tensors and pre-drop demand
FFI-->>JAX: total_recv_tokens
JAX-->>User: outputs plus overflow signal
User->>JAX: ep_finalize()
JAX->>EP: release resources

Loading

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

[jberchtold-nvidia]

When drop-on-overflow aka overflow-detection is disabled, this will now return a "trt" tensor that has an uninitialized value, right? Can we instead update this so ep_prepare will return (token_counts, None, handle_mem) instead when drop-on-overflow=False?

[phu0ngng]

No, it actually has the total_recv_tokens value i.e., the sum of token_counts.

The only difference here is that when drop-on-overflow=False (which is the default), as soon as there is an overflow, the kernel will trap and the program will crash.

[jberchtold-nvidia]

Oh I see, then the current implementation looks good. The PR LGTM once CI passes

[phu0ngng]

/te-ci L1 JAX