[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

[Low risk] Adds Mixture of Experts documentation and examples.

The PR appears safe to merge; no new actionable issue was identified.

Summary

The PR adds a Mixture of Experts guide, PyTorch and JAX examples, API references, and diagrams. The change since the previous review clarifies the required dtype, device, and layout of PyTorch’s compact top-k index buffer.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
  Router --> Dispatch[Token dispatch]
  Dispatch --> Experts[Grouped expert computation]
  Experts --> Combine[Token combine]
  Dispatch -. Expert parallelism .-> EPDispatch[All-to-all dispatch]
  EPDispatch --> Experts
  Experts --> EPCombine[All-to-all combine]

Loading

_{Reviews (8) · Last reviewed commit: "Clarify top-k index buffer requirements"}

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

[vthumbe1503]

I have reviewed the GroupedMLP side of documentation for now. And it looks pretty solid @pggPL !

[vthumbe1503]

Suggested change

	supported block-scaled recipe (MXFP8 or NVFP4). When the configuration is not
	supported, the three operations run separately.
	supported block-scaled recipe (MXFP8 or NVFP4). Fusion is only supported for a specific configuration [link](https://github.com/NVIDIA/TransformerEngine/blob/ea1a165decdd09c1ef98b273d965796c89c47473/transformer_engine/pytorch/ops/fused/grouped_mlp.py#L816-L817). When the configuration is not
	supported, the three operations run separately.

We should specify what configuration means here. Feel free to modify the suggestion.

[pggPL]

It was intentionally left unspecified, since this is going to rapidly change in the future I guess and we usually forget to update the docs. I added link and specified that this is for some release and may change in the future.

[phu0ngng]

I don't think we should expose these API now.
I would rather deprecate them and expose the permutation ops alone instead, in case there are users interested in the permutation alone.

[pggPL]

I'm not sure what you mean by "permutation ops". Can you please elaborate on that?

[phu0ngng]

I mean permute alone, but not token_dispatch/token_combine.
It's confusing for users seeing ep.ep_dispatch vs permutation.token_dispatch