[JAX] Make MoEBlock aware of Dense TP axes - #3462

Draft

jberchtold-nvidia wants to merge 5 commits into

NVIDIA:mainfrom

jberchtold-nvidia:jberchtold/fix-tp-etp1-awareness-in-moeblock

Commits on Sep 2, 2026

[JAX] Support compound EP with ETP1 in MoEBlock
jberchtold-nvidia
committed
0900878
[JAX] Keep TE EP MoE tests gated to Blackwell
jberchtold-nvidia
committed
4dfb6b7
[JAX] Configure compound EP directly in MoEBlock
jberchtold-nvidia
committed
be04a38
Revert compound EP MeshResource changes
jberchtold-nvidia
committed
990b223
[JAX] Support compound EP axes in MoEBlock
jberchtold-nvidia
committed
7514041