Search papers, labs, and topics across Lattice.
This paper introduces CAP (Communication-Aware Assignment and Pruning), a framework designed to optimize the placement and pruning of experts in Mixture-of-Experts (MoE) models by balancing computation, communication, and accuracy. The framework employs co-activation driven expert placement to minimize communication overhead, adjusts the communication-computation trade-off for varying loads, and implements communication-aware expert pruning to selectively reduce routing destinations with minimal accuracy loss. Experimental results demonstrate that CAP achieves throughput improvements of 1.23x to 1.86x over existing methods while maintaining superior model accuracy under accelerated conditions.
CAP achieves up to 86% higher throughput in MoE models while preserving accuracy, revolutionizing how we optimize expert placement and pruning.
As MoE models scale to hundreds of experts, placement and pruning decisions increasingly dictate communication volume, affecting the performance of distributed inference across GPUs and nodes. We propose CAP (Communication-Aware Assignment and Pruning), a framework that considers computation, communication and accuracy together for efficient MoE inference through expert placement and pruning. It consists of three components: (1) Co-activation driven expert placement, which groups frequently co-activated experts to reduce inter-device and inter-node communication; (2) Communicationcomputation trade-off adjustment, which generates placements with different computational load and communication volume; and (3) Communication-aware expert pruning, which selectively removes routing destinations to reduce communication with limited accuracy degradation. By combining these components, CAP selects an efficient operating strategy for different hardware configurations. Across our single-node and multi-node experiments, it achieves 1.23x - 1.86 x throughput improvement over DeepSeek EPLB and sequential placement in vLLM, and preserves better model accuracy at the same target speedup under lossy acceleration.