HybriMoE: Agendamento Híbrido CPU-GPU e Gerenciamento de Cache para Inferência 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

Resumo

A arquitetura Mixture of Experts (MoE) demonstrou vantagens significativas, pois permite aumentar a capacidade do modelo sem um aumento proporcional na computação. No entanto, o grande tamanho dos modelos MoE ainda introduz demandas substanciais de memória, o que geralmente requer o descarregamento de especialistas em plataformas com recursos limitados e incorre em sobrecarga significativa. A inferência híbrida CPU-GPU foi proposta para aproveitar a computação da CPU e reduzir a sobrecarga de carregamento de especialistas, mas enfrenta grandes desafios: por um lado, os padrões de ativação de especialistas dos modelos MoE são altamente instáveis, tornando as estratégias de mapeamento fixo em trabalhos existentes ineficientes; por outro lado, o agendamento híbrido CPU-GPU para MoE é inerentemente complexo devido aos diversos tamanhos de especialistas, estruturas, distribuição desigual de carga de trabalho, etc. Para enfrentar esses desafios, neste artigo, propomos o HybriMoE, um framework de inferência híbrida CPU-GPU que melhora a utilização de recursos por meio de um novo sistema de agendamento e gerenciamento de cache CPU-GPU. O HybriMoE introduz (i) uma estratégia de agendamento dinâmico intra-camada para equilibrar as cargas de trabalho entre CPU e GPU, (ii) um algoritmo de pré-busca inter-camada orientado por impacto, e (iii) um algoritmo de cache baseado em pontuação para mitigar a instabilidade na ativação de especialistas. Implementamos o HybriMoE sobre o framework kTransformers e o avaliamos em três LLMs baseados em MoE amplamente utilizados. Os resultados experimentais demonstram que o HybriMoE alcança um ganho de velocidade médio de 1,33 vezes na etapa de preenchimento e 1,70 vezes na etapa de decodificação em comparação com o framework de inferência híbrida MoE estado da arte. Nosso código está disponível em: 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.