Add fused_moe_mlp: fuse wi_0 and wi_1 into one grouped GEMM for MoE FFN1#3736

abhinavgoel95 wants to merge 1 commit intoAI-Hypercomputer:mainAI-Hypercomputer/maxtext:mainfrom

abhinavgoel95:abgoel/fused-moe-mlpabhinavgoel95/maxtext:abgoel/fused-moe-mlpCopy head branch name to clipboard