[JAX] MoEBlock tutorial

[jberchtold-nvidia]

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

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

Please list the changes introduced in this PR:

• Change A
• Change B

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

This PR adds a new JAX tutorial for Mixture-of-Experts with TransformerEngine's experimental _MoEBlock, including a native JAX/Flax BF16 baseline (moe_native.py), the TE integration example (moe.py), RST documentation (moe.rst), pre-captured benchmark output (moe.out), and pytest tests (test_moe.py). The hub page (te_jax_integration.rst) is updated to link to the new tutorial.

• The core tutorial demonstrates a 2×2 EP/FSDP mesh setup, correctness comparison between native and TE paths, and a forward+backward timing benchmark showing ~1.43× TE speedup on GB200 hardware.
• A helper function _block_until_ready_tree is introduced for benchmark synchronization but only waits on leaves[0] rather than the whole tree; the same single-leaf wait pattern is repeated inline in setup_demo and test_moe.py, and the main() entrypoint silently skips the native-only benchmark when TE is unavailable.

Confidence Score: 4/5

The tutorial is safe to merge; all issues are non-blocking style and consistency concerns with no impact on correctness of the documented TE path.

The core benchmark helper _block_until_ready_tree only synchronizes the first leaf of the output tree, and the same single-leaf pattern is repeated in setup_demo and in the test. Functionally, JAX/XLA makes all outputs of a single JIT dispatch available together, so benchmarks are not mismeasured today, but the implementation contradicts its name and is misleading tutorial code. Separately, main() bails out before running the native baseline when TE is unsupported, leaving users on pre-Blackwell hardware with no benchmark output at all. The _te_moe_available function in test_moe.py duplicates the capability check from moe.py, creating a maintenance risk if the threshold changes.

docs/examples/jax/moe.py — _block_until_ready_tree, setup_demo sync call, and main() early-return logic; docs/examples/jax/test_moe.py — duplicated capability check and sync pattern.

Important Files Changed

Filename	Overview
docs/examples/jax/moe.py	Main tutorial script; introduces _block_until_ready_tree that only waits on leaves[0], and main() exits early when TE is unsupported, skipping the native baseline entirely.
docs/examples/jax/moe_native.py	Native JAX/Flax MoE baseline; ragged all-to-all + grouped GEMMs are implemented correctly; no significant issues found.
docs/examples/jax/test_moe.py	Pytest tests with appropriate requires_4gpu / requires_te_moe skip guards; duplicates capability-check logic from moe.py and repeats the single-leaf block_until_ready pattern.
docs/examples/jax/moe.rst	RST tutorial document; literalinclude markers align with moe.py; benchmark tables and correctness numbers are well documented.
docs/examples/jax/moe.out	Pre-captured benchmark output used by literalinclude; no code issues.
docs/examples/te_jax_integration.rst	Hub page updated to link to the new MoE tutorial; minimal change, correct placement.

Sequence Diagram

sequenceDiagram
    participant User
    participant moe.py
    participant NativeMoEBlock
    participant TEMoEBlock
    participant JAX/XLA

    User->>moe.py: python moe.py
    moe.py->>moe.py: build_ep_fsdp_mesh()
    moe.py->>NativeMoEBlock: native_model.init(k_init, x)
    NativeMoEBlock-->>moe.py: variables (gate_kernel, wi_0, wi_1, wo)
    moe.py->>moe.py: shard_inputs_and_variables()

    Note over moe.py,JAX/XLA: compare_forward()
    moe.py->>JAX/XLA: jit(native_model.apply)(variables, x)
    JAX/XLA-->>moe.py: native_out
    moe.py->>JAX/XLA: jit(te_apply)(variables, x)
    JAX/XLA-->>moe.py: te_out
    moe.py->>User: "max |native - TE| = 0.0604"

    Note over moe.py,JAX/XLA: run_benchmarks() — fwd+bwd timing
    loop warmup + timing iters
        moe.py->>JAX/XLA: jit(value_and_grad(loss_fn))(vars, x, dy)
        JAX/XLA-->>moe.py: (loss, (grad_vars, grad_x))
        moe.py->>moe.py: _block_until_ready_tree(result)
    end
    moe.py->>User: Mean time (native), Mean time (TE)

Loading

Comments Outside Diff (1)

1. docs/examples/jax/test_moe.py, line 975-991 (link)

   Duplicated capability-check logic

   _te_moe_available() here is an exact copy of te_moe_supported() in moe.py. If the Blackwell compute-capability threshold, the module path, or the attribute name ever changes, both copies need to be updated in sync. Consider importing te_moe_supported from moe directly (since moe is already imported in the TE-guarded tests) or moving the shared check into a small helper module.

   Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

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

[greptile-apps]

main() returns immediately when TE is unavailable (non-Blackwell hardware or missing build), so the native baseline never runs either. Users who just want to validate the native JAX path — or who are running on a pre-Blackwell GPU — see a silent skip with no output. Separating the TE guard from the native baseline makes the script more useful as a standalone tutorial.

Suggested change

	return
	demo = setup_demo()
	compare_forward(demo)
	run_benchmarks(demo)
	if __name__ == "__main__":
	te_supported, te_reason = te_moe_supported()
	if not te_supported:
	print(f"[skipped TE comparison: {te_reason}]")
	demo = setup_demo()
	if te_supported:
	compare_forward(demo)
	run_benchmarks(demo)