[Draft][JAX] MoE block optimization

docs/api/jax.rst

@@ -59,3 +59,45 @@ Modules
   :members: __call__
 
 .. autoapifunction:: transformer_engine.jax.flax.extend_logical_axis_rules
+
+MoE parallelism resources
+-------------------------
+
+The experimental ``transformer_engine.jax.moe.moe`` function and Flax
+``_MoEBlock`` accept ``mesh_resource=None`` and ``quant_before_fsdp_ag=False``.
+Pass a ``MeshResource`` explicitly, or provide one through
+``global_shard_guard``. An explicit resource takes precedence; a missing
+resource raises ``ValueError``. The resource must name an EP axis. DP and
+FSDP are outer batch-sharding axes, in that order, with EP innermost.
+Unset outer axes are omitted, and a shared DP/FSDP axis is included once.
+
+For example, with an active physical mesh containing ``dp``, ``fsdp`` and
+``ep`` axes::
+
+    from transformer_engine.jax.flax import _MoEBlock
+    from transformer_engine.jax.sharding import MeshResource
+
+    block = _MoEBlock(
+        mesh_resource=MeshResource(
+            dp_resource="dp", fsdp_resource="fsdp", ep_resource="ep"
+        ),
+        quant_before_fsdp_ag=True,
+        # Supply an MXFP8 quantization recipe and the model dimensions.
+    )
+
+``quant_before_fsdp_ag=True`` quantizes local expert-weight shards before
+all-gathering their MXFP8 data and scales on the resource's FSDP axis.
+It requires an FSDP resource and MXFP8 kernel quantizers. The default
+gathers full-precision weights before quantization.
+
+With active Flax logical-axis rules, ``wi_kernel_axes`` and ``wo_kernel_axes``
+also select the forward FFN weight input specs when quantized gathering is
+enabled. FSDP can shard complete experts together with EP on the first weight
+dimension, or split a matrix dimension within each expert. Complete-expert
+gathering concatenates the quantized matrices and their existing scale blocks;
+matrix-shard gathering reconstructs each expert's scale blocks. These specs
+are resolved at trace time and do not require JAX Explicit sharding mode.
+
+The old ``ep_axis`` and ``data_parallelism_axes`` arguments remain accepted
+with a ``DeprecationWarning``. They are translated into a resource before
+calling the new API. Conflicting old and new arguments raise ``ValueError``.

tests/jax/conftest.py

@@ -98,7 +98,12 @@ def pytest_addoption(parser):
     """
     parser.addoption("--num-process", action="store", default=0)
     parser.addoption("--process-id", action="store", default=0)
-
+    parser.addoption(
+        "--use-cudnn-fusion",
+        choices=("0", "1"),
+        default="1",
+        help="Pass use_cudnn_fusion to MoE blocks and EP capacity calculation",
+    )
 
 def pytest_configure(config):
     config.addinivalue_line(