BEAM: Binaire Expert Activeringsmaskering voor Dynamische Routering in MoE

Samenvatting

Mengsel-van-Experts (MoE)-architecturen verbeteren de efficiëntie van grote taalmodellen door per token slechts een subset van experts te activeren. Standaard MoE gebruikt echter een vaste Top-K-routeringsstrategie, wat leidt tot overbodige berekeningen en suboptimale inferentielatentie. Bestaande versnellingsmethoden vereisen ofwel kostbare hertraining met architectuurwijzigingen, ofwel lijden ze onder een ernstige prestatieverlies bij hoge schaarste vanwege een mismatch tussen training en inferentie. Om deze beperkingen aan te pakken, stellen we BEAM (Binary Expert Activation Masking) voor, een nieuwe methode die token-adaptieve experts selectie leert via trainbare binaire maskers. Met behulp van een straight-through-schatter en een hulp regularisatieverlies induceert BEAM dynamische expertschaarste door middel van end-to-end training, terwijl de modelcapaciteit behouden blijft. We implementeren verder een efficiënte aangepaste CUDA-kernel voor BEAM, wat zorgt voor naadloze integratie met het vLLM-inferentieraamwerk. Experimenten tonen aan dat BEAM meer dan 98% van de oorspronkelijke modelprestaties behoudt, terwijl het aantal MoE-laag-FLOPs tot 85% wordt verminderd, met een tot 2,5 keer snellere decodering en 1,4 keer hogere doorvoer, wat de effectiviteit als een praktische, plug-and-play-oplossing voor efficiënte MoE-inferentie aantoont.

English

Mixture-of-Experts (MoE) architectures enhance the efficiency of large language models by activating only a subset of experts per token. However, standard MoE employs a fixed Top-K routing strategy, leading to redundant computation and suboptimal inference latency. Existing acceleration methods either require costly retraining with architectural changes or suffer from severe performance drop at high sparsity due to train-inference mismatch. To address these limitations, we propose BEAM (Binary Expert Activation Masking), a novel method that learns token-adaptive expert selection via trainable binary masks. With a straight-through estimator and an auxiliary regularization loss, BEAM induces dynamic expert sparsity through end-to-end training while maintaining model capability. We further implement an efficient custom CUDA kernel for BEAM, ensuring seamless integration with the vLLM inference framework. Experiments show that BEAM retains over 98\% of the original model's performance while reducing MoE layer FLOPs by up to 85\%, achieving up to 2.5times faster decoding and 1.4times higher throughput, demonstrating its effectiveness as a practical, plug-and-play solution for efficient MoE inference.