SlimWise: disaccoppiare il pruning degli esperti tra prefill e decode per un serving efficiente dei modelli MoE

Abstract degli autori

I modelli mixture-of-experts (MoE) attivano pochi esperti per token, ma il decode in batch può accedere a quasi tutti gli esperti disponibili, rendendo il trasferimento dei loro pesi un importante collo di bottiglia. Il pruning degli esperti riduce questo traffico, ma gli approcci convenzionali applicano il pruning anche al prefill, che è limitato dalla capacità di calcolo: sacrificano così la qualità del modello a fronte di un beneficio minimo per il throughput. Presentiamo SlimWise, un framework di serving che adatta l’insieme degli esperti a ciascuna fase dell’inferenza. SlimWise esegue il prefill con il modello completo e il decode con un modello sottoposto a pruning, che riutilizza direttamente la cache KV generata durante il prefill senza convertirla. Su due modelli MoE di base e con tre criteri di pruning, questo passaggio della cache KV, che non richiede addestramento, riduce sensibilmente il divario di accuratezza rispetto al modello completo in molte configurazioni. Mostriamo inoltre che l’accuratezza misurata nei benchmark può nascondere variazioni sostanziali della lunghezza del testo generato dovute al pruning. Per affrontare queste distorsioni e la perdita residua di accuratezza, SlimWise introduce una fase di distillazione a basso costo che addestra il decoder a proseguire a partire dalle cache KV del modello completo, aggiornando solo un piccolo sottoinsieme dei parametri. Implementato in vLLM, SlimWise supporta sia il serving con prefill e decode disaggregati (PD) sia quello con prefill e decode co-localizzati. Su Qwen3.6-35B-A3B, SlimWise aumenta il throughput del decode fino a 1,81 volte con il pruning del 50% degli esperti e una perdita minima di accuratezza.