BlockFFN: В направлении дружественной к ускорению на стороне конечного устройства смеси экспертов с разреженностью активаций на уровне блоков

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

July 11, 2025

Авторы: Chenyang Song, Weilin Zhao, Xu Han, Chaojun Xiao, Yingfa Chen, Yuxuan Li, Zhiyuan Liu, Maosong Sun

cs.AI

Аннотация

Для снижения вычислительной нагрузки больших языковых моделей (LLM) архитектуры с разреженностью активаций, представленные смесью экспертов (MoE), привлекают все больше внимания. Однако недифференцируемая и негибкая маршрутизация в классической MoE ухудшает производительность модели. Кроме того, хотя каждый токен активирует лишь небольшое количество параметров, такие разреженно-активируемые архитектуры демонстрируют низкую разреженность на уровне блоков, что указывает на то, что объединение нескольких последовательных токенов активирует значительную долю параметров. Такой паттерн разреженности не подходит для ускорения в условиях ограниченных ресурсов (например, на устройствах конечного пользователя) и несовместим с основными методами ускорения (например, спекулятивным декодированием). Для решения этих проблем мы представляем новую архитектуру MoE, BlockFFN, а также эффективные методы ее обучения и развертывания. В частности, мы используем маршрутизатор, интегрирующий активацию ReLU и RMSNorm, для дифференцируемой и гибкой маршрутизации. Далее, для повышения как разреженности на уровне токенов (TLS), так и разреженности на уровне блоков (CLS), разработаны цели обучения, учитывающие CLS, что делает BlockFFN более подходящим для ускорения. Наконец, мы реализуем эффективные ядра ускорения, впервые объединяя разреженность активаций и спекулятивное декодирование. Результаты экспериментов демонстрируют превосходство BlockFFN над другими базовыми MoE, достигая более 80% TLS и 70% CLS для 8 токенов. Наши ядра обеспечивают ускорение до 3,67 раз на реальных устройствах конечного пользователя по сравнению с плотными моделями. Все коды и контрольные точки доступны публично (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).