BlockFFN: Hacia una Mezcla de Expertos Amigable con la Aceleración en el Lado del Extremo mediante Esparcidad de Activación a Nivel de Fragmentos

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

July 11, 2025

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

cs.AI

Resumen

Para aliviar la carga computacional de los modelos de lenguaje de gran escala (LLMs, por sus siglas en inglés), las arquitecturas con esparcidad de activación, representadas por la mezcla de expertos (MoE, por sus siglas en inglés), han atraído una atención creciente. Sin embargo, el enrutamiento no diferenciable e inflexible del MoE básico perjudica el rendimiento del modelo. Además, aunque cada token activa solo unos pocos parámetros, estas arquitecturas de activación dispersa exhiben una baja esparcidad a nivel de fragmento, lo que indica que la unión de múltiples tokens consecutivos activa una gran proporción de parámetros. Este patrón de esparcidad no es favorable para la aceleración en condiciones de recursos limitados (por ejemplo, dispositivos del lado del usuario) y es incompatible con las técnicas de aceleración más comunes (por ejemplo, decodificación especulativa). Para abordar estos desafíos, presentamos una nueva arquitectura MoE, BlockFFN, junto con sus técnicas eficientes de entrenamiento y despliegue. Específicamente, utilizamos un enrutador que integra la activación ReLU y RMSNorm para un enrutamiento diferenciable y flexible. A continuación, para promover tanto la esparcidad a nivel de token (TLS, por sus siglas en inglés) como la esparcidad a nivel de fragmento (CLS, por sus siglas en inglés), se diseñan objetivos de entrenamiento conscientes de CLS, haciendo que BlockFFN sea más amigable para la aceleración. Finalmente, implementamos núcleos de aceleración eficientes, combinando por primera vez la esparcidad de activación y la decodificación especulativa. Los resultados experimentales demuestran el rendimiento superior de BlockFFN sobre otras líneas base de MoE, logrando más del 80% de TLS y 70% de CLS para 8 tokens. Nuestros núcleos alcanzan una aceleración de hasta 3.67 veces en dispositivos reales del lado del usuario en comparación con modelos densos. Todos los códigos y puntos de control están disponibles públicamente (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).