[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]

Confidence Score: 5/5

The PR appears safe to merge; no actionable defects remain, and the prior expert-parallel masking finding was manually resolved after the example was corrected.

Summary

• Documents routing, load balancing, token permutation, grouped expert computation, and expert parallelism.
• Adds concise framework-specific examples and diagrams, including dark-theme styling.
• Improves public router and JAX MoE docstring formatting.
• Revises the grouped-MLP explanation to distinguish fused FC1 activation from the separate FC2 grouped GEMM.
• Corrects the PyTorch expert-parallel example by clearing unused receive-buffer tail rows before expert computation and combine.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  T[Input tokens] --> R[Router and top-k]
  R --> D{Expert placement}
  D -->|Single device| P[Token permutation]
  D -->|Expert parallel| EP[All-to-all dispatch]
  P --> G[Grouped expert MLP]
  EP --> G
  G --> C{Combine path}
  C -->|Single device| U[Token unpermute and weighted combine]
  C -->|Expert parallel| EC[All-to-all combine]
  U --> O[MoE output]
  EC --> O

Loading

[pggPL]

Hi @vthumbe1503 @phu0ngng this is PR with docs for MoE operations and EP. Can you have a look?

[jberchtold-nvidia]

Small update, Hopper BF16 is now supported since this PR merged yesterday: #3083

[jberchtold-nvidia]

@pggPL How do we want to handle group alignment in the docs? This does support MXFP8 but we require group sizes to be aligned to a multiple of 128 for MXFP8, so if a user passes in un-aligned groups we can get IMA, illegal instruction, or incorrect results

For now, I've been guiding users to our more monolithic MoE block to avoid this complexity. But the same constraints apply to TE/PyTorch, so if you've found a better way to explain the nuances of this alignment, let me know and I'm open to adding it.

The on-device group size alignment is difficult since we can't assert it on Host without introducing runtime overhead

[jberchtold-nvidia]

At the moment, I'd prefer guiding users to our monolithic MoE layer here:

TransformerEngine/transformer_engine/jax/moe.py

Line 1208 in 1634a5a

def moe(

It's simpler and handles things like group size alignment automatically and hides it from the user. We do also want to highlight the lower-level APIs like this at some point in the future, but when we do so we need to significantly mark all the caveats like group size alignment and other constraints, which may be changing as we add new fused kernel support. Additionally, te_permutation.token_dispatch/combine are only recommended for non-EP, with EP enabled the TE EP APIs should be used. This is handled automatically by the TE MoEBlock, so my preference is to guide users to this monolithic MoE block as a starting point rather than these lower-level APIs

If Phuong wants to include a dedicated doc on TE EP APIs specifically, I'm okay with that lower-level API being highlighted. For other things, like grouped GEMM, I'd prefer to hide the alignment complexity in the TE MoEBlock and limit usage of lower-level APIs like grouped_dense to only users who are not doing a standard MoE architecture and are okay with handling the alignment constraints.

[jberchtold-nvidia]

Thanks for adding JAX documentation as well @pggPL! 🙌

I've reviewed and left a few comments. Let me know what you think. Thanks!