[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 opt-in EP overflow dropping and exposes pre-drop receive demand throughout the JAX EP/MoE stack.

• Plumbs drop_on_overflow through Python bootstrap, pybind, and the C++ EP group configuration.
• Adds total_recv_tokens outputs to prepare, dispatch, fused MoE, and the Flax MoE wrapper.
• Adds explicit EP teardown/re-bootstrap support and multi-process overflow coverage.

Confidence Score: 3/5

The public return-contract break should be resolved before merging because existing EP and MoE callers will fail even without opting into overflow dropping.

The low-level overflow plumbing and custom-VJP structures are coordinated, but ep_prepare, ep_dispatch, moe, and the Flax wrapper all unconditionally add positional outputs, causing existing fixed-arity unpacking to raise at runtime.

Files Needing Attention: transformer_engine/jax/ep.py, transformer_engine/jax/cpp_extensions/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 bootstrap plumbing, teardown support, and a new dispatch result, but unconditionally breaks the existing positional return contract.
transformer_engine/jax/cpp_extensions/ep.py	Extends the prepare primitive and its sharding metadata with pre-drop receive totals; the public primitive now has incompatible tuple arity.
transformer_engine/jax/csrc/extensions/ep.cpp	Correctly forwards the overflow policy into EP initialization and returns the new telemetry buffer through FFI.
transformer_engine/jax/moe.py	Threads telemetry through the fused custom VJP correctly, but changes the public MoE result from two to three positional values.
transformer_engine/jax/flax/moe.py	Exposes MoE telemetry from the Flax module while inheriting the public return-contract break.
tests/jax/test_multi_process_ep.py	Updates EP callers and adds multi-process coverage for dropping overflow and reporting pre-drop demand.

Sequence Diagram

sequenceDiagram
  participant User
  participant Bootstrap as ep_bootstrap
  participant EP as EP backend
  participant Prepare as ep_prepare
  participant Dispatch as ep_dispatch / moe
  User->>Bootstrap: drop_on_overflow
  Bootstrap->>EP: NVTEEpGroupConfig
  User->>Prepare: routing indices
  Prepare->>EP: exchange routing metadata
  EP-->>Prepare: token counts + total receive demand
  User->>Dispatch: tokens and routing
  Dispatch->>EP: dispatch with configured capacity
  EP-->>Dispatch: received tokens + pre-drop demand
  Dispatch-->>User: output tuple including total_recv_tokens

Loading

_{Reviews (1): Last reviewed commit: "[JAX] Add NCCL EP drop-on-overflow polic..." | Re-trigger Greptile}