[Docs] Add Mixture of Experts guide

[pggPL]

Adds a Mixture of Experts guide covering routing, token permutation, grouped expert computation, and expert parallelism.

Includes concise PyTorch and JAX examples.

[greptile-apps]

Greptile Summary

Adds a comprehensive Mixture of Experts guide covering routing, token permutation, grouped expert computation, and expert parallelism across PyTorch and JAX.

• Adds runnable-style PyTorch and JAX snippets for MoE building blocks and complete layer composition.
• Extends the public API references and documentation navigation.
• Adds diagrams with light- and dark-theme styling.
• Updates router and JAX MoE docstrings to clarify API contracts.
• Correctly masks unused fixed-capacity PyTorch EP receive slots before combine, resolving the previous finding.

Confidence Score: 5/5

The PR appears safe to merge; no outstanding correctness, security, or repository-rule issues remain.

The only change since the previous review constructs a per-expert validity mask from the fixed slot capacity and token counts, then zeros padded expert outputs while applying routing weights. This fully addresses the previously reported risk of undefined padded rows reaching expert combine, and the resolved thread is not outstanding.

Important Files Changed

Filename	Overview
docs/features/mixture_of_experts/mixture_of_experts.rst	Adds the central MoE guide and accurately connects routing, permutation, grouped computation, and expert-parallel APIs.
docs/features/mixture_of_experts/moe_expert_parallel_pytorch.py	Demonstrates PyTorch expert parallelism and now masks padded receive slots before combining expert outputs.
docs/features/mixture_of_experts/moe_expert_parallel_jax.py	Demonstrates the required eager bootstrap, mesh-resource setup, signature recording, and functional JAX MoE call.
docs/features/mixture_of_experts/moe_layer_pytorch.py	Adds a concise differentiable single-device PyTorch MoE composition example.
docs/features/mixture_of_experts/moe_layer_jax.py	Adds a jitted single-device JAX MoE composition using dispatch, grouped dense, and weighted combine.
transformer_engine/jax/moe.py	Improves documentation of JAX MoE behavior and public contracts without changing runtime behavior.
transformer_engine/jax/router.py	Clarifies JAX router API documentation without changing routing behavior.
transformer_engine/pytorch/router.py	Clarifies PyTorch router API documentation without changing routing behavior.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  A[Input tokens] --> B[Router scoring and top-k]
  B --> C[Token or EP dispatch]
  C --> D[Expert-contiguous tokens]
  D --> E[Grouped expert computation]
  E --> F[Token or EP combine]
  F --> G[Output in original token order]
  B -. routing weights .-> F

Loading

_{Reviews (2): Last reviewed commit: "[Docs] MoE: mask padded receive slots be..." | Re-trigger Greptile}