[DRAFT][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 substantially optimizes the JAX MoE path by combining FC1 projections, adding distributed grouped-quantize/GEMM partitioning, expanding EP diagnostics and receive-capacity handling, and adding distributed coverage. It also changes the exported Flax MoE parameter layout without checkpoint migration.

• Reworks JAX MoE forward/backward execution and expert-parallel sharding.
• Adds grouped quantization and grouped GEMM custom partitioning for DP/FSDP/EP layouts.
• Adds EP handle tracing and cache-update diagnostics.
• Expands single- and multi-process distributed MoE/grouped-GEMM tests.

Confidence Score: 4/5

The Flax MoE checkpoint incompatibility should be addressed before merging so existing parameter and optimizer states remain restorable.

The changed module replaces two previously exported parameter leaves with a differently named, differently shaped leaf and provides no migration path, causing old checkpoints to fail when applied.

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

Important Files Changed

Filename	Overview
transformer_engine/jax/moe.py	Rewrites the MoE custom-VJP and distributed grouped-FFN path; inspected gradient and sharding branches were internally consistent.
transformer_engine/jax/flax/moe.py	Combines FC1 parameters into wi, breaking restoration of parameter and optimizer trees created with wi_0 and wi_1.
transformer_engine/jax/cpp_extensions/gemm.py	Adds custom partitioning and local-shape derivation for grouped GEMM across supported DP/FSDP/EP axes.
transformer_engine/jax/cpp_extensions/quantization.py	Preserves logical grouped tensor ranks and adds sharding-aware MXFP8 scale carriers and partitioning.
transformer_engine/common/ep/ep_backend.cpp	Adds handle-update tracking and opt-in diagnostics without an established default-path security failure.
tests/jax/test_te_ep_moe.py	Greatly expands MoE numerical and distributed coverage but does not cover restoration of the previous Flax parameter tree.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Input tokens] --> B[Gate and top-k routing]
  B --> C[EP prepare and dispatch]
  C --> D[Grouped FC1 quantize and GEMM]
  D --> E[Activation and grouped FC2 GEMM]
  E --> F[EP combine]
  F --> G[MoE output]

Loading

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