The mixture-of-experts (MoE) architecture has proven indispensable for scaling transformer-based artificial intelligence (AI) models. MoEs use a router network to activate specialized subnetworks, called experts, within each transformer layer on a per-token basis. Thus, MoEs increase model capacity without proportionally increasing the computational requirements of training and inference. However, a key challenge is that MoE inference throughput scales poorly with batching, as each input in a batch activates a different subset of experts, thereby increasing memory intensity.

To improve batched MoE inference throughput, we propose TokenMix, a novel technique that mixes token representations within the router across inputs in a batch to promote expert reuse. In effect, TokenMix encourages coordinated expert selection across inputs. TokenMix requires neither training nor fine-tuning and can be applied directly to off-the-shelf MoE models. We evaluate TokenMix across five state-of-the-art MoE models with parameter counts ranging from 16B to 235B. TokenMix achieves a 25–60% improvement in expert reuse, leading to a 12–28% reduction in inter-token latency with negligible or no loss in accuracy.