BlockFFN: Op weg naar Mixture-of-Experts die vriendelijk is voor versnelling aan de eindzijde met Chunk-niveau Activatie Sparsity

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

July 11, 2025

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

cs.AI

Samenvatting

Om de computationele belasting van grote taalmodelen (LLM's) te verlichten, hebben architecturen met activatiesparsheid, vertegenwoordigd door mixture-of-experts (MoE), steeds meer aandacht getrokken. Echter, de niet-differentieerbare en onflexibele routering van standaard MoE schaadt de modelprestaties. Bovendien, hoewel elk token slechts een paar parameters activeert, vertonen deze spaarzaam geactiveerde architecturen een lage chunk-level sparsheid, wat aangeeft dat de vereniging van meerdere opeenvolgende tokens een groot deel van de parameters activeert. Zo'n sparsheidspatroon is onvriendelijk voor versnelling onder laagresource-omstandigheden (bijv. end-side apparaten) en incompatibel met hoofdstroom versnellingsmethoden (bijv. speculatieve decodering). Om deze uitdagingen aan te pakken, introduceren we een nieuwe MoE-architectuur, BlockFFN, evenals efficiënte trainings- en implementatietechnieken. Specifiek gebruiken we een router die ReLU-activering en RMSNorm integreert voor differentieerbare en flexibele routering. Vervolgens, om zowel token-level sparsheid (TLS) als chunk-level sparsheid (CLS) te bevorderen, zijn CLS-bewuste trainingsdoelen ontworpen, waardoor BlockFFN meer versnellingsvriendelijk wordt. Ten slotte implementeren we efficiënte versnellingskernels, waarbij voor het eerst activatiesparsheid en speculatieve decodering worden gecombineerd. De experimentele resultaten tonen de superieure prestaties van BlockFFN aan ten opzichte van andere MoE-baselines, met meer dan 80% TLS en 70% 8-token CLS. Onze kernels behalen tot 3,67 keer versnelling op echte end-side apparaten in vergelijking met dichte modellen. Alle codes en checkpoints zijn openbaar beschikbaar (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).