Skip to content

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

Open
tdophung wants to merge 19 commits into
NVIDIA:mainfrom
tdophung:teddy/moe_block
Open

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

Commits

Commits on May 5, 2026

Commits on May 6, 2026

Commits on May 7, 2026

Commits on May 12, 2026

Commits on May 21, 2026

Commits on May 22, 2026