HybriMoE: Planificación Híbrida CPU-GPU y Gestión de Caché para una Inferencia Eficiente de MoE

HybriMoE: Hybrid CPU-GPU Scheduling and Cache Management for Efficient MoE Inference

April 8, 2025

Autores: Shuzhang Zhong, Yanfan Sun, Ling Liang, Runsheng Wang, Ru Huang, Meng Li

cs.AI

Resumen

La arquitectura Mixture of Experts (MoE) ha demostrado ventajas significativas, ya que permite aumentar la capacidad del modelo sin un incremento proporcional en el cómputo. Sin embargo, el gran tamaño de los modelos MoE aún introduce demandas sustanciales de memoria, lo que generalmente requiere la descarga de expertos en plataformas con recursos limitados y conlleva un sobrecosto considerable. Se ha propuesto la inferencia híbrida CPU-GPU para aprovechar el cómputo de la CPU y reducir el sobrecosto de carga de expertos, pero enfrenta grandes desafíos: por un lado, los patrones de activación de expertos en los modelos MoE son altamente inestables, lo que hace ineficientes las estrategias de mapeo fijo en trabajos existentes; por otro lado, la programación híbrida CPU-GPU para MoE es inherentemente compleja debido a los diversos tamaños de expertos, estructuras, distribución desigual de cargas de trabajo, etc. Para abordar estos desafíos, en este artículo proponemos HybriMoE, un marco de inferencia híbrida CPU-GPU que mejora la utilización de recursos mediante un novedoso sistema de programación y gestión de caché CPU-GPU. HybriMoE introduce (i) una estrategia de programación dinámica intra-capa para equilibrar las cargas de trabajo entre la CPU y la GPU, (ii) un algoritmo de prebúsqueda inter-capa basado en impacto, y (iii) un algoritmo de almacenamiento en caché basado en puntuaciones para mitigar la inestabilidad en la activación de expertos. Implementamos HybriMoE sobre el marco kTransformers y lo evaluamos en tres modelos de lenguaje basados en MoE ampliamente utilizados. Los resultados experimentales demuestran que HybriMoE logra una aceleración promedio de 1.33 veces en la etapa de prefilling y 1.70 veces en la etapa de decodificación en comparación con el marco de inferencia híbrida MoE más avanzado. Nuestro código está disponible en: https://github.com/PKU-SEC-Lab/HybriMoE.

English

The Mixture of Experts (MoE) architecture has demonstrated significant advantages as it enables to increase the model capacity without a proportional increase in computation. However, the large MoE model size still introduces substantial memory demands, which usually requires expert offloading on resource-constrained platforms and incurs significant overhead. Hybrid CPU-GPU inference has been proposed to leverage CPU computation to reduce expert loading overhead but faces major challenges: on one hand, the expert activation patterns of MoE models are highly unstable, rendering the fixed mapping strategies in existing works inefficient; on the other hand, the hybrid CPU-GPU schedule for MoE is inherently complex due to the diverse expert sizes, structures, uneven workload distribution, etc. To address these challenges, in this paper, we propose HybriMoE, a hybrid CPU-GPU inference framework that improves resource utilization through a novel CPU-GPU scheduling and cache management system. HybriMoE introduces (i) a dynamic intra-layer scheduling strategy to balance workloads across CPU and GPU, (ii) an impact-driven inter-layer prefetching algorithm, and (iii) a score-based caching algorithm to mitigate expert activation instability. We implement HybriMoE on top of the kTransformers framework and evaluate it on three widely used MoE-based LLMs. Experimental results demonstrate that HybriMoE achieves an average speedup of 1.33times in the prefill stage and 1.70times in the decode stage compared to state-of-the-art hybrid MoE inference framework. Our code is available at: https://github.com/PKU-SEC-Lab/HybriMoE.