[JAX] Optimize MoE block

[jberchtold-nvidia]

Description

Improves performance of the MoE block by exposing support for quantization, removal of unnecessary masking overheads, and support for less memory usage via a reduced receive capacity in TE EP

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

• Direct support for MXFP8 quantization in the MoE block along with corresponding tests
• Removal of additional overheads like jnp.where masking that wasn't required as TE EP and grouped GEMM are group-aware
• Support for a reduced receive capacity and integration with TE EP's overflow detection

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

The PR optimizes the JAX expert-parallel MoE path by adding MXFP8 grouped quantization, reducing receive-capacity memory usage, removing redundant padding masks, and storing the gated FC1 weights contiguously.

• Adds configurable receive capacity and overflow reporting to the functional and Flax MoE interfaces.
• Threads grouped quantizer sets through the MoE forward and backward paths.
• Expands distributed tests to cover BF16 and MXFP8 numerical behavior.

Confidence Score: 4/5

The PR is not yet safe to merge because existing _MoEBlock checkpoints and optimizer states cannot be restored after the FC1 parameter-tree replacement.

The current module requests one wi leaf where prior checkpoints contain wi_0 and wi_1, and current HEAD provides no migration or compatibility path.

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

Important Files Changed

Filename	Overview
transformer_engine/jax/flax/moe.py	Adds recipe-driven grouped quantizers, receive-capacity configuration, and a contiguous FC1 parameter layout to the Flax MoE adapter.
transformer_engine/jax/moe.py	Reworks the expert FFN forward and backward paths for grouped quantization, explicit capacity sizing, overflow reporting, and global sharding.
transformer_engine/jax/cpp_extensions/quantization.py	Allows global stateless MXFP8 grouped-quantizer descriptors to cover shard-local group counts.
transformer_engine/jax/quantize/tensor.py	Updates quantized tensor checkpoint handling used by the MoE custom differentiation path.
tests/jax/test_te_ep_moe.py	Extends distributed MoE numerical tests to BF16 and MXFP8 configurations and the contiguous FC1 layout.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Sharded hidden states] --> B[Gate and top-k routing]
  B --> C[EP prepare and dispatch]
  C --> D[Grouped FC1 quantization and GEMM]
  D --> E[Activation]
  E --> F[Grouped FC2 quantization and GEMM]
  F --> G[EP combine]
  C -. receive demand .-> H[Overflow detection]

Loading

_{Reviews (9): Last reviewed commit: "Merge branch 'main' into jberchtold/moeb..." | Re-trigger Greptile}

[jberchtold-nvidia]

/te-ci L1 jax

[jberchtold-nvidia]

/te-ci L1 jax

[jberchtold-nvidia]

/te-ci L1 jax

[tdophung]

LGTM pending CI

[jberchtold-nvidia]

/te-ci L1 jax

[jberchtold-nvidia]

/te-ci L1 jax