[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 block by adding grouped MXFP8 quantization, configurable receive capacity, and a contiguous FC1 kernel layout.

• Threads independent FC1/FC2 quantizer sets through the custom VJP.
• Adds aligned receive-capacity calculation and overflow integration.
• Reworks grouped-GEMM sharding and residual handling.
• Extends distributed BF16 and MXFP8 numerical tests.

Confidence Score: 4/5

The PR is not yet safe to merge because legacy _MoEBlock checkpoints and optimizer states cannot be restored into the new parameter tree without migration.

Current HEAD still replaces wi_0 and wi_1 with a single differently shaped wi parameter and provides no adapter for previously persisted Flax variables or optimizer state.

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 and receive-capacity plumbing while replacing two FC1 parameter leaves with one contiguous kernel; legacy checkpoint compatibility remains outstanding.
transformer_engine/jax/moe.py	Reworks the MoE custom VJP for quantized grouped GEMMs, global sharding, and reduced receive capacity.
transformer_engine/jax/cpp_extensions/quantization.py	Permits stateless MXFP8 grouped quantizers whose global descriptor contains at least the shard-local group count.
transformer_engine/jax/quantize/tensor.py	Adjusts grouped quantized-tensor handling needed by the new MoE sharding path.
tests/jax/test_te_ep_moe.py	Expands distributed MoE parity coverage to BF16 and Blackwell MXFP8 using the contiguous FC1 layout.
tests/jax/run_te_ep_moe.sh	Allows callers to override the distributed MoE test file through TEST_FILE.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Flax _MoEBlock] --> B[Gate and top-k routing]
  B --> C[EP dispatch]
  C --> D[Grouped FC1 quantize and GEMM]
  D --> E[SwiGLU]
  E --> F[Grouped FC2 quantize and GEMM]
  F --> G[EP combine]
  H[Receive capacity] --> C
  I[FC1 and FC2 quantizer sets] --> D
  I --> F

Loading

_{Reviews (2): Last reviewed commit: "[pre-commit.ci] auto fixes from pre-comm..." | Re-trigger Greptile}

[jberchtold-nvidia]

/te-ci L1 jax

[jberchtold-nvidia]

/te-ci L1 jax

[jberchtold-nvidia]

/te-ci L1 jax