BlockFFN: Auf Endgerätebeschleunigung optimierte Mixture-of-Experts mit Chunk-Level-Aktivierungs-Sparsamkeit

BlockFFN: Towards End-Side Acceleration-Friendly Mixture-of-Experts with Chunk-Level Activation Sparsity

July 11, 2025

papers.authors: Chenyang Song, Weilin Zhao, Xu Han, Chaojun Xiao, Yingfa Chen, Yuxuan Li, Zhiyuan Liu, Maosong Sun

cs.AI

papers.abstract

Um die rechenintensive Belastung großer Sprachmodelle (LLMs) zu verringern, haben Architekturen mit Aktivierungssparsität, repräsentiert durch Mixture-of-Experts (MoE), zunehmend Aufmerksamkeit erregt. Allerdings beeinträchtigt das nicht-differenzierbare und unflexible Routing von herkömmlichen MoE-Modellen die Modellleistung. Darüber hinaus aktivieren, obwohl jedes Token nur wenige Parameter aktiviert, diese spärlich aktivierten Architekturen eine geringe Chunk-Level-Sparsity, was darauf hinweist, dass die Vereinigung mehrerer aufeinanderfolgender Token einen großen Anteil der Parameter aktiviert. Ein solches Sparsity-Muster ist für die Beschleunigung unter ressourcenarmen Bedingungen (z. B. Endgeräte) ungeeignet und inkompatibel mit gängigen Beschleunigungstechniken (z. B. spekulatives Decodieren). Um diese Herausforderungen zu bewältigen, führen wir eine neuartige MoE-Architektur, BlockFFN, sowie effiziente Trainings- und Bereitstellungstechniken ein. Insbesondere verwenden wir einen Router, der ReLU-Aktivierung und RMSNorm integriert, um differenzierbares und flexibles Routing zu ermöglichen. Um sowohl Token-Level-Sparsity (TLS) als auch Chunk-Level-Sparsity (CLS) zu fördern, werden CLS-bewusste Trainingsziele entworfen, wodurch BlockFFN beschleunigungsfreundlicher wird. Schließlich implementieren wir effiziente Beschleunigungskerne, die erstmals Aktivierungssparsität und spekulatives Decodieren kombinieren. Die experimentellen Ergebnisse demonstrieren die überlegene Leistung von BlockFFN gegenüber anderen MoE-Baselines, mit über 80 % TLS und 70 % 8-Token-CLS. Unsere Kerne erreichen auf realen Endgeräten eine bis zu 3,67-fache Beschleunigung im Vergleich zu dichten Modellen. Alle Codes und Checkpoints sind öffentlich verfügbar (https://github.com/thunlp/BlockFFN).

English

To alleviate the computational burden of large language models (LLMs), architectures with activation sparsity, represented by mixture-of-experts (MoE), have attracted increasing attention. However, the non-differentiable and inflexible routing of vanilla MoE hurts model performance. Moreover, while each token activates only a few parameters, these sparsely-activated architectures exhibit low chunk-level sparsity, indicating that the union of multiple consecutive tokens activates a large ratio of parameters. Such a sparsity pattern is unfriendly for acceleration under low-resource conditions (e.g., end-side devices) and incompatible with mainstream acceleration techniques (e.g., speculative decoding). To address these challenges, we introduce a novel MoE architecture, BlockFFN, as well as its efficient training and deployment techniques. Specifically, we use a router integrating ReLU activation and RMSNorm for differentiable and flexible routing. Next, to promote both token-level sparsity (TLS) and chunk-level sparsity (CLS), CLS-aware training objectives are designed, making BlockFFN more acceleration-friendly. Finally, we implement efficient acceleration kernels, combining activation sparsity and speculative decoding for the first time. The experimental results demonstrate the superior performance of BlockFFN over other MoE baselines, achieving over 80% TLS and 70% 8-token CLS. Our kernels achieve up to 3.67times speedup on real end-side devices than dense models. All codes and checkpoints are available publicly (https://github.com/thunlp/BlockFFN).