BEAM: Binäre Expertenaktivierungsmaskierung für dynamisches Routing in MoE

Zusammenfassung

Mixture-of-Experts (MoE)-Architekturen verbessern die Effizienz großer Sprachmodelle, indem sie pro Token nur eine Teilmenge der Experten aktivieren. Das standardmäßige MoE verwendet jedoch eine feste Top-K-Routingstrategie, was zu redundanter Berechnung und suboptimaler Inferenzlatenz führt. Bestehende Beschleunigungsmethoden erfordern entweder kostspieliges Neutraining mit Architekturänderungen oder leiden aufgrund von Trainings-Inferenz-Diskrepanz bei hoher Sparsity unter einem starken Leistungsabfall. Um diese Einschränkungen zu überwinden, schlagen wir BEAM (Binary Expert Activation Masking) vor, eine neuartige Methode, die token-adaptive Expertenauswahl durch trainierbare binäre Masken lernt. Mit einem Straight-Through-Schätzer und einem zusätzlichen Regularisierungsverlust induziert BEAM dynamische Expertensparsity durch End-to-End-Training, während die Modellfähigkeit erhalten bleibt. Zudem implementieren wir einen effizienten benutzerdefinierten CUDA-Kernel für BEAM, der eine nahtlose Integration mit dem vLLM-Inferenzframework gewährleistet. Experimente zeigen, dass BEAM über 98% der ursprünglichen Modellleistung beibehält, während die FLOPs der MoE-Schicht um bis zu 85% reduziert werden, was eine bis zu 2,5-mal schnellere Dekodierung und einen 1,4-mal höheren Durchsatz ermöglicht. Dies belegt seine Wirksamkeit als praktische Plug-and-Play-Lösung für effiziente MoE-Inferenz.

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.