[JAX] Add an MoE Block (Layer) that compound router, permutation, groupedGEMM and communication

docs/envvars.rst

@@ -458,6 +458,12 @@ JAX Triton Extensions
    :Default: ``0``
    :Description: Raise a ``RuntimeError`` when the installed JAX is too old to safely run ``TritonAutotunedKernelCall`` (`jax-ml/jax#35218 <https://github.com/jax-ml/jax/pull/35218>`_) instead of silently falling back to non-autotuned dispatch. Useful for CI or debugging to ensure Triton autotuning is active. When set to ``0`` (default), old JAX versions silently fall back to single-config (non-autotuned) kernel dispatch for compatibility.
 
+.. envvar:: NVTE_TRITON_PERMUTATION_BLOCK_SIZES
+
+   :Type: comma-separated list of ``int`` (e.g. ``"128"`` or ``"64,128,256"``)
+   :Default: ``"64,128,256,512,1024,2048,4096"`` (the full sweep)
+   :Description: Override the ``BLOCK_SIZE`` configs evaluated by ``triton.autotune`` for the MoE permutation kernels in ``transformer_engine/common/triton/permutation.py`` (``_permute_kernel``, ``_unpermute_kernel``, ``_unpermute_bwd_with_merging_probs_kernel``, ``_sort_chunks_by_map_kernel``). The default 7-config sweep yields the best runtime on production shapes but costs ~1-5 s of MLIR→LLVM→PTX→cubin compile per config-per-kernel on a cold start (≈2-5 min total per backend, serialized on a single GPU). Set to a single value (e.g. ``"128"``) to skip autotuning entirely for tests / CI where correctness -- not throughput -- is the goal. Must be a comma-separated list of positive ints; malformed values raise ``ValueError`` at kernel-registration time. **Do NOT set this in production runs** -- you will lose autotuned performance.
+
 Examples
 --------
 

qa/L0_jax_distributed_unittest/test.sh

@@ -37,6 +37,33 @@ wait
 TE_PATH=$TE_PATH bash $TE_PATH/examples/jax/collective_gemm/run_test_cgemm.sh || test_fail "run_test_cgemm.sh"
 wait
 
+# MoE custom_vjp distributed (Level 2 smoke + Level 3 perf). Single-host
+# multi-GPU; requires >=4 visible GPUs.
+#
+# Flags required for this file (mirrored in tests/jax/run_distributed_moe_vjp.sh):
+#
+# * ``-p no:typeguard`` — jaxtyping's pytest plugin auto-loads typeguard,
+#   whose @typechecked import hook materialises JAX tracers via isinstance()
+#   checks during shard_map tracing. We disable it only here (other jax tests
+#   need it for type-hint validation).
+# * ``XLA_PYTHON_CLIENT_PREALLOCATE=false`` + ``MEM_FRACTION=0.5`` —
+#   prevents NCCL OOM during EP all-to-all communicator setup (default 90%
+#   preallocation leaves no room).
+# * ``CUDA_LAUNCH_BLOCKING=1`` — workaround for an async-dispatch hang
+#   between Triton custom_calls with ``input_output_aliases`` and the
+#   downstream NCCL ragged_all_to_all in this test's bwd path. Without it,
+#   MainThread parks in _pjit_call_impl_python and one GPU pins at 100%
+#   forever. With it, the smoke suite passes in <1 min. See
+#   ``tests/jax/test_distributed_moe_vjp.py`` module docstring for the
+#   bisection record and TODO for the proper fix.
+XLA_PYTHON_CLIENT_PREALLOCATE=false XLA_PYTHON_CLIENT_MEM_FRACTION=0.5 \
+    CUDA_LAUNCH_BLOCKING=1 \
+    python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v -s \
+    -p no:typeguard \
+    --junitxml=$XML_LOG_DIR/pytest_test_distributed_moe_vjp.xml \
+    $TE_PATH/tests/jax/test_distributed_moe_vjp.py || test_fail "test_distributed_moe_vjp.py"
+wait
+
 if [ $RET -ne 0 ]; then
     echo "Error: some sub-tests failed: $FAILED_CASES"
     exit 1

tests/jax/conftest.py

@@ -86,6 +86,20 @@ def pytest_sessionfinish(self, session, exitstatus):
         print("=" * 80)
 
 
+def pytest_addoption(parser):
+    """CLI options for multiprocess JAX tests.
+
+    Mirrors examples/jax/encoder/conftest.py so multiprocess tests in
+    tests/jax/ can be launched one-process-per-GPU via a sibling shell
+    script. Required by tests/jax/test_multiprocess_moe_vjp.py to work
+    around the JAX/XLA + lazy Triton kernel load + active NCCL deadlock
+    documented in past_JAX_XLA_deadlock.txt and nvbug/5564750. Harmless
+    for other tests; defaults to 0 (= "not a multiprocess launch").
+    """
+    parser.addoption("--num-process", action="store", default=0)
+    parser.addoption("--process-id", action="store", default=0)
+
+
 def pytest_configure(config):
     config.addinivalue_line(
         "markers",