Skip to content

[JAX] Add an MoE Block (Layer) that compound router, permutation, groupedGEMM and communication#2912

Open
tdophung wants to merge 8 commits intoNVIDIA:mainfrom
tdophung:teddy/moe_block
Open

[JAX] Add an MoE Block (Layer) that compound router, permutation, groupedGEMM and communication#2912
tdophung wants to merge 8 commits intoNVIDIA:mainfrom
tdophung:teddy/moe_block

Commits

Commits on May 5, 2026

Commits on May 6, 2026