Make MoE dispatch/MLP expert-axis batch sharding configurable (fix Mixtral EP throughput)#4179

Open

gulsumgudukbay wants to merge 5 commits into

AI-Hypercomputer:mainfrom

ROCm:fix-moe-expert-parallel-sharding

Commits on Jun 18, 2026

Make MoE dispatch/MLP expert-axis batch sharding configurable
gulsumgudukbay
committed
Enable moe_dispatch_no_expert_sharding for mixtral-8x22b. Same 8-large-expert geometry as 8x7b, so it benefits from the same expert-parallel MoE dispatch/MLP sharding.
gulsumgudukbay
committed
Add MoE dispatch expert-sharding regression test. Asserts that with moe_dispatch_no_expert_sharding the expert dim is sharded by 'expert' and the batch dim is not, guarding the expert-parallel disp…
gulsumgudukbay
committed
Address review: peel 'expert' from MoE dispatch/MLP batch dim instead of a new logical rule
gulsumgudukbay
committed
Test both flag states for MoE dispatch sharding; drop deepseek3-671b
gulsumgudukbay
committed